From de9ddbe08ae4231b26deda1f04eca14d1a34bacd Mon Sep 17 00:00:00 2001 From: Greg Clark Date: Thu, 16 Jul 2026 19:39:06 -0400 Subject: [PATCH 01/14] feat(switchyard-translation): encode / decode helpers Signed-off-by: Greg Clark --- Cargo.lock | 379 +++++++------------ crates/switchyard-translation/Cargo.toml | 2 + crates/switchyard-translation/src/helpers.rs | 274 ++++++++++++++ crates/switchyard-translation/src/lib.rs | 4 + crates/switchyard-translation/src/sse.rs | 141 +++++++ 5 files changed, 567 insertions(+), 233 deletions(-) create mode 100644 crates/switchyard-translation/src/helpers.rs create mode 100644 crates/switchyard-translation/src/sse.rs diff --git a/Cargo.lock b/Cargo.lock index 18749f29..642b93ae 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -217,27 +217,27 @@ checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" [[package]] name = "bitflags" -version = "2.11.1" +version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" [[package]] name = "bumpalo" -version = "3.20.2" +version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" [[package]] name = "bytes" -version = "1.11.1" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "cc" -version = "1.2.62" +version = "1.2.67" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1dce859f0832a7d088c4f1119888ab94ef4b5d6795d1ce05afb7fe159d79f98" +checksum = "e17dd265a7d0f31ef544e1b20e03add05d3b45b491b633b10d67145d2acc1a38" dependencies = [ "find-msvc-tools", "jobserver", @@ -253,15 +253,26 @@ checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" [[package]] name = "cfg_aliases" -version = "0.2.1" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" + +[[package]] +name = "chacha20" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core 0.10.1", +] [[package]] name = "clap" -version = "4.6.1" +version = "4.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +checksum = "dd059f9da4f5c36b3787f65d38ccaab1cc315f07b01f89abc8359ee6a8205011" dependencies = [ "clap_builder", "clap_derive", @@ -269,9 +280,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.6.0" +version = "4.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +checksum = "f09628afdcc538b57f3c6341e9c8e9970f18e4a481690a64974d7023bd33548b" dependencies = [ "anstream", "anstyle", @@ -328,6 +339,15 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "diff" version = "0.1.13" @@ -336,9 +356,9 @@ checksum = "56254986775e3233ffa9c4d7d3faaf6d36a2c09d30b20687e9f88bc8bafc16c8" [[package]] name = "displaydoc" -version = "0.2.5" +version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" dependencies = [ "proc-macro2", "quote", @@ -517,20 +537,6 @@ dependencies = [ "wasm-bindgen", ] -[[package]] -name = "getrandom" -version = "0.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" -dependencies = [ - "cfg-if", - "js-sys", - "libc", - "r-efi 5.3.0", - "wasip2", - "wasm-bindgen", -] - [[package]] name = "getrandom" version = "0.4.3" @@ -538,8 +544,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", + "js-sys", "libc", - "r-efi 6.0.0", + "r-efi", + "rand_core 0.10.1", + "wasm-bindgen", ] [[package]] @@ -586,9 +595,9 @@ checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" [[package]] name = "http" -version = "1.4.0" +version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" dependencies = [ "bytes", "itoa", @@ -596,9 +605,9 @@ dependencies = [ [[package]] name = "http-body" -version = "1.0.1" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" dependencies = [ "bytes", "http", @@ -606,9 +615,9 @@ dependencies = [ [[package]] name = "http-body-util" -version = "0.1.3" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" dependencies = [ "bytes", "futures-core", @@ -631,9 +640,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "hyper" -version = "1.9.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6299f016b246a94207e63da54dbe807655bf9e00044f73ded42c3ac5305fbcca" +checksum = "55281c53a1894c864990125767da440a4e630446785086f52523b20033b74498" dependencies = [ "atomic-waker", "bytes", @@ -833,13 +842,12 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.98" +version = "0.3.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67df7112613f8bfd9150013a0314e196f4800d3201ae742489d999db2f979f08" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" dependencies = [ "cfg-if", "futures-util", - "once_cell", "wasm-bindgen", ] @@ -855,7 +863,7 @@ version = "0.1.0" dependencies = [ "async-trait", "futures", - "rand 0.8.6", + "rand 0.8.7", "serde_json", "switchyard-protocol", "tokio", @@ -897,9 +905,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.29" +version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" [[package]] name = "lru" @@ -924,9 +932,9 @@ checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" [[package]] name = "memchr" -version = "2.8.0" +version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "mime" @@ -936,9 +944,9 @@ checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" [[package]] name = "mio" -version = "1.2.0" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" dependencies = [ "libc", "wasi", @@ -1131,9 +1139,9 @@ dependencies = [ [[package]] name = "quinn" -version = "0.11.9" +version = "0.11.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" dependencies = [ "bytes", "cfg_aliases", @@ -1151,14 +1159,15 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.14" +version = "0.11.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" dependencies = [ "bytes", - "getrandom 0.3.4", + "getrandom 0.4.3", "lru-slab", - "rand 0.9.4", + "rand 0.10.2", + "rand_pcg", "ring", "rustc-hash", "rustls", @@ -1172,33 +1181,27 @@ dependencies = [ [[package]] name = "quinn-udp" -version = "0.5.14" +version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" dependencies = [ "cfg_aliases", "libc", "once_cell", "socket2", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] name = "quote" -version = "1.0.45" +version = "1.0.46" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" dependencies = [ "proc-macro2", ] -[[package]] -name = "r-efi" -version = "5.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" - [[package]] name = "r-efi" version = "6.0.0" @@ -1207,23 +1210,24 @@ checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] name = "rand" -version = "0.8.6" +version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" +checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", - "rand_chacha 0.3.1", + "rand_chacha", "rand_core 0.6.4", ] [[package]] name = "rand" -version = "0.9.4" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" dependencies = [ - "rand_chacha 0.9.0", - "rand_core 0.9.5", + "chacha20", + "getrandom 0.4.3", + "rand_core 0.10.1", ] [[package]] @@ -1236,16 +1240,6 @@ dependencies = [ "rand_core 0.6.4", ] -[[package]] -name = "rand_chacha" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" -dependencies = [ - "ppv-lite86", - "rand_core 0.9.5", -] - [[package]] name = "rand_core" version = "0.6.4" @@ -1257,11 +1251,17 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.9.5" +version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" dependencies = [ - "getrandom 0.3.4", + "rand_core 0.10.1", ] [[package]] @@ -1330,15 +1330,15 @@ dependencies = [ [[package]] name = "rustc-hash" -version = "2.1.2" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" [[package]] name = "rustls" -version = "0.23.40" +version = "0.23.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" +checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" dependencies = [ "aws-lc-rs", "log", @@ -1352,9 +1352,9 @@ dependencies = [ [[package]] name = "rustls-native-certs" -version = "0.8.3" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" dependencies = [ "openssl-probe", "rustls-pki-types", @@ -1364,9 +1364,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.1" +version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" dependencies = [ "web-time", "zeroize", @@ -1386,9 +1386,9 @@ dependencies = [ [[package]] name = "rustversion" -version = "1.0.22" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" [[package]] name = "ryu" @@ -1466,9 +1466,9 @@ dependencies = [ [[package]] name = "serde_json" -version = "1.0.149" +version = "1.0.150" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" dependencies = [ "itoa", "memchr", @@ -1511,9 +1511,9 @@ dependencies = [ [[package]] name = "shlex" -version = "1.3.0" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" [[package]] name = "signal-hook-registry" @@ -1533,15 +1533,15 @@ checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" [[package]] name = "smallvec" -version = "1.15.1" +version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" [[package]] name = "socket2" -version = "0.6.3" +version = "0.6.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", "windows-sys 0.61.2", @@ -1573,7 +1573,7 @@ dependencies = [ "async-trait", "futures-util", "parking_lot", - "rand 0.8.6", + "rand 0.8.7", "reqwest", "serde", "serde_json", @@ -1592,7 +1592,7 @@ dependencies = [ "futures-core", "futures-util", "parking_lot", - "rand 0.8.6", + "rand 0.8.7", "reqwest", "serde", "serde_json", @@ -1621,7 +1621,7 @@ dependencies = [ "async-trait", "futures-core", "lru", - "rand 0.8.6", + "rand 0.8.7", "serde", "serde_json", "thiserror", @@ -1696,6 +1696,8 @@ dependencies = [ name = "switchyard-translation" version = "0.1.0" dependencies = [ + "async-stream", + "futures", "pretty_assertions", "serde", "serde_json", @@ -1705,9 +1707,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.117" +version = "2.0.119" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" dependencies = [ "proc-macro2", "quote", @@ -1772,9 +1774,9 @@ dependencies = [ [[package]] name = "tinyvec" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" dependencies = [ "tinyvec_macros", ] @@ -1787,9 +1789,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.52.3" +version = "1.52.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "317fafbbe3f02fc663dad00ea6186197de963cd4190e86a26d8d0fae095539af" dependencies = [ "bytes", "libc", @@ -2011,20 +2013,11 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" -[[package]] -name = "wasip2" -version = "1.0.3+wasi-0.2.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" -dependencies = [ - "wit-bindgen", -] - [[package]] name = "wasm-bindgen" -version = "0.2.121" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49ace1d07c165b0864824eee619580c4689389afa9dc9ed3a4c75040d82e6790" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" dependencies = [ "cfg-if", "once_cell", @@ -2035,9 +2028,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.71" +version = "0.4.76" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96492d0d3ffba25305a7dc88720d250b1401d7edca02cc3bcd50633b424673b8" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" dependencies = [ "js-sys", "wasm-bindgen", @@ -2045,9 +2038,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.121" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e68e6f4afd367a562002c05637acb8578ff2dea1943df76afb9e83d177c8578" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -2055,9 +2048,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.121" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d95a9ec35c64b2a7cb35d3fead40c4238d0940c86d107136999567a4703259f2" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" dependencies = [ "bumpalo", "proc-macro2", @@ -2068,9 +2061,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.121" +version = "0.2.126" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4e0100b01e9f0d03189a92b96772a1fb998639d981193d7dbab487302513441" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" dependencies = [ "unicode-ident", ] @@ -2090,9 +2083,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.98" +version = "0.3.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b572dff8bcf38bad0fa19729c89bb5748b2b9b1d8be70cf90df697e3a8f32aa" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" dependencies = [ "js-sys", "wasm-bindgen", @@ -2120,16 +2113,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", -] - -[[package]] -name = "windows-sys" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" -dependencies = [ - "windows-targets 0.53.5", + "windows-targets", ] [[package]] @@ -2147,31 +2131,14 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", - "windows_i686_gnullvm 0.52.6", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", -] - -[[package]] -name = "windows-targets" -version = "0.53.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" -dependencies = [ - "windows-link", - "windows_aarch64_gnullvm 0.53.1", - "windows_aarch64_msvc 0.53.1", - "windows_i686_gnu 0.53.1", - "windows_i686_gnullvm 0.53.1", - "windows_i686_msvc 0.53.1", - "windows_x86_64_gnu 0.53.1", - "windows_x86_64_gnullvm 0.53.1", - "windows_x86_64_msvc 0.53.1", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] [[package]] @@ -2180,96 +2147,48 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_aarch64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" - [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" -[[package]] -name = "windows_i686_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" - [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_i686_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "windows_x86_64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" - [[package]] name = "winnow" version = "0.7.15" @@ -2279,12 +2198,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "wit-bindgen" -version = "0.57.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" - [[package]] name = "writeable" version = "0.6.3" @@ -2312,9 +2225,9 @@ checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" [[package]] name = "yoke" -version = "0.8.2" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "abe8c5fda708d9ca3df187cae8bfb9ceda00dd96231bed36e445a1a48e66f9ca" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" dependencies = [ "stable_deref_trait", "yoke-derive", @@ -2335,18 +2248,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.48" +version = "0.8.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eed437bf9d6692032087e337407a86f04cd8d6a16a37199ed57949d415bd68e9" +checksum = "b7cbbc0a705a0fd05cc3676525980d2bf5a9bc4adac6d6475209a7887cf59d19" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.48" +version = "0.8.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4" +checksum = "e2e817b7b52d0c7358d3246da9d69935ebb18116b2b102b4230dac079b4862f5" dependencies = [ "proc-macro2", "quote", @@ -2376,9 +2289,9 @@ dependencies = [ [[package]] name = "zeroize" -version = "1.8.2" +version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" [[package]] name = "zerotrie" @@ -2415,6 +2328,6 @@ dependencies = [ [[package]] name = "zmij" -version = "1.0.21" +version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/crates/switchyard-translation/Cargo.toml b/crates/switchyard-translation/Cargo.toml index 6f762904..9e55e7e9 100644 --- a/crates/switchyard-translation/Cargo.toml +++ b/crates/switchyard-translation/Cargo.toml @@ -21,6 +21,8 @@ serde = { version = "1", features = ["derive"] } serde_json = "1" switchyard-protocol = { path = "../libsy-protocol", version = "0.1.0" } thiserror = "2" +futures = "0.3" +async-stream = "0.3" [dev-dependencies] pretty_assertions = "1" diff --git a/crates/switchyard-translation/src/helpers.rs b/crates/switchyard-translation/src/helpers.rs new file mode 100644 index 00000000..a5e644e0 --- /dev/null +++ b/crates/switchyard-translation/src/helpers.rs @@ -0,0 +1,274 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Convenience wrappers over the default [`TranslationEngine`] — decode a wire +//! request/response to the neutral IR, encode the IR back, and decode/encode a +//! streamed response — so callers can translate without threading an engine and +//! policy through every call. + +use std::pin::Pin; +use std::sync::LazyLock; + +use async_stream::try_stream; +use futures::{Stream, StreamExt}; +use serde_json::Value; + +use crate::sse; +use crate::{ + AggLlmResponse, LlmRequest, LlmResponseStream, Result, StreamCodecRegistry, + StreamTranslationState, TranslationEngine, TranslationPolicy, WireFormat, +}; + +static DEFAULT_TRANSLATION_POLICY: LazyLock = + LazyLock::new(TranslationPolicy::default); +static DEFAULT_TRANSLATION_ENGINE: LazyLock = + LazyLock::new(TranslationEngine::default); + +/// Decodes a `wire_format` request body into the neutral IR. +pub fn decode_request(wire_format: WireFormat, body: &Value) -> Result { + Ok(DEFAULT_TRANSLATION_ENGINE + .decode_request(wire_format, body, &DEFAULT_TRANSLATION_POLICY)? + .request) +} + +/// Encodes a normalized request into `wire_format`'s JSON body. +pub fn encode_request(request: &LlmRequest, wire_format: WireFormat) -> Result { + Ok(DEFAULT_TRANSLATION_ENGINE + .encode_request(wire_format, request, &DEFAULT_TRANSLATION_POLICY)? + .body) +} + +/// Decodes a buffered `wire_format` response body into the neutral aggregate. +pub fn decode_buffered_response(body: &Value, wire_format: WireFormat) -> Result { + Ok(DEFAULT_TRANSLATION_ENGINE + .decode_response(wire_format, body, &DEFAULT_TRANSLATION_POLICY)? + .response) +} + +/// Encodes a buffered aggregate into `wire_format`'s JSON body, stamping +/// `requested_model` over the upstream id so the caller sees the model it asked for. +pub fn encode_buffered_response( + agg: &AggLlmResponse, + wire_format: WireFormat, + requested_model: Option<&str>, +) -> Result { + let mut body = DEFAULT_TRANSLATION_ENGINE + .encode_response(wire_format, agg, &DEFAULT_TRANSLATION_POLICY)? + .body; + if let (Some(model), Value::Object(object)) = (requested_model, &mut body) { + object.insert("model".to_string(), Value::String(model.to_string())); + } + Ok(body) +} + +/// A stream of wire-format event objects in one format — the unframed body of an +/// SSE response. The serving layer frames each `Value` (e.g. as an SSE +/// `data:`/`event:` block). +pub type RawEventStream = Pin< + Box< + dyn Stream>> + + Send, + >, +>; + +/// Encodes a stream of IR chunks into a stream of target-format wire events. +/// +/// `requested_model` is exposed as the response model (via the stream state's +/// `target_model`). The target stream codec is resolved once and reused per chunk; +/// terminal events (`message_stop` / `response.completed`) come from `finish`. +pub fn encode_stream( + chunks: LlmResponseStream, + target: WireFormat, + requested_model: Option, +) -> RawEventStream { + // The target is always a built-in wire format, so this lookup cannot fail; a + // failure surfaces as a single error item rather than a panic. + let codec = match StreamCodecRegistry::with_builtins().codec(target) { + Ok(codec) => codec, + Err(error) => { + let message = error.to_string(); + return Box::pin(futures::stream::once(async move { + Err(Box::new(std::io::Error::other(message)) + as Box) + })); + } + }; + + let events = try_stream! { + let mut state = StreamTranslationState { + target: Some(target.into()), + target_model: requested_model, + ..Default::default() + }; + + let mut chunks = chunks; + while let Some(item) = chunks.next().await { + let chunk = item?; + for value in codec.encode_event(&mut state, chunk) { + yield value; + } + } + for value in codec.finish(&mut state) { + yield value; + } + }; + + Box::pin(events) +} + +/// Decodes a byte stream of `source`-format SSE frames into neutral IR chunks. +/// +/// Operates on raw bytes, not any HTTP client type: the caller adapts its +/// transport's body stream into `Stream, _>>`. Frames are +/// buffered across chunks (a partial frame waits for its boundary); the source +/// stream codec is resolved once and reused for every frame. +pub fn decode_stream(bytes: S, source: WireFormat) -> LlmResponseStream +where + S: Stream, Box>> + + Send + + 'static, +{ + let marker = sse::done_marker(source); + // The source is always a built-in wire format, so this lookup cannot fail; a + // failure still surfaces as a single error item rather than a panic. + let codec = match StreamCodecRegistry::with_builtins().codec(source) { + Ok(codec) => codec, + Err(error) => { + let boxed: Box = error.to_string().into(); + return Box::pin(futures::stream::once(async move { Err(boxed) })); + } + }; + Box::pin(try_stream! { + let mut state = StreamTranslationState::default(); + let mut buffer = Vec::new(); + futures::pin_mut!(bytes); + + while let Some(chunk) = bytes.next().await { + buffer.extend_from_slice(&chunk?); + while let Some(frame) = sse::drain_next_sse_frame(&mut buffer)? { + match sse::parse_json_sse_frame(&frame, marker)? { + sse::ParsedSseFrame::Json(value) => { + for event in codec.decode_event(&mut state, &value) { + yield event; + } + } + sse::ParsedSseFrame::Done => return, + sse::ParsedSseFrame::Empty => {} + } + } + } + + // A non-standard upstream might omit the final blank line; parse a trailing + // complete frame instead of losing its last chunk. + if sse::has_non_whitespace_bytes(&buffer) { + let frame = sse::decode_sse_frame(&buffer)?; + if let sse::ParsedSseFrame::Json(value) = sse::parse_json_sse_frame(&frame, marker)? { + for event in codec.decode_event(&mut state, &value) { + yield event; + } + } + } + }) +} + +#[cfg(test)] +mod tests { + use futures::executor::block_on; + use futures::{stream, StreamExt}; + use serde_json::{json, Value}; + use switchyard_protocol::{completion_text, LlmResponseChunk}; + + use super::{ + decode_buffered_response, decode_request, decode_stream, encode_buffered_response, + encode_request, encode_stream, + }; + use crate::{LlmResponseStream, WireFormat}; + + // A boxed stream item error, matching the streamed IR contract. + type BoxError = Box; + + #[test] + fn request_round_trips_through_openai_chat() { + let body = json!({"model": "gpt", "messages": [{"role": "user", "content": "hi"}]}); + let request = decode_request(WireFormat::OpenAiChat, &body).unwrap(); + assert_eq!(request.model.as_deref(), Some("gpt")); + + let encoded = encode_request(&request, WireFormat::OpenAiChat).unwrap(); + assert_eq!(encoded["model"], "gpt"); + assert_eq!(encoded["messages"][0]["content"], "hi"); + } + + #[test] + fn buffered_response_round_trips_and_restamps_model() { + let body = json!({ + "id": "1", + "model": "upstream", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hi there"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} + }); + let agg = decode_buffered_response(&body, WireFormat::OpenAiChat).unwrap(); + assert_eq!(completion_text(&agg), "Hi there"); + + // `requested_model` overrides the encoded model. + let encoded = + encode_buffered_response(&agg, WireFormat::OpenAiChat, Some("client-facing")).unwrap(); + assert_eq!(encoded["model"], "client-facing"); + assert_eq!(encoded["choices"][0]["message"]["content"], "Hi there"); + } + + #[test] + fn encode_stream_emits_openai_chunks_whose_deltas_reassemble() { + let chunks: LlmResponseStream = stream::iter(vec![ + Ok(LlmResponseChunk::TextDelta { + index: 0, + text: "Hello".to_string(), + }), + Ok(LlmResponseChunk::TextDelta { + index: 0, + text: " world".to_string(), + }), + Ok(LlmResponseChunk::MessageStop { + reason: Some("stop".to_string()), + }), + ]) + .boxed(); + + let events: Vec = block_on( + encode_stream(chunks, WireFormat::OpenAiChat, Some("m".to_string())) + .map(|item| item.unwrap()) + .collect(), + ); + let content: String = events + .iter() + .filter_map(|event| event["choices"][0]["delta"]["content"].as_str()) + .collect(); + assert_eq!(content, "Hello world"); + } + + #[test] + fn decode_stream_parses_sse_bytes_into_ir_chunks() { + let sse = b"data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\ + data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\ + data: [DONE]\n\n" + .to_vec(); + let bytes = stream::once(async move { Ok::, BoxError>(sse) }); + + let chunks: Vec = block_on( + decode_stream(bytes, WireFormat::OpenAiChat) + .map(|item| item.unwrap()) + .collect(), + ); + let content: String = chunks + .iter() + .filter_map(|chunk| match chunk { + LlmResponseChunk::TextDelta { text, .. } => Some(text.as_str()), + _ => None, + }) + .collect(); + assert_eq!(content, "Hello world"); + } +} diff --git a/crates/switchyard-translation/src/lib.rs b/crates/switchyard-translation/src/lib.rs index 3eaaf21a..b54296d8 100644 --- a/crates/switchyard-translation/src/lib.rs +++ b/crates/switchyard-translation/src/lib.rs @@ -11,17 +11,21 @@ pub mod codecs; pub mod diagnostic; pub mod engine; pub mod error; +mod helpers; pub mod policy; +mod sse; pub mod stream; pub mod util; pub use switchyard_protocol::stream::LlmResponseChunk; +pub use switchyard_protocol::stream::LlmResponseStream; pub use switchyard_protocol::{format, llm}; pub use diagnostic::*; pub use engine::*; pub use error::*; pub use format::*; +pub use helpers::*; pub use llm::*; pub use policy::*; pub use stream::*; diff --git a/crates/switchyard-translation/src/sse.rs b/crates/switchyard-translation/src/sse.rs new file mode 100644 index 00000000..443f062e --- /dev/null +++ b/crates/switchyard-translation/src/sse.rs @@ -0,0 +1,141 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Minimal SSE frame parser for decoding streamed provider responses. +//! +//! One copy backs all neutral-IR stream decoding ([`decode_stream`](crate::decode_stream)). +//! Errors are boxed `std::error::Error`s — the item error type of a streamed +//! response — so this module stays free of any HTTP client or server types. + +use serde_json::Value; + +use crate::WireFormat; + +/// Boxed, thread-safe error carried by a streamed item. +pub(crate) type BoxError = Box; + +/// The outcome of parsing one SSE frame's data lines. +pub(crate) enum ParsedSseFrame { + /// Frame contained a JSON payload. + Json(Value), + /// Frame contained the provider's terminal marker (e.g. `[DONE]`). + Done, + /// Frame had no data payload. + Empty, +} + +/// The terminal SSE marker for `format`, if any. OpenAI Chat and Responses use +/// `[DONE]`; Anthropic ends on its `message_stop` event with no marker. +pub(crate) fn done_marker(format: WireFormat) -> Option<&'static str> { + match format { + WireFormat::OpenAiChat | WireFormat::OpenAiResponses => Some("[DONE]"), + WireFormat::AnthropicMessages => None, + } +} + +/// Drains one complete SSE frame from the buffer when a boundary is present. +pub(crate) fn drain_next_sse_frame(buffer: &mut Vec) -> Result, BoxError> { + let Some((index, separator_len)) = next_sse_boundary(buffer) else { + return Ok(None); + }; + let frame = decode_sse_frame(&buffer[..index])?; + buffer.drain(..index + separator_len); + Ok(Some(frame)) +} + +/// Decodes one raw SSE frame as UTF-8. +pub(crate) fn decode_sse_frame(frame: &[u8]) -> Result { + std::str::from_utf8(frame) + .map(str::to_string) + .map_err(|error| format!("stream emitted invalid UTF-8 frame: {error}").into()) +} + +/// Returns whether the buffer has any non-whitespace bytes. +pub(crate) fn has_non_whitespace_bytes(buffer: &[u8]) -> bool { + buffer.iter().any(|byte| !byte.is_ascii_whitespace()) +} + +/// Parses data lines from one SSE frame into JSON, terminal, or empty states. +pub(crate) fn parse_json_sse_frame( + frame: &str, + done_marker: Option<&str>, +) -> Result { + let mut data_lines = Vec::new(); + for line in frame.lines() { + // SSE comments and blank lines do not contribute data. + if line.is_empty() || line.starts_with(':') { + continue; + } + if let Some(data) = line.strip_prefix("data:") { + data_lines.push(data.trim_start().to_string()); + } + } + + if data_lines.is_empty() { + return Ok(ParsedSseFrame::Empty); + } + + let data = data_lines.join("\n"); + if done_marker.is_some_and(|marker| data.trim() == marker) { + return Ok(ParsedSseFrame::Done); + } + + let value = serde_json::from_str::(&data) + .map_err(|error| format!("stream emitted invalid JSON frame: {error}"))?; + Ok(ParsedSseFrame::Json(value)) +} + +/// Finds the next CRLF or LF SSE frame boundary. +fn next_sse_boundary(buffer: &[u8]) -> Option<(usize, usize)> { + match (find_bytes(buffer, b"\r\n\r\n"), find_bytes(buffer, b"\n\n")) { + (Some(crlf), Some(lf)) if crlf < lf => Some((crlf, 4)), + (Some(_), Some(lf)) => Some((lf, 2)), + (Some(crlf), None) => Some((crlf, 4)), + (None, Some(lf)) => Some((lf, 2)), + (None, None) => None, + } +} + +/// Finds a byte needle inside a byte haystack. +fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .position(|window| window == needle) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + // Multi-byte UTF-8 split across network chunks should wait for a full frame. + #[test] + fn buffers_incomplete_utf8_until_a_complete_sse_frame_arrives() -> Result<(), BoxError> { + let mut buffer = b"data: {\"text\":\"".to_vec(); + let multibyte = "é".as_bytes(); + buffer.extend_from_slice(&multibyte[..1]); + assert!(drain_next_sse_frame(&mut buffer)?.is_none()); + + buffer.extend_from_slice(&multibyte[1..]); + buffer.extend_from_slice(b"\"}\n\n"); + + let Some(frame) = drain_next_sse_frame(&mut buffer)? else { + return Err("complete SSE frame should be drained".into()); + }; + let ParsedSseFrame::Json(value) = parse_json_sse_frame(&frame, None)? else { + return Err("SSE frame should parse as JSON".into()); + }; + assert_eq!(value, json!({"text": "é"})); + assert!(buffer.is_empty()); + Ok(()) + } + + #[test] + fn recognizes_done_marker() -> Result<(), BoxError> { + let ParsedSseFrame::Done = parse_json_sse_frame("data: [DONE]\n", Some("[DONE]"))? else { + return Err("DONE frame should stop".into()); + }; + Ok(()) + } +} From 30d8979839d08b22d97a85bc6569da7aa52f27e6 Mon Sep 17 00:00:00 2001 From: Greg Clark Date: Thu, 16 Jul 2026 19:48:37 -0400 Subject: [PATCH 02/14] feat(switchyard-llm-client): add HTTP LLM client crate speaking the neutral IR Signed-off-by: Greg Clark --- Cargo.lock | 128 ++++ Cargo.toml | 1 + crates/libsy-llm-client/Cargo.toml | 27 + crates/libsy-llm-client/README.md | 174 ++++++ crates/libsy-llm-client/src/backend.rs | 246 ++++++++ crates/libsy-llm-client/src/client.rs | 812 +++++++++++++++++++++++++ crates/libsy-llm-client/src/error.rs | 157 +++++ crates/libsy-llm-client/src/lib.rs | 30 + crates/libsy-llm-client/src/raw.rs | 22 + 9 files changed, 1597 insertions(+) create mode 100644 crates/libsy-llm-client/Cargo.toml create mode 100644 crates/libsy-llm-client/README.md create mode 100644 crates/libsy-llm-client/src/backend.rs create mode 100644 crates/libsy-llm-client/src/client.rs create mode 100644 crates/libsy-llm-client/src/error.rs create mode 100644 crates/libsy-llm-client/src/lib.rs create mode 100644 crates/libsy-llm-client/src/raw.rs diff --git a/Cargo.lock b/Cargo.lock index 642b93ae..2dda8961 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,15 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -67,6 +76,16 @@ dependencies = [ "rustversion", ] +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "async-stream" version = "0.3.6" @@ -348,6 +367,24 @@ dependencies = [ "libc", ] +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" + [[package]] name = "diff" version = "0.1.13" @@ -593,6 +630,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "http" version = "1.4.2" @@ -851,6 +894,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + [[package]] name = "libc" version = "0.2.186" @@ -953,6 +1002,16 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -1273,6 +1332,35 @@ dependencies = [ "bitflags", ] +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + [[package]] name = "reqwest" version = "0.12.28" @@ -1629,6 +1717,23 @@ dependencies = [ "tracing", ] +[[package]] +name = "switchyard-llm-client" +version = "0.1.0" +dependencies = [ + "async-trait", + "futures", + "futures-util", + "pretty_assertions", + "reqwest", + "serde_json", + "switchyard-protocol", + "switchyard-translation", + "thiserror", + "tokio", + "wiremock", +] + [[package]] name = "switchyard-protocol" version = "0.1.0" @@ -2198,6 +2303,29 @@ dependencies = [ "memchr", ] +[[package]] +name = "wiremock" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031" +dependencies = [ + "assert-json-diff", + "base64", + "deadpool", + "futures", + "http", + "http-body-util", + "hyper", + "hyper-util", + "log", + "once_cell", + "regex", + "serde", + "serde_json", + "tokio", + "url", +] + [[package]] name = "writeable" version = "0.6.3" diff --git a/Cargo.toml b/Cargo.toml index 582ea8df..f39f72ed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ resolver = "2" members = [ "crates/libsy", "crates/libsy-examples", + "crates/libsy-llm-client", "crates/switchyard-components", "crates/switchyard-components-v2", "crates/switchyard-components-v2-macros", diff --git a/crates/libsy-llm-client/Cargo.toml b/crates/libsy-llm-client/Cargo.toml new file mode 100644 index 00000000..b930adcb --- /dev/null +++ b/crates/libsy-llm-client/Cargo.toml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "switchyard-llm-client" +version = "0.1.0" +description = "HTTP LLM client speaking Switchyard's neutral IR directly" +authors.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +rust-version.workspace = true + +[dependencies] +switchyard-protocol = { path = "../libsy-protocol", version = "0.1.0" } +switchyard-translation = { path = "../switchyard-translation", version = "0.1.0" } +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-native-roots", "stream"] } +async-trait = "0.1" +futures = "0.3" +futures-util = "0.3" +serde_json = "1" +thiserror = "2" + +[dev-dependencies] +pretty_assertions = "1" +tokio = { version = "1", features = ["macros", "rt-multi-thread"] } +wiremock = "0.6" diff --git a/crates/libsy-llm-client/README.md b/crates/libsy-llm-client/README.md new file mode 100644 index 00000000..5bc9afee --- /dev/null +++ b/crates/libsy-llm-client/README.md @@ -0,0 +1,174 @@ + + +# libsy-llm-client + +An HTTP client that speaks Switchyard's neutral IR directly. You hand it a +[`libsy_protocol::Request`] and a model name; it looks up the configured backend, +encodes the request to that backend's wire format, adds auth and forwards your +headers, makes the call with a shared `reqwest::Client`, and decodes the reply +back into a [`libsy_protocol::Response`] — buffered or streamed. + +It depends only on `libsy-protocol` and `switchyard-translation`; no server, no +provider SDK. For an HTTP front door built on top of it, see [`libsy-proxy`]. + +## Concepts + +- **Two-layer map.** A client is built from `model → wire format → Backend`, so + one model can be served over several upstream formats. `call(request, model, format)` + picks the backend by `(model, format)`. +- **Backends.** A [`Backend`] is one of `OpenAiChat`, `OpenAiResponses`, or + `Anthropic`, each wrapping an [`HttpBackendConfig`] (`base_url`, `api_key`, + static `extra_headers`). The variant fixes the URL path and auth scheme + (Bearer vs `x-api-key` + `anthropic-version`). +- **Model rewrite.** The resolved model name is both the map key and the model id + sent upstream — it overwrites whatever `model` the request arrived with. +- **Streaming is chosen by the request.** If the encoded body has `stream: true` + (i.e. `request.llm_request.stream`), you get `LlmResponse::Stream`; otherwise + `LlmResponse::Agg`. + +## Add the dependency + +Within this workspace: + +```toml +[dependencies] +switchyard-llm-client = { path = "../libsy-llm-client" } +switchyard-protocol = { path = "../libsy-protocol" } +switchyard-translation = { path = "../switchyard-translation" } # for WireFormat +``` + +## Quickstart + +### Build a client + +```rust +use std::collections::{BTreeMap, HashMap}; +use switchyard_llm_client::{Backend, HttpBackendConfig, LlmModelClient}; +use switchyard_translation::WireFormat; + +fn build_client() -> libsy_llm_client::Result { + let openai = HttpBackendConfig { + base_url: "https://api.openai.com/v1".to_string(), + api_key: std::env::var("OPENAI_API_KEY").ok(), + extra_headers: BTreeMap::new(), + }; + + // model → format → backend + let map = HashMap::from([( + "gpt-4o-mini".to_string(), + HashMap::from([(WireFormat::OpenAiChat, Backend::OpenAiChat(openai))]), + )]); + + LlmModelClient::new(map) +} +``` + +### Buffered call + +```rust +use libsy_protocol::{completion_text, text_request, Request}; +use switchyard_translation::WireFormat; + +async fn ask(client: &LlmModelClient) -> libsy_llm_client::Result { + let request = Request { + llm_request: text_request(None, "Say hello in five words."), + raw_request: None, + metadata: None, + }; + + // model_name wins over request.llm_request.model; it is also sent upstream. + let response = client + .call(request, Some("gpt-4o-mini"), WireFormat::OpenAiChat) + .await?; + + // Buffered backends return `Agg`; a stream can be folded with `.aggregate()`. + let agg = response.llm_response.aggregate().await.unwrap(); + Ok(completion_text(&agg)) +} +``` + +### Streaming call + +Set `stream` on the IR request and drive the returned chunk stream: + +```rust +use futures_util::StreamExt; +use libsy_protocol::{text_request, LlmResponse, LlmResponseChunk, Request}; +use switchyard_translation::WireFormat; + +async fn stream(client: &LlmModelClient) -> libsy_llm_client::Result<()> { + let mut llm_request = text_request(None, "Count to five."); + llm_request.stream = true; + let request = Request { llm_request, raw_request: None, metadata: None }; + + let response = client + .call(request, Some("gpt-4o-mini"), WireFormat::OpenAiChat) + .await?; + + if let LlmResponse::Stream(mut chunks) = response.llm_response { + while let Some(item) = chunks.next().await { + match item { + Ok(LlmResponseChunk::TextDelta { text, .. }) => print!("{text}"), + Ok(_) => {} // usage, tool-call deltas, message start/stop + Err(error) => return Err(libsy_llm_client::LlmClientError::Stream(error.to_string())), + } + } + } + Ok(()) +} +``` + +## Cross-format translation + +The request/response are translated through the neutral IR, so the inbound shape +you build and the backend's wire format are independent. Pointing a +`WireFormat::AnthropicMessages` backend at an Anthropic endpoint while building +requests with the same helpers works the same way — pass the matching format to +`call`. Register several formats under one model to serve it over more than one +upstream API: + +```rust +let backends = HashMap::from([ + (WireFormat::OpenAiChat, Backend::OpenAiChat(openai_chat_cfg)), + (WireFormat::OpenAiResponses, Backend::OpenAiResponses(openai_resp_cfg)), + (WireFormat::AnthropicMessages, Backend::Anthropic(anthropic_cfg)), +]); +let map = HashMap::from([("my-model".to_string(), backends)]); +// client.formats_for("my-model") -> the configured formats +``` + +## Headers & auth + +- The `Backend` variant sets auth: OpenAI formats send `Authorization: Bearer `; + Anthropic sends `x-api-key: ` plus `anthropic-version`. +- `request.metadata.http_headers` are forwarded upstream, **except** reserved ones: + `host`, `content-length`, `connection`, and the backend-owned + `authorization` / `x-api-key` / `anthropic-version` / `content-type`. So a + caller's placeholder credential never overrides the backend's real key. +- Per-backend static headers go in `HttpBackendConfig::extra_headers`. + +## Errors + +`call` returns [`LlmClientError`]: + +| Variant | When | +|---------|------| +| `MissingModel` | no `model_name` arg and no `request.llm_request.model` | +| `UnknownModel(model)` | model has no backends configured | +| `UnknownModelFormat { model, format }` | model has no backend for that format | +| `Translation(msg)` | encode/decode failed in the translation engine | +| `Transport(msg)` | connect/timeout/transport failure | +| `ContextWindowExceeded { model, message }` | upstream 400 detected as a context overflow (checked before `UpstreamHttp`, so callers can evict-and-retry) | +| `UpstreamHttp { status, body }` | any other non-2xx upstream response | +| `Stream(msg)` | mid-stream read / malformed frame | + +[`libsy_protocol::Request`]: ../libsy-protocol +[`libsy_protocol::Response`]: ../libsy-protocol +[`libsy-proxy`]: ../libsy-proxy +[`Backend`]: src/backend.rs +[`HttpBackendConfig`]: src/backend.rs +[`LlmModelClient`]: src/client.rs +[`LlmClientError`]: src/error.rs diff --git a/crates/libsy-llm-client/src/backend.rs b/crates/libsy-llm-client/src/backend.rs new file mode 100644 index 00000000..0512f2ac --- /dev/null +++ b/crates/libsy-llm-client/src/backend.rs @@ -0,0 +1,246 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Per-provider backend configuration: wire format, upstream URL, and auth. + +use std::collections::BTreeMap; + +use reqwest::RequestBuilder; +use switchyard_protocol::WireFormat; + +use crate::error::is_overflow_body; + +const ANTHROPIC_VERSION: &str = "2023-06-01"; + +// Canonical OpenAI phrase plus NVIDIA/LiteLLM wrap variants. Adding a new +// provider-wrap is a one-line entry here, not a fork of the parsing logic. +const OPENAI_OVERFLOW_PHRASES: &[&str] = &[ + "maximum context length", + "context length exceeded", + "context window", + "context length is only", + "please reduce the length of the input", +]; + +// Anthropic has no structured `error.code`, so detection is phrase-based only. +const ANTHROPIC_OVERFLOW_PHRASES: &[&str] = &[ + "prompt is too long", + "maximum number of tokens", + "context window", + "context length", +]; + +/// Shared HTTP configuration for one upstream backend. +#[derive(Clone, Debug)] +pub struct HttpBackendConfig { + /// Base URL of the provider API (e.g. `https://api.openai.com/v1`). + pub base_url: String, + /// API key for the provider, loaded by the caller. `None` sends no auth. + pub api_key: Option, + /// Static headers added to every outbound call to this backend. + pub extra_headers: BTreeMap, +} + +/// A configured upstream backend, one variant per built-in wire format. +/// +/// The variant fixes the wire format, URL path, and auth scheme together so no +/// invalid combination can be constructed. +#[derive(Clone, Debug)] +pub enum Backend { + /// OpenAI-compatible Chat Completions API. + OpenAiChat(HttpBackendConfig), + /// OpenAI Responses API. + OpenAiResponses(HttpBackendConfig), + /// Anthropic Messages API. + Anthropic(HttpBackendConfig), +} + +impl Backend { + /// The wire format the request IR is encoded to for this backend. + pub fn wire_format(&self) -> WireFormat { + match self { + Backend::OpenAiChat(_) => WireFormat::OpenAiChat, + Backend::OpenAiResponses(_) => WireFormat::OpenAiResponses, + Backend::Anthropic(_) => WireFormat::AnthropicMessages, + } + } + + // Shared HTTP config, regardless of variant. + fn config(&self) -> &HttpBackendConfig { + match self { + Backend::OpenAiChat(config) + | Backend::OpenAiResponses(config) + | Backend::Anthropic(config) => config, + } + } + + /// The fully resolved upstream URL for this backend's endpoint. + /// + /// Tolerates base URLs that already include the provider path (or a bare + /// `/v1`), matching the join rules of the existing native backends. + pub fn url(&self) -> String { + let base_url = self.config().base_url.trim_end_matches('/'); + match self { + Backend::OpenAiChat(_) => openai_url(base_url, "/chat/completions"), + Backend::OpenAiResponses(_) => openai_url(base_url, "/responses"), + Backend::Anthropic(_) => anthropic_url(base_url), + } + } + + /// Applies this backend's auth and version headers to a request builder. + /// + /// OpenAI variants use `Authorization: Bearer `; Anthropic uses + /// `x-api-key: ` plus the required `anthropic-version` header. + pub fn apply_auth(&self, mut builder: RequestBuilder) -> RequestBuilder { + let api_key = self.config().api_key.as_deref(); + match self { + Backend::OpenAiChat(_) | Backend::OpenAiResponses(_) => { + if let Some(api_key) = api_key { + builder = builder.bearer_auth(api_key); + } + } + Backend::Anthropic(_) => { + builder = builder.header("anthropic-version", ANTHROPIC_VERSION); + if let Some(api_key) = api_key { + builder = builder.header("x-api-key", api_key); + } + } + } + builder + } + + /// Static per-backend headers to forward on every call. + pub fn extra_headers(&self) -> &BTreeMap { + &self.config().extra_headers + } + + /// Whether an upstream 400 `body` looks like a context-window overflow for + /// this backend's provider. + pub(crate) fn is_context_overflow(&self, body: &str) -> bool { + match self { + Backend::OpenAiChat(_) | Backend::OpenAiResponses(_) => is_overflow_body( + body, + |value| { + value + .get("error") + .and_then(|err| err.get("code")) + .and_then(serde_json::Value::as_str) + == Some("context_length_exceeded") + }, + OPENAI_OVERFLOW_PHRASES, + ), + Backend::Anthropic(_) => is_overflow_body(body, |_| false, ANTHROPIC_OVERFLOW_PHRASES), + } + } +} + +// Accept either a root `/v1` URL or an already-specific OpenAI endpoint URL. +fn openai_url(base_url: &str, suffix: &str) -> String { + let base_root = base_url + .strip_suffix("/chat/completions") + .or_else(|| base_url.strip_suffix("/responses")) + .unwrap_or(base_url); + format!("{base_root}{suffix}") +} + +// Accept a bare host, a `/v1` root, or an already-specific `/v1/messages` URL. +fn anthropic_url(base_url: &str) -> String { + if base_url.ends_with("/v1/messages") { + base_url.to_string() + } else if base_url.ends_with("/v1") { + format!("{base_url}/messages") + } else { + format!("{base_url}/v1/messages") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn config(base_url: &str) -> HttpBackendConfig { + HttpBackendConfig { + base_url: base_url.to_string(), + api_key: Some("secret".to_string()), + extra_headers: BTreeMap::new(), + } + } + + #[test] + fn openai_chat_url_joins_bare_v1() { + let backend = Backend::OpenAiChat(config("https://api.openai.com/v1")); + assert_eq!(backend.url(), "https://api.openai.com/v1/chat/completions"); + } + + #[test] + fn openai_chat_url_tolerates_trailing_slash_and_existing_suffix() { + assert_eq!( + Backend::OpenAiChat(config("https://api.openai.com/v1/")).url(), + "https://api.openai.com/v1/chat/completions" + ); + assert_eq!( + Backend::OpenAiChat(config("https://api.openai.com/v1/chat/completions")).url(), + "https://api.openai.com/v1/chat/completions" + ); + } + + #[test] + fn openai_responses_url_uses_responses_path() { + assert_eq!( + Backend::OpenAiResponses(config("https://api.openai.com/v1")).url(), + "https://api.openai.com/v1/responses" + ); + } + + #[test] + fn anthropic_url_join_cases() { + assert_eq!( + Backend::Anthropic(config("https://api.anthropic.com")).url(), + "https://api.anthropic.com/v1/messages" + ); + assert_eq!( + Backend::Anthropic(config("https://api.anthropic.com/v1")).url(), + "https://api.anthropic.com/v1/messages" + ); + assert_eq!( + Backend::Anthropic(config("https://api.anthropic.com/v1/messages")).url(), + "https://api.anthropic.com/v1/messages" + ); + } + + #[test] + fn wire_format_matches_variant() { + assert_eq!( + Backend::OpenAiChat(config("x")).wire_format(), + WireFormat::OpenAiChat + ); + assert_eq!( + Backend::OpenAiResponses(config("x")).wire_format(), + WireFormat::OpenAiResponses + ); + assert_eq!( + Backend::Anthropic(config("x")).wire_format(), + WireFormat::AnthropicMessages + ); + } + + #[test] + fn openai_detects_canonical_and_wrapped_overflow() { + let backend = Backend::OpenAiChat(config("x")); + assert!(backend + .is_context_overflow(r#"{"error":{"code":"context_length_exceeded","message":"x"}}"#)); + // NVIDIA/LiteLLM message wrap with no structured code. + assert!(backend.is_context_overflow( + r#"{"error":{"message":"the model's context length is only 131072 tokens"}}"# + )); + assert!(!backend.is_context_overflow(r#"{"error":{"code":"invalid_api_key"}}"#)); + } + + #[test] + fn anthropic_detects_prompt_too_long() { + let backend = Backend::Anthropic(config("x")); + assert!(backend + .is_context_overflow(r#"{"error":{"message":"prompt is too long: 200000 tokens"}}"#)); + assert!(!backend.is_context_overflow(r#"{"error":{"message":"overloaded"}}"#)); + } +} diff --git a/crates/libsy-llm-client/src/client.rs b/crates/libsy-llm-client/src/client.rs new file mode 100644 index 00000000..13407b35 --- /dev/null +++ b/crates/libsy-llm-client/src/client.rs @@ -0,0 +1,812 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! [`TranslatingLlmClient`] — the crate's single public entry point: encode a neutral +//! request, call the configured backend over HTTP, decode the neutral response. + +use std::collections::{BTreeMap, HashMap}; + +use async_trait::async_trait; +use futures_util::StreamExt; +use reqwest::RequestBuilder; +use serde_json::Value; +use switchyard_protocol::{ + Context, Decision, LlmResponse, Metadata, Request, Response, RoutedLlmClient, +}; +use switchyard_translation::{ + decode_buffered_response, decode_request, decode_stream, encode_buffered_response, + encode_request, encode_stream, WireFormat, +}; + +use crate::backend::Backend; +use crate::error::{LlmClientError, Result}; +use crate::raw::RawResponse; + +// Headers this client owns or that are hop-by-hop; never forwarded from the +// caller's metadata. Auth/version/content-type are set by the backend or the +// JSON body, so a forwarded copy would either be ignored or conflict. Compared +// case-insensitively. Aligns with `_SENSITIVE_HEADERS` in the Python +// `switchyard/lib/request_metadata.py` forwarding logic. +const RESERVED_HEADERS: &[&str] = &[ + "host", + "content-length", + "connection", + "authorization", + "x-api-key", + "anthropic-version", + "content-type", +]; + +/// How one model is served: the `default_backend` used when the request does not +/// pin a wire format, plus any `other_backends` reachable over additional formats. +#[derive(Clone, Debug)] +pub struct ModelConfig { + model_name: String, + default_backend: Backend, + other_backends: Option>, +} + +impl ModelConfig { + /// A model named `model_name` served by `default_backend`, optionally reachable + /// over additional wire formats via `other_backends`. + pub fn new( + model_name: impl Into, + default_backend: Backend, + other_backends: Option>, + ) -> Self { + Self { + model_name: model_name.into(), + default_backend, + other_backends, + } + } +} + +/// A client that dispatches neutral-IR requests to per-model HTTP backends. +/// +/// Construct it with a list of [`ModelConfig`]s — one per model, each naming a +/// default [`Backend`] and any additional per-format backends. Each call resolves +/// the model and wire format, encodes the request to that backend's wire format, +/// applies auth and forwarded headers, sends the HTTP request with a shared +/// [`reqwest::Client`], and decodes the response back to the neutral IR (buffered +/// or streamed). +pub struct TranslatingLlmClient { + model_to_config: HashMap, + client: reqwest::Client, +} + +impl TranslatingLlmClient { + /// Builds a client over the given [`ModelConfig`]s, with a fresh shared HTTP + /// client and the built-in translation codecs. + pub fn new(model_configs: &[ModelConfig]) -> Result { + let client = reqwest::Client::builder().build().map_err(|error| { + LlmClientError::Transport(format!("failed to build HTTP client: {error}")) + })?; + let model_to_config = model_configs + .iter() + .map(|config| (config.model_name.clone(), config.clone())) + .collect(); + + Ok(Self { + model_to_config, + client, + }) + } + + /// The backend serving `model` over `format` — the default backend when its + /// format matches, otherwise a matching entry in `other_backends`; `None` when + /// the model is unknown or has no backend for `format`. + pub fn backend_for(&self, model: &str, format: WireFormat) -> Option<&Backend> { + self.model_to_config.get(model).and_then(|config| { + if config.default_backend.wire_format() == format { + Some(&config.default_backend) + } else { + config + .other_backends + .as_ref() + .and_then(|backends| backends.iter().find(|b| b.wire_format() == format)) + } + }) + } + + /// Calls the backend for `model_name` (or the request's own model), over the + /// wire format the request pins in its metadata (else the model's default + /// backend), and returns the neutral response. + /// + /// Resolution: `model_name` wins over `request.llm_request.model`; the + /// resolved name is both the outer map key and the model id written into the + /// request before translation. Errors with [`LlmClientError::MissingModel`] + /// when neither is set, [`LlmClientError::UnknownModel`] when the model has + /// no backends, and [`LlmClientError::UnknownModelFormat`] when the model has + /// no backend for `format`. + pub async fn call_rewrite_model( + &self, + _ctx: Context, + request: Request, + model_name: Option<&str>, + ) -> Result { + // Own the request's parts so the model can be set without a `mut` param + // and without cloning the messages. `raw_request` is unused here. + let Request { + mut llm_request, + metadata, + .. + } = request; + + let model = model_name + .map(str::to_string) + .or_else(|| llm_request.model.clone()) + .ok_or(LlmClientError::MissingModel)?; + + let orig_format = metadata.as_ref().and_then(|m| m.wire_format); + let wire_format = orig_format.unwrap_or( + self.model_to_config + .get(&model) + .map(|config| config.default_backend.wire_format()) + .ok_or(LlmClientError::UnknownModel(model.clone()))?, + ); + let backend = + self.backend_for(&model, wire_format) + .ok_or(LlmClientError::UnknownModelFormat { + model: model.clone(), + format: wire_format, + })?; + + // The resolved name is the upstream model id (per the crate contract). + llm_request.model = Some(model.clone()); + + let mut body = encode_request(&llm_request, wire_format) + .map_err(|error| LlmClientError::Translation(error.to_string()))?; + // `encode_request` round-trips a preserved same-format body verbatim, + // which keeps the caller's original `model`; force the resolved model so + // the upstream always sees the target id. + set_json_model(&mut body, &model); + let streaming = body.get("stream").and_then(Value::as_bool).unwrap_or(false); + + let builder = self.client.post(backend.url()).json(&body); + let builder = forward_metadata_headers(builder, metadata.as_ref()); + let builder = apply_extra_headers(builder, backend); + let builder = backend.apply_auth(builder); + + let http_response = builder + .send() + .await + .map_err(|error| LlmClientError::Transport(error.to_string()))?; + let status = http_response.status(); + if !status.is_success() { + let body = http_response + .text() + .await + .unwrap_or_else(|error| format!("")); + if status == reqwest::StatusCode::BAD_REQUEST && backend.is_context_overflow(&body) { + return Err(LlmClientError::ContextWindowExceeded { + model, + message: body, + }); + } + return Err(LlmClientError::UpstreamHttp { + status: status.as_u16(), + body, + }); + } + + let llm_response = if streaming { + // Adapt the reqwest body stream to plain bytes; the SSE-decode itself is + // transport-agnostic and lives in `switchyard-translation`. + let bytes = http_response.bytes_stream().map(|chunk| { + chunk.map(|bytes| bytes.to_vec()).map_err(|error| { + Box::new(LlmClientError::Stream(format!( + "stream read failed: {error}" + ))) as Box + }) + }); + LlmResponse::Stream(decode_stream(bytes, wire_format)) + } else { + let body = http_response.json::().await.map_err(|error| { + LlmClientError::Transport(format!("invalid upstream JSON: {error}")) + })?; + let agg = decode_buffered_response(&body, wire_format) + .map_err(|error| LlmClientError::Translation(error.to_string()))?; + LlmResponse::Agg(agg) + }; + + Ok(Response { + llm_response, + metadata, + }) + } + + /// The whole decode → call → encode path a wire endpoint needs, in one call. + /// + /// Decodes `raw_http_request` from `wire_format` to the neutral IR, serves it via + /// [`call_rewrite_model`](Self::call_rewrite_model) — the *upstream* wire format is + /// resolved there from the model's backend, independently of `wire_format` — then + /// encodes the neutral response back into `wire_format`. The result is a buffered + /// [`RawResponse::Buffered`] JSON body or a streamed [`RawResponse::Stream`] of + /// wire events (the caller frames the stream as SSE). The response's `model` is + /// restamped with the model the request asked for, never the upstream id. + /// + /// `http_headers` are carried through as the request's + /// [`Metadata::http_headers`] and forwarded to the upstream (minus the reserved + /// set); pass `None` to forward nothing. + pub async fn call_rewrite_model_raw( + &self, + ctx: Context, + raw_http_request: Value, + http_headers: Option>, + model: Option<&str>, + wire_format: WireFormat, + ) -> Result { + let llm_request = decode_request(wire_format, &raw_http_request) + .map_err(|error| LlmClientError::Translation(error.to_string()))?; + // The model the client asked for; restamped onto the response so it never + // leaks the upstream id. + let requested_model = llm_request.model.clone(); + + let request = Request { + llm_request, + raw_request: None, + metadata: Some(Metadata { + session_id: None, + agent_id: None, + task_id: None, + correlation_id: None, + extra_metadata: None, + http_headers, + wire_format: None, + }), + }; + let response = self.call_rewrite_model(ctx, request, model).await?; + + match response.llm_response { + LlmResponse::Agg(agg) => { + let body = encode_buffered_response(&agg, wire_format, requested_model.as_deref()) + .map_err(|error| LlmClientError::Translation(error.to_string()))?; + Ok(RawResponse::Buffered(body)) + } + LlmResponse::Stream(chunks) => Ok(RawResponse::Stream(encode_stream( + chunks, + wire_format, + requested_model, + ))), + } + } +} + +#[async_trait] +impl RoutedLlmClient for TranslatingLlmClient { + async fn call( + &self, + ctx: Context, + request: Request, + decision: std::sync::Arc, + ) -> std::result::Result> { + let model_name = Some(decision.selected_model()); + self.call_rewrite_model(ctx, request, model_name) + .await + .map_err(|error| Box::new(error) as Box) + } +} + +// Forwards caller-supplied metadata headers, skipping the reserved set. +fn forward_metadata_headers( + mut builder: RequestBuilder, + metadata: Option<&Metadata>, +) -> RequestBuilder { + let Some(headers) = metadata.and_then(|metadata| metadata.http_headers.as_ref()) else { + return builder; + }; + for (name, value) in headers { + if is_reserved_header(name) { + continue; + } + builder = builder.header(name, value); + } + builder +} + +// Adds the backend's static per-call headers. +fn apply_extra_headers(mut builder: RequestBuilder, backend: &Backend) -> RequestBuilder { + for (name, value) in backend.extra_headers() { + builder = builder.header(name, value); + } + builder +} + +// Overwrites the outbound body's `model` field with the resolved model id. +fn set_json_model(body: &mut Value, model: &str) { + if let Value::Object(object) = body { + object.insert("model".to_string(), Value::String(model.to_string())); + } +} + +// Case-insensitive membership test against RESERVED_HEADERS. +fn is_reserved_header(name: &str) -> bool { + RESERVED_HEADERS + .iter() + .any(|reserved| name.eq_ignore_ascii_case(reserved)) +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use serde_json::json; + use switchyard_protocol::{completion_text, text_request, LlmRequest}; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::*; + use crate::backend::HttpBackendConfig; + + fn config(base_url: &str) -> HttpBackendConfig { + HttpBackendConfig { + base_url: base_url.to_string(), + api_key: Some("secret".to_string()), + extra_headers: BTreeMap::new(), + } + } + + // A one-model config list: "gpt" served over OpenAI Chat at base_url. + fn chat_map(base_url: &str) -> Vec { + vec![ModelConfig::new( + "gpt", + Backend::OpenAiChat(config(base_url)), + None, + )] + } + + fn request_for(model: Option<&str>, stream: bool) -> Request { + let mut llm_request = text_request(model.map(str::to_string), "hi"); + llm_request.stream = stream; + Request { + llm_request, + raw_request: None, + metadata: None, + } + } + + // A request that pins `format` in its metadata, so the client resolves that + // wire format instead of the model's default backend. + fn request_with_wire_format(model: &str, format: WireFormat) -> Request { + let mut request = request_for(Some(model), false); + request.metadata = Some(Metadata { + session_id: None, + agent_id: None, + task_id: None, + correlation_id: None, + extra_metadata: None, + http_headers: None, + wire_format: Some(format), + }); + request + } + + #[test] + fn reserved_headers_are_case_insensitive() { + assert!(is_reserved_header("Authorization")); + assert!(is_reserved_header("content-type")); + assert!(is_reserved_header("X-Api-Key")); + assert!(!is_reserved_header("x-request-id")); + } + + #[tokio::test] + async fn missing_model_errors() { + let client = TranslatingLlmClient::new(&[]).unwrap(); + let Err(error) = client + .call_rewrite_model(Context::default(), request_for(None, false), None) + .await + else { + panic!("expected an error"); + }; + assert!(matches!(error, LlmClientError::MissingModel)); + } + + #[tokio::test] + async fn unknown_model_errors() { + let client = TranslatingLlmClient::new(&[]).unwrap(); + let Err(error) = client + .call_rewrite_model(Context::default(), request_for(Some("gpt"), false), None) + .await + else { + panic!("expected an error"); + }; + assert!(matches!(error, LlmClientError::UnknownModel(model) if model == "gpt")); + } + + #[tokio::test] + async fn unknown_model_format_errors() { + // "gpt" exists but only over OpenAI Chat; the request pins Anthropic. + let client = TranslatingLlmClient::new(&chat_map("https://example.test/v1")).unwrap(); + let Err(error) = client + .call_rewrite_model( + Context::default(), + request_with_wire_format("gpt", WireFormat::AnthropicMessages), + None, + ) + .await + else { + panic!("expected an error"); + }; + assert!(matches!( + error, + LlmClientError::UnknownModelFormat { model, format } + if model == "gpt" && format == WireFormat::AnthropicMessages + )); + } + + #[test] + fn backend_for_resolves_configured_format() { + let client = TranslatingLlmClient::new(&chat_map("https://example.test/v1")).unwrap(); + // "gpt" is served over OpenAI Chat only; other formats and models miss. + assert!(client.backend_for("gpt", WireFormat::OpenAiChat).is_some()); + assert!(client + .backend_for("gpt", WireFormat::AnthropicMessages) + .is_none()); + assert!(client + .backend_for("missing", WireFormat::OpenAiChat) + .is_none()); + } + + #[tokio::test] + async fn model_name_arg_wins_over_request_model() { + let client = TranslatingLlmClient::new(&[]).unwrap(); + // Arg "b" is looked up (and reported), not the request's "a". + let Err(error) = client + .call_rewrite_model(Context::default(), request_for(Some("a"), false), Some("b")) + .await + else { + panic!("expected an error"); + }; + assert!(matches!(error, LlmClientError::UnknownModel(model) if model == "b")); + } + + #[tokio::test] + async fn buffered_openai_chat_round_trips() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "chatcmpl-1", + "model": "gpt", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hi there"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} + }))) + .mount(&server) + .await; + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri()))).unwrap(); + + let response = client + .call_rewrite_model(Context::default(), request_for(Some("gpt"), false), None) + .await + .unwrap(); + let agg = response + .llm_response + .into_agg() + .await + .expect("buffered response"); + assert_eq!(completion_text(&agg), "Hi there"); + } + + #[tokio::test] + async fn rewrites_model_to_resolved_upstream_id() { + // Inbound body says "switchyard"; the upstream must receive "gpt". + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .and(wiremock::matchers::body_partial_json(json!({"model": "gpt"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "1", "model": "gpt", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {} + }))) + .mount(&server) + .await; + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri()))).unwrap(); + // Inbound model differs from the map key / resolved model. + client + .call_rewrite_model( + Context::default(), + request_for(Some("switchyard"), false), + Some("gpt"), + ) + .await + .unwrap(); + // The body_partial_json matcher asserts the upstream saw model "gpt". + } + + #[tokio::test] + async fn streaming_openai_chat_aggregates() { + let server = MockServer::start().await; + let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\ + data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\ + data: [DONE]\n\n"; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")) + .mount(&server) + .await; + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri()))).unwrap(); + + let response = client + .call_rewrite_model(Context::default(), request_for(Some("gpt"), true), None) + .await + .unwrap(); + assert!(matches!(response.llm_response, LlmResponse::Stream(_))); + let agg = response.llm_response.into_agg().await.unwrap(); + assert_eq!(completion_text(&agg), "Hello world"); + } + + #[tokio::test] + async fn upstream_500_is_upstream_http() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(500).set_body_string("boom")) + .mount(&server) + .await; + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri()))).unwrap(); + + let Err(error) = client + .call_rewrite_model(Context::default(), request_for(Some("gpt"), false), None) + .await + else { + panic!("expected an error"); + }; + assert!(matches!( + error, + LlmClientError::UpstreamHttp { status: 500, .. } + )); + } + + #[tokio::test] + async fn context_overflow_400_is_mapped() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(400).set_body_json(json!({ + "error": {"code": "context_length_exceeded", "message": "too big"} + }))) + .mount(&server) + .await; + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri()))).unwrap(); + + let Err(error) = client + .call_rewrite_model(Context::default(), request_for(Some("gpt"), false), None) + .await + else { + panic!("expected an error"); + }; + assert!(matches!( + error, + LlmClientError::ContextWindowExceeded { model, .. } if model == "gpt" + )); + } + + #[tokio::test] + async fn forwards_metadata_headers_except_reserved() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(wiremock::matchers::header("x-request-id", "abc")) + // A forwarded Authorization must NOT override the backend's bearer key. + .and(wiremock::matchers::header("authorization", "Bearer secret")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "1", "model": "gpt", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {} + }))) + .mount(&server) + .await; + + let mut headers = BTreeMap::new(); + headers.insert("x-request-id".to_string(), "abc".to_string()); + headers.insert("authorization".to_string(), "Bearer client-key".to_string()); + let request = Request { + llm_request: LlmRequest { + model: Some("gpt".to_string()), + ..LlmRequest::default() + }, + raw_request: None, + metadata: Some(Metadata { + session_id: None, + agent_id: None, + task_id: None, + correlation_id: None, + extra_metadata: None, + http_headers: Some(headers), + wire_format: None, + }), + }; + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri()))).unwrap(); + + // Matchers assert forwarded x-request-id survives and reserved + // authorization is the backend's, not the client's. + client + .call_rewrite_model(Context::default(), request, None) + .await + .unwrap(); + } + + // Minimal `Decision` for driving the client through the `RoutedLlmClient` trait. + struct FixedDecision(&'static str); + + impl Decision for FixedDecision { + fn selected_model(&self) -> &str { + self.0 + } + fn reasoning(&self) -> Option<&str> { + None + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + // Exercises the `RoutedLlmClient` impl: `call` resolves the upstream model from the + // decision (the request carries none) and round-trips a buffered response. + #[tokio::test] + async fn routed_llm_client_serves_the_decision_model() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "chatcmpl-1", + "model": "gpt", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "routed hi"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} + }))) + .mount(&server) + .await; + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri()))).unwrap(); + let decision: std::sync::Arc = std::sync::Arc::new(FixedDecision("gpt")); + // Called through the trait; the request has no model, so "gpt" comes from the decision. + let response = client + .call(Context::default(), request_for(None, false), decision) + .await + .unwrap(); + let agg = response + .llm_response + .into_agg() + .await + .expect("buffered response"); + assert_eq!(completion_text(&agg), "routed hi"); + } + + // Raw path, buffered: decode an OpenAI Chat body -> call -> encode back to OpenAI + // Chat JSON, with the client-facing `model` restamped over the upstream id. + #[tokio::test] + async fn call_rewrite_model_raw_round_trips_buffered_json() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "chatcmpl-1", + "model": "gpt", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hi there"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} + }))) + .mount(&server) + .await; + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri()))).unwrap(); + let raw = json!({ + "model": "client-facing", + "messages": [{"role": "user", "content": "hi"}] + }); + let RawResponse::Buffered(body) = client + .call_rewrite_model_raw( + Context::default(), + raw, + None, + Some("gpt"), + WireFormat::OpenAiChat, + ) + .await + .unwrap() + else { + panic!("expected a buffered response"); + }; + + assert_eq!(body["choices"][0]["message"]["content"], "Hi there"); + // The client sees the model it asked for, not the upstream "gpt". + assert_eq!(body["model"], "client-facing"); + } + + // Raw path, streaming: an inbound `stream: true` request yields an unframed stream + // of OpenAI Chat chunk objects whose deltas reassemble the completion. + #[tokio::test] + async fn call_rewrite_model_raw_streams_wire_events() { + use futures::StreamExt; + + let server = MockServer::start().await; + let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\ + data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\n\ + data: [DONE]\n\n"; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream")) + .mount(&server) + .await; + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri()))).unwrap(); + let raw = json!({ + "model": "client-facing", + "messages": [{"role": "user", "content": "hi"}], + "stream": true + }); + let RawResponse::Stream(stream) = client + .call_rewrite_model_raw( + Context::default(), + raw, + None, + Some("gpt"), + WireFormat::OpenAiChat, + ) + .await + .unwrap() + else { + panic!("expected a streamed response"); + }; + + let events: Vec = stream.map(|item| item.unwrap()).collect().await; + assert!(!events.is_empty(), "expected at least one wire event"); + let content: String = events + .iter() + .filter_map(|event| event["choices"][0]["delta"]["content"].as_str()) + .collect(); + assert_eq!(content, "Hello world"); + } + + // Raw path forwards caller headers (minus the reserved set) to the upstream. + #[tokio::test] + async fn call_rewrite_model_raw_forwards_headers() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(wiremock::matchers::header("x-request-id", "abc")) + // A forwarded authorization must NOT override the backend's bearer key. + .and(wiremock::matchers::header("authorization", "Bearer secret")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "1", "model": "gpt", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {} + }))) + .mount(&server) + .await; + + let mut headers = BTreeMap::new(); + headers.insert("x-request-id".to_string(), "abc".to_string()); + headers.insert("authorization".to_string(), "Bearer client-key".to_string()); + + let client = TranslatingLlmClient::new(&chat_map(&format!("{}/v1", server.uri()))).unwrap(); + let raw = json!({"model": "gpt", "messages": [{"role": "user", "content": "hi"}]}); + // Matchers assert the forwarded x-request-id survives and reserved + // authorization is the backend's, not the client's. + client + .call_rewrite_model_raw( + Context::default(), + raw, + Some(headers), + Some("gpt"), + WireFormat::OpenAiChat, + ) + .await + .unwrap(); + } +} diff --git a/crates/libsy-llm-client/src/error.rs b/crates/libsy-llm-client/src/error.rs new file mode 100644 index 00000000..d02c0523 --- /dev/null +++ b/crates/libsy-llm-client/src/error.rs @@ -0,0 +1,157 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Error type for the LLM client, plus shared context-window-overflow detection. +//! +//! The overflow detection is ported from +//! `switchyard-components/src/backends/context_overflow.rs`; that helper is +//! crate-private there and this crate cannot depend on `switchyard-components`, +//! so the small, self-contained logic is vendored here. + +use serde_json::Value; +use switchyard_translation::WireFormat; +use thiserror::Error; + +/// Result alias for LLM client operations. +pub type Result = std::result::Result; + +/// Failures surfaced while resolving, translating, sending, or decoding a call. +#[derive(Debug, Error)] +pub enum LlmClientError { + /// The resolved model name has no entry in the client's backend map. + #[error("no backend configured for model {0:?}")] + UnknownModel(String), + + /// The model exists but has no backend configured for the requested format. + #[error("model {model:?} has no backend for format {format}")] + UnknownModelFormat { + /// Model that was resolved. + model: String, + /// Wire format requested for the call. + format: WireFormat, + }, + + /// Neither an explicit `model_name` nor a model on the request was given. + #[error("no model given: pass model_name or set request.llm_request.model")] + MissingModel, + + /// Request encoding or response decoding failed in the translation engine. + #[error("translation failed: {0}")] + Translation(String), + + /// The HTTP request could not be sent (connect/timeout/transport failure). + #[error("upstream transport error: {0}")] + Transport(String), + + /// An upstream 400 whose body is detected as a context-window overflow. + /// + /// Kept distinct from [`LlmClientError::UpstreamHttp`] and checked first so a + /// caller can implement evict-and-retry by matching this variant instead of + /// sniffing error strings. + #[error("context window exceeded for model {model}: {message}")] + ContextWindowExceeded { + /// Model that overflowed, as sent upstream. + model: String, + /// Raw upstream error body. + message: String, + }, + + /// A non-2xx upstream response that is not a context-window overflow. + #[error("upstream returned HTTP {status}: {body}")] + UpstreamHttp { + /// Upstream HTTP status code. + status: u16, + /// Raw upstream error body. + body: String, + }, + + /// A streamed response failed mid-flight (read error or malformed frame). + #[error("stream error: {0}")] + Stream(String), +} + +/// Detects a context-overflow body using a provider-supplied structured check +/// and a substring phrase list. +/// +/// Parses the body once, runs the structured check (e.g. against `error.code`), +/// then falls back to matching phrases against `error.message` or — when the +/// body is not JSON — the raw body. Centralizing the shape means each new +/// provider-wrap of the canonical error is a one-line phrase entry, not a fork +/// of the parsing logic. +pub(crate) fn is_overflow_body(body: &str, structured_check: F, phrases: &[&str]) -> bool +where + F: Fn(&Value) -> bool, +{ + if let Ok(value) = serde_json::from_str::(body) { + if structured_check(&value) { + return true; + } + if let Some(message) = value + .get("error") + .and_then(|err| err.get("message")) + .and_then(Value::as_str) + { + if contains_any(message, phrases) { + return true; + } + } + } + // Some upstream proxies return plain-text bodies; fall through to a string + // match on the raw body. + contains_any(body, phrases) +} + +// Case-insensitive substring match of any phrase against the message. +fn contains_any(message: &str, phrases: &[&str]) -> bool { + let lower = message.to_ascii_lowercase(); + phrases.iter().any(|phrase| lower.contains(phrase)) +} + +#[cfg(test)] +mod tests { + use super::*; + + const PHRASES: &[&str] = &["context window", "too long"]; + + fn never(_value: &Value) -> bool { + false + } + + #[test] + fn structured_check_short_circuits() { + let body = r#"{"error":{"code":"context_length_exceeded","message":"unrelated"}}"#; + let matched = is_overflow_body( + body, + |value| { + value + .get("error") + .and_then(|err| err.get("code")) + .and_then(Value::as_str) + == Some("context_length_exceeded") + }, + &[], + ); + assert!(matched); + } + + #[test] + fn falls_back_to_message_phrase_match() { + let body = r#"{"error":{"message":"prompt too long"}}"#; + assert!(is_overflow_body(body, never, PHRASES)); + } + + #[test] + fn matches_plain_text_body() { + assert!(is_overflow_body( + "plain text mentioning context window", + never, + PHRASES + )); + } + + #[test] + fn non_match_returns_false() { + let body = r#"{"error":{"message":"rate limit exceeded"}}"#; + assert!(!is_overflow_body(body, never, PHRASES)); + } +} diff --git a/crates/libsy-llm-client/src/lib.rs b/crates/libsy-llm-client/src/lib.rs new file mode 100644 index 00000000..97fe5b3a --- /dev/null +++ b/crates/libsy-llm-client/src/lib.rs @@ -0,0 +1,30 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! HTTP LLM client that speaks Switchyard's neutral IR directly. +//! +//! [`TranslatingLlmClient`] maps a model name (and the wire format resolved from +//! the request) to a [`Backend`], +//! encodes a [`switchyard_protocol::Request`] to that backend's wire format via +//! `switchyard-translation`, applies auth and forwards caller headers, makes the +//! HTTP call with a shared [`reqwest::Client`], and decodes the wire response +//! back to a [`switchyard_protocol::Response`] — supporting both buffered and +//! streamed responses. +//! +//! The crate depends only on `libsy-protocol` and `switchyard-translation`. Neutral +//! IR encode/decode — including SSE stream decoding — lives in +//! `switchyard-translation`; this crate is the HTTP transport around it. The +//! context-overflow detection in [`mod@error`] is vendored from +//! `switchyard-components` (whose copy is crate-private and unavailable here); a +//! future refactor could promote it to a shared location. + +pub mod backend; +pub mod client; +pub mod error; +pub mod raw; + +pub use backend::{Backend, HttpBackendConfig}; +pub use client::{ModelConfig, TranslatingLlmClient}; +pub use error::{LlmClientError, Result}; +pub use raw::RawResponse; +pub use switchyard_translation::RawEventStream; diff --git a/crates/libsy-llm-client/src/raw.rs b/crates/libsy-llm-client/src/raw.rs new file mode 100644 index 00000000..7391c1a0 --- /dev/null +++ b/crates/libsy-llm-client/src/raw.rs @@ -0,0 +1,22 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! The raw wire result of a call. +//! +//! The encoders that produce it — buffered JSON and the streamed wire-event +//! encoder — live in `switchyard-translation` alongside the rest of the neutral-IR +//! codecs; this module only holds the result type they feed. + +use serde_json::Value; +use switchyard_translation::RawEventStream; + +/// The wire result of +/// [`call_rewrite_model_raw`](crate::TranslatingLlmClient::call_rewrite_model_raw): +/// a buffered JSON body, or a live stream of wire events — both already in the +/// requested wire format. +pub enum RawResponse { + /// A complete JSON response body. + Buffered(Value), + /// A live stream of wire-format event objects, ready for SSE framing. + Stream(RawEventStream), +} From 5b6b07bc544f4af2e66c984ced5582727b2b140d Mon Sep 17 00:00:00 2001 From: Greg Clark Date: Thu, 16 Jul 2026 21:26:05 -0400 Subject: [PATCH 03/14] feat(libsy-server): proxy binary Signed-off-by: Greg Clark --- Cargo.lock | 18 ++ Cargo.toml | 1 + crates/libsy-llm-client/src/client.rs | 4 +- crates/libsy-server/Cargo.toml | 28 +++ crates/libsy-server/README.md | 58 +++++ crates/libsy-server/src/cli.rs | 141 +++++++++++ crates/libsy-server/src/lib.rs | 236 +++++++++++++++++++ crates/libsy-server/src/main.rs | 20 ++ crates/libsy-server/src/sse.rs | 87 +++++++ crates/libsy-server/tests/server.rs | 322 ++++++++++++++++++++++++++ 10 files changed, 913 insertions(+), 2 deletions(-) create mode 100644 crates/libsy-server/Cargo.toml create mode 100644 crates/libsy-server/README.md create mode 100644 crates/libsy-server/src/cli.rs create mode 100644 crates/libsy-server/src/lib.rs create mode 100644 crates/libsy-server/src/main.rs create mode 100644 crates/libsy-server/src/sse.rs create mode 100644 crates/libsy-server/tests/server.rs diff --git a/Cargo.lock b/Cargo.lock index 2dda8961..df26224a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -931,6 +931,24 @@ dependencies = [ "tokio-stream", ] +[[package]] +name = "libsy-server" +version = "0.1.0" +dependencies = [ + "async-stream", + "axum", + "clap", + "futures-util", + "http-body-util", + "serde_json", + "switchyard-llm-client", + "switchyard-protocol", + "switchyard-translation", + "tokio", + "tower", + "wiremock", +] + [[package]] name = "libyaml-rs" version = "0.3.0" diff --git a/Cargo.toml b/Cargo.toml index f39f72ed..e2a989de 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ members = [ "crates/libsy", "crates/libsy-examples", "crates/libsy-llm-client", + "crates/libsy-server", "crates/switchyard-components", "crates/switchyard-components-v2", "crates/switchyard-components-v2-macros", diff --git a/crates/libsy-llm-client/src/client.rs b/crates/libsy-llm-client/src/client.rs index 13407b35..bc6e8420 100644 --- a/crates/libsy-llm-client/src/client.rs +++ b/crates/libsy-llm-client/src/client.rs @@ -245,7 +245,7 @@ impl TranslatingLlmClient { let request = Request { llm_request, - raw_request: None, + raw_request: Some(raw_http_request), metadata: Some(Metadata { session_id: None, agent_id: None, @@ -253,7 +253,7 @@ impl TranslatingLlmClient { correlation_id: None, extra_metadata: None, http_headers, - wire_format: None, + wire_format: Some(wire_format), }), }; let response = self.call_rewrite_model(ctx, request, model).await?; diff --git a/crates/libsy-server/Cargo.toml b/crates/libsy-server/Cargo.toml new file mode 100644 index 00000000..608eedfa --- /dev/null +++ b/crates/libsy-server/Cargo.toml @@ -0,0 +1,28 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "libsy-server" +version = "0.1.0" +description = "Minimal axum server exposing LLM wire APIs over switchyard-llm-client" +authors.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +rust-version.workspace = true + +[dependencies] +switchyard-protocol = { path = "../libsy-protocol" } +switchyard-llm-client = { path = "../libsy-llm-client" } +switchyard-translation = { path = "../switchyard-translation" } +axum = "0.8" +clap = { version = "4", features = ["derive", "env"] } +tokio = { version = "1", features = ["macros", "net", "rt-multi-thread", "signal"] } +async-stream = "0.3" +futures-util = "0.3" +serde_json = "1" + +[dev-dependencies] +wiremock = "0.6" +tower = { version = "0.5", features = ["util"] } +http-body-util = "0.1" diff --git a/crates/libsy-server/README.md b/crates/libsy-server/README.md new file mode 100644 index 00000000..648765b1 --- /dev/null +++ b/crates/libsy-server/README.md @@ -0,0 +1,58 @@ +# libsy-server + +A minimal [axum](https://github.com/tokio-rs/axum) server that exposes the OpenAI +Chat Completions, OpenAI Responses, and Anthropic Messages wire APIs and forwards +every request to a single upstream via +[`switchyard-llm-client`](../libsy-llm-client). + +This crate is only the HTTP surface: routing, header normalization, SSE framing, +and error mapping. The actual work — decoding an inbound request to Switchyard's +neutral IR, calling the upstream, and encoding the response back into the same +wire format (buffered JSON or a stream of wire events) — lives in +`switchyard-llm-client`'s `TranslatingLlmClient::call_rewrite_model_raw`. It +mirrors `switchyard-server` but swaps the profile-chain executor for one client. + +## Endpoints + +| Method | Path | Purpose | +|--------|------------------------|----------------------------------| +| POST | `/v1/chat/completions` | OpenAI Chat Completions inbound | +| POST | `/v1/messages` | Anthropic Messages inbound | +| POST | `/v1/responses` | OpenAI Responses inbound | +| GET | `/v1/models` | Lists the single served model | +| GET | `/health` | Liveness check | + +Clients address the server's served model id, `switchyard`; the response is +always restamped with that id rather than the real upstream model. + +## Upstream and served formats + +Serving is same-format: each inbound endpoint is forwarded to an upstream of the +same wire format. A backend is configured for every format whose provider key is +set — OpenAI Chat + Responses from `OPENAI_API_KEY`, Anthropic Messages from +`ANTHROPIC_API_KEY` — and the server serves exactly those endpoints. An inbound +request for a format with no configured backend gets a `400` with an +`unsupported_format` error. + +## Running + +```bash +export OPENAI_API_KEY="sk-..." # and/or ANTHROPIC_API_KEY + +cargo run -p libsy-server -- \ + --base-url https://api.openai.com/v1 \ + --model-name gpt-4o-mini \ + --port 4000 +``` + +Flags: `--base-url` (required), `--model-name` (required upstream id), +`--host` (default `127.0.0.1`), `--port` (default `4000`). + +## Testing + +```bash +cargo test -p libsy-server +``` + +The integration tests drive `build_router` with `tower::oneshot` against a +`wiremock` upstream — no sockets, no credentials. diff --git a/crates/libsy-server/src/cli.rs b/crates/libsy-server/src/cli.rs new file mode 100644 index 00000000..a602eb28 --- /dev/null +++ b/crates/libsy-server/src/cli.rs @@ -0,0 +1,141 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! CLI: parse flags, build the client from env credentials, and serve. + +use std::collections::BTreeMap; +use std::env; +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + +use clap::Parser; +use switchyard_llm_client::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; +use switchyard_translation::WireFormat; +use tokio::net::TcpListener; + +use crate::{build_router, ProxyState, SERVED_MODEL}; + +const DEFAULT_HOST: IpAddr = IpAddr::V4(Ipv4Addr::LOCALHOST); +const DEFAULT_PORT: u16 = 4000; +const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; +const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY"; + +/// Command-line arguments for the server binary. +#[derive(Debug, Parser)] +#[command( + name = "libsy-server", + about = "Minimal LLM API server over switchyard-llm-client", + version +)] +pub struct Args { + /// Upstream base URL the server forwards to (e.g. https://api.openai.com/v1). + #[arg(long)] + pub base_url: String, + + /// Upstream model id sent on every call. + #[arg(long)] + pub model_name: String, + + /// Host address to bind. + #[arg(long, default_value_t = DEFAULT_HOST)] + pub host: IpAddr, + + /// Port to bind. + #[arg(long, default_value_t = DEFAULT_PORT)] + pub port: u16, +} + +impl Args { + /// Builds server state (and the bind address) from the args and env keys. + /// + /// The client's raw path is same-format: each inbound format is served by a + /// backend of the *same* format. A backend is configured for each format whose + /// provider key is set — OpenAI Chat + Responses from `OPENAI_API_KEY`, + /// Anthropic from `ANTHROPIC_API_KEY` — so the server serves exactly those + /// inbound endpoints. Errors if neither key is present. + fn build(self) -> Result<(ProxyState, SocketAddr, Vec), String> { + let openai_key = env_key(OPENAI_API_KEY_ENV); + let anthropic_key = env_key(ANTHROPIC_API_KEY_ENV); + + let mut backends: Vec = Vec::new(); + if let Some(key) = &openai_key { + backends.push(Backend::OpenAiChat(self.config(key))); + backends.push(Backend::OpenAiResponses(self.config(key))); + } + if let Some(key) = &anthropic_key { + backends.push(Backend::Anthropic(self.config(key))); + } + if backends.is_empty() { + return Err(format!( + "no upstream credentials: set {OPENAI_API_KEY_ENV} and/or {ANTHROPIC_API_KEY_ENV}" + )); + } + + let served: Vec = backends.iter().map(Backend::wire_format).collect(); + // Which backend is "default" vs "other" does not affect routing: the raw + // path resolves a backend by the inbound format, matching either slot. + let default_backend = backends.remove(0); + let other_backends = (!backends.is_empty()).then_some(backends); + + let config = ModelConfig::new(self.model_name.clone(), default_backend, other_backends); + let client = TranslatingLlmClient::new(&[config]).map_err(|error| error.to_string())?; + let addr = SocketAddr::new(self.host, self.port); + let state = ProxyState::new(client, self.model_name.as_str()); + Ok((state, addr, served)) + } + + fn config(&self, api_key: &str) -> HttpBackendConfig { + HttpBackendConfig { + base_url: self.base_url.clone(), + api_key: Some(api_key.to_string()), + extra_headers: BTreeMap::new(), + } + } +} + +/// Parses args, binds the listener, and serves until Ctrl-C. +pub async fn run(args: Args) -> Result<(), String> { + let base_url = args.base_url.clone(); + let model_name = args.model_name.clone(); + let (state, addr, served) = args.build()?; + + let listener = TcpListener::bind(addr) + .await + .map_err(|error| format!("failed to bind {addr}: {error}"))?; + let bound = listener.local_addr().map_err(|error| error.to_string())?; + eprintln!("{}", banner(bound, &base_url, &model_name, &served)); + + axum::serve(listener, build_router(state)) + .with_graceful_shutdown(shutdown_signal()) + .await + .map_err(|error| error.to_string()) +} + +fn env_key(name: &str) -> Option { + env::var(name) + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +fn banner(addr: SocketAddr, base_url: &str, model_name: &str, served: &[WireFormat]) -> String { + let mut formats: Vec = served.iter().map(|format| format.to_string()).collect(); + formats.sort(); + format!( + "libsy-server\n \ + listening: http://{addr}\n \ + upstream: {base_url} (model {model_name})\n \ + served fmts: {}\n \ + serving model: {SERVED_MODEL}\n \ + endpoints: GET /health, GET /v1/models, POST /v1/chat/completions, \ + POST /v1/messages, POST /v1/responses\n \ + stop: Ctrl-C", + formats.join(", ") + ) +} + +async fn shutdown_signal() { + if let Err(error) = tokio::signal::ctrl_c().await { + eprintln!("ctrl-c handler unavailable; running without shutdown trigger: {error}"); + std::future::pending::<()>().await; + } +} diff --git a/crates/libsy-server/src/lib.rs b/crates/libsy-server/src/lib.rs new file mode 100644 index 00000000..2a5184de --- /dev/null +++ b/crates/libsy-server/src/lib.rs @@ -0,0 +1,236 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Minimal axum server exposing the LLM wire APIs over [`TranslatingLlmClient`]. +//! +//! Each inbound request is handed to the client's raw path +//! ([`TranslatingLlmClient::call_rewrite_model_raw`]), which decodes it to +//! Switchyard's neutral IR, dispatches to the model's upstream backend, and +//! encodes the response back into the inbound wire format — buffered JSON or a +//! stream of wire events. This crate is the HTTP surface around that: routing, +//! header normalization, SSE framing, and error mapping. It mirrors +//! `switchyard-server` but swaps the profile-chain executor for a single client. +//! See [`build_router`]. + +mod sse; + +pub mod cli; + +use std::collections::BTreeMap; +use std::sync::Arc; + +use axum::extract::rejection::JsonRejection; +use axum::extract::State; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use serde_json::{json, Value}; +use switchyard_llm_client::{LlmClientError, RawResponse, TranslatingLlmClient}; +use switchyard_protocol::Context; +use switchyard_translation::WireFormat; + +use crate::sse::frame_stream; + +/// The single model name this server advertises to clients. +pub const SERVED_MODEL: &str = "switchyard"; + +/// Shared state for all endpoint handlers. +#[derive(Clone)] +pub struct ProxyState { + client: Arc, + /// Model id sent upstream and used as the client's map key. + upstream_model: Arc, +} + +impl ProxyState { + /// Builds server state around a configured client. + /// + /// `upstream_model` is the model id the client resolves to its backend and + /// sends upstream; the response is always restamped with the model the + /// caller asked for (the served name), never this id. + pub fn new(client: TranslatingLlmClient, upstream_model: impl Into>) -> Self { + Self { + client: Arc::new(client), + upstream_model: upstream_model.into(), + } + } +} + +/// Builds the server router. Public so integration tests can drive it without a socket. +pub fn build_router(state: ProxyState) -> Router { + Router::new() + .route("/v1/chat/completions", post(openai_chat_completions)) + .route("/v1/messages", post(anthropic_messages)) + .route("/v1/responses", post(openai_responses)) + .route("/v1/models", get(models)) + .route("/health", get(health)) + .fallback(not_found) + .with_state(state) +} + +async fn openai_chat_completions( + State(state): State, + headers: HeaderMap, + body: Result, JsonRejection>, +) -> Response { + dispatch(state, headers, body, WireFormat::OpenAiChat).await +} + +async fn anthropic_messages( + State(state): State, + headers: HeaderMap, + body: Result, JsonRejection>, +) -> Response { + dispatch(state, headers, body, WireFormat::AnthropicMessages).await +} + +async fn openai_responses( + State(state): State, + headers: HeaderMap, + body: Result, JsonRejection>, +) -> Response { + dispatch(state, headers, body, WireFormat::OpenAiResponses).await +} + +// Extracts and validates the JSON body, then runs the shared request path. +async fn dispatch( + state: ProxyState, + headers: HeaderMap, + body: Result, JsonRejection>, + inbound: WireFormat, +) -> Response { + let body = match body { + Ok(Json(value)) if value.is_object() => value, + Ok(_) => return invalid_body("Request body must be a JSON object"), + Err(error) => return invalid_body(format!("Request body must be valid JSON: {error}")), + }; + handle(state, headers, body, inbound).await +} + +// The client's raw path owns decode → call upstream → encode-back-to-`inbound`; +// this only forwards headers, picks the upstream model, and frames the result. +async fn handle( + state: ProxyState, + headers: HeaderMap, + body: Value, + inbound: WireFormat, +) -> Response { + // Caller headers ride along as metadata; the client drops reserved/auth + // headers and injects the backend's real credential, so the sentinel key a + // coding agent sends is filtered automatically. + let http_headers = Some(normalized_headers(&headers)); + let response = state + .client + .call_rewrite_model_raw( + Context::default(), + body, + http_headers, + Some(&state.upstream_model), + inbound, + ) + .await; + + match response { + Ok(RawResponse::Buffered(body)) => Json(body).into_response(), + Ok(RawResponse::Stream(events)) => frame_stream(events, inbound).into_response(), + Err(error) => map_client_error(error), + } +} + +// Lowercases header names and keeps the first UTF-8 value for each. +fn normalized_headers(headers: &HeaderMap) -> BTreeMap { + let mut normalized = BTreeMap::new(); + for (name, value) in headers { + let Ok(value) = value.to_str() else { + continue; + }; + normalized + .entry(name.as_str().to_ascii_lowercase()) + .or_insert_with(|| value.to_string()); + } + normalized +} + +// Maps client errors onto OpenAI-style HTTP error envelopes. +fn map_client_error(error: LlmClientError) -> Response { + match error { + LlmClientError::UnknownModel(model) => error_body( + StatusCode::NOT_FOUND, + format!("no backend configured for model {model:?}"), + "model_not_found", + ), + // The inbound wire format has no configured upstream backend — a client + // problem (they hit an endpoint this proxy is not serving), not a fault. + LlmClientError::UnknownModelFormat { format, .. } => error_body( + StatusCode::BAD_REQUEST, + format!("no upstream backend configured for {format} requests"), + "unsupported_format", + ), + LlmClientError::MissingModel | LlmClientError::Translation(_) => error_body( + StatusCode::BAD_REQUEST, + error.to_string(), + "invalid_request_error", + ), + LlmClientError::ContextWindowExceeded { .. } => error_body( + StatusCode::BAD_REQUEST, + error.to_string(), + "context_window_exceeded", + ), + LlmClientError::UpstreamHttp { status, body } => error_body( + StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_GATEWAY), + body, + "upstream_error", + ), + LlmClientError::Transport(message) => { + error_body(StatusCode::BAD_GATEWAY, message, "upstream_error") + } + LlmClientError::Stream(message) => server_error(message), + } +} + +async fn models() -> Json { + Json(json!({ + "object": "list", + "data": [{ + "id": SERVED_MODEL, + "object": "model", + "created": 0, + "owned_by": "switchyard", + "capabilities": { + "streaming": true, + "supported_inbound_formats": [ + "openai-chat-completions", + "openai-responses", + "anthropic-messages", + ], + }, + }], + })) +} + +async fn health() -> Json { + Json(json!({"status": "ok"})) +} + +async fn not_found() -> Response { + error_body(StatusCode::NOT_FOUND, "Not Found".to_string(), "not_found") +} + +fn invalid_body(message: impl Into) -> Response { + error_body(StatusCode::BAD_REQUEST, message.into(), "invalid_body") +} + +fn server_error(message: String) -> Response { + error_body(StatusCode::INTERNAL_SERVER_ERROR, message, "server_error") +} + +fn error_body(status: StatusCode, message: String, code: &str) -> Response { + ( + status, + Json(json!({ + "error": {"message": message, "type": code, "code": code} + })), + ) + .into_response() +} diff --git a/crates/libsy-server/src/main.rs b/crates/libsy-server/src/main.rs new file mode 100644 index 00000000..df3386a7 --- /dev/null +++ b/crates/libsy-server/src/main.rs @@ -0,0 +1,20 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Binary entrypoint for `libsy-server`. + +use std::process::ExitCode; + +use clap::Parser; +use libsy_server::cli::{run, Args}; + +#[tokio::main(flavor = "multi_thread")] +async fn main() -> ExitCode { + match run(Args::parse()).await { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("{error}"); + ExitCode::FAILURE + } + } +} diff --git a/crates/libsy-server/src/sse.rs b/crates/libsy-server/src/sse.rs new file mode 100644 index 00000000..9e9117fe --- /dev/null +++ b/crates/libsy-server/src/sse.rs @@ -0,0 +1,87 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! SSE framing for the OpenAI Chat, Anthropic Messages, and Responses formats. +//! +//! Ported from `switchyard-server/src/sse.rs`; that copy is crate-private and +//! typed against `SwitchyardError`, so this re-types the same framing against a +//! boxed error — the exact item type carried by +//! [`switchyard_llm_client::RawEventStream`], so the client's streamed wire +//! events frame without an adapter. Contract: OpenAI Chat emits `data: {json}` +//! frames plus a terminal `data: [DONE]`; Anthropic and Responses emit named +//! `event: \ndata: {json}` frames with no `[DONE]` (the stream ends on +//! `message_stop` / `response.completed`). + +use std::pin::Pin; + +use axum::response::sse::{Event, Sse}; +use futures_util::{Stream, StreamExt}; +use serde_json::{json, Value}; +use switchyard_translation::WireFormat; + +/// Boxed, thread-safe error carried by a framed SSE item. +pub(crate) type BoxError = Box; + +/// Boxed stream type accepted by Axum's SSE response wrapper. +pub(crate) type SseFrameStream = Pin> + Send>>; + +/// Frames a stream of target-format wire events as endpoint-specific SSE. +pub(crate) fn frame_stream( + stream: impl Stream> + Send + 'static, + target_format: WireFormat, +) -> Sse { + let framed = async_stream::stream! { + let mut stream = Box::pin(stream); + while let Some(item) = stream.next().await { + match item { + Ok(value) => yield frame_event(target_format, value), + Err(error) => { + // The 200 + headers are already committed, so a mid-stream + // failure surfaces as a final format-specific error frame. + yield Ok(error_event(target_format, error.to_string())); + return; + } + } + } + + if target_format == WireFormat::OpenAiChat { + yield Ok(Event::default().data("[DONE]")); + } + }; + + Sse::new(Box::pin(framed) as SseFrameStream) +} + +// Frames one wire event `Value` for the target format. +fn frame_event(target_format: WireFormat, value: Value) -> Result { + match target_format { + WireFormat::OpenAiChat => Event::default() + .json_data(value) + .map_err(|error| Box::new(error) as BoxError), + WireFormat::AnthropicMessages | WireFormat::OpenAiResponses => { + let event_type = value + .get("type") + .and_then(Value::as_str) + .unwrap_or("message") + .to_string(); + Event::default() + .event(event_type) + .json_data(value) + .map_err(|error| Box::new(error) as BoxError) + } + } +} + +// Builds a terminal error frame in the target format's shape. +fn error_event(target_format: WireFormat, message: String) -> Event { + match target_format { + WireFormat::OpenAiChat => Event::default() + .data(json!({"error": {"message": message, "type": "proxy_error"}}).to_string()), + WireFormat::AnthropicMessages | WireFormat::OpenAiResponses => { + Event::default().event("error").data( + json!({"type": "error", "error": {"message": message, "type": "proxy_error"}}) + .to_string(), + ) + } + } +} diff --git a/crates/libsy-server/tests/server.rs b/crates/libsy-server/tests/server.rs new file mode 100644 index 00000000..89b59ed5 --- /dev/null +++ b/crates/libsy-server/tests/server.rs @@ -0,0 +1,322 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! End-to-end router tests: drive `build_router` with `tower::oneshot` against a +//! `wiremock` upstream, covering the buffered and streaming paths per served wire +//! format, error mapping, and discovery endpoints. +//! +//! The client's raw path is same-format — each inbound format is served by a +//! backend of the same format — so a request routes to the matching upstream +//! endpoint, and an inbound format with no backend is a client error. + +use std::collections::BTreeMap; + +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use http_body_util::BodyExt; +use libsy_server::{build_router, ProxyState}; +use serde_json::{json, Value}; +use switchyard_llm_client::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; +use switchyard_translation::WireFormat; +use tower::ServiceExt; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const UPSTREAM_MODEL: &str = "upstream-model"; + +fn backend_for(format: WireFormat, base_url: &str) -> Backend { + let config = HttpBackendConfig { + base_url: base_url.to_string(), + api_key: Some("secret".to_string()), + extra_headers: BTreeMap::new(), + }; + match format { + WireFormat::OpenAiChat => Backend::OpenAiChat(config), + WireFormat::OpenAiResponses => Backend::OpenAiResponses(config), + WireFormat::AnthropicMessages => Backend::Anthropic(config), + } +} + +// Builds server state serving each of `formats` via a same-format upstream at +// `base_url`. The first format is the default backend, the rest are others; the +// raw path resolves either slot by the inbound format. +fn state(base_url: &str, formats: &[WireFormat]) -> ProxyState { + let mut backends: Vec = formats + .iter() + .map(|&format| backend_for(format, base_url)) + .collect(); + let default_backend = backends.remove(0); + let other_backends = (!backends.is_empty()).then_some(backends); + let config = ModelConfig::new(UPSTREAM_MODEL, default_backend, other_backends); + let client = TranslatingLlmClient::new(&[config]).unwrap(); + ProxyState::new(client, UPSTREAM_MODEL) +} + +async fn post(state: ProxyState, uri: &str, body: Value) -> (StatusCode, String) { + let request = Request::builder() + .method("POST") + .uri(uri) + .header("content-type", "application/json") + .body(Body::from(serde_json::to_vec(&body).unwrap())) + .unwrap(); + let response = build_router(state).oneshot(request).await.unwrap(); + let status = response.status(); + let bytes = response.into_body().collect().await.unwrap().to_bytes(); + (status, String::from_utf8(bytes.to_vec()).unwrap()) +} + +async fn mock(server: &MockServer, endpoint: &str, template: ResponseTemplate) { + Mock::given(method("POST")) + .and(path(endpoint)) + .respond_with(template) + .mount(server) + .await; +} + +fn chat_completion_body() -> Value { + json!({ + "id": "chatcmpl-1", + "model": UPSTREAM_MODEL, + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "Hi there"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3} + }) +} + +fn anthropic_message_body() -> Value { + json!({ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": UPSTREAM_MODEL, + "content": [{"type": "text", "text": "Hi there"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 2} + }) +} + +#[tokio::test] +async fn buffered_chat_round_trips() { + let server = MockServer::start().await; + mock( + &server, + "/v1/chat/completions", + ResponseTemplate::new(200).set_body_json(chat_completion_body()), + ) + .await; + + let (status, body) = post( + state(&format!("{}/v1", server.uri()), &[WireFormat::OpenAiChat]), + "/v1/chat/completions", + json!({"model": "switchyard", "messages": [{"role": "user", "content": "hi"}]}), + ) + .await; + + assert_eq!(status, StatusCode::OK); + let value: Value = serde_json::from_str(&body).unwrap(); + assert_eq!(value["choices"][0]["message"]["content"], "Hi there"); + // The response echoes the requested model, not the upstream id. + assert_eq!(value["model"], "switchyard"); +} + +#[tokio::test] +async fn buffered_messages_round_trips() { + let server = MockServer::start().await; + mock( + &server, + "/v1/messages", + ResponseTemplate::new(200).set_body_json(anthropic_message_body()), + ) + .await; + + let (status, body) = post( + state( + &format!("{}/v1", server.uri()), + &[WireFormat::AnthropicMessages], + ), + "/v1/messages", + json!({ + "model": "switchyard", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}] + }), + ) + .await; + + assert_eq!(status, StatusCode::OK); + let value: Value = serde_json::from_str(&body).unwrap(); + assert_eq!(value["type"], "message"); + assert_eq!(value["content"][0]["text"], "Hi there"); + assert_eq!(value["model"], "switchyard"); +} + +#[tokio::test] +async fn streaming_chat_ends_with_done() { + let server = MockServer::start().await; + let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n\ + data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n\ + data: [DONE]\n\n"; + mock( + &server, + "/v1/chat/completions", + ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"), + ) + .await; + + let (status, body) = post( + state(&format!("{}/v1", server.uri()), &[WireFormat::OpenAiChat]), + "/v1/chat/completions", + json!({ + "model": "switchyard", + "stream": true, + "messages": [{"role": "user", "content": "hi"}] + }), + ) + .await; + + assert_eq!(status, StatusCode::OK); + assert!( + body.contains("[DONE]"), + "openai stream must send [DONE]: {body}" + ); +} + +#[tokio::test] +async fn streaming_messages_ends_with_message_stop() { + let server = MockServer::start().await; + // A minimal Anthropic upstream stream: start, one text delta, then stop. + let sse = "event: message_start\n\ + data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_1\",\"type\":\"message\",\ + \"role\":\"assistant\",\"model\":\"upstream-model\",\"content\":[],\ + \"stop_reason\":null,\"usage\":{\"input_tokens\":1,\"output_tokens\":0}}}\n\n\ + event: content_block_delta\n\ + data: {\"type\":\"content_block_delta\",\"index\":0,\ + \"delta\":{\"type\":\"text_delta\",\"text\":\"Hi there\"}}\n\n\ + event: message_delta\n\ + data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\ + \"usage\":{\"output_tokens\":2}}\n\n\ + event: message_stop\n\ + data: {\"type\":\"message_stop\"}\n\n"; + mock( + &server, + "/v1/messages", + ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream"), + ) + .await; + + let (status, body) = post( + state( + &format!("{}/v1", server.uri()), + &[WireFormat::AnthropicMessages], + ), + "/v1/messages", + json!({ + "model": "switchyard", + "max_tokens": 64, + "stream": true, + "messages": [{"role": "user", "content": "hi"}] + }), + ) + .await; + + assert_eq!(status, StatusCode::OK); + assert!(body.contains("event: message_stop"), "body: {body}"); + assert!( + !body.contains("[DONE]"), + "anthropic stream must not send [DONE]: {body}" + ); +} + +#[tokio::test] +async fn upstream_500_passes_status_through() { + let server = MockServer::start().await; + mock( + &server, + "/v1/chat/completions", + ResponseTemplate::new(500).set_body_string("boom"), + ) + .await; + + let (status, body) = post( + state(&format!("{}/v1", server.uri()), &[WireFormat::OpenAiChat]), + "/v1/chat/completions", + json!({"model": "switchyard", "messages": [{"role": "user", "content": "hi"}]}), + ) + .await; + + assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR); + assert!(body.contains("boom"), "body: {body}"); +} + +#[tokio::test] +async fn unsupported_inbound_format_is_400() { + let server = MockServer::start().await; + // Only OpenAI Chat is served; an inbound Anthropic request has no backend. + let (status, body) = post( + state(&format!("{}/v1", server.uri()), &[WireFormat::OpenAiChat]), + "/v1/messages", + json!({"model": "switchyard", "max_tokens": 8, "messages": [{"role": "user", "content": "hi"}]}), + ) + .await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert!(body.contains("no upstream backend"), "body: {body}"); +} + +#[tokio::test] +async fn non_object_body_is_400() { + let server = MockServer::start().await; + let (status, body) = post( + state(&format!("{}/v1", server.uri()), &[WireFormat::OpenAiChat]), + "/v1/chat/completions", + json!([1, 2, 3]), + ) + .await; + + assert_eq!(status, StatusCode::BAD_REQUEST); + assert!(body.contains("JSON object"), "body: {body}"); +} + +#[tokio::test] +async fn models_lists_served_model() { + let server = MockServer::start().await; + let router = build_router(state( + &format!("{}/v1", server.uri()), + &[WireFormat::OpenAiChat], + )); + let response = router + .oneshot( + Request::builder() + .uri("/v1/models") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let bytes = response.into_body().collect().await.unwrap().to_bytes(); + let value: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(value["data"][0]["id"], "switchyard"); +} + +#[tokio::test] +async fn health_ok() { + let server = MockServer::start().await; + let router = build_router(state( + &format!("{}/v1", server.uri()), + &[WireFormat::OpenAiChat], + )); + let response = router + .oneshot( + Request::builder() + .uri("/health") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); +} From 6c835da11d4e101f376c705f7fd972bde0de8ad1 Mon Sep 17 00:00:00 2001 From: Greg Clark Date: Fri, 17 Jul 2026 09:48:41 -0400 Subject: [PATCH 04/14] feat(libsy-server): random routing proxy Signed-off-by: Greg Clark --- Cargo.lock | 1 + crates/libsy-server/Cargo.toml | 1 + crates/libsy-server/README.md | 44 ++++---- crates/libsy-server/src/cli.rs | 160 +++++++++++++++++++--------- crates/libsy-server/src/lib.rs | 135 ++++++++++++++++------- crates/libsy-server/tests/server.rs | 83 +++++++++------ 6 files changed, 281 insertions(+), 143 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index df26224a..b6d1448a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -940,6 +940,7 @@ dependencies = [ "clap", "futures-util", "http-body-util", + "libsy", "serde_json", "switchyard-llm-client", "switchyard-protocol", diff --git a/crates/libsy-server/Cargo.toml b/crates/libsy-server/Cargo.toml index 608eedfa..399e3035 100644 --- a/crates/libsy-server/Cargo.toml +++ b/crates/libsy-server/Cargo.toml @@ -12,6 +12,7 @@ repository.workspace = true rust-version.workspace = true [dependencies] +libsy = { path = "../libsy" } switchyard-protocol = { path = "../libsy-protocol" } switchyard-llm-client = { path = "../libsy-llm-client" } switchyard-translation = { path = "../switchyard-translation" } diff --git a/crates/libsy-server/README.md b/crates/libsy-server/README.md index 648765b1..2ff11d6f 100644 --- a/crates/libsy-server/README.md +++ b/crates/libsy-server/README.md @@ -1,16 +1,17 @@ # libsy-server A minimal [axum](https://github.com/tokio-rs/axum) server that exposes the OpenAI -Chat Completions, OpenAI Responses, and Anthropic Messages wire APIs and forwards -every request to a single upstream via +Chat Completions, OpenAI Responses, and Anthropic Messages wire APIs and routes +every request through a libsy [`RandomAlgo`](../libsy), which picks a weak- or +strong-tier upstream uniformly at random and serves it via [`switchyard-llm-client`](../libsy-llm-client). This crate is only the HTTP surface: routing, header normalization, SSE framing, -and error mapping. The actual work — decoding an inbound request to Switchyard's -neutral IR, calling the upstream, and encoding the response back into the same -wire format (buffered JSON or a stream of wire events) — lives in -`switchyard-llm-client`'s `TranslatingLlmClient::call_rewrite_model_raw`. It -mirrors `switchyard-server` but swaps the profile-chain executor for one client. +and error mapping. Each request is decoded to Switchyard's neutral IR, run +through the routing algorithm — which selects a target and calls the upstream — +and the response is encoded back into the same wire format (buffered JSON or a +stream of wire events). It mirrors `switchyard-server` but swaps the +profile-chain executor for a libsy algorithm. ## Endpoints @@ -25,28 +26,35 @@ mirrors `switchyard-server` but swaps the profile-chain executor for one client. Clients address the server's served model id, `switchyard`; the response is always restamped with that id rather than the real upstream model. -## Upstream and served formats +## Tiers, formats, and translation -Serving is same-format: each inbound endpoint is forwarded to an upstream of the -same wire format. A backend is configured for every format whose provider key is -set — OpenAI Chat + Responses from `OPENAI_API_KEY`, Anthropic Messages from -`ANTHROPIC_API_KEY` — and the server serves exactly those endpoints. An inbound -request for a format with no configured backend gets a `400` with an -`unsupported_format` error. +Two tiers — `weak` and `strong` — are the random-routing targets. Each has its +own upstream wire format (`--weak-format` / `--strong-format`); every inbound +request (any of the three served endpoints) is decoded to the neutral IR, +re-encoded to the chosen tier's format for the upstream call, and translated back +to the inbound format for the client. Each tier authenticates with its format's +provider key — OpenAI formats from `OPENAI_API_KEY`, Anthropic Messages from +`ANTHROPIC_API_KEY` — at the shared `--base-url`; startup errors if a tier's key +is unset. + +Wire-format values: `openai-chat`, `openai-responses`, `anthropic-messages`. ## Running ```bash -export OPENAI_API_KEY="sk-..." # and/or ANTHROPIC_API_KEY +export OPENAI_API_KEY="sk-..." # for openai-chat / openai-responses tiers +export ANTHROPIC_API_KEY="sk-..." # for anthropic-messages tiers cargo run -p libsy-server -- \ --base-url https://api.openai.com/v1 \ - --model-name gpt-4o-mini \ + --weak gpt-4o-mini --weak-format openai-chat \ + --strong gpt-4o --strong-format openai-chat \ --port 4000 ``` -Flags: `--base-url` (required), `--model-name` (required upstream id), -`--host` (default `127.0.0.1`), `--port` (default `4000`). +Flags: `--base-url` (required), `--weak` / `--strong` (required tier model ids), +`--weak-format` / `--strong-format` (required tier wire formats), `--host` +(default `127.0.0.1`), `--port` (default `4000`). ## Testing diff --git a/crates/libsy-server/src/cli.rs b/crates/libsy-server/src/cli.rs index a602eb28..05a0798f 100644 --- a/crates/libsy-server/src/cli.rs +++ b/crates/libsy-server/src/cli.rs @@ -6,8 +6,10 @@ use std::collections::BTreeMap; use std::env; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; use clap::Parser; +use libsy::{Algorithm, LlmTarget, LlmTargetSet, RandomAlgo, RoutedLlmClient}; use switchyard_llm_client::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; use switchyard_translation::WireFormat; use tokio::net::TcpListener; @@ -31,9 +33,27 @@ pub struct Args { #[arg(long)] pub base_url: String, - /// Upstream model id sent on every call. + /// Weak-tier upstream model id — one of the two random-routing targets. #[arg(long)] - pub model_name: String, + pub weak: String, + + /// Wire format the weak tier's upstream speaks (openai-chat, openai-responses, + /// anthropic-messages). Its provider key selects the credential. + #[arg(long, value_parser = parse_wire_format)] + pub weak_format: WireFormat, + + /// Strong-tier upstream model id — one of the two random-routing targets. + #[arg(long)] + pub strong: String, + + /// Wire format the strong tier's upstream speaks (openai-chat, openai-responses, + /// anthropic-messages). Its provider key selects the credential. + #[arg(long, value_parser = parse_wire_format)] + pub strong_format: WireFormat, + + /// Log each request's routing decision (the selected tier) to stderr. + #[arg(long)] + pub log_routing: bool, /// Host address to bind. #[arg(long, default_value_t = DEFAULT_HOST)] @@ -47,62 +67,52 @@ pub struct Args { impl Args { /// Builds server state (and the bind address) from the args and env keys. /// - /// The client's raw path is same-format: each inbound format is served by a - /// backend of the *same* format. A backend is configured for each format whose - /// provider key is set — OpenAI Chat + Responses from `OPENAI_API_KEY`, - /// Anthropic from `ANTHROPIC_API_KEY` — so the server serves exactly those - /// inbound endpoints. Errors if neither key is present. - fn build(self) -> Result<(ProxyState, SocketAddr, Vec), String> { - let openai_key = env_key(OPENAI_API_KEY_ENV); - let anthropic_key = env_key(ANTHROPIC_API_KEY_ENV); - - let mut backends: Vec = Vec::new(); - if let Some(key) = &openai_key { - backends.push(Backend::OpenAiChat(self.config(key))); - backends.push(Backend::OpenAiResponses(self.config(key))); - } - if let Some(key) = &anthropic_key { - backends.push(Backend::Anthropic(self.config(key))); - } - if backends.is_empty() { - return Err(format!( - "no upstream credentials: set {OPENAI_API_KEY_ENV} and/or {ANTHROPIC_API_KEY_ENV}" - )); - } - - let served: Vec = backends.iter().map(Backend::wire_format).collect(); - // Which backend is "default" vs "other" does not affect routing: the raw - // path resolves a backend by the inbound format, matching either slot. - let default_backend = backends.remove(0); - let other_backends = (!backends.is_empty()).then_some(backends); - - let config = ModelConfig::new(self.model_name.clone(), default_backend, other_backends); - let client = TranslatingLlmClient::new(&[config]).map_err(|error| error.to_string())?; - let addr = SocketAddr::new(self.host, self.port); - let state = ProxyState::new(client, self.model_name.as_str()); - Ok((state, addr, served)) - } + /// Each tier (`weak`, `strong`) is served by a backend of its own + /// `--*-format`, at the shared `--base-url`, authenticated with that format's + /// provider key — OpenAI formats from `OPENAI_API_KEY`, Anthropic from + /// `ANTHROPIC_API_KEY`. The tiers become the [`RandomAlgo`]'s routing targets: + /// each request picks one uniformly at random. Any inbound format is decoded + /// to the neutral IR, re-encoded to the chosen tier's format for the upstream + /// call, and translated back to the inbound format for the client. Errors if a + /// tier's provider key is unset. + fn build(self) -> Result<(ProxyState, SocketAddr), String> { + let weak_backend = backend_for_format(self.weak_format, &self.base_url)?; + let strong_backend = backend_for_format(self.strong_format, &self.base_url)?; + + // One config per tier — its own format's backend, keyed by the tier's + // model id. The client resolves a target's name to that config. + let client = Arc::new( + TranslatingLlmClient::new(&[ + ModelConfig::new(self.weak.clone(), weak_backend, None), + ModelConfig::new(self.strong.clone(), strong_backend, None), + ]) + .map_err(|error| error.to_string())?, + ); + + // The two tiers as routing targets, sharing the client that serves them. + let targets = LlmTargetSet::new(vec![ + target(&self.weak, client.clone()), + target(&self.strong, client.clone()), + ]); + let algorithm: Arc = Arc::new(RandomAlgo::new(targets)); - fn config(&self, api_key: &str) -> HttpBackendConfig { - HttpBackendConfig { - base_url: self.base_url.clone(), - api_key: Some(api_key.to_string()), - extra_headers: BTreeMap::new(), - } + let addr = SocketAddr::new(self.host, self.port); + Ok((ProxyState::new(algorithm, self.log_routing), addr)) } } /// Parses args, binds the listener, and serves until Ctrl-C. pub async fn run(args: Args) -> Result<(), String> { let base_url = args.base_url.clone(); - let model_name = args.model_name.clone(); - let (state, addr, served) = args.build()?; + let weak = (args.weak.clone(), args.weak_format); + let strong = (args.strong.clone(), args.strong_format); + let (state, addr) = args.build()?; let listener = TcpListener::bind(addr) .await .map_err(|error| format!("failed to bind {addr}: {error}"))?; let bound = listener.local_addr().map_err(|error| error.to_string())?; - eprintln!("{}", banner(bound, &base_url, &model_name, &served)); + eprintln!("{}", banner(bound, &base_url, &weak, &strong)); axum::serve(listener, build_router(state)) .with_graceful_shutdown(shutdown_signal()) @@ -110,6 +120,49 @@ pub async fn run(args: Args) -> Result<(), String> { .map_err(|error| error.to_string()) } +// Parses a `--*-format` flag into a built-in wire format, accepting the canonical +// ids plus common hyphenated aliases. +fn parse_wire_format(value: &str) -> Result { + match value.to_ascii_lowercase().replace('-', "_").as_str() { + "openai_chat" | "openai_chat_completions" | "chat" => Ok(WireFormat::OpenAiChat), + "openai_responses" | "responses" => Ok(WireFormat::OpenAiResponses), + "anthropic_messages" | "anthropic" | "messages" => Ok(WireFormat::AnthropicMessages), + _ => Err(format!( + "unknown wire format '{value}': expected one of \ + openai-chat, openai-responses, anthropic-messages" + )), + } +} + +// Builds the backend serving `format` at `base_url`, drawing the API key from the +// format's provider env var. Errors if that key is unset. +fn backend_for_format(format: WireFormat, base_url: &str) -> Result { + let env_name = match format { + WireFormat::OpenAiChat | WireFormat::OpenAiResponses => OPENAI_API_KEY_ENV, + WireFormat::AnthropicMessages => ANTHROPIC_API_KEY_ENV, + }; + let api_key = env_key(env_name) + .ok_or_else(|| format!("no credential for {format} tier: set {env_name}"))?; + let config = HttpBackendConfig { + base_url: base_url.to_string(), + api_key: Some(api_key), + extra_headers: BTreeMap::new(), + }; + Ok(match format { + WireFormat::OpenAiChat => Backend::OpenAiChat(config), + WireFormat::OpenAiResponses => Backend::OpenAiResponses(config), + WireFormat::AnthropicMessages => Backend::Anthropic(config), + }) +} + +// Builds a routing target named `model` that serves its calls through `client`. +fn target(model: &str, client: Arc) -> LlmTarget { + LlmTarget { + semantic_name: model.to_string(), + llm_client: Some(client as Arc), + } +} + fn env_key(name: &str) -> Option { env::var(name) .ok() @@ -117,19 +170,22 @@ fn env_key(name: &str) -> Option { .filter(|value| !value.is_empty()) } -fn banner(addr: SocketAddr, base_url: &str, model_name: &str, served: &[WireFormat]) -> String { - let mut formats: Vec = served.iter().map(|format| format.to_string()).collect(); - formats.sort(); +fn banner( + addr: SocketAddr, + base_url: &str, + weak: &(String, WireFormat), + strong: &(String, WireFormat), +) -> String { format!( "libsy-server\n \ listening: http://{addr}\n \ - upstream: {base_url} (model {model_name})\n \ - served fmts: {}\n \ + upstream: {base_url}\n \ + random routing: weak {} ({}), strong {} ({})\n \ serving model: {SERVED_MODEL}\n \ endpoints: GET /health, GET /v1/models, POST /v1/chat/completions, \ POST /v1/messages, POST /v1/responses\n \ stop: Ctrl-C", - formats.join(", ") + weak.0, weak.1, strong.0, strong.1 ) } diff --git a/crates/libsy-server/src/lib.rs b/crates/libsy-server/src/lib.rs index 2a5184de..370d3bdd 100644 --- a/crates/libsy-server/src/lib.rs +++ b/crates/libsy-server/src/lib.rs @@ -1,16 +1,17 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Minimal axum server exposing the LLM wire APIs over [`TranslatingLlmClient`]. +//! Minimal axum server exposing the LLM wire APIs over a libsy [`Algorithm`]. //! -//! Each inbound request is handed to the client's raw path -//! ([`TranslatingLlmClient::call_rewrite_model_raw`]), which decodes it to -//! Switchyard's neutral IR, dispatches to the model's upstream backend, and -//! encodes the response back into the inbound wire format — buffered JSON or a -//! stream of wire events. This crate is the HTTP surface around that: routing, +//! Each inbound request is decoded to Switchyard's neutral IR and run through the +//! routing [`Algorithm`] (a [`RandomAlgo`](libsy::RandomAlgo) over weak/strong +//! targets). The algorithm picks a target and serves the call through the +//! target's client, which dispatches to the upstream backend; the neutral +//! response is then encoded back into the inbound wire format — buffered JSON or +//! a stream of wire events. This crate is the HTTP surface around that: routing, //! header normalization, SSE framing, and error mapping. It mirrors -//! `switchyard-server` but swaps the profile-chain executor for a single client. -//! See [`build_router`]. +//! `switchyard-server` but swaps the profile-chain executor for a libsy +//! algorithm. See [`build_router`]. mod sse; @@ -25,10 +26,11 @@ use axum::http::{HeaderMap, StatusCode}; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::{Json, Router}; +use libsy::Algorithm; use serde_json::{json, Value}; -use switchyard_llm_client::{LlmClientError, RawResponse, TranslatingLlmClient}; -use switchyard_protocol::Context; -use switchyard_translation::WireFormat; +use switchyard_llm_client::LlmClientError; +use switchyard_protocol::{Context, LlmResponse, Metadata, Request}; +use switchyard_translation::{decode_request, encode_buffered_response, encode_stream, WireFormat}; use crate::sse::frame_stream; @@ -38,21 +40,24 @@ pub const SERVED_MODEL: &str = "switchyard"; /// Shared state for all endpoint handlers. #[derive(Clone)] pub struct ProxyState { - client: Arc, - /// Model id sent upstream and used as the client's map key. - upstream_model: Arc, + /// The routing algorithm run once per request; it picks a target and serves + /// the call through the target's client. + algorithm: Arc, + /// When set, log each request's routing decision (the selected tier) to stderr. + log_routing: bool, } impl ProxyState { - /// Builds server state around a configured client. + /// Builds server state around a routing algorithm. /// - /// `upstream_model` is the model id the client resolves to its backend and - /// sends upstream; the response is always restamped with the model the - /// caller asked for (the served name), never this id. - pub fn new(client: TranslatingLlmClient, upstream_model: impl Into>) -> Self { + /// Each request is decoded to the neutral IR and handed to `algorithm`, whose + /// chosen target resolves its own upstream model id. The response is always + /// restamped with the model the caller asked for, never the upstream id. With + /// `log_routing`, each request's chosen tier is logged to stderr. + pub fn new(algorithm: Arc, log_routing: bool) -> Self { Self { - client: Arc::new(client), - upstream_model: upstream_model.into(), + algorithm, + log_routing, } } } @@ -108,33 +113,83 @@ async fn dispatch( handle(state, headers, body, inbound).await } -// The client's raw path owns decode → call upstream → encode-back-to-`inbound`; -// this only forwards headers, picks the upstream model, and frames the result. +// Decodes the inbound body to the neutral IR, runs the routing algorithm, then +// encodes its response back into `inbound`. The chosen target owns the upstream +// call; this frames the neutral result and restamps the caller's model. async fn handle( state: ProxyState, headers: HeaderMap, body: Value, inbound: WireFormat, ) -> Response { + let llm_request = match decode_request(inbound, &body) { + Ok(request) => request, + Err(error) => return invalid_body(format!("Request body could not be decoded: {error}")), + }; + // The model the caller asked for; restamped onto the response so it never + // leaks the upstream id the algorithm's target resolved. + let requested_model = llm_request.model.clone(); + // Caller headers ride along as metadata; the client drops reserved/auth // headers and injects the backend's real credential, so the sentinel key a - // coding agent sends is filtered automatically. - let http_headers = Some(normalized_headers(&headers)); - let response = state - .client - .call_rewrite_model_raw( - Context::default(), - body, - http_headers, - Some(&state.upstream_model), - inbound, - ) - .await; - - match response { - Ok(RawResponse::Buffered(body)) => Json(body).into_response(), - Ok(RawResponse::Stream(events)) => frame_stream(events, inbound).into_response(), - Err(error) => map_client_error(error), + // coding agent sends is filtered automatically. `wire_format` is left unset so + // the upstream call uses the chosen tier's own format, not the inbound one — + // the response is translated back to `inbound` below. + let request = Request { + llm_request, + raw_request: Some(body), + metadata: Some(Metadata { + session_id: None, + agent_id: None, + task_id: None, + correlation_id: None, + extra_metadata: None, + http_headers: Some(normalized_headers(&headers)), + wire_format: None, + }), + }; + + let (trace, response) = match state + .algorithm + .clone() + .run(Context::default(), request) + .await + { + Ok(result) => result, + Err(error) => return map_run_error(error), + }; + + // The trace's first decision is the routing choice; log which tier served. + if state.log_routing { + if let Some(decision) = trace.first() { + eprintln!("[route] inbound={inbound} -> {}", decision.selected_model()); + } + } + + match response.llm_response { + LlmResponse::Agg(agg) => { + match encode_buffered_response(&agg, inbound, requested_model.as_deref()) { + Ok(body) => Json(body).into_response(), + Err(error) => error_body( + StatusCode::BAD_REQUEST, + error.to_string(), + "invalid_request_error", + ), + } + } + LlmResponse::Stream(chunks) => { + frame_stream(encode_stream(chunks, inbound, requested_model), inbound).into_response() + } + } +} + +// Maps an algorithm run failure onto an HTTP envelope. A failed upstream model +// call surfaces as a boxed [`LlmClientError`] (mapped like the direct path); +// anything else is an internal proxy fault. +fn map_run_error(error: Box) -> Response { + match error.downcast::() { + Ok(client_error) => map_client_error(*client_error), + Err(other) => server_error(other.to_string()), } } diff --git a/crates/libsy-server/tests/server.rs b/crates/libsy-server/tests/server.rs index 89b59ed5..ee367717 100644 --- a/crates/libsy-server/tests/server.rs +++ b/crates/libsy-server/tests/server.rs @@ -2,18 +2,20 @@ // SPDX-License-Identifier: Apache-2.0 //! End-to-end router tests: drive `build_router` with `tower::oneshot` against a -//! `wiremock` upstream, covering the buffered and streaming paths per served wire -//! format, error mapping, and discovery endpoints. +//! `wiremock` upstream, covering the buffered and streaming paths per wire format, +//! cross-format routing, error mapping, and discovery endpoints. //! -//! The client's raw path is same-format — each inbound format is served by a -//! backend of the same format — so a request routes to the matching upstream -//! endpoint, and an inbound format with no backend is a client error. +//! Each request is decoded to the neutral IR, routed by the `RandomAlgo` to a +//! tier, translated to that tier's upstream format for the call, and translated +//! back to the inbound format for the client. use std::collections::BTreeMap; +use std::sync::Arc; use axum::body::Body; use axum::http::{Request, StatusCode}; use http_body_util::BodyExt; +use libsy::{Algorithm, LlmTarget, LlmTargetSet, RandomAlgo, RoutedLlmClient}; use libsy_server::{build_router, ProxyState}; use serde_json::{json, Value}; use switchyard_llm_client::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; @@ -37,19 +39,19 @@ fn backend_for(format: WireFormat, base_url: &str) -> Backend { } } -// Builds server state serving each of `formats` via a same-format upstream at -// `base_url`. The first format is the default backend, the rest are others; the -// raw path resolves either slot by the inbound format. -fn state(base_url: &str, formats: &[WireFormat]) -> ProxyState { - let mut backends: Vec = formats - .iter() - .map(|&format| backend_for(format, base_url)) - .collect(); - let default_backend = backends.remove(0); - let other_backends = (!backends.is_empty()).then_some(backends); - let config = ModelConfig::new(UPSTREAM_MODEL, default_backend, other_backends); - let client = TranslatingLlmClient::new(&[config]).unwrap(); - ProxyState::new(client, UPSTREAM_MODEL) +// Builds server state whose single random-routing target serves `UPSTREAM_MODEL` +// over an upstream speaking `upstream_format` at `base_url`. Inbound requests are +// translated to `upstream_format` for the upstream call; when it matches the +// inbound format this is a same-format round trip. +fn state(base_url: &str, upstream_format: WireFormat) -> ProxyState { + let config = ModelConfig::new(UPSTREAM_MODEL, backend_for(upstream_format, base_url), None); + let client = Arc::new(TranslatingLlmClient::new(&[config]).unwrap()); + let targets = LlmTargetSet::new(vec![LlmTarget { + semantic_name: UPSTREAM_MODEL.to_string(), + llm_client: Some(client as Arc), + }]); + let algorithm: Arc = Arc::new(RandomAlgo::new(targets)); + ProxyState::new(algorithm, false) } async fn post(state: ProxyState, uri: &str, body: Value) -> (StatusCode, String) { @@ -109,7 +111,7 @@ async fn buffered_chat_round_trips() { .await; let (status, body) = post( - state(&format!("{}/v1", server.uri()), &[WireFormat::OpenAiChat]), + state(&format!("{}/v1", server.uri()), WireFormat::OpenAiChat), "/v1/chat/completions", json!({"model": "switchyard", "messages": [{"role": "user", "content": "hi"}]}), ) @@ -135,7 +137,7 @@ async fn buffered_messages_round_trips() { let (status, body) = post( state( &format!("{}/v1", server.uri()), - &[WireFormat::AnthropicMessages], + WireFormat::AnthropicMessages, ), "/v1/messages", json!({ @@ -167,7 +169,7 @@ async fn streaming_chat_ends_with_done() { .await; let (status, body) = post( - state(&format!("{}/v1", server.uri()), &[WireFormat::OpenAiChat]), + state(&format!("{}/v1", server.uri()), WireFormat::OpenAiChat), "/v1/chat/completions", json!({ "model": "switchyard", @@ -210,7 +212,7 @@ async fn streaming_messages_ends_with_message_stop() { let (status, body) = post( state( &format!("{}/v1", server.uri()), - &[WireFormat::AnthropicMessages], + WireFormat::AnthropicMessages, ), "/v1/messages", json!({ @@ -241,7 +243,7 @@ async fn upstream_500_passes_status_through() { .await; let (status, body) = post( - state(&format!("{}/v1", server.uri()), &[WireFormat::OpenAiChat]), + state(&format!("{}/v1", server.uri()), WireFormat::OpenAiChat), "/v1/chat/completions", json!({"model": "switchyard", "messages": [{"role": "user", "content": "hi"}]}), ) @@ -252,25 +254,40 @@ async fn upstream_500_passes_status_through() { } #[tokio::test] -async fn unsupported_inbound_format_is_400() { +async fn routes_inbound_openai_to_anthropic_upstream() { let server = MockServer::start().await; - // Only OpenAI Chat is served; an inbound Anthropic request has no backend. - let (status, body) = post( - state(&format!("{}/v1", server.uri()), &[WireFormat::OpenAiChat]), + // Inbound is OpenAI Chat, but the tier's upstream speaks Anthropic: the + // request is translated to Anthropic Messages, and the Anthropic response is + // translated back to the inbound OpenAI Chat shape. + mock( + &server, "/v1/messages", - json!({"model": "switchyard", "max_tokens": 8, "messages": [{"role": "user", "content": "hi"}]}), + ResponseTemplate::new(200).set_body_json(anthropic_message_body()), ) .await; - assert_eq!(status, StatusCode::BAD_REQUEST); - assert!(body.contains("no upstream backend"), "body: {body}"); + let (status, body) = post( + state( + &format!("{}/v1", server.uri()), + WireFormat::AnthropicMessages, + ), + "/v1/chat/completions", + json!({"model": "switchyard", "messages": [{"role": "user", "content": "hi"}]}), + ) + .await; + + assert_eq!(status, StatusCode::OK); + let value: Value = serde_json::from_str(&body).unwrap(); + // Client-facing shape is OpenAI Chat, restamped to the served model id. + assert_eq!(value["choices"][0]["message"]["content"], "Hi there"); + assert_eq!(value["model"], "switchyard"); } #[tokio::test] async fn non_object_body_is_400() { let server = MockServer::start().await; let (status, body) = post( - state(&format!("{}/v1", server.uri()), &[WireFormat::OpenAiChat]), + state(&format!("{}/v1", server.uri()), WireFormat::OpenAiChat), "/v1/chat/completions", json!([1, 2, 3]), ) @@ -285,7 +302,7 @@ async fn models_lists_served_model() { let server = MockServer::start().await; let router = build_router(state( &format!("{}/v1", server.uri()), - &[WireFormat::OpenAiChat], + WireFormat::OpenAiChat, )); let response = router .oneshot( @@ -307,7 +324,7 @@ async fn health_ok() { let server = MockServer::start().await; let router = build_router(state( &format!("{}/v1", server.uri()), - &[WireFormat::OpenAiChat], + WireFormat::OpenAiChat, )); let response = router .oneshot( From 90e3f57d83150e1cca894e00c5c9b6277cfcaf49 Mon Sep 17 00:00:00 2001 From: Greg Clark Date: Fri, 17 Jul 2026 09:49:29 -0400 Subject: [PATCH 05/14] fix(switchyard-tranaslation): anthropic -> responses reasoning fixes Signed-off-by: Greg Clark --- .../src/codecs/responses/buffered.rs | 38 ++++--- .../tests/request_translation.rs | 107 ++++++++++++++++++ 2 files changed, 130 insertions(+), 15 deletions(-) diff --git a/crates/switchyard-translation/src/codecs/responses/buffered.rs b/crates/switchyard-translation/src/codecs/responses/buffered.rs index e92ec45b..224e5c38 100644 --- a/crates/switchyard-translation/src/codecs/responses/buffered.rs +++ b/crates/switchyard-translation/src/codecs/responses/buffered.rs @@ -763,23 +763,32 @@ fn encode_responses_input( } let mut encoded = Vec::new(); for message in messages { - if message.content.iter().any(|block| { + // Reasoning has no valid Responses *input* representation — the API rejects + // a reasoning input item that carries content, and a cross-provider turn has + // no prior response id to reference — so drop reasoning blocks, mirroring the + // OpenAI Chat encoder's handling of Anthropic thinking. A message that was + // only reasoning contributes no input item. + let content: Vec = message + .content + .iter() + .filter(|block| !matches!(block, ContentBlock::Reasoning { .. })) + .cloned() + .collect(); + if content.is_empty() { + continue; + } + if content.iter().any(|block| { matches!( block, ContentBlock::ToolCall(_) | ContentBlock::ToolResult(_) ) }) { - encoded.extend( - message - .content - .iter() - .filter_map(encode_responses_special_input), - ); + encoded.extend(content.iter().filter_map(encode_responses_special_input)); continue; } let mut visible_content = Vec::new(); let mut emitted_special = false; - for block in &message.content { + for block in &content { if let Some(item) = encode_responses_special_input(block) { encoded.push(item); emitted_special = true; @@ -799,19 +808,18 @@ fn encode_responses_input( Ok(Value::Array(encoded)) } -// Encodes IR blocks that Responses represents as top-level input items. +// Encodes IR blocks that Responses represents as top-level input items. Reasoning +// is intentionally not handled here — it is dropped by the caller, since a +// reasoning input item carrying content is rejected by the Responses API. fn encode_responses_special_input(block: &ContentBlock) -> Option { match block { - ContentBlock::Reasoning { text, .. } => Some(json!({ - "type": "reasoning", - "content": [{"type": "reasoning_text", "text": text}], - "summary": [], - })), ContentBlock::ToolCall(call) => Some(json!({ "type": "function_call", "call_id": call.id, "name": call.name, - "arguments": call.arguments, + // Responses requires `arguments` as a JSON-encoded string, not an + // object (matching the OpenAI Chat request encoder). + "arguments": json_string(&call.arguments), })), ContentBlock::ToolResult(result) => Some(json!({ "type": "function_call_output", diff --git a/crates/switchyard-translation/tests/request_translation.rs b/crates/switchyard-translation/tests/request_translation.rs index ea15c63d..e00fdd3e 100644 --- a/crates/switchyard-translation/tests/request_translation.rs +++ b/crates/switchyard-translation/tests/request_translation.rs @@ -1101,3 +1101,110 @@ fn responses_to_chat_preserves_tool_choice_when_tools_survive() -> TestResult { ); Ok(()) } + +// Verifies an Anthropic tool_use encodes into a Responses function_call whose +// `arguments` is a JSON-encoded string, not an object. The Responses API rejects +// an object here ("expected a string"), which broke agent tool turns routed to a +// Responses backend. +#[test] +fn anthropic_tool_use_encodes_responses_arguments_as_json_string() -> TestResult { + let engine = TranslationEngine::default(); + let body = json!({ + "model": "claude-sonnet-4-20250514", + "messages": [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": [{ + "type": "tool_use", + "id": "toolu_1", + "name": "get_weather", + "input": {"city": "SF"} + }] + } + ], + "tools": [{"name": "get_weather", "input_schema": {"type": "object"}}], + "max_tokens": 64 + }); + + let output = engine + .translate_request( + WireFormat::AnthropicMessages, + WireFormat::OpenAiResponses, + &body, + &TranslationPolicy::default(), + )? + .body; + + let call = output["input"] + .as_array() + .and_then(|items| items.iter().find(|item| item["type"] == "function_call")) + .ok_or("expected a function_call input item")?; + assert_eq!(call["name"], "get_weather"); + // Must be a string, and that string must parse back to the original object. + let arguments = call["arguments"] + .as_str() + .ok_or("function_call arguments must be a JSON string")?; + assert_eq!( + serde_json::from_str::(arguments)?, + json!({"city": "SF"}) + ); + Ok(()) +} + +// Verifies an Anthropic thinking block is dropped when encoding a request to +// OpenAI Responses — never emitted as a reasoning input item carrying content, +// which the Responses API rejects ("input[..].content: expected max length 0"). +// Mirrors the OpenAI Chat encoder, which also drops thinking. +#[test] +fn anthropic_thinking_block_is_dropped_for_openai_responses_input() -> TestResult { + let engine = TranslationEngine::default(); + let body = json!({ + "model": "claude-sonnet-4-20250514", + "max_tokens": 64, + "tools": [{"name": "read_file", "input_schema": {"type": "object"}}], + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "read foo.py"}]}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "private chain of thought", "signature": "sig"}, + {"type": "text", "text": "Reading it."}, + {"type": "tool_use", "id": "tu_1", "name": "read_file", "input": {"path": "foo.py"}} + ]}, + {"role": "user", "content": [ + {"type": "tool_result", "tool_use_id": "tu_1", "content": "print(1)"} + ]} + ] + }); + + let output = engine + .translate_request( + WireFormat::AnthropicMessages, + WireFormat::OpenAiResponses, + &body, + &TranslationPolicy::default(), + )? + .body; + + let input = output["input"] + .as_array() + .ok_or("Responses input should be an array")?; + // No reasoning input item at all, and no private reasoning text leaks anywhere. + assert!( + input.iter().all(|item| item["type"] != "reasoning"), + "reasoning input item must be dropped: {input:?}" + ); + assert!( + !json_contains_content_type(&output, "reasoning_text"), + "reasoning text must not leak into Responses input" + ); + assert!( + !output.to_string().contains("private chain of thought"), + "private reasoning must not leak into the request" + ); + // The rest of the turn survives: the tool call and its output are still there. + assert!(input.iter().any(|item| item["type"] == "function_call")); + assert!(input + .iter() + .any(|item| item["type"] == "function_call_output")); + Ok(()) +} From dd87be35801cfe190000a8785278fddf311ab6c7 Mon Sep 17 00:00:00 2001 From: ayushag Date: Fri, 10 Jul 2026 01:01:48 -0700 Subject: [PATCH 06/14] feat(libsy): add agent-aware model pool routing Signed-off-by: ayushag --- Cargo.lock | 18 + Cargo.toml | 1 + crates/libsy-protocol/src/envelope.rs | 19 +- crates/libsy/Cargo.toml | 1 + crates/libsy/README.md | 32 ++ crates/libsy/src/agentic.rs | 728 ++++++++++++++++++++++++++ crates/libsy/src/lib.rs | 8 +- demo/libsy-proxy/Cargo.toml | 26 + demo/libsy-proxy/README.md | 36 ++ demo/libsy-proxy/src/main.rs | 415 +++++++++++++++ 10 files changed, 1280 insertions(+), 4 deletions(-) create mode 100644 crates/libsy/src/agentic.rs create mode 100644 demo/libsy-proxy/Cargo.toml create mode 100644 demo/libsy-proxy/README.md create mode 100644 demo/libsy-proxy/src/main.rs diff --git a/Cargo.lock b/Cargo.lock index 18749f29..5a06fd77 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -856,6 +856,7 @@ dependencies = [ "async-trait", "futures", "rand 0.8.6", + "serde", "serde_json", "switchyard-protocol", "tokio", @@ -874,6 +875,23 @@ dependencies = [ "tokio-stream", ] +[[package]] +name = "libsy-proxy" +version = "0.1.0" +dependencies = [ + "async-trait", + "axum", + "libsy", + "serde_json", + "switchyard-components", + "switchyard-components-v2", + "switchyard-core", + "switchyard-protocol", + "switchyard-server", + "switchyard-translation", + "tokio", +] + [[package]] name = "libyaml-rs" version = "0.3.0" diff --git a/Cargo.toml b/Cargo.toml index 582ea8df..a7847a15 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ members = [ "crates/switchyard-server", "crates/switchyard-skill-distillation", "crates/switchyard-translation", + "demo/libsy-proxy", ] [workspace.package] diff --git a/crates/libsy-protocol/src/envelope.rs b/crates/libsy-protocol/src/envelope.rs index 9906bd86..eef02e89 100644 --- a/crates/libsy-protocol/src/envelope.rs +++ b/crates/libsy-protocol/src/envelope.rs @@ -20,7 +20,7 @@ pub struct Context { /// All fields are optional; algorithms and observers use whichever are present /// (e.g. to key per-session state or emit correlated telemetry). `extra_metadata` /// is a free-form escape hatch for host-specific keys. -#[derive(Clone)] +#[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct Metadata { /// Stable id for a multi-request session/conversation. pub session_id: Option, @@ -28,6 +28,8 @@ pub struct Metadata { pub agent_id: Option, /// Id of the task the request belongs to. pub task_id: Option, + /// Agent-specific lineage and semantic routing signals. + pub agent_context: Option>, /// External trace/request id for joining with the host's telemetry. pub correlation_id: Option, /// Arbitrary host-defined key/value metadata. @@ -38,6 +40,21 @@ pub struct Metadata { pub wire_format: Option, } +/// Optional lineage and semantic signals for agent-aware routing. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct AgentContext { + /// Id of the parent agent, when this request comes from a child agent. + pub parent_agent_id: Option, + /// Harness-defined kind of agent call, such as `collab_spawn` or `review`. + pub agent_kind: Option, + /// Semantic agent role, such as `explorer`, `worker`, or `reviewer`. + pub agent_role: Option, + /// Semantic task class supplied by the harness or a prior classifier. + pub task_kind: Option, + /// Id of the current agent turn. + pub turn_id: Option, +} + /// A request an algorithm routes: the normalized [`LlmRequest`] plus the original /// provider payload and correlation [`Metadata`]. #[derive(Clone)] diff --git a/crates/libsy/Cargo.toml b/crates/libsy/Cargo.toml index df830383..810629cf 100644 --- a/crates/libsy/Cargo.toml +++ b/crates/libsy/Cargo.toml @@ -16,6 +16,7 @@ async-trait = "0.1" serde_json = "1" futures = "0.3" rand = "0.8" +serde = { version = "1", features = ["derive"] } switchyard-protocol = { path = "../libsy-protocol" } tokio = { version = "1", features = ["full"] } tokio-stream = "0.1" diff --git a/crates/libsy/README.md b/crates/libsy/README.md index 36f28495..463775bf 100644 --- a/crates/libsy/README.md +++ b/crates/libsy/README.md @@ -40,6 +40,38 @@ println!("answer: {}", completion_text(&response.llm_response.aggregate().await? Runnable: [`research_agent`](../libsy-examples/examples/research_agent.rs) (in the `libsy-examples` crate). +## Agent- and subtask-aware routing + +`agentic::AgentAwareOrchAlgo` classifies the first request for a stable agent/task +against an arbitrary model pool, then reuses that assignment. This keeps a child +agent and its prompt-cache locality on one model while allowing sibling agents in +the same root session to use different models. Invalid or failed classifier calls +fall back without caching the failure. + +```rust +use libsy::agentic::{AgentAwareOrchAlgo, AgentRoutingCandidate}; + +let algo: Arc = Arc::new(AgentAwareOrchAlgo::new( + "classifier", + vec![ + AgentRoutingCandidate::new("frontier", "planning, synthesis, and review"), + AgentRoutingCandidate::new("fast", "bounded research and mechanical edits"), + ], + "frontier", // fail-open target + targets, +)); +``` + +`metadata_from_headers` converts harness-specific identity into neutral `Metadata`. +Supported inputs include Codex `session-id`, `thread-id`, +`x-codex-parent-thread-id`, `x-openai-subagent`, and `x-codex-turn-metadata`; +NeMo Relay `x-nemo-relay-session-id` / `x-nemo-relay-subagent-id`; Dynamo +`x-dynamo-session-id` / `x-dynamo-parent-session-id`; and explicit +`x-switchyard-*` overrides. The adapter stays local and dependency-free because +Relay's matching gateway normalizer is not exposed through its public Rust library API. + +See [`demo/libsy-proxy`](../../demo/libsy-proxy) for a runnable Switchyard HTTP proxy. + ## Requests & responses ```rust diff --git a/crates/libsy/src/agentic.rs b/crates/libsy/src/agentic.rs new file mode 100644 index 00000000..1125b64a --- /dev/null +++ b/crates/libsy/src/agentic.rs @@ -0,0 +1,728 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Agent- and subtask-aware model routing. +//! +//! [`metadata_from_headers`] normalizes identity signals emitted by Codex, NeMo Relay, +//! Dynamo, or explicit Switchyard headers into neutral [`Metadata`]. +//! [`AgentAwareOrchAlgo`] classifies a stable agent/task once, then reuses that model +//! assignment to retain model and prompt-cache locality. + +use std::collections::{BTreeMap, HashMap}; +use std::error::Error; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use serde::Deserialize; +use switchyard_protocol::{completion_text, prompt_text, text_request, AgentContext}; + +use crate::{ + Algorithm, Context, Decision, Driver, LlmTargetSet, Metadata, Request, Response, Signals, +}; + +const CODEX_TURN_METADATA_HEADER: &str = "x-codex-turn-metadata"; +const MAX_ASSIGNMENTS: usize = 4096; + +type BoxErr = Box; + +/// One model available to the agent-aware classifier. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct AgentRoutingCandidate { + /// Semantic target name used to look up the model in [`LlmTargetSet`]. + pub target: String, + /// Capability/cost guidance shown to the classifier. + pub description: String, +} + +impl AgentRoutingCandidate { + /// Creates a model-pool entry from a target name and classifier guidance. + pub fn new(target: impl Into, description: impl Into) -> Self { + Self { + target: target.into(), + description: description.into(), + } + } +} + +/// Inspectable decision emitted for classifier and routed calls. +pub struct AgentRoutingDecision { + /// Target selected from the configured model pool. + pub selected_model: String, + /// Human-readable classifier, cache, or fallback explanation. + pub reason: String, + /// Classifier-inferred task class, when supplied. + pub task_kind: Option, + /// Classifier confidence in `[0, 1]`, when supplied. + pub confidence: Option, + /// Stable agent id used for the assignment, when available. + pub agent_id: Option, + /// Whether the selection came from an existing per-agent assignment. + pub cache_hit: bool, +} + +impl Decision for AgentRoutingDecision { + fn selected_model(&self) -> &str { + &self.selected_model + } + + fn reasoning(&self) -> Option<&str> { + Some(&self.reason) + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +/// LLM classifier over a model pool with per-agent/task assignment affinity. +pub struct AgentAwareOrchAlgo { + classifier_model: String, + candidates: Vec, + fallback_model: String, + target_set: LlmTargetSet, + assignments: Mutex>, +} + +impl AgentAwareOrchAlgo { + /// Configures the classifier target, candidate pool, fail-open target, and targets. + pub fn new( + classifier_model: impl Into, + candidates: Vec, + fallback_model: impl Into, + target_set: LlmTargetSet, + ) -> Self { + Self { + classifier_model: classifier_model.into(), + candidates, + fallback_model: fallback_model.into(), + target_set, + assignments: Mutex::new(HashMap::new()), + } + } + + /// Classifies a request, failing open without caching invalid or failed results. + async fn classify( + &self, + ctx: &Context, + driver: &Driver, + request: &Request, + ) -> Result<(Assignment, bool), BoxErr> { + let classifier_decision: Arc = Arc::new(AgentRoutingDecision { + selected_model: self.classifier_model.clone(), + reason: format!("classifying agent/subtask via {}", self.classifier_model), + task_kind: request + .metadata + .as_ref() + .and_then(|metadata| metadata.agent_context.as_deref()) + .and_then(|agent| agent.task_kind.clone()), + confidence: None, + agent_id: request + .metadata + .as_ref() + .and_then(|metadata| metadata.agent_id.clone()), + cache_hit: false, + }); + driver + .info(ctx.clone(), Arc::clone(&classifier_decision)) + .await?; + + let classify_request = Request { + llm_request: text_request( + request.llm_request.model.clone(), + classifier_prompt( + &self.candidates, + request.metadata.as_ref(), + &prompt_text(&request.llm_request), + ), + ), + // Classifier calls receive neutral text and metadata, never the provider + // payload with its tools or continuation state. + raw_request: None, + metadata: request.metadata.clone(), + }; + + let output = match self.target_set.get_target(&self.classifier_model) { + Ok(target) => { + driver + .call_llm_target(ctx.clone(), &target, classify_request, classifier_decision) + .await + } + Err(error) => Err(error), + }; + + match output { + Ok(response) => match response.llm_response.into_agg().await { + Ok(response) => { + match parse_classifier_output(&completion_text(&response), &self.candidates) { + Some(assignment) => Ok((assignment, true)), + None => Ok(( + self.fallback_assignment( + "classifier returned an invalid model-pool decision", + ), + false, + )), + } + } + Err(error) => Ok(( + self.fallback_assignment(&format!("classifier response failed: {error}")), + false, + )), + }, + Err(error) => Ok(( + self.fallback_assignment(&format!("classifier call failed: {error}")), + false, + )), + } + } + + fn fallback_assignment(&self, reason: &str) -> Assignment { + Assignment { + model: self.fallback_model.clone(), + reason: format!("{reason}; fell back to {}", self.fallback_model), + task_kind: None, + confidence: None, + } + } + + fn cached_assignment(&self, key: &AssignmentKey) -> Option { + self.assignments + .lock() + .ok() + .and_then(|assignments| assignments.get(key).cloned()) + } + + fn store_assignment(&self, key: AssignmentKey, assignment: Assignment) { + if let Ok(mut assignments) = self.assignments.lock() { + if !assignments.contains_key(&key) && assignments.len() >= MAX_ASSIGNMENTS { + if let Some(evicted) = assignments.keys().next().cloned() { + assignments.remove(&evicted); + } + } + assignments.insert(key, assignment); + } + } +} + +#[async_trait] +impl Algorithm for AgentAwareOrchAlgo { + async fn create_run_task( + self: Arc, + ctx: Context, + driver: Driver, + request: Request, + ) -> Result { + let key = assignment_key(request.metadata.as_ref()); + let cached = key.as_ref().and_then(|key| self.cached_assignment(key)); + let (assignment, cache_hit) = match cached { + Some(assignment) => (assignment, true), + None => { + let (assignment, cacheable) = self.classify(&ctx, &driver, &request).await?; + if cacheable { + if let Some(key) = key { + self.store_assignment(key, assignment.clone()); + } + } + (assignment, false) + } + }; + + let routed_target = self.target_set.get_target(&assignment.model)?; + let route_decision: Arc = Arc::new(AgentRoutingDecision { + selected_model: assignment.model.clone(), + reason: if cache_hit { + format!("reused stable agent/task assignment: {}", assignment.reason) + } else { + assignment.reason.clone() + }, + task_kind: assignment.task_kind, + confidence: assignment.confidence, + agent_id: request + .metadata + .as_ref() + .and_then(|metadata| metadata.agent_id.clone()), + cache_hit, + }); + driver + .info(ctx.clone(), Arc::clone(&route_decision)) + .await?; + driver + .call_llm_target(ctx, &routed_target, request, route_decision) + .await + } + + async fn process_signals(self: Arc, _signals: Signals) -> Result<(), BoxErr> { + Ok(()) + } +} + +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +struct AssignmentKey { + session_id: Option, + agent_id: String, + task_id: Option, +} + +#[derive(Clone, Debug)] +struct Assignment { + model: String, + reason: String, + task_kind: Option, + confidence: Option, +} + +#[derive(Deserialize)] +struct ClassifierOutput { + model: String, + #[serde(default)] + task_kind: Option, + #[serde(default)] + confidence: Option, + #[serde(default)] + reason: Option, +} + +#[derive(Default, Deserialize)] +struct CodexTurnMetadata { + session_id: Option, + thread_id: Option, + parent_thread_id: Option, + turn_id: Option, + subagent_kind: Option, + agent_role: Option, + task_id: Option, + task_kind: Option, +} + +/// Normalizes harness-specific request headers into libsy metadata. +/// +/// Explicit `x-switchyard-*` headers win. NeMo Relay and Dynamo correlation +/// headers are accepted without linking either runtime. Codex's structured turn +/// metadata is preferred over its compatibility projections. +pub fn metadata_from_headers(headers: &BTreeMap>) -> Metadata { + let codex = header(headers, CODEX_TURN_METADATA_HEADER) + .and_then(|value| serde_json::from_str::(&value).ok()) + .unwrap_or_default(); + + let agent_context = AgentContext { + parent_agent_id: first_some([ + header(headers, "x-switchyard-parent-agent-id"), + header(headers, "x-dynamo-parent-session-id"), + codex.parent_thread_id, + header(headers, "x-codex-parent-thread-id"), + ]), + agent_kind: first_some([ + header(headers, "x-switchyard-agent-kind"), + codex.subagent_kind, + header(headers, "x-openai-subagent"), + ]), + agent_role: first_some([header(headers, "x-switchyard-agent-role"), codex.agent_role]), + task_kind: first_some([header(headers, "x-switchyard-task-kind"), codex.task_kind]), + turn_id: first_some([header(headers, "x-switchyard-turn-id"), codex.turn_id]), + }; + let agent_context = (agent_context != AgentContext::default()).then(|| Box::new(agent_context)); + + Metadata { + session_id: first_some([ + header(headers, "x-switchyard-session-id"), + header(headers, "x-nemo-relay-session-id"), + codex.session_id, + header(headers, "session-id"), + ]), + agent_id: first_some([ + header(headers, "x-switchyard-agent-id"), + header(headers, "x-nemo-relay-subagent-id"), + header(headers, "x-dynamo-session-id"), + codex.thread_id, + header(headers, "thread-id"), + ]), + task_id: first_some([ + header(headers, "x-switchyard-task-id"), + codex.task_id, + header(headers, "x-task-id"), + ]), + agent_context, + correlation_id: first_some([ + header(headers, "x-request-id"), + header(headers, "x-client-request-id"), + ]), + ..Metadata::default() + } +} + +fn assignment_key(metadata: Option<&Metadata>) -> Option { + let metadata = metadata?; + Some(AssignmentKey { + session_id: metadata.session_id.clone(), + agent_id: metadata.agent_id.clone()?, + task_id: metadata.task_id.clone(), + }) +} + +fn classifier_prompt( + candidates: &[AgentRoutingCandidate], + metadata: Option<&Metadata>, + prompt: &str, +) -> String { + let metadata = metadata.cloned().unwrap_or_default(); + let agent = metadata + .agent_context + .as_deref() + .cloned() + .unwrap_or_default(); + let candidate_lines = candidates + .iter() + .map(|candidate| format!("- {}: {}", candidate.target, candidate.description)) + .collect::>() + .join("\n"); + format!( + "You route one model call inside an agent hierarchy. Select exactly one target from the \n\ + model pool. Prefer the cheapest candidate likely to complete the subtask correctly. \n\ + Return JSON only: {{\"model\":\"\",\"task_kind\":\"\",\ + \"confidence\":<0..1>,\"reason\":\"\"}}.\n\n\ + Agent context:\n\ + session_id={:?}\nagent_id={:?}\nparent_agent_id={:?}\nagent_kind={:?}\n\ + agent_role={:?}\ntask_id={:?}\ntask_kind={:?}\nturn_id={:?}\n\n\ + Model pool:\n{}\n\nSubtask:\n\n{}\n", + metadata.session_id, + metadata.agent_id, + agent.parent_agent_id, + agent.agent_kind, + agent.agent_role, + metadata.task_id, + agent.task_kind, + agent.turn_id, + candidate_lines, + prompt, + ) +} + +fn parse_classifier_output( + output: &str, + candidates: &[AgentRoutingCandidate], +) -> Option { + let trimmed = output.trim(); + if let Some(candidate) = candidates + .iter() + .find(|candidate| candidate.target == trimmed) + { + return Some(Assignment { + model: candidate.target.clone(), + reason: "classifier selected the target by name".to_string(), + task_kind: None, + confidence: None, + }); + } + + let start = trimmed.find('{')?; + let end = trimmed.rfind('}')?; + let parsed = serde_json::from_str::(&trimmed[start..=end]).ok()?; + let model = parsed.model.trim(); + if !candidates.iter().any(|candidate| candidate.target == model) { + return None; + } + Some(Assignment { + model: model.to_string(), + reason: parsed + .reason + .filter(|reason| !reason.trim().is_empty()) + .unwrap_or_else(|| "classifier selected the target".to_string()), + task_kind: parsed.task_kind.filter(|kind| !kind.trim().is_empty()), + confidence: parsed + .confidence + .filter(|value| value.is_finite()) + .map(|value| value.clamp(0.0, 1.0)), + }) +} + +fn header(headers: &BTreeMap>, name: &str) -> Option { + headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + .and_then(|(_, values)| values.iter().find(|value| !value.trim().is_empty())) + .map(|value| value.trim().to_string()) +} + +fn first_some(values: [Option; N]) -> Option { + values.into_iter().flatten().next() +} + +#[cfg(test)] +mod tests { + use std::collections::VecDeque; + + use super::*; + use crate::{LlmClient, LlmResponse, LlmTarget, RoutedRequest}; + use switchyard_protocol::{text_response, Metadata}; + + struct RoutingClient { + classifier_calls: Arc>, + responses: Arc>>, + } + + #[async_trait] + impl LlmClient for RoutingClient { + async fn call(&self, routed: RoutedRequest) -> Result { + let selected = routed.decision.selected_model().to_string(); + let completion = if selected == "classifier" { + let mut calls = self.classifier_calls.lock().map_err(|_| "lock poisoned")?; + *calls += 1; + self.responses + .lock() + .map_err(|_| "lock poisoned")? + .pop_front() + .ok_or("missing classifier response")? + } else { + selected + }; + Ok(Response { + llm_response: LlmResponse::Agg(text_response(None, completion)), + metadata: routed.request.metadata, + }) + } + } + + fn algo(responses: Vec<&str>) -> (Arc, Arc>) { + let classifier_calls = Arc::new(Mutex::new(0)); + let client = Arc::new(RoutingClient { + classifier_calls: Arc::clone(&classifier_calls), + responses: Arc::new(Mutex::new( + responses.into_iter().map(str::to_string).collect(), + )), + }) as Arc; + let target = |name: &str| LlmTarget { + semantic_name: name.to_string(), + llm_client: Some(Arc::clone(&client)), + }; + let targets = LlmTargetSet::new(vec![ + target("classifier"), + target("frontier"), + target("fast"), + ]); + let algo: Arc = Arc::new(AgentAwareOrchAlgo::new( + "classifier", + vec![ + AgentRoutingCandidate::new("frontier", "complex planning and final review"), + AgentRoutingCandidate::new("fast", "bounded research and mechanical edits"), + ], + "frontier", + targets, + )); + (algo, classifier_calls) + } + + fn request(agent_id: Option<&str>, task_id: Option<&str>, prompt: &str) -> Request { + Request { + llm_request: text_request(Some("auto".to_string()), prompt), + raw_request: Some(serde_json::json!({"input": prompt})), + metadata: Some(Metadata { + session_id: Some("session-1".to_string()), + agent_id: agent_id.map(str::to_string), + task_id: task_id.map(str::to_string), + agent_context: Some(Box::new(AgentContext { + agent_role: Some("explorer".to_string()), + ..AgentContext::default() + })), + ..Metadata::default() + }), + } + } + + fn decision(model: &str, task_kind: &str, confidence: f64, reason: &str) -> String { + serde_json::json!({ + "model": model, + "task_kind": task_kind, + "confidence": confidence, + "reason": reason, + }) + .to_string() + } + + fn response_text(response: &Response) -> String { + response + .llm_response + .as_agg() + .map(completion_text) + .unwrap_or_default() + } + + #[test] + fn normalizes_codex_turn_metadata_and_compatibility_headers() { + let mut headers = BTreeMap::new(); + headers.insert("session-id".to_string(), vec!["compat-session".to_string()]); + headers.insert("thread-id".to_string(), vec!["compat-agent".to_string()]); + headers.insert( + CODEX_TURN_METADATA_HEADER.to_string(), + vec![serde_json::json!({ + "session_id": "root-session", + "thread_id": "child-agent", + "parent_thread_id": "root-agent", + "turn_id": "turn-7", + "subagent_kind": "collab_spawn", + }) + .to_string()], + ); + + let metadata = metadata_from_headers(&headers); + assert_eq!(metadata.session_id.as_deref(), Some("root-session")); + assert_eq!(metadata.agent_id.as_deref(), Some("child-agent")); + let agent = metadata.agent_context.as_deref(); + assert_eq!( + agent.and_then(|value| value.parent_agent_id.as_deref()), + Some("root-agent") + ); + assert_eq!( + agent.and_then(|value| value.agent_kind.as_deref()), + Some("collab_spawn") + ); + assert_eq!( + agent.and_then(|value| value.turn_id.as_deref()), + Some("turn-7") + ); + } + + #[test] + fn explicit_switchyard_headers_override_relay_and_codex() { + let headers = BTreeMap::from([ + ( + "x-switchyard-session-id".to_string(), + vec!["canonical-session".to_string()], + ), + ( + "x-switchyard-agent-id".to_string(), + vec!["canonical-agent".to_string()], + ), + ( + "x-nemo-relay-session-id".to_string(), + vec!["relay-session".to_string()], + ), + ( + "x-dynamo-session-id".to_string(), + vec!["relay-agent".to_string()], + ), + ]); + + let metadata = metadata_from_headers(&headers); + assert_eq!(metadata.session_id.as_deref(), Some("canonical-session")); + assert_eq!(metadata.agent_id.as_deref(), Some("canonical-agent")); + } + + #[test] + fn normalizes_relay_and_dynamo_headers_without_runtime_dependency() { + let headers = BTreeMap::from([ + ( + "x-nemo-relay-session-id".to_string(), + vec!["relay-session".to_string()], + ), + ( + "x-nemo-relay-subagent-id".to_string(), + vec!["relay-child".to_string()], + ), + ( + "x-dynamo-parent-session-id".to_string(), + vec!["relay-parent".to_string()], + ), + ]); + + let metadata = metadata_from_headers(&headers); + assert_eq!(metadata.session_id.as_deref(), Some("relay-session")); + assert_eq!(metadata.agent_id.as_deref(), Some("relay-child")); + assert_eq!( + metadata + .agent_context + .as_deref() + .and_then(|agent| agent.parent_agent_id.as_deref()), + Some("relay-parent") + ); + } + + #[tokio::test] + async fn valid_classifier_output_routes_from_the_pool() -> Result<(), BoxErr> { + let (algo, calls) = algo(vec![&decision("fast", "research", 0.9, "bounded lookup")]); + let (trace, response) = algo + .run( + Context::default(), + request(Some("agent-a"), None, "survey the API"), + ) + .await?; + + assert_eq!(response_text(&response), "fast"); + assert_eq!(trace.len(), 2); + let routed = trace[1] + .as_any() + .downcast_ref::() + .ok_or("expected AgentRoutingDecision")?; + assert_eq!(routed.task_kind.as_deref(), Some("research")); + assert_eq!(routed.confidence, Some(0.9)); + assert!(!routed.cache_hit); + assert_eq!(*calls.lock().map_err(|_| "lock poisoned")?, 1); + Ok(()) + } + + #[tokio::test] + async fn stable_agent_assignment_avoids_reclassification() -> Result<(), BoxErr> { + let (algo, calls) = algo(vec![&decision("fast", "research", 0.8, "explorer")]); + Arc::clone(&algo) + .run( + Context::default(), + request(Some("agent-a"), None, "first turn"), + ) + .await?; + let (trace, response) = algo + .run( + Context::default(), + request(Some("agent-a"), None, "second turn"), + ) + .await?; + + assert_eq!(response_text(&response), "fast"); + assert_eq!(*calls.lock().map_err(|_| "lock poisoned")?, 1); + assert_eq!(trace.len(), 1); + let routed = trace[0] + .as_any() + .downcast_ref::() + .ok_or("expected AgentRoutingDecision")?; + assert!(routed.cache_hit); + Ok(()) + } + + #[tokio::test] + async fn explicit_task_change_reclassifies_the_same_agent() -> Result<(), BoxErr> { + let first = decision("fast", "research", 0.8, "lookup"); + let second = decision("frontier", "review", 0.9, "final review"); + let (algo, calls) = algo(vec![&first, &second]); + let (_, first_response) = Arc::clone(&algo) + .run( + Context::default(), + request(Some("agent-a"), Some("task-1"), "collect evidence"), + ) + .await?; + let (_, second_response) = algo + .run( + Context::default(), + request(Some("agent-a"), Some("task-2"), "adjudicate findings"), + ) + .await?; + + assert_eq!(response_text(&first_response), "fast"); + assert_eq!(response_text(&second_response), "frontier"); + assert_eq!(*calls.lock().map_err(|_| "lock poisoned")?, 2); + Ok(()) + } + + #[tokio::test] + async fn invalid_decision_falls_back_without_poisoning_affinity() -> Result<(), BoxErr> { + let valid = decision("fast", "research", 0.7, "retry succeeded"); + let (algo, calls) = algo(vec!["not-json", &valid]); + let (_, first) = Arc::clone(&algo) + .run(Context::default(), request(Some("agent-a"), None, "first")) + .await?; + let (_, second) = algo + .run(Context::default(), request(Some("agent-a"), None, "second")) + .await?; + + assert_eq!(response_text(&first), "frontier"); + assert_eq!(response_text(&second), "fast"); + assert_eq!(*calls.lock().map_err(|_| "lock poisoned")?, 2); + Ok(()) + } +} diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index 9d107233..0176342a 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -55,12 +55,14 @@ //! ## Algorithms //! //! [`RandomAlgo`] provides uniform random routing in the core crate. Worked -//! implementations of an LLM classifier and a stateful ensemble, plus runnable -//! agents, live in the `libsy-examples` crate. +//! [`agentic::AgentAwareOrchAlgo`] provides sub-agent-aware model-pool routing and +//! neutral harness-header normalization. Worked implementations of an LLM classifier +//! and a stateful ensemble, plus runnable agents, live in the `libsy-examples` crate. mod algorithms; pub use algorithms::noop::{NoopAlgo, NoopDecision}; pub use algorithms::rand::{RandomAlgo, RandomDecision}; +pub mod agentic; mod driver; @@ -74,7 +76,7 @@ use futures::{Stream, StreamExt}; /// [`LlmResponseChunk`] is one streaming event; [`LlmResponse`] is the streamed response /// (a live [`LlmResponseStream`] or the terminal aggregate). pub use switchyard_protocol::{ - AggLlmResponse, Context, Decision, LlmRequest, LlmResponse, LlmResponseChunk, + AgentContext, AggLlmResponse, Context, Decision, LlmRequest, LlmResponse, LlmResponseChunk, LlmResponseStream, Metadata, Request, Response, RoutedLlmClient, }; diff --git a/demo/libsy-proxy/Cargo.toml b/demo/libsy-proxy/Cargo.toml new file mode 100644 index 00000000..cb09ab40 --- /dev/null +++ b/demo/libsy-proxy/Cargo.toml @@ -0,0 +1,26 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "libsy-proxy" +version = "0.1.0" +description = "Demo LLM proxy with libsy agent-aware routing" +authors.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +rust-version.workspace = true + +[dependencies] +libsy = { path = "../../crates/libsy" } +switchyard-components = { path = "../../crates/switchyard-components" } +switchyard-components-v2 = { path = "../../crates/switchyard-components-v2" } +switchyard-core = { path = "../../crates/switchyard-core" } +switchyard-protocol = { path = "../../crates/libsy-protocol" } +switchyard-server = { path = "../../crates/switchyard-server" } +switchyard-translation = { path = "../../crates/switchyard-translation" } + +async-trait = "0.1" +axum = "0.8" +serde_json = "1" +tokio = { version = "1", features = ["macros", "rt-multi-thread", "net", "signal"] } diff --git a/demo/libsy-proxy/README.md b/demo/libsy-proxy/README.md new file mode 100644 index 00000000..0b2aa39f --- /dev/null +++ b/demo/libsy-proxy/README.md @@ -0,0 +1,36 @@ +# Agent-aware libsy proxy POC + +This demo embeds `libsy`'s agent-aware model-pool router in the PR #17 Switchyard server. It accepts +OpenAI Chat, Anthropic Messages, and OpenAI Responses request shapes, normalizes Codex/Relay lineage +headers, asks the classifier for a model assignment, and reuses valid assignments per stable +`(session, agent, explicit task)` key. + +Run it with an NVIDIA Inference Hub bearer token: + +```bash +ANTHROPIC_API_KEY="$INFERENCE_HUB_SY_API_KEY" cargo run -p libsy-proxy +``` + +Send two buffered calls for the same Codex child thread. The first call classifies; the second reuses +the assignment: + +```bash +curl http://127.0.0.1:4000/v1/chat/completions \ + -H 'content-type: application/json' \ + -H 'session-id: root-session-1' \ + -H 'thread-id: child-thread-1' \ + -H 'x-codex-parent-thread-id: root-thread-1' \ + -H 'x-openai-subagent: collab_spawn' \ + -H 'x-codex-turn-metadata: {"session_id":"root-session-1","thread_id":"child-thread-1","parent_thread_id":"root-thread-1","turn_id":"turn-1","subagent_kind":"collab_spawn"}' \ + -d '{"model":"libsy-agent-aware","messages":[{"role":"user","content":"Survey the request protocol and report evidence."}],"stream":false}' +``` + +The response exposes the chosen logical target and rationale in `x-model-router-*` headers. Explicit +`x-switchyard-session-id`, `x-switchyard-agent-id`, `x-switchyard-parent-agent-id`, +`x-switchyard-agent-role`, `x-switchyard-task-id`, and `x-switchyard-task-kind` headers override +harness-derived values. + +The demo currently requires buffered upstream responses (`"stream": false`). Header normalization +and route selection work for Responses-shaped requests, but running an unmodified streaming Codex +CLI through this demo requires a follow-up adapter between Switchyard's SSE response and libsy's +neutral response stream. diff --git a/demo/libsy-proxy/src/main.rs b/demo/libsy-proxy/src/main.rs new file mode 100644 index 00000000..2992271a --- /dev/null +++ b/demo/libsy-proxy/src/main.rs @@ -0,0 +1,415 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Demo proxy combining Switchyard HTTP/translation with libsy agent-aware routing. +//! +//! The server accepts OpenAI Chat, Anthropic Messages, and OpenAI Responses requests. +//! libsy normalizes harness identity headers, assigns a stable agent/task to a model +//! pool target, and makes calls through Switchyard's OpenAI-compatible backend. + +use std::collections::BTreeMap; +use std::error::Error; +use std::net::SocketAddr; +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::Value; + +use libsy::agentic::{ + metadata_from_headers, AgentAwareOrchAlgo, AgentRoutingCandidate, AgentRoutingDecision, +}; +use libsy::{ + Algorithm, Context, Decision, LlmClient, LlmResponse, LlmTarget, LlmTargetSet, Request, + Response, RoutedRequest, +}; +use switchyard_components::OpenAiPassthroughBackend; +use switchyard_components_v2::{Profile, ProfileInput, ProfileResponse, RoutingMetadata}; +use switchyard_core::{ + ChatRequest, ChatRequestType, ChatResponse, EndpointConfig, LlmBackend, ModelId, ProxyContext, + Result, SwitchyardError, +}; +use switchyard_server::{build_switchyard_router, ProfileRegistry, ServerState}; +use switchyard_translation::{TranslationEngine, TranslationPolicy, WireFormat}; + +const CLASSIFIER_MODEL: &str = "nvidia/deepseek-ai/deepseek-v4-flash"; +const FRONTIER_MODEL: &str = "aws/anthropic/bedrock-claude-opus-4-7"; +const FAST_MODEL: &str = "nvidia/deepseek-ai/deepseek-v4-flash"; +const CLASSIFIER_TARGET: &str = "classifier"; +const FRONTIER_TARGET: &str = "frontier"; +const FAST_TARGET: &str = "fast"; + +const DEFAULT_BASE_URL: &str = "https://inference-api.nvidia.com/v1"; +const DEFAULT_ADDR: &str = "127.0.0.1:4000"; +const PROFILE_MODEL_ID: &str = "libsy-agent-aware"; + +type BoxErr = Box; + +/// libsy client backed by Switchyard's OpenAI-compatible backend. +struct SwitchyardBackendClient { + backend: Arc, + model_ids: BTreeMap, + translation: Arc, +} + +#[async_trait] +impl LlmClient for SwitchyardBackendClient { + async fn call(&self, routed: RoutedRequest) -> std::result::Result { + let target = routed.decision.selected_model(); + let model = self + .model_ids + .get(target) + .ok_or_else(|| format!("no provider model configured for target '{target}'"))?; + let chat_request = + chat_request_for_call(&routed.request, model, self.translation.as_ref())?; + + let mut ctx = ProxyContext::new(); + let response = self + .backend + .call(&mut ctx, &chat_request) + .await + .map_err(|error| BoxErr::from(error.to_string()))?; + let body = response.body().cloned().ok_or_else(|| { + BoxErr::from("the libsy proxy demo currently requires buffered upstream responses") + })?; + let decoded = self + .translation + .decode_response(WireFormat::OpenAiChat, &body, &TranslationPolicy::default()) + .map_err(|error| BoxErr::from(error.to_string()))?; + + Ok(Response { + llm_response: LlmResponse::Agg(decoded.response), + metadata: routed.request.metadata, + }) + } +} + +/// Switchyard profile that delegates all routing decisions to libsy. +struct LibsyAgentAwareProfile { + algorithm: Arc, + translation: Arc, +} + +#[async_trait] +impl Profile for LibsyAgentAwareProfile { + async fn run(&self, input: ProfileInput) -> Result { + let request = libsy_request(&input, self.translation.as_ref())?; + let (trace, response) = Arc::clone(&self.algorithm) + .run(Context::default(), request) + .await + .map_err(|error| SwitchyardError::Other(error.to_string()))?; + let aggregate = response + .llm_response + .into_agg() + .await + .map_err(|error| SwitchyardError::Other(error.to_string()))?; + let body = self + .translation + .encode_response( + WireFormat::OpenAiChat, + &aggregate, + &TranslationPolicy::default(), + ) + .map_err(|error| SwitchyardError::Other(error.to_string()))? + .body; + + Ok(ProfileResponse::with_routing_metadata( + ChatResponse::openai_completion(body), + routing_metadata(&trace), + )) + } +} + +/// Decode the inbound wire body and attach normalized harness metadata. +fn libsy_request(input: &ProfileInput, translation: &TranslationEngine) -> Result { + let format = wire_format(input.request.request_type()); + let decoded = translation + .decode_request(format, input.request.body(), &TranslationPolicy::default()) + .map_err(|error| SwitchyardError::InvalidRequest(error.to_string()))?; + let mut metadata = metadata_from_headers(&input.metadata.headers); + metadata + .extra_metadata + .get_or_insert_with(Default::default) + .insert( + "inbound_format".to_string(), + inbound_format(input.request.request_type()).to_string(), + ); + + Ok(Request { + llm_request: decoded.request, + raw_request: Some(input.request.body().clone()), + metadata: Some(metadata), + }) +} + +/// Preserve routed provider payloads; encode classifier calls from the neutral IR. +fn chat_request_for_call( + request: &Request, + model: &str, + translation: &TranslationEngine, +) -> std::result::Result { + let Some(mut body) = request.raw_request.clone() else { + let mut classifier_request = request.llm_request.clone(); + classifier_request.model = Some(model.to_string()); + classifier_request.stream = false; + let body = translation + .encode_request( + WireFormat::OpenAiChat, + &classifier_request, + &TranslationPolicy::default(), + ) + .map_err(|error| BoxErr::from(error.to_string()))? + .body; + return Ok(ChatRequest::openai_chat(body)); + }; + + if let Some(object) = body.as_object_mut() { + object.insert("model".to_string(), Value::String(model.to_string())); + } + let request = match request + .metadata + .as_ref() + .and_then(|metadata| metadata.extra_metadata.as_ref()) + .and_then(|extra| extra.get("inbound_format")) + .map(String::as_str) + { + Some("anthropic") => ChatRequest::anthropic(body), + Some("openai_responses") => ChatRequest::openai_responses(body), + _ => ChatRequest::openai_chat(body), + }; + Ok(request) +} + +fn inbound_format(request_type: ChatRequestType) -> &'static str { + match request_type { + ChatRequestType::OpenAiChat => "openai_chat", + ChatRequestType::Anthropic => "anthropic", + ChatRequestType::OpenAiResponses => "openai_responses", + } +} + +fn wire_format(request_type: ChatRequestType) -> WireFormat { + match request_type { + ChatRequestType::OpenAiChat => WireFormat::OpenAiChat, + ChatRequestType::Anthropic => WireFormat::AnthropicMessages, + ChatRequestType::OpenAiResponses => WireFormat::OpenAiResponses, + } +} + +/// Surface libsy's routing decision as `x-model-router-*` response headers. +fn routing_metadata(trace: &[Arc]) -> RoutingMetadata { + let decision = trace.last(); + let agent = + decision.and_then(|decision| decision.as_any().downcast_ref::()); + let selected_target = decision.map(|decision| decision.selected_model()); + RoutingMetadata { + selected_model: selected_target.map(|target| match target { + FRONTIER_TARGET => FRONTIER_MODEL.to_string(), + FAST_TARGET => FAST_MODEL.to_string(), + other => other.to_string(), + }), + selected_tier: selected_target.map(str::to_string), + confidence: agent.and_then(|decision| decision.confidence), + router_version: Some("libsy-agent-aware-v1".to_string()), + tolerance: None, + rationale: decision.and_then(|decision| decision.reasoning().map(str::to_string)), + } +} + +fn build_algorithm() -> Result> { + let base_url = + std::env::var("LIBSY_PROXY_BASE_URL").unwrap_or_else(|_| DEFAULT_BASE_URL.to_string()); + let api_key = std::env::var("ANTHROPIC_API_KEY").map_err(|_| { + SwitchyardError::InvalidConfig( + "ANTHROPIC_API_KEY must be set to the upstream bearer key".to_string(), + ) + })?; + let backend = Arc::new(OpenAiPassthroughBackend::new(EndpointConfig { + base_url: Some(base_url), + api_key: Some(api_key), + timeout_secs: Some(120.0), + })?); + let model_ids = BTreeMap::from([ + (CLASSIFIER_TARGET.to_string(), CLASSIFIER_MODEL.to_string()), + (FRONTIER_TARGET.to_string(), FRONTIER_MODEL.to_string()), + (FAST_TARGET.to_string(), FAST_MODEL.to_string()), + ]); + let translation = Arc::new(TranslationEngine::default()); + let client = Arc::new(SwitchyardBackendClient { + backend, + model_ids, + translation, + }) as Arc; + let target = |name: &str| LlmTarget { + semantic_name: name.to_string(), + llm_client: Some(Arc::clone(&client)), + }; + let targets = LlmTargetSet::new(vec![ + target(CLASSIFIER_TARGET), + target(FRONTIER_TARGET), + target(FAST_TARGET), + ]); + + Ok(Arc::new(AgentAwareOrchAlgo::new( + CLASSIFIER_TARGET, + vec![ + AgentRoutingCandidate::new( + FRONTIER_TARGET, + "frontier model for planning, ambiguous implementation, synthesis, and review", + ), + AgentRoutingCandidate::new( + FAST_TARGET, + "efficient model for bounded exploration, retrieval, and mechanical edits", + ), + ], + FRONTIER_TARGET, + targets, + ))) +} + +#[tokio::main] +async fn main() -> Result<()> { + let translation = Arc::new(TranslationEngine::default()); + let profile = Arc::new(LibsyAgentAwareProfile { + algorithm: build_algorithm()?, + translation, + }) as Arc; + let registry = ProfileRegistry::from_profiles([( + ModelId::new(PROFILE_MODEL_ID)?, + profile, + PROFILE_MODEL_ID.to_string(), + )])?; + let state = ServerState::new(registry); + let addr: SocketAddr = std::env::var("LIBSY_PROXY_ADDR") + .unwrap_or_else(|_| DEFAULT_ADDR.to_string()) + .parse() + .map_err(|error: std::net::AddrParseError| { + SwitchyardError::InvalidConfig(error.to_string()) + })?; + let listener = tokio::net::TcpListener::bind(addr) + .await + .map_err(|error| SwitchyardError::Other(error.to_string()))?; + let bound_addr = listener + .local_addr() + .map_err(|error| SwitchyardError::Other(error.to_string()))?; + + println!("libsy-proxy listening on http://{bound_addr}"); + println!(" routing (libsy agent-aware): classifier={CLASSIFIER_MODEL}"); + println!(" frontier={FRONTIER_MODEL}"); + println!(" fast={FAST_MODEL}"); + println!( + " send model \"{PROFILE_MODEL_ID}\" to /v1/chat/completions, /v1/messages, or /v1/responses" + ); + + axum::serve(listener, build_switchyard_router(state)) + .with_graceful_shutdown(async { + let _ = tokio::signal::ctrl_c().await; + }) + .await + .map_err(|error| SwitchyardError::Other(error.to_string())) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use switchyard_components_v2::RequestMetadata; + use switchyard_protocol::{prompt_text, text_request, Metadata}; + + #[test] + fn routed_call_preserves_provider_body_and_rewrites_model() -> std::result::Result<(), BoxErr> { + let request = Request { + llm_request: text_request(Some("libsy-agent-aware".to_string()), "inspect"), + raw_request: Some(json!({ + "model": "libsy-agent-aware", + "input": "inspect", + "tools": [{"type": "function", "name": "shell"}], + "stream": false, + })), + metadata: Some(Metadata { + extra_metadata: Some(BTreeMap::from([( + "inbound_format".to_string(), + "openai_responses".to_string(), + )])), + ..Metadata::default() + }), + }; + + let routed = + chat_request_for_call(&request, "provider/model", &TranslationEngine::default())?; + assert_eq!(routed.request_type(), ChatRequestType::OpenAiResponses); + assert_eq!( + routed.body().get("model").and_then(Value::as_str), + Some("provider/model") + ); + assert!(routed.body().get("tools").is_some()); + Ok(()) + } + + #[test] + fn decodes_responses_input_and_normalizes_headers() -> Result<()> { + let input = ProfileInput { + request: ChatRequest::openai_responses(json!({ + "model": "libsy-agent-aware", + "input": [{ + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "current subtask"}] + }], + "stream": false, + })), + metadata: RequestMetadata { + headers: BTreeMap::from([( + "thread-id".to_string(), + vec!["child-agent".to_string()], + )]), + ..RequestMetadata::default() + }, + }; + + let request = libsy_request(&input, &TranslationEngine::default())?; + assert_eq!(prompt_text(&request.llm_request), "current subtask"); + assert_eq!( + request + .metadata + .as_ref() + .and_then(|metadata| metadata.agent_id.as_deref()), + Some("child-agent") + ); + Ok(()) + } + + #[test] + fn classifier_call_uses_a_buffered_synthetic_request() -> std::result::Result<(), BoxErr> { + let request = Request { + llm_request: text_request(Some("auto".to_string()), "classify this"), + raw_request: None, + metadata: None, + }; + + let classifier = + chat_request_for_call(&request, "classifier/model", &TranslationEngine::default())?; + assert_eq!(classifier.request_type(), ChatRequestType::OpenAiChat); + assert_ne!( + classifier.body().get("stream").and_then(Value::as_bool), + Some(true) + ); + Ok(()) + } + + #[test] + fn response_metadata_exposes_provider_model_and_logical_tier() { + let decision: Arc = Arc::new(AgentRoutingDecision { + selected_model: FAST_TARGET.to_string(), + reason: "bounded lookup".to_string(), + task_kind: Some("research".to_string()), + confidence: Some(0.9), + agent_id: Some("child-1".to_string()), + cache_hit: false, + }); + + let metadata = routing_metadata(&[decision]); + assert_eq!(metadata.selected_model.as_deref(), Some(FAST_MODEL)); + assert_eq!(metadata.selected_tier.as_deref(), Some(FAST_TARGET)); + assert_eq!(metadata.confidence, Some(0.9)); + } +} From 519797ecfb0ba368195d55e73f25eb4ebe3db959 Mon Sep 17 00:00:00 2001 From: ayushag Date: Thu, 16 Jul 2026 15:21:39 -0700 Subject: [PATCH 07/14] feat(libsy): add pluggable routing affinity --- crates/libsy-protocol/src/envelope.rs | 2 + crates/libsy/README.md | 44 +- crates/libsy/src/affinity.rs | 430 +++++++++++++++ crates/libsy/src/agentic.rs | 728 -------------------------- crates/libsy/src/algorithms/rand.rs | 236 ++++++++- crates/libsy/src/lib.rs | 10 +- crates/libsy/tests/affinity.rs | 45 ++ demo/libsy-proxy/README.md | 18 +- demo/libsy-proxy/src/main.rs | 116 ++-- 9 files changed, 786 insertions(+), 843 deletions(-) create mode 100644 crates/libsy/src/affinity.rs delete mode 100644 crates/libsy/src/agentic.rs create mode 100644 crates/libsy/tests/affinity.rs diff --git a/crates/libsy-protocol/src/envelope.rs b/crates/libsy-protocol/src/envelope.rs index eef02e89..36a39054 100644 --- a/crates/libsy-protocol/src/envelope.rs +++ b/crates/libsy-protocol/src/envelope.rs @@ -43,6 +43,8 @@ pub struct Metadata { /// Optional lineage and semantic signals for agent-aware routing. #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct AgentContext { + /// Whether the harness explicitly identified this request as coming from a child agent. + pub is_subagent: bool, /// Id of the parent agent, when this request comes from a child agent. pub parent_agent_id: Option, /// Harness-defined kind of agent call, such as `collab_spawn` or `review`. diff --git a/crates/libsy/README.md b/crates/libsy/README.md index 463775bf..6d9a52b6 100644 --- a/crates/libsy/README.md +++ b/crates/libsy/README.md @@ -40,35 +40,35 @@ println!("answer: {}", completion_text(&response.llm_response.aggregate().await? Runnable: [`research_agent`](../libsy-examples/examples/research_agent.rs) (in the `libsy-examples` crate). -## Agent- and subtask-aware routing +## Composable affinity -`agentic::AgentAwareOrchAlgo` classifies the first request for a stable agent/task -against an arbitrary model pool, then reuses that assignment. This keeps a child -agent and its prompt-cache locality on one model while allowing sibling agents in -the same root session to use different models. Invalid or failed classifier calls -fall back without caching the failure. +Routing algorithms can consume any `affinity::Affinity` policy at the point where +their final model is known. `SessionAffinity` retains one model per session; +`SubAgentAffinity` retains one model per explicitly identified child agent while +ordinary requests continue through the algorithm on every turn. ```rust -use libsy::agentic::{AgentAwareOrchAlgo, AgentRoutingCandidate}; - -let algo: Arc = Arc::new(AgentAwareOrchAlgo::new( - "classifier", - vec![ - AgentRoutingCandidate::new("frontier", "planning, synthesis, and review"), - AgentRoutingCandidate::new("fast", "bounded research and mechanical edits"), - ], - "frontier", // fail-open target - targets, -)); +use libsy::affinity::SubAgentAffinity; +use libsy::RandomAlgo; + +let algo: Arc = Arc::new( + RandomAlgo::new(targets) + .with_affinity(Arc::new(SubAgentAffinity::new())), +); ``` -`metadata_from_headers` converts harness-specific identity into neutral `Metadata`. +The first request from a child agent runs the random algorithm and atomically retains +its choice; later requests for the same `(session_id, agent_id)` reuse that target. +With one configured target, that target is retained directly. `metadata_from_headers` +converts harness-specific identity into neutral `Metadata`. Supported inputs include Codex `session-id`, `thread-id`, `x-codex-parent-thread-id`, `x-openai-subagent`, and `x-codex-turn-metadata`; NeMo Relay `x-nemo-relay-session-id` / `x-nemo-relay-subagent-id`; Dynamo `x-dynamo-session-id` / `x-dynamo-parent-session-id`; and explicit -`x-switchyard-*` overrides. The adapter stays local and dependency-free because -Relay's matching gateway normalizer is not exposed through its public Rust library API. +`x-switchyard-*` overrides, including `x-switchyard-is-subagent`. Header parsing is +separate from affinity keying, so non-HTTP embedders can populate the same metadata +directly. The adapter stays local and dependency-free because Relay's matching gateway +normalizer is not exposed through its public Rust library API. See [`demo/libsy-proxy`](../../demo/libsy-proxy) for a runnable Switchyard HTTP proxy. @@ -289,8 +289,8 @@ with `cargo test -p libsy-examples`). **Reference algorithms** — implementations to read and route with: -- [`RandomAlgo`](src/algorithms/rand.rs) — uniform random over the set - (one call). +- [`RandomAlgo`](src/algorithms/rand.rs) — uniform random over the set, optionally + retaining selections through an `Affinity` policy (one call). - [`LlmClassifierOrchAlgo`](../libsy-examples/src/llm_class.rs) — classify, then route strong/weak; fail open to strong. - [`EnsembleOrchAlgo`](../libsy-examples/src/ensemble.rs) — stateful: fan out to diff --git a/crates/libsy/src/affinity.rs b/crates/libsy/src/affinity.rs new file mode 100644 index 00000000..7cc6f92f --- /dev/null +++ b/crates/libsy/src/affinity.rs @@ -0,0 +1,430 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Reusable model-affinity policies and harness metadata normalization. +//! +//! Algorithms remain responsible for choosing a model. An [`Affinity`] policy only +//! decides whether that choice should be retained and which stable request identity +//! owns the assignment. + +use std::collections::{BTreeMap, HashMap}; +use std::sync::Mutex; + +use serde::Deserialize; +use switchyard_protocol::{AgentContext, Metadata, Request}; + +const CODEX_TURN_METADATA_HEADER: &str = "x-codex-turn-metadata"; +const MAX_ASSIGNMENTS: usize = 4096; + +/// Namespaced stable identity used to retain a model assignment. +#[derive(Clone, Debug, Hash, PartialEq, Eq)] +pub struct AffinityKey { + namespace: String, + components: Vec, +} + +impl AffinityKey { + /// Creates a key whose namespace prevents collisions with other affinity policies. + pub fn new(namespace: impl Into, components: Vec) -> Self { + Self { + namespace: namespace.into(), + components, + } + } +} + +/// Bounded, process-local storage shared by an affinity policy. +#[derive(Default)] +pub struct AffinityState { + assignments: Mutex>, +} + +impl AffinityState { + /// Returns the retained model for `key`, failing open if the state lock is poisoned. + fn get(&self, key: &AffinityKey) -> Option { + self.assignments + .lock() + .ok() + .and_then(|assignments| assignments.get(key).cloned()) + } + + /// Atomically retains the first model assigned to `key` and returns that model. + fn get_or_insert(&self, key: AffinityKey, proposed: String) -> String { + let Ok(mut assignments) = self.assignments.lock() else { + return proposed; + }; + if let Some(assigned) = assignments.get(&key) { + return assigned.clone(); + } + if assignments.len() >= MAX_ASSIGNMENTS { + if let Some(evicted) = assignments.keys().next().cloned() { + assignments.remove(&evicted); + } + } + assignments.insert(key, proposed.clone()); + proposed + } +} + +/// Policy that maps eligible requests to stable assignment keys. +/// +/// Algorithms consume this trait at the point where their final model is known: +/// check [`assignment`](Self::assignment), run normal selection on a miss, then call +/// [`retain`](Self::retain) before routing the request. +pub trait Affinity: Send + Sync { + /// Returns the stable key for an eligible request, or `None` when affinity does + /// not apply. + fn key(&self, request: &Request) -> Option; + + /// Returns this policy's process-local assignment state. + fn state(&self) -> &AffinityState; + + /// Returns an existing model assignment for `request`. + fn assignment(&self, request: &Request) -> Option { + self.key(request).and_then(|key| self.state().get(&key)) + } + + /// Retains `proposed` for an eligible request and returns the canonical model. + /// + /// When concurrent requests propose different models, the first insertion wins. + /// Ineligible requests simply receive `proposed` without storing it. + fn retain(&self, request: &Request, proposed: String) -> String { + let Some(key) = self.key(request) else { + return proposed; + }; + self.state().get_or_insert(key, proposed) + } +} + +/// Retains one model for every request carrying the same session id. +#[derive(Default)] +pub struct SessionAffinity { + state: AffinityState, +} + +impl SessionAffinity { + /// Creates an empty process-local session-affinity policy. + pub fn new() -> Self { + Self::default() + } +} + +impl Affinity for SessionAffinity { + fn key(&self, request: &Request) -> Option { + Some(AffinityKey::new( + "session", + vec![request.metadata.as_ref()?.session_id.clone()?], + )) + } + + fn state(&self) -> &AffinityState { + &self.state + } +} + +/// Retains one model for each explicitly identified child agent in a session. +#[derive(Default)] +pub struct SubAgentAffinity { + state: AffinityState, +} + +impl SubAgentAffinity { + /// Creates an empty process-local child-agent-affinity policy. + pub fn new() -> Self { + Self::default() + } +} + +impl Affinity for SubAgentAffinity { + fn key(&self, request: &Request) -> Option { + let metadata = request.metadata.as_ref()?; + if !metadata.agent_context.as_deref()?.is_subagent { + return None; + } + Some(AffinityKey::new( + "subagent", + vec![metadata.session_id.clone()?, metadata.agent_id.clone()?], + )) + } + + fn state(&self) -> &AffinityState { + &self.state + } +} + +#[derive(Default, Deserialize)] +struct CodexTurnMetadata { + session_id: Option, + thread_id: Option, + parent_thread_id: Option, + turn_id: Option, + subagent_kind: Option, + agent_role: Option, + task_id: Option, + task_kind: Option, +} + +/// Normalizes harness-specific request headers into libsy metadata. +/// +/// Explicit `x-switchyard-*` headers win. NeMo Relay and Dynamo correlation +/// headers are accepted without linking either runtime. Codex's structured turn +/// metadata is preferred over its compatibility projections. +pub fn metadata_from_headers(headers: &BTreeMap>) -> Metadata { + let codex = header(headers, CODEX_TURN_METADATA_HEADER) + .and_then(|value| serde_json::from_str::(&value).ok()) + .unwrap_or_default(); + + let switchyard_parent = header(headers, "x-switchyard-parent-agent-id"); + let relay_subagent = header(headers, "x-nemo-relay-subagent-id"); + let dynamo_parent = header(headers, "x-dynamo-parent-session-id"); + let codex_parent = header(headers, "x-codex-parent-thread-id"); + let openai_subagent = header(headers, "x-openai-subagent"); + let inferred_subagent = switchyard_parent.is_some() + || relay_subagent.is_some() + || dynamo_parent.is_some() + || codex.parent_thread_id.is_some() + || codex.subagent_kind.is_some() + || codex_parent.is_some() + || openai_subagent.is_some(); + let is_subagent = header(headers, "x-switchyard-is-subagent") + .as_deref() + .and_then(parse_bool) + .unwrap_or(inferred_subagent); + + let agent_context = AgentContext { + is_subagent, + parent_agent_id: first_some([ + switchyard_parent, + dynamo_parent, + codex.parent_thread_id, + codex_parent, + ]), + agent_kind: first_some([ + header(headers, "x-switchyard-agent-kind"), + codex.subagent_kind, + openai_subagent, + ]), + agent_role: first_some([header(headers, "x-switchyard-agent-role"), codex.agent_role]), + task_kind: first_some([header(headers, "x-switchyard-task-kind"), codex.task_kind]), + turn_id: first_some([header(headers, "x-switchyard-turn-id"), codex.turn_id]), + }; + let agent_context = (agent_context != AgentContext::default()).then(|| Box::new(agent_context)); + + Metadata { + session_id: first_some([ + header(headers, "x-switchyard-session-id"), + header(headers, "x-nemo-relay-session-id"), + codex.session_id, + header(headers, "session-id"), + ]), + agent_id: first_some([ + header(headers, "x-switchyard-agent-id"), + relay_subagent, + header(headers, "x-dynamo-session-id"), + codex.thread_id, + header(headers, "thread-id"), + ]), + task_id: first_some([ + header(headers, "x-switchyard-task-id"), + codex.task_id, + header(headers, "x-task-id"), + ]), + agent_context, + correlation_id: first_some([ + header(headers, "x-request-id"), + header(headers, "x-client-request-id"), + ]), + ..Metadata::default() + } +} + +fn header(headers: &BTreeMap>, name: &str) -> Option { + headers + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + .and_then(|(_, values)| values.iter().find(|value| !value.trim().is_empty())) + .map(|value| value.trim().to_string()) +} + +fn parse_bool(value: &str) -> Option { + match value.trim().to_ascii_lowercase().as_str() { + "1" | "true" | "yes" | "on" => Some(true), + "0" | "false" | "no" | "off" => Some(false), + _ => None, + } +} + +fn first_some(values: [Option; N]) -> Option { + values.into_iter().flatten().next() +} + +#[cfg(test)] +mod tests { + use super::*; + use switchyard_protocol::{text_request, AgentContext}; + + fn request(metadata: Metadata) -> Request { + Request { + llm_request: text_request(Some("auto".to_string()), "hi"), + raw_request: None, + metadata: Some(metadata), + } + } + + fn child_request(task_id: &str) -> Request { + request(Metadata { + session_id: Some("session-1".to_string()), + agent_id: Some("child-1".to_string()), + task_id: Some(task_id.to_string()), + agent_context: Some(Box::new(AgentContext { + is_subagent: true, + ..AgentContext::default() + })), + ..Metadata::default() + }) + } + + #[test] + fn session_affinity_keys_all_requests_by_session() { + let affinity = SessionAffinity::new(); + let first = request(Metadata { + session_id: Some("session-1".to_string()), + agent_id: Some("agent-a".to_string()), + ..Metadata::default() + }); + let second = request(Metadata { + session_id: Some("session-1".to_string()), + agent_id: Some("agent-b".to_string()), + ..Metadata::default() + }); + + assert_eq!(affinity.retain(&first, "model-a".to_string()), "model-a"); + assert_eq!(affinity.assignment(&second).as_deref(), Some("model-a")); + } + + #[test] + fn subagent_affinity_ignores_task_changes() { + let affinity = SubAgentAffinity::new(); + assert_eq!( + affinity.retain(&child_request("task-1"), "model-a".to_string()), + "model-a" + ); + assert_eq!( + affinity.assignment(&child_request("task-2")).as_deref(), + Some("model-a") + ); + } + + #[test] + fn subagent_affinity_does_not_apply_to_root_agents() { + let affinity = SubAgentAffinity::new(); + let root = request(Metadata { + session_id: Some("session-1".to_string()), + agent_id: Some("root-1".to_string()), + agent_context: Some(Box::new(AgentContext::default())), + ..Metadata::default() + }); + + assert_eq!(affinity.retain(&root, "model-a".to_string()), "model-a"); + assert!(affinity.assignment(&root).is_none()); + } + + #[test] + fn normalizes_codex_child_metadata() { + let headers = BTreeMap::from([( + CODEX_TURN_METADATA_HEADER.to_string(), + vec![serde_json::json!({ + "session_id": "root-session", + "thread_id": "child-agent", + "parent_thread_id": "root-agent", + "turn_id": "turn-7", + "subagent_kind": "collab_spawn", + }) + .to_string()], + )]); + + let metadata = metadata_from_headers(&headers); + assert_eq!(metadata.session_id.as_deref(), Some("root-session")); + assert_eq!(metadata.agent_id.as_deref(), Some("child-agent")); + let agent = metadata.agent_context.as_deref(); + assert_eq!(agent.map(|value| value.is_subagent), Some(true)); + assert_eq!( + agent.and_then(|value| value.parent_agent_id.as_deref()), + Some("root-agent") + ); + } + + #[test] + fn root_codex_metadata_is_not_inferred_as_a_subagent() { + let headers = BTreeMap::from([( + CODEX_TURN_METADATA_HEADER.to_string(), + vec![serde_json::json!({ + "session_id": "root-session", + "thread_id": "root-agent", + "turn_id": "turn-1", + }) + .to_string()], + )]); + + let metadata = metadata_from_headers(&headers); + assert_eq!( + metadata + .agent_context + .as_deref() + .map(|agent| agent.is_subagent), + Some(false) + ); + } + + #[test] + fn explicit_switchyard_subagent_flag_overrides_inference() { + let headers = BTreeMap::from([ + ( + "x-switchyard-is-subagent".to_string(), + vec!["false".to_string()], + ), + ( + "x-switchyard-parent-agent-id".to_string(), + vec!["parent".to_string()], + ), + ]); + + let metadata = metadata_from_headers(&headers); + assert_eq!( + metadata + .agent_context + .as_deref() + .map(|agent| agent.is_subagent), + Some(false) + ); + } + + #[test] + fn normalizes_relay_and_dynamo_child_headers() { + let headers = BTreeMap::from([ + ( + "x-nemo-relay-session-id".to_string(), + vec!["relay-session".to_string()], + ), + ( + "x-nemo-relay-subagent-id".to_string(), + vec!["relay-child".to_string()], + ), + ( + "x-dynamo-parent-session-id".to_string(), + vec!["relay-parent".to_string()], + ), + ]); + + let metadata = metadata_from_headers(&headers); + assert_eq!(metadata.session_id.as_deref(), Some("relay-session")); + assert_eq!(metadata.agent_id.as_deref(), Some("relay-child")); + assert_eq!( + metadata + .agent_context + .as_deref() + .map(|agent| agent.is_subagent), + Some(true) + ); + } +} diff --git a/crates/libsy/src/agentic.rs b/crates/libsy/src/agentic.rs deleted file mode 100644 index 1125b64a..00000000 --- a/crates/libsy/src/agentic.rs +++ /dev/null @@ -1,728 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Agent- and subtask-aware model routing. -//! -//! [`metadata_from_headers`] normalizes identity signals emitted by Codex, NeMo Relay, -//! Dynamo, or explicit Switchyard headers into neutral [`Metadata`]. -//! [`AgentAwareOrchAlgo`] classifies a stable agent/task once, then reuses that model -//! assignment to retain model and prompt-cache locality. - -use std::collections::{BTreeMap, HashMap}; -use std::error::Error; -use std::sync::{Arc, Mutex}; - -use async_trait::async_trait; -use serde::Deserialize; -use switchyard_protocol::{completion_text, prompt_text, text_request, AgentContext}; - -use crate::{ - Algorithm, Context, Decision, Driver, LlmTargetSet, Metadata, Request, Response, Signals, -}; - -const CODEX_TURN_METADATA_HEADER: &str = "x-codex-turn-metadata"; -const MAX_ASSIGNMENTS: usize = 4096; - -type BoxErr = Box; - -/// One model available to the agent-aware classifier. -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct AgentRoutingCandidate { - /// Semantic target name used to look up the model in [`LlmTargetSet`]. - pub target: String, - /// Capability/cost guidance shown to the classifier. - pub description: String, -} - -impl AgentRoutingCandidate { - /// Creates a model-pool entry from a target name and classifier guidance. - pub fn new(target: impl Into, description: impl Into) -> Self { - Self { - target: target.into(), - description: description.into(), - } - } -} - -/// Inspectable decision emitted for classifier and routed calls. -pub struct AgentRoutingDecision { - /// Target selected from the configured model pool. - pub selected_model: String, - /// Human-readable classifier, cache, or fallback explanation. - pub reason: String, - /// Classifier-inferred task class, when supplied. - pub task_kind: Option, - /// Classifier confidence in `[0, 1]`, when supplied. - pub confidence: Option, - /// Stable agent id used for the assignment, when available. - pub agent_id: Option, - /// Whether the selection came from an existing per-agent assignment. - pub cache_hit: bool, -} - -impl Decision for AgentRoutingDecision { - fn selected_model(&self) -> &str { - &self.selected_model - } - - fn reasoning(&self) -> Option<&str> { - Some(&self.reason) - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } -} - -/// LLM classifier over a model pool with per-agent/task assignment affinity. -pub struct AgentAwareOrchAlgo { - classifier_model: String, - candidates: Vec, - fallback_model: String, - target_set: LlmTargetSet, - assignments: Mutex>, -} - -impl AgentAwareOrchAlgo { - /// Configures the classifier target, candidate pool, fail-open target, and targets. - pub fn new( - classifier_model: impl Into, - candidates: Vec, - fallback_model: impl Into, - target_set: LlmTargetSet, - ) -> Self { - Self { - classifier_model: classifier_model.into(), - candidates, - fallback_model: fallback_model.into(), - target_set, - assignments: Mutex::new(HashMap::new()), - } - } - - /// Classifies a request, failing open without caching invalid or failed results. - async fn classify( - &self, - ctx: &Context, - driver: &Driver, - request: &Request, - ) -> Result<(Assignment, bool), BoxErr> { - let classifier_decision: Arc = Arc::new(AgentRoutingDecision { - selected_model: self.classifier_model.clone(), - reason: format!("classifying agent/subtask via {}", self.classifier_model), - task_kind: request - .metadata - .as_ref() - .and_then(|metadata| metadata.agent_context.as_deref()) - .and_then(|agent| agent.task_kind.clone()), - confidence: None, - agent_id: request - .metadata - .as_ref() - .and_then(|metadata| metadata.agent_id.clone()), - cache_hit: false, - }); - driver - .info(ctx.clone(), Arc::clone(&classifier_decision)) - .await?; - - let classify_request = Request { - llm_request: text_request( - request.llm_request.model.clone(), - classifier_prompt( - &self.candidates, - request.metadata.as_ref(), - &prompt_text(&request.llm_request), - ), - ), - // Classifier calls receive neutral text and metadata, never the provider - // payload with its tools or continuation state. - raw_request: None, - metadata: request.metadata.clone(), - }; - - let output = match self.target_set.get_target(&self.classifier_model) { - Ok(target) => { - driver - .call_llm_target(ctx.clone(), &target, classify_request, classifier_decision) - .await - } - Err(error) => Err(error), - }; - - match output { - Ok(response) => match response.llm_response.into_agg().await { - Ok(response) => { - match parse_classifier_output(&completion_text(&response), &self.candidates) { - Some(assignment) => Ok((assignment, true)), - None => Ok(( - self.fallback_assignment( - "classifier returned an invalid model-pool decision", - ), - false, - )), - } - } - Err(error) => Ok(( - self.fallback_assignment(&format!("classifier response failed: {error}")), - false, - )), - }, - Err(error) => Ok(( - self.fallback_assignment(&format!("classifier call failed: {error}")), - false, - )), - } - } - - fn fallback_assignment(&self, reason: &str) -> Assignment { - Assignment { - model: self.fallback_model.clone(), - reason: format!("{reason}; fell back to {}", self.fallback_model), - task_kind: None, - confidence: None, - } - } - - fn cached_assignment(&self, key: &AssignmentKey) -> Option { - self.assignments - .lock() - .ok() - .and_then(|assignments| assignments.get(key).cloned()) - } - - fn store_assignment(&self, key: AssignmentKey, assignment: Assignment) { - if let Ok(mut assignments) = self.assignments.lock() { - if !assignments.contains_key(&key) && assignments.len() >= MAX_ASSIGNMENTS { - if let Some(evicted) = assignments.keys().next().cloned() { - assignments.remove(&evicted); - } - } - assignments.insert(key, assignment); - } - } -} - -#[async_trait] -impl Algorithm for AgentAwareOrchAlgo { - async fn create_run_task( - self: Arc, - ctx: Context, - driver: Driver, - request: Request, - ) -> Result { - let key = assignment_key(request.metadata.as_ref()); - let cached = key.as_ref().and_then(|key| self.cached_assignment(key)); - let (assignment, cache_hit) = match cached { - Some(assignment) => (assignment, true), - None => { - let (assignment, cacheable) = self.classify(&ctx, &driver, &request).await?; - if cacheable { - if let Some(key) = key { - self.store_assignment(key, assignment.clone()); - } - } - (assignment, false) - } - }; - - let routed_target = self.target_set.get_target(&assignment.model)?; - let route_decision: Arc = Arc::new(AgentRoutingDecision { - selected_model: assignment.model.clone(), - reason: if cache_hit { - format!("reused stable agent/task assignment: {}", assignment.reason) - } else { - assignment.reason.clone() - }, - task_kind: assignment.task_kind, - confidence: assignment.confidence, - agent_id: request - .metadata - .as_ref() - .and_then(|metadata| metadata.agent_id.clone()), - cache_hit, - }); - driver - .info(ctx.clone(), Arc::clone(&route_decision)) - .await?; - driver - .call_llm_target(ctx, &routed_target, request, route_decision) - .await - } - - async fn process_signals(self: Arc, _signals: Signals) -> Result<(), BoxErr> { - Ok(()) - } -} - -#[derive(Clone, Debug, Hash, PartialEq, Eq)] -struct AssignmentKey { - session_id: Option, - agent_id: String, - task_id: Option, -} - -#[derive(Clone, Debug)] -struct Assignment { - model: String, - reason: String, - task_kind: Option, - confidence: Option, -} - -#[derive(Deserialize)] -struct ClassifierOutput { - model: String, - #[serde(default)] - task_kind: Option, - #[serde(default)] - confidence: Option, - #[serde(default)] - reason: Option, -} - -#[derive(Default, Deserialize)] -struct CodexTurnMetadata { - session_id: Option, - thread_id: Option, - parent_thread_id: Option, - turn_id: Option, - subagent_kind: Option, - agent_role: Option, - task_id: Option, - task_kind: Option, -} - -/// Normalizes harness-specific request headers into libsy metadata. -/// -/// Explicit `x-switchyard-*` headers win. NeMo Relay and Dynamo correlation -/// headers are accepted without linking either runtime. Codex's structured turn -/// metadata is preferred over its compatibility projections. -pub fn metadata_from_headers(headers: &BTreeMap>) -> Metadata { - let codex = header(headers, CODEX_TURN_METADATA_HEADER) - .and_then(|value| serde_json::from_str::(&value).ok()) - .unwrap_or_default(); - - let agent_context = AgentContext { - parent_agent_id: first_some([ - header(headers, "x-switchyard-parent-agent-id"), - header(headers, "x-dynamo-parent-session-id"), - codex.parent_thread_id, - header(headers, "x-codex-parent-thread-id"), - ]), - agent_kind: first_some([ - header(headers, "x-switchyard-agent-kind"), - codex.subagent_kind, - header(headers, "x-openai-subagent"), - ]), - agent_role: first_some([header(headers, "x-switchyard-agent-role"), codex.agent_role]), - task_kind: first_some([header(headers, "x-switchyard-task-kind"), codex.task_kind]), - turn_id: first_some([header(headers, "x-switchyard-turn-id"), codex.turn_id]), - }; - let agent_context = (agent_context != AgentContext::default()).then(|| Box::new(agent_context)); - - Metadata { - session_id: first_some([ - header(headers, "x-switchyard-session-id"), - header(headers, "x-nemo-relay-session-id"), - codex.session_id, - header(headers, "session-id"), - ]), - agent_id: first_some([ - header(headers, "x-switchyard-agent-id"), - header(headers, "x-nemo-relay-subagent-id"), - header(headers, "x-dynamo-session-id"), - codex.thread_id, - header(headers, "thread-id"), - ]), - task_id: first_some([ - header(headers, "x-switchyard-task-id"), - codex.task_id, - header(headers, "x-task-id"), - ]), - agent_context, - correlation_id: first_some([ - header(headers, "x-request-id"), - header(headers, "x-client-request-id"), - ]), - ..Metadata::default() - } -} - -fn assignment_key(metadata: Option<&Metadata>) -> Option { - let metadata = metadata?; - Some(AssignmentKey { - session_id: metadata.session_id.clone(), - agent_id: metadata.agent_id.clone()?, - task_id: metadata.task_id.clone(), - }) -} - -fn classifier_prompt( - candidates: &[AgentRoutingCandidate], - metadata: Option<&Metadata>, - prompt: &str, -) -> String { - let metadata = metadata.cloned().unwrap_or_default(); - let agent = metadata - .agent_context - .as_deref() - .cloned() - .unwrap_or_default(); - let candidate_lines = candidates - .iter() - .map(|candidate| format!("- {}: {}", candidate.target, candidate.description)) - .collect::>() - .join("\n"); - format!( - "You route one model call inside an agent hierarchy. Select exactly one target from the \n\ - model pool. Prefer the cheapest candidate likely to complete the subtask correctly. \n\ - Return JSON only: {{\"model\":\"\",\"task_kind\":\"\",\ - \"confidence\":<0..1>,\"reason\":\"\"}}.\n\n\ - Agent context:\n\ - session_id={:?}\nagent_id={:?}\nparent_agent_id={:?}\nagent_kind={:?}\n\ - agent_role={:?}\ntask_id={:?}\ntask_kind={:?}\nturn_id={:?}\n\n\ - Model pool:\n{}\n\nSubtask:\n\n{}\n", - metadata.session_id, - metadata.agent_id, - agent.parent_agent_id, - agent.agent_kind, - agent.agent_role, - metadata.task_id, - agent.task_kind, - agent.turn_id, - candidate_lines, - prompt, - ) -} - -fn parse_classifier_output( - output: &str, - candidates: &[AgentRoutingCandidate], -) -> Option { - let trimmed = output.trim(); - if let Some(candidate) = candidates - .iter() - .find(|candidate| candidate.target == trimmed) - { - return Some(Assignment { - model: candidate.target.clone(), - reason: "classifier selected the target by name".to_string(), - task_kind: None, - confidence: None, - }); - } - - let start = trimmed.find('{')?; - let end = trimmed.rfind('}')?; - let parsed = serde_json::from_str::(&trimmed[start..=end]).ok()?; - let model = parsed.model.trim(); - if !candidates.iter().any(|candidate| candidate.target == model) { - return None; - } - Some(Assignment { - model: model.to_string(), - reason: parsed - .reason - .filter(|reason| !reason.trim().is_empty()) - .unwrap_or_else(|| "classifier selected the target".to_string()), - task_kind: parsed.task_kind.filter(|kind| !kind.trim().is_empty()), - confidence: parsed - .confidence - .filter(|value| value.is_finite()) - .map(|value| value.clamp(0.0, 1.0)), - }) -} - -fn header(headers: &BTreeMap>, name: &str) -> Option { - headers - .iter() - .find(|(key, _)| key.eq_ignore_ascii_case(name)) - .and_then(|(_, values)| values.iter().find(|value| !value.trim().is_empty())) - .map(|value| value.trim().to_string()) -} - -fn first_some(values: [Option; N]) -> Option { - values.into_iter().flatten().next() -} - -#[cfg(test)] -mod tests { - use std::collections::VecDeque; - - use super::*; - use crate::{LlmClient, LlmResponse, LlmTarget, RoutedRequest}; - use switchyard_protocol::{text_response, Metadata}; - - struct RoutingClient { - classifier_calls: Arc>, - responses: Arc>>, - } - - #[async_trait] - impl LlmClient for RoutingClient { - async fn call(&self, routed: RoutedRequest) -> Result { - let selected = routed.decision.selected_model().to_string(); - let completion = if selected == "classifier" { - let mut calls = self.classifier_calls.lock().map_err(|_| "lock poisoned")?; - *calls += 1; - self.responses - .lock() - .map_err(|_| "lock poisoned")? - .pop_front() - .ok_or("missing classifier response")? - } else { - selected - }; - Ok(Response { - llm_response: LlmResponse::Agg(text_response(None, completion)), - metadata: routed.request.metadata, - }) - } - } - - fn algo(responses: Vec<&str>) -> (Arc, Arc>) { - let classifier_calls = Arc::new(Mutex::new(0)); - let client = Arc::new(RoutingClient { - classifier_calls: Arc::clone(&classifier_calls), - responses: Arc::new(Mutex::new( - responses.into_iter().map(str::to_string).collect(), - )), - }) as Arc; - let target = |name: &str| LlmTarget { - semantic_name: name.to_string(), - llm_client: Some(Arc::clone(&client)), - }; - let targets = LlmTargetSet::new(vec![ - target("classifier"), - target("frontier"), - target("fast"), - ]); - let algo: Arc = Arc::new(AgentAwareOrchAlgo::new( - "classifier", - vec![ - AgentRoutingCandidate::new("frontier", "complex planning and final review"), - AgentRoutingCandidate::new("fast", "bounded research and mechanical edits"), - ], - "frontier", - targets, - )); - (algo, classifier_calls) - } - - fn request(agent_id: Option<&str>, task_id: Option<&str>, prompt: &str) -> Request { - Request { - llm_request: text_request(Some("auto".to_string()), prompt), - raw_request: Some(serde_json::json!({"input": prompt})), - metadata: Some(Metadata { - session_id: Some("session-1".to_string()), - agent_id: agent_id.map(str::to_string), - task_id: task_id.map(str::to_string), - agent_context: Some(Box::new(AgentContext { - agent_role: Some("explorer".to_string()), - ..AgentContext::default() - })), - ..Metadata::default() - }), - } - } - - fn decision(model: &str, task_kind: &str, confidence: f64, reason: &str) -> String { - serde_json::json!({ - "model": model, - "task_kind": task_kind, - "confidence": confidence, - "reason": reason, - }) - .to_string() - } - - fn response_text(response: &Response) -> String { - response - .llm_response - .as_agg() - .map(completion_text) - .unwrap_or_default() - } - - #[test] - fn normalizes_codex_turn_metadata_and_compatibility_headers() { - let mut headers = BTreeMap::new(); - headers.insert("session-id".to_string(), vec!["compat-session".to_string()]); - headers.insert("thread-id".to_string(), vec!["compat-agent".to_string()]); - headers.insert( - CODEX_TURN_METADATA_HEADER.to_string(), - vec![serde_json::json!({ - "session_id": "root-session", - "thread_id": "child-agent", - "parent_thread_id": "root-agent", - "turn_id": "turn-7", - "subagent_kind": "collab_spawn", - }) - .to_string()], - ); - - let metadata = metadata_from_headers(&headers); - assert_eq!(metadata.session_id.as_deref(), Some("root-session")); - assert_eq!(metadata.agent_id.as_deref(), Some("child-agent")); - let agent = metadata.agent_context.as_deref(); - assert_eq!( - agent.and_then(|value| value.parent_agent_id.as_deref()), - Some("root-agent") - ); - assert_eq!( - agent.and_then(|value| value.agent_kind.as_deref()), - Some("collab_spawn") - ); - assert_eq!( - agent.and_then(|value| value.turn_id.as_deref()), - Some("turn-7") - ); - } - - #[test] - fn explicit_switchyard_headers_override_relay_and_codex() { - let headers = BTreeMap::from([ - ( - "x-switchyard-session-id".to_string(), - vec!["canonical-session".to_string()], - ), - ( - "x-switchyard-agent-id".to_string(), - vec!["canonical-agent".to_string()], - ), - ( - "x-nemo-relay-session-id".to_string(), - vec!["relay-session".to_string()], - ), - ( - "x-dynamo-session-id".to_string(), - vec!["relay-agent".to_string()], - ), - ]); - - let metadata = metadata_from_headers(&headers); - assert_eq!(metadata.session_id.as_deref(), Some("canonical-session")); - assert_eq!(metadata.agent_id.as_deref(), Some("canonical-agent")); - } - - #[test] - fn normalizes_relay_and_dynamo_headers_without_runtime_dependency() { - let headers = BTreeMap::from([ - ( - "x-nemo-relay-session-id".to_string(), - vec!["relay-session".to_string()], - ), - ( - "x-nemo-relay-subagent-id".to_string(), - vec!["relay-child".to_string()], - ), - ( - "x-dynamo-parent-session-id".to_string(), - vec!["relay-parent".to_string()], - ), - ]); - - let metadata = metadata_from_headers(&headers); - assert_eq!(metadata.session_id.as_deref(), Some("relay-session")); - assert_eq!(metadata.agent_id.as_deref(), Some("relay-child")); - assert_eq!( - metadata - .agent_context - .as_deref() - .and_then(|agent| agent.parent_agent_id.as_deref()), - Some("relay-parent") - ); - } - - #[tokio::test] - async fn valid_classifier_output_routes_from_the_pool() -> Result<(), BoxErr> { - let (algo, calls) = algo(vec![&decision("fast", "research", 0.9, "bounded lookup")]); - let (trace, response) = algo - .run( - Context::default(), - request(Some("agent-a"), None, "survey the API"), - ) - .await?; - - assert_eq!(response_text(&response), "fast"); - assert_eq!(trace.len(), 2); - let routed = trace[1] - .as_any() - .downcast_ref::() - .ok_or("expected AgentRoutingDecision")?; - assert_eq!(routed.task_kind.as_deref(), Some("research")); - assert_eq!(routed.confidence, Some(0.9)); - assert!(!routed.cache_hit); - assert_eq!(*calls.lock().map_err(|_| "lock poisoned")?, 1); - Ok(()) - } - - #[tokio::test] - async fn stable_agent_assignment_avoids_reclassification() -> Result<(), BoxErr> { - let (algo, calls) = algo(vec![&decision("fast", "research", 0.8, "explorer")]); - Arc::clone(&algo) - .run( - Context::default(), - request(Some("agent-a"), None, "first turn"), - ) - .await?; - let (trace, response) = algo - .run( - Context::default(), - request(Some("agent-a"), None, "second turn"), - ) - .await?; - - assert_eq!(response_text(&response), "fast"); - assert_eq!(*calls.lock().map_err(|_| "lock poisoned")?, 1); - assert_eq!(trace.len(), 1); - let routed = trace[0] - .as_any() - .downcast_ref::() - .ok_or("expected AgentRoutingDecision")?; - assert!(routed.cache_hit); - Ok(()) - } - - #[tokio::test] - async fn explicit_task_change_reclassifies_the_same_agent() -> Result<(), BoxErr> { - let first = decision("fast", "research", 0.8, "lookup"); - let second = decision("frontier", "review", 0.9, "final review"); - let (algo, calls) = algo(vec![&first, &second]); - let (_, first_response) = Arc::clone(&algo) - .run( - Context::default(), - request(Some("agent-a"), Some("task-1"), "collect evidence"), - ) - .await?; - let (_, second_response) = algo - .run( - Context::default(), - request(Some("agent-a"), Some("task-2"), "adjudicate findings"), - ) - .await?; - - assert_eq!(response_text(&first_response), "fast"); - assert_eq!(response_text(&second_response), "frontier"); - assert_eq!(*calls.lock().map_err(|_| "lock poisoned")?, 2); - Ok(()) - } - - #[tokio::test] - async fn invalid_decision_falls_back_without_poisoning_affinity() -> Result<(), BoxErr> { - let valid = decision("fast", "research", 0.7, "retry succeeded"); - let (algo, calls) = algo(vec!["not-json", &valid]); - let (_, first) = Arc::clone(&algo) - .run(Context::default(), request(Some("agent-a"), None, "first")) - .await?; - let (_, second) = algo - .run(Context::default(), request(Some("agent-a"), None, "second")) - .await?; - - assert_eq!(response_text(&first), "frontier"); - assert_eq!(response_text(&second), "fast"); - assert_eq!(*calls.lock().map_err(|_| "lock poisoned")?, 2); - Ok(()) - } -} diff --git a/crates/libsy/src/algorithms/rand.rs b/crates/libsy/src/algorithms/rand.rs index adc2b366..f75c51b2 100644 --- a/crates/libsy/src/algorithms/rand.rs +++ b/crates/libsy/src/algorithms/rand.rs @@ -14,6 +14,7 @@ use std::sync::Arc; use async_trait::async_trait; use rand::seq::SliceRandom; +use crate::affinity::Affinity; use crate::{Algorithm, Context, Decision, Driver, LlmTargetSet, Request, Response}; /// Decision produced by [`RandomAlgo`]: which target was chosen and why. @@ -41,6 +42,7 @@ impl Decision for RandomDecision { /// Uniform random router over a target set. pub struct RandomAlgo { target_set: LlmTargetSet, + affinity: Option>, } impl RandomAlgo { @@ -49,7 +51,35 @@ impl RandomAlgo { /// Wrap it in an [`Arc`] and drive it with [`Algorithm::run`] or /// [`Algorithm::run_stream`]. pub fn new(target_set: LlmTargetSet) -> Self { - Self { target_set } + Self { + target_set, + affinity: None, + } + } + + /// Retains selections according to `affinity` while preserving per-request + /// random routing for requests the policy does not key. + pub fn with_affinity(mut self, affinity: Arc) -> Self { + self.affinity = Some(affinity); + self + } + + /// Selects one configured target name uniformly at random. + fn choose_target(&self) -> Result> { + if let [target] = self.target_set.targets() { + return Ok(target.semantic_name.clone()); + } + // Scope the non-Send ThreadRng so callers may await after this returns. + let selected = { + let mut rng = rand::thread_rng(); + self.target_set + .targets() + .choose(&mut rng) + .ok_or("no targets available")? + .semantic_name + .clone() + }; + Ok(selected) } } @@ -61,19 +91,35 @@ impl Algorithm for RandomAlgo { driver: Driver, request: Request, ) -> Result> { - // Scope the non-Send ThreadRng before the await so the future remains Send. - let target = { - let mut rng = rand::thread_rng(); - self.target_set - .targets() - .choose(&mut rng) - .ok_or("no targets available")? - .clone() + let (selected, reasoning) = match self.affinity.as_ref() { + Some(affinity) => match affinity.assignment(&request) { + Some(selected) => { + let reasoning = format!("affinity reused target '{selected}'"); + (selected, reasoning) + } + None => { + let proposed = self.choose_target()?; + let applies = affinity.key(&request).is_some(); + let selected = affinity.retain(&request, proposed.clone()); + let reasoning = if !applies { + format!("random routing selected target '{selected}'") + } else if selected == proposed { + format!("random routing selected and retained target '{selected}'") + } else { + format!("affinity retained concurrent target '{selected}'") + }; + (selected, reasoning) + } + }, + None => { + let selected = self.choose_target()?; + let reasoning = format!("random routing selected target '{selected}'"); + (selected, reasoning) + } }; - - let selected = target.semantic_name.clone(); + let target = self.target_set.get_target(&selected)?; let decision: Arc = Arc::new(RandomDecision { - reasoning: format!("random routing selected target '{selected}'"), + reasoning, selected_model: selected, }); @@ -91,7 +137,8 @@ mod tests { use switchyard_protocol::{completion_text, text_request, text_response}; - use crate::{LlmResponse, LlmTarget, Request, RoutedLlmClient, Signals}; + use crate::affinity::{Affinity, SessionAffinity, SubAgentAffinity}; + use crate::{AgentContext, LlmResponse, LlmTarget, Metadata, Request, RoutedLlmClient, Signals}; /// Echoes the selected target so tests can inspect which target was called. struct EchoClient; @@ -135,6 +182,30 @@ mod tests { Arc::new(algorithm(names)) } + fn shared_algorithm_with_affinity( + names: &[&str], + affinity: Arc, + ) -> Arc { + Arc::new(algorithm(names).with_affinity(affinity)) + } + + fn affinity_request(is_subagent: bool, agent_id: &str, task_id: &str) -> Request { + Request { + llm_request: text_request(Some("auto".to_string()), "hi"), + raw_request: None, + metadata: Some(Metadata { + session_id: Some("session-1".to_string()), + agent_id: Some(agent_id.to_string()), + task_id: Some(task_id.to_string()), + agent_context: Some(Box::new(AgentContext { + is_subagent, + ..AgentContext::default() + })), + ..Metadata::default() + }), + } + } + #[tokio::test] async fn single_target_is_always_selected_and_called( ) -> Result<(), Box> { @@ -234,4 +305,143 @@ mod tests { assert_eq!(concrete.selected_model, "only/model"); Ok(()) } + + #[tokio::test] + async fn subagent_affinity_reuses_selection_across_task_changes( + ) -> Result<(), Box> { + let algorithm = shared_algorithm_with_affinity( + &["a/model", "b/model"], + Arc::new(SubAgentAffinity::new()), + ); + let (_, first) = Arc::clone(&algorithm) + .run( + Context::default(), + affinity_request(true, "child-1", "task-1"), + ) + .await?; + let (trace, second) = algorithm + .run( + Context::default(), + affinity_request(true, "child-1", "task-2"), + ) + .await?; + + assert_eq!(first.selected_model(), second.selected_model()); + assert!(trace[0] + .reasoning() + .unwrap_or_default() + .contains("affinity reused")); + Ok(()) + } + + #[tokio::test] + async fn single_target_subagent_is_retained_without_reselection( + ) -> Result<(), Box> { + let algorithm = + shared_algorithm_with_affinity(&["only/model"], Arc::new(SubAgentAffinity::new())); + Arc::clone(&algorithm) + .run( + Context::default(), + affinity_request(true, "child-1", "task-1"), + ) + .await?; + let (trace, response) = algorithm + .run( + Context::default(), + affinity_request(true, "child-1", "task-2"), + ) + .await?; + + assert_eq!( + response + .llm_response + .as_agg() + .map(completion_text) + .unwrap_or_default(), + "only/model" + ); + assert!(trace[0] + .reasoning() + .unwrap_or_default() + .contains("affinity reused")); + Ok(()) + } + + #[tokio::test] + async fn subagent_affinity_does_not_retain_root_selection( + ) -> Result<(), Box> { + let affinity = Arc::new(SubAgentAffinity::new()); + let algorithm = shared_algorithm_with_affinity( + &["a/model", "b/model"], + Arc::clone(&affinity) as Arc, + ); + Arc::clone(&algorithm) + .run( + Context::default(), + affinity_request(false, "root-1", "task-1"), + ) + .await?; + + assert!(affinity + .assignment(&affinity_request(false, "root-1", "task-2")) + .is_none()); + let (trace, _) = algorithm + .run( + Context::default(), + affinity_request(false, "root-1", "task-2"), + ) + .await?; + assert!(trace[0] + .reasoning() + .unwrap_or_default() + .starts_with("random routing selected")); + Ok(()) + } + + #[tokio::test] + async fn session_affinity_reuses_selection_for_different_agents( + ) -> Result<(), Box> { + let algorithm = shared_algorithm_with_affinity( + &["a/model", "b/model"], + Arc::new(SessionAffinity::new()), + ); + let (_, first) = Arc::clone(&algorithm) + .run( + Context::default(), + affinity_request(false, "agent-a", "task-1"), + ) + .await?; + let (_, second) = algorithm + .run( + Context::default(), + affinity_request(false, "agent-b", "task-2"), + ) + .await?; + + assert_eq!(first.selected_model(), second.selected_model()); + Ok(()) + } + + #[tokio::test] + async fn concurrent_first_subagent_turns_use_one_assignment( + ) -> Result<(), Box> { + let algorithm = shared_algorithm_with_affinity( + &["a/model", "b/model"], + Arc::new(SubAgentAffinity::new()), + ); + let first = Arc::clone(&algorithm).run( + Context::default(), + affinity_request(true, "child-1", "task-1"), + ); + let second = algorithm.run( + Context::default(), + affinity_request(true, "child-1", "task-1"), + ); + let (first, second) = tokio::join!(first, second); + let (_, first) = first?; + let (_, second) = second?; + + assert_eq!(first.selected_model(), second.selected_model()); + Ok(()) + } } diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index 0176342a..0cd7b751 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -54,15 +54,15 @@ //! //! ## Algorithms //! -//! [`RandomAlgo`] provides uniform random routing in the core crate. Worked -//! [`agentic::AgentAwareOrchAlgo`] provides sub-agent-aware model-pool routing and -//! neutral harness-header normalization. Worked implementations of an LLM classifier -//! and a stateful ensemble, plus runnable agents, live in the `libsy-examples` crate. +//! [`RandomAlgo`] provides uniform random routing in the core crate and can retain +//! selections through an [`affinity::Affinity`] policy. Worked implementations of an +//! LLM classifier and a stateful ensemble, plus runnable agents, live in the +//! `libsy-examples` crate. mod algorithms; pub use algorithms::noop::{NoopAlgo, NoopDecision}; pub use algorithms::rand::{RandomAlgo, RandomDecision}; -pub mod agentic; +pub mod affinity; mod driver; diff --git a/crates/libsy/tests/affinity.rs b/crates/libsy/tests/affinity.rs new file mode 100644 index 00000000..6226bacb --- /dev/null +++ b/crates/libsy/tests/affinity.rs @@ -0,0 +1,45 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Public affinity extension-point contract tests. + +use libsy::affinity::{Affinity, AffinityKey, AffinityState}; +use libsy::{Metadata, Request}; +use switchyard_protocol::text_request; + +/// Example policy implemented outside the `libsy` crate. +#[derive(Default)] +struct TaskAffinity { + state: AffinityState, +} + +impl Affinity for TaskAffinity { + fn key(&self, request: &Request) -> Option { + let metadata = request.metadata.as_ref()?; + Some(AffinityKey::new( + "task", + vec![metadata.session_id.clone()?, metadata.task_id.clone()?], + )) + } + + fn state(&self) -> &AffinityState { + &self.state + } +} + +#[test] +fn external_policy_can_reuse_the_public_affinity_state() { + let affinity = TaskAffinity::default(); + let request = Request { + llm_request: text_request(Some("auto".to_string()), "hi"), + raw_request: None, + metadata: Some(Metadata { + session_id: Some("session-1".to_string()), + task_id: Some("task-1".to_string()), + ..Metadata::default() + }), + }; + + assert_eq!(affinity.retain(&request, "model-a".to_string()), "model-a"); + assert_eq!(affinity.assignment(&request).as_deref(), Some("model-a")); +} diff --git a/demo/libsy-proxy/README.md b/demo/libsy-proxy/README.md index 0b2aa39f..f5962304 100644 --- a/demo/libsy-proxy/README.md +++ b/demo/libsy-proxy/README.md @@ -1,9 +1,9 @@ -# Agent-aware libsy proxy POC +# Affinity-aware libsy proxy POC -This demo embeds `libsy`'s agent-aware model-pool router in the PR #17 Switchyard server. It accepts +This demo embeds `libsy`'s `RandomAlgo` with `SubAgentAffinity` in the Switchyard server. It accepts OpenAI Chat, Anthropic Messages, and OpenAI Responses request shapes, normalizes Codex/Relay lineage -headers, asks the classifier for a model assignment, and reuses valid assignments per stable -`(session, agent, explicit task)` key. +headers, randomly routes ordinary requests on every turn, and retains the first selection for each +stable `(session_id, agent_id)` child-agent key. Run it with an NVIDIA Inference Hub bearer token: @@ -11,8 +11,8 @@ Run it with an NVIDIA Inference Hub bearer token: ANTHROPIC_API_KEY="$INFERENCE_HUB_SY_API_KEY" cargo run -p libsy-proxy ``` -Send two buffered calls for the same Codex child thread. The first call classifies; the second reuses -the assignment: +Send two buffered calls for the same Codex child thread. The first call randomly selects a target; +the second reuses the assignment: ```bash curl http://127.0.0.1:4000/v1/chat/completions \ @@ -22,13 +22,13 @@ curl http://127.0.0.1:4000/v1/chat/completions \ -H 'x-codex-parent-thread-id: root-thread-1' \ -H 'x-openai-subagent: collab_spawn' \ -H 'x-codex-turn-metadata: {"session_id":"root-session-1","thread_id":"child-thread-1","parent_thread_id":"root-thread-1","turn_id":"turn-1","subagent_kind":"collab_spawn"}' \ - -d '{"model":"libsy-agent-aware","messages":[{"role":"user","content":"Survey the request protocol and report evidence."}],"stream":false}' + -d '{"model":"libsy-random-affinity","messages":[{"role":"user","content":"Survey the request protocol and report evidence."}],"stream":false}' ``` The response exposes the chosen logical target and rationale in `x-model-router-*` headers. Explicit `x-switchyard-session-id`, `x-switchyard-agent-id`, `x-switchyard-parent-agent-id`, -`x-switchyard-agent-role`, `x-switchyard-task-id`, and `x-switchyard-task-kind` headers override -harness-derived values. +`x-switchyard-is-subagent`, `x-switchyard-agent-role`, and `x-switchyard-task-kind` headers override +harness-derived values. A root-agent request is not retained and runs random routing on every turn. The demo currently requires buffered upstream responses (`"stream": false`). Header normalization and route selection work for Responses-shaped requests, but running an unmodified streaming Codex diff --git a/demo/libsy-proxy/src/main.rs b/demo/libsy-proxy/src/main.rs index 2992271a..a735fb43 100644 --- a/demo/libsy-proxy/src/main.rs +++ b/demo/libsy-proxy/src/main.rs @@ -1,11 +1,11 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Demo proxy combining Switchyard HTTP/translation with libsy agent-aware routing. +//! Demo proxy combining Switchyard HTTP/translation with libsy random routing and affinity. //! //! The server accepts OpenAI Chat, Anthropic Messages, and OpenAI Responses requests. -//! libsy normalizes harness identity headers, assigns a stable agent/task to a model -//! pool target, and makes calls through Switchyard's OpenAI-compatible backend. +//! libsy normalizes harness identity headers, randomly routes ordinary requests, and +//! retains the first model selected for each stable child agent. use std::collections::BTreeMap; use std::error::Error; @@ -15,12 +15,10 @@ use std::sync::Arc; use async_trait::async_trait; use serde_json::Value; -use libsy::agentic::{ - metadata_from_headers, AgentAwareOrchAlgo, AgentRoutingCandidate, AgentRoutingDecision, -}; +use libsy::affinity::{metadata_from_headers, SubAgentAffinity}; use libsy::{ - Algorithm, Context, Decision, LlmClient, LlmResponse, LlmTarget, LlmTargetSet, Request, - Response, RoutedRequest, + Algorithm, Context, Decision, LlmClient, LlmResponse, LlmTarget, LlmTargetSet, RandomAlgo, + Request, Response, RoutedRequest, }; use switchyard_components::OpenAiPassthroughBackend; use switchyard_components_v2::{Profile, ProfileInput, ProfileResponse, RoutingMetadata}; @@ -31,16 +29,14 @@ use switchyard_core::{ use switchyard_server::{build_switchyard_router, ProfileRegistry, ServerState}; use switchyard_translation::{TranslationEngine, TranslationPolicy, WireFormat}; -const CLASSIFIER_MODEL: &str = "nvidia/deepseek-ai/deepseek-v4-flash"; const FRONTIER_MODEL: &str = "aws/anthropic/bedrock-claude-opus-4-7"; const FAST_MODEL: &str = "nvidia/deepseek-ai/deepseek-v4-flash"; -const CLASSIFIER_TARGET: &str = "classifier"; const FRONTIER_TARGET: &str = "frontier"; const FAST_TARGET: &str = "fast"; const DEFAULT_BASE_URL: &str = "https://inference-api.nvidia.com/v1"; const DEFAULT_ADDR: &str = "127.0.0.1:4000"; -const PROFILE_MODEL_ID: &str = "libsy-agent-aware"; +const PROFILE_MODEL_ID: &str = "libsy-random-affinity"; type BoxErr = Box; @@ -84,13 +80,13 @@ impl LlmClient for SwitchyardBackendClient { } /// Switchyard profile that delegates all routing decisions to libsy. -struct LibsyAgentAwareProfile { +struct LibsyAffinityProfile { algorithm: Arc, translation: Arc, } #[async_trait] -impl Profile for LibsyAgentAwareProfile { +impl Profile for LibsyAffinityProfile { async fn run(&self, input: ProfileInput) -> Result { let request = libsy_request(&input, self.translation.as_ref())?; let (trace, response) = Arc::clone(&self.algorithm) @@ -141,20 +137,20 @@ fn libsy_request(input: &ProfileInput, translation: &TranslationEngine) -> Resul }) } -/// Preserve routed provider payloads; encode classifier calls from the neutral IR. +/// Preserve routed provider payloads; encode neutral requests when no raw body exists. fn chat_request_for_call( request: &Request, model: &str, translation: &TranslationEngine, ) -> std::result::Result { let Some(mut body) = request.raw_request.clone() else { - let mut classifier_request = request.llm_request.clone(); - classifier_request.model = Some(model.to_string()); - classifier_request.stream = false; + let mut neutral_request = request.llm_request.clone(); + neutral_request.model = Some(model.to_string()); + neutral_request.stream = false; let body = translation .encode_request( WireFormat::OpenAiChat, - &classifier_request, + &neutral_request, &TranslationPolicy::default(), ) .map_err(|error| BoxErr::from(error.to_string()))? @@ -198,8 +194,6 @@ fn wire_format(request_type: ChatRequestType) -> WireFormat { /// Surface libsy's routing decision as `x-model-router-*` response headers. fn routing_metadata(trace: &[Arc]) -> RoutingMetadata { let decision = trace.last(); - let agent = - decision.and_then(|decision| decision.as_any().downcast_ref::()); let selected_target = decision.map(|decision| decision.selected_model()); RoutingMetadata { selected_model: selected_target.map(|target| match target { @@ -208,8 +202,8 @@ fn routing_metadata(trace: &[Arc]) -> RoutingMetadata { other => other.to_string(), }), selected_tier: selected_target.map(str::to_string), - confidence: agent.and_then(|decision| decision.confidence), - router_version: Some("libsy-agent-aware-v1".to_string()), + confidence: None, + router_version: Some("libsy-random-affinity-v1".to_string()), tolerance: None, rationale: decision.and_then(|decision| decision.reasoning().map(str::to_string)), } @@ -229,7 +223,6 @@ fn build_algorithm() -> Result> { timeout_secs: Some(120.0), })?); let model_ids = BTreeMap::from([ - (CLASSIFIER_TARGET.to_string(), CLASSIFIER_MODEL.to_string()), (FRONTIER_TARGET.to_string(), FRONTIER_MODEL.to_string()), (FAST_TARGET.to_string(), FAST_MODEL.to_string()), ]); @@ -243,33 +236,16 @@ fn build_algorithm() -> Result> { semantic_name: name.to_string(), llm_client: Some(Arc::clone(&client)), }; - let targets = LlmTargetSet::new(vec![ - target(CLASSIFIER_TARGET), - target(FRONTIER_TARGET), - target(FAST_TARGET), - ]); + let targets = LlmTargetSet::new(vec![target(FRONTIER_TARGET), target(FAST_TARGET)]); + let algorithm = RandomAlgo::new(targets).with_affinity(Arc::new(SubAgentAffinity::new())); - Ok(Arc::new(AgentAwareOrchAlgo::new( - CLASSIFIER_TARGET, - vec![ - AgentRoutingCandidate::new( - FRONTIER_TARGET, - "frontier model for planning, ambiguous implementation, synthesis, and review", - ), - AgentRoutingCandidate::new( - FAST_TARGET, - "efficient model for bounded exploration, retrieval, and mechanical edits", - ), - ], - FRONTIER_TARGET, - targets, - ))) + Ok(Arc::new(algorithm)) } #[tokio::main] async fn main() -> Result<()> { let translation = Arc::new(TranslationEngine::default()); - let profile = Arc::new(LibsyAgentAwareProfile { + let profile = Arc::new(LibsyAffinityProfile { algorithm: build_algorithm()?, translation, }) as Arc; @@ -293,9 +269,9 @@ async fn main() -> Result<()> { .map_err(|error| SwitchyardError::Other(error.to_string()))?; println!("libsy-proxy listening on http://{bound_addr}"); - println!(" routing (libsy agent-aware): classifier={CLASSIFIER_MODEL}"); - println!(" frontier={FRONTIER_MODEL}"); - println!(" fast={FAST_MODEL}"); + println!(" routing (libsy random + sub-agent affinity):"); + println!(" frontier={FRONTIER_MODEL}"); + println!(" fast={FAST_MODEL}"); println!( " send model \"{PROFILE_MODEL_ID}\" to /v1/chat/completions, /v1/messages, or /v1/responses" ); @@ -318,9 +294,9 @@ mod tests { #[test] fn routed_call_preserves_provider_body_and_rewrites_model() -> std::result::Result<(), BoxErr> { let request = Request { - llm_request: text_request(Some("libsy-agent-aware".to_string()), "inspect"), + llm_request: text_request(Some(PROFILE_MODEL_ID.to_string()), "inspect"), raw_request: Some(json!({ - "model": "libsy-agent-aware", + "model": PROFILE_MODEL_ID, "input": "inspect", "tools": [{"type": "function", "name": "shell"}], "stream": false, @@ -349,7 +325,7 @@ mod tests { fn decodes_responses_input_and_normalizes_headers() -> Result<()> { let input = ProfileInput { request: ChatRequest::openai_responses(json!({ - "model": "libsy-agent-aware", + "model": PROFILE_MODEL_ID, "input": [{ "type": "message", "role": "user", @@ -358,10 +334,14 @@ mod tests { "stream": false, })), metadata: RequestMetadata { - headers: BTreeMap::from([( - "thread-id".to_string(), - vec!["child-agent".to_string()], - )]), + headers: BTreeMap::from([ + ("session-id".to_string(), vec!["session-1".to_string()]), + ("thread-id".to_string(), vec!["child-agent".to_string()]), + ( + "x-openai-subagent".to_string(), + vec!["collab_spawn".to_string()], + ), + ]), ..RequestMetadata::default() }, }; @@ -375,22 +355,30 @@ mod tests { .and_then(|metadata| metadata.agent_id.as_deref()), Some("child-agent") ); + assert_eq!( + request + .metadata + .as_ref() + .and_then(|metadata| metadata.agent_context.as_deref()) + .map(|agent| agent.is_subagent), + Some(true) + ); Ok(()) } #[test] - fn classifier_call_uses_a_buffered_synthetic_request() -> std::result::Result<(), BoxErr> { + fn neutral_call_uses_a_buffered_synthetic_request() -> std::result::Result<(), BoxErr> { let request = Request { llm_request: text_request(Some("auto".to_string()), "classify this"), raw_request: None, metadata: None, }; - let classifier = - chat_request_for_call(&request, "classifier/model", &TranslationEngine::default())?; - assert_eq!(classifier.request_type(), ChatRequestType::OpenAiChat); + let neutral = + chat_request_for_call(&request, "provider/model", &TranslationEngine::default())?; + assert_eq!(neutral.request_type(), ChatRequestType::OpenAiChat); assert_ne!( - classifier.body().get("stream").and_then(Value::as_bool), + neutral.body().get("stream").and_then(Value::as_bool), Some(true) ); Ok(()) @@ -398,18 +386,14 @@ mod tests { #[test] fn response_metadata_exposes_provider_model_and_logical_tier() { - let decision: Arc = Arc::new(AgentRoutingDecision { + let decision: Arc = Arc::new(libsy::RandomDecision { selected_model: FAST_TARGET.to_string(), - reason: "bounded lookup".to_string(), - task_kind: Some("research".to_string()), - confidence: Some(0.9), - agent_id: Some("child-1".to_string()), - cache_hit: false, + reasoning: "random routing selected target 'fast'".to_string(), }); let metadata = routing_metadata(&[decision]); assert_eq!(metadata.selected_model.as_deref(), Some(FAST_MODEL)); assert_eq!(metadata.selected_tier.as_deref(), Some(FAST_TARGET)); - assert_eq!(metadata.confidence, Some(0.9)); + assert_eq!(metadata.confidence, None); } } From e61b9c6358048b8699a583725649529718256b7d Mon Sep 17 00:00:00 2001 From: ayushag Date: Thu, 16 Jul 2026 15:24:05 -0700 Subject: [PATCH 08/14] docs(libsy): refresh affinity terminology Signed-off-by: ayushag --- crates/libsy-protocol/src/envelope.rs | 4 ++-- demo/libsy-proxy/src/main.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/libsy-protocol/src/envelope.rs b/crates/libsy-protocol/src/envelope.rs index 36a39054..aa8f0591 100644 --- a/crates/libsy-protocol/src/envelope.rs +++ b/crates/libsy-protocol/src/envelope.rs @@ -40,7 +40,7 @@ pub struct Metadata { pub wire_format: Option, } -/// Optional lineage and semantic signals for agent-aware routing. +/// Optional lineage and semantic signals for affinity-aware routing. #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct AgentContext { /// Whether the harness explicitly identified this request as coming from a child agent. @@ -51,7 +51,7 @@ pub struct AgentContext { pub agent_kind: Option, /// Semantic agent role, such as `explorer`, `worker`, or `reviewer`. pub agent_role: Option, - /// Semantic task class supplied by the harness or a prior classifier. + /// Semantic task class supplied by the harness. pub task_kind: Option, /// Id of the current agent turn. pub turn_id: Option, diff --git a/demo/libsy-proxy/src/main.rs b/demo/libsy-proxy/src/main.rs index a735fb43..4886fcb2 100644 --- a/demo/libsy-proxy/src/main.rs +++ b/demo/libsy-proxy/src/main.rs @@ -369,7 +369,7 @@ mod tests { #[test] fn neutral_call_uses_a_buffered_synthetic_request() -> std::result::Result<(), BoxErr> { let request = Request { - llm_request: text_request(Some("auto".to_string()), "classify this"), + llm_request: text_request(Some("auto".to_string()), "route this"), raw_request: None, metadata: None, }; From c2cefcd66f4ec60e28d1285176cc069cc7725a57 Mon Sep 17 00:00:00 2001 From: ayushag Date: Thu, 16 Jul 2026 15:28:03 -0700 Subject: [PATCH 09/14] ci: rerun checks after GitHub outage Signed-off-by: ayushag From 55f9d5ab426140d348a992035db96bc3d0b1f504 Mon Sep 17 00:00:00 2001 From: ayushag Date: Thu, 16 Jul 2026 15:29:44 -0700 Subject: [PATCH 10/14] ci: pin uv setup version Signed-off-by: ayushag --- .github/workflows/ci.yml | 4 ++++ .github/workflows/perf.yml | 1 + 2 files changed, 5 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5f8d5aa7..50f79114 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,6 +22,7 @@ jobs: - uses: actions/checkout@v4 - uses: astral-sh/setup-uv@v6 with: + version: "0.11.29" python-version: "3.12" enable-cache: true cache-dependency-glob: "uv.lock" @@ -92,6 +93,7 @@ jobs: - uses: actions/checkout@v4 - uses: astral-sh/setup-uv@v6 with: + version: "0.11.29" python-version: "3.12" enable-cache: true cache-dependency-glob: "uv.lock" @@ -113,6 +115,7 @@ jobs: - uses: actions/checkout@v4 - uses: astral-sh/setup-uv@v6 with: + version: "0.11.29" python-version: ${{ matrix.python-version }} enable-cache: true cache-dependency-glob: "uv.lock" @@ -136,6 +139,7 @@ jobs: - uses: actions/checkout@v4 - uses: astral-sh/setup-uv@v6 with: + version: "0.11.29" python-version: "3.12" enable-cache: true cache-dependency-glob: "uv.lock" diff --git a/.github/workflows/perf.yml b/.github/workflows/perf.yml index a7131216..1ece41e2 100644 --- a/.github/workflows/perf.yml +++ b/.github/workflows/perf.yml @@ -41,6 +41,7 @@ jobs: - uses: astral-sh/setup-uv@v6 with: + version: "0.11.29" python-version: "3.12" enable-cache: true cache-dependency-glob: "uv.lock" From adfc7bc67e1aabffb82df7a2d9dddfebbf201795 Mon Sep 17 00:00:00 2001 From: Greg Clark Date: Fri, 17 Jul 2026 10:07:33 -0400 Subject: [PATCH 11/14] fix(libsy-proxy): adapt demo client to renamed RoutedLlmClient trait Signed-off-by: Greg Clark --- crates/libsy/src/algorithms/rand.rs | 4 +++- demo/libsy-proxy/src/main.rs | 22 +++++++++++++--------- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/crates/libsy/src/algorithms/rand.rs b/crates/libsy/src/algorithms/rand.rs index f75c51b2..f4f1bdf5 100644 --- a/crates/libsy/src/algorithms/rand.rs +++ b/crates/libsy/src/algorithms/rand.rs @@ -138,7 +138,9 @@ mod tests { use switchyard_protocol::{completion_text, text_request, text_response}; use crate::affinity::{Affinity, SessionAffinity, SubAgentAffinity}; - use crate::{AgentContext, LlmResponse, LlmTarget, Metadata, Request, RoutedLlmClient, Signals}; + use crate::{ + AgentContext, LlmResponse, LlmTarget, Metadata, Request, RoutedLlmClient, Signals, + }; /// Echoes the selected target so tests can inspect which target was called. struct EchoClient; diff --git a/demo/libsy-proxy/src/main.rs b/demo/libsy-proxy/src/main.rs index 4886fcb2..1a10150d 100644 --- a/demo/libsy-proxy/src/main.rs +++ b/demo/libsy-proxy/src/main.rs @@ -17,8 +17,8 @@ use serde_json::Value; use libsy::affinity::{metadata_from_headers, SubAgentAffinity}; use libsy::{ - Algorithm, Context, Decision, LlmClient, LlmResponse, LlmTarget, LlmTargetSet, RandomAlgo, - Request, Response, RoutedRequest, + Algorithm, Context, Decision, LlmResponse, LlmTarget, LlmTargetSet, RandomAlgo, Request, + Response, RoutedLlmClient, }; use switchyard_components::OpenAiPassthroughBackend; use switchyard_components_v2::{Profile, ProfileInput, ProfileResponse, RoutingMetadata}; @@ -48,15 +48,19 @@ struct SwitchyardBackendClient { } #[async_trait] -impl LlmClient for SwitchyardBackendClient { - async fn call(&self, routed: RoutedRequest) -> std::result::Result { - let target = routed.decision.selected_model(); +impl RoutedLlmClient for SwitchyardBackendClient { + async fn call( + &self, + _ctx: Context, + request: Request, + decision: Arc, + ) -> std::result::Result { + let target = decision.selected_model(); let model = self .model_ids .get(target) .ok_or_else(|| format!("no provider model configured for target '{target}'"))?; - let chat_request = - chat_request_for_call(&routed.request, model, self.translation.as_ref())?; + let chat_request = chat_request_for_call(&request, model, self.translation.as_ref())?; let mut ctx = ProxyContext::new(); let response = self @@ -74,7 +78,7 @@ impl LlmClient for SwitchyardBackendClient { Ok(Response { llm_response: LlmResponse::Agg(decoded.response), - metadata: routed.request.metadata, + metadata: request.metadata, }) } } @@ -231,7 +235,7 @@ fn build_algorithm() -> Result> { backend, model_ids, translation, - }) as Arc; + }) as Arc; let target = |name: &str| LlmTarget { semantic_name: name.to_string(), llm_client: Some(Arc::clone(&client)), From 2318fec2c2f20dbe98e6d27329f72a84b146077b Mon Sep 17 00:00:00 2001 From: Greg Clark Date: Fri, 17 Jul 2026 10:51:22 -0400 Subject: [PATCH 12/14] feat(libsy-server): enable affinitiy Signed-off-by: Greg Clark --- crates/libsy-server/src/cli.rs | 50 ++++++++++++++++++++++++++++++++-- crates/libsy-server/src/lib.rs | 48 ++++++++++++++++++++------------ crates/libsy/src/affinity.rs | 17 ++++++++++++ 3 files changed, 95 insertions(+), 20 deletions(-) diff --git a/crates/libsy-server/src/cli.rs b/crates/libsy-server/src/cli.rs index 05a0798f..5e4bb5f2 100644 --- a/crates/libsy-server/src/cli.rs +++ b/crates/libsy-server/src/cli.rs @@ -8,7 +8,8 @@ use std::env; use std::net::{IpAddr, Ipv4Addr, SocketAddr}; use std::sync::Arc; -use clap::Parser; +use clap::{Parser, ValueEnum}; +use libsy::affinity::{Affinity, SessionAffinity, SubAgentAffinity}; use libsy::{Algorithm, LlmTarget, LlmTargetSet, RandomAlgo, RoutedLlmClient}; use switchyard_llm_client::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient}; use switchyard_translation::WireFormat; @@ -55,6 +56,14 @@ pub struct Args { #[arg(long)] pub log_routing: bool, + /// Enable model affinity so related requests reuse a tier instead of routing + /// each one independently. Bare `--affinity` uses session affinity; pass + /// `subagent` to pin sub-agents only. Off when omitted. Requires the harness to + /// send identifying headers (session/agent ids); otherwise it falls open to + /// random routing. + #[arg(long, value_enum, num_args = 0..=1, default_missing_value = "session")] + pub affinity: Option, + /// Host address to bind. #[arg(long, default_value_t = DEFAULT_HOST)] pub host: IpAddr, @@ -64,6 +73,26 @@ pub struct Args { pub port: u16, } +/// Which affinity policy `--affinity` enables. +#[derive(Clone, Copy, Debug, ValueEnum)] +pub enum AffinityKind { + /// Pin every request sharing a session id to one tier. + Session, + /// Pin each identified sub-agent to one tier across its turns; root-agent + /// turns keep routing randomly. + Subagent, +} + +impl AffinityKind { + /// Builds the affinity policy this kind selects. + fn policy(self) -> Arc { + match self { + AffinityKind::Session => Arc::new(SessionAffinity::new()), + AffinityKind::Subagent => Arc::new(SubAgentAffinity::new()), + } + } +} + impl Args { /// Builds server state (and the bind address) from the args and env keys. /// @@ -94,7 +123,14 @@ impl Args { target(&self.weak, client.clone()), target(&self.strong, client.clone()), ]); - let algorithm: Arc = Arc::new(RandomAlgo::new(targets)); + // Optionally retain selections per the affinity policy; without it every + // request routes independently. + let router = RandomAlgo::new(targets); + let router = match self.affinity { + Some(kind) => router.with_affinity(kind.policy()), + None => router, + }; + let algorithm: Arc = Arc::new(router); let addr = SocketAddr::new(self.host, self.port); Ok((ProxyState::new(algorithm, self.log_routing), addr)) @@ -106,13 +142,14 @@ pub async fn run(args: Args) -> Result<(), String> { let base_url = args.base_url.clone(); let weak = (args.weak.clone(), args.weak_format); let strong = (args.strong.clone(), args.strong_format); + let affinity = args.affinity; let (state, addr) = args.build()?; let listener = TcpListener::bind(addr) .await .map_err(|error| format!("failed to bind {addr}: {error}"))?; let bound = listener.local_addr().map_err(|error| error.to_string())?; - eprintln!("{}", banner(bound, &base_url, &weak, &strong)); + eprintln!("{}", banner(bound, &base_url, &weak, &strong, affinity)); axum::serve(listener, build_router(state)) .with_graceful_shutdown(shutdown_signal()) @@ -175,12 +212,19 @@ fn banner( base_url: &str, weak: &(String, WireFormat), strong: &(String, WireFormat), + affinity: Option, ) -> String { + let affinity = match affinity { + Some(AffinityKind::Session) => "session", + Some(AffinityKind::Subagent) => "subagent", + None => "off", + }; format!( "libsy-server\n \ listening: http://{addr}\n \ upstream: {base_url}\n \ random routing: weak {} ({}), strong {} ({})\n \ + affinity: {affinity}\n \ serving model: {SERVED_MODEL}\n \ endpoints: GET /health, GET /v1/models, POST /v1/chat/completions, \ POST /v1/messages, POST /v1/responses\n \ diff --git a/crates/libsy-server/src/lib.rs b/crates/libsy-server/src/lib.rs index 2483fd27..c21304f1 100644 --- a/crates/libsy-server/src/lib.rs +++ b/crates/libsy-server/src/lib.rs @@ -26,10 +26,11 @@ use axum::http::{HeaderMap, StatusCode}; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::{Json, Router}; +use libsy::affinity::metadata_from_headers; use libsy::Algorithm; use serde_json::{json, Value}; use switchyard_llm_client::LlmClientError; -use switchyard_protocol::{Context, LlmResponse, Metadata, Request}; +use switchyard_protocol::{Context, LlmResponse, Request}; use switchyard_translation::{decode_request, encode_buffered_response, encode_stream, WireFormat}; use crate::sse::frame_stream; @@ -130,24 +131,19 @@ async fn handle( // leaks the upstream id the algorithm's target resolved. let requested_model = llm_request.model.clone(); - // Caller headers ride along as metadata; the client drops reserved/auth - // headers and injects the backend's real credential, so the sentinel key a - // coding agent sends is filtered automatically. `wire_format` is left unset so - // the upstream call uses the chosen tier's own format, not the inbound one — - // the response is translated back to `inbound` below. + // Normalize harness headers into routing metadata (session/agent/task ids and + // sub-agent context) so an affinity policy can key on them. The forwarded + // headers ride along too; the client drops reserved/auth headers and injects + // the backend's real credential, so the sentinel key a coding agent sends is + // filtered automatically. `wire_format` stays unset so the upstream call uses + // the chosen tier's own format — the response is translated back to `inbound`. + let mut metadata = metadata_from_headers(&multi_headers(&headers)); + let session = metadata.session_id.clone().unwrap_or("none".to_string()); + metadata.http_headers = Some(normalized_headers(&headers)); let request = Request { llm_request, raw_request: Some(body), - metadata: Some(Metadata { - session_id: None, - agent_id: None, - task_id: None, - correlation_id: None, - extra_metadata: None, - http_headers: Some(normalized_headers(&headers)), - wire_format: None, - agent_context: None, - }), + metadata: Some(metadata), }; let (trace, response) = match state @@ -163,7 +159,10 @@ async fn handle( // The trace's first decision is the routing choice; log which tier served. if state.log_routing { if let Some(decision) = trace.first() { - eprintln!("[route] inbound={inbound} -> {}", decision.selected_model()); + eprintln!( + "[route][session={session}] inbound={inbound} -> {}", + decision.selected_model() + ); } } @@ -194,6 +193,21 @@ fn map_run_error(error: Box) -> Response { } } +// Collects every UTF-8 header value per lowercased name — the multi-valued shape +// `metadata_from_headers` expects for harness metadata normalization. +fn multi_headers(headers: &HeaderMap) -> BTreeMap> { + let mut collected: BTreeMap> = BTreeMap::new(); + for (name, value) in headers { + if let Ok(value) = value.to_str() { + collected + .entry(name.as_str().to_ascii_lowercase()) + .or_default() + .push(value.to_string()); + } + } + collected +} + // Lowercases header names and keeps the first UTF-8 value for each. fn normalized_headers(headers: &HeaderMap) -> BTreeMap { let mut normalized = BTreeMap::new(); diff --git a/crates/libsy/src/affinity.rs b/crates/libsy/src/affinity.rs index 7cc6f92f..12c94d7f 100644 --- a/crates/libsy/src/affinity.rs +++ b/crates/libsy/src/affinity.rs @@ -213,6 +213,7 @@ pub fn metadata_from_headers(headers: &BTreeMap>) -> Metadat Metadata { session_id: first_some([ header(headers, "x-switchyard-session-id"), + header(headers, "x-claude-code-session-id"), header(headers, "x-nemo-relay-session-id"), codex.session_id, header(headers, "session-id"), @@ -399,6 +400,22 @@ mod tests { ); } + #[test] + fn normalizes_claude_code_session_header() { + // Claude Code identifies a session with `x-claude-code-session-id`; session + // affinity keys on it so a whole CLI session pins to one tier. + let headers = BTreeMap::from([( + "x-claude-code-session-id".to_string(), + vec!["fb46caae-eac6-4f5f-83fd-8fc8f5743abb".to_string()], + )]); + + let metadata = metadata_from_headers(&headers); + assert_eq!( + metadata.session_id.as_deref(), + Some("fb46caae-eac6-4f5f-83fd-8fc8f5743abb") + ); + } + #[test] fn normalizes_relay_and_dynamo_child_headers() { let headers = BTreeMap::from([ From c3909f63d5874ed3af979d9d60c0279ce7298493 Mon Sep 17 00:00:00 2001 From: Greg Clark Date: Fri, 17 Jul 2026 10:55:56 -0400 Subject: [PATCH 13/14] docs(libsy-server): readme for proxy and cc usage Signed-off-by: Greg Clark --- crates/libsy-server/README.md | 155 ++++++++++++++++++++++++++-------- 1 file changed, 118 insertions(+), 37 deletions(-) diff --git a/crates/libsy-server/README.md b/crates/libsy-server/README.md index 2ff11d6f..34f173ff 100644 --- a/crates/libsy-server/README.md +++ b/crates/libsy-server/README.md @@ -2,59 +2,140 @@ A minimal [axum](https://github.com/tokio-rs/axum) server that exposes the OpenAI Chat Completions, OpenAI Responses, and Anthropic Messages wire APIs and routes -every request through a libsy [`RandomAlgo`](../libsy), which picks a weak- or -strong-tier upstream uniformly at random and serves it via -[`switchyard-llm-client`](../libsy-llm-client). +every request through a libsy [`RandomAlgo`](../libsy), which picks a **weak-** or +**strong-tier** upstream and serves it via +[`switchyard-llm-client`](../libsy-llm-client). It is designed to sit in front of a +coding agent (e.g. Claude Code) for A/B-style routing between two models. This crate is only the HTTP surface: routing, header normalization, SSE framing, -and error mapping. Each request is decoded to Switchyard's neutral IR, run -through the routing algorithm — which selects a target and calls the upstream — -and the response is encoded back into the same wire format (buffered JSON or a -stream of wire events). It mirrors `switchyard-server` but swaps the -profile-chain executor for a libsy algorithm. +and error mapping. Each request is decoded to Switchyard's neutral IR, run through +the routing algorithm — which selects a tier and calls the upstream — and the +response is encoded back into the inbound wire format (buffered JSON or a stream of +wire events). It mirrors `switchyard-server` but swaps the profile-chain executor +for a libsy algorithm. -## Endpoints +## Quickstart -| Method | Path | Purpose | -|--------|------------------------|----------------------------------| -| POST | `/v1/chat/completions` | OpenAI Chat Completions inbound | -| POST | `/v1/messages` | Anthropic Messages inbound | -| POST | `/v1/responses` | OpenAI Responses inbound | -| GET | `/v1/models` | Lists the single served model | -| GET | `/health` | Liveness check | +```bash +# 1. Provide upstream credentials (whichever providers your tiers use). +export OPENAI_API_KEY="sk-..." # for openai-chat / openai-responses tiers +export ANTHROPIC_API_KEY="sk-ant-..." # for anthropic-messages tiers + +# 2. Run the proxy: a weak tier and a strong tier over one upstream gateway. +cargo run -p libsy-server -- \ + --base-url https://api.openai.com/v1 \ + --weak gpt-4o-mini --weak-format openai-chat \ + --strong gpt-4o --strong-format openai-chat \ + --port 4000 + +# 3. Sanity-check it. +curl -s http://localhost:4000/health +curl -s http://localhost:4000/v1/chat/completions \ + -H 'content-type: application/json' \ + -d '{"model":"anything","messages":[{"role":"user","content":"say hi"}]}' +``` + +The server advertises a single served model id, `switchyard`. The `model` field a +client sends is **ignored for routing** (the algorithm picks the tier) and the +response `model` is always restamped to the id the client asked for — the real +upstream model id never leaks. + +## Point Claude Code at it + +Claude Code speaks the Anthropic Messages API. Point it at the proxy with two +environment variables and run as normal: -Clients address the server's served model id, `switchyard`; the response is -always restamped with that id rather than the real upstream model. +```bash +export ANTHROPIC_BASE_URL=http://localhost:4000 +export ANTHROPIC_API_KEY=sentinel # any non-empty value; see below + +claude # interactive +claude -p "refactor this function" # one-shot / scriptable +``` + +- **The API key is a sentinel.** The proxy drops the caller's `Authorization` / + `x-api-key` header and injects the *backend's* real credential (from + `OPENAI_API_KEY` / `ANTHROPIC_API_KEY`), so the key Claude Code sends never + reaches an upstream and can be any non-empty placeholder. +- **The model name doesn't matter.** Claude Code sends e.g. `claude-3-5-sonnet`; + the proxy routes by tier and restamps the response to that same name. +- **Streaming, tool use, and multi-turn all work** — inbound Anthropic requests are + translated to each tier's upstream format and the response is translated back. ## Tiers, formats, and translation -Two tiers — `weak` and `strong` — are the random-routing targets. Each has its -own upstream wire format (`--weak-format` / `--strong-format`); every inbound -request (any of the three served endpoints) is decoded to the neutral IR, -re-encoded to the chosen tier's format for the upstream call, and translated back -to the inbound format for the client. Each tier authenticates with its format's -provider key — OpenAI formats from `OPENAI_API_KEY`, Anthropic Messages from -`ANTHROPIC_API_KEY` — at the shared `--base-url`; startup errors if a tier's key -is unset. +Two tiers — `weak` and `strong` — are the random-routing targets. Each has its own +upstream wire format (`--weak-format` / `--strong-format`); every inbound request +(any of the three served endpoints) is decoded to the neutral IR, re-encoded to the +chosen tier's format for the upstream call, and translated back to the inbound +format for the client. So a tier can speak a *different* wire format than the client +— e.g. Claude Code (Anthropic) routed to an OpenAI Responses model. + +Each tier authenticates with its format's provider key — OpenAI formats from +`OPENAI_API_KEY`, Anthropic Messages from `ANTHROPIC_API_KEY` — at the shared +`--base-url`; startup errors if a tier's key is unset. Wire-format values: `openai-chat`, `openai-responses`, `anthropic-messages`. -## Running +## Affinity + +By default every request is routed independently, so a single agent session mixes +tiers request-by-request. Pass `--affinity` to instead **pin related requests to one +tier**: ```bash -export OPENAI_API_KEY="sk-..." # for openai-chat / openai-responses tiers -export ANTHROPIC_API_KEY="sk-..." # for anthropic-messages tiers +# Bare flag = session affinity: a whole session sticks to one tier. +cargo run -p libsy-server -- ... --affinity -cargo run -p libsy-server -- \ - --base-url https://api.openai.com/v1 \ - --weak gpt-4o-mini --weak-format openai-chat \ - --strong gpt-4o --strong-format openai-chat \ - --port 4000 +# Or pin only sub-agents (root-agent turns keep routing randomly): +cargo run -p libsy-server -- ... --affinity subagent +``` + +Affinity keys on identifying request headers, normalized into request metadata: + +- **`session`** keys on the session id. Claude Code's `X-Claude-Code-Session-Id` is + recognized, so a whole Claude Code session (and concurrent sessions independently) + pins to one tier. This makes A/B outcomes per-session and all-or-nothing rather + than a per-request blend. +- **`subagent`** keys on session + agent id for requests identified as sub-agents + (via Switchyard / Codex / Relay / Dynamo headers); root-agent turns stay random. + +If a request carries no identifying header, affinity **falls open** to normal random +routing. + +## Observability + +Add `--log-routing` to print the chosen tier (and session id) for each request to +stderr: + +``` +[route][session=fb46caae-…] inbound=anthropic_messages -> gpt-4o +[route][session=fb46caae-…] inbound=anthropic_messages -> gpt-4o ``` -Flags: `--base-url` (required), `--weak` / `--strong` (required tier model ids), -`--weak-format` / `--strong-format` (required tier wire formats), `--host` -(default `127.0.0.1`), `--port` (default `4000`). +With `--affinity`, this is how you confirm a session stayed pinned to one tier. + +## Flags + +| Flag | Required | Description | +|------|----------|-------------| +| `--base-url` | yes | Shared upstream base URL (e.g. `https://api.openai.com/v1`). | +| `--weak` / `--strong` | yes | Upstream model id for each tier. | +| `--weak-format` / `--strong-format` | yes | Wire format each tier's upstream speaks: `openai-chat`, `openai-responses`, `anthropic-messages`. | +| `--affinity []` | no | Pin related requests to a tier. Bare = `session`; or `subagent`. Off when omitted. | +| `--log-routing` | no | Log each request's routing decision to stderr. | +| `--host` | no | Bind address (default `127.0.0.1`). | +| `--port` | no | Bind port (default `4000`). | + +## Endpoints + +| Method | Path | Purpose | +|--------|------------------------|----------------------------------| +| POST | `/v1/chat/completions` | OpenAI Chat Completions inbound | +| POST | `/v1/messages` | Anthropic Messages inbound | +| POST | `/v1/responses` | OpenAI Responses inbound | +| GET | `/v1/models` | Lists the single served model | +| GET | `/health` | Liveness check | ## Testing From 7d2f6da14ead381dfcf4bcde943339154064f3db Mon Sep 17 00:00:00 2001 From: ayushag Date: Fri, 17 Jul 2026 10:27:44 -0700 Subject: [PATCH 14/14] chore: added claude code headers + e2e testing with claude code Signed-off-by: ayushag --- crates/libsy-server/README.md | 8 +++---- crates/libsy-server/src/lib.rs | 17 ++++++++++++++- crates/libsy/README.md | 4 +++- crates/libsy/src/affinity.rs | 40 +++++++++++++++++++++++++++++++--- 4 files changed, 60 insertions(+), 9 deletions(-) diff --git a/crates/libsy-server/README.md b/crates/libsy-server/README.md index 34f173ff..9d41b35a 100644 --- a/crates/libsy-server/README.md +++ b/crates/libsy-server/README.md @@ -98,7 +98,7 @@ Affinity keys on identifying request headers, normalized into request metadata: pins to one tier. This makes A/B outcomes per-session and all-or-nothing rather than a per-request blend. - **`subagent`** keys on session + agent id for requests identified as sub-agents - (via Switchyard / Codex / Relay / Dynamo headers); root-agent turns stay random. + (via Claude Code / Switchyard / Codex / Relay / Dynamo headers); root-agent turns stay random. If a request carries no identifying header, affinity **falls open** to normal random routing. @@ -109,11 +109,11 @@ Add `--log-routing` to print the chosen tier (and session id) for each request t stderr: ``` -[route][session=fb46caae-…] inbound=anthropic_messages -> gpt-4o -[route][session=fb46caae-…] inbound=anthropic_messages -> gpt-4o +[route][session=fb46caae-…][actor=root] inbound=anthropic_messages -> gpt-4o +[route][session=fb46caae-…][actor=subagent:agent-123] inbound=anthropic_messages -> gpt-4o ``` -With `--affinity`, this is how you confirm a session stayed pinned to one tier. +The actor label shows whether the request came from the root path or a specific child agent. ## Flags diff --git a/crates/libsy-server/src/lib.rs b/crates/libsy-server/src/lib.rs index c21304f1..c2e2c7bb 100644 --- a/crates/libsy-server/src/lib.rs +++ b/crates/libsy-server/src/lib.rs @@ -139,6 +139,12 @@ async fn handle( // the chosen tier's own format — the response is translated back to `inbound`. let mut metadata = metadata_from_headers(&multi_headers(&headers)); let session = metadata.session_id.clone().unwrap_or("none".to_string()); + let subagent_id = metadata + .agent_context + .as_deref() + .is_some_and(|context| context.is_subagent) + .then(|| metadata.agent_id.clone()) + .flatten(); metadata.http_headers = Some(normalized_headers(&headers)); let request = Request { llm_request, @@ -159,8 +165,9 @@ async fn handle( // The trace's first decision is the routing choice; log which tier served. if state.log_routing { if let Some(decision) = trace.first() { + let actor = actor_label(subagent_id.as_deref()); eprintln!( - "[route][session={session}] inbound={inbound} -> {}", + "[route][session={session}][actor={actor}] inbound={inbound} -> {}", decision.selected_model() ); } @@ -183,6 +190,14 @@ async fn handle( } } +// Labels the request source for the demo's routing log. Claude Code only sends a +// child-agent id for sub-agent requests, so every other request is the root path. +fn actor_label(subagent_id: Option<&str>) -> String { + subagent_id + .map(|id| format!("subagent:{id}")) + .unwrap_or_else(|| "root".to_string()) +} + // Maps an algorithm run failure onto an HTTP envelope. A failed upstream model // call surfaces as a boxed [`LlmClientError`] (mapped like the direct path); // anything else is an internal proxy fault. diff --git a/crates/libsy/README.md b/crates/libsy/README.md index 6d9a52b6..0ba6e6e4 100644 --- a/crates/libsy/README.md +++ b/crates/libsy/README.md @@ -61,7 +61,9 @@ The first request from a child agent runs the random algorithm and atomically re its choice; later requests for the same `(session_id, agent_id)` reuse that target. With one configured target, that target is retained directly. `metadata_from_headers` converts harness-specific identity into neutral `Metadata`. -Supported inputs include Codex `session-id`, `thread-id`, +Supported inputs include Claude Code `x-claude-code-session-id`, +`x-claude-code-agent-id`, and `x-claude-code-parent-agent-id`; Codex `session-id`, +`thread-id`, `x-codex-parent-thread-id`, `x-openai-subagent`, and `x-codex-turn-metadata`; NeMo Relay `x-nemo-relay-session-id` / `x-nemo-relay-subagent-id`; Dynamo `x-dynamo-session-id` / `x-dynamo-parent-session-id`; and explicit diff --git a/crates/libsy/src/affinity.rs b/crates/libsy/src/affinity.rs index 12c94d7f..9d79db42 100644 --- a/crates/libsy/src/affinity.rs +++ b/crates/libsy/src/affinity.rs @@ -166,20 +166,24 @@ struct CodexTurnMetadata { /// Normalizes harness-specific request headers into libsy metadata. /// -/// Explicit `x-switchyard-*` headers win. NeMo Relay and Dynamo correlation -/// headers are accepted without linking either runtime. Codex's structured turn -/// metadata is preferred over its compatibility projections. +/// Explicit `x-switchyard-*` headers win. Claude Code, NeMo Relay, and Dynamo +/// correlation headers are accepted without linking their runtimes. Codex's structured +/// turn metadata is preferred over its compatibility projections. pub fn metadata_from_headers(headers: &BTreeMap>) -> Metadata { let codex = header(headers, CODEX_TURN_METADATA_HEADER) .and_then(|value| serde_json::from_str::(&value).ok()) .unwrap_or_default(); let switchyard_parent = header(headers, "x-switchyard-parent-agent-id"); + let claude_agent = header(headers, "x-claude-code-agent-id"); + let claude_parent = header(headers, "x-claude-code-parent-agent-id"); let relay_subagent = header(headers, "x-nemo-relay-subagent-id"); let dynamo_parent = header(headers, "x-dynamo-parent-session-id"); let codex_parent = header(headers, "x-codex-parent-thread-id"); let openai_subagent = header(headers, "x-openai-subagent"); let inferred_subagent = switchyard_parent.is_some() + || claude_agent.is_some() + || claude_parent.is_some() || relay_subagent.is_some() || dynamo_parent.is_some() || codex.parent_thread_id.is_some() @@ -195,6 +199,7 @@ pub fn metadata_from_headers(headers: &BTreeMap>) -> Metadat is_subagent, parent_agent_id: first_some([ switchyard_parent, + claude_parent, dynamo_parent, codex.parent_thread_id, codex_parent, @@ -220,6 +225,7 @@ pub fn metadata_from_headers(headers: &BTreeMap>) -> Metadat ]), agent_id: first_some([ header(headers, "x-switchyard-agent-id"), + claude_agent, relay_subagent, header(headers, "x-dynamo-session-id"), codex.thread_id, @@ -416,6 +422,34 @@ mod tests { ); } + #[test] + fn normalizes_claude_code_subagent_headers() { + let headers = BTreeMap::from([ + ( + "x-claude-code-session-id".to_string(), + vec!["root-session".to_string()], + ), + ( + "x-claude-code-agent-id".to_string(), + vec!["child-agent".to_string()], + ), + ( + "x-claude-code-parent-agent-id".to_string(), + vec!["root-agent".to_string()], + ), + ]); + + let metadata = metadata_from_headers(&headers); + assert_eq!(metadata.session_id.as_deref(), Some("root-session")); + assert_eq!(metadata.agent_id.as_deref(), Some("child-agent")); + let agent = metadata.agent_context.as_deref(); + assert_eq!(agent.map(|value| value.is_subagent), Some(true)); + assert_eq!( + agent.and_then(|value| value.parent_agent_id.as_deref()), + Some("root-agent") + ); + } + #[test] fn normalizes_relay_and_dynamo_child_headers() { let headers = BTreeMap::from([