diff --git a/.gitignore b/.gitignore index 525a9c40..61814f76 100644 --- a/.gitignore +++ b/.gitignore @@ -85,3 +85,4 @@ CHANGELOG.md # Relay /.nemo-relay +/artifacts/ diff --git a/ATTRIBUTIONS-Rust.md b/ATTRIBUTIONS-Rust.md index 4bdfa809..9a773274 100644 --- a/ATTRIBUTIONS-Rust.md +++ b/ATTRIBUTIONS-Rust.md @@ -37612,6 +37612,87 @@ SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ``` +## switchyard-translation - 0.1.0 +**Repository URL**: https://github.com/NVIDIA-NeMo/Switchyard +**License Type(s)**: Apache-2.0 +### License: https://spdx.org/licenses/Apache-2.0.html +``` +Apache License +Version 2.0, January 2004 +http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + +"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. + +"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. + +"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. + +"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. + +"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. + +"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. + +"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). + +"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. + +"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." + +"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: + + (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. + + You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + +To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. + +Copyright [yyyy] [name of copyright owner] + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +``` + ## syn - 2.0.117 **Repository URL**: https://github.com/dtolnay/syn **License Type(s)**: Apache-2.0 diff --git a/Cargo.lock b/Cargo.lock index f7398ba7..c1e6dada 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1547,6 +1547,7 @@ dependencies = [ "nemo-relay", "nemo-relay-adaptive", "nemo-relay-pii-redaction", + "nemo-relay-switchyard", "opentelemetry", "opentelemetry_sdk", "percent-encoding", @@ -1646,6 +1647,23 @@ dependencies = [ "uuid", ] +[[package]] +name = "nemo-relay-switchyard" +version = "0.6.0" +dependencies = [ + "async-stream", + "axum", + "futures-util", + "nemo-relay", + "reqwest", + "schemars", + "serde", + "serde_json", + "switchyard-translation", + "tokio", + "uuid", +] + [[package]] name = "nemo-relay-types" version = "0.6.0" @@ -2924,6 +2942,16 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "switchyard-translation" +version = "0.1.0" +source = "git+https://github.com/NVIDIA-NeMo/Switchyard.git?rev=8f9db9a6a47f848cdff1d262276ba25a8ae9cbc8#8f9db9a6a47f848cdff1d262276ba25a8ae9cbc8" +dependencies = [ + "serde", + "serde_json", + "thiserror 2.0.18", +] + [[package]] name = "syn" version = "2.0.117" diff --git a/Cargo.toml b/Cargo.toml index 73c5ea53..a4f4eeef 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,6 +10,7 @@ members = [ "crates/worker", "crates/adaptive", "crates/pii-redaction", + "crates/switchyard", "crates/cli", # Language Bindings "crates/python", @@ -33,6 +34,8 @@ nemo-relay-worker-proto = { version = "0.6.0", path = "crates/worker-proto" } nemo-relay-worker = { version = "0.6.0", path = "crates/worker" } nemo-relay-adaptive = { version = "0.6.0", path = "crates/adaptive" } nemo-relay-pii-redaction = { version = "0.6.0", path = "crates/pii-redaction" } +nemo-relay-switchyard = { version = "0.6.0", path = "crates/switchyard" } +switchyard-translation = { git = "https://github.com/NVIDIA-NeMo/Switchyard.git", rev = "8f9db9a6a47f848cdff1d262276ba25a8ae9cbc8" } nemo-relay-ffi = { version = "0.6.0", path = "crates/ffi" } nemo-relay-cli = { version = "0.6.0", path = "crates/cli" } opentelemetry = { version = "0.31", default-features = false } diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index b38d5741..b4da6573 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -23,11 +23,13 @@ workspace = true [features] default = ["atof-streaming"] atof-streaming = ["nemo-relay/atof-streaming"] +switchyard = ["dep:nemo-relay-switchyard"] [dependencies] nemo-relay = { workspace = true, features = ["guardrails-remote", "object-store", "otel", "openinference", "worker-grpc"] } nemo-relay-adaptive = { workspace = true, features = ["redis-backend"] } nemo-relay-pii-redaction.workspace = true +nemo-relay-switchyard = { workspace = true, optional = true, features = ["schema"] } async-stream = "0.3" axum = "0.8" base64 = "0.22" @@ -61,3 +63,8 @@ opentelemetry = { workspace = true, features = ["trace"] } opentelemetry_sdk = { workspace = true, features = ["trace", "testing"] } tempfile = "3" tower = { version = "0.5", features = ["util"] } + +[[test]] +name = "switchyard_process_e2e" +path = "tests/switchyard_process_e2e.rs" +required-features = ["switchyard"] diff --git a/crates/cli/src/doctor.rs b/crates/cli/src/doctor.rs index 60edb4f4..f8699962 100644 --- a/crates/cli/src/doctor.rs +++ b/crates/cli/src/doctor.rs @@ -20,6 +20,10 @@ use nemo_relay::observability::plugin_component::OBSERVABILITY_PLUGIN_KIND; use nemo_relay::plugin::{DiagnosticLevel, PluginConfig, validate_plugin_config}; use nemo_relay_adaptive::plugin_component::register_adaptive_component; use nemo_relay_pii_redaction::component::register_pii_redaction_component; +#[cfg(feature = "switchyard")] +use nemo_relay_switchyard::{ + register_switchyard_component, validate_switchyard_atof_configuration, +}; use serde::Serialize; use serde_json::{Value, json}; use tokio::time::timeout; @@ -678,6 +682,24 @@ async fn collect_observability(gateway: &GatewayConfig) -> Vec { }); return checks; } + #[cfg(feature = "switchyard")] + if let Err(error) = register_switchyard_component() { + checks.push(Check { + name: "Switchyard plugin", + status: Status::Fail, + details: format!("registration failed: {error}"), + }); + return checks; + } + #[cfg(feature = "switchyard")] + if let Err(error) = validate_switchyard_atof_configuration(&plugin_config) { + checks.push(Check { + name: "Switchyard ATOF", + status: Status::Fail, + details: error, + }); + return checks; + } let report = validate_plugin_config(&plugin_config); if report.diagnostics.is_empty() { checks.push(Check { @@ -982,18 +1004,45 @@ async fn probe_atof_endpoint(index: usize, endpoint: &Value) -> Check { } fn endpoint_headers(endpoint: &Value) -> Result, String> { - let Some(headers) = endpoint.get("headers") else { - return Ok(Vec::new()); - }; - let Some(object) = headers.as_object() else { - return Err("headers must be an object of string values".into()); - }; - let mut out = Vec::with_capacity(object.len()); - for (key, value) in object { - let Some(value) = value.as_str() else { - return Err(format!("headers.{key} must be a string")); + let mut out = Vec::new(); + let mut names = std::collections::HashSet::new(); + if let Some(headers) = endpoint.get("headers") { + let Some(object) = headers.as_object() else { + return Err("headers must be an object of string values".into()); + }; + for (key, value) in object { + let name = reqwest::header::HeaderName::from_bytes(key.as_bytes()) + .map_err(|error| error.to_string())?; + let Some(value) = value.as_str() else { + return Err(format!("headers.{key} must be a string")); + }; + names.insert(name); + out.push((key.clone(), value.to_string())); + } + } + if let Some(header_env) = endpoint.get("header_env") { + let Some(object) = header_env.as_object() else { + return Err("header_env must be an object of string values".into()); }; - out.push((key.clone(), value.to_string())); + for (key, variable) in object { + let name = reqwest::header::HeaderName::from_bytes(key.as_bytes()) + .map_err(|error| error.to_string())?; + if names.contains(&name) { + return Err(format!( + "header {key:?} cannot appear in both headers and header_env" + )); + } + let Some(variable) = variable.as_str() else { + return Err(format!("header_env.{key} must be a string")); + }; + let value = std::env::var(variable) + .map_err(|_| format!("environment variable {variable:?} is not set"))?; + if value.trim().is_empty() { + return Err(format!("environment variable {variable:?} is blank")); + } + names.insert(name); + out.push((key.clone(), value)); + } } Ok(out) } diff --git a/crates/cli/src/error.rs b/crates/cli/src/error.rs index ba019ae7..0d5dbb3b 100644 --- a/crates/cli/src/error.rs +++ b/crates/cli/src/error.rs @@ -4,7 +4,7 @@ use axum::Json; use axum::http::StatusCode; use axum::response::{IntoResponse, Response}; -use nemo_relay::error::FlowError; +use nemo_relay::error::{FlowError, UpstreamFailure}; use serde::Serialize; use serde_json::{Map, Value, json}; use strum::Display; @@ -36,6 +36,8 @@ pub(crate) enum CliError { PayloadTooLarge(String), #[error("gateway upstream error: {0}")] Upstream(#[from] reqwest::Error), + #[error("{0}")] + ProviderFailure(UpstreamFailure), #[error("http error: {0}")] Http(#[from] http::Error), #[error("io error: {0}")] @@ -87,16 +89,21 @@ impl CliError { impl IntoResponse for CliError { // Maps gateway errors into a compact JSON HTTP response. Bad hook payloads are client errors, - // upstream gateway failures are bad gateway responses, and local install/config/runtime faults - // remain internal errors so callers do not mistake them for agent policy decisions. + // network-level upstream failures are bad gateway responses, provider failures mirror the + // upstream status when available, and local install/config/runtime faults remain internal + // errors so callers do not mistake them for agent policy decisions. fn into_response(self) -> Response { let message = self.to_string(); let guardrail_reason = self.guardrail_rejection_reason().map(ToOwned::to_owned); - let status = match (guardrail_reason.is_some(), self) { + let status = match (guardrail_reason.is_some(), &self) { (true, _) => StatusCode::FORBIDDEN, (false, Self::PayloadTooLarge(_)) => StatusCode::PAYLOAD_TOO_LARGE, (false, Self::InvalidPayload(_)) => StatusCode::BAD_REQUEST, (false, Self::Upstream(_)) => StatusCode::BAD_GATEWAY, + (false, Self::ProviderFailure(failure)) => failure + .status + .and_then(|status| StatusCode::from_u16(status).ok()) + .unwrap_or(StatusCode::BAD_GATEWAY), ( false, Self::Http(_) diff --git a/crates/cli/src/gateway.rs b/crates/cli/src/gateway.rs index ca0b7f03..06f4d188 100644 --- a/crates/cli/src/gateway.rs +++ b/crates/cli/src/gateway.rs @@ -1,6 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use std::collections::BTreeMap; use std::error::Error; use std::sync::{Arc, Mutex}; @@ -23,7 +24,7 @@ use nemo_relay::codec::resolve::{ }; use nemo_relay::codec::streaming::StreamingCodec; use nemo_relay::codec::traits::LlmResponseCodec; -use nemo_relay::error::FlowError; +use nemo_relay::error::{FlowError, UpstreamFailure, UpstreamFailureClass}; use serde_json::{Map, Value, json}; use crate::alignment::{self, GatewayRouteKind}; @@ -32,6 +33,11 @@ use crate::error::CliError; use crate::server::AppState; use crate::session::{GatewayCallPrep, LlmGatewayStart, SessionManager}; +const INTERNAL_DISPATCH_URL_HEADER: &str = "x-nemo-relay-internal-dispatch-url"; +const INTERNAL_DISPATCH_ROUTE_HEADER: &str = "x-nemo-relay-internal-dispatch-route"; +const INTERNAL_RETRY_AWARE_HEADER: &str = "x-nemo-relay-internal-retry-aware"; +const MAX_UPSTREAM_ERROR_BODY_BYTES: usize = 64 * 1024; + /// Proxies supported LLM API requests through NeMo Relay's managed execution pipeline. /// /// The gateway buffers the inbound body once, opens a managed LLM call against the resolved @@ -95,9 +101,11 @@ async fn prepare_gateway_request( .get("stream") .and_then(Value::as_bool) .unwrap_or(false); + let mut headers = parts.headers; + strip_internal_dispatch_headers(&mut headers); Ok(PreparedGatewayRequest { method: parts.method, - headers: parts.headers, + headers, path: parts.uri.path().to_string(), provider, upstream_url, @@ -359,6 +367,7 @@ fn build_buffered_func( let upstream_error = upstream_error.clone(); let response_bytes = response_bytes.clone(); Box::pin(async move { + let retry_aware = retry_aware_dispatch(&request); let response = match forward_upstream_request( &http, &method, @@ -373,6 +382,9 @@ fn build_buffered_func( Ok(response) => response, Err(error) => { let message = error.to_string(); + if retry_aware { + return Err(FlowError::Upstream(transport_failure(&error))); + } *upstream_error.lock().expect("upstream error lock poisoned") = Some(error); return Err(FlowError::Internal(message)); } @@ -383,10 +395,20 @@ fn build_buffered_func( Ok(bytes) => bytes, Err(error) => { let message = error.to_string(); + if retry_aware { + return Err(FlowError::Upstream(transport_failure(&error))); + } *upstream_error.lock().expect("upstream error lock poisoned") = Some(error); return Err(FlowError::Internal(message)); } }; + if retry_aware && !status.is_success() { + return Err(FlowError::Upstream(http_failure( + status, + &response_headers, + &bytes, + ))); + } let json = serde_json::from_slice::(&bytes) .unwrap_or_else(|_| json!({ "body_bytes": bytes.len() })); *upstream_info.lock().expect("upstream info lock poisoned") = @@ -521,6 +543,7 @@ fn build_streaming_func( let upstream_info = upstream_info.clone(); let upstream_error = upstream_error.clone(); Box::pin(async move { + let retry_aware = retry_aware_dispatch(&request); let response = match forward_upstream_request( &http, &method, @@ -535,12 +558,26 @@ fn build_streaming_func( Ok(response) => response, Err(error) => { let message = error.to_string(); + if retry_aware { + return Err(FlowError::Upstream(transport_failure(&error))); + } *upstream_error.lock().expect("upstream error lock poisoned") = Some(error); return Err(FlowError::Internal(message)); } }; let status = response.status(); let response_headers = response_headers(response.headers()); + if retry_aware && !status.is_success() { + let bytes = response + .bytes() + .await + .map_err(|error| FlowError::Upstream(transport_failure(&error)))?; + return Err(FlowError::Upstream(http_failure( + status, + &response_headers, + &bytes, + ))); + } *upstream_info.lock().expect("upstream info lock poisoned") = Some((status, response_headers)); let json_stream = sse_json_stream(response); @@ -725,25 +762,60 @@ async fn forward_upstream_request( effective_request: Option<&LlmRequest>, route: ProviderRoute, ) -> Result { - let (body_bytes, headers) = effective_upstream_request(body_bytes, headers, effective_request); - let sanitized = strip_replaceable_agent_auth_headers(&headers, route); - let mut upstream = http.request(method.clone(), url).body(body_bytes.clone()); + let effective = effective_dispatch_request(body_bytes, headers, effective_request, url, route); + let sanitized = strip_replaceable_agent_auth_headers(&effective.headers, effective.route); + let mut upstream = http + .request(method.clone(), &effective.url) + .body(effective.body_bytes.clone()); for (name, value) in &sanitized { if should_forward_request_header(name) { upstream = upstream.header(name, value); } } - upstream = inject_provider_auth(upstream, route, &sanitized); + upstream = inject_provider_auth(upstream, effective.route, &sanitized); upstream.send().await } +#[derive(Clone)] +struct EffectiveUpstreamRequest { + body_bytes: Bytes, + headers: HeaderMap, + url: String, + route: ProviderRoute, +} + +#[cfg(test)] fn effective_upstream_request( body_bytes: &Bytes, headers: &HeaderMap, effective_request: Option<&LlmRequest>, ) -> (Bytes, HeaderMap) { + let effective = effective_dispatch_request( + body_bytes, + headers, + effective_request, + "", + ProviderRoute::OpenAiChatCompletions, + ); + (effective.body_bytes, effective.headers) +} + +fn effective_dispatch_request( + body_bytes: &Bytes, + headers: &HeaderMap, + effective_request: Option<&LlmRequest>, + url: &str, + route: ProviderRoute, +) -> EffectiveUpstreamRequest { + let mut headers = headers.clone(); + strip_internal_dispatch_headers(&mut headers); let Some(request) = effective_request else { - return (body_bytes.clone(), headers.clone()); + return EffectiveUpstreamRequest { + body_bytes: body_bytes.clone(), + headers, + url: url.to_string(), + route, + }; }; let body_bytes = if request.content.is_null() { @@ -755,21 +827,65 @@ fn effective_upstream_request( eprintln!( "nemo-relay CLI gateway: failed to serialize rewritten LLM request body; forwarding original request: {error}" ); - return (body_bytes.clone(), headers.clone()); + return EffectiveUpstreamRequest { + body_bytes: body_bytes.clone(), + headers, + url: url.to_string(), + route, + }; } } }; - let mut headers = headers.clone(); + let mut override_url = None; + let mut override_route = None; for (name, value) in &request.headers { + if name.eq_ignore_ascii_case(INTERNAL_DISPATCH_URL_HEADER) { + override_url = json_header_string(value); + continue; + } + if name.eq_ignore_ascii_case(INTERNAL_DISPATCH_ROUTE_HEADER) { + override_route = json_header_string(value) + .and_then(|value| ProviderRoute::from_dispatch_override(&value)); + continue; + } let Ok(name) = HeaderName::from_bytes(name.as_bytes()) else { continue; }; + if is_internal_dispatch_header(&name) { + continue; + } let Some(value) = json_header_value(value) else { continue; }; headers.insert(name, value); } - (body_bytes, headers) + EffectiveUpstreamRequest { + body_bytes, + headers, + url: override_url.unwrap_or_else(|| url.to_string()), + route: override_route.unwrap_or(route), + } +} + +fn json_header_string(value: &Value) -> Option { + value + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn strip_internal_dispatch_headers(headers: &mut HeaderMap) { + headers.remove(INTERNAL_DISPATCH_URL_HEADER); + headers.remove(INTERNAL_DISPATCH_ROUTE_HEADER); + headers.remove(INTERNAL_RETRY_AWARE_HEADER); +} + +fn is_internal_dispatch_header(name: &HeaderName) -> bool { + matches!( + name.as_str(), + INTERNAL_DISPATCH_URL_HEADER | INTERNAL_DISPATCH_ROUTE_HEADER | INTERNAL_RETRY_AWARE_HEADER + ) } fn json_header_value(value: &Value) -> Option { @@ -878,10 +994,99 @@ fn translate_runtime_error(error: FlowError, upstream_error: &UpstreamErrorSlot) } match error { FlowError::GuardrailRejected(reason) => CliError::GuardrailRejected(reason), + FlowError::Upstream(failure) => CliError::ProviderFailure(failure), other => CliError::InvalidPayload(other.to_string()), } } +fn retry_aware_dispatch(request: &LlmRequest) -> bool { + request + .headers + .get(INTERNAL_RETRY_AWARE_HEADER) + .and_then(Value::as_str) + .is_some_and(|value| value.eq_ignore_ascii_case("true")) +} + +fn transport_failure(error: &reqwest::Error) -> UpstreamFailure { + UpstreamFailure { + status: error.status().map(|status| status.as_u16()), + body: bounded_error_body(error.to_string().as_bytes()), + headers: BTreeMap::new(), + class: if error.is_timeout() { + UpstreamFailureClass::Timeout + } else { + UpstreamFailureClass::Connection + }, + } +} + +fn http_failure(status: StatusCode, headers: &HeaderMap, body: &[u8]) -> UpstreamFailure { + let body = bounded_error_body(body); + let normalized = body.to_ascii_lowercase(); + let class = if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + UpstreamFailureClass::Authentication + } else if normalized.contains("context_length_exceeded") + || normalized.contains("context window") + || normalized.contains("too many tokens") + { + UpstreamFailureClass::ContextWindow + } else if normalized.contains("model_not_found") + || normalized.contains("model unavailable") + || normalized.contains("model_overloaded") + { + UpstreamFailureClass::ModelUnavailable + } else if matches!(status.as_u16(), 408 | 429 | 500 | 502 | 503 | 504) { + UpstreamFailureClass::RetryableStatus + } else if status.is_client_error() { + UpstreamFailureClass::InvalidRequest + } else { + UpstreamFailureClass::Other + }; + UpstreamFailure { + status: Some(status.as_u16()), + body, + headers: failure_headers(headers), + class, + } +} + +// Captures only safe provider metadata for retry and fallback diagnostics. This is intentionally +// separate from `response_headers`, which preserves response headers for ordinary downstream +// forwarding and therefore must not apply this failure-specific credential filter. +fn failure_headers(headers: &HeaderMap) -> BTreeMap { + headers + .iter() + .filter(|(name, _)| { + !is_hop_by_hop(name) + && *name != http::header::CONTENT_LENGTH + && !is_sensitive_response_header(name) + }) + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.as_str().to_string(), value.to_string())) + }) + .collect() +} + +fn is_sensitive_response_header(name: &HeaderName) -> bool { + matches!( + name.as_str(), + "set-cookie" + | "www-authenticate" + | "authorization" + | "cookie" + | "x-api-key" + | "api-key" + | "anthropic-api-key" + ) +} + +fn bounded_error_body(body: &[u8]) -> String { + String::from_utf8_lossy(&body[..body.len().min(MAX_UPSTREAM_ERROR_BODY_BYTES)]).into_owned() +} + /// Proxies OpenAI model-list requests without creating LLM runtime events. /// /// The route is registered as GET-only but still verifies the method so direct tests or future @@ -948,6 +1153,22 @@ impl ProviderRoute { } } + fn from_dispatch_override(value: &str) -> Option { + match value { + "openai_chat" + | "openai_chat_completions" + | "openai.chat_completions" + | "/v1/chat/completions" => Some(Self::OpenAiChatCompletions), + "openai_responses" | "openai.responses" | "/v1/responses" => { + Some(Self::OpenAiResponses) + } + "anthropic_messages" | "anthropic.messages" | "/v1/messages" => { + Some(Self::AnthropicMessages) + } + _ => None, + } + } + const fn provider_surface(self) -> Option { match self { Self::OpenAiResponses => Some(ProviderSurface::OpenAIResponses), diff --git a/crates/cli/src/launcher.rs b/crates/cli/src/launcher.rs index 4a9fa13d..40954bf5 100644 --- a/crates/cli/src/launcher.rs +++ b/crates/cli/src/launcher.rs @@ -344,9 +344,12 @@ impl PreparedRun { CodingAgent::Codex => run.prepare_codex(gateway_url), CodingAgent::Hermes => { if dry_run { - run.prepare_hermes_dry(resolved.agents.hermes.hooks_path.as_deref())?; + run.prepare_hermes_dry( + gateway_url, + resolved.agents.hermes.hooks_path.as_deref(), + )?; } else { - run.prepare_hermes(resolved.agents.hermes.hooks_path.as_deref())?; + run.prepare_hermes(gateway_url, resolved.agents.hermes.hooks_path.as_deref())?; } } } @@ -458,10 +461,18 @@ impl PreparedRun { // Hermes discovers hooks from `.hermes/config.yaml` instead of command-line flags. For // transparent runs, temporarily merge gateway hook-forward entries into the configured Hermes // hook file, then restore it after the child exits. - fn prepare_hermes(&mut self, hooks_path: Option<&std::path::Path>) -> Result<(), CliError> { + fn prepare_hermes( + &mut self, + gateway_url: &str, + hooks_path: Option<&std::path::Path>, + ) -> Result<(), CliError> { let path = hermes_hooks_path(hooks_path)?; let (had_original, backup_path) = backup_existing_hermes_hooks(&path)?; - write_merged_hermes_hooks(&path)?; + write_merged_hermes_config(&path, gateway_url)?; + self.env.push(( + "OPENAI_BASE_URL".into(), + hermes_openai_base_url(gateway_url), + )); self.env.push(("HERMES_ACCEPT_HOOKS".into(), "1".into())); self.notes.push(format!( "temporarily merged NeMo Relay hooks into {}", @@ -477,8 +488,16 @@ impl PreparedRun { // Records the Hermes hook file that would be patched during a real run without touching the // filesystem, preserving dry-run as an inspection-only operation. - fn prepare_hermes_dry(&mut self, hooks_path: Option<&std::path::Path>) -> Result<(), CliError> { + fn prepare_hermes_dry( + &mut self, + gateway_url: &str, + hooks_path: Option<&std::path::Path>, + ) -> Result<(), CliError> { let path = hermes_hooks_path(hooks_path)?; + self.env.push(( + "OPENAI_BASE_URL".into(), + hermes_openai_base_url(gateway_url), + )); self.env.push(("HERMES_ACCEPT_HOOKS".into(), "1".into())); self.notes.push(format!( "would temporarily merge NeMo Relay hooks into {}", @@ -882,9 +901,14 @@ fn backup_existing_hermes_hooks(path: &Path) -> Result<(bool, Option), Ok((true, Some(backup))) } +fn hermes_openai_base_url(gateway_url: &str) -> String { + format!("{}/v1", gateway_url.trim_end_matches('/')) +} + // Creates the Hermes config parent directory when needed, merges generated gateway hooks with any -// existing YAML config, and writes the patched YAML used for this transparent run. -fn write_merged_hermes_hooks(path: &Path) -> Result<(), CliError> { +// existing YAML config, and temporarily selects Relay as Hermes's custom OpenAI-compatible +// provider. The caller restores the complete original config after the transparent run. +fn write_merged_hermes_config(path: &Path, gateway_url: &str) -> Result<(), CliError> { if let Some(parent) = path.parent() { std::fs::create_dir_all(parent)?; } @@ -900,7 +924,41 @@ fn write_merged_hermes_hooks(path: &Path) -> Result<(), CliError> { &hook_forward_command(&transparent_hook_executable(), CodingAgent::Hermes), ), )?; - std::fs::write(path, contents)?; + let mut document = serde_yaml::from_str::(&contents).map_err(|error| { + CliError::Launch(format!("could not parse merged Hermes config: {error}")) + })?; + let root = document.as_mapping_mut().ok_or_else(|| { + CliError::Launch("merged Hermes config must be a YAML mapping".to_string()) + })?; + let model_key = serde_yaml::Value::String("model".into()); + let mut model = match root.remove(&model_key) { + Some(serde_yaml::Value::Mapping(model)) => model, + Some(serde_yaml::Value::String(default)) => serde_yaml::Mapping::from_iter([( + serde_yaml::Value::String("default".into()), + serde_yaml::Value::String(default), + )]), + Some(serde_yaml::Value::Null) | None => serde_yaml::Mapping::new(), + Some(_) => { + return Err(CliError::Launch( + "Hermes model config must be a string or mapping".into(), + )); + } + }; + model.insert( + serde_yaml::Value::String("provider".into()), + serde_yaml::Value::String("custom".into()), + ); + model.insert( + serde_yaml::Value::String("base_url".into()), + serde_yaml::Value::String(hermes_openai_base_url(gateway_url)), + ); + root.insert(model_key, serde_yaml::Value::Mapping(model)); + std::fs::write( + path, + serde_yaml::to_string(&document).map_err(|error| { + CliError::Launch(format!("could not render Hermes config: {error}")) + })?, + )?; Ok(()) } diff --git a/crates/cli/src/plugins.rs b/crates/cli/src/plugins.rs index be3676c5..f4cc7b8d 100644 --- a/crates/cli/src/plugins.rs +++ b/crates/cli/src/plugins.rs @@ -351,6 +351,11 @@ fn edit_component_field( edit_config_field(theme, &mut state.config, field)?; state.mark_config_touched(); } + #[cfg(feature = "switchyard")] + EditableComponent::Switchyard(state) => { + edit_config_field(theme, &mut state.config, field)?; + state.mark_config_touched(); + } } Ok(()) } diff --git a/crates/cli/src/plugins/config_io.rs b/crates/cli/src/plugins/config_io.rs index d3c55fc3..c34acd19 100644 --- a/crates/cli/src/plugins/config_io.rs +++ b/crates/cli/src/plugins/config_io.rs @@ -10,6 +10,10 @@ use nemo_relay::plugin::dynamic::DynamicPluginManifest; use nemo_relay::plugin::{ConfigPolicy, PluginConfig, validate_plugin_config}; use nemo_relay_adaptive::plugin_component::register_adaptive_component; use nemo_relay_pii_redaction::component::register_pii_redaction_component; +#[cfg(feature = "switchyard")] +use nemo_relay_switchyard::{ + register_switchyard_component, validate_switchyard_atof_configuration, +}; use serde::Serialize; use serde_json::{Map, Value}; @@ -718,6 +722,13 @@ pub(crate) fn validate_config(config: &PluginConfig) -> Result<(), CliError> { register_pii_redaction_component().map_err(|error| { CliError::Config(format!("PII redaction plugin registration failed: {error}")) })?; + #[cfg(feature = "switchyard")] + register_switchyard_component().map_err(|error| { + CliError::Config(format!("Switchyard plugin registration failed: {error}")) + })?; + #[cfg(feature = "switchyard")] + validate_switchyard_atof_configuration(config) + .map_err(|error| CliError::Config(format!("Switchyard ATOF validation failed: {error}")))?; let report = validate_plugin_config(config); if report.has_errors() { let messages = report diff --git a/crates/cli/src/plugins/editor_model.rs b/crates/cli/src/plugins/editor_model.rs index 91c42fad..50eb2fa5 100644 --- a/crates/cli/src/plugins/editor_model.rs +++ b/crates/cli/src/plugins/editor_model.rs @@ -14,6 +14,8 @@ use nemo_relay::plugins::nemo_guardrails::component::{ use nemo_relay_adaptive::AdaptiveConfig; use nemo_relay_adaptive::plugin_component::ADAPTIVE_PLUGIN_KIND; use nemo_relay_pii_redaction::component::{PII_REDACTION_PLUGIN_KIND, PiiRedactionConfig}; +#[cfg(feature = "switchyard")] +use nemo_relay_switchyard::{SWITCHYARD_PLUGIN_KIND, SwitchyardConfig}; use serde::Serialize; use serde::de::DeserializeOwned; use serde_json::{Map, Value, json}; @@ -38,6 +40,8 @@ pub(super) enum EditableComponent { Adaptive(Box>), NemoGuardrails(Box>), PiiRedaction(Box>), + #[cfg(feature = "switchyard")] + Switchyard(Box>), } impl EditableComponent { @@ -47,6 +51,8 @@ impl EditableComponent { Self::Adaptive(_) => "Adaptive", Self::NemoGuardrails(_) => "NeMo Guardrails", Self::PiiRedaction(_) => "PII Redaction", + #[cfg(feature = "switchyard")] + Self::Switchyard(_) => "Switchyard Decision API", } } @@ -56,6 +62,8 @@ impl EditableComponent { Self::Adaptive(_) => AdaptiveConfig::editor_schema().fields, Self::NemoGuardrails(_) => NeMoGuardrailsConfig::editor_schema().fields, Self::PiiRedaction(_) => PiiRedactionConfig::editor_schema().fields, + #[cfg(feature = "switchyard")] + Self::Switchyard(_) => SwitchyardConfig::editor_schema().fields, } } @@ -65,6 +73,8 @@ impl EditableComponent { Self::Adaptive(state) => state.enabled, Self::NemoGuardrails(state) => state.enabled, Self::PiiRedaction(state) => state.enabled, + #[cfg(feature = "switchyard")] + Self::Switchyard(state) => state.enabled, } } @@ -74,6 +84,8 @@ impl EditableComponent { Self::Adaptive(state) => state.toggle_enabled(), Self::NemoGuardrails(state) => state.toggle_enabled(), Self::PiiRedaction(state) => state.toggle_enabled(), + #[cfg(feature = "switchyard")] + Self::Switchyard(state) => state.toggle_enabled(), } } @@ -83,6 +95,8 @@ impl EditableComponent { Self::Adaptive(state) => state.set_enabled(enabled), Self::NemoGuardrails(state) => state.set_enabled(enabled), Self::PiiRedaction(state) => state.set_enabled(enabled), + #[cfg(feature = "switchyard")] + Self::Switchyard(state) => state.set_enabled(enabled), } } @@ -92,6 +106,8 @@ impl EditableComponent { Self::Adaptive(state) => state.reset_enabled(), Self::NemoGuardrails(state) => state.reset_enabled(), Self::PiiRedaction(state) => state.reset_enabled(), + #[cfg(feature = "switchyard")] + Self::Switchyard(state) => state.reset_enabled(), } } @@ -101,6 +117,8 @@ impl EditableComponent { Self::Adaptive(state) => adaptive_summary(state), Self::NemoGuardrails(state) => nemo_guardrails_summary(state), Self::PiiRedaction(state) => pii_redaction_summary(state), + #[cfg(feature = "switchyard")] + Self::Switchyard(state) => switchyard_summary(state), } } @@ -114,6 +132,10 @@ impl EditableComponent { Self::PiiRedaction(state) => { config_field_configured(&state.config, field).unwrap_or(false) } + #[cfg(feature = "switchyard")] + Self::Switchyard(state) => { + config_field_configured(&state.config, field).unwrap_or(false) + } } } @@ -135,6 +157,11 @@ impl EditableComponent { reset_config_field(&mut state.config, field)?; state.mark_config_touched(); } + #[cfg(feature = "switchyard")] + Self::Switchyard(state) => { + reset_config_field(&mut state.config, field)?; + state.mark_config_touched(); + } } Ok(()) } @@ -160,6 +187,11 @@ impl EditableComponent { remove_struct_field(&mut state.config, field.name)?; state.mark_config_touched(); } + #[cfg(feature = "switchyard")] + Self::Switchyard(state) => { + remove_struct_field(&mut state.config, field.name)?; + state.mark_config_touched(); + } } Ok(true) } @@ -170,6 +202,8 @@ impl EditableComponent { Self::Adaptive(state) => store_adaptive_state(config, state), Self::NemoGuardrails(state) => store_nemo_guardrails_state(config, state), Self::PiiRedaction(state) => store_pii_redaction_state(config, state), + #[cfg(feature = "switchyard")] + Self::Switchyard(state) => store_switchyard_state(config, state), } } } @@ -193,12 +227,21 @@ pub(super) enum ComponentMenuAction { pub(super) fn editable_components( config: &PluginConfig, ) -> Result, CliError> { - Ok(vec![ + let components = vec![ EditableComponent::Observability(Box::new(component_observability_state(config)?)), EditableComponent::Adaptive(Box::new(component_adaptive_state(config)?)), EditableComponent::NemoGuardrails(Box::new(component_nemo_guardrails_state(config)?)), EditableComponent::PiiRedaction(Box::new(component_pii_redaction_state(config)?)), - ]) + ]; + #[cfg(feature = "switchyard")] + let components = { + let mut components = components; + components.push(EditableComponent::Switchyard(Box::new( + component_switchyard_state(config)?, + ))); + components + }; + Ok(components) } pub(super) fn plugin_menu_items( @@ -403,6 +446,13 @@ pub(super) fn component_pii_redaction_state( component_editor_state(config, PII_REDACTION_PLUGIN_KIND, false) } +#[cfg(feature = "switchyard")] +pub(super) fn component_switchyard_state( + config: &PluginConfig, +) -> Result, CliError> { + component_editor_state(config, SWITCHYARD_PLUGIN_KIND, false) +} + pub(super) fn store_observability_state( config: &mut PluginConfig, state: &ComponentEditorState, @@ -467,6 +517,23 @@ pub(super) fn store_pii_redaction_state( Ok(()) } +#[cfg(feature = "switchyard")] +pub(super) fn store_switchyard_state( + config: &mut PluginConfig, + state: &ComponentEditorState, +) -> Result<(), CliError> { + if state.should_store(state.config_touched || switchyard_configured(&state.config)) { + store_component_editor_config( + config, + SWITCHYARD_PLUGIN_KIND, + state.enabled, + switchyard_config_map(&state.config)?, + merge_switchyard_editor_config, + ); + } + Ok(()) +} + fn store_component_editor_config( config: &mut PluginConfig, kind: &str, @@ -815,6 +882,24 @@ pub(super) fn pii_redaction_config_map( } } +#[cfg(feature = "switchyard")] +pub(super) fn switchyard_config_map( + config: &SwitchyardConfig, +) -> Result, CliError> { + let value = serde_json::to_value(config).map_err(serde_error)?; + match value { + Value::Object(mut map) => { + if is_version_one(map.get("version")) { + map.remove("version"); + } + Ok(map) + } + _ => Err(CliError::Config( + "switchyard config must serialize to an object".into(), + )), + } +} + pub(super) fn merge_observability_editor_config( existing: &mut Map, edited: Map, @@ -872,6 +957,22 @@ pub(super) fn merge_pii_redaction_editor_config( ); } +#[cfg(feature = "switchyard")] +pub(super) fn merge_switchyard_editor_config( + existing: &mut Map, + edited: Map, +) { + if is_version_one(existing.get("version")) { + existing.remove("version"); + } + merge_known_editor_object( + existing, + edited, + &nested_editor_keys(SwitchyardConfig::editor_schema()), + SwitchyardConfig::editor_schema(), + ); +} + fn is_version_one(value: Option<&Value>) -> bool { value.and_then(Value::as_u64) == Some(1) } @@ -1055,3 +1156,23 @@ pub(super) fn pii_redaction_summary(state: &ComponentEditorState bool { + !config.decision_profile_id.is_empty() || !config.targets.is_empty() +} + +#[cfg(feature = "switchyard")] +pub(super) fn switchyard_summary(state: &ComponentEditorState) -> String { + let profile = if state.config.decision_profile_id.is_empty() { + "unconfigured" + } else { + state.config.decision_profile_id.as_str() + }; + format!( + "component {}, profile {}, targets {}", + if state.enabled { "enabled" } else { "disabled" }, + profile, + state.config.targets.len() + ) +} diff --git a/crates/cli/src/server.rs b/crates/cli/src/server.rs index 386ab8a2..3fe0e73f 100644 --- a/crates/cli/src/server.rs +++ b/crates/cli/src/server.rs @@ -21,6 +21,10 @@ use nemo_relay::plugin::{ }; use nemo_relay_adaptive::plugin_component::register_adaptive_component; use nemo_relay_pii_redaction::component::register_pii_redaction_component; +#[cfg(feature = "switchyard")] +use nemo_relay_switchyard::{ + register_switchyard_component, validate_switchyard_atof_configuration, +}; use reqwest::Client; use serde_json::Value; use tokio::net::TcpListener; @@ -341,6 +345,10 @@ impl PluginActivation { register_pii_redaction_component().map_err(|error| { CliError::Config(format!("PII redaction plugin registration failed: {error}")) })?; + #[cfg(feature = "switchyard")] + register_switchyard_component().map_err(|error| { + CliError::Config(format!("Switchyard plugin registration failed: {error}")) + })?; let native_specs = dynamic_plugins .iter() .filter(|plugin| plugin.kind == DynamicPluginKind::RustDynamic) @@ -408,6 +416,10 @@ impl PluginActivation { config: plugin.config, }), ); + #[cfg(feature = "switchyard")] + validate_switchyard_atof_configuration(&plugin_config).map_err(|error| { + CliError::Config(format!("Switchyard ATOF validation failed: {error}")) + })?; initialize_plugins_exact(plugin_config) .await .map_err(|error| CliError::Config(format!("plugin activation failed: {error}")))?; diff --git a/crates/cli/src/session.rs b/crates/cli/src/session.rs index a479a5c0..130a72f7 100644 --- a/crates/cli/src/session.rs +++ b/crates/cli/src/session.rs @@ -39,6 +39,18 @@ const TOOL_HINT_TTL: Duration = Duration::from_secs(300); const LAST_OWNER_TTL: Duration = Duration::from_secs(300); const AGENT_IDLE_TIMEOUT: Duration = Duration::from_secs(30); const AGENT_IDLE_SWEEP_INTERVAL: Duration = Duration::from_secs(5); +const ROUTING_IDENTITY_HEADERS: &[&str] = &[ + "x-nemo-relay-session-id", + "x-nemo-relay-agent-kind", + "x-nemo-relay-turn-id", + "x-nemo-relay-request-id", + "x-nemo-relay-owner-id", + "x-nemo-relay-subagent-id", + "x-nemo-relay-parent-scope-id", + "x-nemo-relay-root-scope-id", + "x-nemo-relay-identity-quality", + "x-nemo-relay-source", +]; #[derive(Clone)] pub(crate) struct SessionManager { @@ -97,6 +109,86 @@ pub(crate) struct GatewayCallPrep { pub(crate) prune_empty_session_on_finish: bool, } +struct RoutingIdentityHeaderContext<'a> { + session_id: &'a str, + agent_kind: AgentKind, + turn_index: u64, + request_id: Option<&'a str>, + owner_id: Option<&'a str>, + parent: Option<&'a ScopeHandle>, + root: Option<&'a ScopeHandle>, + metadata: &'a Value, +} + +fn enrich_routing_identity_headers( + request: &mut LlmRequest, + context: RoutingIdentityHeaderContext<'_>, +) { + request.headers.retain(|name, _| { + !ROUTING_IDENTITY_HEADERS + .iter() + .any(|reserved| name.eq_ignore_ascii_case(reserved)) + }); + insert_routing_identity_header( + &mut request.headers, + "x-nemo-relay-session-id", + context.session_id, + ); + insert_routing_identity_header( + &mut request.headers, + "x-nemo-relay-agent-kind", + context.agent_kind.as_str(), + ); + insert_routing_identity_header( + &mut request.headers, + "x-nemo-relay-turn-id", + &context.turn_index.to_string(), + ); + let request_id = context + .request_id + .map(ToOwned::to_owned) + .or_else(|| { + json_string_at( + context.metadata, + &[ + &["llm_correlation_request_id"][..], + &["request_id"][..], + &["requestId"][..], + ], + ) + }) + .unwrap_or_else(|| format!("relay-request-{}", uuid::Uuid::now_v7())); + insert_routing_identity_header(&mut request.headers, "x-nemo-relay-request-id", &request_id); + if let Some(owner_id) = context.owner_id { + insert_routing_identity_header(&mut request.headers, "x-nemo-relay-owner-id", owner_id); + insert_routing_identity_header(&mut request.headers, "x-nemo-relay-subagent-id", owner_id); + } + if let Some(parent) = context.parent { + insert_routing_identity_header( + &mut request.headers, + "x-nemo-relay-parent-scope-id", + &parent.uuid.to_string(), + ); + } + if let Some(root) = context.root { + insert_routing_identity_header( + &mut request.headers, + "x-nemo-relay-root-scope-id", + &root.uuid.to_string(), + ); + } + insert_routing_identity_header( + &mut request.headers, + "x-nemo-relay-identity-quality", + "native", + ); + insert_routing_identity_header(&mut request.headers, "x-nemo-relay-source", "gateway"); +} + +fn insert_routing_identity_header(headers: &mut Map, name: &str, value: &str) { + headers.insert(name.to_string(), json!(value)); +} + struct Session { agent_kind: AgentKind, session_id: String, @@ -1063,11 +1155,25 @@ impl Session { ), owner.metadata, ); + let mut request = start.request; + enrich_routing_identity_headers( + &mut request, + RoutingIdentityHeaderContext { + session_id: &self.session_id, + agent_kind: self.agent_kind, + turn_index: self.turn_index, + request_id: start.request_id.as_deref(), + owner_id: owner.subagent_id.as_deref(), + parent: owner.parent.as_ref(), + root: self.agent_scope.as_ref().or(self.turn_scope.as_ref()), + metadata: &metadata, + }, + ); Ok(GatewayCallPrep { scope_stack: stack.clone(), session_id: self.session_id.clone(), provider_name: start.provider, - request: start.request, + request, parent: owner.parent, attributes, metadata, @@ -1229,6 +1335,22 @@ impl Session { ) } + // Tool hook payloads do not consistently repeat the harness session id. + // Mirror the stable managed identity onto each tool event so external + // consumers can correlate it without reconstructing the parent scope tree. + fn event_identity_metadata(&self, event_metadata: Value) -> Value { + merge_metadata( + event_metadata, + json!({ + "session_id": self.session_id, + "turn_id": self.turn_index.to_string(), + "harness": self.agent_kind.as_str(), + "source": "hook", + "identity_quality": "native", + }), + ) + } + async fn end_turn(&mut self, event: SessionEvent) -> Result<(), CliError> { if let Some(subagent_id) = alignment::aliased_turn_subagent_id(&event) { self.close_subagent_scope(&subagent_id, event.payload) @@ -1627,7 +1749,7 @@ impl Session { let active_tool_owner_subagent_id = owner.subagent_id.clone(); tool_conditional_execution(event.tool_name.as_str(), &arguments)?; let metadata = tool_correlation_metadata( - event.metadata, + self.event_identity_metadata(event.metadata), owner.status, owner.source.as_deref(), owner.subagent_id.as_deref(), @@ -1659,6 +1781,7 @@ impl Session { // hooks observable and preserves the final result/status instead of dropping orphaned endings. async fn end_tool(&mut self, event: ToolEvent) -> Result<(), CliError> { self.ensure_turn_started(event.metadata.clone())?; + let event_metadata = self.event_identity_metadata(event.metadata.clone()); let completed_agent_subagent_id = alignment::completed_subagent_from_tool(&event); let explicit_subagent_id = event .subagent_id @@ -1678,7 +1801,7 @@ impl Session { event.arguments }; let metadata = tool_correlation_metadata( - event.metadata.clone(), + event_metadata.clone(), owner.status, owner.source.as_deref(), owner.subagent_id.as_deref(), @@ -1701,7 +1824,7 @@ impl Session { .handle(&handle) .result(event.result.clone()) .metadata(merge_metadata( - event.metadata, + event_metadata, json!({ "status": event.status }), )) .build(), diff --git a/crates/cli/tests/coverage/doctor_tests.rs b/crates/cli/tests/coverage/doctor_tests.rs index a93cab3c..fb3f751c 100644 --- a/crates/cli/tests/coverage/doctor_tests.rs +++ b/crates/cli/tests/coverage/doctor_tests.rs @@ -1238,8 +1238,24 @@ async fn atof_endpoint_validation_rejects_missing_url_headers_timeout_and_transp .contains("headers.x-test must be a string") ); - let unsupported = probe_atof_endpoint( + let mixed_case_duplicate = probe_atof_endpoint( 4, + &serde_json::json!({ + "url": "http://127.0.0.1:1/events", + "headers": {"Authorization": "Bearer literal"}, + "header_env": {"authorization": "PATH"} + }), + ) + .await; + assert_eq!(mixed_case_duplicate.status, Status::Fail); + assert!( + mixed_case_duplicate + .details + .contains("cannot appear in both headers and header_env") + ); + + let unsupported = probe_atof_endpoint( + 5, &serde_json::json!({ "url": "http://127.0.0.1:1/events", "transport": "grpc" diff --git a/crates/cli/tests/coverage/gateway_tests.rs b/crates/cli/tests/coverage/gateway_tests.rs index 17d55f8c..18511e67 100644 --- a/crates/cli/tests/coverage/gateway_tests.rs +++ b/crates/cli/tests/coverage/gateway_tests.rs @@ -11,6 +11,7 @@ use axum::extract::State; use axum::http::{HeaderMap, HeaderValue, Method, Request, StatusCode, header}; use http_body_util::BodyExt; use reqwest::Client; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; fn test_http_client() -> Client { Client::new() @@ -279,6 +280,201 @@ fn effective_upstream_request_skips_invalid_runtime_headers() { assert!(headers.keys().all(|name| name.as_str() != "bad header")); } +#[test] +fn internal_dispatch_controls_are_consumed_and_never_forwarded() { + let original_body = Bytes::from_static(br#"{"model":"original"}"#); + let mut original_headers = HeaderMap::new(); + original_headers.insert( + INTERNAL_DISPATCH_URL_HEADER, + HeaderValue::from_static("http://attacker.invalid"), + ); + original_headers.insert( + INTERNAL_RETRY_AWARE_HEADER, + HeaderValue::from_static("true"), + ); + let request = LlmRequest { + headers: Map::from_iter([ + ( + INTERNAL_DISPATCH_URL_HEADER.to_string(), + json!("http://127.0.0.1:9000/v1/responses"), + ), + ( + INTERNAL_DISPATCH_ROUTE_HEADER.to_string(), + json!("openai_responses"), + ), + (INTERNAL_RETRY_AWARE_HEADER.to_string(), json!("true")), + ("x-backend".to_string(), json!("selected")), + ]), + content: json!({"model": "selected"}), + }; + + let effective = effective_dispatch_request( + &original_body, + &original_headers, + Some(&request), + "http://default.invalid/v1/chat/completions", + ProviderRoute::OpenAiChatCompletions, + ); + assert_eq!(effective.url, "http://127.0.0.1:9000/v1/responses"); + assert_eq!(effective.route, ProviderRoute::OpenAiResponses); + assert_eq!(effective.headers.get("x-backend").unwrap(), "selected"); + assert!( + effective + .headers + .get(INTERNAL_DISPATCH_URL_HEADER) + .is_none() + ); + assert!( + effective + .headers + .get(INTERNAL_DISPATCH_ROUTE_HEADER) + .is_none() + ); + assert!(effective.headers.get(INTERNAL_RETRY_AWARE_HEADER).is_none()); + assert!(retry_aware_dispatch(&request)); +} + +#[test] +fn structured_upstream_failure_classification_matches_retry_policy() { + let mut headers = HeaderMap::new(); + headers.insert("set-cookie", HeaderValue::from_static("session=secret")); + headers.insert( + "www-authenticate", + HeaderValue::from_static("Bearer realm=provider"), + ); + headers.insert( + "proxy-authenticate", + HeaderValue::from_static("Basic realm=proxy"), + ); + headers.insert( + "proxy-authorization", + HeaderValue::from_static("Basic secret"), + ); + headers.insert("authorization", HeaderValue::from_static("Bearer secret")); + headers.insert("cookie", HeaderValue::from_static("session=secret")); + headers.insert("x-api-key", HeaderValue::from_static("secret")); + headers.insert("api-key", HeaderValue::from_static("secret")); + headers.insert("anthropic-api-key", HeaderValue::from_static("secret")); + headers.insert("connection", HeaderValue::from_static("close")); + headers.insert("content-length", HeaderValue::from_static("12")); + headers.insert("retry-after", HeaderValue::from_static("3")); + headers.insert("x-request-id", HeaderValue::from_static("request-123")); + for status in [408, 429, 500, 502, 503, 504] { + let failure = http_failure( + StatusCode::from_u16(status).unwrap(), + &headers, + b"temporary", + ); + assert!(failure.is_retryable(), "status={status}"); + } + let failure = http_failure(StatusCode::BAD_GATEWAY, &headers, b"temporary"); + for name in [ + "set-cookie", + "www-authenticate", + "proxy-authenticate", + "proxy-authorization", + "authorization", + "cookie", + "x-api-key", + "api-key", + "anthropic-api-key", + "connection", + "content-length", + ] { + assert!( + !failure.headers.contains_key(name), + "sensitive header: {name}" + ); + } + assert_eq!(failure.headers.get("retry-after"), Some(&"3".to_string())); + assert_eq!( + failure.headers.get("x-request-id"), + Some(&"request-123".to_string()) + ); + assert_eq!( + http_failure( + StatusCode::BAD_REQUEST, + &headers, + b"context_length_exceeded" + ) + .class, + UpstreamFailureClass::ContextWindow + ); + assert_eq!( + http_failure( + StatusCode::SERVICE_UNAVAILABLE, + &headers, + b"model unavailable" + ) + .class, + UpstreamFailureClass::ModelUnavailable + ); + assert!(!http_failure(StatusCode::UNAUTHORIZED, &headers, b"bad token").is_retryable()); + assert!(!http_failure(StatusCode::BAD_REQUEST, &headers, b"invalid request").is_retryable()); + assert_eq!( + bounded_error_body(&vec![b'x'; MAX_UPSTREAM_ERROR_BODY_BYTES + 10]).len(), + MAX_UPSTREAM_ERROR_BODY_BYTES + ); +} + +#[tokio::test] +async fn retry_aware_buffered_body_read_failure_stays_structured() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0_u8; 1024]; + let _ = socket.read(&mut request).await.unwrap(); + socket + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: 64\r\nconnection: close\r\n\r\n{}", + ) + .await + .unwrap(); + }); + + let config = GatewayConfig::default(); + let state = AppState { + config: config.clone(), + http: test_http_client(), + sessions: SessionManager::new(config), + last_activity: Arc::new(Mutex::new(std::time::Instant::now())), + }; + let prepared = PreparedGatewayRequest { + method: Method::POST, + headers: HeaderMap::new(), + path: "/v1/chat/completions".into(), + provider: ProviderRoute::OpenAiChatCompletions, + upstream_url: format!("http://{address}/v1/chat/completions"), + body_bytes: Bytes::from_static(b"{}"), + request_json: json!({}), + streaming: false, + }; + let upstream_info = Arc::new(Mutex::new(None)); + let upstream_error = Arc::new(Mutex::new(None)); + let response_bytes = Arc::new(Mutex::new(None)); + let func = build_buffered_func( + state, + &prepared, + upstream_info, + upstream_error.clone(), + response_bytes, + ); + let error = func(LlmRequest { + headers: Map::from_iter([(INTERNAL_RETRY_AWARE_HEADER.into(), json!("true"))]), + content: json!({}), + }) + .await + .unwrap_err(); + + let FlowError::Upstream(failure) = error else { + panic!("expected structured upstream failure, got {error:?}"); + }; + assert_eq!(failure.class, UpstreamFailureClass::Connection); + assert!(upstream_error.lock().unwrap().is_none()); + server.await.unwrap(); +} + #[test] fn gateway_session_id_prefers_headers_and_has_fallbacks() { let mut headers = HeaderMap::new(); diff --git a/crates/cli/tests/coverage/launcher_tests.rs b/crates/cli/tests/coverage/launcher_tests.rs index 1615642c..d0231a54 100644 --- a/crates/cli/tests/coverage/launcher_tests.rs +++ b/crates/cli/tests/coverage/launcher_tests.rs @@ -607,6 +607,11 @@ fn prepares_hermes_hook_environment() { .env .contains(&("HERMES_ACCEPT_HOOKS".into(), "1".into())) ); + assert!( + prepared + .env + .contains(&("OPENAI_BASE_URL".into(), "http://127.0.0.1:1234/v1".into())) + ); assert_eq!( prepared .hermes_restore @@ -616,6 +621,8 @@ fn prepares_hermes_hook_environment() { ); let hooks = std::fs::read_to_string(&hooks_path).unwrap(); assert!(hooks.contains("hook-forward hermes")); + assert!(hooks.contains("provider: custom")); + assert!(hooks.contains("base_url: http://127.0.0.1:1234/v1")); assert!(prepared.notes[0].contains("temporarily merged")); prepared.restore().unwrap(); @@ -623,6 +630,14 @@ fn prepares_hermes_hook_environment() { std::env::set_current_dir(previous).unwrap(); } +#[test] +fn hermes_openai_base_url_normalizes_trailing_slashes() { + assert_eq!( + hermes_openai_base_url("http://127.0.0.1:1234///"), + "http://127.0.0.1:1234/v1" + ); +} + #[test] fn prepares_hermes_dry_uses_home_path_without_writing_hooks() { let _guard = current_dir_lock().lock().unwrap(); @@ -655,6 +670,11 @@ fn prepares_hermes_dry_uses_home_path_without_writing_hooks() { .env .contains(&("HERMES_ACCEPT_HOOKS".into(), "1".into())) ); + assert!( + prepared + .env + .contains(&("OPENAI_BASE_URL".into(), "http://127.0.0.1:1234/v1".into())) + ); assert!(!hook_path.exists()); } diff --git a/crates/cli/tests/coverage/session_tests.rs b/crates/cli/tests/coverage/session_tests.rs index 117831a3..9e2f5cae 100644 --- a/crates/cli/tests/coverage/session_tests.rs +++ b/crates/cli/tests/coverage/session_tests.rs @@ -21,6 +21,80 @@ use crate::test_support::PLUGIN_CONFIG_TEST_LOCK; const HERMES_ROUTED_TEST_SESSION_KEY: &str = "hermes_routed_test_session_id"; +#[test] +fn routing_identity_enrichment_replaces_untrusted_reserved_headers() { + let mut request = LlmRequest { + headers: Map::from_iter([ + ("X-Nemo-Relay-Session-Id".into(), json!("spoofed-session")), + ("x-nemo-relay-agent-kind".into(), json!("spoofed-agent")), + ("x-nemo-relay-turn-id".into(), json!("spoofed-turn")), + ("x-nemo-relay-request-id".into(), json!("spoofed-request")), + ("x-nemo-relay-owner-id".into(), json!("spoofed-owner")), + ("x-nemo-relay-subagent-id".into(), json!("spoofed-subagent")), + ( + "x-nemo-relay-parent-scope-id".into(), + json!("spoofed-parent"), + ), + ("x-nemo-relay-root-scope-id".into(), json!("spoofed-root")), + ( + "x-nemo-relay-identity-quality".into(), + json!("spoofed-quality"), + ), + ("x-nemo-relay-source".into(), json!("spoofed-source")), + ("x-nemo-relay-config-profile".into(), json!("preserved")), + ]), + content: json!({}), + }; + + enrich_routing_identity_headers( + &mut request, + RoutingIdentityHeaderContext { + session_id: "trusted-session", + agent_kind: AgentKind::Hermes, + turn_index: 7, + request_id: Some("trusted-request"), + owner_id: None, + parent: None, + root: None, + metadata: &json!({}), + }, + ); + + assert_eq!( + request.headers["x-nemo-relay-session-id"], + json!("trusted-session") + ); + assert_eq!( + request.headers["x-nemo-relay-request-id"], + json!("trusted-request") + ); + assert_eq!(request.headers["x-nemo-relay-agent-kind"], json!("hermes")); + assert_eq!(request.headers["x-nemo-relay-turn-id"], json!("7")); + assert_eq!( + request.headers["x-nemo-relay-identity-quality"], + json!("native") + ); + assert_eq!(request.headers["x-nemo-relay-source"], json!("gateway")); + for absent in [ + "x-nemo-relay-owner-id", + "x-nemo-relay-subagent-id", + "x-nemo-relay-parent-scope-id", + "x-nemo-relay-root-scope-id", + ] { + assert!(!request.headers.contains_key(absent)); + } + assert_eq!( + request.headers["x-nemo-relay-config-profile"], + json!("preserved") + ); + assert!( + request + .headers + .keys() + .all(|name| name != "X-Nemo-Relay-Session-Id") + ); +} + async fn install_test_atif_plugin(output_directory: &Path) { let _ = clear_plugin_configuration(); std::fs::create_dir_all(output_directory).unwrap(); @@ -5118,6 +5192,15 @@ async fn claude_agent_tool_async_launch_keeps_subagent_open_for_later_hooks() { tool.metadata.as_ref().unwrap()["tool_correlation_subagent_id"], json!("worker") ); + assert_eq!( + tool.metadata.as_ref().unwrap()["session_id"], + json!("agent-tool-async") + ); + assert_eq!(tool.metadata.as_ref().unwrap()["turn_id"], json!("1")); + assert_eq!( + tool.metadata.as_ref().unwrap()["identity_quality"], + json!("native") + ); } #[tokio::test] diff --git a/crates/cli/tests/switchyard_process_e2e.rs b/crates/cli/tests/switchyard_process_e2e.rs new file mode 100644 index 00000000..31f01819 --- /dev/null +++ b/crates/cli/tests/switchyard_process_e2e.rs @@ -0,0 +1,306 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! CI-safe process-boundary coverage for the Switchyard plugin. + +use std::process::{Child, Command, Stdio}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use axum::body::Body; +use axum::extract::State; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::Response; +use axum::routing::post; +use axum::{Json, Router}; +use serde_json::{Value, json}; + +fn gateway_bin() -> &'static str { + env!("CARGO_BIN_EXE_nemo-relay") +} + +struct ChildGuard(Child); + +impl Drop for ChildGuard { + fn drop(&mut self) { + let _ = self.0.kill(); + let _ = self.0.wait(); + } +} + +#[derive(Clone, Default)] +struct DecisionState { + requests: Arc>>, +} + +async fn decide( + State(state): State, + headers: HeaderMap, + Json(request): Json, +) -> Response { + let call = { + let mut requests = state.requests.lock().unwrap(); + requests.push((headers, request)); + requests.len() + }; + if call == 3 { + return Response::builder() + .status(StatusCode::SERVICE_UNAVAILABLE) + .body(Body::from("decision API unavailable")) + .unwrap(); + } + let body = json!({ + "schema_version": "switchyard.routing_decision.v1", + "decision_id": format!("decision-{call}"), + "router": {"name": "fake-ci-router", "version": "1"}, + "route": { + "tier": "strong", + "target_model": "provider/selected", + "backend_id": "selected-chat", + "target_protocol_profile": "openai_chat", + "target_endpoint": "/v1/chat/completions" + }, + "confidence": 0.99, + "reason_code": "ci_fixture", + "reason_summary": "deterministic process E2E decision" + }); + Response::builder() + .status(StatusCode::OK) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap() +} + +#[derive(Clone, Default)] +struct ProviderState { + requests: Arc>>, +} + +async fn provide( + State(state): State, + headers: HeaderMap, + Json(request): Json, +) -> Response { + let stream = request["stream"].as_bool().unwrap_or(false); + let model = request["model"].as_str().unwrap_or("unknown").to_string(); + state.requests.lock().unwrap().push((headers, request)); + if stream { + let first = json!({ + "id": "chat-ci", "object": "chat.completion.chunk", "model": model, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "streamed"}, "finish_reason": null}] + }); + let last = json!({ + "id": "chat-ci", "object": "chat.completion.chunk", "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 4, "completion_tokens": 1, "total_tokens": 5} + }); + let body = format!("data: {first}\n\ndata: {last}\n\ndata: [DONE]\n\n"); + return Response::builder() + .status(StatusCode::OK) + .header("content-type", "text/event-stream") + .body(Body::from(body)) + .unwrap(); + } + let body = json!({ + "id": "chat-ci", "object": "chat.completion", "model": model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": format!("served by {model}")}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7} + }); + Response::builder() + .status(StatusCode::OK) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap() +} + +async fn start_server(router: Router) -> (String, tokio::task::JoinHandle<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let task = tokio::spawn(async move { + axum::serve(listener, router).await.unwrap(); + }); + (format!("http://{address}"), task) +} + +fn unused_address() -> std::net::SocketAddr { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap() +} + +async fn wait_for_gateway(client: &reqwest::Client, url: &str, child: &mut Child) { + for _ in 0..120 { + if let Some(status) = child.try_wait().unwrap() { + panic!("gateway exited before readiness with {status}"); + } + if client + .get(format!("{url}/healthz")) + .send() + .await + .is_ok_and(|response| response.status().is_success()) + { + return; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + panic!("gateway did not become ready at {url}"); +} + +#[tokio::test(flavor = "multi_thread")] +async fn switchyard_plugin_routes_buffered_and_streaming_then_fails_open() { + let decision_state = DecisionState::default(); + let decision_requests = Arc::clone(&decision_state.requests); + let (decision_url, decision_task) = start_server( + Router::new() + .route("/v1/routing/decision", post(decide)) + .with_state(decision_state), + ) + .await; + + let provider_state = ProviderState::default(); + let provider_requests = Arc::clone(&provider_state.requests); + let (provider_url, provider_task) = start_server( + Router::new() + .route("/v1/chat/completions", post(provide)) + .with_state(provider_state), + ) + .await; + + let temp = tempfile::tempdir().unwrap(); + let config_path = temp.path().join("plugins.toml"); + let config = format!( + r#"version = 1 + +[[components]] +kind = "switchyard" +enabled = true + +[components.config] +mode = "enforce" +decision_api_url = "{decision_url}/v1/routing/decision" +decision_profile_id = "ci-process-e2e" +request_materialization = "full_body" +context_mode = "payload_only" +decision_timeout_millis = 1000 +max_retries = 0 + +[components.config.default_targets] +openai_chat = "fallback-chat" +openai_responses = "fallback-responses" +anthropic_messages = "fallback-anthropic" + +[components.config.targets.selected-chat] +model = "provider/selected" +protocol = "openai_chat" +endpoint = "/v1/chat/completions" +base_url = "{provider_url}" + +[components.config.targets.fallback-chat] +model = "provider/fallback" +protocol = "openai_chat" +endpoint = "/v1/chat/completions" +base_url = "{provider_url}" + +[components.config.targets.fallback-responses] +model = "provider/fallback" +protocol = "openai_responses" +endpoint = "/v1/responses" +base_url = "{provider_url}" + +[components.config.targets.fallback-anthropic] +model = "provider/fallback" +protocol = "anthropic_messages" +endpoint = "/v1/messages" +base_url = "{provider_url}" +"# + ); + std::fs::write(&config_path, config).unwrap(); + + let address = unused_address(); + let gateway_url = format!("http://{address}"); + let stderr = std::fs::File::create(temp.path().join("gateway.log")).unwrap(); + let child = Command::new(gateway_bin()) + .arg("--plugin-config-path") + .arg(&config_path) + .arg("--bind") + .arg(address.to_string()) + .stdout(Stdio::null()) + .stderr(Stdio::from(stderr)) + .spawn() + .unwrap(); + let mut gateway = ChildGuard(child); + let client = reqwest::Client::new(); + wait_for_gateway(&client, &gateway_url, &mut gateway.0).await; + + let send = |request_id: &'static str, stream: bool| { + client + .post(format!("{gateway_url}/v1/chat/completions")) + .header("x-nemo-relay-session-id", "ci-process-session") + .header("x-nemo-relay-request-id", request_id) + .header( + "x-nemo-relay-internal-dispatch-url", + "http://attacker.invalid", + ) + .header("x-nemo-relay-internal-dispatch-route", "attacker-route") + .json(&json!({ + "model": "client/model", + "stream": stream, + "messages": [{"role": "user", "content": "process boundary test"}] + })) + .send() + }; + + let buffered = send("buffered-request", false).await.unwrap(); + assert!(buffered.status().is_success()); + let buffered: Value = buffered.json().await.unwrap(); + assert_eq!(buffered["model"], "provider/selected"); + + let streaming = send("stream-request", true).await.unwrap(); + assert!(streaming.status().is_success()); + let streaming = streaming.text().await.unwrap(); + assert!(streaming.contains("streamed")); + assert!(streaming.contains("[DONE]")); + + let fallback = send("fallback-request", false).await.unwrap(); + assert!(fallback.status().is_success()); + let fallback: Value = fallback.json().await.unwrap(); + assert_eq!(fallback["model"], "provider/fallback"); + + let decisions = decision_requests.lock().unwrap(); + assert_eq!(decisions.len(), 3); + for (headers, body) in decisions.iter() { + assert!(!headers.contains_key("x-nemo-relay-internal-dispatch-url")); + assert!(!headers.contains_key("x-nemo-relay-internal-dispatch-route")); + assert_eq!( + headers + .get("x-nemo-relay-session-id") + .unwrap() + .to_str() + .unwrap(), + "ci-process-session" + ); + assert_eq!(body["schema_version"], "switchyard.routing_request.v1"); + assert_eq!(body["decision_profile"]["profile_id"], "ci-process-e2e"); + } + drop(decisions); + + let providers = provider_requests.lock().unwrap(); + let models = providers + .iter() + .map(|(_, body)| body["model"].as_str().unwrap()) + .collect::>(); + assert_eq!( + models, + vec![ + "provider/selected", + "provider/selected", + "provider/fallback" + ] + ); + for (headers, _) in providers.iter() { + assert!(!headers.contains_key("x-nemo-relay-internal-dispatch-url")); + assert!(!headers.contains_key("x-nemo-relay-internal-dispatch-route")); + } + + decision_task.abort(); + provider_task.abort(); +} diff --git a/crates/core/src/api/scope.rs b/crates/core/src/api/scope.rs index 7546dc92..60a1aa53 100644 --- a/crates/core/src/api/scope.rs +++ b/crates/core/src/api/scope.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use crate::api::event::{BaseEvent, MarkEvent}; +use crate::api::event::{BaseEvent, CategoryProfile, DataSchema, EventCategory, MarkEvent}; use crate::api::runtime::NemoRelayContextState; use crate::api::runtime::global_context; use crate::api::runtime::{ @@ -155,9 +155,18 @@ pub struct EmitMarkEventParams<'a> { /// Optional JSON payload recorded as the mark data. #[builder(default)] pub data: Option, + /// Optional schema identifier for the mark data. + #[builder(default)] + pub data_schema: Option, /// Optional JSON metadata recorded on the emitted event. #[builder(default)] pub metadata: Option, + /// Optional semantic category for the mark. + #[builder(default)] + pub category: Option, + /// Optional category-specific mark profile. + #[builder(default)] + pub category_profile: Option, /// Optional timestamp recorded on the emitted mark event. When omitted, the /// current UTC time is used. #[builder(default)] @@ -353,10 +362,11 @@ pub fn event(params: EmitMarkEventParams<'_>) -> Result<()> { .parent_uuid_opt(parent_uuid) .timestamp(params.timestamp.unwrap_or_else(Utc::now)) .data_opt(params.data) + .data_schema_opt(params.data_schema) .metadata_opt(params.metadata) .build(), - None, - None, + params.category, + params.category_profile, )); (event, subscribers) }; diff --git a/crates/core/src/codec/streaming.rs b/crates/core/src/codec/streaming.rs index db60ab46..92bf2a08 100644 --- a/crates/core/src/codec/streaming.rs +++ b/crates/core/src/codec/streaming.rs @@ -25,6 +25,47 @@ use crate::api::runtime::{LlmCollectorFn, LlmFinalizerFn}; use crate::error::{FlowError, Result}; use crate::json::Json; +use serde::{Deserialize, Serialize}; + +/// Provider-neutral incremental stream item used by cross-protocol transcoders. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum NormalizedStreamEvent { + /// Assistant text delta. + TextDelta { + /// Newly generated text. + text: String, + }, + /// Start of a function tool call. + ToolCallStart { + /// Stable tool-call identifier. + id: String, + /// Function name. + name: String, + }, + /// Incremental JSON arguments for a function tool call. + ToolCallArgumentsDelta { + /// Stable tool-call identifier. + id: String, + /// Newly generated argument bytes. + delta: String, + }, + /// Provider-reported terminal reason. + Finish { + /// Normalized or provider-native finish label. + reason: Option, + }, + /// Incremental or terminal usage object. + Usage { + /// Provider-neutral usage JSON. + usage: Json, + }, + /// Provider stream error payload. + Error { + /// Structured provider error. + error: Json, + }, +} /// Per-provider streaming codec used with [`crate::api::llm::llm_stream_call_execute`]. /// diff --git a/crates/core/src/codec/traits.rs b/crates/core/src/codec/traits.rs index 2f7bef13..1c62647c 100644 --- a/crates/core/src/codec/traits.rs +++ b/crates/core/src/codec/traits.rs @@ -77,3 +77,12 @@ pub trait LlmResponseCodec: Send + Sync { /// Implementations should return `Err` only for genuinely unparseable input. fn decode_response(&self, response: &Json) -> Result; } + +/// Additive encoder for materializing a normalized response in a provider wire format. +/// +/// Existing response codecs remain decode-only. Components that perform cross-protocol dispatch +/// can supply this companion trait without changing the managed observability pipeline. +pub trait LlmResponseEncoder: Send + Sync { + /// Encode a normalized response into the target provider's buffered JSON representation. + fn encode_response(&self, response: &AnnotatedLlmResponse) -> Result; +} diff --git a/crates/core/src/error.rs b/crates/core/src/error.rs index 8c7e24f9..70d49aa4 100644 --- a/crates/core/src/error.rs +++ b/crates/core/src/error.rs @@ -7,8 +7,77 @@ //! [`FlowError`] as the error type. Errors are categorized by cause //! (duplicate registration, missing entity, guardrail rejection, etc.). +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; use thiserror::Error; +/// Stable classification for an upstream provider failure captured by retry-aware dispatch. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum UpstreamFailureClass { + /// Provider connection could not be established or was interrupted. + Connection, + /// Provider request timed out. + Timeout, + /// Retryable HTTP status without a more specific provider classification. + RetryableStatus, + /// Provider rejected the request because its context window was exceeded. + ContextWindow, + /// Requested provider model is temporarily unavailable. + ModelUnavailable, + /// Provider authentication or authorization failed. + Authentication, + /// Provider rejected an invalid request. + InvalidRequest, + /// Other non-retryable provider failure. + Other, +} + +/// Structured provider failure surfaced only for explicitly retry-aware dispatches. +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct UpstreamFailure { + /// HTTP status when a provider response was received. + pub status: Option, + /// Bounded response body or transport error message. + pub body: String, + /// Safe response headers captured from the provider. + pub headers: BTreeMap, + /// Retry classification. + pub class: UpstreamFailureClass, +} + +impl UpstreamFailure { + /// Whether Switchyard may be consulted for another bounded provider attempt. + pub fn is_retryable(&self) -> bool { + matches!( + self.class, + UpstreamFailureClass::Connection + | UpstreamFailureClass::Timeout + | UpstreamFailureClass::RetryableStatus + | UpstreamFailureClass::ContextWindow + | UpstreamFailureClass::ModelUnavailable + ) + } +} + +impl std::fmt::Display for UpstreamFailure { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self.status { + Some(status) => write!( + formatter, + "upstream provider returned HTTP {status} ({:?}): {}", + self.class, self.body + ), + None => write!( + formatter, + "upstream provider transport failure ({:?}): {}", + self.class, self.body + ), + } + } +} + /// The error type for all NeMo Relay runtime operations. /// /// Each variant represents a distinct failure mode that callers can match on @@ -53,6 +122,10 @@ pub enum FlowError { #[error("guardrail rejected: {0}")] GuardrailRejected(String), + /// Structured upstream provider failure from retry-aware gateway dispatch. + #[error("{0}")] + Upstream(UpstreamFailure), + /// An internal runtime error (e.g., lock poisoning). #[error("internal error: {0}")] Internal(String), diff --git a/crates/core/src/observability/atof.rs b/crates/core/src/observability/atof.rs index 75b40665..ffa2013f 100644 --- a/crates/core/src/observability/atof.rs +++ b/crates/core/src/observability/atof.rs @@ -177,6 +177,9 @@ pub struct AtofEndpointConfig { /// Headers applied to endpoint requests or handshakes. #[serde(default, skip_serializing_if = "HashMap::is_empty")] pub headers: HashMap, + /// Header names mapped to environment variables containing their values. + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub header_env: HashMap, /// Per-endpoint timeout in milliseconds. #[serde(default = "default_endpoint_timeout_millis")] pub timeout_millis: u64, @@ -192,6 +195,7 @@ impl AtofEndpointConfig { url: url.into(), transport, headers: HashMap::new(), + header_env: HashMap::new(), timeout_millis: default_endpoint_timeout_millis(), field_name_policy: AtofEndpointFieldNamePolicy::Preserve, } @@ -203,6 +207,17 @@ impl AtofEndpointConfig { self } + /// Add a header whose value is resolved from an environment variable at activation. + pub fn with_header_env( + mut self, + key: impl Into, + environment_variable: impl Into, + ) -> Self { + self.header_env + .insert(key.into(), environment_variable.into()); + self + } + /// Override the endpoint timeout. pub fn with_timeout_millis(mut self, timeout_millis: u64) -> Self { self.timeout_millis = timeout_millis; @@ -585,12 +600,15 @@ fn validate_endpoint_config(config: &AtofEndpointConfig) -> Result<()> { url.scheme() ))); } - build_header_map(&config.headers)?; + resolved_header_map(&config.headers, &config.header_env)?; Ok(()) } #[cfg(feature = "atof-streaming")] -fn build_header_map(headers: &HashMap) -> Result { +fn resolved_header_map( + headers: &HashMap, + header_env: &HashMap, +) -> Result { let mut out = reqwest::header::HeaderMap::new(); for (key, value) in headers { let name = reqwest::header::HeaderName::from_bytes(key.as_bytes()) @@ -599,6 +617,28 @@ fn build_header_map(headers: &HashMap) -> Result headers, - Err(error) => { - eprintln!("nemo_relay: ATOF endpoint[{index}] disabled: {error}"); - drain_closed(rx).await; - return; - } - }) + .default_headers( + match resolved_header_map(&config.headers, &config.header_env) { + Ok(headers) => headers, + Err(error) => { + eprintln!("nemo_relay: ATOF endpoint[{index}] disabled: {error}"); + drain_closed(rx).await; + return; + } + }, + ) .build() { Ok(client) => client, @@ -795,12 +837,12 @@ async fn connect_websocket( .as_str() .into_client_request() .map_err(|error| error.to_string())?; - for (key, value) in &config.headers { - let name = - tokio_tungstenite::tungstenite::http::header::HeaderName::from_bytes(key.as_bytes()) - .map_err(|error| error.to_string())?; - let value = tokio_tungstenite::tungstenite::http::header::HeaderValue::from_str(value) - .map_err(|error| error.to_string())?; + let headers = resolved_header_map(&config.headers, &config.header_env) + .map_err(|error| error.to_string())?; + for (name, value) in headers { + let Some(name) = name else { + continue; + }; request.headers_mut().insert(name, value); } tokio::time::timeout( @@ -859,8 +901,8 @@ async fn run_ndjson_endpoint( fn build_ndjson_client( config: &AtofEndpointConfig, ) -> std::result::Result { - let headers = - build_header_map(&config.headers).map_err(|error| format!("disabled: {error}"))?; + let headers = resolved_header_map(&config.headers, &config.header_env) + .map_err(|error| format!("disabled: {error}"))?; reqwest::Client::builder() .connect_timeout(Duration::from_millis(config.timeout_millis)) .default_headers(headers) diff --git a/crates/core/src/observability/plugin_component.rs b/crates/core/src/observability/plugin_component.rs index 8481d064..a2758419 100644 --- a/crates/core/src/observability/plugin_component.rs +++ b/crates/core/src/observability/plugin_component.rs @@ -199,6 +199,9 @@ pub struct AtofEndpointSectionConfig { /// Headers applied to endpoint requests or handshakes. #[serde(default)] pub headers: HashMap, + /// Header names mapped to environment variables containing their values. + #[serde(default)] + pub header_env: HashMap, /// Per-endpoint timeout in milliseconds. #[serde(default = "default_timeout_millis")] pub timeout_millis: u64, @@ -696,6 +699,9 @@ fn build_atof_endpoint_config( for (key, value) in endpoint.headers { config = config.with_header(key, value); } + for (key, variable) in endpoint.header_env { + config = config.with_header_env(key, variable); + } Ok(config) } diff --git a/crates/core/src/plugin/dynamic/manifest.rs b/crates/core/src/plugin/dynamic/manifest.rs index 339bf89a..643bc586 100644 --- a/crates/core/src/plugin/dynamic/manifest.rs +++ b/crates/core/src/plugin/dynamic/manifest.rs @@ -143,6 +143,10 @@ pub struct DynamicPluginManifestRustDynamicLoad { #[cfg(feature = "schema")] #[derive(schemars::JsonSchema)] +#[allow( + dead_code, + reason = "variants are consumed by schemars derive metadata" +)] #[serde(untagged)] enum DynamicPluginManifestLoadSchema { Worker(DynamicPluginManifestWorkerLoad), diff --git a/crates/core/src/plugin/dynamic/native.rs b/crates/core/src/plugin/dynamic/native.rs index cfc409bf..578e8cf3 100644 --- a/crates/core/src/plugin/dynamic/native.rs +++ b/crates/core/src/plugin/dynamic/native.rs @@ -992,7 +992,7 @@ fn status_from_flow_error(err: FlowError) -> NemoRelayStatus { FlowError::InvalidArgument(_) => NemoRelayStatus::InvalidArg, FlowError::ScopeStackEmpty => NemoRelayStatus::ScopeStackEmpty, FlowError::GuardrailRejected(_) => NemoRelayStatus::GuardrailRejected, - FlowError::Internal(_) => NemoRelayStatus::Internal, + FlowError::Upstream(_) | FlowError::Internal(_) => NemoRelayStatus::Internal, } } diff --git a/crates/core/tests/unit/observability/atof_tests.rs b/crates/core/tests/unit/observability/atof_tests.rs index ded5e7f7..e89d34a6 100644 --- a/crates/core/tests/unit/observability/atof_tests.rs +++ b/crates/core/tests/unit/observability/atof_tests.rs @@ -1388,16 +1388,42 @@ fn endpoint_validation_rejects_empty_timeout_and_invalid_headers() { url: "http://127.0.0.1:9/events".into(), transport: AtofEndpointTransport::HttpPost, headers: headers.clone(), + header_env: std::collections::HashMap::new(), timeout_millis: 1, field_name_policy: AtofEndpointFieldNamePolicy::Preserve, }) .unwrap(); - assert_eq!(build_header_map(&headers).unwrap().len(), 1); + assert_eq!( + resolved_header_map(&headers, &Default::default()) + .unwrap() + .len(), + 1 + ); + let env_headers = std::collections::HashMap::from([("x-path".into(), "PATH".into())]); + assert_eq!( + resolved_header_map(&Default::default(), &env_headers) + .unwrap() + .len(), + 1 + ); + let missing_env = std::collections::HashMap::from([( + "authorization".into(), + "NEMO_RELAY_TEST_MISSING_ATOF_SECRET".into(), + )]); + assert!(resolved_header_map(&Default::default(), &missing_env).is_err()); + let conflict = std::collections::HashMap::from([("x-test".into(), "PATH".into())]); + assert!(resolved_header_map(&headers, &conflict).is_err()); + let mixed_case_headers = + std::collections::HashMap::from([("Authorization".into(), "Bearer literal".into())]); + let mixed_case_conflict = + std::collections::HashMap::from([("authorization".into(), "PATH".into())]); + assert!(resolved_header_map(&mixed_case_headers, &mixed_case_conflict).is_err()); let empty_url = AtofEndpointConfig { url: " ".into(), transport: AtofEndpointTransport::HttpPost, headers: std::collections::HashMap::new(), + header_env: std::collections::HashMap::new(), timeout_millis: 1, field_name_policy: AtofEndpointFieldNamePolicy::Preserve, }; @@ -1412,6 +1438,7 @@ fn endpoint_validation_rejects_empty_timeout_and_invalid_headers() { url: "http://127.0.0.1:9/events".into(), transport: AtofEndpointTransport::HttpPost, headers: std::collections::HashMap::new(), + header_env: std::collections::HashMap::new(), timeout_millis: 0, field_name_policy: AtofEndpointFieldNamePolicy::Preserve, }; @@ -1424,16 +1451,17 @@ fn endpoint_validation_rejects_empty_timeout_and_invalid_headers() { let mut bad_header_name = std::collections::HashMap::new(); bad_header_name.insert("bad header".to_string(), "ok".to_string()); - assert!(build_header_map(&bad_header_name).is_err()); + assert!(resolved_header_map(&bad_header_name, &Default::default()).is_err()); let mut bad_header_value = std::collections::HashMap::new(); bad_header_value.insert("x-test".to_string(), "bad\nvalue".to_string()); - assert!(build_header_map(&bad_header_value).is_err()); + assert!(resolved_header_map(&bad_header_value, &Default::default()).is_err()); assert!( build_ndjson_client(&AtofEndpointConfig { url: "http://127.0.0.1:9/events".into(), transport: AtofEndpointTransport::Ndjson, headers: bad_header_value, + header_env: std::collections::HashMap::new(), timeout_millis: 1, field_name_policy: AtofEndpointFieldNamePolicy::Preserve, }) @@ -1534,6 +1562,7 @@ fn http_endpoint_worker_disables_invalid_headers_and_drains_control_messages() { url: "http://127.0.0.1:9/events".into(), transport: AtofEndpointTransport::HttpPost, headers, + header_env: std::collections::HashMap::new(), timeout_millis: 1, field_name_policy: AtofEndpointFieldNamePolicy::Preserve, }, @@ -1567,6 +1596,7 @@ fn websocket_helpers_cover_invalid_headers_and_timeout_reconnect_path() { url: "ws://127.0.0.1:9/events".into(), transport: AtofEndpointTransport::Websocket, headers, + header_env: std::collections::HashMap::new(), timeout_millis: 1, field_name_policy: AtofEndpointFieldNamePolicy::Preserve, }; diff --git a/crates/core/tests/unit/observability/plugin_component_tests.rs b/crates/core/tests/unit/observability/plugin_component_tests.rs index c43d3162..318b1b41 100644 --- a/crates/core/tests/unit/observability/plugin_component_tests.rs +++ b/crates/core/tests/unit/observability/plugin_component_tests.rs @@ -619,6 +619,10 @@ fn build_atof_endpoint_config_maps_headers_timeout_and_rejects_transport() { url: "ws://127.0.0.1:47632/events".into(), transport: "websocket".into(), headers: headers.clone(), + header_env: std::collections::HashMap::from([( + "x-api-key".into(), + "SWITCHYARD_API_KEY".into(), + )]), timeout_millis: 123, field_name_policy: "replace_dots".into(), }, @@ -631,6 +635,10 @@ fn build_atof_endpoint_config_maps_headers_timeout_and_rejects_transport() { crate::observability::atof::AtofEndpointTransport::Websocket ); assert_eq!(config.headers, headers); + assert_eq!( + config.header_env.get("x-api-key").map(String::as_str), + Some("SWITCHYARD_API_KEY") + ); assert_eq!(config.timeout_millis, 123); assert_eq!( config.field_name_policy, @@ -643,6 +651,7 @@ fn build_atof_endpoint_config_maps_headers_timeout_and_rejects_transport() { url: "http://127.0.0.1:47632/events".into(), transport: "smtp".into(), headers: std::collections::HashMap::new(), + header_env: std::collections::HashMap::new(), timeout_millis: 3_000, field_name_policy: "preserve".into(), }, @@ -656,6 +665,7 @@ fn build_atof_endpoint_config_maps_headers_timeout_and_rejects_transport() { url: "http://127.0.0.1:47632/events".into(), transport: "http_post".into(), headers: std::collections::HashMap::new(), + header_env: std::collections::HashMap::new(), timeout_millis: 3_000, field_name_policy: "bogus".into(), }, diff --git a/crates/ffi/src/error.rs b/crates/ffi/src/error.rs index dc1f195c..20719712 100644 --- a/crates/ffi/src/error.rs +++ b/crates/ffi/src/error.rs @@ -120,7 +120,7 @@ impl From<&FlowError> for NemoRelayStatus { FlowError::InvalidArgument(_) => NemoRelayStatus::InvalidArg, FlowError::ScopeStackEmpty => NemoRelayStatus::ScopeStackEmpty, FlowError::GuardrailRejected(_) => NemoRelayStatus::GuardrailRejected, - FlowError::Internal(_) => NemoRelayStatus::Internal, + FlowError::Upstream(_) | FlowError::Internal(_) => NemoRelayStatus::Internal, } } } diff --git a/crates/switchyard/Cargo.toml b/crates/switchyard/Cargo.toml new file mode 100644 index 00000000..8bbb5808 --- /dev/null +++ b/crates/switchyard/Cargo.toml @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[package] +name = "nemo-relay-switchyard" +version.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "First-party Switchyard Decision API routing plugin for NeMo Relay." +publish = false +readme = "README.md" + +[lints] +workspace = true + +[features] +default = [] +schema = ["dep:schemars", "nemo-relay/schema"] + +[dependencies] +nemo-relay.workspace = true +serde = { version = "1", features = ["derive"] } +serde_json = "1" +switchyard-translation.workspace = true +schemars = { version = "0.8", optional = true } +reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-native-roots"] } +tokio = { version = "1", default-features = false, features = ["time"] } +uuid = { workspace = true, features = ["v7"] } +futures-util = "0.3" +async-stream = "0.3" + +[dev-dependencies] +axum = "0.8" +tokio = { version = "1", features = ["net", "rt", "macros", "sync", "test-util", "rt-multi-thread", "time"] } diff --git a/crates/switchyard/README.md b/crates/switchyard/README.md new file mode 100644 index 00000000..82778f37 --- /dev/null +++ b/crates/switchyard/README.md @@ -0,0 +1,109 @@ + + +[![License](https://img.shields.io/github/license/NVIDIA/NeMo-Relay)](https://github.com/NVIDIA/NeMo-Relay/blob/main/LICENSE) +[![GitHub](https://img.shields.io/badge/github-repo-blue?logo=github)](https://github.com/NVIDIA/NeMo-Relay/) +[![Release](https://img.shields.io/github/v/release/NVIDIA/NeMo-Relay?color=green)](https://github.com/NVIDIA/NeMo-Relay/releases) + +# NeMo Relay Switchyard Plugin + +`nemo-relay-switchyard` is NeMo Relay's experimental, source-only integration +with the [NVIDIA NeMo Switchyard](https://github.com/NVIDIA-NeMo/Switchyard) +Decision API. It adds routing-aware LLM execution intercepts to the Relay +runtime while preserving Relay ownership of provider credentials, target +bindings, dispatch, retries, fallbacks, and observability. + +This crate is intentionally not published to crates.io yet. Build it from the +NeMo Relay source checkout with the optional CLI feature while the Switchyard +Decision API contract and service/library boundary are still evolving. + +## Why use it? + +- **Route through Switchyard decisions**: Select an exact Relay-owned target + using a versioned Decision API contract. +- **Keep provider protocols stable**: Use Switchyard's translation library for + OpenAI Chat, OpenAI Responses, and Anthropic Messages request and response + translation. +- **Preserve Relay execution semantics**: Keep retries, trusted fallbacks, + credentials, streaming behavior, and optimization accounting in Relay. +- **Support staged rollout**: Run in enforce or observe-only mode with + explicit target bindings and protocol defaults. + +## What you get + +- `SwitchyardConfig`: the typed plugin configuration contract. +- `SwitchyardRuntime`: buffered and streaming routing intercepts. +- Decision and target validation for exact backend, model, protocol, and + endpoint bindings. +- ATOF-backed or payload-only routing context modes. +- Routing marks and model-routing optimization contributions for Relay's + cumulative accounting pipeline. +- Switchyard-owned protocol translation through the pinned + `switchyard-translation` dependency. + +## Source build + +The crate is a private workspace member and must be built from source: + +```bash +cargo build -p nemo-relay-cli --features switchyard +cargo test -p nemo-relay-switchyard +``` + +The resulting CLI includes the Switchyard component only when the `switchyard` +feature is enabled. A default Relay build does not include this experimental +integration. + +## Runtime boundary + +The current integration calls Switchyard's HTTP Decision API at runtime. Relay +does not start or supervise the Switchyard service. For ATOF-backed profiles, +Switchyard also provides the `/v1/atof/events` ingestion and accumulator +runtime. A reachable service must therefore be configured before routed traffic +is sent. + +The current service setup is documented in +[`examples/switchyard/README.md`](../../examples/switchyard/README.md), including +the pinned topic-branch commit, local configuration, compatibility smoke test, +and trajectory workflow. + +Translation is already in-process through Switchyard's Rust translation +library. A future in-process DecisionProvider may replace the HTTP Decision API +call without changing the Relay-owned dispatch and observability boundary. + +## Configuration and registration + +The CLI registers the component when built with `--features switchyard` and +accepts a `[[components]]` entry with `kind = "switchyard"`. A minimal +configuration selects the Decision API and trusted protocol defaults: + +```toml +[[components]] +kind = "switchyard" +enabled = true + +[components.config] +mode = "enforce" +decision_api_url = "http://127.0.0.1:4000/v1/routing/decision" +decision_profile_id = "my-profile" +context_mode = "payload_only" +request_materialization = "summary_only" + +[components.config.default_targets] +openai_chat = "my-openai-target" +openai_responses = "my-responses-target" +anthropic_messages = "my-anthropic-target" +``` + +For ATOF-backed profiles, configure an enabled Relay ATOF HTTP exporter that +targets the Switchyard `/v1/atof/events` endpoint and use environment-referenced +authentication headers. Keep provider and Decision API credentials outside +tracked configuration files. + +## Documentation + +- [Switchyard integration examples](../../examples/switchyard/README.md) +- [NeMo Relay documentation](https://docs.nvidia.com/nemo/relay) +- [Switchyard repository](https://github.com/NVIDIA-NeMo/Switchyard) diff --git a/crates/switchyard/src/component.rs b/crates/switchyard/src/component.rs new file mode 100644 index 00000000..ab48e258 --- /dev/null +++ b/crates/switchyard/src/component.rs @@ -0,0 +1,1526 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Switchyard plugin configuration and Relay execution integration. + +use std::collections::{BTreeMap, BTreeSet}; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use async_stream::stream; +use futures_util::{StreamExt, stream as futures_stream}; +use nemo_relay::api::event::{CategoryProfile, DataSchema, EventCategory}; +use nemo_relay::api::llm::LlmRequest; +use nemo_relay::api::optimization::record_llm_optimization_contribution; +use nemo_relay::api::runtime::{LlmExecutionFn, LlmJsonStream, LlmStreamExecutionFn}; +use nemo_relay::api::scope::{EmitMarkEventParams, event}; +use nemo_relay::codec::optimization::{ + LlmOptimizationContribution, LlmOptimizationKind, LlmOptimizationModel, + LlmOptimizationModelTransition, +}; +use nemo_relay::error::{FlowError, Result as FlowResult}; +use nemo_relay::plugin::{ + ConfigDiagnostic, DiagnosticLevel, Plugin, PluginComponentSpec, PluginConfig, PluginError, + PluginRegistrationContext, Result as PluginResult, deregister_plugin, register_plugin, +}; +use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value as Json, json}; +use uuid::Uuid; + +use crate::contract::{ + DecisionAttempt, DecisionProfile, ROUTING_DECISION_SCHEMA_VERSION, + ROUTING_REQUEST_SCHEMA_VERSION, RequestIdentity, RequestMaterialization, RequestProtocol, + RequestSummary, RoutingDecision, RoutingRequest, RoutingTarget, +}; +use crate::stream_translation::StreamTranscoder; +use crate::translation::{ + decode_request, encode_request, latest_user_prompt, recent_message_window, translate_response, + translation_engine, validate_portable_request, +}; + +/// Plugin kind used in Relay plugin configuration. +pub const SWITCHYARD_PLUGIN_KIND: &str = "switchyard"; + +const INTERNAL_DISPATCH_URL_HEADER: &str = "x-nemo-relay-internal-dispatch-url"; +const INTERNAL_DISPATCH_ROUTE_HEADER: &str = "x-nemo-relay-internal-dispatch-route"; +const INTERNAL_RETRY_AWARE_HEADER: &str = "x-nemo-relay-internal-retry-aware"; +const ROUTING_MARK_SCHEMA: &str = "switchyard.routing_mark"; +const ROUTING_CONTRIBUTION_SCHEMA: &str = "nvidia.switchyard.routing_optimization"; + +/// Supported provider wire protocols. +#[derive(Clone, Copy, Debug, Deserialize, Eq, Ord, PartialEq, PartialOrd, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum WireProtocol { + /// OpenAI Chat Completions. + OpenaiChat, + /// OpenAI Responses. + OpenaiResponses, + /// Anthropic Messages. + AnthropicMessages, +} + +impl WireProtocol { + fn label(self) -> &'static str { + match self { + Self::OpenaiChat => "openai_chat", + Self::OpenaiResponses => "openai_responses", + Self::AnthropicMessages => "anthropic_messages", + } + } + + fn endpoint(self) -> &'static str { + match self { + Self::OpenaiChat => "/v1/chat/completions", + Self::OpenaiResponses => "/v1/responses", + Self::AnthropicMessages => "/v1/messages", + } + } + + fn from_call(name: &str, request: &LlmRequest) -> Option { + match name { + "openai.chat_completions" | "openai_chat" | "openai_chat_completions" => { + Some(Self::OpenaiChat) + } + "openai.responses" | "openai_responses" => Some(Self::OpenaiResponses), + "anthropic.messages" | "anthropic" | "anthropic_messages" => { + Some(Self::AnthropicMessages) + } + _ if request.content.get("input").is_some() => Some(Self::OpenaiResponses), + _ if request.content.get("system").is_some() => Some(Self::AnthropicMessages), + _ if request.content.get("messages").is_some() => Some(Self::OpenaiChat), + _ => None, + } + } +} + +/// Routing rollout mode. +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum RoutingMode { + /// Apply Switchyard decisions. + #[default] + Enforce, + /// Record decisions but dispatch trusted defaults. + ObserveOnly, +} + +impl RoutingMode { + fn label(self) -> &'static str { + match self { + Self::Enforce => "enforce", + Self::ObserveOnly => "observe_only", + } + } +} + +/// Whether the selected Switchyard profile depends on ATOF-derived history. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum ContextMode { + /// The router uses only current request material. + PayloadOnly, + /// Stable identity and a configured ATOF endpoint are required. + AtofRequired, +} + +/// Exact Relay-owned backend binding for one Switchyard backend ID. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct TargetBinding { + /// Exact model expected in the Switchyard decision. + pub model: String, + /// Exact protocol expected in the Switchyard decision. + pub protocol: WireProtocol, + /// Exact endpoint expected in the Switchyard decision. + pub endpoint: String, + /// Relay-owned backend base URL. + pub base_url: String, + /// Static non-sensitive backend headers. + #[serde(default)] + pub headers: BTreeMap, + /// Backend headers resolved from environment variables. + #[serde(default)] + pub header_env: BTreeMap, +} + +/// Trusted fallback target IDs for each inbound protocol. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct ProtocolDefaults { + /// OpenAI Chat fallback target. + #[serde(default)] + pub openai_chat: String, + /// OpenAI Responses fallback target. + #[serde(default)] + pub openai_responses: String, + /// Anthropic Messages fallback target. + #[serde(default)] + pub anthropic_messages: String, +} + +impl ProtocolDefaults { + fn target(&self, protocol: WireProtocol) -> &str { + match protocol { + WireProtocol::OpenaiChat => &self.openai_chat, + WireProtocol::OpenaiResponses => &self.openai_responses, + WireProtocol::AnthropicMessages => &self.anthropic_messages, + } + } +} + +/// Versioned Switchyard plugin configuration. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +pub struct SwitchyardConfig { + /// Config schema version. + #[serde(default = "default_version")] + pub version: u32, + /// Enforce or observe-only rollout mode. + #[serde(default)] + pub mode: RoutingMode, + /// Execution-intercept priority. + #[serde(default)] + pub priority: i32, + /// Switchyard Decision API URL. + pub decision_api_url: String, + /// Switchyard profile ID. + pub decision_profile_id: String, + /// Current-request materialization. + pub request_materialization: RequestMaterialization, + /// Profile context requirement. + pub context_mode: ContextMode, + /// Decision call timeout. + #[serde(default = "default_decision_timeout_millis")] + pub decision_timeout_millis: u64, + /// Provider retries after the initial attempt. + #[serde(default = "default_max_retries")] + pub max_retries: u32, + /// Number of messages in recent-message materialization. + #[serde(default = "default_recent_message_count")] + pub recent_message_count: usize, + /// Static non-sensitive Decision API headers. + #[serde(default)] + pub decision_headers: BTreeMap, + /// Decision API headers resolved from environment variables. + #[serde(default)] + pub decision_header_env: BTreeMap, + /// Enabled inbound protocols. + #[serde(default = "default_enabled_protocols")] + pub enabled_inbound_profiles: BTreeSet, + /// Exact backend bindings keyed by Switchyard backend ID. + pub targets: BTreeMap, + /// Trusted per-protocol fallbacks. + pub default_targets: ProtocolDefaults, + /// Optional explicit ATOF endpoint used by CLI cross-validation. + #[serde(default)] + pub atof_endpoint_url: Option, +} + +impl Default for SwitchyardConfig { + fn default() -> Self { + Self { + version: default_version(), + mode: RoutingMode::default(), + priority: 0, + decision_api_url: "http://127.0.0.1:8080/v1/routing/decision".into(), + decision_profile_id: String::new(), + request_materialization: RequestMaterialization::SummaryOnly, + context_mode: ContextMode::PayloadOnly, + decision_timeout_millis: default_decision_timeout_millis(), + max_retries: default_max_retries(), + recent_message_count: default_recent_message_count(), + decision_headers: BTreeMap::new(), + decision_header_env: BTreeMap::new(), + enabled_inbound_profiles: default_enabled_protocols(), + targets: BTreeMap::new(), + default_targets: ProtocolDefaults { + openai_chat: String::new(), + openai_responses: String::new(), + anthropic_messages: String::new(), + }, + atof_endpoint_url: None, + } + } +} + +nemo_relay::editor_config! { + impl SwitchyardConfig { + mode => { label: "Rollout mode", kind: Enum, values: ["enforce", "observe_only"] }, + priority => { label: "Intercept priority", kind: Integer }, + decision_api_url => { label: "Decision API URL", kind: String }, + decision_profile_id => { label: "Decision profile ID", kind: String }, + request_materialization => { + label: "Request materialization", + kind: Enum, + values: ["none", "summary_only", "latest_user_prompt", "recent_message_window", "annotated_request", "full_body"] + }, + context_mode => { label: "Context mode", kind: Enum, values: ["payload_only", "atof_required"] }, + decision_timeout_millis => { label: "Decision timeout (ms)", kind: Integer }, + max_retries => { label: "Maximum provider retries", kind: Integer }, + recent_message_count => { label: "Recent message count", kind: Integer }, + decision_headers => { label: "Decision API static headers", kind: StringMap }, + decision_header_env => { label: "Decision API environment headers", kind: StringMap }, + enabled_inbound_profiles => { label: "Enabled inbound profiles", kind: Json }, + targets => { label: "Backend target bindings", kind: Json }, + default_targets => { label: "Trusted protocol defaults", kind: Json }, + atof_endpoint_url => { label: "ATOF endpoint URL", kind: String, optional: true } + } +} + +impl From for PluginComponentSpec { + fn from(value: SwitchyardConfig) -> Self { + let Json::Object(config) = + serde_json::to_value(value).expect("Switchyard config should serialize to an object") + else { + unreachable!("Switchyard config must serialize to an object") + }; + Self { + kind: SWITCHYARD_PLUGIN_KIND.into(), + enabled: true, + config, + } + } +} + +fn default_version() -> u32 { + 1 +} +fn default_decision_timeout_millis() -> u64 { + 25 +} +fn default_max_retries() -> u32 { + 3 +} +fn default_recent_message_count() -> usize { + 8 +} +fn default_enabled_protocols() -> BTreeSet { + BTreeSet::from([ + WireProtocol::OpenaiChat, + WireProtocol::OpenaiResponses, + WireProtocol::AnthropicMessages, + ]) +} + +struct SwitchyardPlugin; + +impl Plugin for SwitchyardPlugin { + fn plugin_kind(&self) -> &str { + SWITCHYARD_PLUGIN_KIND + } + + fn allows_multiple_components(&self) -> bool { + false + } + + fn validate(&self, plugin_config: &Map) -> Vec { + match parse_config(plugin_config).and_then(SwitchyardRuntime::new) { + Ok(_) => Vec::new(), + Err(error) => vec![ConfigDiagnostic { + level: DiagnosticLevel::Error, + code: "switchyard.invalid_config".into(), + component: Some(SWITCHYARD_PLUGIN_KIND.into()), + field: None, + message: error, + }], + } + } + + fn register<'a>( + &'a self, + plugin_config: &Map, + ctx: &'a mut PluginRegistrationContext, + ) -> Pin> + Send + 'a>> { + let parsed = parse_config(plugin_config); + Box::pin(async move { + let runtime = Arc::new( + parsed + .and_then(SwitchyardRuntime::new) + .map_err(PluginError::InvalidConfig)?, + ); + let buffered = Arc::clone(&runtime); + let buffered_intercept: LlmExecutionFn = Arc::new(move |name, request, next| { + let runtime = Arc::clone(&buffered); + let name = name.to_string(); + Box::pin(async move { runtime.execute_buffered(&name, request, next).await }) + }); + ctx.register_llm_execution_intercept( + "decision", + runtime.config.priority, + buffered_intercept, + )?; + + let streaming = Arc::clone(&runtime); + let stream_intercept: LlmStreamExecutionFn = Arc::new(move |name, request, next| { + let runtime = Arc::clone(&streaming); + let name = name.to_string(); + Box::pin(async move { runtime.execute_stream(&name, request, next).await }) + }); + ctx.register_llm_stream_execution_intercept( + "decision_stream", + runtime.config.priority, + stream_intercept, + )?; + Ok(()) + }) + } +} + +/// Register the first-party Switchyard component kind. +pub fn register_switchyard_component() -> PluginResult<()> { + match register_plugin(Arc::new(SwitchyardPlugin)) { + Ok(()) => Ok(()), + Err(PluginError::RegistrationFailed(message)) if message.contains("already registered") => { + Ok(()) + } + Err(error) => Err(error), + } +} + +/// Deregister the first-party Switchyard component kind. +pub fn deregister_switchyard_component() -> bool { + deregister_plugin(SWITCHYARD_PLUGIN_KIND) +} + +/// Validate the cross-component ATOF requirement for enabled history-backed profiles. +pub fn validate_switchyard_atof_configuration(config: &PluginConfig) -> Result<(), String> { + let Some(component) = config + .components + .iter() + .find(|component| component.enabled && component.kind == SWITCHYARD_PLUGIN_KIND) + else { + return Ok(()); + }; + let switchyard = parse_config(&component.config)?; + if switchyard.context_mode != ContextMode::AtofRequired { + return Ok(()); + } + let required_url = match &switchyard.atof_endpoint_url { + Some(url) => url.clone(), + None => derived_atof_url(&switchyard.decision_api_url)?, + }; + let observability = config + .components + .iter() + .find(|component| component.enabled && component.kind == "observability") + .ok_or_else(|| "atof_required Switchyard profiles require observability".to_string())?; + let endpoints = observability + .config + .get("atof") + .filter(|atof| atof.get("enabled").and_then(Json::as_bool) == Some(true)) + .and_then(|atof| atof.get("endpoints")) + .and_then(Json::as_array) + .ok_or_else(|| { + "atof_required Switchyard profiles require an enabled ATOF endpoint".to_string() + })?; + let endpoint = endpoints + .iter() + .find(|endpoint| { + endpoint.get("url").and_then(Json::as_str) == Some(required_url.as_str()) + && endpoint + .get("transport") + .and_then(Json::as_str) + .unwrap_or("http_post") + == "http_post" + }) + .ok_or_else(|| { + format!("atof_required Switchyard profile requires HTTP ATOF endpoint {required_url}") + })?; + if endpoint + .get("field_name_policy") + .and_then(Json::as_str) + .unwrap_or("preserve") + != "preserve" + { + return Err("Switchyard ATOF endpoint must use field_name_policy = preserve".into()); + } + if endpoint + .get("header_env") + .and_then(Json::as_object) + .is_none_or(Map::is_empty) + { + return Err( + "Switchyard ATOF endpoint authentication must use at least one environment-referenced header" + .into(), + ); + } + Ok(()) +} + +fn derived_atof_url(decision_api_url: &str) -> Result { + let mut url = reqwest::Url::parse(decision_api_url) + .map_err(|error| format!("decision_api_url is invalid: {error}"))?; + url.set_path("/v1/atof/events"); + url.set_query(None); + url.set_fragment(None); + Ok(url.to_string().trim_end_matches('/').to_string()) +} + +fn parse_config(config: &Map) -> Result { + serde_json::from_value(Json::Object(config.clone())) + .map_err(|error| format!("invalid Switchyard plugin config: {error}")) +} + +struct SwitchyardRuntime { + config: SwitchyardConfig, + client: reqwest::Client, + target_headers: BTreeMap>, + translation: switchyard_translation::TranslationEngine, +} + +impl SwitchyardRuntime { + fn new(config: SwitchyardConfig) -> Result { + validate_config(&config)?; + let headers = resolve_headers(&config.decision_headers, &config.decision_header_env)?; + let client = reqwest::Client::builder() + .default_headers(headers) + .timeout(Duration::from_millis(config.decision_timeout_millis)) + .build() + .map_err(|error| format!("failed to build Decision API client: {error}"))?; + let target_headers = config + .targets + .iter() + .map(|(id, target)| { + let headers = resolve_json_headers(&target.headers, &target.header_env)?; + Ok((id.clone(), headers)) + }) + .collect::>()?; + Ok(Self { + config, + client, + target_headers, + translation: translation_engine(), + }) + } + + async fn execute_buffered( + &self, + name: &str, + original: LlmRequest, + next: nemo_relay::api::runtime::LlmExecutionNextFn, + ) -> FlowResult { + let Some(inbound) = WireProtocol::from_call(name, &original) else { + return next(original).await; + }; + if !self.config.enabled_inbound_profiles.contains(&inbound) { + return next(original).await; + } + if let Err(error) = validate_portable_request(&self.translation, inbound, &original) { + self.emit_error( + None, + 0, + "unsupported_provider_extension", + &error.to_string(), + ); + return self + .dispatch_fallback_buffered( + inbound, + original, + next, + "unsupported_provider_extension", + ) + .await; + } + + if self.config.mode == RoutingMode::ObserveOnly { + match self.decided_request(inbound, &original, 1, None).await { + Ok((_, decision, _)) => { + self.record_routing_contribution(&decision, 1, false); + } + Err(error) => self.emit_error(None, 1, "decision_api", &error), + } + return self + .dispatch_fallback_buffered(inbound, original, next, "observe_only") + .await; + } + + let max_attempts = self.config.max_retries.saturating_add(1); + let mut previous = None; + for attempt in 1..=max_attempts { + let decided = self + .decided_request(inbound, &original, attempt, previous.clone()) + .await; + let (routing_request, decision, routed) = match decided { + Ok(value) => value, + Err(error) => { + self.emit_error(None, attempt, "decision_api", &error); + return self + .dispatch_fallback_buffered(inbound, original, next, "decision_error") + .await; + } + }; + let target_protocol = protocol_from_label(&decision.route.target_protocol_profile)?; + match next(routed).await { + Ok(response) => { + match translate_response(&self.translation, target_protocol, inbound, &response) + { + Ok(response) => { + self.record_routing_contribution(&decision, attempt, true); + return Ok(response); + } + Err(error) => { + self.emit_error( + Some(&routing_request), + attempt, + "response_translation", + &error.to_string(), + ); + return self + .dispatch_fallback_buffered( + inbound, + original, + next, + "translation_error", + ) + .await; + } + } + } + Err(error) if error_is_retryable(&error) && attempt < max_attempts => { + let retry_reason = provider_error_summary(&error); + self.emit_error(Some(&routing_request), attempt, "provider", &retry_reason); + self.emit_retry(&routing_request, &decision, attempt, &retry_reason); + previous = Some((decision.route.backend_id, retry_reason)); + } + Err(error) => { + let summary = provider_error_summary(&error); + self.emit_error(Some(&routing_request), attempt, "provider", &summary); + return self + .dispatch_fallback_buffered( + inbound, + original, + next, + if error_is_retryable(&error) { + "retry_exhausted" + } else { + "non_retryable_provider_error" + }, + ) + .await; + } + } + } + unreachable!("routing attempt loop always returns") + } + + async fn execute_stream( + &self, + name: &str, + original: LlmRequest, + next: nemo_relay::api::runtime::LlmStreamExecutionNextFn, + ) -> FlowResult { + let Some(inbound) = WireProtocol::from_call(name, &original) else { + return next(original).await; + }; + if !self.config.enabled_inbound_profiles.contains(&inbound) { + return next(original).await; + } + if let Err(error) = validate_portable_request(&self.translation, inbound, &original) { + self.emit_error( + None, + 0, + "unsupported_provider_extension", + &error.to_string(), + ); + return self + .dispatch_fallback_stream(inbound, original, next, "unsupported_provider_extension") + .await; + } + if self.config.mode == RoutingMode::ObserveOnly { + match self.decided_request(inbound, &original, 1, None).await { + Ok((_, decision, _)) => { + self.record_routing_contribution(&decision, 1, false); + } + Err(error) => self.emit_error(None, 1, "decision_api", &error), + } + return self + .dispatch_fallback_stream(inbound, original, next, "observe_only") + .await; + } + + let max_attempts = self.config.max_retries.saturating_add(1); + let mut previous = None; + for attempt in 1..=max_attempts { + let (routing_request, decision, routed) = match self + .decided_request(inbound, &original, attempt, previous.clone()) + .await + { + Ok(value) => value, + Err(error) => { + self.emit_error(None, attempt, "decision_api", &error); + return self + .dispatch_fallback_stream(inbound, original, next, "decision_error") + .await; + } + }; + let target_protocol = protocol_from_label(&decision.route.target_protocol_profile)?; + match next(routed).await { + Ok(mut upstream) => match upstream.next().await { + Some(Ok(first)) => { + self.record_routing_contribution(&decision, attempt, true); + let committed = Box::pin( + futures_stream::once(async move { Ok(first) }).chain(upstream), + ) as LlmJsonStream; + let output = if target_protocol == inbound { + committed + } else { + translated_stream( + target_protocol, + inbound, + decision.route.target_model.clone(), + committed, + ) + }; + return Ok(mark_terminal_stream( + output, + "provider_stream_committed", + self.config.mode.label(), + identity_metadata(&routing_request), + )); + } + Some(Err(error)) if error_is_retryable(&error) && attempt < max_attempts => { + let retry_reason = provider_error_summary(&error); + self.emit_error( + Some(&routing_request), + attempt, + "provider_stream_open", + &retry_reason, + ); + self.emit_retry(&routing_request, &decision, attempt, &retry_reason); + previous = Some((decision.route.backend_id, retry_reason)); + } + None if attempt < max_attempts => { + self.emit_retry(&routing_request, &decision, attempt, "empty_stream"); + previous = Some((decision.route.backend_id, "empty_stream".into())); + } + Some(Err(error)) => { + let summary = provider_error_summary(&error); + self.emit_error( + Some(&routing_request), + attempt, + "provider_stream_open", + &summary, + ); + return self + .dispatch_fallback_stream( + inbound, + original, + next, + if error_is_retryable(&error) { + "retry_exhausted" + } else { + "non_retryable_provider_error" + }, + ) + .await; + } + None => { + return self + .dispatch_fallback_stream(inbound, original, next, "empty_stream") + .await; + } + }, + Err(error) if error_is_retryable(&error) && attempt < max_attempts => { + let retry_reason = provider_error_summary(&error); + self.emit_error( + Some(&routing_request), + attempt, + "provider_stream_setup", + &retry_reason, + ); + self.emit_retry(&routing_request, &decision, attempt, &retry_reason); + previous = Some((decision.route.backend_id, retry_reason)); + } + Err(error) => { + let summary = provider_error_summary(&error); + self.emit_error( + Some(&routing_request), + attempt, + "provider_stream_setup", + &summary, + ); + return self + .dispatch_fallback_stream( + inbound, + original, + next, + if error_is_retryable(&error) { + "retry_exhausted" + } else { + "non_retryable_provider_error" + }, + ) + .await; + } + } + } + unreachable!("stream routing attempt loop always returns") + } + + async fn decided_request( + &self, + inbound: WireProtocol, + original: &LlmRequest, + attempt: u32, + previous: Option<(String, String)>, + ) -> Result<(RoutingRequest, RoutingDecision, LlmRequest), String> { + let request = self.routing_request(inbound, original, attempt, previous)?; + self.emit_requested(&request); + let started = Instant::now(); + let response = self + .client + .post(&self.config.decision_api_url) + .header("x-nemo-relay-session-id", &request.identity.session_id) + .json(&request) + .send() + .await + .map_err(|error| format!("Decision API request failed: {error}"))?; + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(format!("Decision API returned HTTP {status}: {body}")); + } + let decision = response + .json::() + .await + .map_err(|error| format!("Decision API returned invalid JSON: {error}"))?; + self.validate_decision(&decision)?; + if let Some(baseline) = decision.baseline_route.as_ref() + && let Err(error) = self.validate_target(baseline) + { + self.emit_error(Some(&request), attempt, "baseline_binding", &error); + } + let routed = self.apply_target(inbound, original.clone(), &decision)?; + let latency = started.elapsed().as_millis() as u64; + self.emit_decision( + &request, + &decision, + attempt, + self.config.mode == RoutingMode::ObserveOnly, + latency, + ); + Ok((request, decision, routed)) + } + + fn routing_request( + &self, + inbound: WireProtocol, + request: &LlmRequest, + attempt: u32, + previous: Option<(String, String)>, + ) -> Result { + let session = header(request, "x-nemo-relay-session-id"); + let stable_request_id = header(request, "x-nemo-relay-request-id"); + if self.config.context_mode == ContextMode::AtofRequired + && (session.is_none() || stable_request_id.is_none()) + { + return Err("stable session and request identity are required for this profile".into()); + } + let identity_is_stable = session.is_some() && stable_request_id.is_some(); + let synthetic_session = format!("request-{}", Uuid::now_v7()); + let session_id = session.unwrap_or_else(|| synthetic_session.clone()); + let request_id = stable_request_id.unwrap_or_else(|| format!("request-{}", Uuid::now_v7())); + let annotated = decode_request(&self.translation, inbound, request) + .map_err(|error| format!("request translation decode failed: {error}"))?; + let current_request = self.materialize(inbound, request, &annotated)?; + let (previous_route, retry_reason) = previous.unzip(); + Ok(RoutingRequest { + schema_version: ROUTING_REQUEST_SCHEMA_VERSION.into(), + decision_profile: DecisionProfile { + profile_id: self.config.decision_profile_id.clone(), + request_materialization: self.config.request_materialization, + }, + identity: RequestIdentity { + session_id, + request_id, + turn_id: header(request, "x-nemo-relay-turn-id"), + parent_scope_id: header(request, "x-nemo-relay-parent-scope-id"), + root_scope_id: header(request, "x-nemo-relay-root-scope-id"), + harness: header(request, "x-nemo-relay-agent-kind") + .unwrap_or_else(|| "unknown".into()), + source: header(request, "x-nemo-relay-source") + .unwrap_or_else(|| "nemo-relay".into()), + owner_id: header(request, "x-nemo-relay-owner-id"), + quality: header(request, "x-nemo-relay-identity-quality").unwrap_or_else(|| { + if identity_is_stable { + "explicit".into() + } else { + "synthetic".into() + } + }), + }, + protocol: RequestProtocol { + inbound_profile: inbound.label().into(), + inbound_endpoint: inbound.endpoint().into(), + desired_response_profile: inbound.label().into(), + }, + request_summary: RequestSummary { + client_requested_model: request + .content + .get("model") + .and_then(Json::as_str) + .map(ToOwned::to_owned), + prompt_token_estimate: None, + tool_count_in_payload: request + .content + .get("tools") + .and_then(Json::as_array) + .map(|tools| tools.len() as u64), + has_system_prompt: Some( + annotated.instructions.iter().any(|instruction| { + instruction.role == switchyard_translation::Role::System + }) || annotated + .messages + .iter() + .any(|message| message.role == switchyard_translation::Role::System), + ), + }, + current_request, + attempt: DecisionAttempt { + routing_attempt: attempt, + max_routing_attempts: self.config.max_retries.saturating_add(1), + previous_route, + retry_reason, + }, + }) + } + + fn materialize( + &self, + inbound: WireProtocol, + request: &LlmRequest, + annotated: &switchyard_translation::ConversationRequest, + ) -> Result, String> { + match self.config.request_materialization { + RequestMaterialization::None | RequestMaterialization::SummaryOnly => Ok(None), + RequestMaterialization::FullBody => Ok(Some(json!({"body": request.content}))), + RequestMaterialization::AnnotatedRequest => Ok(Some(json!({ + "body": request.content, + "annotated_request": annotated, + }))), + RequestMaterialization::LatestUserPrompt => { + let prompt = latest_user_prompt(annotated) + .ok_or_else(|| "latest_user_prompt requires a user message".to_string())?; + let latest = recent_message_window(annotated, 1); + let body = encode_request(&self.translation, inbound, &latest, Map::new()) + .map_err(|error| format!("latest user prompt encode failed: {error}"))? + .content; + Ok(Some(json!({"body": body, "latest_user_prompt": prompt}))) + } + RequestMaterialization::RecentMessageWindow => { + let window = recent_message_window(annotated, self.config.recent_message_count); + let body = encode_request(&self.translation, inbound, &window, Map::new()) + .map_err(|error| format!("recent window encode failed: {error}"))? + .content; + Ok(Some(json!({"body": body, "annotated_request": window}))) + } + } + } + + fn validate_decision(&self, decision: &RoutingDecision) -> Result<(), String> { + if decision.schema_version != ROUTING_DECISION_SCHEMA_VERSION { + return Err(format!( + "unsupported decision schema {:?}", + decision.schema_version + )); + } + self.validate_target(&decision.route).map(|_| ()) + } + + fn validate_target(&self, target: &RoutingTarget) -> Result<&TargetBinding, String> { + let binding = self + .config + .targets + .get(&target.backend_id) + .ok_or_else(|| format!("unknown backend_id {:?}", target.backend_id))?; + if binding.model != target.target_model + || binding.protocol.label() != target.target_protocol_profile + || binding.endpoint != target.target_endpoint + { + return Err(format!( + "decision target {:?} does not match its exact Relay binding", + target.backend_id + )); + } + Ok(binding) + } + + fn record_routing_contribution(&self, decision: &RoutingDecision, attempt: u32, applied: bool) { + let Some(contribution) = self.routing_contribution(decision, attempt, applied) else { + return; + }; + let _ = record_llm_optimization_contribution(contribution); + } + + fn routing_contribution( + &self, + decision: &RoutingDecision, + attempt: u32, + applied: bool, + ) -> Option { + let baseline = decision + .baseline_route + .as_ref() + .filter(|baseline| self.validate_target(baseline).is_ok())?; + let mut contribution = LlmOptimizationContribution::new( + SWITCHYARD_PLUGIN_KIND, + LlmOptimizationKind::model_routing(), + ); + contribution.applied = applied; + contribution.model_transition = Some(LlmOptimizationModelTransition { + baseline: Some(LlmOptimizationModel::new(&baseline.target_model)), + effective: Some(LlmOptimizationModel::new(&decision.route.target_model)), + }); + contribution.payload_schema = Some(DataSchema { + name: ROUTING_CONTRIBUTION_SCHEMA.to_string(), + version: "1".to_string(), + }); + contribution.payload = Some(json!({ + "decision_id": decision.decision_id, + "selected_backend_id": decision.route.backend_id, + "selected_tier": decision.route.tier, + "baseline_backend_id": baseline.backend_id, + "baseline_tier": baseline.tier, + "routing_attempt": attempt, + "rollout_mode": self.config.mode.label(), + "reason_code": decision.reason_code, + "reason_summary": decision.reason_summary, + "router_metadata": decision.metadata, + })); + Some(contribution) + } + + fn apply_target( + &self, + inbound: WireProtocol, + request: LlmRequest, + decision: &RoutingDecision, + ) -> Result { + let binding = self + .config + .targets + .get(&decision.route.backend_id) + .ok_or_else(|| format!("unknown backend_id {:?}", decision.route.backend_id))?; + let annotated = decode_request(&self.translation, inbound, &request) + .map_err(|error| format!("request decode failed: {error}"))?; + let mut routed = if inbound == binding.protocol { + request + } else { + encode_request( + &self.translation, + binding.protocol, + &annotated, + request.headers, + ) + .map_err(|error| format!("request translation failed: {error}"))? + }; + let object = routed + .content + .as_object_mut() + .ok_or_else(|| "translated request body is not an object".to_string())?; + object.insert("model".into(), Json::String(binding.model.clone())); + if let Some(headers) = self.target_headers.get(&decision.route.backend_id) { + routed.headers.extend(headers.clone()); + } + routed.headers.insert( + INTERNAL_DISPATCH_ROUTE_HEADER.into(), + Json::String(binding.protocol.label().into()), + ); + routed.headers.insert( + INTERNAL_DISPATCH_URL_HEADER.into(), + Json::String(dispatch_url(&binding.base_url, &binding.endpoint)), + ); + routed.headers.insert( + INTERNAL_RETRY_AWARE_HEADER.into(), + Json::String("true".into()), + ); + Ok(routed) + } + + fn fallback_request( + &self, + inbound: WireProtocol, + request: LlmRequest, + ) -> Result { + let id = self.config.default_targets.target(inbound); + let binding = self + .config + .targets + .get(id) + .ok_or_else(|| format!("unknown fallback target {id:?}"))?; + let decision = RoutingDecision { + schema_version: ROUTING_DECISION_SCHEMA_VERSION.into(), + decision_id: "relay-fallback".into(), + router: crate::contract::DecisionProvider { + name: "relay-fallback".into(), + version: "1".into(), + }, + route: crate::contract::RoutingTarget { + tier: "fallback".into(), + target_model: binding.model.clone(), + backend_id: id.to_string(), + target_protocol_profile: binding.protocol.label().into(), + target_endpoint: binding.endpoint.clone(), + }, + baseline_route: None, + confidence: None, + reason_code: Some("relay_trusted_fallback".into()), + reason_summary: None, + metadata: BTreeMap::new(), + extra: BTreeMap::new(), + }; + self.apply_target(inbound, request, &decision) + } + + async fn dispatch_fallback_buffered( + &self, + inbound: WireProtocol, + original: LlmRequest, + next: nemo_relay::api::runtime::LlmExecutionNextFn, + reason: &str, + ) -> FlowResult { + self.emit_fallback(inbound, reason, &original); + let metadata = identity_metadata_from_request(&original); + let request = self + .fallback_request(inbound, original) + .map_err(FlowError::Internal)?; + match next(request).await { + Ok(response) => Ok(response), + Err(error) => { + emit_terminal_error( + &error, + "fallback_buffered", + self.config.mode.label(), + metadata, + ); + Err(error) + } + } + } + + async fn dispatch_fallback_stream( + &self, + inbound: WireProtocol, + original: LlmRequest, + next: nemo_relay::api::runtime::LlmStreamExecutionNextFn, + reason: &str, + ) -> FlowResult { + self.emit_fallback(inbound, reason, &original); + let metadata = identity_metadata_from_request(&original); + let request = self + .fallback_request(inbound, original) + .map_err(FlowError::Internal)?; + match next(request).await { + Ok(stream) => Ok(mark_terminal_stream( + stream, + "fallback_stream", + self.config.mode.label(), + metadata.clone(), + )), + Err(error) => { + emit_terminal_error( + &error, + "fallback_stream_setup", + self.config.mode.label(), + metadata, + ); + Err(error) + } + } + } + + fn emit_requested(&self, request: &RoutingRequest) { + emit_mark( + "switchyard.routing.requested", + json!({ + "session_id": request.identity.session_id, + "request_id": request.identity.request_id, + "routing_attempt": request.attempt.routing_attempt, + "profile_id": request.decision_profile.profile_id, + "rollout_mode": self.config.mode.label(), + }), + identity_metadata(request), + ); + } + + fn emit_decision( + &self, + request: &RoutingRequest, + decision: &RoutingDecision, + attempt: u32, + observe_only: bool, + latency_ms: u64, + ) { + emit_mark( + "switchyard.routing.decision", + json!({ + "decision_id": decision.decision_id, + "profile_id": request.decision_profile.profile_id, + "router": decision.router.name, + "router_version": decision.router.version, + "routing_attempt": attempt, + "backend_id": decision.route.backend_id, + "selected_tier": decision.route.tier, + "selected_model": decision.route.target_model, + "target_protocol_profile": decision.route.target_protocol_profile, + "target_endpoint": decision.route.target_endpoint, + "confidence": decision.confidence, + "reason_code": decision.reason_code, + "reason_summary": decision.reason_summary, + "router_metadata": decision.metadata, + "latency_ms": latency_ms, + "observe_only": observe_only, + "rollout_mode": self.config.mode.label(), + }), + identity_metadata(request), + ); + } + + fn emit_retry( + &self, + request: &RoutingRequest, + decision: &RoutingDecision, + attempt: u32, + reason: &str, + ) { + emit_mark( + "switchyard.routing.retry", + json!({"routing_attempt": attempt, "previous_route": decision.route.backend_id, "retry_reason": reason, "rollout_mode": self.config.mode.label()}), + identity_metadata(request), + ); + } + + fn emit_error(&self, request: Option<&RoutingRequest>, attempt: u32, class: &str, error: &str) { + emit_mark( + "switchyard.routing.error", + json!({"routing_attempt": attempt, "error_class": class, "error": error, "rollout_mode": self.config.mode.label()}), + request.map(identity_metadata).unwrap_or_else(|| json!({})), + ); + } + + fn emit_fallback(&self, inbound: WireProtocol, reason: &str, request: &LlmRequest) { + emit_mark( + "switchyard.routing.fallback", + json!({ + "fallback_reason": reason, + "fallback_route": self.config.default_targets.target(inbound), + "inbound_profile": inbound.label(), + "rollout_mode": self.config.mode.label(), + }), + identity_metadata_from_request(request), + ); + } +} + +fn validate_config(config: &SwitchyardConfig) -> Result<(), String> { + if config.version != 1 { + return Err(format!( + "unsupported Switchyard config version {}", + config.version + )); + } + if config.decision_profile_id.trim().is_empty() { + return Err("decision_profile_id must be non-empty".into()); + } + if config.decision_timeout_millis == 0 { + return Err("decision_timeout_millis must be greater than zero".into()); + } + if config.max_retries > 10 { + return Err("max_retries must not exceed 10".into()); + } + if config.recent_message_count == 0 { + return Err("recent_message_count must be greater than zero".into()); + } + let url = reqwest::Url::parse(&config.decision_api_url) + .map_err(|error| format!("decision_api_url is invalid: {error}"))?; + if !matches!(url.scheme(), "http" | "https") { + return Err("decision_api_url must use http or https".into()); + } + if config.targets.is_empty() { + return Err("targets must not be empty".into()); + } + if config.enabled_inbound_profiles.is_empty() { + return Err("enabled_inbound_profiles must not be empty".into()); + } + let mut exact_bindings = BTreeSet::new(); + for (id, target) in &config.targets { + if id.trim().is_empty() + || target.model.trim().is_empty() + || target.endpoint.trim().is_empty() + { + return Err("target IDs, models, and endpoints must be non-empty".into()); + } + let base_url = reqwest::Url::parse(&target.base_url) + .map_err(|error| format!("target {id:?} base_url is invalid: {error}"))?; + if !matches!(base_url.scheme(), "http" | "https") { + return Err(format!("target {id:?} base_url must use http or https")); + } + if target.endpoint != target.protocol.endpoint() { + return Err(format!( + "target {id:?} endpoint must be {:?} for {}", + target.protocol.endpoint(), + target.protocol.label() + )); + } + if !exact_bindings.insert(( + target.model.clone(), + target.protocol, + target.endpoint.clone(), + target.base_url.trim_end_matches('/').to_string(), + )) { + return Err(format!( + "target {id:?} conflicts with another exact backend binding" + )); + } + } + for &protocol in &config.enabled_inbound_profiles { + let id = config.default_targets.target(protocol); + let target = config + .targets + .get(id) + .ok_or_else(|| format!("default target {id:?} is not configured"))?; + if target.protocol != protocol { + return Err(format!( + "default target {id:?} must use protocol {}", + protocol.label() + )); + } + } + Ok(()) +} + +fn resolve_headers( + static_headers: &BTreeMap, + environment_headers: &BTreeMap, +) -> Result { + let mut headers = HeaderMap::new(); + for (name, value) in static_headers { + insert_http_header(&mut headers, name, value)?; + } + for (name, variable) in environment_headers { + if static_headers.contains_key(name) { + return Err(format!( + "header {name:?} cannot appear in both headers and header_env" + )); + } + let value = std::env::var(variable) + .map_err(|_| format!("environment variable {variable:?} is not set"))?; + if value.trim().is_empty() { + return Err(format!("environment variable {variable:?} is blank")); + } + insert_http_header(&mut headers, name, &value)?; + } + Ok(headers) +} + +fn resolve_json_headers( + static_headers: &BTreeMap, + environment_headers: &BTreeMap, +) -> Result, String> { + let mut headers = Map::new(); + for (name, value) in static_headers { + headers.insert(name.clone(), Json::String(value.clone())); + } + for (name, variable) in environment_headers { + if static_headers.contains_key(name) { + return Err(format!( + "target header {name:?} cannot appear in both headers and header_env" + )); + } + let value = std::env::var(variable) + .map_err(|_| format!("environment variable {variable:?} is not set"))?; + if value.trim().is_empty() { + return Err(format!("environment variable {variable:?} is blank")); + } + headers.insert(name.clone(), Json::String(value)); + } + Ok(headers) +} + +fn insert_http_header(headers: &mut HeaderMap, name: &str, value: &str) -> Result<(), String> { + let name = HeaderName::from_bytes(name.as_bytes()) + .map_err(|error| format!("invalid header name: {error}"))?; + let value = + HeaderValue::from_str(value).map_err(|error| format!("invalid header value: {error}"))?; + headers.insert(name, value); + Ok(()) +} + +fn protocol_from_label(label: &str) -> FlowResult { + match label { + "openai_chat" | "openai_chat_completions" | "openai_chat_completions.v1" => { + Ok(WireProtocol::OpenaiChat) + } + "openai_responses" | "openai_responses.v1" => Ok(WireProtocol::OpenaiResponses), + "anthropic_messages" | "anthropic_messages.v1" => Ok(WireProtocol::AnthropicMessages), + value => Err(FlowError::InvalidArgument(format!( + "unsupported Switchyard target protocol {value:?}" + ))), + } +} + +fn header(request: &LlmRequest, name: &str) -> Option { + request + .headers + .get(name) + .and_then(Json::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn dispatch_url(base_url: &str, endpoint: &str) -> String { + let base = base_url.trim_end_matches('/'); + let endpoint = if base.ends_with("/v1") && endpoint.starts_with("/v1/") { + &endpoint[3..] + } else { + endpoint + }; + format!("{base}{endpoint}") +} + +fn identity_metadata(request: &RoutingRequest) -> Json { + json!({ + "session_id": request.identity.session_id, + "request_id": request.identity.request_id, + "turn_id": request.identity.turn_id, + "owner_id": request.identity.owner_id, + }) +} + +fn identity_metadata_from_request(request: &LlmRequest) -> Json { + json!({ + "session_id": header(request, "x-nemo-relay-session-id"), + "request_id": header(request, "x-nemo-relay-request-id"), + "turn_id": header(request, "x-nemo-relay-turn-id"), + "owner_id": header(request, "x-nemo-relay-owner-id"), + }) +} + +fn error_is_retryable(error: &FlowError) -> bool { + matches!(error, FlowError::Upstream(failure) if failure.is_retryable()) +} + +fn emit_mark(name: &str, data: Json, metadata: Json) { + if let Err(error) = event( + EmitMarkEventParams::builder() + .name(name) + .data(data) + .data_schema( + DataSchema::builder() + .name(ROUTING_MARK_SCHEMA) + .version("1") + .build(), + ) + .metadata(metadata) + .category(EventCategory::custom()) + .category_profile(CategoryProfile::builder().subtype(name).build()) + .build(), + ) { + eprintln!("nemo-relay switchyard: failed to emit {name}: {error}"); + } +} + +fn emit_terminal_error(error: &FlowError, phase: &str, rollout_mode: &str, metadata: Json) { + emit_mark( + "switchyard.routing.terminal_error", + json!({"error_class": provider_error_class(error), "error": provider_error_summary(error), "phase": phase, "rollout_mode": rollout_mode}), + metadata, + ); +} + +fn provider_error_class(error: &FlowError) -> &'static str { + match error { + FlowError::Upstream(failure) => match failure.class { + nemo_relay::error::UpstreamFailureClass::Connection => "connection", + nemo_relay::error::UpstreamFailureClass::Timeout => "timeout", + nemo_relay::error::UpstreamFailureClass::RetryableStatus => "retryable_status", + nemo_relay::error::UpstreamFailureClass::ContextWindow => "context_window", + nemo_relay::error::UpstreamFailureClass::ModelUnavailable => "model_unavailable", + nemo_relay::error::UpstreamFailureClass::Authentication => "authentication", + nemo_relay::error::UpstreamFailureClass::InvalidRequest => "invalid_request", + nemo_relay::error::UpstreamFailureClass::Other => "other", + }, + _ => "relay", + } +} + +fn provider_error_summary(error: &FlowError) -> String { + match error { + FlowError::Upstream(failure) => match failure.status { + Some(status) => format!("{}:http_{status}", provider_error_class(error)), + None => provider_error_class(error).to_string(), + }, + _ => error.to_string(), + } +} + +fn mark_terminal_stream( + mut upstream: LlmJsonStream, + phase: &'static str, + rollout_mode: &'static str, + metadata: Json, +) -> LlmJsonStream { + Box::pin(stream! { + while let Some(item) = upstream.next().await { + match item { + Ok(chunk) => yield Ok(chunk), + Err(error) => { + emit_terminal_error(&error, phase, rollout_mode, metadata.clone()); + yield Err(error); + return; + } + } + } + }) +} + +fn translated_stream( + source: WireProtocol, + target: WireProtocol, + effective_model: String, + mut upstream: LlmJsonStream, +) -> LlmJsonStream { + let mut transcoder = StreamTranscoder::new(source, target, effective_model); + Box::pin(stream! { + while let Some(item) = upstream.next().await { + match item { + Ok(chunk) => { + match transcoder.transcode(&chunk) { + Ok(chunks) => { + for chunk in chunks { + yield Ok(chunk); + } + } + Err(error) => { + yield Err(error); + return; + } + } + } + Err(error) => { + yield Err(error); + return; + } + } + } + match transcoder.finish() { + Ok(chunks) => { + for chunk in chunks { + yield Ok(chunk); + } + } + Err(error) => yield Err(error), + } + }) +} + +#[cfg(test)] +#[path = "../tests/unit/component_tests.rs"] +mod tests; diff --git a/crates/switchyard/src/contract.rs b/crates/switchyard/src/contract.rs new file mode 100644 index 00000000..675234ca --- /dev/null +++ b/crates/switchyard/src/contract.rs @@ -0,0 +1,187 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! JSON compatibility types for the Switchyard Decision API. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +/// Routing request schema supported by this plugin. +pub const ROUTING_REQUEST_SCHEMA_VERSION: &str = "switchyard.routing_request.v1"; +/// Routing decision schema supported by this plugin. +pub const ROUTING_DECISION_SCHEMA_VERSION: &str = "switchyard.routing_decision.v1"; + +/// Request-time materialization supplied to Switchyard. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +pub enum RequestMaterialization { + /// Identity, protocol, summary, and attempt only. + None, + /// Baseline summary without current request material. + SummaryOnly, + /// Latest user prompt. + LatestUserPrompt, + /// Bounded recent message window. + RecentMessageWindow, + /// Relay-normalized request plus its provider body. + AnnotatedRequest, + /// Complete provider request body. + FullBody, +} + +/// Switchyard profile selection. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct DecisionProfile { + /// Profile ID loaded by Switchyard. + pub profile_id: String, + /// Request materialization mode. + pub request_materialization: RequestMaterialization, +} + +/// Normalized Relay identity. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct RequestIdentity { + /// Stable session identifier. + pub session_id: String, + /// Per-request identifier. + pub request_id: String, + /// Optional turn identifier. + #[serde(skip_serializing_if = "Option::is_none")] + pub turn_id: Option, + /// Optional parent scope identifier. + #[serde(skip_serializing_if = "Option::is_none")] + pub parent_scope_id: Option, + /// Optional root scope identifier. + #[serde(skip_serializing_if = "Option::is_none")] + pub root_scope_id: Option, + /// Harness name. + pub harness: String, + /// Request source. + pub source: String, + /// Optional resolved work owner. + #[serde(skip_serializing_if = "Option::is_none")] + pub owner_id: Option, + /// Native, explicit, or synthetic identity quality. + pub quality: String, +} + +/// Inbound protocol context. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct RequestProtocol { + /// Inbound protocol profile. + pub inbound_profile: String, + /// Inbound endpoint. + pub inbound_endpoint: String, + /// Response profile expected by the harness. + pub desired_response_profile: String, +} + +/// Cheap provider-request summary. +#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)] +pub struct RequestSummary { + /// Client-requested model. + #[serde(skip_serializing_if = "Option::is_none")] + pub client_requested_model: Option, + /// Optional prompt-token estimate. + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_token_estimate: Option, + /// Number of tools in the request. + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_count_in_payload: Option, + /// Whether a system prompt is present. + #[serde(skip_serializing_if = "Option::is_none")] + pub has_system_prompt: Option, +} + +/// Routing-attempt context. Additive fields are ignored by older Switchyard servers. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct DecisionAttempt { + /// One-indexed routing attempt. + pub routing_attempt: u32, + /// Maximum Decision API attempts. + pub max_routing_attempts: u32, + /// Previously selected backend. + #[serde(skip_serializing_if = "Option::is_none")] + pub previous_route: Option, + /// Reason another decision is requested. + #[serde(skip_serializing_if = "Option::is_none")] + pub retry_reason: Option, +} + +/// Canonical routing request sent by Relay. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct RoutingRequest { + /// Schema identifier. + pub schema_version: String, + /// Switchyard profile selection. + pub decision_profile: DecisionProfile, + /// Normalized identity. + pub identity: RequestIdentity, + /// Inbound protocol. + pub protocol: RequestProtocol, + /// Cheap request summary. + pub request_summary: RequestSummary, + /// Optional request materialization. + #[serde(skip_serializing_if = "Option::is_none")] + pub current_request: Option, + /// Attempt metadata. + pub attempt: DecisionAttempt, +} + +/// Router metadata returned by Switchyard. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct DecisionProvider { + /// Router name. + pub name: String, + /// Router version. + pub version: String, +} + +/// Selected target returned by Switchyard. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct RoutingTarget { + /// Tier label. + pub tier: String, + /// Selected model. + pub target_model: String, + /// Backend binding ID. + pub backend_id: String, + /// Target protocol profile. + pub target_protocol_profile: String, + /// Target endpoint. + pub target_endpoint: String, +} + +/// Canonical Switchyard routing decision. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct RoutingDecision { + /// Schema identifier. + pub schema_version: String, + /// Decision identifier. + pub decision_id: String, + /// Router metadata. + pub router: DecisionProvider, + /// Selected target. + pub route: RoutingTarget, + /// Explicit counterfactual route supplied by routers that define one. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub baseline_route: Option, + /// Optional confidence. + #[serde(default)] + pub confidence: Option, + /// Optional reason code. + #[serde(default)] + pub reason_code: Option, + /// Optional reason summary. + #[serde(default)] + pub reason_summary: Option, + /// Additive router metadata. + #[serde(default)] + pub metadata: BTreeMap, + /// Unknown additive response fields. + #[serde(flatten)] + pub extra: BTreeMap, +} diff --git a/crates/switchyard/src/lib.rs b/crates/switchyard/src/lib.rs new file mode 100644 index 00000000..238d4337 --- /dev/null +++ b/crates/switchyard/src/lib.rs @@ -0,0 +1,17 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#![deny(rustdoc::broken_intra_doc_links, rustdoc::private_intra_doc_links)] + +//! First-party Switchyard Decision API routing plugin for NeMo Relay. + +pub mod component; +pub mod contract; +mod stream_translation; +mod translation; + +pub use component::{ + ContextMode, ProtocolDefaults, RoutingMode, SWITCHYARD_PLUGIN_KIND, SwitchyardConfig, + TargetBinding, WireProtocol, deregister_switchyard_component, register_switchyard_component, + validate_switchyard_atof_configuration, +}; diff --git a/crates/switchyard/src/stream_translation.rs b/crates/switchyard/src/stream_translation.rs new file mode 100644 index 00000000..808834a2 --- /dev/null +++ b/crates/switchyard/src/stream_translation.rs @@ -0,0 +1,94 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Relay stream adapter for Switchyard's provider translation engine. + +use nemo_relay::error::{FlowError, Result}; +use serde_json::Value as Json; +use switchyard_translation::{StreamTranslationState, TranslationEngine}; + +use crate::component::WireProtocol; +use crate::translation::wire_format; + +pub(crate) struct StreamTranscoder { + engine: TranslationEngine, + source: WireProtocol, + target: WireProtocol, + state: StreamTranslationState, +} + +impl StreamTranscoder { + pub(crate) fn new( + source: WireProtocol, + target: WireProtocol, + effective_model: impl Into, + ) -> Self { + let effective_model = effective_model.into(); + let mut state = StreamTranslationState::new(wire_format(source), wire_format(target)); + state.model = Some(effective_model.clone()); + state.target_model = Some(effective_model); + Self { + engine: TranslationEngine::default(), + source, + target, + state, + } + } + + pub(crate) fn transcode(&mut self, chunk: &Json) -> Result> { + if unsupported_stream_chunk(self.source, chunk) { + return Err(FlowError::InvalidArgument( + "provider-specific streaming extension cannot be translated safely".into(), + )); + } + self.engine + .translate_event( + &mut self.state, + wire_format(self.source), + wire_format(self.target), + chunk, + ) + .map_err(|error| { + FlowError::InvalidArgument(format!("Switchyard stream translation failed: {error}")) + }) + } + + pub(crate) fn finish(&mut self) -> Result> { + self.engine + .finish_stream(&mut self.state, wire_format(self.target)) + .map_err(|error| { + FlowError::InvalidArgument(format!( + "Switchyard stream finalization failed: {error}" + )) + }) + } +} + +fn unsupported_stream_chunk(source: WireProtocol, chunk: &Json) -> bool { + match source { + WireProtocol::OpenaiChat => { + chunk["choices"] + .as_array() + .into_iter() + .flatten() + .any(|choice| { + choice["delta"].get("audio").is_some() + || choice["delta"].get("reasoning_content").is_some() + }) + } + WireProtocol::OpenaiResponses => { + chunk.get("type").and_then(Json::as_str) == Some("response.output_item.added") + && matches!( + chunk["item"].get("type").and_then(Json::as_str), + Some("reasoning" | "computer_call" | "web_search_call") + ) + } + WireProtocol::AnthropicMessages => { + chunk.get("type").and_then(Json::as_str) == Some("content_block_start") + && !matches!( + chunk["content_block"].get("type").and_then(Json::as_str), + Some("text" | "tool_use") + ) + } + } +} diff --git a/crates/switchyard/src/translation.rs b/crates/switchyard/src/translation.rs new file mode 100644 index 00000000..65600033 --- /dev/null +++ b/crates/switchyard/src/translation.rs @@ -0,0 +1,316 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use nemo_relay::api::llm::LlmRequest; +use nemo_relay::error::{FlowError, Result}; +use serde_json::{Map, Value as Json}; +use switchyard_translation::{ + ContentBlock, ConversationRequest, DeterministicIdPolicy, ImageSource, LossyConversionPolicy, + PreservationPolicy, Role, TargetCapabilities, TranslationDiagnostic, TranslationEngine, + TranslationPolicy, UnknownFieldPolicy, WireFormat, +}; + +use crate::component::WireProtocol; + +pub(crate) fn translation_engine() -> TranslationEngine { + TranslationEngine::default() +} + +pub(crate) fn decode_request( + engine: &TranslationEngine, + protocol: WireProtocol, + request: &LlmRequest, +) -> Result { + let output = engine + .decode_request( + wire_format(protocol), + &request.content, + &translation_policy(), + ) + .map_err(translation_error)?; + ensure_no_diagnostics(&output.diagnostics)?; + Ok(output.request) +} + +pub(crate) fn encode_request( + engine: &TranslationEngine, + protocol: WireProtocol, + request: &ConversationRequest, + headers: Map, +) -> Result { + let output = engine + .encode_request(wire_format(protocol), request, &translation_policy()) + .map_err(translation_error)?; + ensure_no_diagnostics(&output.diagnostics)?; + Ok(LlmRequest { + headers, + content: output.body, + }) +} + +pub(crate) fn validate_portable_request( + engine: &TranslationEngine, + protocol: WireProtocol, + request: &LlmRequest, +) -> Result<()> { + const RESTRICTED_KEYS: &[&str] = &[ + "cache_control", + "audio", + "thinking", + "computer_use", + "server_tool_use", + ]; + if protocol == WireProtocol::AnthropicMessages + && contains_invalid_anthropic_image_source(&request.content) + { + return Err(FlowError::InvalidArgument( + "request uses an unsupported or malformed Anthropic image source".into(), + )); + } + if contains_any_key_recursive(&request.content, RESTRICTED_KEYS) { + return Err(FlowError::InvalidArgument( + "request uses a provider-specific extension that requires same-protocol fail-open" + .into(), + )); + } + let decoded = decode_request(engine, protocol, request)?; + if decoded.reasoning.effort.is_some() + || decoded.reasoning.raw.is_some() + || decoded + .extensions + .fields + .iter() + .any(|(key, value)| key != "stream_options" || !portable_stream_options(value)) + || request_contains_unsupported_content(&decoded) + { + return Err(FlowError::InvalidArgument( + "request uses provider-specific fields that cannot be translated safely".into(), + )); + } + Ok(()) +} + +pub(crate) fn latest_user_prompt(request: &ConversationRequest) -> Option { + request + .messages + .iter() + .rev() + .find(|message| message.role == Role::User) + .and_then(|message| message.text_content("\n")) +} + +pub(crate) fn recent_message_window( + request: &ConversationRequest, + count: usize, +) -> ConversationRequest { + let mut window = request.clone(); + let split = window.messages.len().saturating_sub(count); + window.messages = window.messages.split_off(split); + window +} + +pub(crate) fn translate_response( + engine: &TranslationEngine, + source: WireProtocol, + target: WireProtocol, + response: &Json, +) -> Result { + if source == target { + return Ok(response.clone()); + } + ensure_portable_response(source, response)?; + let output = engine + .translate_response( + wire_format(source), + wire_format(target), + response, + &translation_policy(), + ) + .map_err(translation_error)?; + ensure_no_diagnostics(&output.diagnostics)?; + Ok(output.body) +} + +pub(crate) const fn wire_format(protocol: WireProtocol) -> WireFormat { + match protocol { + WireProtocol::OpenaiChat => WireFormat::OpenAiChat, + WireProtocol::OpenaiResponses => WireFormat::OpenAiResponses, + WireProtocol::AnthropicMessages => WireFormat::AnthropicMessages, + } +} + +fn translation_policy() -> TranslationPolicy { + TranslationPolicy { + unknown_field_policy: UnknownFieldPolicy::Reject, + lossy_conversion_policy: LossyConversionPolicy::Reject, + deterministic_ids: DeterministicIdPolicy::GenerateStable { + prefix: "relay".into(), + }, + preservation: PreservationPolicy::Disabled, + target_capabilities: TargetCapabilities::default(), + } +} + +fn translation_error(error: switchyard_translation::TranslationError) -> FlowError { + FlowError::InvalidArgument(format!("Switchyard translation failed: {error}")) +} + +fn ensure_no_diagnostics(diagnostics: &[TranslationDiagnostic]) -> Result<()> { + if diagnostics.is_empty() { + Ok(()) + } else { + Err(FlowError::InvalidArgument(format!( + "Switchyard translation was not lossless: {diagnostics:?}" + ))) + } +} + +fn portable_stream_options(value: &Json) -> bool { + let Some(options) = value.as_object() else { + return false; + }; + options.len() == 1 && options.get("include_usage").is_some_and(Json::is_boolean) +} + +fn request_contains_unsupported_content(request: &ConversationRequest) -> bool { + request + .instructions + .iter() + .flat_map(|instruction| instruction.content.iter()) + .chain( + request + .messages + .iter() + .flat_map(|message| message.content.iter()), + ) + .any(unsupported_content_block) +} + +fn unsupported_content_block(block: &ContentBlock) -> bool { + match block { + ContentBlock::Text { .. } | ContentBlock::Refusal { .. } | ContentBlock::ToolCall(_) => { + false + } + ContentBlock::Image { source } => invalid_image_source(source), + ContentBlock::ToolResult(result) => result.content.iter().any(unsupported_content_block), + ContentBlock::Reasoning { .. } + | ContentBlock::Audio { .. } + | ContentBlock::Video { .. } + | ContentBlock::File { .. } + | ContentBlock::Unknown { .. } => true, + } +} + +fn invalid_image_source(source: &ImageSource) -> bool { + match source { + ImageSource::Url { url, .. } => { + url.starts_with("data:") && base64_data_uri_parts(url).is_none() + } + ImageSource::Base64 { media_type, data } => { + media_type.as_deref().is_none_or(str::is_empty) || data.is_empty() + } + ImageSource::Raw(_) => true, + } +} + +fn base64_data_uri_parts(url: &str) -> Option<(&str, &str)> { + let (metadata, data) = url.strip_prefix("data:")?.split_once(',')?; + let media_type = metadata.strip_suffix(";base64")?; + (!media_type.is_empty() && !data.is_empty()).then_some((media_type, data)) +} + +fn contains_invalid_anthropic_image_source(value: &Json) -> bool { + match value { + Json::Object(object) => { + if object.get("type").and_then(Json::as_str) == Some("image") { + let Some(source) = object.get("source").and_then(Json::as_object) else { + return true; + }; + match source.get("type").and_then(Json::as_str) { + Some("url") => source + .get("url") + .and_then(Json::as_str) + .is_none_or(str::is_empty), + Some("base64") => { + source + .get("media_type") + .and_then(Json::as_str) + .is_none_or(str::is_empty) + || source + .get("data") + .and_then(Json::as_str) + .is_none_or(str::is_empty) + } + _ => true, + } + } else { + object.values().any(contains_invalid_anthropic_image_source) + } + } + Json::Array(items) => items.iter().any(contains_invalid_anthropic_image_source), + _ => false, + } +} + +fn contains_any_key_recursive(value: &Json, keys: &[&str]) -> bool { + match value { + Json::Object(object) => { + object.keys().any(|key| keys.contains(&key.as_str())) + || object + .values() + .any(|value| contains_any_key_recursive(value, keys)) + } + Json::Array(items) => items + .iter() + .any(|value| contains_any_key_recursive(value, keys)), + _ => false, + } +} + +fn ensure_portable_response(protocol: WireProtocol, response: &Json) -> Result<()> { + let unsupported = + match protocol { + WireProtocol::OpenaiChat => { + response["choices"] + .as_array() + .into_iter() + .flatten() + .any(|choice| { + choice["message"].get("audio").is_some() + || choice["message"].get("reasoning_content").is_some() + }) + } + WireProtocol::OpenaiResponses => response["output"] + .as_array() + .into_iter() + .flatten() + .any(|item| { + matches!( + item.get("type").and_then(Json::as_str), + Some( + "reasoning" + | "computer_call" + | "computer_call_output" + | "web_search_call" + ) + ) + }), + WireProtocol::AnthropicMessages => response["content"] + .as_array() + .into_iter() + .flatten() + .any(|block| { + !matches!( + block.get("type").and_then(Json::as_str), + Some("text" | "tool_use") + ) + }), + }; + if unsupported { + Err(FlowError::InvalidArgument( + "provider-specific response extension cannot be translated safely".into(), + )) + } else { + Ok(()) + } +} diff --git a/crates/switchyard/tests/contract.rs b/crates/switchyard/tests/contract.rs new file mode 100644 index 00000000..7051f6a5 --- /dev/null +++ b/crates/switchyard/tests/contract.rs @@ -0,0 +1,79 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Integration tests for the public Switchyard Decision API contract. + +use nemo_relay_switchyard::contract::{DecisionAttempt, RoutingDecision}; +use serde_json::json; + +#[test] +fn current_switchyard_decision_contract_accepts_additive_fields() { + let decision: RoutingDecision = serde_json::from_value(json!({ + "schema_version": "switchyard.routing_decision.v1", + "decision_id": "decision-1", + "router": {"name": "stage_router", "version": "1"}, + "route": { + "tier": "capable", + "target_model": "model-a", + "backend_id": "backend-a", + "target_protocol_profile": "openai_chat", + "target_endpoint": "/v1/chat/completions" + }, + "confidence": 0.8, + "future_field": {"safe": true} + })) + .unwrap(); + assert_eq!( + decision.extra.get("future_field"), + Some(&json!({"safe": true})) + ); + assert!(decision.baseline_route.is_none()); +} + +#[test] +fn current_switchyard_decision_contract_accepts_an_explicit_baseline() { + let decision: RoutingDecision = serde_json::from_value(json!({ + "schema_version": "switchyard.routing_decision.v1", + "decision_id": "decision-1", + "router": {"name": "stage", "version": "1"}, + "route": { + "tier": "efficient", + "target_model": "model-small", + "backend_id": "backend-small", + "target_protocol_profile": "openai_chat", + "target_endpoint": "/v1/chat/completions" + }, + "baseline_route": { + "tier": "capable", + "target_model": "model-large", + "backend_id": "backend-large", + "target_protocol_profile": "anthropic_messages", + "target_endpoint": "/v1/messages" + } + })) + .unwrap(); + let baseline = decision.baseline_route.unwrap(); + assert_eq!(baseline.backend_id, "backend-large"); + assert_eq!(baseline.target_protocol_profile, "anthropic_messages"); +} + +#[test] +fn malformed_decisions_are_rejected_by_serde() { + let missing_route = json!({ + "schema_version": "switchyard.routing_decision.v1", + "decision_id": "decision-1", + "router": {"name": "stage_router", "version": "1"} + }); + assert!(serde_json::from_value::(missing_route).is_err()); +} + +#[test] +fn additive_retry_metadata_stays_optional_on_the_wire() { + let attempt: DecisionAttempt = serde_json::from_value(json!({ + "routing_attempt": 1, + "max_routing_attempts": 4 + })) + .unwrap(); + assert!(attempt.previous_route.is_none()); + assert!(attempt.retry_reason.is_none()); +} diff --git a/crates/switchyard/tests/unit/component_tests.rs b/crates/switchyard/tests/unit/component_tests.rs new file mode 100644 index 00000000..a6e7d5e9 --- /dev/null +++ b/crates/switchyard/tests/unit/component_tests.rs @@ -0,0 +1,1448 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Unit tests for the Switchyard Relay plugin component. + +use std::sync::Mutex; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use axum::{Json as AxumJson, Router, extract::State, routing::post}; +use nemo_relay::api::event::{Event, ScopeCategory}; +use nemo_relay::api::llm::{ + LlmCallExecuteParams, LlmStreamCallExecuteParams, llm_call_execute, llm_stream_call_execute, +}; +use nemo_relay::api::runtime::{LlmExecutionNextFn, LlmStreamExecutionNextFn}; +use nemo_relay::api::subscriber::{deregister_subscriber, flush_subscribers, register_subscriber}; +use nemo_relay::codec::optimization::LlmOptimizationSummaryStatus; +use nemo_relay::error::{UpstreamFailure, UpstreamFailureClass}; +use nemo_relay::plugin::rollback_registrations; + +use super::*; + +fn binding(protocol: WireProtocol, model: &str) -> TargetBinding { + TargetBinding { + model: model.into(), + protocol, + endpoint: protocol.endpoint().into(), + base_url: "http://127.0.0.1:9999".into(), + headers: BTreeMap::new(), + header_env: BTreeMap::new(), + } +} + +fn config(decision_api_url: String) -> SwitchyardConfig { + SwitchyardConfig { + decision_api_url, + decision_profile_id: "stage_router".into(), + request_materialization: RequestMaterialization::SummaryOnly, + context_mode: ContextMode::PayloadOnly, + decision_timeout_millis: 1_000, + targets: BTreeMap::from([ + ( + "selected-chat".into(), + binding(WireProtocol::OpenaiChat, "selected"), + ), + ( + "baseline-chat".into(), + binding(WireProtocol::OpenaiChat, "baseline"), + ), + ( + "fallback-chat".into(), + binding(WireProtocol::OpenaiChat, "fallback"), + ), + ( + "fallback-responses".into(), + binding(WireProtocol::OpenaiResponses, "fallback"), + ), + ( + "fallback-anthropic".into(), + binding(WireProtocol::AnthropicMessages, "fallback"), + ), + ]), + default_targets: ProtocolDefaults { + openai_chat: "fallback-chat".into(), + openai_responses: "fallback-responses".into(), + anthropic_messages: "fallback-anthropic".into(), + }, + ..SwitchyardConfig::default() + } +} + +fn decision() -> RoutingDecision { + RoutingDecision { + schema_version: ROUTING_DECISION_SCHEMA_VERSION.into(), + decision_id: "decision-1".into(), + router: crate::contract::DecisionProvider { + name: "stage_router".into(), + version: "1".into(), + }, + route: crate::contract::RoutingTarget { + tier: "efficient".into(), + target_model: "selected".into(), + backend_id: "selected-chat".into(), + target_protocol_profile: "openai_chat".into(), + target_endpoint: "/v1/chat/completions".into(), + }, + baseline_route: Some(crate::contract::RoutingTarget { + tier: "capable".into(), + target_model: "baseline".into(), + backend_id: "baseline-chat".into(), + target_protocol_profile: "openai_chat".into(), + target_endpoint: "/v1/chat/completions".into(), + }), + confidence: Some(0.9), + reason_code: Some("test".into()), + reason_summary: None, + metadata: BTreeMap::from([ + ("feature_state".into(), json!("fresh")), + ("snapshot_age_millis".into(), json!(37)), + ("snapshot_max_age_millis".into(), json!(300_000)), + ]), + extra: BTreeMap::new(), + } +} + +fn chat_request() -> LlmRequest { + LlmRequest { + headers: Map::new(), + content: json!({ + "model": "inbound", + "messages": [ + {"role": "system", "content": "system"}, + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "answer"}, + {"role": "user", "content": "latest"} + ] + }), + } +} + +fn request(protocol: WireProtocol) -> LlmRequest { + let content = match protocol { + WireProtocol::OpenaiChat => chat_request().content, + WireProtocol::OpenaiResponses => { + json!({"model": "inbound", "instructions": "system", "input": "latest"}) + } + WireProtocol::AnthropicMessages => { + json!({"model": "inbound", "system": "system", "max_tokens": 32, "messages": [{"role": "user", "content": "latest"}]}) + } + }; + LlmRequest { + headers: Map::new(), + content, + } +} + +fn chat_response() -> Json { + json!({ + "id": "chat-1", "object": "chat.completion", "model": "selected", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }) +} + +fn chat_chunk(text: &str, finish_reason: Json) -> Json { + json!({ + "id": "chat-1", "object": "chat.completion.chunk", "model": "selected", + "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": finish_reason}] + }) +} + +#[test] +fn wire_protocol_and_plugin_lifecycle_contracts_are_stable() { + for (protocol, label, endpoint) in [ + ( + WireProtocol::OpenaiChat, + "openai_chat", + "/v1/chat/completions", + ), + ( + WireProtocol::OpenaiResponses, + "openai_responses", + "/v1/responses", + ), + ( + WireProtocol::AnthropicMessages, + "anthropic_messages", + "/v1/messages", + ), + ] { + assert_eq!(protocol.label(), label); + assert_eq!(protocol.endpoint(), endpoint); + assert_eq!( + WireProtocol::from_call(label, &request(protocol)), + Some(protocol) + ); + assert_eq!( + WireProtocol::from_call("unknown", &request(protocol)), + Some(protocol) + ); + } + assert_eq!( + WireProtocol::from_call( + "unknown", + &LlmRequest { + headers: Map::new(), + content: json!({}) + } + ), + None + ); + assert_eq!(RoutingMode::Enforce.label(), "enforce"); + assert_eq!(RoutingMode::ObserveOnly.label(), "observe_only"); + + let valid = config("http://127.0.0.1:1/v1/routing/decision".into()); + let component: PluginComponentSpec = valid.clone().into(); + assert_eq!(component.kind, SWITCHYARD_PLUGIN_KIND); + assert!(component.enabled); + assert_eq!(component.config["decision_profile_id"], "stage_router"); + + let plugin = SwitchyardPlugin; + assert_eq!(plugin.plugin_kind(), SWITCHYARD_PLUGIN_KIND); + assert!(!plugin.allows_multiple_components()); + assert!(plugin.validate(&component.config).is_empty()); + let diagnostics = plugin.validate(json!({"version": "invalid"}).as_object().unwrap()); + assert_eq!(diagnostics.len(), 1); + assert_eq!(diagnostics[0].level, DiagnosticLevel::Error); + assert_eq!(diagnostics[0].code, "switchyard.invalid_config"); +} + +#[tokio::test] +async fn plugin_registration_installs_and_rolls_back_both_execution_intercepts() { + let plugin = SwitchyardPlugin; + let component: PluginComponentSpec = + config("http://127.0.0.1:1/v1/routing/decision".into()).into(); + let mut context = + PluginRegistrationContext::with_namespace(format!("switchyard-test-{}-", Uuid::now_v7())); + + plugin + .register(&component.config, &mut context) + .await + .unwrap(); + + let mut registrations = context.into_registrations(); + assert_eq!(registrations.len(), 2); + rollback_registrations(&mut registrations); + assert!(registrations.is_empty()); +} + +#[test] +fn configuration_validation_rejects_unsafe_or_ambiguous_bindings() { + let url = "http://127.0.0.1:1/v1/routing/decision".to_string(); + let assert_invalid = |config: SwitchyardConfig, expected: &str| { + let error = SwitchyardRuntime::new(config) + .err() + .expect("config must fail"); + assert!( + error.contains(expected), + "{error:?} did not contain {expected:?}" + ); + }; + + let mut candidate = config(url.clone()); + candidate.version = 2; + assert_invalid(candidate, "unsupported Switchyard config version"); + let mut candidate = config(url.clone()); + candidate.decision_profile_id.clear(); + assert_invalid(candidate, "decision_profile_id"); + let mut candidate = config(url.clone()); + candidate.decision_timeout_millis = 0; + assert_invalid(candidate, "decision_timeout_millis"); + let mut candidate = config(url.clone()); + candidate.max_retries = 11; + assert_invalid(candidate, "max_retries"); + let mut candidate = config(url.clone()); + candidate.recent_message_count = 0; + assert_invalid(candidate, "recent_message_count"); + let candidate = config("file:///tmp/decision".into()); + assert_invalid(candidate, "must use http or https"); + let mut candidate = config(url.clone()); + candidate.targets.clear(); + assert_invalid(candidate, "targets must not be empty"); + let mut candidate = config(url.clone()); + candidate.enabled_inbound_profiles.clear(); + assert_invalid(candidate, "enabled_inbound_profiles"); + + let mut candidate = config(url.clone()); + candidate.targets.get_mut("selected-chat").unwrap().endpoint = "/wrong".into(); + assert_invalid(candidate, "endpoint must be"); + let mut candidate = config(url.clone()); + candidate.targets.get_mut("selected-chat").unwrap().base_url = "file:///tmp/provider".into(); + assert_invalid(candidate, "base_url must use http or https"); + let mut candidate = config(url.clone()); + candidate.targets.insert( + "duplicate-chat".into(), + candidate.targets["selected-chat"].clone(), + ); + assert_invalid(candidate, "conflicts with another exact backend binding"); + let mut candidate = config(url.clone()); + candidate.default_targets.openai_chat = "missing".into(); + assert_invalid(candidate, "default target \"missing\" is not configured"); + let mut candidate = config(url.clone()); + candidate.default_targets.openai_chat = "fallback-responses".into(); + assert_invalid(candidate, "must use protocol openai_chat"); + + let mut candidate = config(url.clone()); + candidate + .decision_headers + .insert("bad header".into(), "value".into()); + assert_invalid(candidate, "invalid header name"); + let mut candidate = config(url.clone()); + candidate + .decision_headers + .insert("x-test".into(), "bad\nvalue".into()); + assert_invalid(candidate, "invalid header value"); + let mut candidate = config(url.clone()); + candidate.decision_header_env.insert( + "authorization".into(), + "SWITCHYARD_TEST_MISSING_DECISION_SECRET".into(), + ); + assert_invalid(candidate, "is not set"); + let mut candidate = config(url); + candidate + .targets + .get_mut("selected-chat") + .unwrap() + .header_env + .insert( + "authorization".into(), + "SWITCHYARD_TEST_MISSING_TARGET_SECRET".into(), + ); + assert_invalid(candidate, "is not set"); +} + +#[test] +fn disabled_protocol_defaults_are_optional() { + let defaults: ProtocolDefaults = serde_json::from_value(json!({ + "anthropic_messages": "fallback-anthropic" + })) + .expect("omitted disabled protocol defaults should deserialize"); + assert!(defaults.openai_chat.is_empty()); + assert!(defaults.openai_responses.is_empty()); + + let mut candidate = config("http://127.0.0.1:1/v1/routing/decision".into()); + candidate.enabled_inbound_profiles = BTreeSet::from([WireProtocol::AnthropicMessages]); + candidate.default_targets.openai_chat.clear(); + candidate.default_targets.openai_responses.clear(); + + SwitchyardRuntime::new(candidate).expect("Anthropic-only configuration should validate"); +} + +#[test] +fn atof_cross_component_validation_reports_each_activation_mismatch() { + assert!(validate_switchyard_atof_configuration(&PluginConfig::default()).is_ok()); + + let payload = PluginConfig { + components: vec![config("http://switchyard.test/v1/routing/decision".into()).into()], + ..PluginConfig::default() + }; + assert!(validate_switchyard_atof_configuration(&payload).is_ok()); + + let mut switchyard = config("http://switchyard.test/v1/routing/decision".into()); + switchyard.context_mode = ContextMode::AtofRequired; + switchyard.atof_endpoint_url = Some("http://events.test/v1/atof/events".into()); + let mut plugin_config = PluginConfig { + components: vec![switchyard.into()], + ..PluginConfig::default() + }; + plugin_config.components.push(PluginComponentSpec { + kind: "observability".into(), + enabled: true, + config: json!({"atof": {"enabled": true, "endpoints": [{ + "url": "http://wrong.test/v1/atof/events", + "header_env": {"authorization": "TOKEN"} + }]}}) + .as_object() + .unwrap() + .clone(), + }); + assert!( + validate_switchyard_atof_configuration(&plugin_config) + .unwrap_err() + .contains("requires HTTP ATOF endpoint") + ); + + plugin_config.components[1].config["atof"]["endpoints"][0]["url"] = + json!("http://events.test/v1/atof/events"); + plugin_config.components[1].config["atof"]["endpoints"][0]["field_name_policy"] = + json!("snake_case"); + assert!( + validate_switchyard_atof_configuration(&plugin_config) + .unwrap_err() + .contains("field_name_policy = preserve") + ); + let endpoint = &mut plugin_config.components[1].config["atof"]["endpoints"][0]; + endpoint["field_name_policy"] = json!("preserve"); + endpoint.as_object_mut().unwrap().remove("header_env"); + assert!( + validate_switchyard_atof_configuration(&plugin_config) + .unwrap_err() + .contains("environment-referenced header") + ); +} + +#[tokio::test] +async fn execution_bypasses_inapplicable_calls_and_fails_open_on_extensions() { + let runtime = + SwitchyardRuntime::new(config("http://127.0.0.1:1/v1/routing/decision".into())).unwrap(); + let passthrough: LlmExecutionNextFn = + Arc::new(|request| Box::pin(async move { Ok(request.content) })); + let unknown = LlmRequest { + headers: Map::new(), + content: json!({"opaque": true}), + }; + assert_eq!( + runtime + .execute_buffered("custom.provider", unknown, Arc::clone(&passthrough)) + .await + .unwrap()["opaque"], + true + ); + + let mut disabled_config = config("http://127.0.0.1:1/v1/routing/decision".into()); + disabled_config + .enabled_inbound_profiles + .remove(&WireProtocol::OpenaiChat); + let disabled = SwitchyardRuntime::new(disabled_config).unwrap(); + assert_eq!( + disabled + .execute_buffered( + "openai.chat_completions", + chat_request(), + Arc::clone(&passthrough), + ) + .await + .unwrap()["model"], + "inbound" + ); + + let mut unsupported = chat_request(); + unsupported.content["thinking"] = json!({"type": "enabled"}); + let response = runtime + .execute_buffered( + "openai.chat_completions", + unsupported, + Arc::new(|request| { + Box::pin(async move { + assert_eq!(request.content["model"], "fallback"); + Ok(chat_response()) + }) + }), + ) + .await + .unwrap(); + assert_eq!(response["choices"][0]["message"]["content"], "ok"); +} + +#[tokio::test] +async fn stream_setup_retries_and_empty_streams_have_one_bounded_fallback() { + let (url, decisions) = decision_server().await; + let runtime = SwitchyardRuntime::new(config(url)).unwrap(); + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let next: LlmStreamExecutionNextFn = Arc::new(move |_| { + let seen = Arc::clone(&seen); + Box::pin(async move { + if seen.fetch_add(1, Ordering::SeqCst) == 0 { + return Err(FlowError::Upstream(UpstreamFailure { + status: None, + body: "connect".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::Connection, + })); + } + Ok(Box::pin(futures_stream::iter(vec![Ok(chat_chunk( + "ok", + json!("stop"), + ))])) as LlmJsonStream) + }) + }); + let output = runtime + .execute_stream("openai.chat_completions", chat_request(), next) + .await + .unwrap() + .collect::>() + .await; + assert_eq!(output.len(), 1); + assert_eq!(dispatches.load(Ordering::SeqCst), 2); + assert_eq!(decisions.lock().unwrap().len(), 2); + + let (url, decisions) = decision_server().await; + let runtime = SwitchyardRuntime::new(config(url)).unwrap(); + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let next: LlmStreamExecutionNextFn = Arc::new(move |request| { + let seen = Arc::clone(&seen); + Box::pin(async move { + let attempt = seen.fetch_add(1, Ordering::SeqCst); + let items = if attempt < 4 { + Vec::new() + } else { + assert_eq!(request.content["model"], "fallback"); + vec![Ok(chat_chunk("fallback", json!("stop")))] + }; + Ok(Box::pin(futures_stream::iter(items)) as LlmJsonStream) + }) + }); + let output = runtime + .execute_stream("openai.chat_completions", chat_request(), next) + .await + .unwrap() + .collect::>() + .await; + assert_eq!(output.len(), 1); + assert_eq!(dispatches.load(Ordering::SeqCst), 5); + assert_eq!(decisions.lock().unwrap().len(), 4); +} + +#[tokio::test] +async fn fallback_setup_failures_preserve_the_provider_error() { + let runtime = + SwitchyardRuntime::new(config("http://127.0.0.1:1/v1/routing/decision".into())).unwrap(); + let buffered: LlmExecutionNextFn = Arc::new(|_| { + Box::pin(async { Err(FlowError::Internal("buffered fallback failed".into())) }) + }); + let error = runtime + .dispatch_fallback_buffered(WireProtocol::OpenaiChat, chat_request(), buffered, "test") + .await + .unwrap_err(); + assert!(error.to_string().contains("buffered fallback failed")); + + let streaming: LlmStreamExecutionNextFn = + Arc::new(|_| Box::pin(async { Err(FlowError::Internal("stream fallback failed".into())) })); + let error = runtime + .dispatch_fallback_stream(WireProtocol::OpenaiChat, chat_request(), streaming, "test") + .await + .err() + .expect("stream fallback setup must fail"); + assert!(error.to_string().contains("stream fallback failed")); +} + +#[tokio::test] +async fn translated_stream_preserves_success_and_propagates_both_error_sources() { + let source = Box::pin(futures_stream::iter(vec![ + Ok(chat_chunk("hello", Json::Null)), + Ok(chat_chunk("", json!("stop"))), + ])) as LlmJsonStream; + let output = translated_stream( + WireProtocol::OpenaiChat, + WireProtocol::AnthropicMessages, + "selected".into(), + source, + ) + .collect::>() + .await; + assert!(output.iter().all(Result::is_ok)); + assert!(output.iter().any(|item| { + item.as_ref() + .is_ok_and(|chunk| chunk.to_string().contains("hello")) + })); + + let upstream_error = FlowError::Internal("upstream stream failed".into()); + let source = Box::pin(futures_stream::iter(vec![Err(upstream_error)])) as LlmJsonStream; + let output = translated_stream( + WireProtocol::OpenaiChat, + WireProtocol::AnthropicMessages, + "selected".into(), + source, + ) + .collect::>() + .await; + assert!(output[0].is_err()); + + let malformed = Box::pin(futures_stream::iter(vec![Ok(json!({ + "choices": [{"delta": {"reasoning_content": "private"}}] + }))])) as LlmJsonStream; + let output = translated_stream( + WireProtocol::OpenaiChat, + WireProtocol::AnthropicMessages, + "selected".into(), + malformed, + ) + .collect::>() + .await; + assert!(output[0].is_err()); +} + +#[test] +fn provider_failure_reporting_covers_every_retry_class() { + for (class, label, retryable) in [ + (UpstreamFailureClass::Connection, "connection", true), + (UpstreamFailureClass::Timeout, "timeout", true), + ( + UpstreamFailureClass::RetryableStatus, + "retryable_status", + true, + ), + (UpstreamFailureClass::ContextWindow, "context_window", true), + ( + UpstreamFailureClass::ModelUnavailable, + "model_unavailable", + true, + ), + ( + UpstreamFailureClass::Authentication, + "authentication", + false, + ), + ( + UpstreamFailureClass::InvalidRequest, + "invalid_request", + false, + ), + (UpstreamFailureClass::Other, "other", false), + ] { + let error = FlowError::Upstream(UpstreamFailure { + status: Some(503), + body: "failure".into(), + headers: BTreeMap::new(), + class, + }); + assert_eq!(provider_error_class(&error), label); + assert_eq!(error_is_retryable(&error), retryable); + assert_eq!(provider_error_summary(&error), format!("{label}:http_503")); + } + let relay = FlowError::Internal("failure".into()); + assert_eq!(provider_error_class(&relay), "relay"); + assert_eq!(provider_error_summary(&relay), relay.to_string()); +} + +#[test] +fn all_materialization_modes_are_bounded_and_provider_valid() { + for protocol in [ + WireProtocol::OpenaiChat, + WireProtocol::OpenaiResponses, + WireProtocol::AnthropicMessages, + ] { + for mode in [ + RequestMaterialization::None, + RequestMaterialization::SummaryOnly, + RequestMaterialization::LatestUserPrompt, + RequestMaterialization::RecentMessageWindow, + RequestMaterialization::AnnotatedRequest, + RequestMaterialization::FullBody, + ] { + let mut config = config("http://127.0.0.1:1/v1/routing/decision".into()); + config.request_materialization = mode; + config.recent_message_count = 2; + let runtime = SwitchyardRuntime::new(config).unwrap(); + let routing = runtime + .routing_request(protocol, &request(protocol), 1, None) + .unwrap(); + match mode { + RequestMaterialization::None | RequestMaterialization::SummaryOnly => { + assert!(routing.current_request.is_none()) + } + RequestMaterialization::LatestUserPrompt => { + let current = routing.current_request.unwrap(); + assert_eq!(current["latest_user_prompt"], "latest"); + let body = current["body"].clone(); + decode_request( + &runtime.translation, + protocol, + &LlmRequest { + headers: Map::new(), + content: body, + }, + ) + .unwrap(); + } + RequestMaterialization::RecentMessageWindow => { + let current = routing.current_request.unwrap(); + decode_request( + &runtime.translation, + protocol, + &LlmRequest { + headers: Map::new(), + content: current["body"].clone(), + }, + ) + .unwrap(); + } + RequestMaterialization::AnnotatedRequest | RequestMaterialization::FullBody => { + assert!(routing.current_request.is_some()) + } + } + } + } +} + +#[test] +fn identity_policy_requires_stable_request_scope_only_for_atof_profiles() { + let mut config = config("http://127.0.0.1:1/v1/routing/decision".into()); + let payload_runtime = SwitchyardRuntime::new(config.clone()).unwrap(); + let synthetic = payload_runtime + .routing_request(WireProtocol::OpenaiChat, &chat_request(), 1, None) + .unwrap(); + assert_eq!(synthetic.identity.quality, "synthetic"); + + config.context_mode = ContextMode::AtofRequired; + let atof_runtime = SwitchyardRuntime::new(config).unwrap(); + assert!( + atof_runtime + .routing_request(WireProtocol::OpenaiChat, &chat_request(), 1, None) + .is_err() + ); + let mut stable = chat_request(); + stable + .headers + .insert("x-nemo-relay-session-id".into(), json!("session-1")); + stable + .headers + .insert("x-nemo-relay-request-id".into(), json!("request-1")); + let routed = atof_runtime + .routing_request(WireProtocol::OpenaiChat, &stable, 1, None) + .unwrap(); + assert_eq!(routed.identity.quality, "explicit"); +} + +#[test] +fn exact_target_validation_rejects_any_switchyard_drift() { + let runtime = + SwitchyardRuntime::new(config("http://127.0.0.1:1/v1/routing/decision".into())).unwrap(); + assert!(runtime.validate_decision(&decision()).is_ok()); + let mut drifted = decision(); + drifted.route.target_model = "unbound-model".into(); + assert!(runtime.validate_decision(&drifted).is_err()); + drifted = decision(); + drifted.route.backend_id = "unknown".into(); + assert!(runtime.validate_decision(&drifted).is_err()); +} + +#[test] +fn routing_contribution_requires_an_exact_independent_baseline_binding() { + let runtime = + SwitchyardRuntime::new(config("http://127.0.0.1:1/v1/routing/decision".into())).unwrap(); + + let contribution = runtime.routing_contribution(&decision(), 2, true).unwrap(); + assert!(contribution.applied); + assert_eq!( + contribution.kind.as_str(), + LlmOptimizationKind::MODEL_ROUTING + ); + let transition = contribution.model_transition.unwrap(); + assert_eq!(transition.baseline.unwrap().model, "baseline"); + assert_eq!(transition.effective.unwrap().model, "selected"); + assert_eq!(contribution.payload.as_ref().unwrap()["routing_attempt"], 2); + assert_eq!( + contribution.payload.as_ref().unwrap()["router_metadata"]["feature_state"], + "fresh" + ); + assert_eq!( + contribution.payload.as_ref().unwrap()["router_metadata"]["snapshot_age_millis"], + 37 + ); + assert_eq!( + contribution.payload_schema.as_ref().unwrap().name, + ROUTING_CONTRIBUTION_SCHEMA + ); + + let observed = runtime.routing_contribution(&decision(), 1, false).unwrap(); + assert!(!observed.applied); + + let mut missing = decision(); + missing.baseline_route = None; + assert!(runtime.routing_contribution(&missing, 1, true).is_none()); + + let mut drifted = decision(); + drifted.baseline_route.as_mut().unwrap().target_model = "drifted".into(); + assert!(runtime.routing_contribution(&drifted, 1, true).is_none()); + assert!(runtime.validate_decision(&drifted).is_ok()); +} + +#[test] +fn routing_decision_mark_has_canonical_shape_and_mirrored_identity() { + let subscriber_name = format!("switchyard-mark-shape-{}", uuid::Uuid::now_v7()); + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + &subscriber_name, + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + + let runtime = + SwitchyardRuntime::new(config("http://127.0.0.1:1/v1/routing/decision".into())).unwrap(); + let routing_request = runtime + .routing_request(WireProtocol::OpenaiChat, &chat_request(), 1, None) + .unwrap(); + runtime.emit_decision(&routing_request, &decision(), 1, false, 17); + flush_subscribers().unwrap(); + deregister_subscriber(&subscriber_name).unwrap(); + + let event = events + .lock() + .unwrap() + .iter() + .map(Event::to_json_value) + .find(|event| { + event["name"] == "switchyard.routing.decision" + && event["metadata"]["session_id"] == routing_request.identity.session_id + && event["metadata"]["request_id"] == routing_request.identity.request_id + }) + .expect("decision mark should be captured"); + assert_eq!(event["kind"], "mark"); + assert_eq!(event["category"], "custom"); + assert_eq!( + event["category_profile"]["subtype"], + "switchyard.routing.decision" + ); + assert_eq!(event["data_schema"]["name"], ROUTING_MARK_SCHEMA); + assert_eq!(event["data_schema"]["version"], "1"); + assert_eq!(event["data"]["profile_id"], "stage_router"); + assert_eq!(event["data"]["selected_model"], "selected"); + assert_eq!(event["data"]["latency_ms"], 17); + assert_eq!(event["data"]["router_metadata"]["feature_state"], "fresh"); + assert_eq!(event["data"]["router_metadata"]["snapshot_age_millis"], 37); + assert_eq!( + event["metadata"]["session_id"], + routing_request.identity.session_id + ); + assert_eq!( + event["metadata"]["request_id"], + routing_request.identity.request_id + ); +} + +#[test] +fn atof_required_cross_component_validation_is_context_sensitive() { + let mut switchyard = config("http://switchyard.test:8080/v1/routing/decision".into()); + switchyard.context_mode = ContextMode::AtofRequired; + let mut plugin_config = PluginConfig { + components: vec![switchyard.into()], + ..PluginConfig::default() + }; + assert!(validate_switchyard_atof_configuration(&plugin_config).is_err()); + plugin_config.components.push(PluginComponentSpec { + kind: "observability".into(), + enabled: true, + config: json!({"atof": { + "enabled": true, + "endpoints": [{ + "url": "http://switchyard.test:8080/v1/atof/events", + "transport": "http_post", + "field_name_policy": "preserve", + "header_env": {"authorization": "SWITCHYARD_TOKEN"} + }] + }}) + .as_object() + .unwrap() + .clone(), + }); + assert!(validate_switchyard_atof_configuration(&plugin_config).is_ok()); +} + +#[derive(Clone)] +struct DecisionState { + requests: Arc>>, + decision: RoutingDecision, +} + +async fn decision_handler( + State(state): State, + AxumJson(request): AxumJson, +) -> AxumJson { + state.requests.lock().unwrap().push(request); + AxumJson(state.decision) +} + +async fn decision_server() -> (String, Arc>>) { + decision_server_for(decision()).await +} + +async fn decision_server_for( + decision: RoutingDecision, +) -> (String, Arc>>) { + let requests = Arc::new(Mutex::new(Vec::new())); + let state = DecisionState { + requests: Arc::clone(&requests), + decision, + }; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve( + listener, + Router::new() + .route("/v1/routing/decision", post(decision_handler)) + .with_state(state), + ) + .await + .unwrap(); + }); + (format!("http://{address}/v1/routing/decision"), requests) +} + +async fn managed_buffered_events( + runtime: SwitchyardRuntime, + next: LlmExecutionNextFn, +) -> Vec { + let subscriber_name = format!("switchyard-accounting-{}", uuid::Uuid::now_v7()); + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + &subscriber_name, + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + let runtime = Arc::new(runtime); + let func: LlmExecutionNextFn = Arc::new(move |request| { + let runtime = Arc::clone(&runtime); + let next = Arc::clone(&next); + Box::pin(async move { + runtime + .execute_buffered("openai.chat_completions", request, next) + .await + }) + }); + llm_call_execute( + LlmCallExecuteParams::builder() + .name("openai.chat_completions") + .request(chat_request()) + .func(func) + .build(), + ) + .await + .unwrap(); + flush_subscribers().unwrap(); + deregister_subscriber(&subscriber_name).unwrap(); + Arc::try_unwrap(events).unwrap().into_inner().unwrap() +} + +async fn managed_stream_events( + runtime: SwitchyardRuntime, + next: LlmStreamExecutionNextFn, +) -> Vec { + let subscriber_name = format!("switchyard-stream-accounting-{}", uuid::Uuid::now_v7()); + let events = Arc::new(Mutex::new(Vec::::new())); + let captured = Arc::clone(&events); + register_subscriber( + &subscriber_name, + Arc::new(move |event| captured.lock().unwrap().push(event.clone())), + ) + .unwrap(); + let runtime = Arc::new(runtime); + let func: LlmStreamExecutionNextFn = Arc::new(move |request| { + let runtime = Arc::clone(&runtime); + let next = Arc::clone(&next); + Box::pin(async move { + runtime + .execute_stream("openai.chat_completions", request, next) + .await + }) + }); + let mut stream = llm_stream_call_execute( + LlmStreamCallExecuteParams::builder() + .name("openai.chat_completions") + .request(chat_request()) + .func(func) + .collector(Box::new(|_| Ok(()))) + .finalizer(Box::new(|| json!({"done": true}))) + .build(), + ) + .await + .unwrap(); + while stream.next().await.is_some() {} + drop(stream); + flush_subscribers().unwrap(); + deregister_subscriber(&subscriber_name).unwrap(); + Arc::try_unwrap(events).unwrap().into_inner().unwrap() +} + +#[tokio::test] +async fn buffered_accounting_records_only_the_terminal_committed_route() { + let (url, _) = decision_server().await; + let successful = managed_buffered_events( + SwitchyardRuntime::new(config(url)).unwrap(), + Arc::new(|_| Box::pin(async { Ok(chat_response()) })), + ) + .await; + let marks = successful + .iter() + .filter(|event| event.name() == "nemo_relay.llm.optimization") + .collect::>(); + assert_eq!(marks.len(), 1); + assert_eq!(marks[0].data().unwrap()["applied"], true); + assert_eq!( + marks[0].data().unwrap()["model_transition"]["baseline"]["model"], + "baseline" + ); + let summary = successful + .iter() + .find_map(|event| { + event + .annotated_response() + .and_then(|response| response.optimization_summary.as_ref()) + }) + .unwrap(); + assert_eq!(summary.contributions.len(), 1); + assert!(summary.contributions[0].applied); + + let (url, _) = decision_server().await; + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let fallback = managed_buffered_events( + SwitchyardRuntime::new(config(url)).unwrap(), + Arc::new(move |_| { + let seen = Arc::clone(&seen); + Box::pin(async move { + if seen.fetch_add(1, Ordering::SeqCst) == 0 { + Err(FlowError::Upstream(UpstreamFailure { + status: Some(401), + body: "unauthorized".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::Authentication, + })) + } else { + Ok(chat_response()) + } + }) + }), + ) + .await; + assert_eq!(dispatches.load(Ordering::SeqCst), 2); + assert!( + fallback + .iter() + .all(|event| event.name() != "nemo_relay.llm.optimization") + ); +} + +#[tokio::test] +async fn retry_then_success_records_one_terminal_route_with_matching_mark_identity() { + let (url, decisions) = decision_server().await; + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let events = managed_buffered_events( + SwitchyardRuntime::new(config(url)).unwrap(), + Arc::new(move |_| { + let seen = Arc::clone(&seen); + Box::pin(async move { + if seen.fetch_add(1, Ordering::SeqCst) == 0 { + Err(FlowError::Upstream(UpstreamFailure { + status: Some(503), + body: "retry once".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::RetryableStatus, + })) + } else { + Ok(chat_response()) + } + }) + }), + ) + .await; + + assert_eq!(dispatches.load(Ordering::SeqCst), 2); + let requests = decisions.lock().unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[1].attempt.routing_attempt, 2); + drop(requests); + + let start = events + .iter() + .find(|event| { + event.name() == "openai.chat_completions" + && event.scope_category() == Some(ScopeCategory::Start) + }) + .unwrap(); + let marks = events + .iter() + .filter(|event| event.name() == "nemo_relay.llm.optimization") + .collect::>(); + assert_eq!(marks.len(), 1); + assert_eq!(marks[0].parent_uuid(), Some(start.uuid())); + assert_eq!(marks[0].data().unwrap()["payload"]["routing_attempt"], 2); + + let summary = events + .iter() + .find_map(|event| { + event + .annotated_response() + .and_then(|response| response.optimization_summary.as_ref()) + }) + .unwrap(); + assert_eq!(summary.contributions.len(), 1); + let contribution = &summary.contributions[0]; + assert!(contribution.applied); + assert_eq!(contribution.payload.as_ref().unwrap()["routing_attempt"], 2); + assert_eq!(marks[0].data().unwrap()["id"], json!(contribution.id)); + assert_eq!( + marks[0].data().unwrap()["sequence"], + json!(contribution.sequence) + ); +} + +#[tokio::test] +async fn oversized_decision_metadata_limits_accounting_without_failing_provider_success() { + let mut oversized = decision(); + oversized + .metadata + .insert("oversized_evidence".into(), json!("x".repeat(20_000))); + let (url, _) = decision_server_for(oversized).await; + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let events = managed_buffered_events( + SwitchyardRuntime::new(config(url)).unwrap(), + Arc::new(move |_| { + let seen = Arc::clone(&seen); + Box::pin(async move { + seen.fetch_add(1, Ordering::SeqCst); + Ok(chat_response()) + }) + }), + ) + .await; + + assert_eq!(dispatches.load(Ordering::SeqCst), 1); + assert!( + events + .iter() + .all(|event| event.name() != "nemo_relay.llm.optimization") + ); + let summary = events + .iter() + .find_map(|event| { + event + .annotated_response() + .and_then(|response| response.optimization_summary.as_ref()) + }) + .unwrap(); + assert_eq!(summary.status, LlmOptimizationSummaryStatus::Partial); + assert!(summary.contributions.is_empty()); + assert!( + summary + .limitations + .iter() + .any(|limitation| limitation == "contribution_limit_exceeded") + ); +} + +#[tokio::test] +async fn malformed_baseline_is_telemetry_only_and_does_not_change_the_selected_route() { + let mut drifted = decision(); + drifted.baseline_route.as_mut().unwrap().target_model = "drifted".into(); + let (url, _) = decision_server_for(drifted).await; + let events = managed_buffered_events( + SwitchyardRuntime::new(config(url)).unwrap(), + Arc::new(|request| { + Box::pin(async move { + assert_eq!(request.content["model"], "selected"); + Ok(chat_response()) + }) + }), + ) + .await; + assert!(events.iter().any(|event| { + event.name() == "switchyard.routing.error" + && event.data().is_some_and(|data| { + data["error_class"] == "baseline_binding" + && data["error"] + .as_str() + .is_some_and(|error| error.contains("exact Relay binding")) + }) + })); + assert!( + events + .iter() + .all(|event| event.name() != "nemo_relay.llm.optimization") + ); +} + +#[tokio::test] +async fn observe_only_accounting_is_visible_but_not_applied() { + let (url, _) = decision_server().await; + let mut observe = config(url); + observe.mode = RoutingMode::ObserveOnly; + let events = managed_buffered_events( + SwitchyardRuntime::new(observe).unwrap(), + Arc::new(|request| { + Box::pin(async move { + assert_eq!(request.content["model"], "fallback"); + Ok(chat_response()) + }) + }), + ) + .await; + let contribution = events + .iter() + .find(|event| event.name() == "nemo_relay.llm.optimization") + .and_then(Event::data) + .unwrap(); + assert_eq!(contribution["applied"], false); + let summary = events + .iter() + .find_map(|event| { + event + .annotated_response() + .and_then(|response| response.optimization_summary.as_ref()) + }) + .unwrap(); + assert_eq!(summary.contributions.len(), 1); + assert!(!summary.contributions[0].applied); + assert!(summary.tokens_saved.total_tokens.is_none()); +} + +#[tokio::test] +async fn streaming_accounting_commits_on_the_first_successful_item_only() { + let (url, _) = decision_server().await; + let committed = managed_stream_events( + SwitchyardRuntime::new(config(url)).unwrap(), + Arc::new(|_| { + Box::pin(async { + Ok(Box::pin(futures_stream::iter(vec![Ok(chat_chunk( + "ok", + json!("stop"), + ))])) as LlmJsonStream) + }) + }), + ) + .await; + assert_eq!( + committed + .iter() + .filter(|event| event.name() == "nemo_relay.llm.optimization") + .count(), + 1 + ); + + let (url, _) = decision_server().await; + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let fallback = managed_stream_events( + SwitchyardRuntime::new(config(url)).unwrap(), + Arc::new(move |_| { + let seen = Arc::clone(&seen); + Box::pin(async move { + let items = if seen.fetch_add(1, Ordering::SeqCst) == 0 { + vec![Err(FlowError::Upstream(UpstreamFailure { + status: Some(401), + body: "unauthorized".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::Authentication, + }))] + } else { + vec![Ok(chat_chunk("fallback", json!("stop")))] + }; + Ok(Box::pin(futures_stream::iter(items)) as LlmJsonStream) + }) + }), + ) + .await; + assert_eq!(dispatches.load(Ordering::SeqCst), 2); + assert!( + fallback + .iter() + .all(|event| event.name() != "nemo_relay.llm.optimization") + ); +} + +#[tokio::test] +async fn committed_stream_error_keeps_one_route_and_never_redispatches() { + let (url, decisions) = decision_server().await; + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let events = managed_stream_events( + SwitchyardRuntime::new(config(url)).unwrap(), + Arc::new(move |_| { + let seen = Arc::clone(&seen); + Box::pin(async move { + seen.fetch_add(1, Ordering::SeqCst); + Ok(Box::pin(futures_stream::iter(vec![ + Ok(chat_chunk("partial", Json::Null)), + Err(FlowError::Upstream(UpstreamFailure { + status: None, + body: "connection closed after commit".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::Connection, + })), + ])) as LlmJsonStream) + }) + }), + ) + .await; + + assert_eq!(dispatches.load(Ordering::SeqCst), 1); + assert_eq!(decisions.lock().unwrap().len(), 1); + assert_eq!( + events + .iter() + .filter(|event| event.name() == "nemo_relay.llm.optimization") + .count(), + 1 + ); + let summary = events + .iter() + .find_map(|event| { + event + .annotated_response() + .and_then(|response| response.optimization_summary.as_ref()) + }) + .unwrap(); + assert_eq!(summary.contributions.len(), 1); + assert!(summary.contributions[0].applied); + assert!( + summary + .limitations + .iter() + .any(|limitation| limitation == "stream_interrupted") + ); +} + +#[tokio::test] +async fn retry_exhaustion_redecides_four_times_then_dispatches_fallback_once() { + let (url, decisions) = decision_server().await; + let runtime = SwitchyardRuntime::new(config(url)).unwrap(); + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let next: LlmExecutionNextFn = Arc::new(move |request| { + let seen = Arc::clone(&seen); + Box::pin(async move { + let attempt = seen.fetch_add(1, Ordering::SeqCst) + 1; + if attempt <= 4 { + return Err(FlowError::Upstream(UpstreamFailure { + status: Some(503), + body: "temporarily unavailable".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::RetryableStatus, + })); + } + assert_eq!(request.content["model"], "fallback"); + Ok(chat_response()) + }) + }); + let response = runtime + .execute_buffered("openai.chat_completions", chat_request(), next) + .await + .unwrap(); + assert_eq!(response["choices"][0]["message"]["content"], "ok"); + assert_eq!(dispatches.load(Ordering::SeqCst), 5); + let requests = decisions.lock().unwrap(); + assert_eq!(requests.len(), 4); + assert_eq!(requests[3].attempt.routing_attempt, 4); + assert_eq!( + requests[3].attempt.previous_route.as_deref(), + Some("selected-chat") + ); +} + +#[tokio::test] +async fn non_retryable_provider_failure_bypasses_retry_loop() { + let (url, decisions) = decision_server().await; + let runtime = SwitchyardRuntime::new(config(url)).unwrap(); + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let next: LlmExecutionNextFn = Arc::new(move |_| { + let seen = Arc::clone(&seen); + Box::pin(async move { + let attempt = seen.fetch_add(1, Ordering::SeqCst); + if attempt == 0 { + return Err(FlowError::Upstream(UpstreamFailure { + status: Some(401), + body: "unauthorized".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::Authentication, + })); + } + Ok(chat_response()) + }) + }); + runtime + .execute_buffered("openai.chat_completions", chat_request(), next) + .await + .unwrap(); + assert_eq!(dispatches.load(Ordering::SeqCst), 2); + assert_eq!(decisions.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn observe_only_records_one_decision_and_dispatches_only_the_trusted_default() { + let (url, decisions) = decision_server().await; + let mut config = config(url); + config.mode = RoutingMode::ObserveOnly; + let runtime = SwitchyardRuntime::new(config).unwrap(); + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let next: LlmExecutionNextFn = Arc::new(move |request| { + let seen = Arc::clone(&seen); + Box::pin(async move { + seen.fetch_add(1, Ordering::SeqCst); + assert_eq!(request.content["model"], "fallback"); + Ok(chat_response()) + }) + }); + runtime + .execute_buffered("openai.chat_completions", chat_request(), next) + .await + .unwrap(); + assert_eq!(dispatches.load(Ordering::SeqCst), 1); + assert_eq!(decisions.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn streaming_retries_before_first_item() { + let (url, decisions) = decision_server().await; + let runtime = SwitchyardRuntime::new(config(url)).unwrap(); + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let next: LlmStreamExecutionNextFn = Arc::new(move |_| { + let seen = Arc::clone(&seen); + Box::pin(async move { + let attempt = seen.fetch_add(1, Ordering::SeqCst); + let items = if attempt == 0 { + vec![Err(FlowError::Upstream(UpstreamFailure { + status: Some(503), + body: "retry".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::RetryableStatus, + }))] + } else { + vec![Ok(chat_chunk("ok", json!("stop")))] + }; + Ok(Box::pin(futures_stream::iter(items)) as LlmJsonStream) + }) + }); + let stream = runtime + .execute_stream("openai.chat_completions", chat_request(), next) + .await + .unwrap(); + let output = stream.collect::>().await; + assert_eq!(output.len(), 1); + assert!(output[0].is_ok()); + assert_eq!(dispatches.load(Ordering::SeqCst), 2); + assert_eq!(decisions.lock().unwrap().len(), 2); +} + +#[tokio::test] +async fn streaming_never_retries_after_first_item() { + let (url, decisions) = decision_server().await; + let runtime = SwitchyardRuntime::new(config(url)).unwrap(); + let dispatches = Arc::new(AtomicUsize::new(0)); + let seen = Arc::clone(&dispatches); + let next: LlmStreamExecutionNextFn = Arc::new(move |_| { + let seen = Arc::clone(&seen); + Box::pin(async move { + seen.fetch_add(1, Ordering::SeqCst); + let items = vec![ + Ok(chat_chunk("partial", Json::Null)), + Err(FlowError::Upstream(UpstreamFailure { + status: None, + body: "connection closed".into(), + headers: BTreeMap::new(), + class: UpstreamFailureClass::Connection, + })), + ]; + Ok(Box::pin(futures_stream::iter(items)) as LlmJsonStream) + }) + }); + let stream = runtime + .execute_stream("openai.chat_completions", chat_request(), next) + .await + .unwrap(); + let output = stream.collect::>().await; + assert_eq!(output.len(), 2); + assert!(output[0].is_ok()); + assert!(output[1].is_err()); + assert_eq!(dispatches.load(Ordering::SeqCst), 1); + assert_eq!(decisions.lock().unwrap().len(), 1); +} diff --git a/examples/switchyard/README.md b/examples/switchyard/README.md new file mode 100644 index 00000000..353e8dd1 --- /dev/null +++ b/examples/switchyard/README.md @@ -0,0 +1,93 @@ + + +# Switchyard integration examples + +These examples exercise the experimental Relay integration with a separately running Switchyard +Decision API service and the in-process Switchyard translation library. They are manual, local +validation workflows rather than production startup orchestration. + +## Required Switchyard revision + +The scripts default to the latest commit currently pinned for the public topic branch: + +```text +https://github.com/NVIDIA-NeMo/Switchyard/tree/topic/nemo-relay-integration +8f9db9a6a47f848cdff1d262276ba25a8ae9cbc8 +``` + +Create the adjacent checkout from the Relay repository root: + +```bash +git fetch upstream topic/nemo-relay-integration +git worktree add --detach ../Switchyard-topic-nemo-relay-integration \ + 8f9db9a6a47f848cdff1d262276ba25a8ae9cbc8 +``` + +Every real-service script verifies this commit before launching `switchyard-server`. To test a +deliberately different checkout, set both variables explicitly: + +```bash +SWITCHYARD_ROOT=/path/to/Switchyard \ +SWITCHYARD_EXPECTED_COMMIT= \ + examples/switchyard/run-real-e2e.sh +``` + +## Examples + +Run these commands from the root of the NeMo Relay checkout. + +The Relay CLI's Switchyard support is compile-time optional. The scripts enable the `switchyard` +feature automatically; custom builds must pass `--features switchyard`. + +### Manual Switchyard compatibility smoke test + +`run-real-e2e.sh` is a manual compatibility smoke test. It starts the pinned Switchyard server, +Relay, and a fake provider, then verifies cold and warm StageRouter decisions, buffered routing, +SSE routing, and the selected model sequence. It is intended to catch cross-repository service +contract regressions; the CI-safe Relay process regression test covers the faster local behavior +checks. The script requires Rust tooling, Python, and `curl`; its temporary logs are removed after +a successful run. + +```bash +examples/switchyard/run-real-e2e.sh +``` + +### Hermes and Ollama trajectory + +`run-hermes-ollama-smoke.sh` runs a fixed multi-query trajectory through Hermes, Relay, Ollama, +and Switchyard. It requires Docker, Hermes, and the configured local Ollama models. The script +produces ATOF, ATIF, and OTEL artifacts and can leave Phoenix running with +`SWITCHYARD_KEEP_PHOENIX=1`. + +```bash +examples/switchyard/run-hermes-ollama-smoke.sh +``` + +## Configuration files + +- `plugins.toml`: minimal plugin configuration example. +- `real-e2e-plugins.toml` and `real-e2e-profiles.yaml`: deterministic fake-provider E2E. +- `hermes-ollama-plugins.toml` and `hermes-ollama-profiles.yaml`: local Ollama trajectory. +- `fake_upstream.py`: deterministic provider used by the service E2E. +- `otel-collector.yaml`: local OTEL artifact export configuration. + +## Runtime model + +The scripts launch Switchyard as a separate local process on port `4000`. Relay sends routing +requests to `/v1/routing/decision` and, for ATOF-backed profiles, sends events to +`/v1/atof/events`. Relay owns provider credentials, target bindings, dispatch, retries, and +fallback behavior; Switchyard owns ATOF accumulation, routing decisions, and provider-protocol +translation. + +The service is not started automatically by Relay outside these examples. A production deployment +must provide a reachable Decision API and configure the Relay plugin with its URL. + +## Artifacts and troubleshooting + +Trajectory scripts write to `artifacts/` by default. Set `SWITCHYARD_TRAJECTORY_DIR` to choose a +shareable output directory. On failure, logs are preserved and include the verified Switchyard +revision. Do not place API keys or bearer tokens in configuration files; use environment variables +or an untracked secrets file. diff --git a/examples/switchyard/e2e-common.sh b/examples/switchyard/e2e-common.sh new file mode 100755 index 00000000..b5e999e0 --- /dev/null +++ b/examples/switchyard/e2e-common.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Shared process and readiness helpers for the manually-invoked Switchyard E2E scripts. + +e2e_pids=() + +e2e_add_pid() { + e2e_pids+=("$1") +} + +e2e_stop_tree() { + local pid="$1" + local child + for child in $(pgrep -P "$pid" 2>/dev/null || true); do + e2e_stop_tree "$child" + done + kill "$pid" 2>/dev/null || true +} + +e2e_stop_processes() { + local pid + for pid in "${e2e_pids[@]}"; do + e2e_stop_tree "$pid" + done + for pid in "${e2e_pids[@]}"; do + wait "$pid" 2>/dev/null || true + done + e2e_pids=() +} + +e2e_wait_for() { + local url="$1" + local attempts="${2:-120}" + local delay="${3:-0.25}" + local attempt + for attempt in $(seq 1 "$attempts"); do + if curl --fail --silent "$url" >/dev/null 2>&1; then + return 0 + fi + sleep "$delay" + done + echo "timed out waiting for $url" >&2 + return 1 +} + +e2e_random_token() { + python3 -c 'import secrets; print(secrets.token_hex(24))' +} + +# The experimental Relay integration is tested against a specific public +# Switchyard topic-branch commit. Override this only when intentionally testing +# another revision. +e2e_verify_switchyard_checkout() { + local root="$1" + local expected="${2:-8f9db9a6a47f848cdff1d262276ba25a8ae9cbc8}" + local actual + actual="$(git -C "$root" rev-parse HEAD 2>/dev/null)" || { + echo "Switchyard worktree is not a git checkout: $root" >&2 + return 1 + } + if [[ "$actual" != "$expected" ]]; then + echo "Switchyard checkout mismatch: expected $expected, found $actual" >&2 + echo "Set SWITCHYARD_EXPECTED_COMMIT only when intentionally testing another revision." >&2 + return 1 + fi + printf 'Switchyard revision: %s (%s)\n' "$actual" "$(git -C "$root" show -s --format=%s HEAD)" +} + +e2e_tail_logs() { + local directory="$1" + local log + for log in "$directory"/*.log; do + [[ -f "$log" ]] || continue + echo "--- $log" >&2 + tail -100 "$log" >&2 + done +} diff --git a/examples/switchyard/fake_upstream.py b/examples/switchyard/fake_upstream.py new file mode 100644 index 00000000..f27527d1 --- /dev/null +++ b/examples/switchyard/fake_upstream.py @@ -0,0 +1,85 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic local OpenAI-compatible upstream for Switchyard plugin smoke tests.""" + +import argparse +import json +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + + +class Handler(BaseHTTPRequestHandler): + log_path: Path + + def do_POST(self) -> None: # noqa: N802 + length = int(self.headers.get("content-length", "0")) + body = json.loads(self.rfile.read(length) or b"{}") + with self.log_path.open("a", encoding="utf-8") as output: + output.write(json.dumps({"path": self.path, "body": body}) + "\n") + model = body.get("model", "fake-model") + if body.get("stream"): + chunks = [ + { + "id": "chatcmpl-fake", + "object": "chat.completion.chunk", + "model": model, + "choices": [ + { + "index": 0, + "delta": {"role": "assistant", "content": "fake"}, + "finish_reason": None, + } + ], + }, + { + "id": "chatcmpl-fake", + "object": "chat.completion.chunk", + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + ] + payload = "".join(f"data: {json.dumps(chunk)}\n\n" for chunk in chunks) + payload += "data: [DONE]\n\n" + encoded = payload.encode() + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.send_header("content-length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + return + response = { + "id": "chatcmpl-fake", + "object": "chat.completion", + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "fake"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + encoded = json.dumps(response).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, _format: str, *_args: object) -> None: + return + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--port", type=int, default=4101) + parser.add_argument("--log", type=Path, required=True) + args = parser.parse_args() + Handler.log_path = args.log + ThreadingHTTPServer(("127.0.0.1", args.port), Handler).serve_forever() + + +if __name__ == "__main__": + main() diff --git a/examples/switchyard/hermes-ollama-plugins.toml b/examples/switchyard/hermes-ollama-plugins.toml new file mode 100644 index 00000000..6a7d128b --- /dev/null +++ b/examples/switchyard/hermes-ollama-plugins.toml @@ -0,0 +1,90 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +version = 1 + +[[components]] +kind = "switchyard" +enabled = true + +[components.config] +mode = "enforce" +decision_api_url = "http://127.0.0.1:4000/v1/routing/decision" +decision_profile_id = "hermes-ollama-stage-router" +request_materialization = "full_body" +context_mode = "atof_required" +decision_timeout_millis = 120000 +max_retries = 3 +recent_message_count = 8 + +[components.config.decision_header_env] +authorization = "SWITCHYARD_AUTHORIZATION" + +[components.config.default_targets] +openai_chat = "ollama-weak" +openai_responses = "ollama-responses" +anthropic_messages = "ollama-anthropic" + +[components.config.targets.ollama-weak] +model = "llama3.2:latest" +protocol = "openai_chat" +endpoint = "/v1/chat/completions" +base_url = "http://127.0.0.1:11434" + +[components.config.targets.ollama-strong] +model = "qwen3.6:35b" +protocol = "openai_chat" +endpoint = "/v1/chat/completions" +base_url = "http://127.0.0.1:11434" + +[components.config.targets.ollama-responses] +model = "llama3.2:latest" +protocol = "openai_responses" +endpoint = "/v1/responses" +base_url = "http://127.0.0.1:11434" + +[components.config.targets.ollama-anthropic] +model = "llama3.2:latest" +protocol = "anthropic_messages" +endpoint = "/v1/messages" +base_url = "http://127.0.0.1:11434" + +[[components]] +kind = "observability" +enabled = true + +[components.config.atof] +enabled = true +mode = "append" +output_directory = "." +filename = "trajectory.atof.jsonl" + +[[components.config.atof.endpoints]] +url = "http://127.0.0.1:4000/v1/atof/events" +transport = "http_post" +field_name_policy = "preserve" +timeout_millis = 1000 + +[components.config.atof.endpoints.header_env] +authorization = "SWITCHYARD_AUTHORIZATION" + +[components.config.atif] +enabled = true +agent_name = "Hermes" +agent_version = "switchyard-stage-router-smoke" +model_name = "switchyard-stage-router" +output_directory = "." +filename_template = "trajectory-{session_id}.atif.json" + +[components.config.opentelemetry] +enabled = true +transport = "http_binary" +endpoint = "http://127.0.0.1:4318/v1/traces" +service_name = "nemo-relay-switchyard-e2e" +service_namespace = "switchyard-integration" +instrumentation_scope = "nemo-relay-switchyard" +timeout_millis = 5000 + +[components.config.opentelemetry.resource_attributes] +"deployment.environment.name" = "local-e2e" +"switchyard.profile.id" = "hermes-ollama-stage-router" diff --git a/examples/switchyard/hermes-ollama-profiles.yaml b/examples/switchyard/hermes-ollama-profiles.yaml new file mode 100644 index 00000000..75120185 --- /dev/null +++ b/examples/switchyard/hermes-ollama-profiles.yaml @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +targets: + ollama-weak: + model: llama3.2:latest + format: openai + base_url: http://127.0.0.1:11434/v1 + ollama-strong: + model: qwen3.6:35b + format: openai + base_url: http://127.0.0.1:11434/v1 + +profiles: + hermes-ollama-stage-router: + type: stage_router + capable: ollama-strong + efficient: ollama-weak + fallback_target_on_evict: ollama-strong + picker: efficient_first + # Force the classifier path for ordinary prompts so this smoke demonstrates + # request-complexity routing rather than relying on incidental tool signals. + confidence_threshold: 1.0 + # Each fixed trajectory transition is driven by the latest canonical tool + # result so older critical signals do not leak into the simple follow-up. + signal_recent_window: 1 + classifier: + # Keep classification small and deterministic; qwen3.6:35b remains the + # capable backend selected for the complex trajectory turn. + model: llama3.2:latest + api_key: ${OLLAMA_CLASSIFIER_API_KEY} + base_url: http://127.0.0.1:11434/v1 + timeout_secs: 120.0 + recent_turn_window: 1 + max_tokens: 128 + system_prompt: >- + Select exactly one model tier for the latest user request. Return only + JSON in the form {"tier":"efficient"} or {"tier":"capable"}. Choose efficient + for arithmetic, extraction, formatting, concise factual responses, or + other direct low-risk work. Choose capable for architecture, concurrency, + debugging, multi-step reasoning, subtle correctness constraints, or + requests requiring a defensible technical analysis. In the supplied + State line, severity at least 1.0 always means capable; tests_passed=true + with severity 0.0 means efficient. Classify the latest state even when earlier + conversation turns were more complex. diff --git a/examples/switchyard/otel-collector.yaml b/examples/switchyard/otel-collector.yaml new file mode 100644 index 00000000..330a8707 --- /dev/null +++ b/examples/switchyard/otel-collector.yaml @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +receivers: + otlp: + protocols: + http: + endpoint: 0.0.0.0:4318 + +exporters: + file/trajectory: + path: /artifacts/trajectory.otel.json + otlp/phoenix: + endpoint: phoenix:4317 + tls: + insecure: true + +service: + pipelines: + traces: + receivers: [otlp] + exporters: [file/trajectory, otlp/phoenix] diff --git a/examples/switchyard/plugins.toml b/examples/switchyard/plugins.toml new file mode 100644 index 00000000..f99d9e4c --- /dev/null +++ b/examples/switchyard/plugins.toml @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +version = 1 + +[[components]] +kind = "switchyard" +enabled = true + +[components.config] +mode = "enforce" +priority = 0 +decision_api_url = "http://127.0.0.1:4000/v1/routing/decision" +decision_profile_id = "smart-stage-router" +request_materialization = "recent_message_window" +context_mode = "atof_required" +decision_timeout_millis = 25 +max_retries = 3 +recent_message_count = 8 +enabled_inbound_profiles = ["openai_chat", "openai_responses", "anthropic_messages"] + +[components.config.decision_header_env] +authorization = "SWITCHYARD_AUTHORIZATION" + +[components.config.default_targets] +openai_chat = "openai-default" +openai_responses = "responses-default" +anthropic_messages = "anthropic-default" + +[components.config.targets.openai-default] +model = "weak-model" +protocol = "openai_chat" +endpoint = "/v1/chat/completions" +base_url = "http://127.0.0.1:4101" + +[components.config.targets.responses-default] +model = "weak-model" +protocol = "openai_responses" +endpoint = "/v1/responses" +base_url = "http://127.0.0.1:4101" + +[components.config.targets.anthropic-default] +model = "weak-model" +protocol = "anthropic_messages" +endpoint = "/v1/messages" +base_url = "http://127.0.0.1:4101" + +[components.config.targets.strong] +model = "strong-model" +protocol = "openai_chat" +endpoint = "/v1/chat/completions" +base_url = "http://127.0.0.1:4102" + +[[components]] +kind = "observability" +enabled = true + +[components.config.atof] +enabled = true +mode = "append" + +[[components.config.atof.endpoints]] +url = "http://127.0.0.1:4000/v1/atof/events" +transport = "http_post" +field_name_policy = "preserve" +timeout_millis = 3000 + +[components.config.atof.endpoints.header_env] +authorization = "SWITCHYARD_AUTHORIZATION" diff --git a/examples/switchyard/real-e2e-plugins.toml b/examples/switchyard/real-e2e-plugins.toml new file mode 100644 index 00000000..9c6351a6 --- /dev/null +++ b/examples/switchyard/real-e2e-plugins.toml @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +version = 1 + +[[components]] +kind = "switchyard" +enabled = true + +[components.config] +mode = "enforce" +decision_api_url = "http://127.0.0.1:4000/v1/routing/decision" +decision_profile_id = "remote-stage-router" +request_materialization = "summary_only" +context_mode = "atof_required" +decision_timeout_millis = 1000 +max_retries = 3 +recent_message_count = 8 + +[components.config.decision_header_env] +authorization = "SWITCHYARD_AUTHORIZATION" + +[components.config.default_targets] +openai_chat = "weak" +openai_responses = "responses-default" +anthropic_messages = "anthropic-default" + +[components.config.targets.strong] +model = "provider/strong" +protocol = "openai_chat" +endpoint = "/v1/chat/completions" +base_url = "http://127.0.0.1:4101" + +[components.config.targets.weak] +model = "provider/weak" +protocol = "openai_chat" +endpoint = "/v1/chat/completions" +base_url = "http://127.0.0.1:4101" + +[components.config.targets.responses-default] +model = "provider/weak-responses" +protocol = "openai_responses" +endpoint = "/v1/responses" +base_url = "http://127.0.0.1:4101" + +[components.config.targets.anthropic-default] +model = "provider/weak-anthropic" +protocol = "anthropic_messages" +endpoint = "/v1/messages" +base_url = "http://127.0.0.1:4101" + +[[components]] +kind = "observability" +enabled = true + +[components.config.atof] +enabled = true +mode = "append" + +[[components.config.atof.endpoints]] +url = "http://127.0.0.1:4000/v1/atof/events" +transport = "http_post" +field_name_policy = "preserve" +timeout_millis = 1000 + +[components.config.atof.endpoints.header_env] +authorization = "SWITCHYARD_AUTHORIZATION" diff --git a/examples/switchyard/real-e2e-profiles.yaml b/examples/switchyard/real-e2e-profiles.yaml new file mode 100644 index 00000000..6127b46e --- /dev/null +++ b/examples/switchyard/real-e2e-profiles.yaml @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +targets: + strong: + model: provider/strong + format: openai + base_url: http://127.0.0.1:4101/v1 + weak: + model: provider/weak + format: openai + base_url: http://127.0.0.1:4101/v1 + +profiles: + remote-stage-router: + type: stage_router + capable: strong + efficient: weak + fallback_target_on_evict: strong + picker: efficient_first + confidence_threshold: 0.7 diff --git a/examples/switchyard/run-hermes-ollama-smoke.sh b/examples/switchyard/run-hermes-ollama-smoke.sh new file mode 100755 index 00000000..f390eb75 --- /dev/null +++ b/examples/switchyard/run-hermes-ollama-smoke.sh @@ -0,0 +1,409 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +relay_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +source "$relay_root/examples/switchyard/e2e-common.sh" +switchyard_root="${SWITCHYARD_ROOT:-$(cd "$relay_root/.." && pwd)/Switchyard-topic-nemo-relay-integration}" +switchyard_expected_commit="${SWITCHYARD_EXPECTED_COMMIT:-8f9db9a6a47f848cdff1d262276ba25a8ae9cbc8}" +run_id="$(date -u +%Y%m%dT%H%M%SZ)-$$" +artifact_dir="${SWITCHYARD_TRAJECTORY_DIR:-$relay_root/artifacts/hermes-switchyard-$run_id}" +token="$(e2e_random_token)" +docker_network="switchyard-e2e-$run_id" +phoenix_container="switchyard-phoenix-$run_id" +collector_container="switchyard-otel-$run_id" +phoenix_port="${SWITCHYARD_PHOENIX_PORT:-6006}" +keep_phoenix="${SWITCHYARD_KEEP_PHOENIX:-0}" +collector_running=0 +phoenix_running=0 +network_created=0 + +mkdir -p "$artifact_dir/phoenix" +[[ -d "$switchyard_root" ]] || { echo "Switchyard worktree not found: $switchyard_root" >&2; exit 1; } +e2e_verify_switchyard_checkout "$switchyard_root" "$switchyard_expected_commit" >"$artifact_dir/switchyard-revision.txt" + +cleanup() { + local status=$? + e2e_stop_processes + if [[ $collector_running -eq 1 ]]; then + docker rm -f "$collector_container" >/dev/null 2>&1 || true + fi + if [[ $phoenix_running -eq 1 && ( $status -ne 0 || "$keep_phoenix" != "1" ) ]]; then + docker rm -f "$phoenix_container" >/dev/null 2>&1 || true + phoenix_running=0 + fi + if [[ $network_created -eq 1 && $phoenix_running -eq 0 ]]; then + docker network rm "$docker_network" >/dev/null 2>&1 || true + fi + if [[ $status -ne 0 ]]; then + echo "Hermes/StageRouter smoke failed; artifacts preserved in $artifact_dir" >&2 + e2e_tail_logs "$artifact_dir" + fi +} +trap cleanup EXIT + +for dependency in cargo curl docker hermes jq python3 tar; do + command -v "$dependency" >/dev/null || { + echo "missing required command: $dependency" >&2 + exit 1 + } +done + +docker info >/dev/null +for model in llama3.2:latest qwen3.6:35b; do + curl --fail --silent http://127.0.0.1:11434/api/tags \ + | jq -e --arg model "$model" '.models[] | select(.name == $model)' >/dev/null || { + echo "required Ollama model is not installed: $model" >&2 + exit 1 + } +done + +docker network create "$docker_network" >/dev/null +network_created=1 +docker run --detach --rm \ + --name "$phoenix_container" \ + --network "$docker_network" \ + --network-alias phoenix \ + --publish "127.0.0.1:$phoenix_port:6006" \ + --env PHOENIX_WORKING_DIR=/mnt/data \ + --volume "$artifact_dir/phoenix:/mnt/data" \ + arizephoenix/phoenix:13.22 >"$artifact_dir/phoenix.container-id" +phoenix_running=1 +e2e_wait_for "http://127.0.0.1:$phoenix_port/" 240 0.5 + +docker run --detach --rm \ + --name "$collector_container" \ + --network "$docker_network" \ + --publish 127.0.0.1:4318:4318 \ + --volume "$relay_root/examples/switchyard/otel-collector.yaml:/etc/otelcol-contrib/config.yaml:ro" \ + --volume "$artifact_dir:/artifacts" \ + otel/opentelemetry-collector-contrib:0.135.0 \ + --config=/etc/otelcol-contrib/config.yaml >"$artifact_dir/collector.container-id" +collector_running=1 + +( + cd "$switchyard_root" + OLLAMA_CLASSIFIER_API_KEY="$token" \ + SWITCHYARD_ATOF_BEARER_TOKEN="$token" \ + cargo run -p switchyard-server -- \ + --config "$relay_root/examples/switchyard/hermes-ollama-profiles.yaml" --port 4000 +) >"$artifact_dir/switchyard.log" 2>&1 & +e2e_add_pid "$!" +e2e_wait_for http://127.0.0.1:4000/health 240 0.5 + +run_query() { + local sequence="$1" + local label="$2" + local query="$3" + local resume_id="${4:-}" + local -a resume_args=() + local before_lines=0 + local after_lines + local atif_path + if [[ -n "$resume_id" ]]; then + resume_args=(--resume "$resume_id") + fi + if [[ -f "$artifact_dir/trajectory.atof.jsonl" ]]; then + before_lines="$(wc -l < "$artifact_dir/trajectory.atof.jsonl" | tr -d ' ')" + fi + ( + cd "$artifact_dir" + HERMES_HOME="$artifact_dir/hermes" \ + OPENAI_API_KEY=ollama \ + SWITCHYARD_AUTHORIZATION="Bearer $token" \ + cargo run --manifest-path "$relay_root/Cargo.toml" -p nemo-relay-cli \ + --features switchyard -- \ + run --agent hermes \ + --plugin-config-path "$relay_root/examples/switchyard/hermes-ollama-plugins.toml" \ + -- chat --provider custom --model llama3.2:latest \ + --query "$query" ${resume_args[@]+"${resume_args[@]}"} \ + --toolsets terminal --quiet --max-turns 2 --ignore-rules + ) >"$artifact_dir/query-$sequence-$label.log" 2>&1 + after_lines="$(wc -l < "$artifact_dir/trajectory.atof.jsonl" | tr -d ' ')" + printf '%s\t%s\t%s\t%s\n' "$sequence" "$label" "$before_lines" "$after_lines" \ + >> "$artifact_dir/query-event-ranges.tsv" + # Each query is a separate Relay process. Replay its persisted ATOF segment + # as a completion barrier before the next process starts; ingestion is + # idempotent, so events already delivered by the best-effort live exporter + # are reported as duplicates rather than applied twice. + sed -n "$((before_lines + 1)),${after_lines}p" "$artifact_dir/trajectory.atof.jsonl" \ + > "$artifact_dir/trajectory-$sequence-$label.atof.jsonl" + curl --fail --silent http://127.0.0.1:4000/v1/atof/events \ + -H "authorization: Bearer $token" \ + -H 'content-type: application/x-ndjson' \ + --data-binary "@$artifact_dir/trajectory-$sequence-$label.atof.jsonl" \ + > "$artifact_dir/trajectory-$sequence-$label.atof-ingest.json" + jq -e '(.batch.ingested_events + .batch.duplicate_events) >= 1' \ + "$artifact_dir/trajectory-$sequence-$label.atof-ingest.json" >/dev/null + atif_path="$(find "$artifact_dir" -maxdepth 1 -name 'trajectory-*.atif.json' -print \ + | grep -E '/trajectory-[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\.atif\.json$' \ + | head -1 || true)" + if [[ -z "$atif_path" ]]; then + echo "ATIF exporter did not produce a trajectory for query $sequence" >&2 + exit 1 + fi + mv "$atif_path" "$artifact_dir/trajectory-$sequence-$label.atif.json" +} + +emit_stage_router_signal() { + local sequence="$1" + local label="$2" + local output="$3" + local path="$artifact_dir/trajectory-signal-$sequence-$label.atof.jsonl" + python3 - "$path" "$session_id" "$label" "$output" <<'PY' +import datetime +import json +import pathlib +import sys +import uuid + +path = pathlib.Path(sys.argv[1]) +session_id, label, output = sys.argv[2:] +event_uuid = str(uuid.uuid4()) +base = { + "atof_version": "0.1", + "kind": "scope", + "uuid": event_uuid, + "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(), + "name": "trajectory_fixture", + "category": "tool", + "category_profile": {"tool_call_id": f"fixture-{label}"}, + "metadata": { + "session_id": session_id, + "trajectory_fixture": True, + "trajectory_fixture_label": label, + }, +} +events = [ + {**base, "scope_category": "start", "data": {"label": label}}, + {**base, "scope_category": "end", "data": {"output": output}}, +] +path.write_text("".join(json.dumps(event, separators=(",", ":")) + "\n" for event in events)) +PY + cat "$path" >> "$artifact_dir/trajectory.atof.jsonl" + curl --fail --silent http://127.0.0.1:4000/v1/atof/events \ + -H "authorization: Bearer $token" \ + -H 'content-type: application/x-ndjson' \ + --data-binary "@$path" > "$artifact_dir/trajectory-signal-$sequence-$label.atof-ingest.json" + jq -e '.batch.ingested_events == 2' \ + "$artifact_dir/trajectory-signal-$sequence-$label.atof-ingest.json" >/dev/null +} + +simple_query='Return exactly the integer result of 17 + 25, with no explanation.' +complex_query='Do not call tools. Act as a principal concurrency engineer. Analyze a bounded lock-free MPMC queue that uses compare-and-swap on head and tail but no generation counters. Give a concrete ABA failure interleaving, identify the required C++ memory order on each publication and consumption edge, and propose the smallest defensible correction. Be precise about the linearization points.' +followup_query='Ignore the earlier technical topic. Reply with exactly SIMPLE_DONE and nothing else.' + +run_query 01 simple "$simple_query" +session_id="$(jq -r 'select(.name == "switchyard.routing.requested") | .data.session_id' "$artifact_dir/trajectory.atof.jsonl" | head -1)" +if [[ -z "$session_id" || "$session_id" == "null" ]]; then + echo "could not recover the Hermes session ID from the first routing mark" >&2 + exit 1 +fi +emit_stage_router_signal 02 complex \ + 'CUDA out of memory while analyzing the concurrent queue; critical failure requires careful recovery and a capable model.' +run_query 02 complex "$complex_query" "$session_id" +emit_stage_router_signal 03 simple-followup \ + 'All tests passed. The next request is a direct low-risk formatting response: return exactly SIMPLE_DONE.' +run_query 03 simple-followup "$followup_query" "$session_id" + +# Give the asynchronous OTLP exporter a short flush window, then stop the +# collector cleanly so its file exporter closes the shareable OTLP JSON file. +sleep 2 +docker stop --time 10 "$collector_container" >/dev/null +collector_running=0 + +python3 - "$artifact_dir" "$session_id" "$phoenix_port" "$simple_query" "$complex_query" "$followup_query" <<'PY' +import json +import pathlib +import sys + +root = pathlib.Path(sys.argv[1]) +session_id = sys.argv[2] +phoenix_port = sys.argv[3] +queries = sys.argv[4:] +atof_path = root / "trajectory.atof.jsonl" +events = [json.loads(line) for line in atof_path.read_text().splitlines() if line.strip()] +marks = [event for event in events if event.get("name", "").startswith("switchyard.routing.")] +decisions = [event for event in marks if event.get("name") == "switchyard.routing.decision"] + +expected_models = ["llama3.2:latest", "qwen3.6:35b", "llama3.2:latest"] +event_ranges = [] +for line in (root / "query-event-ranges.tsv").read_text().splitlines(): + sequence, label, start, end = line.split("\t") + event_ranges.append((sequence, label, int(start), int(end))) +representative_decisions = [] +for sequence, label, start, end in event_ranges: + segment = events[start:end] + decision = next( + (event for event in segment if event.get("name") == "switchyard.routing.decision"), + None, + ) + if decision is None: + raise SystemExit(f"query {sequence} ({label}) produced no successful routing decision") + representative_decisions.append(decision) +actual_models = [event.get("data", {}).get("selected_model") for event in representative_decisions] +if actual_models != expected_models: + raise SystemExit(f"unexpected StageRouter route sequence: {actual_models}; expected {expected_models}") + +required_mark_names = {"switchyard.routing.requested", "switchyard.routing.decision"} +if not required_mark_names.issubset({event.get("name") for event in marks}): + raise SystemExit("routing requested/decision marks were not both emitted") +for event in marks: + name = event.get("name") + if event.get("category") != "custom": + raise SystemExit(f"{name} did not use category=custom") + if event.get("category_profile", {}).get("subtype") != name: + raise SystemExit(f"{name} category_profile.subtype was not canonical") + schema = event.get("data_schema", {}) + if schema != {"name": "switchyard.routing_mark", "version": "1"}: + raise SystemExit(f"{name} had unexpected data_schema: {schema}") + metadata = event.get("metadata", {}) + if metadata.get("session_id") != session_id: + raise SystemExit(f"{name} did not mirror session identity") +for event in decisions: + data = event["data"] + for key in ("decision_id", "router", "routing_attempt", "backend_id", "selected_tier", "selected_model", "latency_ms", "rollout_mode"): + if key not in data or data[key] is None: + raise SystemExit(f"decision mark missing {key}: {data}") + +atif_paths = sorted(root.glob("trajectory-*.atif.json")) +if len(atif_paths) != 3: + raise SystemExit(f"expected three ATIF trajectories, found {len(atif_paths)}") +for path in atif_paths: + payload = json.loads(path.read_text()) + if not payload.get("steps"): + raise SystemExit(f"ATIF trajectory has no steps: {path.name}") + +otel_path = root / "trajectory.otel.json" +if not otel_path.exists() or not otel_path.read_text().strip(): + raise SystemExit("OTLP file exporter did not produce trajectory.otel.json") +otel_batches = [json.loads(line) for line in otel_path.read_text().splitlines() if line.strip()] + +summary = { + "session_id": session_id, + "queries": [{ + "sequence": index + 1, + "input": query, + "selected_model": actual_models[index], + "selected_tier": representative_decisions[index]["data"].get("selected_tier"), + "reason_code": representative_decisions[index]["data"].get("reason_code"), + "reason_summary": representative_decisions[index]["data"].get("reason_summary"), + } for index, query in enumerate(queries)], + "expected_route_sequence": expected_models, + "actual_route_sequence": actual_models, + "routing_basis": [ + "cold StageRouter efficient default", + "canonical ATOF critical-error tool result (capable override)", + "canonical ATOF clean-tests tool result (efficient classifier decision)", + ], + "atof": {"file": atof_path.name, "event_count": len(events), "routing_mark_count": len(marks)}, + "atif": {"files": [path.name for path in atif_paths], "trajectory_count": len(atif_paths)}, + "otel": {"file": otel_path.name, "export_batch_count": len(otel_batches)}, + "phoenix_url": f"http://127.0.0.1:{phoenix_port}", +} +(root / "trajectory-summary.json").write_text(json.dumps(summary, indent=2) + "\n") +readme = f"""# Hermes / Ollama / Switchyard StageRouter trajectory + +This bundle captures one fixed three-query Hermes session routed through NeMo +Relay and the Switchyard Decision API. The verified representative route is: + +1. `llama3.2:latest` (efficient) — cold StageRouter default +2. `qwen3.6:35b` (capable) — critical-signal StageRouter override +3. `llama3.2:latest` (efficient) — clean-state classifier decision + +Session ID: `{session_id}` + +## Important fixture note + +The `CUDA out of memory` text in `trajectory-signal-02-complex.atof.jsonl` is an +intentional, synthetic ATOF tool-result fixture. The machine did not run out of +memory. Fixture events carry `metadata.trajectory_fixture = true` and a fixture +label so they cannot be confused with organic Hermes events. + +The fixtures are necessary for this demonstration because the current +Switchyard StageRouter Decision API classifies from its accumulated ATOF snapshot, +not directly from `current_request.body`. The critical fixture exercises the +real capable override. The clean-tests fixture removes the prior critical signal +from the one-result window and exercises the efficient classifier path. + +## File map + +| Files | Contents | Test coverage | +| --- | --- | --- | +| `trajectory-summary.json` | Machine-readable queries, selected models, reasons, counts, and Phoenix URL | Confirms expected and actual representative routes match | +| `trajectory.atof.jsonl` | Complete Relay ATOF stream for all three queries and the labeled fixtures | Identity propagation, lifecycle events, routing marks, and accumulator input | +| `trajectory-01-simple.atof.jsonl` | ATOF emitted by the first Hermes invocation | Cold-start efficient default | +| `trajectory-02-complex.atof.jsonl` | ATOF emitted by the complex Hermes invocation | Dispatch through the selected capable backend | +| `trajectory-03-simple-followup.atof.jsonl` | ATOF emitted by the final Hermes invocation | Return to the efficient backend | +| `trajectory-signal-*.atof.jsonl` | Canonical, labeled tool start/end fixtures | Capable critical-error override and efficient clean-state classification | +| `*.atof-ingest.json` | Switchyard ingestion reports for query segments and fixtures | Successful or idempotent ATOF accumulation | +| `trajectory-01-simple.atif.json` | ATIF representation of query 1 | Efficient-model trajectory structure | +| `trajectory-02-complex.atif.json` | ATIF representation of query 2 | Capable-model trajectory structure | +| `trajectory-03-simple-followup.atif.json` | ATIF representation of query 3 | Efficient follow-up trajectory structure | +| `trajectory.otel.json` | OTLP JSON batches written by the OpenTelemetry Collector | Relay spans exported to the collector and forwarded to Phoenix | +| `query-*.log` | Hermes/Relay stdout and stderr for each invocation | Human-readable harness responses and execution diagnostics | +| `query-event-ranges.tsv` | Query label and ATOF line-count boundaries | Separates representative user-query decisions from extra Hermes calls | + +## Routing-mark assertions + +The smoke validates every `switchyard.routing.*` mark in the cumulative ATOF +stream. Each mark must have: + +- `category: "custom"` +- `category_profile.subtype` equal to the mark name +- `data_schema.name: "switchyard.routing_mark"` +- `data_schema.version: "1"` +- the expected session identity in metadata + +Decision marks must also include a decision ID, router, attempt, backend, tier, +model, latency, and rollout mode. This run produced {len(marks)} routing marks +across {len(events)} total ATOF events. + +## Phoenix and OTLP + +During the smoke, Relay sends OTLP/HTTP to a local OpenTelemetry Collector. The +collector writes `trajectory.otel.json` and forwards the same spans to Phoenix. +The run produced {len(otel_batches)} OTLP export batches. When the smoke is run +with `SWITCHYARD_KEEP_PHOENIX=1`, open the `phoenix_url` from +`trajectory-summary.json` before stopping the printed Phoenix container. + +## Reproduce + +Install both `llama3.2:latest` and `qwen3.6:35b` in Ollama, ensure the cumulative +Switchyard checkout is available, then run: + +```bash +SWITCHYARD_KEEP_PHOENIX=1 examples/switchyard/run-hermes-ollama-smoke.sh +``` + +The script validates the route, mark shape, ATIF contents, and OTLP output before +creating this bundle. +""" +(root / "TRAJECTORY_README.md").write_text(readme) +print(json.dumps(summary, indent=2)) +PY + +( + cd "$artifact_dir" + tar -czf trajectory-bundle.tar.gz \ + TRAJECTORY_README.md \ + trajectory-summary.json \ + trajectory.atof.jsonl \ + trajectory.otel.json \ + trajectory-*.atif.json \ + trajectory-*.atof.jsonl \ + trajectory-*.atof-ingest.json \ + query-*.log \ + query-event-ranges.tsv +) + +echo "Hermes/Ollama StageRouter trajectory passed: llama3.2 -> qwen3.6:35b -> llama3.2" +echo "Artifacts: $artifact_dir" +echo "Bundle: $artifact_dir/trajectory-bundle.tar.gz" +if [[ "$keep_phoenix" == "1" ]]; then + echo "Phoenix: http://127.0.0.1:$phoenix_port (container $phoenix_container left running)" +else + echo "Set SWITCHYARD_KEEP_PHOENIX=1 to leave Phoenix running after the smoke." +fi diff --git a/examples/switchyard/run-real-e2e.sh b/examples/switchyard/run-real-e2e.sh new file mode 100755 index 00000000..6ec8f35f --- /dev/null +++ b/examples/switchyard/run-real-e2e.sh @@ -0,0 +1,91 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +relay_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +source "$relay_root/examples/switchyard/e2e-common.sh" +switchyard_root="${SWITCHYARD_ROOT:-$(cd "$relay_root/.." && pwd)/Switchyard-topic-nemo-relay-integration}" +switchyard_expected_commit="${SWITCHYARD_EXPECTED_COMMIT:-8f9db9a6a47f848cdff1d262276ba25a8ae9cbc8}" +work_dir="$(mktemp -d)" +upstream_log="$work_dir/upstream.jsonl" +token="$(e2e_random_token)" + +[[ -d "$switchyard_root" ]] || { echo "Switchyard worktree not found: $switchyard_root" >&2; exit 1; } +e2e_verify_switchyard_checkout "$switchyard_root" "$switchyard_expected_commit" + +cleanup() { + local status=$? + e2e_stop_processes + if [[ $status -eq 0 ]]; then + rm -rf "$work_dir" + else + echo "E2E logs preserved in $work_dir" >&2 + e2e_tail_logs "$work_dir" + fi +} +trap cleanup EXIT + +python3 "$relay_root/examples/switchyard/fake_upstream.py" \ + --port 4101 --log "$upstream_log" >"$work_dir/upstream.log" 2>&1 & +e2e_add_pid "$!" + +( + cd "$switchyard_root" + SWITCHYARD_ATOF_BEARER_TOKEN="$token" cargo run -p switchyard-server -- \ + --config "$relay_root/examples/switchyard/real-e2e-profiles.yaml" --port 4000 +) >"$work_dir/switchyard.log" 2>&1 & +e2e_add_pid "$!" + +( + cd "$work_dir" + SWITCHYARD_AUTHORIZATION="Bearer $token" cargo run \ + --manifest-path "$relay_root/Cargo.toml" -p nemo-relay-cli --features switchyard -- \ + --plugin-config-path "$relay_root/examples/switchyard/real-e2e-plugins.toml" \ + --bind 127.0.0.1:4041 +) >"$work_dir/relay.log" 2>&1 & +e2e_add_pid "$!" + +e2e_wait_for http://127.0.0.1:4000/health +e2e_wait_for http://127.0.0.1:4041/healthz + +request() { + local request_id="$1" + local stream="$2" + curl --fail --silent --no-buffer http://127.0.0.1:4041/v1/chat/completions \ + -H 'content-type: application/json' \ + -H 'x-nemo-relay-session-id: e2e-session' \ + -H "x-nemo-relay-request-id: $request_id" \ + --data-binary "{\"model\":\"client-model\",\"stream\":$stream,\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}]}" +} + +request cold-request false >"$work_dir/cold.json" + +for payload in \ + '{"hook_event_name":"on_session_start","session_id":"e2e-session"}' \ + '{"hook_event_name":"pre_tool_call","session_id":"e2e-session","tool_name":"Bash","tool_input":{"command":"test"},"extra":{"task_id":"task-1","tool_call_id":"call-1"}}' \ + '{"hook_event_name":"post_tool_call","session_id":"e2e-session","tool_name":"Bash","tool_input":{"command":"test"},"tool_response":{"output":"CUDA out of memory"},"extra":{"task_id":"task-1","tool_call_id":"call-1"}}' +do + curl --fail --silent http://127.0.0.1:4041/hooks/hermes \ + -H 'content-type: application/json' --data-binary "$payload" >/dev/null +done + +sleep 1 +request warm-request false >"$work_dir/warm.json" +request stream-request true >"$work_dir/stream.sse" + +python3 - "$upstream_log" "$work_dir/stream.sse" <<'PY' +import json +import pathlib +import sys + +records = [json.loads(line) for line in pathlib.Path(sys.argv[1]).read_text().splitlines()] +models = [record["body"]["model"] for record in records] +if models != ["provider/weak", "provider/strong", "provider/strong"]: + raise SystemExit(f"unexpected cold/warm/stream route sequence: {models}") +stream = pathlib.Path(sys.argv[2]).read_text() +if "fake" not in stream or "[DONE]" not in stream: + raise SystemExit(f"unexpected SSE output: {stream}") +print(f"real Switchyard E2E passed: {models}") +PY