From 688d8c264c99e8fe48904c7e2184703d10e92e41 Mon Sep 17 00:00:00 2001 From: skeptrune Date: Wed, 26 Aug 2026 09:51:30 -0700 Subject: [PATCH 1/8] feat(global-routing): add signed request envelope The envelope is the trust boundary between global processing (authentication, canonical normalization, native tokenization, one regional routing decision) and regional execution (pod proxy admission, serving, billing finalization). Detached-JWS wire form: env1... Verifiers check the signature over the received payload bytes and only then parse, so no cross-language canonicalization exists. The payload carries immutable account identifiers (never a raw API key), the canonical normalized body, the native prompt token digest that the serving frontend must reproduce, the selected region and pool, and a bounded validity window. Claude-Session: https://claude.ai/code/session_01Gq5jeTYZ2SDusjNscnyVF5 Signed-off-by: skeptrune --- Cargo.lock | 5 + Cargo.toml | 1 + lib/llm/Cargo.toml | 1 + lib/llm/src/global_routing_envelope.rs | 394 +++++++++++++++++++++++++ lib/llm/src/lib.rs | 1 + 5 files changed, 402 insertions(+) create mode 100644 lib/llm/src/global_routing_envelope.rs diff --git a/Cargo.lock b/Cargo.lock index 0175ab9009b4..2436a7857a74 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2054,6 +2054,7 @@ dependencies = [ "fiat-crypto", "rustc_version", "subtle", + "zeroize", ] [[package]] @@ -2734,6 +2735,7 @@ dependencies = [ "dynamo-runtime", "dynamo-tokenizers", "dynamo-tokens", + "ed25519-dalek", "either", "ffmpeg-next", "flate2", @@ -3248,6 +3250,7 @@ version = "2.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" dependencies = [ + "pkcs8 0.10.2", "signature 2.2.0", ] @@ -3259,9 +3262,11 @@ checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" dependencies = [ "curve25519-dalek", "ed25519", + "serde", "sha2 0.10.9", "signature 2.2.0", "subtle", + "zeroize", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 4a4e21a0619b..35d6f189feb3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -110,6 +110,7 @@ chrono = { version = "0.4", default-features = false, features = [ "serde", ] } cudarc = { version = "=0.19.8", features = ["cuda-version-from-build-system", "fallback-latest"] } +ed25519-dalek = { version = "2.2", features = ["pkcs8", "pem"] } dashmap = { version = "6.1" } moka = { version = "0.12", features = ["future"] } derive_builder = { version = "0.20" } diff --git a/lib/llm/Cargo.toml b/lib/llm/Cargo.toml index 7f086bfdcd57..5b7a280d6b4a 100644 --- a/lib/llm/Cargo.toml +++ b/lib/llm/Cargo.toml @@ -132,6 +132,7 @@ modelexpress-common = { workspace = true } bitflags = { version = "2.4", features = ["serde"] } blake3 = { version = "1.8", features = ["mmap", "rayon"] } +ed25519-dalek = { workspace = true } bytemuck = "1.22" derive-getters = "0.5" offset-allocator = "0.2" diff --git a/lib/llm/src/global_routing_envelope.rs b/lib/llm/src/global_routing_envelope.rs new file mode 100644 index 000000000000..8de46621b5d4 --- /dev/null +++ b/lib/llm/src/global_routing_envelope.rs @@ -0,0 +1,394 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Signed internal request envelope for global routing. +//! +//! The envelope is the trust boundary between global processing +//! (authentication, canonical normalization, native tokenization, the single +//! regional routing decision) and regional execution (pod proxy admission, +//! serving, billing finalization). The dispatcher mints one envelope per +//! request; the selected region's pod proxy verifies it and then trusts the +//! authenticated account identity and canonical request without repeating +//! public authentication or normalization. +//! +//! # Wire format +//! +//! ```text +//! env1.. +//! ``` +//! +//! The signature covers the exact payload bytes as transmitted. Verifiers in +//! any language MUST verify the signature over the received base64-decoded +//! payload bytes and only then parse the JSON — never re-serialize. This +//! removes every cross-language canonicalization hazard by construction. +//! +//! Downstream of the pod proxy, the execute-exact contract is carried by +//! `request.prompt_token_digest`: the serving frontend recomputes the digest +//! from its own native preprocessing and rejects any mismatch, so a request +//! can never be routed with one token sequence and executed with another. + +use serde::{Deserialize, Serialize}; + +use crate::protocols::TokenIdType; + +/// Envelope wire prefix. Bump only with a coordinated verifier rollout. +const WIRE_PREFIX: &str = "env1"; +pub const ENVELOPE_VERSION: u32 = 1; + +/// Domain separation context for the native prompt token digest. +const TOKEN_DIGEST_CONTEXT: &str = "dynamo/global-routing/prompt-tokens/v1"; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Envelope { + pub version: u32, + pub request_id: String, + pub decision_id: String, + /// Unix seconds. + pub issued_at: u64, + /// Unix seconds. Envelopes are short-lived dispatch instructions. + pub expires_at: u64, + pub account: Account, + pub request: Request, + pub routing: Routing, + pub integrity: Integrity, +} + +/// Immutable internal identifiers resolved by global authentication. +/// Never a raw API key. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Account { + pub api_key_id: String, + pub user_id: String, + pub billing_owner_id: String, + pub service_tier: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Request { + /// The public endpoint the client called, e.g. `/v1/chat/completions`. + pub endpoint: String, + pub canonical_model: String, + /// The canonically normalized OpenAI request body. This exact body is + /// what the selected region executes. + pub normalized_body: serde_json::Value, + /// Digest of the exact native token sequence used for routing + /// (see [`prompt_token_digest`]). The serving frontend must reproduce it. + pub prompt_token_digest: String, + pub prompt_token_count: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Routing { + /// Region key of the selected data center, e.g. `us-east5-financial`. + pub selected_region: String, + pub selected_pool: String, + pub policy_version: String, + pub ckf_version: String, + pub cache_overlap_blocks: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct Integrity { + /// `sha256:` over this payload's own serialization of + /// `request.normalized_body`. Informational belt-and-braces inside an + /// already-signed payload; the authoritative execution contract is + /// `prompt_token_digest`. + pub body_digest: String, + pub nonce: String, +} + +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub enum EnvelopeError { + #[error("envelope is malformed")] + Malformed, + #[error("envelope version is unsupported")] + UnsupportedVersion, + #[error("envelope signature is invalid")] + BadSignature, + #[error("envelope is expired or not yet valid")] + Expired, + #[error("envelope is addressed to region {selected}, not {local}")] + WrongRegion { selected: String, local: String }, +} + +pub struct EnvelopeSigner { + key: ed25519_dalek::SigningKey, +} + +impl EnvelopeSigner { + /// Load a PKCS#8 PEM Ed25519 private key + /// (`openssl genpkey -algorithm ed25519`). + pub fn from_pkcs8_pem(pem: &str) -> Result { + use ed25519_dalek::pkcs8::DecodePrivateKey; + Ok(Self { + key: ed25519_dalek::SigningKey::from_pkcs8_pem(pem)?, + }) + } + + pub fn verifier(&self) -> EnvelopeVerifier { + EnvelopeVerifier { + key: self.key.verifying_key(), + } + } + + /// Serialize and sign the envelope into its wire form. + pub fn sign(&self, envelope: &Envelope) -> Result { + use base64::Engine; + use ed25519_dalek::Signer; + + let engine = &base64::engine::general_purpose::URL_SAFE_NO_PAD; + let payload = serde_json::to_vec(envelope)?; + let signature = self.key.sign(&payload); + Ok(format!( + "{WIRE_PREFIX}.{}.{}", + engine.encode(&payload), + engine.encode(signature.to_bytes()) + )) + } +} + +pub struct EnvelopeVerifier { + key: ed25519_dalek::VerifyingKey, +} + +impl EnvelopeVerifier { + /// Load a SPKI PEM Ed25519 public key (`openssl pkey -pubout`). + pub fn from_public_key_pem(pem: &str) -> Result { + use ed25519_dalek::pkcs8::DecodePublicKey; + Ok(Self { + key: ed25519_dalek::VerifyingKey::from_public_key_pem(pem)?, + }) + } + + /// Verify signature, version, validity window, and addressing. + /// + /// `local_region` is the verifier's own region key; an envelope selected + /// for any other region is rejected. Replay protection beyond the validity + /// window (the `nonce`) is the caller's responsibility, since it needs + /// shared state. + pub fn verify( + &self, + wire: &str, + local_region: &str, + now_unix_seconds: u64, + ) -> Result { + use base64::Engine; + + let engine = &base64::engine::general_purpose::URL_SAFE_NO_PAD; + let mut parts = wire.split('.'); + let (Some(WIRE_PREFIX), Some(payload), Some(signature), None) = + (parts.next(), parts.next(), parts.next(), parts.next()) + else { + return Err(EnvelopeError::Malformed); + }; + let payload = engine.decode(payload).map_err(|_| EnvelopeError::Malformed)?; + let signature = engine + .decode(signature) + .ok() + .and_then(|bytes| ed25519_dalek::Signature::from_slice(&bytes).ok()) + .ok_or(EnvelopeError::Malformed)?; + self.key + .verify_strict(&payload, &signature) + .map_err(|_| EnvelopeError::BadSignature)?; + + let envelope: Envelope = + serde_json::from_slice(&payload).map_err(|_| EnvelopeError::Malformed)?; + if envelope.version != ENVELOPE_VERSION { + return Err(EnvelopeError::UnsupportedVersion); + } + if now_unix_seconds < envelope.issued_at || now_unix_seconds > envelope.expires_at { + return Err(EnvelopeError::Expired); + } + if envelope.routing.selected_region != local_region { + return Err(EnvelopeError::WrongRegion { + selected: envelope.routing.selected_region, + local: local_region.to_owned(), + }); + } + Ok(envelope) + } +} + +/// Digest of the exact native token sequence, as `blake3:`. +/// +/// Domain separated and length prefixed; token ids are hashed as little +/// endian `u32`. Both the dispatcher (over the routing tokens) and the +/// serving frontend (over its own preprocessing output) must use this +/// function so the execute-exact comparison is meaningful. +pub fn prompt_token_digest(token_ids: &[TokenIdType]) -> String { + let mut hasher = blake3::Hasher::new_derive_key(TOKEN_DIGEST_CONTEXT); + hasher.update(&(token_ids.len() as u64).to_le_bytes()); + for token_id in token_ids { + hasher.update(&token_id.to_le_bytes()); + } + format!("blake3:{}", hasher.finalize().to_hex()) +} + +/// `sha256:` over the payload's serialization of the normalized body. +pub fn body_digest(normalized_body: &serde_json::Value) -> Result { + use sha2::Digest; + let bytes = serde_json::to_vec(normalized_body)?; + let mut hex = String::with_capacity(7 + 64); + hex.push_str("sha256:"); + for byte in sha2::Sha256::digest(&bytes) { + use std::fmt::Write; + write!(&mut hex, "{byte:02x}").expect("writing to a String cannot fail"); + } + Ok(hex) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Deterministic key for cross-implementation test vectors only. + fn test_signer() -> EnvelopeSigner { + EnvelopeSigner { + key: ed25519_dalek::SigningKey::from_bytes(&[7u8; 32]), + } + } + + fn envelope() -> Envelope { + let normalized_body = serde_json::json!({ + "model": "morph-dsv4flash", + "messages": [{"role": "user", "content": "hello"}], + "stream": true, + }); + Envelope { + version: ENVELOPE_VERSION, + request_id: "req_1".into(), + decision_id: "decision_1".into(), + issued_at: 1_787_780_000, + expires_at: 1_787_780_030, + account: Account { + api_key_id: "key_1".into(), + user_id: "user_1".into(), + billing_owner_id: "org_1".into(), + service_tier: "normal".into(), + }, + request: Request { + endpoint: "/v1/chat/completions".into(), + canonical_model: "morph-dsv4flash".into(), + normalized_body: normalized_body.clone(), + prompt_token_digest: prompt_token_digest(&[1, 2, 3, 4]), + prompt_token_count: 4, + }, + routing: Routing { + selected_region: "us-east5-financial".into(), + selected_pool: "pool_1".into(), + policy_version: "policy_v1".into(), + ckf_version: "ckf_v1".into(), + cache_overlap_blocks: 3, + }, + integrity: Integrity { + body_digest: body_digest(&normalized_body).unwrap(), + nonce: "nonce_1".into(), + }, + } + } + + #[test] + fn sign_verify_roundtrip() { + let signer = test_signer(); + let envelope = envelope(); + let wire = signer.sign(&envelope).unwrap(); + let verified = signer + .verifier() + .verify(&wire, "us-east5-financial", 1_787_780_010) + .unwrap(); + assert_eq!(verified, envelope); + } + + #[test] + fn verification_rejects_tampering_expiry_and_misaddressing() { + let signer = test_signer(); + let verifier = signer.verifier(); + let wire = signer.sign(&envelope()).unwrap(); + let now = 1_787_780_010; + + // Payload tampering: flip one payload character. + let mut tampered = wire.clone().into_bytes(); + let payload_start = WIRE_PREFIX.len() + 1; + tampered[payload_start] = if tampered[payload_start] == b'A' { b'B' } else { b'A' }; + let tampered = String::from_utf8(tampered).unwrap(); + assert!(matches!( + verifier.verify(&tampered, "us-east5-financial", now), + Err(EnvelopeError::BadSignature | EnvelopeError::Malformed) + )); + + // Signature from a different key. + let other = EnvelopeSigner { + key: ed25519_dalek::SigningKey::from_bytes(&[8u8; 32]), + }; + let foreign = other.sign(&envelope()).unwrap(); + assert_eq!( + verifier.verify(&foreign, "us-east5-financial", now), + Err(EnvelopeError::BadSignature) + ); + + // Expired and not-yet-valid. + assert_eq!( + verifier.verify(&wire, "us-east5-financial", 1_787_780_031), + Err(EnvelopeError::Expired) + ); + assert_eq!( + verifier.verify(&wire, "us-east5-financial", 1_787_779_999), + Err(EnvelopeError::Expired) + ); + + // Addressed to another region. + assert_eq!( + verifier.verify(&wire, "us-west1-financial", now), + Err(EnvelopeError::WrongRegion { + selected: "us-east5-financial".into(), + local: "us-west1-financial".into(), + }) + ); + + // Garbage. + assert_eq!( + verifier.verify("env1.zzz", "us-east5-financial", now), + Err(EnvelopeError::Malformed) + ); + assert_eq!( + verifier.verify("env2.a.b", "us-east5-financial", now), + Err(EnvelopeError::Malformed) + ); + } + + #[test] + fn token_digest_is_order_and_length_sensitive() { + assert_eq!(prompt_token_digest(&[1, 2, 3]), prompt_token_digest(&[1, 2, 3])); + assert_ne!(prompt_token_digest(&[1, 2, 3]), prompt_token_digest(&[3, 2, 1])); + assert_ne!(prompt_token_digest(&[1, 2, 3]), prompt_token_digest(&[1, 2])); + assert!(prompt_token_digest(&[]).starts_with("blake3:")); + } + + /// The wire form is fully deterministic (struct-order serialization plus + /// RFC 8032 deterministic signatures), so cross-implementation verifier + /// vectors can be generated by printing `sign` output for a fixed key. + #[test] + fn wire_form_is_deterministic() { + use base64::Engine; + + let first = test_signer().sign(&envelope()).unwrap(); + let second = test_signer().sign(&envelope()).unwrap(); + assert_eq!(first, second); + + let parts: Vec<&str> = first.split('.').collect(); + assert_eq!(parts.len(), 3); + assert_eq!(parts[0], "env1"); + // Ed25519 signatures are 64 bytes; base64url without padding is 86. + assert_eq!(parts[2].len(), 86); + // The payload is the plain serde serialization of the envelope. + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(parts[1]) + .unwrap(); + assert_eq!(payload, serde_json::to_vec(&envelope()).unwrap()); + } +} diff --git a/lib/llm/src/lib.rs b/lib/llm/src/lib.rs index 37f048a3be6a..49f79d230e28 100644 --- a/lib/llm/src/lib.rs +++ b/lib/llm/src/lib.rs @@ -17,6 +17,7 @@ pub mod first_token; pub mod fpm_publisher; pub mod fpm_trace; pub mod frontend_config; +pub mod global_routing_envelope; pub mod global_routing_shadow; pub mod grpc; pub mod http; From 3102b745055230bdfe2626171b707d8ecd72918e Mon Sep 17 00:00:00 2001 From: skeptrune Date: Wed, 26 Aug 2026 09:58:26 -0700 Subject: [PATCH 2/8] feat(global-routing): enforce execute-exact token digests Add nvext.prompt_token_digest, stamped by the regional pod proxy from the signed routing envelope, and reject any request whose native preprocessing does not reproduce the routed token sequence. This closes the route-one-prompt-execute-another gap: normalization drift between global processing and a serving frontend now fails closed with a 409 instead of serving silently, and both outcomes are counted in morph_global_routing_execute_exact_total. Claude-Session: https://claude.ai/code/session_01Gq5jeTYZ2SDusjNscnyVF5 Signed-off-by: skeptrune --- lib/llm/src/global_routing_envelope.rs | 68 ++++++++++++++++++++++ lib/llm/src/http/service/service_v2.rs | 5 ++ lib/llm/src/preprocessor.rs | 18 ++++++ lib/llm/src/protocols/common/extensions.rs | 12 ++++ 4 files changed, 103 insertions(+) diff --git a/lib/llm/src/global_routing_envelope.rs b/lib/llm/src/global_routing_envelope.rs index 8de46621b5d4..fc72dc1fe822 100644 --- a/lib/llm/src/global_routing_envelope.rs +++ b/lib/llm/src/global_routing_envelope.rs @@ -27,6 +27,9 @@ //! from its own native preprocessing and rejects any mismatch, so a request //! can never be routed with one token sequence and executed with another. +use std::sync::LazyLock; + +use prometheus::{IntCounterVec, Opts, Registry}; use serde::{Deserialize, Serialize}; use crate::protocols::TokenIdType; @@ -214,6 +217,57 @@ impl EnvelopeVerifier { } } +static EXECUTE_EXACT: LazyLock = LazyLock::new(|| { + IntCounterVec::new( + Opts::new( + "morph_global_routing_execute_exact_total", + "Execute-exact digest comparisons on globally routed requests.", + ), + &["outcome"], + ) + .expect("static metric options are valid") +}); + +/// Register this module's frontend-side collectors. +pub fn ensure_metrics_registered_prometheus(registry: &Registry) -> Result<(), prometheus::Error> { + registry.register(Box::new(EXECUTE_EXACT.clone())) +} + +/// Enforce the execute-exact contract on the serving frontend. +/// +/// `expected` is `nvext.prompt_token_digest`, stamped by the regional pod +/// proxy from the signed routing envelope. A mismatch means the frontend's +/// preprocessing produced a different token sequence than the one that drove +/// the routing decision — normalization drift between global processing and +/// this frontend — and the request must be rejected, never served. +pub fn enforce_prompt_token_digest( + expected: Option<&str>, + token_ids: &[TokenIdType], + request_id: &str, +) -> Result<(), crate::http::service::error::HttpError> { + let Some(expected) = expected else { + return Ok(()); + }; + let executed = prompt_token_digest(token_ids); + if executed == expected { + EXECUTE_EXACT.with_label_values(&["match"]).inc(); + return Ok(()); + } + EXECUTE_EXACT.with_label_values(&["mismatch"]).inc(); + tracing::warn!( + request_id, + expected, + executed, + token_count = token_ids.len(), + "execute-exact violation: routed and executed token sequences differ" + ); + Err(crate::http::service::error::HttpError { + code: 409, + message: "prompt token digest mismatch: the routed and executed token sequences differ" + .to_string(), + }) +} + /// Digest of the exact native token sequence, as `blake3:`. /// /// Domain separated and length prefixed; token ids are hashed as little @@ -361,6 +415,20 @@ mod tests { ); } + #[test] + fn execute_exact_enforces_only_when_a_digest_is_expected() { + let tokens = [1, 2, 3, 4]; + assert!(enforce_prompt_token_digest(None, &tokens, "req").is_ok()); + assert!( + enforce_prompt_token_digest(Some(&prompt_token_digest(&tokens)), &tokens, "req") + .is_ok() + ); + let error = + enforce_prompt_token_digest(Some(&prompt_token_digest(&[9])), &tokens, "req") + .unwrap_err(); + assert_eq!(error.code, 409); + } + #[test] fn token_digest_is_order_and_length_sensitive() { assert_eq!(prompt_token_digest(&[1, 2, 3]), prompt_token_digest(&[1, 2, 3])); diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index dee46c0e9139..37b5c066358a 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -1212,6 +1212,11 @@ impl HttpServiceConfigBuilder { if let Err(e) = ensure_frontend_perf_metrics_registered_prometheus(®istry) { tracing::warn!("Failed to register frontend perf metrics: {}", e); } + if let Err(e) = + crate::global_routing_envelope::ensure_metrics_registered_prometheus(®istry) + { + tracing::warn!("Failed to register global routing envelope metrics: {}", e); + } if let Err(e) = ensure_tokio_perf_metrics_registered_prometheus(®istry) { tracing::warn!("Failed to register tokio perf metrics: {}", e); } diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 5f2719b2c078..134c79d5723e 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -5457,6 +5457,15 @@ impl .await?; attach_agent_context_from_context(&mut common_request, &context); + crate::global_routing_envelope::enforce_prompt_token_digest( + request + .nvext + .as_ref() + .and_then(|nvext| nvext.prompt_token_digest.as_deref()), + &common_request.token_ids, + &request_id, + )?; + crate::global_routing_shadow::observe( &common_request.model, &common_request.token_ids, @@ -5663,6 +5672,15 @@ impl Self::validate_preprocessed_token_budget(&common_request, self.token_budget.as_ref())?; attach_agent_context_from_context(&mut common_request, &context); + crate::global_routing_envelope::enforce_prompt_token_digest( + request + .nvext + .as_ref() + .and_then(|nvext| nvext.prompt_token_digest.as_deref()), + &common_request.token_ids, + &request_id, + )?; + crate::global_routing_shadow::observe( &common_request.model, &common_request.token_ids, diff --git a/lib/llm/src/protocols/common/extensions.rs b/lib/llm/src/protocols/common/extensions.rs index d0598b30b04f..77a6146f0de7 100644 --- a/lib/llm/src/protocols/common/extensions.rs +++ b/lib/llm/src/protocols/common/extensions.rs @@ -249,6 +249,16 @@ pub struct NvExt { #[builder(default, setter(strip_option))] #[serde(default, skip_serializing_if = "Option::is_none")] pub router: Option, + + /// Execute-exact contract for globally routed requests: the expected + /// native prompt token digest from the signed routing envelope, stamped + /// by the regional pod proxy. After preprocessing, the frontend must + /// reproduce this digest from its own token sequence or reject the + /// request, so a request can never be routed with one prompt and + /// executed with another. + #[builder(default, setter(strip_option, into))] + #[serde(default, skip_serializing_if = "Option::is_none")] + pub prompt_token_digest: Option, } impl Default for NvExt { @@ -290,6 +300,7 @@ impl NvExt { request_timestamp_ms, routing_constraints, router, + prompt_token_digest, } = self; greed_sampling.is_some() @@ -308,6 +319,7 @@ impl NvExt { || request_timestamp_ms.is_some() || routing_constraints.is_some() || router.is_some() + || prompt_token_digest.is_some() } } From aefad0a62e4aeadc7d9be2fa23dd9f83aaff85d8 Mon Sep 17 00:00:00 2001 From: skeptrune Date: Wed, 26 Aug 2026 10:18:23 -0700 Subject: [PATCH 3/8] feat(global-routing): mint signed envelopes from authoritative decisions Replace shadow observation with an authoritative in-path decision client (DYN_GLOBAL_ROUTER_CONSUMER_URL): the frontend queries the consumer after native preprocessing, fails closed when configured, and mints the signed internal request envelope from the trusted auth metadata that pod-proxy-lifted resolves. The envelope rides request context until dispatch. The decision wire contract moves into dynamo-llm as the canonical types and the consumer re-exports them, adding selected_region so dispatchers can address the target region without a dc-id lookup. Envelope wire form becomes a JSON object with key_id for rotation and a domain-separated Ed25519 signature. Claude-Session: https://claude.ai/code/session_01Gq5jeTYZ2SDusjNscnyVF5 Signed-off-by: skeptrune --- components/global-ckf-consumer/Cargo.toml | 2 +- components/global-ckf-consumer/src/api.rs | 222 ++++---- .../global-ckf-consumer/src/contract.rs | 8 + components/global-ckf-consumer/src/lib.rs | 1 + lib/llm/src/global_routing.rs | 257 +++++++++ lib/llm/src/global_routing_envelope.rs | 506 ++++++++++++++---- lib/llm/src/global_routing_shadow.rs | 147 ----- lib/llm/src/http/service/service_v2.rs | 3 + lib/llm/src/lib.rs | 2 +- lib/llm/src/preprocessor.rs | 100 +++- 10 files changed, 895 insertions(+), 353 deletions(-) create mode 100644 components/global-ckf-consumer/src/contract.rs create mode 100644 lib/llm/src/global_routing.rs delete mode 100644 lib/llm/src/global_routing_shadow.rs diff --git a/components/global-ckf-consumer/Cargo.toml b/components/global-ckf-consumer/Cargo.toml index c84f19eafe2b..04badbb1db3c 100644 --- a/components/global-ckf-consumer/Cargo.toml +++ b/components/global-ckf-consumer/Cargo.toml @@ -9,7 +9,7 @@ repository.workspace = true anyhow.workspace = true axum.workspace = true clap.workspace = true -dynamo-llm = { workspace = true, features = ["kv-dc-relay-proto"] } +dynamo-llm = { workspace = true, default-features = false, features = ["kv-dc-relay-proto"] } dynamo-kv-router.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/components/global-ckf-consumer/src/api.rs b/components/global-ckf-consumer/src/api.rs index bec4a4ba395d..d0d4ee9e7627 100644 --- a/components/global-ckf-consumer/src/api.rs +++ b/components/global-ckf-consumer/src/api.rs @@ -17,6 +17,7 @@ use dynamo_kv_router::protocols::{ }; use serde::{Deserialize, Serialize}; +use crate::contract::{DecisionOutcome, QueryRole, TokenDecisionRequest, TokenDecisionResponse}; use crate::lane::{LaneAvailability, LaneSet, LaneUnavailableReason}; use crate::policy::{ Freshness, LaneFact, OccupancyFact, PolicyInput, PoolFacts, ReadinessFact, select_pool, @@ -199,35 +200,6 @@ struct TokenPrefixMatchesRequest { is_eagle: Option, } -#[derive(Debug, Deserialize)] -#[serde(deny_unknown_fields)] -struct TokenDecisionRequest { - #[serde(flatten)] - query: TokenPrefixMatchesRequest, - local_dc: u64, - stable_tie_key: u64, - readiness_max_age_ms: u64, - load_max_age_ms: u64, -} - -#[derive(Debug, Serialize)] -struct TokenDecisionResponse { - generation: u64, - selected_pool_id: Option, - selected_dc: Option, - outcome: &'static str, - matched_prefix_blocks: Option, - uncached_prefill_tokens: Option, -} - -#[derive(Debug, Clone, Copy, Deserialize)] -#[serde(rename_all = "snake_case")] -enum QueryRole { - Aggregated, - Prefill, - Decode, -} - #[derive(Debug, Serialize)] struct PrefixMatchesResponse { generation: u64, @@ -256,19 +228,19 @@ impl From for QueryRoleResponse { } #[derive(Debug, Serialize)] -struct PoolFact { +pub(crate) struct PoolFact { #[serde(skip)] - native_pool_id: PoolId, + pub(crate) native_pool_id: PoolId, pool_id: String, indexer_domain: IndexerDomainFact, - relay: String, - dc_id: u64, + pub(crate) relay: String, + pub(crate) dc_id: u64, physical_lane: u8, availability: &'static str, unavailable_reason: Option<&'static str>, prefix_depth_blocks: Option, producer_incarnation: Option, - layout_generation: Option, + pub(crate) layout_generation: Option, installed_sequence: Option, freshness_age_ms: Option, readiness: Vec, @@ -417,33 +389,111 @@ async fn token_decision( max_query_blocks: usize, ) -> Response { state.metrics.queries.fetch_add(1, Ordering::Relaxed); - if request.query.model.trim().is_empty() { + if request.model.trim().is_empty() { return decision_error(&state, ApiError::bad_request("model must not be empty")); } - let hashes = match native_hashes(&request.query, max_query_blocks) { + let query = TokenPrefixMatchesRequest { + model: request.model.clone(), + role: request.role, + token_ids: request.token_ids.clone(), + block_size: request.block_size, + block_mm_infos: request.block_mm_infos.clone(), + lora_name: request.lora_name.clone(), + cache_namespace: request.cache_namespace.clone(), + is_eagle: request.is_eagle, + }; + let hashes = match native_hashes(&query, max_query_blocks) { Ok(hashes) => hashes, Err(error) => return decision_error(&state, error), }; - let (generation, pools) = match collect_facts( + let evaluated = match evaluate_decision( &state, - Some(&request.query.model), - request.query.role, + &request.model, + request.role, &hashes, + request.block_size, + request.local_dc, + request.stable_tie_key, + request.readiness_max_age_ms, + request.load_max_age_ms, ) { - Ok(result) => result, + Ok(evaluated) => evaluated, Err(error) => return decision_error(&state, error), }; + let ( + selected_pool_id, + selected_dc, + selected_region, + outcome, + matched_prefix_blocks, + uncached_prefill_tokens, + ) = match evaluated.decision.selected { + Some(selected) => { + let dc = selected.pool_id.dc_id().get(); + let selected_region = evaluated + .pools + .iter() + .find(|pool| pool.native_pool_id == selected.pool_id) + .map(|pool| pool.relay.clone()); + ( + Some(selected.pool_id.to_string()), + Some(dc), + selected_region, + if dc == request.local_dc { + DecisionOutcome::Local + } else { + DecisionOutcome::Remote + }, + Some(selected.matched_prefix_blocks), + Some(selected.uncached_prefill_tokens), + ) + } + None => (None, None, None, DecisionOutcome::None, None, None), + }; + Json(TokenDecisionResponse { + generation: evaluated.generation, + selected_pool_id, + selected_dc, + selected_region, + outcome, + matched_prefix_blocks, + uncached_prefill_tokens, + }) + .into_response() +} + +pub(crate) struct EvaluatedDecision { + pub(crate) generation: u64, + pub(crate) decision: crate::policy::PolicyDecision, + pub(crate) pools: Vec, +} + +/// Evaluate the exact routing policy over the current published facts. +/// The single evaluation path for both the decision API and the dispatcher, +/// so their outcomes and decision metrics can never diverge. +#[allow(clippy::too_many_arguments)] +pub(crate) fn evaluate_decision( + state: &AppState, + model: &str, + role: QueryRole, + hashes: &[u64], + block_size: u32, + local_dc: u64, + stable_tie_key: u64, + readiness_max_age_ms: u64, + load_max_age_ms: u64, +) -> Result { + let (generation, pools) = collect_facts(state, Some(model), role, hashes)?; if pools.is_empty() { - return decision_error( - &state, - ApiError::not_found("no catalog pool matches model and role"), - ); + return Err(ApiError::not_found( + "no catalog pool matches model and role", + )); } let input = PolicyInput { - local_dc: dynamo_kv_router::identity::DcId::new(request.local_dc), + local_dc: dynamo_kv_router::identity::DcId::new(local_dc), query_block_count: hashes.len() as u64, - native_block_size_tokens: request.query.block_size as u64, - stable_tie_key: request.stable_tie_key, + native_block_size_tokens: block_size as u64, + stable_tie_key, }; let candidates = pools.iter().map(|pool| PoolFacts { pool_id: pool.native_pool_id, @@ -454,12 +504,13 @@ async fn token_decision( }, matched_prefix_blocks: u64::from(pool.prefix_depth_blocks.unwrap_or(0)), readiness: pool.readiness_age_ms.map(|age| ReadinessFact { - ready: pool.readiness.iter().any(|signal| { - signal.canonical_model_id == request.query.model && signal.state == 2 - }), + ready: pool + .readiness + .iter() + .any(|signal| signal.canonical_model_id == model && signal.state == 2), freshness: Freshness { age: std::time::Duration::from_millis(age), - maximum_age: std::time::Duration::from_millis(request.readiness_max_age_ms), + maximum_age: std::time::Duration::from_millis(readiness_max_age_ms), }, }), occupancy: match ( @@ -482,57 +533,36 @@ async fn token_decision( expected_ranks, freshness: Freshness { age: std::time::Duration::from_millis(age), - maximum_age: std::time::Duration::from_millis(request.load_max_age_ms), + maximum_age: std::time::Duration::from_millis(load_max_age_ms), }, }), _ => None, }, }); - let decision = match select_pool(input, candidates) { - Ok(decision) => decision, - Err(_) => return decision_error(&state, ApiError::internal("routing policy failed")), - }; - let (selected_pool_id, selected_dc, outcome, matched_prefix_blocks, uncached_prefill_tokens) = - match decision.selected { - Some(selected) => { - let dc = selected.pool_id.dc_id().get(); - if dc == request.local_dc { - state - .metrics - .decisions_local - .fetch_add(1, Ordering::Relaxed); - } else { - state - .metrics - .decisions_remote - .fetch_add(1, Ordering::Relaxed); - } - ( - Some(selected.pool_id.to_string()), - Some(dc), - if dc == request.local_dc { - "local" - } else { - "remote" - }, - Some(selected.matched_prefix_blocks), - Some(selected.uncached_prefill_tokens), - ) - } - None => { - state.metrics.decisions_none.fetch_add(1, Ordering::Relaxed); - (None, None, "none", None, None) - } - }; - Json(TokenDecisionResponse { + let decision = + select_pool(input, candidates).map_err(|_| ApiError::internal("routing policy failed"))?; + match decision.selected { + Some(selected) if selected.pool_id.dc_id().get() == local_dc => { + state + .metrics + .decisions_local + .fetch_add(1, Ordering::Relaxed); + } + Some(_) => { + state + .metrics + .decisions_remote + .fetch_add(1, Ordering::Relaxed); + } + None => { + state.metrics.decisions_none.fetch_add(1, Ordering::Relaxed); + } + } + Ok(EvaluatedDecision { generation, - selected_pool_id, - selected_dc, - outcome, - matched_prefix_blocks, - uncached_prefill_tokens, + decision, + pools, }) - .into_response() } fn decision_error(state: &AppState, error: ApiError) -> Response { @@ -790,7 +820,7 @@ async fn metrics(State(state): State) -> String { } #[derive(Debug)] -struct ApiError(StatusCode, &'static str); +pub(crate) struct ApiError(pub(crate) StatusCode, pub(crate) &'static str); impl ApiError { fn bad_request(message: &'static str) -> Self { diff --git a/components/global-ckf-consumer/src/contract.rs b/components/global-ckf-consumer/src/contract.rs new file mode 100644 index 000000000000..18e71c166f15 --- /dev/null +++ b/components/global-ckf-consumer/src/contract.rs @@ -0,0 +1,8 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Shared decision wire contract. The frontend client owns the canonical types. + +pub use dynamo_llm::global_routing::{ + DecisionOutcome, QueryRole, TokenDecisionRequest, TokenDecisionResponse, +}; diff --git a/components/global-ckf-consumer/src/lib.rs b/components/global-ckf-consumer/src/lib.rs index 222d50b0d52d..fff9dcd122e0 100644 --- a/components/global-ckf-consumer/src/lib.rs +++ b/components/global-ckf-consumer/src/lib.rs @@ -3,6 +3,7 @@ pub mod api; pub mod config; +pub mod contract; pub mod coordinator; pub mod lane; pub mod policy; diff --git a/lib/llm/src/global_routing.rs b/lib/llm/src/global_routing.rs new file mode 100644 index 000000000000..f4180f84b428 --- /dev/null +++ b/lib/llm/src/global_routing.rs @@ -0,0 +1,257 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Authoritative client contract for the global CKF consumer. + +use std::sync::{LazyLock, OnceLock}; +use std::time::Duration; + +use prometheus::{Histogram, HistogramOpts, IntCounterVec, Registry}; +use serde::{Deserialize, Serialize}; + +use crate::protocols::TokenIdType; + +static CLIENT: OnceLock, String>> = OnceLock::new(); + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TokenDecisionRequest { + pub model: String, + pub role: QueryRole, + pub token_ids: Vec, + pub block_size: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub block_mm_infos: Option>>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub lora_name: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cache_namespace: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub is_eagle: Option, + pub local_dc: u64, + pub stable_tie_key: u64, + pub readiness_max_age_ms: u64, + pub load_max_age_ms: u64, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum QueryRole { + Aggregated, + Prefill, + Decode, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TokenDecisionResponse { + pub generation: u64, + pub selected_pool_id: Option, + pub selected_dc: Option, + pub selected_region: Option, + pub outcome: DecisionOutcome, + pub matched_prefix_blocks: Option, + pub uncached_prefill_tokens: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum DecisionOutcome { + Local, + Remote, + None, +} + +#[derive(Debug, thiserror::Error)] +pub enum GlobalRoutingError { + #[error("global routing does not yet support multimodal requests")] + MultimodalUnsupported, + #[error("global routing block size is invalid")] + InvalidBlockSize, + #[error("global routing decision request failed: {0}")] + Request(#[from] reqwest::Error), + #[error("global routing consumer returned no eligible pool")] + NoEligiblePool, + #[error("global routing consumer returned an incomplete decision")] + IncompleteDecision, + #[error("global routing is misconfigured: {0}")] + Config(String), +} + +#[derive(Clone)] +struct Client { + http: reqwest::Client, + endpoint: String, + local_dc: u64, + readiness_max_age_ms: u64, + load_max_age_ms: u64, +} + +impl Client { + fn from_env() -> Result, String> { + let Some(base) = std::env::var("DYN_GLOBAL_ROUTER_CONSUMER_URL") + .ok() + .filter(|value| !value.trim().is_empty()) + else { + return Ok(None); + }; + let parse = |name: &str, default: u64| { + std::env::var(name) + .ok() + .map(|value| value.parse::()) + .transpose() + .map(|value| value.unwrap_or(default)) + .map_err(|error| format!("{name} must be an unsigned integer: {error}")) + }; + let timeout = Duration::from_millis(parse("DYN_GLOBAL_ROUTER_DECISION_TIMEOUT_MS", 1_000)?); + Ok(Some(Self { + http: reqwest::Client::builder() + .timeout(timeout) + .build() + .map_err(|error| format!("failed to build decision HTTP client: {error}"))?, + endpoint: format!("{}/v1/decisions/tokens", base.trim_end_matches('/')), + local_dc: std::env::var("DYN_GLOBAL_ROUTER_LOCAL_DC_ID") + .map_err(|_| "DYN_GLOBAL_ROUTER_LOCAL_DC_ID is required".to_owned())? + .parse() + .map_err(|error| format!("DYN_GLOBAL_ROUTER_LOCAL_DC_ID must be a u64: {error}"))?, + readiness_max_age_ms: parse("DYN_GLOBAL_ROUTER_READINESS_MAX_AGE_MS", 45_000)?, + load_max_age_ms: parse("DYN_GLOBAL_ROUTER_LOAD_MAX_AGE_MS", 15_000)?, + })) + } + + async fn decide( + &self, + request: TokenDecisionRequest, + ) -> Result { + let decision = self + .http + .post(&self.endpoint) + .json(&request) + .send() + .await? + .error_for_status()? + .json::() + .await?; + match ( + decision.outcome, + decision.selected_pool_id.as_ref(), + decision.selected_dc, + decision.selected_region.as_ref(), + ) { + (DecisionOutcome::None, _, _, _) => Err(GlobalRoutingError::NoEligiblePool), + (_, Some(_), Some(_), Some(_)) => Ok(decision), + _ => Err(GlobalRoutingError::IncompleteDecision), + } + } +} + +static DECISIONS: LazyLock = LazyLock::new(|| { + IntCounterVec::new( + prometheus::Opts::new( + "morph_global_routing_authoritative_decisions_total", + "Authoritative global CKF decision outcomes owned by the frontend.", + ), + &["outcome"], + ) + .expect("static metric options are valid") +}); + +static DECISION_LATENCY: LazyLock = LazyLock::new(|| { + Histogram::with_opts(HistogramOpts::new( + "morph_global_routing_decision_duration_seconds", + "Latency of authoritative global CKF decision requests.", + )) + .expect("static metric options are valid") +}); + +pub fn ensure_metrics_registered_prometheus(registry: &Registry) -> Result<(), prometheus::Error> { + registry.register(Box::new(DECISIONS.clone()))?; + registry.register(Box::new(DECISION_LATENCY.clone())) +} + +/// Await and own exactly one global decision after native preprocessing. +/// +/// When global routing is configured, every unsupported or unavailable input +/// fails closed. Dropping the returned decision would reintroduce shadow +/// semantics, so callers must retain it in request context until dispatch. +pub async fn decide( + model: &str, + token_ids: &[TokenIdType], + block_size: usize, + request_id: &str, + has_multimodal_data: bool, +) -> Result, GlobalRoutingError> { + let client = match CLIENT.get_or_init(Client::from_env) { + Ok(Some(client)) => client, + Ok(None) => return Ok(None), + Err(error) => return Err(GlobalRoutingError::Config(error.clone())), + }; + if has_multimodal_data { + DECISIONS.with_label_values(&["error"]).inc(); + return Err(GlobalRoutingError::MultimodalUnsupported); + } + let block_size = u32::try_from(block_size).map_err(|_| GlobalRoutingError::InvalidBlockSize)?; + if block_size == 0 { + return Err(GlobalRoutingError::InvalidBlockSize); + } + let request = TokenDecisionRequest { + model: model.to_owned(), + role: QueryRole::Aggregated, + token_ids: token_ids.to_vec(), + block_size, + block_mm_infos: None, + lora_name: None, + cache_namespace: None, + is_eagle: None, + local_dc: client.local_dc, + stable_tie_key: stable_hash(request_id.as_bytes()), + readiness_max_age_ms: client.readiness_max_age_ms, + load_max_age_ms: client.load_max_age_ms, + }; + let timer = DECISION_LATENCY.start_timer(); + let result = client.decide(request).await; + timer.observe_duration(); + match &result { + Ok(decision) => DECISIONS + .with_label_values(&[match decision.outcome { + DecisionOutcome::Local => "local", + DecisionOutcome::Remote => "remote", + DecisionOutcome::None => "none", + }]) + .inc(), + Err(_) => DECISIONS.with_label_values(&["error"]).inc(), + } + result.map(Some) +} + +fn stable_hash(bytes: &[u8]) -> u64 { + bytes.iter().fold(0xcbf29ce484222325, |hash, byte| { + (hash ^ u64::from(*byte)).wrapping_mul(0x100000001b3) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn decision_contract_roundtrips_and_rejects_unknown_fields() { + let json = r#"{"generation":42,"selected_pool_id":"pool","selected_dc":2,"selected_region":"us-west1-financial","outcome":"remote","matched_prefix_blocks":3,"uncached_prefill_tokens":256}"#; + let decision: TokenDecisionResponse = serde_json::from_str(json).unwrap(); + assert_eq!(decision.outcome, DecisionOutcome::Remote); + assert_eq!(decision.selected_dc, Some(2)); + assert!( + serde_json::from_str::(&format!( + "{}", + json.replace('}', ",\"extra\":1}") + )) + .is_err() + ); + } + + #[test] + fn stable_tie_key_is_repeatable_and_request_specific() { + assert_eq!(stable_hash(b"request"), stable_hash(b"request")); + assert_ne!(stable_hash(b"request"), stable_hash(b"other")); + } +} diff --git a/lib/llm/src/global_routing_envelope.rs b/lib/llm/src/global_routing_envelope.rs index fc72dc1fe822..1af9215b3822 100644 --- a/lib/llm/src/global_routing_envelope.rs +++ b/lib/llm/src/global_routing_envelope.rs @@ -13,8 +13,8 @@ //! //! # Wire format //! -//! ```text -//! env1.. +//! ```json +//! {"version":1,"key_id":"...","payload":"...","signature":"..."} //! ``` //! //! The signature covers the exact payload bytes as transmitted. Verifiers in @@ -23,23 +23,23 @@ //! removes every cross-language canonicalization hazard by construction. //! //! Downstream of the pod proxy, the execute-exact contract is carried by -//! `request.prompt_token_digest`: the serving frontend recomputes the digest +//! `request.token_ids_sha256`: the serving frontend recomputes the digest //! from its own native preprocessing and rejects any mismatch, so a request //! can never be routed with one token sequence and executed with another. -use std::sync::LazyLock; +use std::sync::{LazyLock, OnceLock}; use prometheus::{IntCounterVec, Opts, Registry}; use serde::{Deserialize, Serialize}; use crate::protocols::TokenIdType; -/// Envelope wire prefix. Bump only with a coordinated verifier rollout. -const WIRE_PREFIX: &str = "env1"; pub const ENVELOPE_VERSION: u32 = 1; +pub const SIGNATURE_DOMAIN: &[u8] = b"morph.global-routing.envelope.v1\0"; +pub const TRUSTED_AUTH_METADATA_KEY: &str = "global-routing-auth-v1"; +pub const SIGNED_ENVELOPE_CONTEXT_KEY: &str = "global_routing.signed_envelope"; -/// Domain separation context for the native prompt token digest. -const TOKEN_DIGEST_CONTEXT: &str = "dynamo/global-routing/prompt-tokens/v1"; +static SIGNING_CONFIG: OnceLock> = OnceLock::new(); #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] @@ -47,14 +47,16 @@ pub struct Envelope { pub version: u32, pub request_id: String, pub decision_id: String, + pub issuer: String, + pub audience: String, /// Unix seconds. pub issued_at: u64, /// Unix seconds. Envelopes are short-lived dispatch instructions. pub expires_at: u64, + pub nonce: String, pub account: Account, pub request: Request, pub routing: Routing, - pub integrity: Integrity, } /// Immutable internal identifiers resolved by global authentication. @@ -64,8 +66,12 @@ pub struct Envelope { pub struct Account { pub api_key_id: String, pub user_id: String, - pub billing_owner_id: String, + pub org_id: Option, + pub billing_tier: String, pub service_tier: String, + pub zero_data_retention: bool, + pub lifetime_requests: Option, + pub service_tier_explicit: bool, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -78,9 +84,9 @@ pub struct Request { /// what the selected region executes. pub normalized_body: serde_json::Value, /// Digest of the exact native token sequence used for routing - /// (see [`prompt_token_digest`]). The serving frontend must reproduce it. - pub prompt_token_digest: String, - pub prompt_token_count: u64, + /// The selected frontend must execute this exact normalized request. + pub body_sha256: String, + pub token_ids_sha256: String, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -90,19 +96,7 @@ pub struct Routing { pub selected_region: String, pub selected_pool: String, pub policy_version: String, - pub ckf_version: String, - pub cache_overlap_blocks: u64, -} - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -#[serde(deny_unknown_fields)] -pub struct Integrity { - /// `sha256:` over this payload's own serialization of - /// `request.normalized_body`. Informational belt-and-braces inside an - /// already-signed payload; the authoritative execution contract is - /// `prompt_token_digest`. - pub body_digest: String, - pub nonce: String, + pub ckf_generation: u64, } #[derive(Debug, thiserror::Error, PartialEq, Eq)] @@ -113,6 +107,8 @@ pub enum EnvelopeError { UnsupportedVersion, #[error("envelope signature is invalid")] BadSignature, + #[error("envelope key identifier is unknown")] + UnknownKey, #[error("envelope is expired or not yet valid")] Expired, #[error("envelope is addressed to region {selected}, not {local}")] @@ -120,51 +116,240 @@ pub enum EnvelopeError { } pub struct EnvelopeSigner { + key_id: String, key: ed25519_dalek::SigningKey, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SignedEnvelope { + pub version: u32, + pub key_id: String, + pub payload: String, + pub signature: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct TrustedAuthMetadata { + pub version: u32, + pub api_key_id: String, + pub user_id: String, + pub org_id: Option, + pub billing_tier: String, + pub service_tier: String, + pub zero_data_retention: bool, + pub lifetime_requests: Option, + pub service_tier_explicit: bool, +} + +impl From for Account { + fn from(value: TrustedAuthMetadata) -> Self { + Self { + api_key_id: value.api_key_id, + user_id: value.user_id, + org_id: value.org_id, + billing_tier: value.billing_tier, + service_tier: value.service_tier, + zero_data_retention: value.zero_data_retention, + lifetime_requests: value.lifetime_requests, + service_tier_explicit: value.service_tier_explicit, + } + } +} + +#[derive(Debug, Clone)] +pub struct SignedRoutingDecision { + pub decision: crate::global_routing::TokenDecisionResponse, + pub envelope: SignedEnvelope, +} + +struct SigningConfig { + issuer: String, + signer: EnvelopeSigner, +} + +impl SigningConfig { + fn from_env() -> Result { + let issuer = required_env("DYN_GLOBAL_ROUTER_ISSUER")?; + let key_id = required_env("DYN_GLOBAL_ROUTER_SIGNING_KEY_ID")?; + let path = required_env("DYN_GLOBAL_ROUTER_SIGNING_KEY_PATH")?; + let pem = std::fs::read_to_string(&path) + .map_err(|error| format!("failed to read global routing signing key: {error}"))?; + let signer = EnvelopeSigner::from_pkcs8_pem(key_id, &pem) + .map_err(|error| format!("failed to parse global routing signing key: {error}"))?; + Ok(Self { issuer, signer }) + } +} + +fn required_env(name: &str) -> Result { + std::env::var(name) + .ok() + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| format!("{name} is required when global routing is enabled")) +} + +#[derive(Debug, thiserror::Error)] +pub enum MintError { + #[error("trusted global routing authentication metadata is missing")] + MissingAuthMetadata, + #[error("trusted global routing authentication metadata is malformed")] + MalformedAuthMetadata, + #[error("trusted global routing authentication metadata is incomplete")] + IncompleteAuthMetadata, + #[error("global routing decision is incomplete")] + IncompleteDecision, + #[error("global routing signing is misconfigured: {0}")] + SigningConfig(String), + #[error("failed to serialize signed internal request: {0}")] + Serialize(#[from] serde_json::Error), +} + +pub fn decode_trusted_auth(value: Option<&str>) -> Result { + use base64::Engine; + + let encoded = value.ok_or(MintError::MissingAuthMetadata)?; + let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(encoded) + .map_err(|_| MintError::MalformedAuthMetadata)?; + let value: serde_json::Value = + serde_json::from_slice(&bytes).map_err(|_| MintError::MalformedAuthMetadata)?; + let object = value.as_object().ok_or(MintError::MalformedAuthMetadata)?; + if !object.contains_key("org_id") || !object.contains_key("lifetime_requests") { + return Err(MintError::MalformedAuthMetadata); + } + let metadata: TrustedAuthMetadata = + serde_json::from_value(value).map_err(|_| MintError::MalformedAuthMetadata)?; + if metadata.version != ENVELOPE_VERSION + || metadata.api_key_id.trim().is_empty() + || metadata.user_id.trim().is_empty() + || metadata + .org_id + .as_ref() + .is_some_and(|value| value.trim().is_empty()) + || metadata.billing_tier.trim().is_empty() + || !matches!(metadata.service_tier.as_str(), "default" | "standby") + { + return Err(MintError::IncompleteAuthMetadata); + } + Ok(metadata) +} + +pub fn mint_signed_routing_decision( + decision: crate::global_routing::TokenDecisionResponse, + auth_metadata: Option<&str>, + request_id: &str, + endpoint: &str, + canonical_model: &str, + normalized_body: serde_json::Value, + token_ids: &[TokenIdType], + now_unix_seconds: u64, +) -> Result { + let selected_region = decision + .selected_region + .as_ref() + .ok_or(MintError::IncompleteDecision)?; + let selected_pool = decision + .selected_pool_id + .as_ref() + .ok_or(MintError::IncompleteDecision)?; + let account = decode_trusted_auth(auth_metadata)?.into(); + let config = SIGNING_CONFIG + .get_or_init(SigningConfig::from_env) + .as_ref() + .map_err(|error| MintError::SigningConfig(error.clone()))?; + let model_audience = canonical_model + .strip_prefix("morph-") + .unwrap_or(canonical_model); + let payload = Envelope { + version: ENVELOPE_VERSION, + request_id: request_id.to_owned(), + decision_id: uuid::Uuid::new_v4().to_string(), + issuer: config.issuer.clone(), + audience: format!("{model_audience}-{selected_region}"), + issued_at: now_unix_seconds, + expires_at: now_unix_seconds.saturating_add(30), + nonce: uuid::Uuid::new_v4().to_string(), + account, + request: Request { + endpoint: endpoint.to_owned(), + canonical_model: canonical_model.to_owned(), + body_sha256: body_digest(&normalized_body)?, + token_ids_sha256: prompt_token_digest(token_ids), + normalized_body, + }, + routing: Routing { + selected_region: selected_region.clone(), + selected_pool: selected_pool.clone(), + policy_version: "v1".to_owned(), + ckf_generation: decision.generation, + }, + }; + let envelope = config.signer.sign(&payload)?; + ENVELOPES.with_label_values(&["success"]).inc(); + Ok(SignedRoutingDecision { decision, envelope }) +} + impl EnvelopeSigner { /// Load a PKCS#8 PEM Ed25519 private key /// (`openssl genpkey -algorithm ed25519`). - pub fn from_pkcs8_pem(pem: &str) -> Result { + pub fn from_pkcs8_pem( + key_id: impl Into, + pem: &str, + ) -> Result { use ed25519_dalek::pkcs8::DecodePrivateKey; Ok(Self { + key_id: key_id.into(), key: ed25519_dalek::SigningKey::from_pkcs8_pem(pem)?, }) } pub fn verifier(&self) -> EnvelopeVerifier { EnvelopeVerifier { - key: self.key.verifying_key(), + keys: [(self.key_id.clone(), self.key.verifying_key())] + .into_iter() + .collect(), } } /// Serialize and sign the envelope into its wire form. - pub fn sign(&self, envelope: &Envelope) -> Result { + pub fn sign(&self, envelope: &Envelope) -> Result { use base64::Engine; use ed25519_dalek::Signer; let engine = &base64::engine::general_purpose::URL_SAFE_NO_PAD; let payload = serde_json::to_vec(envelope)?; - let signature = self.key.sign(&payload); - Ok(format!( - "{WIRE_PREFIX}.{}.{}", - engine.encode(&payload), - engine.encode(signature.to_bytes()) - )) + let mut signed_bytes = Vec::with_capacity(SIGNATURE_DOMAIN.len() + payload.len()); + signed_bytes.extend_from_slice(SIGNATURE_DOMAIN); + signed_bytes.extend_from_slice(&payload); + let signature = self.key.sign(&signed_bytes); + Ok(SignedEnvelope { + version: ENVELOPE_VERSION, + key_id: self.key_id.clone(), + payload: engine.encode(payload), + signature: engine.encode(signature.to_bytes()), + }) } } pub struct EnvelopeVerifier { - key: ed25519_dalek::VerifyingKey, + keys: std::collections::HashMap, } impl EnvelopeVerifier { /// Load a SPKI PEM Ed25519 public key (`openssl pkey -pubout`). - pub fn from_public_key_pem(pem: &str) -> Result { + pub fn from_public_key_pem( + key_id: impl Into, + pem: &str, + ) -> Result { use ed25519_dalek::pkcs8::DecodePublicKey; Ok(Self { - key: ed25519_dalek::VerifyingKey::from_public_key_pem(pem)?, + keys: [( + key_id.into(), + ed25519_dalek::VerifyingKey::from_public_key_pem(pem)?, + )] + .into_iter() + .collect(), }) } @@ -176,27 +361,32 @@ impl EnvelopeVerifier { /// shared state. pub fn verify( &self, - wire: &str, + wire: &SignedEnvelope, local_region: &str, now_unix_seconds: u64, ) -> Result { use base64::Engine; let engine = &base64::engine::general_purpose::URL_SAFE_NO_PAD; - let mut parts = wire.split('.'); - let (Some(WIRE_PREFIX), Some(payload), Some(signature), None) = - (parts.next(), parts.next(), parts.next(), parts.next()) - else { - return Err(EnvelopeError::Malformed); - }; - let payload = engine.decode(payload).map_err(|_| EnvelopeError::Malformed)?; + if wire.version != ENVELOPE_VERSION { + return Err(EnvelopeError::UnsupportedVersion); + } + let key = self + .keys + .get(&wire.key_id) + .ok_or(EnvelopeError::UnknownKey)?; + let payload = engine + .decode(&wire.payload) + .map_err(|_| EnvelopeError::Malformed)?; let signature = engine - .decode(signature) + .decode(&wire.signature) .ok() .and_then(|bytes| ed25519_dalek::Signature::from_slice(&bytes).ok()) .ok_or(EnvelopeError::Malformed)?; - self.key - .verify_strict(&payload, &signature) + let mut signed_bytes = Vec::with_capacity(SIGNATURE_DOMAIN.len() + payload.len()); + signed_bytes.extend_from_slice(SIGNATURE_DOMAIN); + signed_bytes.extend_from_slice(&payload); + key.verify_strict(&signed_bytes, &signature) .map_err(|_| EnvelopeError::BadSignature)?; let envelope: Envelope = @@ -228,9 +418,21 @@ static EXECUTE_EXACT: LazyLock = LazyLock::new(|| { .expect("static metric options are valid") }); +static ENVELOPES: LazyLock = LazyLock::new(|| { + IntCounterVec::new( + Opts::new( + "morph_global_routing_signed_envelopes_total", + "Signed global routing envelope mint outcomes.", + ), + &["outcome"], + ) + .expect("static metric options are valid") +}); + /// Register this module's frontend-side collectors. pub fn ensure_metrics_registered_prometheus(registry: &Registry) -> Result<(), prometheus::Error> { - registry.register(Box::new(EXECUTE_EXACT.clone())) + registry.register(Box::new(EXECUTE_EXACT.clone()))?; + registry.register(Box::new(ENVELOPES.clone())) } /// Enforce the execute-exact contract on the serving frontend. @@ -268,32 +470,72 @@ pub fn enforce_prompt_token_digest( }) } -/// Digest of the exact native token sequence, as `blake3:`. +/// SHA256 digest of the exact native token sequence, as lowercase hex. /// -/// Domain separated and length prefixed; token ids are hashed as little -/// endian `u32`. Both the dispatcher (over the routing tokens) and the -/// serving frontend (over its own preprocessing output) must use this -/// function so the execute-exact comparison is meaningful. +/// Token ids are encoded as unsigned big endian `u32`, matching the frozen +/// cross language contract vector. pub fn prompt_token_digest(token_ids: &[TokenIdType]) -> String { - let mut hasher = blake3::Hasher::new_derive_key(TOKEN_DIGEST_CONTEXT); - hasher.update(&(token_ids.len() as u64).to_le_bytes()); + use sha2::Digest; + let mut hasher = sha2::Sha256::new(); for token_id in token_ids { - hasher.update(&token_id.to_le_bytes()); + hasher.update(token_id.to_be_bytes()); } - format!("blake3:{}", hasher.finalize().to_hex()) + hex_digest(hasher.finalize()) } -/// `sha256:` over the payload's serialization of the normalized body. +/// SHA256 over compact JSON with object keys sorted recursively. +/// +/// The payload itself preserves the normalized body's field order. This digest +/// uses a language neutral ordering so Rust and Python can recompute it without +/// relying on their map implementations. pub fn body_digest(normalized_body: &serde_json::Value) -> Result { use sha2::Digest; - let bytes = serde_json::to_vec(normalized_body)?; - let mut hex = String::with_capacity(7 + 64); - hex.push_str("sha256:"); - for byte in sha2::Sha256::digest(&bytes) { + let mut bytes = Vec::new(); + write_canonical_json(normalized_body, &mut bytes)?; + Ok(hex_digest(sha2::Sha256::digest(&bytes))) +} + +fn write_canonical_json( + value: &serde_json::Value, + output: &mut Vec, +) -> Result<(), serde_json::Error> { + match value { + serde_json::Value::Array(values) => { + output.push(b'['); + for (index, value) in values.iter().enumerate() { + if index > 0 { + output.push(b','); + } + write_canonical_json(value, output)?; + } + output.push(b']'); + } + serde_json::Value::Object(values) => { + output.push(b'{'); + let mut entries = values.iter().collect::>(); + entries.sort_unstable_by(|left, right| left.0.cmp(right.0)); + for (index, (key, value)) in entries.into_iter().enumerate() { + if index > 0 { + output.push(b','); + } + serde_json::to_writer(&mut *output, key)?; + output.push(b':'); + write_canonical_json(value, output)?; + } + output.push(b'}'); + } + scalar => serde_json::to_writer(output, scalar)?, + } + Ok(()) +} + +fn hex_digest(bytes: impl AsRef<[u8]>) -> String { + let mut hex = String::with_capacity(64); + for byte in bytes.as_ref() { use std::fmt::Write; write!(&mut hex, "{byte:02x}").expect("writing to a String cannot fail"); } - Ok(hex) + hex } #[cfg(test)] @@ -303,6 +545,7 @@ mod tests { /// Deterministic key for cross-implementation test vectors only. fn test_signer() -> EnvelopeSigner { EnvelopeSigner { + key_id: "test-key-1".into(), key: ed25519_dalek::SigningKey::from_bytes(&[7u8; 32]), } } @@ -317,31 +560,33 @@ mod tests { version: ENVELOPE_VERSION, request_id: "req_1".into(), decision_id: "decision_1".into(), + issuer: "morph-global-router".into(), + audience: "dsv4flash-us-east5-financial".into(), issued_at: 1_787_780_000, expires_at: 1_787_780_030, + nonce: "nonce_1".into(), account: Account { api_key_id: "key_1".into(), user_id: "user_1".into(), - billing_owner_id: "org_1".into(), - service_tier: "normal".into(), + org_id: Some("org_1".into()), + billing_tier: "pro".into(), + service_tier: "default".into(), + zero_data_retention: false, + lifetime_requests: Some(7), + service_tier_explicit: true, }, request: Request { endpoint: "/v1/chat/completions".into(), canonical_model: "morph-dsv4flash".into(), normalized_body: normalized_body.clone(), - prompt_token_digest: prompt_token_digest(&[1, 2, 3, 4]), - prompt_token_count: 4, + body_sha256: body_digest(&normalized_body).unwrap(), + token_ids_sha256: prompt_token_digest(&[1, 2, 3, 4]), }, routing: Routing { selected_region: "us-east5-financial".into(), selected_pool: "pool_1".into(), policy_version: "policy_v1".into(), - ckf_version: "ckf_v1".into(), - cache_overlap_blocks: 3, - }, - integrity: Integrity { - body_digest: body_digest(&normalized_body).unwrap(), - nonce: "nonce_1".into(), + ckf_generation: 42, }, } } @@ -366,10 +611,15 @@ mod tests { let now = 1_787_780_010; // Payload tampering: flip one payload character. - let mut tampered = wire.clone().into_bytes(); - let payload_start = WIRE_PREFIX.len() + 1; - tampered[payload_start] = if tampered[payload_start] == b'A' { b'B' } else { b'A' }; - let tampered = String::from_utf8(tampered).unwrap(); + let mut tampered = wire.clone(); + tampered.payload.replace_range( + ..1, + if tampered.payload.starts_with('A') { + "B" + } else { + "A" + }, + ); assert!(matches!( verifier.verify(&tampered, "us-east5-financial", now), Err(EnvelopeError::BadSignature | EnvelopeError::Malformed) @@ -377,6 +627,7 @@ mod tests { // Signature from a different key. let other = EnvelopeSigner { + key_id: "test-key-1".into(), key: ed25519_dalek::SigningKey::from_bytes(&[8u8; 32]), }; let foreign = other.sign(&envelope()).unwrap(); @@ -404,14 +655,46 @@ mod tests { }) ); - // Garbage. + // Malformed encoding. + let malformed = SignedEnvelope { + version: 1, + key_id: "test-key-1".into(), + payload: "***".into(), + signature: "***".into(), + }; assert_eq!( - verifier.verify("env1.zzz", "us-east5-financial", now), + verifier.verify(&malformed, "us-east5-financial", now), Err(EnvelopeError::Malformed) ); + } + + #[test] + fn matches_cross_language_contract_vector() { + use base64::Engine; + + let engine = &base64::engine::general_purpose::URL_SAFE_NO_PAD; + let seed = engine + .decode("AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8") + .unwrap(); + let signer = EnvelopeSigner { + key_id: "contract-test-key".into(), + key: ed25519_dalek::SigningKey::from_bytes(seed.as_slice().try_into().unwrap()), + }; + let payload = r#"{"version":1,"request_id":"req_contract_1","decision_id":"decision_contract_1","issuer":"morph-global-router","audience":"dsv4flash-us-east5-financial","issued_at":1787780000,"expires_at":1787780030,"nonce":"nonce_contract_1","account":{"api_key_id":"key_1","user_id":"user_1","org_id":"org_1","billing_tier":"pro","service_tier":"default","zero_data_retention":false,"lifetime_requests":7,"service_tier_explicit":true},"request":{"endpoint":"/v1/chat/completions","canonical_model":"morph-dsv4flash","normalized_body":{"model":"morph-dsv4flash","messages":[{"role":"user","content":"hello"}],"stream":true},"body_sha256":"b26a7dbf98b6cda6b6efad40e14948127007a81c1dacc0e18c5c71604a5504d9","token_ids_sha256":"7b0b5ea3ff36958c8e32ccf24b71da9ac68e51d0881bf75e62b837ec9ea6f3a5"},"routing":{"selected_region":"us-east5-financial","selected_pool":"dsv4flash-east","policy_version":"v1","ckf_generation":42}}"#; + let envelope: Envelope = serde_json::from_str(payload).unwrap(); + let signed = signer.sign(&envelope).unwrap(); + assert_eq!(signed.payload, engine.encode(payload.as_bytes())); + assert_eq!( + signed.signature, + "JEEHnV8xOgdVVpdVjgzQ4UjA-DeOeMfT17g-ATCxPEHXJC7KZVfu5hqzZ61WAFgJfiD_husTUfzFmB77vsOiCA" + ); assert_eq!( - verifier.verify("env2.a.b", "us-east5-financial", now), - Err(EnvelopeError::Malformed) + prompt_token_digest(&[1, 2, 3]), + "7b0b5ea3ff36958c8e32ccf24b71da9ac68e51d0881bf75e62b837ec9ea6f3a5" + ); + assert_eq!( + body_digest(&envelope.request.normalized_body).unwrap(), + "b26a7dbf98b6cda6b6efad40e14948127007a81c1dacc0e18c5c71604a5504d9" ); } @@ -423,21 +706,53 @@ mod tests { enforce_prompt_token_digest(Some(&prompt_token_digest(&tokens)), &tokens, "req") .is_ok() ); - let error = - enforce_prompt_token_digest(Some(&prompt_token_digest(&[9])), &tokens, "req") - .unwrap_err(); + let error = enforce_prompt_token_digest(Some(&prompt_token_digest(&[9])), &tokens, "req") + .unwrap_err(); assert_eq!(error.code, 409); } + #[test] + fn trusted_auth_metadata_is_versioned_strict_and_complete() { + use base64::Engine; + + let json = r#"{"version":1,"api_key_id":"key","user_id":"user","org_id":"org","billing_tier":"pro","service_tier":"default","service_tier_explicit":true,"zero_data_retention":false,"lifetime_requests":7}"#; + let encoded = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(json); + let metadata = decode_trusted_auth(Some(&encoded)).unwrap(); + assert_eq!(metadata.org_id.as_deref(), Some("org")); + assert!(metadata.service_tier_explicit); + + let unknown = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(json.replace( + "\"lifetime_requests\":7", + "\"lifetime_requests\":7,\"extra\":true", + )); + assert!(matches!( + decode_trusted_auth(Some(&unknown)), + Err(MintError::MalformedAuthMetadata) + )); + assert!(matches!( + decode_trusted_auth(None), + Err(MintError::MissingAuthMetadata) + )); + } + #[test] fn token_digest_is_order_and_length_sensitive() { - assert_eq!(prompt_token_digest(&[1, 2, 3]), prompt_token_digest(&[1, 2, 3])); - assert_ne!(prompt_token_digest(&[1, 2, 3]), prompt_token_digest(&[3, 2, 1])); - assert_ne!(prompt_token_digest(&[1, 2, 3]), prompt_token_digest(&[1, 2])); - assert!(prompt_token_digest(&[]).starts_with("blake3:")); + assert_eq!( + prompt_token_digest(&[1, 2, 3]), + prompt_token_digest(&[1, 2, 3]) + ); + assert_ne!( + prompt_token_digest(&[1, 2, 3]), + prompt_token_digest(&[3, 2, 1]) + ); + assert_ne!( + prompt_token_digest(&[1, 2, 3]), + prompt_token_digest(&[1, 2]) + ); + assert_eq!(prompt_token_digest(&[]).len(), 64); } - /// The wire form is fully deterministic (struct-order serialization plus + /// The envelope is fully deterministic (struct-order serialization plus /// RFC 8032 deterministic signatures), so cross-implementation verifier /// vectors can be generated by printing `sign` output for a fixed key. #[test] @@ -448,14 +763,13 @@ mod tests { let second = test_signer().sign(&envelope()).unwrap(); assert_eq!(first, second); - let parts: Vec<&str> = first.split('.').collect(); - assert_eq!(parts.len(), 3); - assert_eq!(parts[0], "env1"); + assert_eq!(first.version, ENVELOPE_VERSION); + assert_eq!(first.key_id, "test-key-1"); // Ed25519 signatures are 64 bytes; base64url without padding is 86. - assert_eq!(parts[2].len(), 86); + assert_eq!(first.signature.len(), 86); // The payload is the plain serde serialization of the envelope. let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD - .decode(parts[1]) + .decode(first.payload) .unwrap(); assert_eq!(payload, serde_json::to_vec(&envelope()).unwrap()); } diff --git a/lib/llm/src/global_routing_shadow.rs b/lib/llm/src/global_routing_shadow.rs deleted file mode 100644 index e50c8683e62d..000000000000 --- a/lib/llm/src/global_routing_shadow.rs +++ /dev/null @@ -1,147 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Best effort, bounded observation of native global CKF routing decisions. - -use std::sync::{Arc, OnceLock}; -use std::time::Duration; - -use serde::Serialize; -use tokio::sync::{Semaphore, mpsc}; - -use crate::protocols::TokenIdType; - -const QUEUE_CAPACITY: usize = 1024; -const MAX_IN_FLIGHT: usize = 16; -const REQUEST_TIMEOUT: Duration = Duration::from_secs(1); - -static OBSERVER: OnceLock> = OnceLock::new(); - -#[derive(Clone)] -struct Observer { - sender: mpsc::Sender, - config: Arc, -} - -#[derive(Debug, Serialize)] -struct Observation { - model: String, - role: &'static str, - token_ids: Vec, - block_size: u32, - local_dc: u64, - stable_tie_key: u64, - readiness_max_age_ms: u64, - load_max_age_ms: u64, -} - -#[derive(Debug)] -struct Config { - endpoint: String, - local_dc: u64, - readiness_max_age_ms: u64, - load_max_age_ms: u64, -} - -impl Config { - fn from_env() -> Option { - let endpoint = std::env::var("DYN_GLOBAL_ROUTER_SHADOW_URL").ok()?; - let parse = |name: &str, default: u64| { - std::env::var(name) - .ok() - .map(|value| value.parse::()) - .transpose() - .map(|value| value.unwrap_or(default)) - }; - Some(Self { - endpoint: format!("{}/v1/decisions/tokens", endpoint.trim_end_matches('/')), - local_dc: std::env::var("DYN_GLOBAL_ROUTER_LOCAL_DC_ID") - .ok()? - .parse() - .ok()?, - readiness_max_age_ms: parse("DYN_GLOBAL_ROUTER_READINESS_MAX_AGE_MS", 45_000).ok()?, - load_max_age_ms: parse("DYN_GLOBAL_ROUTER_LOAD_MAX_AGE_MS", 15_000).ok()?, - }) - } -} - -/// Queue a shadow decision after native tokenization. This never waits for network I/O. -pub fn observe( - model: &str, - token_ids: &[TokenIdType], - block_size: usize, - request_id: &str, - has_multimodal_data: bool, -) { - if block_size == 0 || token_ids.len() < block_size || has_multimodal_data { - return; - } - let Some(observer) = OBSERVER.get_or_init(Observer::from_env) else { - return; - }; - let Ok(block_size) = u32::try_from(block_size) else { - return; - }; - let _ = observer.sender.try_send(Observation { - model: model.to_owned(), - role: "aggregated", - token_ids: token_ids.to_vec(), - block_size, - local_dc: observer.config.local_dc, - stable_tie_key: stable_hash(request_id.as_bytes()), - readiness_max_age_ms: observer.config.readiness_max_age_ms, - load_max_age_ms: observer.config.load_max_age_ms, - }); -} - -impl Observer { - fn from_env() -> Option { - let config = Arc::new(Config::from_env()?); - let client = reqwest::Client::builder() - .timeout(REQUEST_TIMEOUT) - .build() - .ok()?; - let (sender, mut receiver) = mpsc::channel::(QUEUE_CAPACITY); - let permits = Arc::new(Semaphore::new(MAX_IN_FLIGHT)); - let worker_config = Arc::clone(&config); - tokio::spawn(async move { - while let Some(observation) = receiver.recv().await { - let Ok(permit) = Arc::clone(&permits).acquire_owned().await else { - break; - }; - let client = client.clone(); - let endpoint = worker_config.endpoint.clone(); - tokio::spawn(async move { - let _permit = permit; - if let Err(error) = client - .post(endpoint) - .json(&observation) - .send() - .await - .and_then(reqwest::Response::error_for_status) - { - tracing::debug!(%error, "global CKF shadow observation failed"); - } - }); - } - }); - Some(Self { sender, config }) - } -} - -fn stable_hash(bytes: &[u8]) -> u64 { - bytes.iter().fold(0xcbf29ce484222325, |hash, byte| { - (hash ^ u64::from(*byte)).wrapping_mul(0x100000001b3) - }) -} - -#[cfg(test)] -mod tests { - use super::stable_hash; - - #[test] - fn stable_hash_is_repeatable_and_request_specific() { - assert_eq!(stable_hash(b"request"), stable_hash(b"request")); - assert_ne!(stable_hash(b"request"), stable_hash(b"other")); - } -} diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index 37b5c066358a..fc46c2fbef99 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -1217,6 +1217,9 @@ impl HttpServiceConfigBuilder { { tracing::warn!("Failed to register global routing envelope metrics: {}", e); } + if let Err(e) = crate::global_routing::ensure_metrics_registered_prometheus(®istry) { + tracing::warn!("Failed to register global routing decision metrics: {}", e); + } if let Err(e) = ensure_tokio_perf_metrics_registered_prometheus(®istry) { tracing::warn!("Failed to register tokio perf metrics: {}", e); } diff --git a/lib/llm/src/lib.rs b/lib/llm/src/lib.rs index 49f79d230e28..eb9f0a3f4a57 100644 --- a/lib/llm/src/lib.rs +++ b/lib/llm/src/lib.rs @@ -17,8 +17,8 @@ pub mod first_token; pub mod fpm_publisher; pub mod fpm_trace; pub mod frontend_config; +pub mod global_routing; pub mod global_routing_envelope; -pub mod global_routing_shadow; pub mod grpc; pub mod http; pub mod hub; diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 134c79d5723e..0f5e6c34251e 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -5386,7 +5386,7 @@ impl >, ) -> Result>, Error> { // unpack the request - let (mut request, context) = request.into_parts(); + let (mut request, mut context) = request.into_parts(); // Preserve original inbound streaming flag before any internal overrides let request_id = context.id().to_string(); @@ -5466,20 +5466,58 @@ impl &request_id, )?; - crate::global_routing_shadow::observe( - &common_request.model, - &common_request.token_ids, - self.kv_cache_block_size, - &request_id, - common_request.multi_modal_data.is_some() || common_request.mm_routing_info.is_some(), - ); - let uses_tool_call_structural_tag = self.apply_tool_choice_guided_decoding( &request, &mut common_request, prompt_injected_reasoning, )?; + if let Some(decision) = crate::global_routing::decide( + &common_request.model, + &common_request.token_ids, + self.kv_cache_block_size, + &request_id, + common_request.multi_modal_data.is_some() || common_request.mm_routing_info.is_some(), + ) + .await + .map_err(|error| crate::http::service::error::HttpError { + code: 503, + message: error.to_string(), + })? { + let mut signed_request = request.clone(); + signed_request.inner.stream = Some(original_stream_flag); + let normalized_body = serde_json::to_value(&signed_request).map_err(|error| { + crate::http::service::error::HttpError { + code: 500, + message: format!("failed to serialize normalized request: {error}"), + } + })?; + let signed = crate::global_routing_envelope::mint_signed_routing_decision( + decision, + context + .metadata() + .get(crate::global_routing_envelope::TRUSTED_AUTH_METADATA_KEY) + .map(String::as_str), + &request_id, + "/v1/chat/completions", + &common_request.model, + normalized_body, + &common_request.token_ids, + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + ) + .map_err(|error| crate::http::service::error::HttpError { + code: 503, + message: error.to_string(), + })?; + context.insert( + crate::global_routing_envelope::SIGNED_ENVELOPE_CONTEXT_KEY, + signed, + ); + } + tracing::trace!(request = ?common_request, prompt_injected_reasoning, "Pre-processed request"); let trace_state = crate::request_trace::build_request_end_trace_state( &common_request, @@ -5617,7 +5655,7 @@ impl let _stage_guard = StageGuard::new(STAGE_PREPROCESS, ""); // unpack the request - let (mut request, context) = request.into_parts(); + let (mut request, mut context) = request.into_parts(); let request_id = context.id().to_string(); // Preserve original streaming flag @@ -5681,13 +5719,51 @@ impl &request_id, )?; - crate::global_routing_shadow::observe( + if let Some(decision) = crate::global_routing::decide( &common_request.model, &common_request.token_ids, self.kv_cache_block_size, &request_id, common_request.multi_modal_data.is_some() || common_request.mm_routing_info.is_some(), - ); + ) + .await + .map_err(|error| crate::http::service::error::HttpError { + code: 503, + message: error.to_string(), + })? { + let mut signed_request = request.clone(); + signed_request.inner.stream = Some(original_stream_flag); + let normalized_body = serde_json::to_value(&signed_request).map_err(|error| { + crate::http::service::error::HttpError { + code: 500, + message: format!("failed to serialize normalized request: {error}"), + } + })?; + let signed = crate::global_routing_envelope::mint_signed_routing_decision( + decision, + context + .metadata() + .get(crate::global_routing_envelope::TRUSTED_AUTH_METADATA_KEY) + .map(String::as_str), + &request_id, + "/v1/completions", + &common_request.model, + normalized_body, + &common_request.token_ids, + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(), + ) + .map_err(|error| crate::http::service::error::HttpError { + code: 503, + message: error.to_string(), + })?; + context.insert( + crate::global_routing_envelope::SIGNED_ENVELOPE_CONTEXT_KEY, + signed, + ); + } let trace_state = crate::request_trace::build_request_end_trace_state( &common_request, From 044862b8f7b568013d82439f859d406b64f57c57 Mon Sep 17 00:00:00 2001 From: skeptrune Date: Wed, 26 Aug 2026 10:28:07 -0700 Subject: [PATCH 4/8] feat(global-routing): dispatch signed envelopes from the consumer The consumer owns both halves of the routing decision: it evaluates the policy and it ships the signed internal request. POST /v1/dispatch verifies the envelope (signature, version, validity window; never forwarding an envelope it cannot verify), resolves the selected region against --dispatch-target region=https://url entries that must cover every configured relay, forwards the identical wire object over HTTPS/mTLS, and relays the response byte stream back unmodified so client cancellation propagates hop-by-hop. Addressing is deliberately left to the target region's own verifier. Dispatch outcomes surface as global_ckf_consumer_dispatch{es,_rejects,_failures}_total. Claude-Session: https://claude.ai/code/session_01Gq5jeTYZ2SDusjNscnyVF5 Signed-off-by: skeptrune --- Cargo.lock | 2 + components/global-ckf-consumer/Cargo.toml | 4 + components/global-ckf-consumer/src/api.rs | 28 +- components/global-ckf-consumer/src/config.rs | 141 ++++++- .../global-ckf-consumer/src/dispatch.rs | 371 ++++++++++++++++++ components/global-ckf-consumer/src/lib.rs | 1 + components/global-ckf-consumer/src/main.rs | 7 +- lib/llm/src/global_routing_envelope.rs | 34 +- 8 files changed, 574 insertions(+), 14 deletions(-) create mode 100644 components/global-ckf-consumer/src/dispatch.rs diff --git a/Cargo.lock b/Cargo.lock index 2436a7857a74..bd889cf1b838 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3968,11 +3968,13 @@ dependencies = [ "clap", "dynamo-kv-router", "dynamo-llm", + "reqwest 0.12.28", "serde", "serde_json", "thiserror 2.0.18", "tokio", "tonic 0.13.1", + "tower 0.5.3", "tracing", "tracing-subscriber", ] diff --git a/components/global-ckf-consumer/Cargo.toml b/components/global-ckf-consumer/Cargo.toml index 04badbb1db3c..ab552bc3cd3c 100644 --- a/components/global-ckf-consumer/Cargo.toml +++ b/components/global-ckf-consumer/Cargo.toml @@ -11,6 +11,7 @@ axum.workspace = true clap.workspace = true dynamo-llm = { workspace = true, default-features = false, features = ["kv-dc-relay-proto"] } dynamo-kv-router.workspace = true +reqwest.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true @@ -18,3 +19,6 @@ tokio.workspace = true tonic = { version = "0.13.1", default-features = false, features = ["channel", "codegen", "prost", "tls-ring", "zstd"] } tracing.workspace = true tracing-subscriber.workspace = true + +[dev-dependencies] +tower = { version = "0.5", features = ["util"] } diff --git a/components/global-ckf-consumer/src/api.rs b/components/global-ckf-consumer/src/api.rs index d0d4ee9e7627..c1348fcc2c49 100644 --- a/components/global-ckf-consumer/src/api.rs +++ b/components/global-ckf-consumer/src/api.rs @@ -75,6 +75,23 @@ pub struct Metrics { decisions_remote: AtomicU64, decisions_none: AtomicU64, decision_errors: AtomicU64, + dispatches: AtomicU64, + dispatch_rejects: AtomicU64, + dispatch_failures: AtomicU64, +} + +impl Metrics { + pub(crate) fn note_dispatch(&self) { + self.dispatches.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn note_dispatch_reject(&self) { + self.dispatch_rejects.fetch_add(1, Ordering::Relaxed); + } + + pub(crate) fn note_dispatch_failure(&self) { + self.dispatch_failures.fetch_add(1, Ordering::Relaxed); + } } #[derive(Clone, Default)] @@ -805,7 +822,13 @@ async fn metrics(State(state): State) -> String { "global_ckf_consumer_decisions_total{{outcome=\"remote\"}} {}\n", "global_ckf_consumer_decisions_total{{outcome=\"none\"}} {}\n", "# TYPE global_ckf_consumer_decision_errors_total counter\n", - "global_ckf_consumer_decision_errors_total {}\n" + "global_ckf_consumer_decision_errors_total {}\n", + "# TYPE global_ckf_consumer_dispatches_total counter\n", + "global_ckf_consumer_dispatches_total {}\n", + "# TYPE global_ckf_consumer_dispatch_rejects_total counter\n", + "global_ckf_consumer_dispatch_rejects_total {}\n", + "# TYPE global_ckf_consumer_dispatch_failures_total counter\n", + "global_ckf_consumer_dispatch_failures_total {}\n" ), ready, state.metrics.queries.load(Ordering::Relaxed), @@ -816,6 +839,9 @@ async fn metrics(State(state): State) -> String { state.metrics.decisions_remote.load(Ordering::Relaxed), state.metrics.decisions_none.load(Ordering::Relaxed), state.metrics.decision_errors.load(Ordering::Relaxed), + state.metrics.dispatches.load(Ordering::Relaxed), + state.metrics.dispatch_rejects.load(Ordering::Relaxed), + state.metrics.dispatch_failures.load(Ordering::Relaxed), ) } diff --git a/components/global-ckf-consumer/src/config.rs b/components/global-ckf-consumer/src/config.rs index 91f19c0fadbe..57ac7c6d4eb0 100644 --- a/components/global-ckf-consumer/src/config.rs +++ b/components/global-ckf-consumer/src/config.rs @@ -31,6 +31,27 @@ pub struct Config { pub freshness_timeout_seconds: u64, #[arg(long, default_value_t = 16_384)] pub max_query_blocks: usize, + /// Region-keyed dispatch targets, `region=https://url`. Setting any + /// target turns on the dispatch endpoint; the region keys must match the + /// configured relay names so every routable decision has a destination. + #[arg(long = "dispatch-target", value_parser = DispatchTarget::from_str)] + pub dispatch_targets: Vec, + /// SPKI PEM Ed25519 public key used to verify envelopes before dispatch. + #[arg(long)] + pub dispatch_envelope_public_key: Option, + /// Key identifier the verified envelopes must carry. + #[arg(long)] + pub dispatch_envelope_key_id: Option, + /// Optional client certificate and key (PEM) for mTLS toward targets. + #[arg(long, requires = "dispatch_client_key")] + pub dispatch_client_cert: Option, + #[arg(long, requires = "dispatch_client_cert")] + pub dispatch_client_key: Option, + /// Optional additional CA bundle (PEM) trusted for target certificates. + #[arg(long)] + pub dispatch_ca: Option, + #[arg(long, default_value_t = 5)] + pub dispatch_connect_timeout_seconds: u64, } impl Config { @@ -55,10 +76,70 @@ impl Config { bail!("expected DC ID {} is duplicated", relay.expected_dc_id); } } + if !self.dispatch_targets.is_empty() { + if self.dispatch_envelope_public_key.is_none() + || self.dispatch_envelope_key_id.is_none() + { + bail!( + "dispatch targets require --dispatch-envelope-public-key and \ + --dispatch-envelope-key-id; the dispatcher never forwards an \ + envelope it cannot verify" + ); + } + if self.dispatch_connect_timeout_seconds == 0 { + bail!("dispatch connect timeout must be greater than zero"); + } + let mut regions = HashSet::new(); + for target in &self.dispatch_targets { + if !regions.insert(target.region.as_str()) { + bail!("dispatch region {:?} is duplicated", target.region); + } + if !names.contains(target.region.as_str()) { + bail!( + "dispatch region {:?} does not match any configured relay name", + target.region + ); + } + } + for relay in &self.relays { + if !regions.contains(relay.name.as_str()) { + bail!( + "relay {:?} has no dispatch target; every routable decision \ + needs a destination", + relay.name + ); + } + } + } Ok(()) } } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DispatchTarget { + pub region: String, + pub url: String, +} + +impl FromStr for DispatchTarget { + type Err = anyhow::Error; + + fn from_str(value: &str) -> Result { + let (region, url) = value + .split_once('=') + .context("dispatch target must be region=https://url")?; + validate_text("dispatch region", region)?; + validate_text("dispatch URL", url)?; + if !url.starts_with("https://") { + bail!("dispatch target URL must use HTTPS"); + } + Ok(Self { + region: region.to_string(), + url: url.to_string(), + }) + } +} + #[derive(Debug, Clone, PartialEq, Eq)] pub struct RelayConfig { pub name: String, @@ -130,11 +211,9 @@ mod tests { assert!(RelayConfig::from_str("ue5,relay:4443,relay").is_err()); } - #[test] - fn config_rejects_duplicate_names_and_dc_ids() { - let relay = RelayConfig::from_str("ue5,relay:4443,relay,17").unwrap(); - let mut config = Config { - relays: vec![relay.clone(), relay], + fn config(relays: Vec) -> Config { + Config { + relays, tls_cert: "cert".into(), tls_key: "key".into(), tls_ca: "ca".into(), @@ -144,9 +223,59 @@ mod tests { subscriber_id: "consumer".into(), freshness_timeout_seconds: 45, max_query_blocks: 16_384, - }; + dispatch_targets: vec![], + dispatch_envelope_public_key: None, + dispatch_envelope_key_id: None, + dispatch_client_cert: None, + dispatch_client_key: None, + dispatch_ca: None, + dispatch_connect_timeout_seconds: 5, + } + } + + #[test] + fn config_rejects_duplicate_names_and_dc_ids() { + let relay = RelayConfig::from_str("ue5,relay:4443,relay,17").unwrap(); + let mut config = config(vec![relay.clone(), relay]); assert!(config.validate().is_err()); config.relays[1].name = "uw1".into(); assert!(config.validate().is_err()); } + + #[test] + fn dispatch_target_parses_and_requires_https() { + let target = DispatchTarget::from_str("ue5=https://pod-proxy.example/global").unwrap(); + assert_eq!(target.region, "ue5"); + assert_eq!(target.url, "https://pod-proxy.example/global"); + assert!(DispatchTarget::from_str("ue5=http://pod-proxy.example").is_err()); + assert!(DispatchTarget::from_str("https://pod-proxy.example").is_err()); + } + + #[test] + fn dispatch_configuration_is_all_or_nothing_and_covers_every_relay() { + let ue5 = RelayConfig::from_str("ue5,relay:4443,relay,17").unwrap(); + let uw1 = RelayConfig::from_str("uw1,relay2:4443,relay2,18").unwrap(); + let mut config = config(vec![ue5, uw1]); + assert!(config.validate().is_ok()); + + // Targets without a verification key are rejected. + config.dispatch_targets = vec![ + DispatchTarget::from_str("ue5=https://ue5.example/global").unwrap(), + DispatchTarget::from_str("uw1=https://uw1.example/global").unwrap(), + ]; + assert!(config.validate().is_err()); + + config.dispatch_envelope_public_key = Some("key.pem".into()); + config.dispatch_envelope_key_id = Some("2026-08".into()); + assert!(config.validate().is_ok()); + + // A relay without a destination or an unknown region is rejected. + let missing = config.dispatch_targets.split_off(1); + assert!(config.validate().is_err()); + config.dispatch_targets.extend(missing); + config + .dispatch_targets + .push(DispatchTarget::from_str("nowhere=https://n.example").unwrap()); + assert!(config.validate().is_err()); + } } diff --git a/components/global-ckf-consumer/src/dispatch.rs b/components/global-ckf-consumer/src/dispatch.rs new file mode 100644 index 000000000000..6fb874e038ff --- /dev/null +++ b/components/global-ckf-consumer/src/dispatch.rs @@ -0,0 +1,371 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Signed-envelope dispatcher. +//! +//! The consumer owns both halves of the routing decision: it evaluates the +//! policy and it ships the resulting signed internal request to the selected +//! region's pod proxy. `POST /v1/dispatch` accepts a [`SignedEnvelope`], +//! verifies it (signature, version, validity window — never forwarding an +//! envelope it cannot verify), resolves `routing.selected_region` against the +//! configured targets, forwards the identical wire object over HTTPS/mTLS, +//! and relays the response byte stream back unmodified. Dropping the inbound +//! connection drops the upstream request, so client cancellation propagates +//! hop-by-hop through the dispatcher. +//! +//! Addressing is deliberately not enforced here: only the target region's own +//! verifier may assert "this envelope is for me". The dispatcher enforces +//! everything else fail-closed. + +use std::collections::HashMap; +use std::sync::Arc; +use std::time::SystemTime; + +use anyhow::Context; +use axum::body::Body; +use axum::extract::State; +use axum::http::StatusCode; +use axum::response::{IntoResponse, Response}; +use axum::routing::post; +use axum::{Json, Router}; +use dynamo_llm::global_routing_envelope::{EnvelopeError, EnvelopeVerifier, SignedEnvelope}; + +use crate::api::AppState; +use crate::config::Config; + +/// Response headers relayed from the target pod proxy. Everything else is +/// hop metadata the dispatcher must not impersonate. +const RELAYED_HEADERS: [&str; 5] = [ + "content-type", + "cache-control", + "x-accel-buffering", + "x-request-id", + "retry-after", +]; + +pub struct Dispatcher { + verifier: EnvelopeVerifier, + targets: HashMap, + client: reqwest::Client, + state: AppState, +} + +impl Dispatcher { + /// Build the dispatcher when dispatch targets are configured; `None` + /// leaves the consumer as a pure decision service. + pub fn from_config(config: &Config, state: AppState) -> anyhow::Result>> { + if config.dispatch_targets.is_empty() { + return Ok(None); + } + let key_path = config + .dispatch_envelope_public_key + .as_ref() + .context("dispatch requires --dispatch-envelope-public-key")?; + let key_id = config + .dispatch_envelope_key_id + .as_ref() + .context("dispatch requires --dispatch-envelope-key-id")?; + let pem = std::fs::read_to_string(key_path) + .with_context(|| format!("cannot read {}", key_path.display()))?; + let verifier = EnvelopeVerifier::from_public_key_pem(key_id.clone(), &pem) + .context("invalid dispatch envelope public key")?; + + let mut builder = reqwest::Client::builder() + .use_rustls_tls() + .connect_timeout(std::time::Duration::from_secs( + config.dispatch_connect_timeout_seconds, + )); + if let (Some(cert), Some(key)) = (&config.dispatch_client_cert, &config.dispatch_client_key) + { + let mut identity = std::fs::read(cert) + .with_context(|| format!("cannot read {}", cert.display()))?; + identity.extend( + std::fs::read(key).with_context(|| format!("cannot read {}", key.display()))?, + ); + builder = builder.identity( + reqwest::Identity::from_pem(&identity) + .context("invalid dispatch client certificate or key")?, + ); + } + if let Some(ca) = &config.dispatch_ca { + let bundle = + std::fs::read(ca).with_context(|| format!("cannot read {}", ca.display()))?; + for certificate in reqwest::Certificate::from_pem_bundle(&bundle) + .context("invalid dispatch CA bundle")? + { + builder = builder.add_root_certificate(certificate); + } + } + Ok(Some(Arc::new(Self { + verifier, + targets: config + .dispatch_targets + .iter() + .map(|target| (target.region.clone(), target.url.clone())) + .collect(), + client: builder.build().context("cannot build dispatch client")?, + state, + }))) + } + + #[cfg(test)] + fn for_tests( + verifier: EnvelopeVerifier, + targets: HashMap, + state: AppState, + ) -> Arc { + Arc::new(Self { + verifier, + targets, + client: reqwest::Client::new(), + state, + }) + } +} + +pub fn dispatch_router(dispatcher: Arc) -> Router { + Router::new() + .route("/v1/dispatch", post(dispatch)) + .with_state(dispatcher) +} + +async fn dispatch( + State(dispatcher): State>, + Json(wire): Json, +) -> Response { + let now = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + let envelope = match dispatcher.verifier.verify_addressed(&wire, now) { + Ok(envelope) => envelope, + Err(error) => { + dispatcher.state.metrics.note_dispatch_reject(); + let status = match error { + EnvelopeError::Malformed | EnvelopeError::UnsupportedVersion => { + StatusCode::BAD_REQUEST + } + EnvelopeError::BadSignature + | EnvelopeError::UnknownKey + | EnvelopeError::Expired + | EnvelopeError::WrongRegion { .. } => StatusCode::FORBIDDEN, + }; + return reject(status, error.to_string()); + } + }; + let Some(url) = dispatcher.targets.get(&envelope.routing.selected_region) else { + dispatcher.state.metrics.note_dispatch_reject(); + return reject( + StatusCode::NOT_FOUND, + format!( + "no dispatch target for region {:?}", + envelope.routing.selected_region + ), + ); + }; + match dispatcher.client.post(url).json(&wire).send().await { + Ok(upstream) => { + dispatcher.state.metrics.note_dispatch(); + relay(upstream) + } + Err(error) => { + dispatcher.state.metrics.note_dispatch_failure(); + tracing::warn!( + %error, + region = %envelope.routing.selected_region, + request_id = %envelope.request_id, + "dispatch to selected region failed" + ); + reject( + StatusCode::BAD_GATEWAY, + "dispatch to the selected region failed".to_string(), + ) + } + } +} + +/// Relay the target's status, allow-listed headers, and body byte stream. +fn relay(upstream: reqwest::Response) -> Response { + let mut builder = Response::builder().status(upstream.status().as_u16()); + for name in RELAYED_HEADERS { + if let Some(value) = upstream.headers().get(name) { + builder = builder.header(name, value); + } + } + builder + .body(Body::from_stream(upstream.bytes_stream())) + .unwrap_or_else(|error| { + tracing::error!(%error, "failed to assemble relayed response"); + StatusCode::BAD_GATEWAY.into_response() + }) +} + +fn reject(status: StatusCode, message: String) -> Response { + (status, Json(serde_json::json!({ "error": message }))).into_response() +} + +#[cfg(test)] +mod tests { + use super::*; + use dynamo_llm::global_routing_envelope::{ + Account, ENVELOPE_VERSION, Envelope, EnvelopeSigner, Request as EnvelopeRequest, Routing, + prompt_token_digest, + }; + + fn signer() -> EnvelopeSigner { + EnvelopeSigner::from_key_bytes("test-key", &[7u8; 32]) + } + + fn now_unix() -> u64 { + SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap() + .as_secs() + } + + /// A currently-valid envelope; tests mutate the window to expire it. + fn envelope(selected_region: &str) -> Envelope { + Envelope { + version: ENVELOPE_VERSION, + request_id: "req_1".into(), + decision_id: "decision_1".into(), + issuer: "consumer-test".into(), + audience: format!("dsv4flash-{selected_region}"), + issued_at: now_unix().saturating_sub(1), + expires_at: now_unix() + 30, + nonce: "nonce_1".into(), + account: Account { + api_key_id: "key_1".into(), + user_id: "user_1".into(), + org_id: None, + billing_tier: "usage".into(), + service_tier: "default".into(), + zero_data_retention: false, + lifetime_requests: None, + service_tier_explicit: false, + }, + request: EnvelopeRequest { + endpoint: "/v1/chat/completions".into(), + canonical_model: "morph-dsv4flash".into(), + normalized_body: serde_json::json!({"model": "morph-dsv4flash"}), + body_sha256: "sha256:0".into(), + token_ids_sha256: prompt_token_digest(&[1, 2, 3]), + }, + routing: Routing { + selected_region: selected_region.into(), + selected_pool: "pool_1".into(), + policy_version: "v1".into(), + ckf_generation: 7, + }, + } + } + + async fn spawn_target( + response: &'static str, + ) -> (String, tokio::sync::mpsc::UnboundedReceiver) { + let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); + let app = Router::new().route( + "/global", + post(move |Json(wire): Json| { + let sender = sender.clone(); + async move { + let _ = sender.send(wire); + ( + [("content-type", "text/event-stream"), ("x-request-id", "req_1")], + response, + ) + } + }), + ); + 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, app).await.unwrap(); + }); + (format!("http://{address}/global"), receiver) + } + + async fn call( + dispatcher: Arc, + wire: &SignedEnvelope, + ) -> (StatusCode, axum::http::HeaderMap, String) { + use tower::ServiceExt; + + let response = dispatch_router(dispatcher) + .oneshot( + axum::http::Request::builder() + .method("POST") + .uri("/v1/dispatch") + .header("content-type", "application/json") + .body(Body::from(serde_json::to_vec(wire).unwrap())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let headers = response.headers().clone(); + let body = axum::body::to_bytes(response.into_body(), 1 << 20) + .await + .unwrap(); + (status, headers, String::from_utf8(body.to_vec()).unwrap()) + } + + #[tokio::test] + async fn verified_envelopes_are_forwarded_verbatim_and_streamed_back() { + let signer = signer(); + let (url, mut received) = spawn_target("data: {\"ok\":true}\n\ndata: [DONE]\n\n").await; + let dispatcher = Dispatcher::for_tests( + signer.verifier(), + HashMap::from([("us-east5-financial".to_string(), url)]), + AppState::default(), + ); + let wire = signer.sign(&envelope("us-east5-financial")).unwrap(); + + let (status, headers, body) = call(dispatcher, &wire).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(headers["content-type"], "text/event-stream"); + assert_eq!(headers["x-request-id"], "req_1"); + assert!(body.contains("[DONE]")); + // The target received the identical signed wire object. + assert_eq!(received.recv().await.unwrap(), wire); + } + + #[tokio::test] + async fn unverifiable_and_unroutable_envelopes_are_rejected() { + let signer = signer(); + let dispatcher = Dispatcher::for_tests( + signer.verifier(), + HashMap::from([( + "us-east5-financial".to_string(), + "http://127.0.0.1:9/never".to_string(), + )]), + AppState::default(), + ); + + // Foreign signature. + let foreign = EnvelopeSigner::from_key_bytes("test-key", &[9u8; 32]) + .sign(&envelope("us-east5-financial")) + .unwrap(); + let (status, _, _) = call(dispatcher.clone(), &foreign).await; + assert_eq!(status, StatusCode::FORBIDDEN); + + // Expired. + let mut stale = envelope("us-east5-financial"); + stale.issued_at = now_unix().saturating_sub(120); + stale.expires_at = now_unix().saturating_sub(60); + let expired = signer.sign(&stale).unwrap(); + let (status, _, _) = call(dispatcher.clone(), &expired).await; + assert_eq!(status, StatusCode::FORBIDDEN); + + // Valid but addressed to a region with no target. + let unroutable = signer.sign(&envelope("nowhere")).unwrap(); + let (status, _, body) = call(dispatcher.clone(), &unroutable).await; + assert_eq!(status, StatusCode::NOT_FOUND); + assert!(body.contains("nowhere")); + + // Valid and routable, but the target is unreachable. + let unreachable = signer.sign(&envelope("us-east5-financial")).unwrap(); + let (status, _, _) = call(dispatcher, &unreachable).await; + assert_eq!(status, StatusCode::BAD_GATEWAY); + } +} diff --git a/components/global-ckf-consumer/src/lib.rs b/components/global-ckf-consumer/src/lib.rs index fff9dcd122e0..97bb8abe856a 100644 --- a/components/global-ckf-consumer/src/lib.rs +++ b/components/global-ckf-consumer/src/lib.rs @@ -5,6 +5,7 @@ pub mod api; pub mod config; pub mod contract; pub mod coordinator; +pub mod dispatch; pub mod lane; pub mod policy; pub mod pool; diff --git a/components/global-ckf-consumer/src/main.rs b/components/global-ckf-consumer/src/main.rs index 28a0f3efaa26..8f797c826304 100644 --- a/components/global-ckf-consumer/src/main.rs +++ b/components/global-ckf-consumer/src/main.rs @@ -6,6 +6,7 @@ use clap::Parser; use global_ckf_consumer::api::{AppState, api_router, system_router}; use global_ckf_consumer::config::Config; use global_ckf_consumer::coordinator; +use global_ckf_consumer::dispatch::{Dispatcher, dispatch_router}; use global_ckf_consumer::supervisor::spawn_relay_supervisors; #[tokio::main] @@ -20,11 +21,15 @@ async fn main() -> Result<()> { let events = spawn_relay_supervisors(&config); let state = AppState::default(); + let mut app = api_router(state.clone(), config.max_query_blocks); + if let Some(dispatcher) = Dispatcher::from_config(&config, state.clone())? { + app = app.merge(dispatch_router(dispatcher)); + } let api = tokio::net::TcpListener::bind(config.listen_address).await?; let system = tokio::net::TcpListener::bind(config.metrics_listen_address).await?; tokio::select! { result = coordinator::run(config.clone(), state.clone(), events) => result?, - result = axum::serve(api, api_router(state.clone(), config.max_query_blocks)) => result?, + result = axum::serve(api, app) => result?, result = axum::serve(system, system_router(state)) => result?, _ = tokio::signal::ctrl_c() => {} } diff --git a/lib/llm/src/global_routing_envelope.rs b/lib/llm/src/global_routing_envelope.rs index 1af9215b3822..8718868da721 100644 --- a/lib/llm/src/global_routing_envelope.rs +++ b/lib/llm/src/global_routing_envelope.rs @@ -304,6 +304,15 @@ impl EnvelopeSigner { }) } + /// Build a signer from raw Ed25519 seed bytes, for keys sourced from a + /// secret manager rather than a PEM file. + pub fn from_key_bytes(key_id: impl Into, seed: &[u8; 32]) -> Self { + Self { + key_id: key_id.into(), + key: ed25519_dalek::SigningKey::from_bytes(seed), + } + } + pub fn verifier(&self) -> EnvelopeVerifier { EnvelopeVerifier { keys: [(self.key_id.clone(), self.key.verifying_key())] @@ -364,6 +373,25 @@ impl EnvelopeVerifier { wire: &SignedEnvelope, local_region: &str, now_unix_seconds: u64, + ) -> Result { + let envelope = self.verify_addressed(wire, now_unix_seconds)?; + if envelope.routing.selected_region != local_region { + return Err(EnvelopeError::WrongRegion { + selected: envelope.routing.selected_region, + local: local_region.to_owned(), + }); + } + Ok(envelope) + } + + /// Verify signature, version, and validity window without asserting the + /// verifier is the addressed region. This is the dispatcher's check: it + /// forwards to whatever region the envelope selects, so only the target + /// region's own verification may enforce addressing. + pub fn verify_addressed( + &self, + wire: &SignedEnvelope, + now_unix_seconds: u64, ) -> Result { use base64::Engine; @@ -397,12 +425,6 @@ impl EnvelopeVerifier { if now_unix_seconds < envelope.issued_at || now_unix_seconds > envelope.expires_at { return Err(EnvelopeError::Expired); } - if envelope.routing.selected_region != local_region { - return Err(EnvelopeError::WrongRegion { - selected: envelope.routing.selected_region, - local: local_region.to_owned(), - }); - } Ok(envelope) } } From 1e1482a7c00483b58139df2ea0d4f79e6e64ee13 Mon Sep 17 00:00:00 2001 From: skeptrune Date: Wed, 26 Aug 2026 10:59:02 -0700 Subject: [PATCH 5/8] refactor(global-routing): move WAN dispatch into frontend Signed-off-by: skeptrune --- Cargo.lock | 2 - components/global-ckf-consumer/Cargo.toml | 4 - components/global-ckf-consumer/src/api.rs | 30 +- components/global-ckf-consumer/src/config.rs | 141 +--- .../global-ckf-consumer/src/dispatch.rs | 371 ----------- components/global-ckf-consumer/src/lib.rs | 1 - components/global-ckf-consumer/src/main.rs | 7 +- lib/llm/src/global_routing_envelope.rs | 34 +- lib/llm/src/global_routing_transport.rs | 621 ++++++++++++++++++ lib/llm/src/http/service/openai.rs | 72 ++ lib/llm/src/lib.rs | 1 + lib/llm/src/preprocessor.rs | 36 + lib/runtime/src/engine.rs | 6 + lib/runtime/src/pipeline/context.rs | 17 + lib/runtime/src/pipeline/registry.rs | 4 + 15 files changed, 771 insertions(+), 576 deletions(-) delete mode 100644 components/global-ckf-consumer/src/dispatch.rs create mode 100644 lib/llm/src/global_routing_transport.rs diff --git a/Cargo.lock b/Cargo.lock index bd889cf1b838..2436a7857a74 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3968,13 +3968,11 @@ dependencies = [ "clap", "dynamo-kv-router", "dynamo-llm", - "reqwest 0.12.28", "serde", "serde_json", "thiserror 2.0.18", "tokio", "tonic 0.13.1", - "tower 0.5.3", "tracing", "tracing-subscriber", ] diff --git a/components/global-ckf-consumer/Cargo.toml b/components/global-ckf-consumer/Cargo.toml index ab552bc3cd3c..04badbb1db3c 100644 --- a/components/global-ckf-consumer/Cargo.toml +++ b/components/global-ckf-consumer/Cargo.toml @@ -11,7 +11,6 @@ axum.workspace = true clap.workspace = true dynamo-llm = { workspace = true, default-features = false, features = ["kv-dc-relay-proto"] } dynamo-kv-router.workspace = true -reqwest.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true @@ -19,6 +18,3 @@ tokio.workspace = true tonic = { version = "0.13.1", default-features = false, features = ["channel", "codegen", "prost", "tls-ring", "zstd"] } tracing.workspace = true tracing-subscriber.workspace = true - -[dev-dependencies] -tower = { version = "0.5", features = ["util"] } diff --git a/components/global-ckf-consumer/src/api.rs b/components/global-ckf-consumer/src/api.rs index c1348fcc2c49..5889eb14e3c3 100644 --- a/components/global-ckf-consumer/src/api.rs +++ b/components/global-ckf-consumer/src/api.rs @@ -75,23 +75,6 @@ pub struct Metrics { decisions_remote: AtomicU64, decisions_none: AtomicU64, decision_errors: AtomicU64, - dispatches: AtomicU64, - dispatch_rejects: AtomicU64, - dispatch_failures: AtomicU64, -} - -impl Metrics { - pub(crate) fn note_dispatch(&self) { - self.dispatches.fetch_add(1, Ordering::Relaxed); - } - - pub(crate) fn note_dispatch_reject(&self) { - self.dispatch_rejects.fetch_add(1, Ordering::Relaxed); - } - - pub(crate) fn note_dispatch_failure(&self) { - self.dispatch_failures.fetch_add(1, Ordering::Relaxed); - } } #[derive(Clone, Default)] @@ -486,8 +469,6 @@ pub(crate) struct EvaluatedDecision { } /// Evaluate the exact routing policy over the current published facts. -/// The single evaluation path for both the decision API and the dispatcher, -/// so their outcomes and decision metrics can never diverge. #[allow(clippy::too_many_arguments)] pub(crate) fn evaluate_decision( state: &AppState, @@ -822,13 +803,7 @@ async fn metrics(State(state): State) -> String { "global_ckf_consumer_decisions_total{{outcome=\"remote\"}} {}\n", "global_ckf_consumer_decisions_total{{outcome=\"none\"}} {}\n", "# TYPE global_ckf_consumer_decision_errors_total counter\n", - "global_ckf_consumer_decision_errors_total {}\n", - "# TYPE global_ckf_consumer_dispatches_total counter\n", - "global_ckf_consumer_dispatches_total {}\n", - "# TYPE global_ckf_consumer_dispatch_rejects_total counter\n", - "global_ckf_consumer_dispatch_rejects_total {}\n", - "# TYPE global_ckf_consumer_dispatch_failures_total counter\n", - "global_ckf_consumer_dispatch_failures_total {}\n" + "global_ckf_consumer_decision_errors_total {}\n" ), ready, state.metrics.queries.load(Ordering::Relaxed), @@ -839,9 +814,6 @@ async fn metrics(State(state): State) -> String { state.metrics.decisions_remote.load(Ordering::Relaxed), state.metrics.decisions_none.load(Ordering::Relaxed), state.metrics.decision_errors.load(Ordering::Relaxed), - state.metrics.dispatches.load(Ordering::Relaxed), - state.metrics.dispatch_rejects.load(Ordering::Relaxed), - state.metrics.dispatch_failures.load(Ordering::Relaxed), ) } diff --git a/components/global-ckf-consumer/src/config.rs b/components/global-ckf-consumer/src/config.rs index 57ac7c6d4eb0..91f19c0fadbe 100644 --- a/components/global-ckf-consumer/src/config.rs +++ b/components/global-ckf-consumer/src/config.rs @@ -31,27 +31,6 @@ pub struct Config { pub freshness_timeout_seconds: u64, #[arg(long, default_value_t = 16_384)] pub max_query_blocks: usize, - /// Region-keyed dispatch targets, `region=https://url`. Setting any - /// target turns on the dispatch endpoint; the region keys must match the - /// configured relay names so every routable decision has a destination. - #[arg(long = "dispatch-target", value_parser = DispatchTarget::from_str)] - pub dispatch_targets: Vec, - /// SPKI PEM Ed25519 public key used to verify envelopes before dispatch. - #[arg(long)] - pub dispatch_envelope_public_key: Option, - /// Key identifier the verified envelopes must carry. - #[arg(long)] - pub dispatch_envelope_key_id: Option, - /// Optional client certificate and key (PEM) for mTLS toward targets. - #[arg(long, requires = "dispatch_client_key")] - pub dispatch_client_cert: Option, - #[arg(long, requires = "dispatch_client_cert")] - pub dispatch_client_key: Option, - /// Optional additional CA bundle (PEM) trusted for target certificates. - #[arg(long)] - pub dispatch_ca: Option, - #[arg(long, default_value_t = 5)] - pub dispatch_connect_timeout_seconds: u64, } impl Config { @@ -76,70 +55,10 @@ impl Config { bail!("expected DC ID {} is duplicated", relay.expected_dc_id); } } - if !self.dispatch_targets.is_empty() { - if self.dispatch_envelope_public_key.is_none() - || self.dispatch_envelope_key_id.is_none() - { - bail!( - "dispatch targets require --dispatch-envelope-public-key and \ - --dispatch-envelope-key-id; the dispatcher never forwards an \ - envelope it cannot verify" - ); - } - if self.dispatch_connect_timeout_seconds == 0 { - bail!("dispatch connect timeout must be greater than zero"); - } - let mut regions = HashSet::new(); - for target in &self.dispatch_targets { - if !regions.insert(target.region.as_str()) { - bail!("dispatch region {:?} is duplicated", target.region); - } - if !names.contains(target.region.as_str()) { - bail!( - "dispatch region {:?} does not match any configured relay name", - target.region - ); - } - } - for relay in &self.relays { - if !regions.contains(relay.name.as_str()) { - bail!( - "relay {:?} has no dispatch target; every routable decision \ - needs a destination", - relay.name - ); - } - } - } Ok(()) } } -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct DispatchTarget { - pub region: String, - pub url: String, -} - -impl FromStr for DispatchTarget { - type Err = anyhow::Error; - - fn from_str(value: &str) -> Result { - let (region, url) = value - .split_once('=') - .context("dispatch target must be region=https://url")?; - validate_text("dispatch region", region)?; - validate_text("dispatch URL", url)?; - if !url.starts_with("https://") { - bail!("dispatch target URL must use HTTPS"); - } - Ok(Self { - region: region.to_string(), - url: url.to_string(), - }) - } -} - #[derive(Debug, Clone, PartialEq, Eq)] pub struct RelayConfig { pub name: String, @@ -211,9 +130,11 @@ mod tests { assert!(RelayConfig::from_str("ue5,relay:4443,relay").is_err()); } - fn config(relays: Vec) -> Config { - Config { - relays, + #[test] + fn config_rejects_duplicate_names_and_dc_ids() { + let relay = RelayConfig::from_str("ue5,relay:4443,relay,17").unwrap(); + let mut config = Config { + relays: vec![relay.clone(), relay], tls_cert: "cert".into(), tls_key: "key".into(), tls_ca: "ca".into(), @@ -223,59 +144,9 @@ mod tests { subscriber_id: "consumer".into(), freshness_timeout_seconds: 45, max_query_blocks: 16_384, - dispatch_targets: vec![], - dispatch_envelope_public_key: None, - dispatch_envelope_key_id: None, - dispatch_client_cert: None, - dispatch_client_key: None, - dispatch_ca: None, - dispatch_connect_timeout_seconds: 5, - } - } - - #[test] - fn config_rejects_duplicate_names_and_dc_ids() { - let relay = RelayConfig::from_str("ue5,relay:4443,relay,17").unwrap(); - let mut config = config(vec![relay.clone(), relay]); + }; assert!(config.validate().is_err()); config.relays[1].name = "uw1".into(); assert!(config.validate().is_err()); } - - #[test] - fn dispatch_target_parses_and_requires_https() { - let target = DispatchTarget::from_str("ue5=https://pod-proxy.example/global").unwrap(); - assert_eq!(target.region, "ue5"); - assert_eq!(target.url, "https://pod-proxy.example/global"); - assert!(DispatchTarget::from_str("ue5=http://pod-proxy.example").is_err()); - assert!(DispatchTarget::from_str("https://pod-proxy.example").is_err()); - } - - #[test] - fn dispatch_configuration_is_all_or_nothing_and_covers_every_relay() { - let ue5 = RelayConfig::from_str("ue5,relay:4443,relay,17").unwrap(); - let uw1 = RelayConfig::from_str("uw1,relay2:4443,relay2,18").unwrap(); - let mut config = config(vec![ue5, uw1]); - assert!(config.validate().is_ok()); - - // Targets without a verification key are rejected. - config.dispatch_targets = vec![ - DispatchTarget::from_str("ue5=https://ue5.example/global").unwrap(), - DispatchTarget::from_str("uw1=https://uw1.example/global").unwrap(), - ]; - assert!(config.validate().is_err()); - - config.dispatch_envelope_public_key = Some("key.pem".into()); - config.dispatch_envelope_key_id = Some("2026-08".into()); - assert!(config.validate().is_ok()); - - // A relay without a destination or an unknown region is rejected. - let missing = config.dispatch_targets.split_off(1); - assert!(config.validate().is_err()); - config.dispatch_targets.extend(missing); - config - .dispatch_targets - .push(DispatchTarget::from_str("nowhere=https://n.example").unwrap()); - assert!(config.validate().is_err()); - } } diff --git a/components/global-ckf-consumer/src/dispatch.rs b/components/global-ckf-consumer/src/dispatch.rs deleted file mode 100644 index 6fb874e038ff..000000000000 --- a/components/global-ckf-consumer/src/dispatch.rs +++ /dev/null @@ -1,371 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Signed-envelope dispatcher. -//! -//! The consumer owns both halves of the routing decision: it evaluates the -//! policy and it ships the resulting signed internal request to the selected -//! region's pod proxy. `POST /v1/dispatch` accepts a [`SignedEnvelope`], -//! verifies it (signature, version, validity window — never forwarding an -//! envelope it cannot verify), resolves `routing.selected_region` against the -//! configured targets, forwards the identical wire object over HTTPS/mTLS, -//! and relays the response byte stream back unmodified. Dropping the inbound -//! connection drops the upstream request, so client cancellation propagates -//! hop-by-hop through the dispatcher. -//! -//! Addressing is deliberately not enforced here: only the target region's own -//! verifier may assert "this envelope is for me". The dispatcher enforces -//! everything else fail-closed. - -use std::collections::HashMap; -use std::sync::Arc; -use std::time::SystemTime; - -use anyhow::Context; -use axum::body::Body; -use axum::extract::State; -use axum::http::StatusCode; -use axum::response::{IntoResponse, Response}; -use axum::routing::post; -use axum::{Json, Router}; -use dynamo_llm::global_routing_envelope::{EnvelopeError, EnvelopeVerifier, SignedEnvelope}; - -use crate::api::AppState; -use crate::config::Config; - -/// Response headers relayed from the target pod proxy. Everything else is -/// hop metadata the dispatcher must not impersonate. -const RELAYED_HEADERS: [&str; 5] = [ - "content-type", - "cache-control", - "x-accel-buffering", - "x-request-id", - "retry-after", -]; - -pub struct Dispatcher { - verifier: EnvelopeVerifier, - targets: HashMap, - client: reqwest::Client, - state: AppState, -} - -impl Dispatcher { - /// Build the dispatcher when dispatch targets are configured; `None` - /// leaves the consumer as a pure decision service. - pub fn from_config(config: &Config, state: AppState) -> anyhow::Result>> { - if config.dispatch_targets.is_empty() { - return Ok(None); - } - let key_path = config - .dispatch_envelope_public_key - .as_ref() - .context("dispatch requires --dispatch-envelope-public-key")?; - let key_id = config - .dispatch_envelope_key_id - .as_ref() - .context("dispatch requires --dispatch-envelope-key-id")?; - let pem = std::fs::read_to_string(key_path) - .with_context(|| format!("cannot read {}", key_path.display()))?; - let verifier = EnvelopeVerifier::from_public_key_pem(key_id.clone(), &pem) - .context("invalid dispatch envelope public key")?; - - let mut builder = reqwest::Client::builder() - .use_rustls_tls() - .connect_timeout(std::time::Duration::from_secs( - config.dispatch_connect_timeout_seconds, - )); - if let (Some(cert), Some(key)) = (&config.dispatch_client_cert, &config.dispatch_client_key) - { - let mut identity = std::fs::read(cert) - .with_context(|| format!("cannot read {}", cert.display()))?; - identity.extend( - std::fs::read(key).with_context(|| format!("cannot read {}", key.display()))?, - ); - builder = builder.identity( - reqwest::Identity::from_pem(&identity) - .context("invalid dispatch client certificate or key")?, - ); - } - if let Some(ca) = &config.dispatch_ca { - let bundle = - std::fs::read(ca).with_context(|| format!("cannot read {}", ca.display()))?; - for certificate in reqwest::Certificate::from_pem_bundle(&bundle) - .context("invalid dispatch CA bundle")? - { - builder = builder.add_root_certificate(certificate); - } - } - Ok(Some(Arc::new(Self { - verifier, - targets: config - .dispatch_targets - .iter() - .map(|target| (target.region.clone(), target.url.clone())) - .collect(), - client: builder.build().context("cannot build dispatch client")?, - state, - }))) - } - - #[cfg(test)] - fn for_tests( - verifier: EnvelopeVerifier, - targets: HashMap, - state: AppState, - ) -> Arc { - Arc::new(Self { - verifier, - targets, - client: reqwest::Client::new(), - state, - }) - } -} - -pub fn dispatch_router(dispatcher: Arc) -> Router { - Router::new() - .route("/v1/dispatch", post(dispatch)) - .with_state(dispatcher) -} - -async fn dispatch( - State(dispatcher): State>, - Json(wire): Json, -) -> Response { - let now = SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(); - let envelope = match dispatcher.verifier.verify_addressed(&wire, now) { - Ok(envelope) => envelope, - Err(error) => { - dispatcher.state.metrics.note_dispatch_reject(); - let status = match error { - EnvelopeError::Malformed | EnvelopeError::UnsupportedVersion => { - StatusCode::BAD_REQUEST - } - EnvelopeError::BadSignature - | EnvelopeError::UnknownKey - | EnvelopeError::Expired - | EnvelopeError::WrongRegion { .. } => StatusCode::FORBIDDEN, - }; - return reject(status, error.to_string()); - } - }; - let Some(url) = dispatcher.targets.get(&envelope.routing.selected_region) else { - dispatcher.state.metrics.note_dispatch_reject(); - return reject( - StatusCode::NOT_FOUND, - format!( - "no dispatch target for region {:?}", - envelope.routing.selected_region - ), - ); - }; - match dispatcher.client.post(url).json(&wire).send().await { - Ok(upstream) => { - dispatcher.state.metrics.note_dispatch(); - relay(upstream) - } - Err(error) => { - dispatcher.state.metrics.note_dispatch_failure(); - tracing::warn!( - %error, - region = %envelope.routing.selected_region, - request_id = %envelope.request_id, - "dispatch to selected region failed" - ); - reject( - StatusCode::BAD_GATEWAY, - "dispatch to the selected region failed".to_string(), - ) - } - } -} - -/// Relay the target's status, allow-listed headers, and body byte stream. -fn relay(upstream: reqwest::Response) -> Response { - let mut builder = Response::builder().status(upstream.status().as_u16()); - for name in RELAYED_HEADERS { - if let Some(value) = upstream.headers().get(name) { - builder = builder.header(name, value); - } - } - builder - .body(Body::from_stream(upstream.bytes_stream())) - .unwrap_or_else(|error| { - tracing::error!(%error, "failed to assemble relayed response"); - StatusCode::BAD_GATEWAY.into_response() - }) -} - -fn reject(status: StatusCode, message: String) -> Response { - (status, Json(serde_json::json!({ "error": message }))).into_response() -} - -#[cfg(test)] -mod tests { - use super::*; - use dynamo_llm::global_routing_envelope::{ - Account, ENVELOPE_VERSION, Envelope, EnvelopeSigner, Request as EnvelopeRequest, Routing, - prompt_token_digest, - }; - - fn signer() -> EnvelopeSigner { - EnvelopeSigner::from_key_bytes("test-key", &[7u8; 32]) - } - - fn now_unix() -> u64 { - SystemTime::now() - .duration_since(SystemTime::UNIX_EPOCH) - .unwrap() - .as_secs() - } - - /// A currently-valid envelope; tests mutate the window to expire it. - fn envelope(selected_region: &str) -> Envelope { - Envelope { - version: ENVELOPE_VERSION, - request_id: "req_1".into(), - decision_id: "decision_1".into(), - issuer: "consumer-test".into(), - audience: format!("dsv4flash-{selected_region}"), - issued_at: now_unix().saturating_sub(1), - expires_at: now_unix() + 30, - nonce: "nonce_1".into(), - account: Account { - api_key_id: "key_1".into(), - user_id: "user_1".into(), - org_id: None, - billing_tier: "usage".into(), - service_tier: "default".into(), - zero_data_retention: false, - lifetime_requests: None, - service_tier_explicit: false, - }, - request: EnvelopeRequest { - endpoint: "/v1/chat/completions".into(), - canonical_model: "morph-dsv4flash".into(), - normalized_body: serde_json::json!({"model": "morph-dsv4flash"}), - body_sha256: "sha256:0".into(), - token_ids_sha256: prompt_token_digest(&[1, 2, 3]), - }, - routing: Routing { - selected_region: selected_region.into(), - selected_pool: "pool_1".into(), - policy_version: "v1".into(), - ckf_generation: 7, - }, - } - } - - async fn spawn_target( - response: &'static str, - ) -> (String, tokio::sync::mpsc::UnboundedReceiver) { - let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); - let app = Router::new().route( - "/global", - post(move |Json(wire): Json| { - let sender = sender.clone(); - async move { - let _ = sender.send(wire); - ( - [("content-type", "text/event-stream"), ("x-request-id", "req_1")], - response, - ) - } - }), - ); - 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, app).await.unwrap(); - }); - (format!("http://{address}/global"), receiver) - } - - async fn call( - dispatcher: Arc, - wire: &SignedEnvelope, - ) -> (StatusCode, axum::http::HeaderMap, String) { - use tower::ServiceExt; - - let response = dispatch_router(dispatcher) - .oneshot( - axum::http::Request::builder() - .method("POST") - .uri("/v1/dispatch") - .header("content-type", "application/json") - .body(Body::from(serde_json::to_vec(wire).unwrap())) - .unwrap(), - ) - .await - .unwrap(); - let status = response.status(); - let headers = response.headers().clone(); - let body = axum::body::to_bytes(response.into_body(), 1 << 20) - .await - .unwrap(); - (status, headers, String::from_utf8(body.to_vec()).unwrap()) - } - - #[tokio::test] - async fn verified_envelopes_are_forwarded_verbatim_and_streamed_back() { - let signer = signer(); - let (url, mut received) = spawn_target("data: {\"ok\":true}\n\ndata: [DONE]\n\n").await; - let dispatcher = Dispatcher::for_tests( - signer.verifier(), - HashMap::from([("us-east5-financial".to_string(), url)]), - AppState::default(), - ); - let wire = signer.sign(&envelope("us-east5-financial")).unwrap(); - - let (status, headers, body) = call(dispatcher, &wire).await; - assert_eq!(status, StatusCode::OK); - assert_eq!(headers["content-type"], "text/event-stream"); - assert_eq!(headers["x-request-id"], "req_1"); - assert!(body.contains("[DONE]")); - // The target received the identical signed wire object. - assert_eq!(received.recv().await.unwrap(), wire); - } - - #[tokio::test] - async fn unverifiable_and_unroutable_envelopes_are_rejected() { - let signer = signer(); - let dispatcher = Dispatcher::for_tests( - signer.verifier(), - HashMap::from([( - "us-east5-financial".to_string(), - "http://127.0.0.1:9/never".to_string(), - )]), - AppState::default(), - ); - - // Foreign signature. - let foreign = EnvelopeSigner::from_key_bytes("test-key", &[9u8; 32]) - .sign(&envelope("us-east5-financial")) - .unwrap(); - let (status, _, _) = call(dispatcher.clone(), &foreign).await; - assert_eq!(status, StatusCode::FORBIDDEN); - - // Expired. - let mut stale = envelope("us-east5-financial"); - stale.issued_at = now_unix().saturating_sub(120); - stale.expires_at = now_unix().saturating_sub(60); - let expired = signer.sign(&stale).unwrap(); - let (status, _, _) = call(dispatcher.clone(), &expired).await; - assert_eq!(status, StatusCode::FORBIDDEN); - - // Valid but addressed to a region with no target. - let unroutable = signer.sign(&envelope("nowhere")).unwrap(); - let (status, _, body) = call(dispatcher.clone(), &unroutable).await; - assert_eq!(status, StatusCode::NOT_FOUND); - assert!(body.contains("nowhere")); - - // Valid and routable, but the target is unreachable. - let unreachable = signer.sign(&envelope("us-east5-financial")).unwrap(); - let (status, _, _) = call(dispatcher, &unreachable).await; - assert_eq!(status, StatusCode::BAD_GATEWAY); - } -} diff --git a/components/global-ckf-consumer/src/lib.rs b/components/global-ckf-consumer/src/lib.rs index 97bb8abe856a..fff9dcd122e0 100644 --- a/components/global-ckf-consumer/src/lib.rs +++ b/components/global-ckf-consumer/src/lib.rs @@ -5,7 +5,6 @@ pub mod api; pub mod config; pub mod contract; pub mod coordinator; -pub mod dispatch; pub mod lane; pub mod policy; pub mod pool; diff --git a/components/global-ckf-consumer/src/main.rs b/components/global-ckf-consumer/src/main.rs index 8f797c826304..28a0f3efaa26 100644 --- a/components/global-ckf-consumer/src/main.rs +++ b/components/global-ckf-consumer/src/main.rs @@ -6,7 +6,6 @@ use clap::Parser; use global_ckf_consumer::api::{AppState, api_router, system_router}; use global_ckf_consumer::config::Config; use global_ckf_consumer::coordinator; -use global_ckf_consumer::dispatch::{Dispatcher, dispatch_router}; use global_ckf_consumer::supervisor::spawn_relay_supervisors; #[tokio::main] @@ -21,15 +20,11 @@ async fn main() -> Result<()> { let events = spawn_relay_supervisors(&config); let state = AppState::default(); - let mut app = api_router(state.clone(), config.max_query_blocks); - if let Some(dispatcher) = Dispatcher::from_config(&config, state.clone())? { - app = app.merge(dispatch_router(dispatcher)); - } let api = tokio::net::TcpListener::bind(config.listen_address).await?; let system = tokio::net::TcpListener::bind(config.metrics_listen_address).await?; tokio::select! { result = coordinator::run(config.clone(), state.clone(), events) => result?, - result = axum::serve(api, app) => result?, + result = axum::serve(api, api_router(state.clone(), config.max_query_blocks)) => result?, result = axum::serve(system, system_router(state)) => result?, _ = tokio::signal::ctrl_c() => {} } diff --git a/lib/llm/src/global_routing_envelope.rs b/lib/llm/src/global_routing_envelope.rs index 8718868da721..1af9215b3822 100644 --- a/lib/llm/src/global_routing_envelope.rs +++ b/lib/llm/src/global_routing_envelope.rs @@ -304,15 +304,6 @@ impl EnvelopeSigner { }) } - /// Build a signer from raw Ed25519 seed bytes, for keys sourced from a - /// secret manager rather than a PEM file. - pub fn from_key_bytes(key_id: impl Into, seed: &[u8; 32]) -> Self { - Self { - key_id: key_id.into(), - key: ed25519_dalek::SigningKey::from_bytes(seed), - } - } - pub fn verifier(&self) -> EnvelopeVerifier { EnvelopeVerifier { keys: [(self.key_id.clone(), self.key.verifying_key())] @@ -373,25 +364,6 @@ impl EnvelopeVerifier { wire: &SignedEnvelope, local_region: &str, now_unix_seconds: u64, - ) -> Result { - let envelope = self.verify_addressed(wire, now_unix_seconds)?; - if envelope.routing.selected_region != local_region { - return Err(EnvelopeError::WrongRegion { - selected: envelope.routing.selected_region, - local: local_region.to_owned(), - }); - } - Ok(envelope) - } - - /// Verify signature, version, and validity window without asserting the - /// verifier is the addressed region. This is the dispatcher's check: it - /// forwards to whatever region the envelope selects, so only the target - /// region's own verification may enforce addressing. - pub fn verify_addressed( - &self, - wire: &SignedEnvelope, - now_unix_seconds: u64, ) -> Result { use base64::Engine; @@ -425,6 +397,12 @@ impl EnvelopeVerifier { if now_unix_seconds < envelope.issued_at || now_unix_seconds > envelope.expires_at { return Err(EnvelopeError::Expired); } + if envelope.routing.selected_region != local_region { + return Err(EnvelopeError::WrongRegion { + selected: envelope.routing.selected_region, + local: local_region.to_owned(), + }); + } Ok(envelope) } } diff --git a/lib/llm/src/global_routing_transport.rs b/lib/llm/src/global_routing_transport.rs new file mode 100644 index 000000000000..e47f035def05 --- /dev/null +++ b/lib/llm/src/global_routing_transport.rs @@ -0,0 +1,621 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Authoritative WAN dispatch to a selected regional pod proxy. +//! +//! A configured global route owns the request. Dispatch is attempted exactly +//! once against the region named by the signed envelope. There is deliberately +//! no local or alternate-region fallback in this module. + +use std::collections::BTreeMap; +use std::pin::Pin; +use std::sync::{Arc, Mutex, OnceLock}; + +use async_trait::async_trait; +use base64::Engine; +use bytes::Bytes; +use futures::{Stream, TryStreamExt}; +use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue}; +use serde::de::{MapAccess, Visitor}; +use url::Url; + +use crate::global_routing_envelope::{ + ENVELOPE_VERSION, Envelope, SignedEnvelope, SignedRoutingDecision, +}; + +const INTERNAL_CHAT_PATH: &str = "/internal/v1/chat/completions"; +const ENV_CERT_PATH: &str = "DYN_GLOBAL_ROUTER_MTLS_CERT_PATH"; +const ENV_KEY_PATH: &str = "DYN_GLOBAL_ROUTER_MTLS_KEY_PATH"; +const ENV_CA_PATH: &str = "DYN_GLOBAL_ROUTER_MTLS_CA_PATH"; +const ENV_ENDPOINTS: &str = "DYN_GLOBAL_ROUTER_REGIONAL_ENDPOINTS_JSON"; + +pub type ResponseBody = Pin> + Send + 'static>>; + +pub struct WanResponse { + pub status: reqwest::StatusCode, + pub headers: HeaderMap, + /// The live upstream body. It is never buffered. Dropping it cancels the + /// HTTP/2 stream and therefore propagates client cancellation upstream. + pub body: ResponseBody, +} + +/// Registry-safe handoff from the preprocessing pipeline to the HTTP source. +/// The mutex is only an ownership cell; response bytes are never buffered. +pub struct WanResponseHandle(Mutex>); + +impl WanResponseHandle { + pub fn new(response: WanResponse) -> Self { + Self(Mutex::new(Some(response))) + } + + pub fn take(&self) -> Option { + self.0.lock().ok()?.take() + } +} + +pub const WAN_RESPONSE_CONTEXT_KEY: &str = "global_routing.wan_response"; + +static DISPATCHER: OnceLock> = OnceLock::new(); + +#[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)] +pub enum DispatchError { + #[error("global routing WAN transport is misconfigured: {0}")] + Config(String), + #[error("signed global routing envelope is malformed")] + MalformedEnvelope, + #[error("signed envelope and authoritative decision select different regions")] + DecisionMismatch, + #[error("no regional pod proxy endpoint is configured for {0}")] + UnknownRegion(String), + #[error("regional pod proxy request failed before a response: {0}")] + BeforeResponse(String), + #[error("regional pod proxy response stream failed: {0}")] + Midstream(String), +} + +#[derive(Clone)] +pub struct RegionalDispatcher { + transport: Arc, +} + +impl RegionalDispatcher { + pub fn from_env() -> Result { + let config = TransportConfig::from_env()?; + Ok(Self { + transport: Arc::new(Http2Transport::new(config)?), + }) + } + + #[cfg(test)] + fn with_transport(transport: Arc) -> Self { + Self { transport } + } + + pub async fn dispatch( + &self, + signed: &SignedRoutingDecision, + ) -> Result { + let payload = decode_payload(&signed.envelope)?; + let decision_region = signed + .decision + .selected_region + .as_deref() + .ok_or(DispatchError::DecisionMismatch)?; + if decision_region != payload.routing.selected_region { + return Err(DispatchError::DecisionMismatch); + } + self.transport + .send( + &payload.routing.selected_region, + &payload.request_id, + &signed.envelope, + ) + .await + } +} + +/// The only non-WAN result means global routing is absent. Once configured, +/// callers receive a WAN response or an error and must never call the local +/// pipeline's `next.generate`. +pub enum DispatchOwnership { + GlobalRoutingDisabled, + Wan(WanResponse), +} + +pub async fn dispatch_authoritatively( + dispatcher: Option<&RegionalDispatcher>, + signed: Option<&SignedRoutingDecision>, +) -> Result { + match dispatcher { + None => Ok(DispatchOwnership::GlobalRoutingDisabled), + Some(dispatcher) => { + let signed = signed.ok_or(DispatchError::MalformedEnvelope)?; + dispatcher + .dispatch(signed) + .await + .map(DispatchOwnership::Wan) + } + } +} + +/// Dispatch a minted request using the process-wide persistent HTTP/2 client. +/// This is called from the preprocessor before its local `next.generate` site. +pub async fn dispatch_signed(signed: &SignedRoutingDecision) -> Result { + let dispatcher = DISPATCHER + .get_or_init(|| RegionalDispatcher::from_env().map_err(|error| error.to_string())) + .as_ref() + .map_err(|error| DispatchError::Config(error.clone()))?; + dispatcher.dispatch(signed).await +} + +#[async_trait] +trait RegionalTransport: Send + Sync { + async fn send( + &self, + region: &str, + request_id: &str, + envelope: &SignedEnvelope, + ) -> Result; +} + +struct Http2Transport { + client: reqwest::Client, + endpoints: BTreeMap, +} + +impl Http2Transport { + fn new(config: TransportConfig) -> Result { + let mut identity_pem = + std::fs::read(&config.cert_path).map_err(|error| config_error(ENV_CERT_PATH, error))?; + identity_pem.push(b'\n'); + identity_pem.extend( + std::fs::read(&config.key_path).map_err(|error| config_error(ENV_KEY_PATH, error))?, + ); + let identity = reqwest::Identity::from_pem(&identity_pem) + .map_err(|error| DispatchError::Config(format!("invalid mTLS identity: {error}")))?; + let ca_pem = + std::fs::read(&config.ca_path).map_err(|error| config_error(ENV_CA_PATH, error))?; + let ca = reqwest::Certificate::from_pem(&ca_pem) + .map_err(|error| DispatchError::Config(format!("invalid mTLS CA: {error}")))?; + let client = reqwest::Client::builder() + .identity(identity) + .add_root_certificate(ca) + .https_only(true) + .http2_prior_knowledge() + .http2_adaptive_window(true) + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|error| { + DispatchError::Config(format!("failed to build regional HTTP/2 client: {error}")) + })?; + Ok(Self { + client, + endpoints: config.endpoints, + }) + } +} + +#[async_trait] +impl RegionalTransport for Http2Transport { + async fn send( + &self, + region: &str, + request_id: &str, + envelope: &SignedEnvelope, + ) -> Result { + let base = self + .endpoints + .get(region) + .ok_or_else(|| DispatchError::UnknownRegion(region.to_owned()))?; + let endpoint = base.join(INTERNAL_CHAT_PATH).map_err(|error| { + DispatchError::Config(format!("invalid endpoint for region {region}: {error}")) + })?; + let request_id = + HeaderValue::from_str(request_id).map_err(|_| DispatchError::MalformedEnvelope)?; + + // One send, one selected region. reqwest does not retry requests. + let response = self + .client + .post(endpoint) + .header(CONTENT_TYPE, "application/json") + .header("x-request-id", request_id) + .header("x-morph-envelope-version", ENVELOPE_VERSION.to_string()) + .json(envelope) + .send() + .await + .map_err(|error| DispatchError::BeforeResponse(error.to_string()))?; + let status = response.status(); + let headers = response.headers().clone(); + let body = response + .bytes_stream() + .map_err(|error| DispatchError::Midstream(error.to_string())); + Ok(WanResponse { + status, + headers, + body: Box::pin(body), + }) + } +} + +#[derive(Debug)] +struct TransportConfig { + cert_path: String, + key_path: String, + ca_path: String, + endpoints: BTreeMap, +} + +impl TransportConfig { + fn from_env() -> Result { + Self::parse( + required_env(ENV_CERT_PATH)?, + required_env(ENV_KEY_PATH)?, + required_env(ENV_CA_PATH)?, + &required_env(ENV_ENDPOINTS)?, + ) + } + + fn parse( + cert_path: String, + key_path: String, + ca_path: String, + endpoints_json: &str, + ) -> Result { + let RawEndpoints(raw) = serde_json::from_str(endpoints_json).map_err(|error| { + DispatchError::Config(format!("{ENV_ENDPOINTS} must be a JSON object: {error}")) + })?; + if raw.is_empty() { + return Err(DispatchError::Config(format!( + "{ENV_ENDPOINTS} must not be empty" + ))); + } + let mut endpoints = BTreeMap::new(); + for (region, raw_url) in raw { + if region.trim().is_empty() || region != region.trim() { + return Err(DispatchError::Config( + "regional endpoint keys must be nonempty exact region names".to_owned(), + )); + } + let url = Url::parse(&raw_url).map_err(|error| { + DispatchError::Config(format!("invalid endpoint for region {region}: {error}")) + })?; + if url.scheme() != "https" + || url.cannot_be_a_base() + || url.host_str().is_none() + || url.query().is_some() + || url.fragment().is_some() + || url.path() != "/" + { + return Err(DispatchError::Config(format!( + "endpoint for region {region} must be an HTTPS base URL without a path, query, or fragment" + ))); + } + endpoints.insert(region, url); + } + Ok(Self { + cert_path, + key_path, + ca_path, + endpoints, + }) + } +} + +struct RawEndpoints(Vec<(String, String)>); + +impl<'de> serde::Deserialize<'de> for RawEndpoints { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + struct EndpointsVisitor; + impl<'de> Visitor<'de> for EndpointsVisitor { + type Value = RawEndpoints; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("an object mapping exact region names to HTTPS base URLs") + } + + fn visit_map(self, mut map: A) -> Result + where + A: MapAccess<'de>, + { + let mut entries = Vec::<(String, String)>::new(); + while let Some((key, value)) = map.next_entry::()? { + if entries.iter().any(|(existing, _)| existing == &key) { + return Err(serde::de::Error::custom(format!( + "duplicate regional endpoint {key}" + ))); + } + entries.push((key, value)); + } + Ok(RawEndpoints(entries)) + } + } + deserializer.deserialize_map(EndpointsVisitor) + } +} + +fn required_env(name: &str) -> Result { + std::env::var(name) + .ok() + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| DispatchError::Config(format!("{name} is required"))) +} + +fn config_error(name: &str, error: std::io::Error) -> DispatchError { + DispatchError::Config(format!("failed to read {name}: {error}")) +} + +fn decode_payload(wire: &SignedEnvelope) -> Result { + if wire.version != ENVELOPE_VERSION { + return Err(DispatchError::MalformedEnvelope); + } + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(&wire.payload) + .map_err(|_| DispatchError::MalformedEnvelope)?; + serde_json::from_slice(&payload).map_err(|_| DispatchError::MalformedEnvelope) +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + use futures::{StreamExt, stream}; + + use super::*; + use crate::global_routing::{DecisionOutcome, TokenDecisionResponse}; + use crate::global_routing_envelope::{Account, Request, Routing}; + + struct MockTransport { + calls: AtomicUsize, + regions: Mutex>, + request_ids: Mutex>, + result: Mutex>>, + } + + #[async_trait] + impl RegionalTransport for MockTransport { + async fn send( + &self, + region: &str, + request_id: &str, + _envelope: &SignedEnvelope, + ) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + self.regions.lock().unwrap().push(region.to_owned()); + self.request_ids.lock().unwrap().push(request_id.to_owned()); + self.result.lock().unwrap().take().unwrap() + } + } + + fn signed(region: &str) -> SignedRoutingDecision { + let payload = Envelope { + version: 1, + request_id: "req_1".into(), + decision_id: "decision_1".into(), + issuer: "global-router".into(), + audience: format!("dsv4flash-{region}"), + issued_at: 1, + expires_at: 2, + nonce: "nonce_1".into(), + account: Account { + api_key_id: "key_1".into(), + user_id: "user_1".into(), + org_id: None, + billing_tier: "pro".into(), + service_tier: "default".into(), + zero_data_retention: false, + lifetime_requests: None, + service_tier_explicit: false, + }, + request: Request { + endpoint: "/v1/chat/completions".into(), + canonical_model: "morph-dsv4flash".into(), + normalized_body: serde_json::json!({"model":"morph-dsv4flash"}), + body_sha256: "body".into(), + token_ids_sha256: "tokens".into(), + }, + routing: Routing { + selected_region: region.into(), + selected_pool: "pool".into(), + policy_version: "v1".into(), + ckf_generation: 1, + }, + }; + SignedRoutingDecision { + decision: TokenDecisionResponse { + generation: 1, + selected_pool_id: Some("pool".into()), + selected_dc: Some(2), + selected_region: Some(region.into()), + outcome: DecisionOutcome::Remote, + matched_prefix_blocks: Some(1), + uncached_prefill_tokens: Some(1), + }, + envelope: SignedEnvelope { + version: 1, + key_id: "key".into(), + payload: base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(serde_json::to_vec(&payload).unwrap()), + signature: "signature".into(), + }, + } + } + + fn response(body: ResponseBody) -> WanResponse { + WanResponse { + status: reqwest::StatusCode::OK, + headers: HeaderMap::new(), + body, + } + } + + fn mock(result: Result) -> Arc { + Arc::new(MockTransport { + calls: AtomicUsize::new(0), + regions: Mutex::new(Vec::new()), + request_ids: Mutex::new(Vec::new()), + result: Mutex::new(Some(result)), + }) + } + + #[tokio::test] + async fn dispatches_only_to_selected_region() { + let transport = mock(Ok(response(Box::pin(stream::empty())))); + let dispatcher = RegionalDispatcher::with_transport(transport.clone()); + dispatcher + .dispatch(&signed("us-west1-financial")) + .await + .unwrap(); + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + assert_eq!(&*transport.regions.lock().unwrap(), &["us-west1-financial"]); + assert_eq!(&*transport.request_ids.lock().unwrap(), &["req_1"]); + } + + #[tokio::test] + async fn rejects_a_decision_envelope_region_mismatch_before_network_io() { + let transport = mock(Ok(response(Box::pin(stream::empty())))); + let dispatcher = RegionalDispatcher::with_transport(transport.clone()); + let mut signed = signed("us-east5-financial"); + signed.decision.selected_region = Some("us-west1-financial".into()); + + assert_eq!( + dispatcher.dispatch(&signed).await.err(), + Some(DispatchError::DecisionMismatch) + ); + assert_eq!(transport.calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn pre_response_failure_is_not_retried() { + let transport = mock(Err(DispatchError::BeforeResponse("connect".into()))); + let dispatcher = RegionalDispatcher::with_transport(transport.clone()); + assert_eq!( + dispatcher + .dispatch(&signed("us-east5-financial")) + .await + .err(), + Some(DispatchError::BeforeResponse("connect".into())) + ); + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + } + + struct DropGuard(Arc); + impl Drop for DropGuard { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + #[tokio::test] + async fn dropping_body_propagates_cancellation_without_buffering() { + let dropped = Arc::new(AtomicBool::new(false)); + let guard = DropGuard(dropped.clone()); + let body = async_stream::stream! { + let _guard = guard; + yield Ok(Bytes::from_static(b"data: first\n\n")); + futures::future::pending::<()>().await; + }; + let transport = mock(Ok(response(Box::pin(body)))); + let dispatcher = RegionalDispatcher::with_transport(transport); + let mut response = dispatcher + .dispatch(&signed("us-east5-financial")) + .await + .unwrap(); + assert_eq!( + response.body.next().await.unwrap().unwrap(), + "data: first\n\n" + ); + drop(response); + assert!(dropped.load(Ordering::SeqCst)); + } + + #[tokio::test] + async fn midstream_failure_preserves_prior_bytes_and_does_not_retry() { + let body = stream::iter(vec![ + Ok(Bytes::from_static(b"data: first\n\n")), + Err(DispatchError::Midstream("reset".into())), + ]); + let transport = mock(Ok(response(Box::pin(body)))); + let dispatcher = RegionalDispatcher::with_transport(transport.clone()); + let mut response = dispatcher + .dispatch(&signed("us-west1-financial")) + .await + .unwrap(); + assert_eq!( + response.body.next().await.unwrap().unwrap(), + "data: first\n\n" + ); + assert_eq!( + response.body.next().await.unwrap().unwrap_err(), + DispatchError::Midstream("reset".into()) + ); + assert_eq!(transport.calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn configured_dispatch_never_returns_local_ownership() { + let transport = mock(Ok(response(Box::pin(stream::empty())))); + let dispatcher = RegionalDispatcher::with_transport(transport); + let outcome = dispatch_authoritatively(Some(&dispatcher), Some(&signed("east"))) + .await + .unwrap(); + assert!(matches!(outcome, DispatchOwnership::Wan(_))); + + assert!( + dispatch_authoritatively(Some(&dispatcher), None) + .await + .is_err() + ); + } + + #[test] + fn endpoint_contract_is_strict() { + let parse = + |json: &str| TransportConfig::parse("cert".into(), "key".into(), "ca".into(), json); + assert!(parse(r#"{"east":"https://east.internal"}"#).is_ok()); + assert!(parse(r#"{"east":"http://east.internal"}"#).is_err()); + assert!(parse(r#"{"east":"https://east.internal/path"}"#).is_err()); + assert!(parse(r#"{"east":"https://one","east":"https://two"}"#).is_err()); + assert!(parse(r#"[]"#).is_err()); + assert!(parse(r#"{}"#).is_err()); + } + + #[test] + fn production_call_graph_returns_wan_ownership_before_local_generate() { + let preprocessor = include_str!("preprocessor.rs"); + let dispatch_sites = preprocessor + .match_indices("global_routing_transport::dispatch_signed(&signed)") + .map(|(offset, _)| offset) + .collect::>(); + assert_eq!( + dispatch_sites.len(), + 2, + "chat and completion must both dispatch" + ); + for dispatch in dispatch_sites { + let tail = &preprocessor[dispatch..]; + let owned_return = tail + .find("return Ok(ResponseStream::new") + .expect("globally routed request must return WAN ownership"); + let local_generate = tail + .find("let response_stream = next.generate") + .expect("test must observe the local generation site"); + assert!( + owned_return < local_generate, + "local next.generate became reachable after authoritative dispatch" + ); + } + + let http_source = include_str!("http/service/openai.rs"); + assert_eq!( + http_source + .matches("take_global_wan_response(&ctx, &request_id)?") + .count(), + 2, + "both public OpenAI sources must consume the live WAN response" + ); + } +} diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 3156fea15e87..6a051901e0b5 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -98,6 +98,70 @@ pub const ANNOTATION_REQUEST_ID: &str = "request_id"; const VALIDATION_PREFIX: &str = "Validation: "; const BATCH_FILE_STORAGE_NOT_IMPLEMENTED: &str = "Batch file storage is not implemented yet."; + +fn take_global_wan_response( + ctx: &Arc, + expected_request_id: &str, +) -> Result, ErrorResponse> { + let Some(extension) = + ctx.get_extension(crate::global_routing_transport::WAN_RESPONSE_CONTEXT_KEY) + else { + return Ok(None); + }; + let handle = extension + .downcast::() + .map_err(|_| { + ErrorMessage::internal_server_error("Global WAN response context has the wrong type") + })?; + let response = handle.take().ok_or_else(|| { + ErrorMessage::internal_server_error("Global WAN response was already consumed") + })?; + let returned_request_id = response.headers.get("x-request-id").ok_or_else(|| { + ErrorMessage::internal_server_error("Regional response is missing x-request-id") + })?; + if returned_request_id.as_bytes() != expected_request_id.as_bytes() { + return Err(ErrorMessage::internal_server_error( + "Regional response request ID does not match the signed request", + )); + } + Ok(Some(response)) +} + +fn relay_global_wan_response( + response: crate::global_routing_transport::WanResponse, + inflight_guard: super::metrics::InflightGuard, + http_queue_guard: super::metrics::HttpQueueGuard, +) -> Result { + let successful_status = response.status.is_success(); + let body = async_stream::stream! { + let mut body = response.body; + let mut inflight_guard = inflight_guard; + let mut http_queue_guard = Some(http_queue_guard); + while let Some(item) = body.next().await { + drop(http_queue_guard.take()); + match item { + Ok(bytes) => yield Ok(bytes), + Err(error) => { + inflight_guard.mark_error(ErrorType::Internal); + yield Err(error); + return; + } + } + } + drop(http_queue_guard.take()); + if successful_status { + inflight_guard.mark_ok(); + } + }; + let mut builder = Response::builder().status(response.status); + *builder.headers_mut().expect("response builder is valid") = response.headers; + builder.body(Body::from_stream(body)).map_err(|error| { + ErrorMessage::internal_server_error_with_details( + "Failed to relay global WAN response", + error.to_string(), + ) + }) +} const BATCH_JOB_STATE_NOT_IMPLEMENTED: &str = "Batch job lifecycle persistence is not implemented yet."; const BATCH_OUTPUT_RETRIEVAL_NOT_IMPLEMENTED: &str = @@ -935,6 +999,10 @@ async fn completions_single( // capture the context to cancel the stream if the client disconnects let ctx = stream.context(); + if let Some(response) = take_global_wan_response(&ctx, &request_id)? { + return relay_global_wan_response(response, inflight_guard, http_queue_guard); + } + let annotations = annotations.map_or(Vec::new(), |annotations| { annotations .iter() @@ -2894,6 +2962,10 @@ async fn chat_completions( // capture the context to cancel the stream if the client disconnects let ctx = stream.context(); + if let Some(response) = take_global_wan_response(&ctx, &request_id)? { + return relay_global_wan_response(response, inflight_guard, http_queue_guard); + } + // prepare any requested annotations let annotations = annotations.map_or(Vec::new(), |annotations| { annotations diff --git a/lib/llm/src/lib.rs b/lib/llm/src/lib.rs index eb9f0a3f4a57..5ef9d2596b7c 100644 --- a/lib/llm/src/lib.rs +++ b/lib/llm/src/lib.rs @@ -19,6 +19,7 @@ pub mod fpm_trace; pub mod frontend_config; pub mod global_routing; pub mod global_routing_envelope; +pub mod global_routing_transport; pub mod grpc; pub mod http; pub mod hub; diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 0f5e6c34251e..98eb556d0201 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -5512,10 +5512,30 @@ impl code: 503, message: error.to_string(), })?; + let wan_response = crate::global_routing_transport::dispatch_signed(&signed) + .await + .map_err(|error| crate::http::service::error::HttpError { + code: 502, + message: error.to_string(), + })?; context.insert( crate::global_routing_envelope::SIGNED_ENVELOPE_CONTEXT_KEY, signed, ); + context.insert( + crate::global_routing_transport::WAN_RESPONSE_CONTEXT_KEY, + crate::global_routing_transport::WanResponseHandle::new(wan_response), + ); + + // Global routing now owns the request. The HTTP source takes the + // live WAN response from context and relays it byte-for-byte. Do + // not reach the local `next.generate` call below. + return Ok(ResponseStream::new( + Box::pin(stream::empty()), + Arc::new(dynamo_runtime::pipeline::context::StreamContext::from( + context, + )), + )); } tracing::trace!(request = ?common_request, prompt_injected_reasoning, "Pre-processed request"); @@ -5759,10 +5779,26 @@ impl code: 503, message: error.to_string(), })?; + let wan_response = crate::global_routing_transport::dispatch_signed(&signed) + .await + .map_err(|error| crate::http::service::error::HttpError { + code: 502, + message: error.to_string(), + })?; context.insert( crate::global_routing_envelope::SIGNED_ENVELOPE_CONTEXT_KEY, signed, ); + context.insert( + crate::global_routing_transport::WAN_RESPONSE_CONTEXT_KEY, + crate::global_routing_transport::WanResponseHandle::new(wan_response), + ); + return Ok(ResponseStream::new( + Box::pin(stream::empty()), + Arc::new(dynamo_runtime::pipeline::context::StreamContext::from( + context, + )), + )); } let trace_state = crate::request_trace::build_request_end_trace_state( diff --git a/lib/runtime/src/engine.rs b/lib/runtime/src/engine.rs index 5d4e00004e0f..c2a1bd77910e 100644 --- a/lib/runtime/src/engine.rs +++ b/lib/runtime/src/engine.rs @@ -118,6 +118,12 @@ pub trait AsyncEngineContext: Send + Sync + Debug { /// Unique ID for the Stream fn id(&self) -> &str; + /// Retrieve request-scoped shared state carried by a pipeline context. + /// Context implementations without an extension registry return `None`. + fn get_extension(&self, _key: &str) -> Option> { + None + } + /// Returns true if `stop_generating()` has been called; otherwise, false. fn is_stopped(&self) -> bool; diff --git a/lib/runtime/src/pipeline/context.rs b/lib/runtime/src/pipeline/context.rs index 484809b4a4fc..c4964789c0fe 100644 --- a/lib/runtime/src/pipeline/context.rs +++ b/lib/runtime/src/pipeline/context.rs @@ -298,6 +298,10 @@ impl AsyncEngineContext for StreamContext { self.controller.id() } + fn get_extension(&self, key: &str) -> Option> { + self.registry.get_shared_erased(key) + } + fn stop(&self) { self.controller.stop(); } @@ -608,6 +612,19 @@ mod tests { ); } + #[test] + fn test_shared_extension_is_visible_through_engine_context() { + let mut ctx = Context::new(Input { + value: "Hello".to_string(), + }); + ctx.insert("wan_response", 42_u64); + + let stream_ctx = StreamContext::from(ctx); + let extension = AsyncEngineContext::get_extension(&stream_ctx, "wan_response") + .expect("shared extension must survive the context handoff"); + assert_eq!(*extension.downcast::().unwrap(), 42); + } + #[test] fn test_transfer() { let ctx = Context::new(Input { diff --git a/lib/runtime/src/pipeline/registry.rs b/lib/runtime/src/pipeline/registry.rs index 8365b025591e..44f81b0b7a54 100644 --- a/lib/runtime/src/pipeline/registry.rs +++ b/lib/runtime/src/pipeline/registry.rs @@ -67,6 +67,10 @@ impl Registry { .ok_or_else(|| format!("Shared key not found: {}", key)) } + pub(crate) fn get_shared_erased(&self, key: &str) -> Option> { + self.shared_storage.get(key).cloned() + } + /// Retrieve an optional shared object from the registry by key and type. pub fn get_shared_optional( &self, From 83d176fdd5fbdb863b8ddaa34a6225b4d0f131fd Mon Sep 17 00:00:00 2001 From: skeptrune Date: Wed, 26 Aug 2026 11:09:04 -0700 Subject: [PATCH 6/8] fix(global-routing): terminate regional WAN ownership Signed-off-by: skeptrune --- lib/llm/src/global_routing_transport.rs | 262 ++++++++++++++++++++++-- lib/llm/src/http/service/openai.rs | 7 +- lib/llm/src/http/service/service_v2.rs | 5 + lib/llm/src/preprocessor.rs | 103 +++------- 4 files changed, 279 insertions(+), 98 deletions(-) diff --git a/lib/llm/src/global_routing_transport.rs b/lib/llm/src/global_routing_transport.rs index e47f035def05..f4cf28a492a9 100644 --- a/lib/llm/src/global_routing_transport.rs +++ b/lib/llm/src/global_routing_transport.rs @@ -9,12 +9,13 @@ use std::collections::BTreeMap; use std::pin::Pin; -use std::sync::{Arc, Mutex, OnceLock}; +use std::sync::{Arc, LazyLock, Mutex, OnceLock}; use async_trait::async_trait; use base64::Engine; use bytes::Bytes; -use futures::{Stream, TryStreamExt}; +use futures::Stream; +use prometheus::{IntCounterVec, Opts, Registry}; use reqwest::header::{CONTENT_TYPE, HeaderMap, HeaderValue}; use serde::de::{MapAccess, Visitor}; use url::Url; @@ -55,8 +56,90 @@ impl WanResponseHandle { pub const WAN_RESPONSE_CONTEXT_KEY: &str = "global_routing.wan_response"; +/// Remove headers that describe the regional HTTP hop rather than the outer +/// client response. `Connection` may nominate additional hop-scoped fields. +pub fn sanitize_response_headers(headers: &mut HeaderMap) { + let connection_scoped = headers + .get_all("connection") + .iter() + .filter_map(|value| value.to_str().ok()) + .flat_map(|value| value.split(',')) + .map(str::trim) + .filter(|name| !name.is_empty()) + .map(str::to_owned) + .collect::>(); + for name in connection_scoped { + headers.remove(name); + } + for name in [ + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + // The outer Hyper response computes framing for its streaming body. + "content-length", + ] { + headers.remove(name); + } +} + static DISPATCHER: OnceLock> = OnceLock::new(); +static WAN_ATTEMPTS: LazyLock = LazyLock::new(|| { + IntCounterVec::new( + Opts::new( + "morph_global_routing_wan_attempts_total", + "Authoritative WAN dispatch attempts by selected region.", + ), + &["selected_region"], + ) + .expect("static metric options are valid") +}); + +static WAN_OUTCOMES: LazyLock = LazyLock::new(|| { + IntCounterVec::new( + Opts::new( + "morph_global_routing_wan_outcomes_total", + "WAN transport outcomes by selected region and bounded phase.", + ), + &["selected_region", "phase", "outcome"], + ) + .expect("static metric options are valid") +}); + +static WAN_RESPONSE_STATUS: LazyLock = LazyLock::new(|| { + IntCounterVec::new( + Opts::new( + "morph_global_routing_wan_response_status_total", + "Regional pod proxy HTTP response statuses.", + ), + &["selected_region", "status"], + ) + .expect("static metric options are valid") +}); + +static WAN_BYTES: LazyLock = LazyLock::new(|| { + IntCounterVec::new( + Opts::new( + "morph_global_routing_wan_bytes_relayed_total", + "Response bytes relayed from regional pod proxies.", + ), + &["selected_region"], + ) + .expect("static metric options are valid") +}); + +pub fn ensure_metrics_registered_prometheus(registry: &Registry) -> Result<(), prometheus::Error> { + registry.register(Box::new(WAN_ATTEMPTS.clone()))?; + registry.register(Box::new(WAN_OUTCOMES.clone()))?; + registry.register(Box::new(WAN_RESPONSE_STATUS.clone()))?; + registry.register(Box::new(WAN_BYTES.clone())) +} + #[derive(Debug, thiserror::Error, Clone, PartialEq, Eq)] pub enum DispatchError { #[error("global routing WAN transport is misconfigured: {0}")] @@ -213,6 +296,8 @@ impl RegionalTransport for Http2Transport { let request_id = HeaderValue::from_str(request_id).map_err(|_| DispatchError::MalformedEnvelope)?; + WAN_ATTEMPTS.with_label_values(&[region]).inc(); + // One send, one selected region. reqwest does not retry requests. let response = self .client @@ -223,12 +308,22 @@ impl RegionalTransport for Http2Transport { .json(envelope) .send() .await - .map_err(|error| DispatchError::BeforeResponse(error.to_string()))?; + .map_err(|error| { + WAN_OUTCOMES + .with_label_values(&[region, "before_response", "failure"]) + .inc(); + DispatchError::BeforeResponse(error.to_string()) + })?; let status = response.status(); let headers = response.headers().clone(); - let body = response - .bytes_stream() - .map_err(|error| DispatchError::Midstream(error.to_string())); + WAN_RESPONSE_STATUS + .with_label_values(&[region, &status.as_u16().to_string()]) + .inc(); + WAN_OUTCOMES + .with_label_values(&[region, "response", "received"]) + .inc(); + let body = + InstrumentedResponseBody::new(region.to_owned(), Box::pin(response.bytes_stream())); Ok(WanResponse { status, headers, @@ -237,6 +332,68 @@ impl RegionalTransport for Http2Transport { } } +struct InstrumentedResponseBody { + selected_region: String, + inner: Pin> + Send>>, + terminal: bool, +} + +impl InstrumentedResponseBody { + fn new( + selected_region: String, + inner: Pin> + Send>>, + ) -> Self { + Self { + selected_region, + inner, + terminal: false, + } + } +} + +impl Stream for InstrumentedResponseBody { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + match self.inner.as_mut().poll_next(cx) { + std::task::Poll::Ready(Some(Ok(bytes))) => { + WAN_BYTES + .with_label_values(&[&self.selected_region]) + .inc_by(bytes.len() as u64); + std::task::Poll::Ready(Some(Ok(bytes))) + } + std::task::Poll::Ready(Some(Err(error))) => { + WAN_OUTCOMES + .with_label_values(&[&self.selected_region, "stream", "failure"]) + .inc(); + self.terminal = true; + std::task::Poll::Ready(Some(Err(DispatchError::Midstream(error.to_string())))) + } + std::task::Poll::Ready(None) => { + WAN_OUTCOMES + .with_label_values(&[&self.selected_region, "stream", "complete"]) + .inc(); + self.terminal = true; + std::task::Poll::Ready(None) + } + std::task::Poll::Pending => std::task::Poll::Pending, + } + } +} + +impl Drop for InstrumentedResponseBody { + fn drop(&mut self) { + if !self.terminal { + WAN_OUTCOMES + .with_label_values(&[&self.selected_region, "stream", "cancelled"]) + .inc(); + } + } +} + #[derive(Debug)] struct TransportConfig { cert_path: String, @@ -583,6 +740,71 @@ mod tests { assert!(parse(r#"{}"#).is_err()); } + #[test] + fn response_headers_drop_hop_and_framing_state() { + let mut headers = HeaderMap::new(); + headers.insert("connection", HeaderValue::from_static("x-hop, keep-alive")); + headers.insert("x-hop", HeaderValue::from_static("private")); + headers.insert("keep-alive", HeaderValue::from_static("timeout=5")); + headers.insert("transfer-encoding", HeaderValue::from_static("chunked")); + headers.insert("content-length", HeaderValue::from_static("123")); + headers.insert( + "content-type", + HeaderValue::from_static("text/event-stream"), + ); + + sanitize_response_headers(&mut headers); + + assert!(!headers.contains_key("connection")); + assert!(!headers.contains_key("x-hop")); + assert!(!headers.contains_key("keep-alive")); + assert!(!headers.contains_key("transfer-encoding")); + assert!(!headers.contains_key("content-length")); + assert_eq!(headers["content-type"], "text/event-stream"); + } + + #[test] + fn wan_metrics_register_with_bounded_label_contracts() { + let registry = Registry::new(); + ensure_metrics_registered_prometheus(®istry).unwrap(); + WAN_ATTEMPTS.with_label_values(&["test-region"]).inc(); + WAN_OUTCOMES + .with_label_values(&["test-region", "stream", "complete"]) + .inc(); + WAN_RESPONSE_STATUS + .with_label_values(&["test-region", "200"]) + .inc(); + WAN_BYTES.with_label_values(&["test-region"]).inc(); + let names = registry + .gather() + .into_iter() + .map(|family| family.name().to_owned()) + .collect::>(); + for expected in [ + "morph_global_routing_wan_attempts_total", + "morph_global_routing_wan_outcomes_total", + "morph_global_routing_wan_response_status_total", + "morph_global_routing_wan_bytes_relayed_total", + ] { + assert!( + names.iter().any(|name| name == expected), + "missing {expected}" + ); + } + } + + #[test] + fn dropping_live_stream_records_observable_cancellation() { + let counter = WAN_OUTCOMES.with_label_values(&["cancel-test", "stream", "cancelled"]); + let before = counter.get(); + let never_polled = stream::pending::>(); + drop(InstrumentedResponseBody::new( + "cancel-test".to_owned(), + Box::pin(never_polled), + )); + assert_eq!(counter.get(), before + 1); + } + #[test] fn production_call_graph_returns_wan_ownership_before_local_generate() { let preprocessor = include_str!("preprocessor.rs"); @@ -590,12 +812,8 @@ mod tests { .match_indices("global_routing_transport::dispatch_signed(&signed)") .map(|(offset, _)| offset) .collect::>(); - assert_eq!( - dispatch_sites.len(), - 2, - "chat and completion must both dispatch" - ); - for dispatch in dispatch_sites { + assert_eq!(dispatch_sites.len(), 1, "only chat may dispatch over WAN"); + for &dispatch in &dispatch_sites { let tail = &preprocessor[dispatch..]; let owned_return = tail .find("return Ok(ResponseStream::new") @@ -614,8 +832,24 @@ mod tests { http_source .matches("take_global_wan_response(&ctx, &request_id)?") .count(), - 2, - "both public OpenAI sources must consume the live WAN response" + 1, + "only the chat source may consume a WAN response" + ); + + assert!( + !preprocessor.contains("\"/v1/completions\",\n &common_request.model"), + "legacy completions must not own WAN dispatch" + ); + assert!( + preprocessor.contains("if supplied_prompt_token_digest.is_none()"), + "a selected regional request must skip the second global decision" + ); + let ownership_guard = preprocessor + .find("if supplied_prompt_token_digest.is_none()") + .unwrap(); + assert!( + ownership_guard < dispatch_sites[0], + "the sole WAN dispatch must be inside the unselected request guard" ); } } diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs index 6a051901e0b5..42be9a7725e1 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -128,10 +128,11 @@ fn take_global_wan_response( } fn relay_global_wan_response( - response: crate::global_routing_transport::WanResponse, + mut response: crate::global_routing_transport::WanResponse, inflight_guard: super::metrics::InflightGuard, http_queue_guard: super::metrics::HttpQueueGuard, ) -> Result { + crate::global_routing_transport::sanitize_response_headers(&mut response.headers); let successful_status = response.status.is_success(); let body = async_stream::stream! { let mut body = response.body; @@ -999,10 +1000,6 @@ async fn completions_single( // capture the context to cancel the stream if the client disconnects let ctx = stream.context(); - if let Some(response) = take_global_wan_response(&ctx, &request_id)? { - return relay_global_wan_response(response, inflight_guard, http_queue_guard); - } - let annotations = annotations.map_or(Vec::new(), |annotations| { annotations .iter() diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index fc46c2fbef99..c9563cdef387 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -1220,6 +1220,11 @@ impl HttpServiceConfigBuilder { if let Err(e) = crate::global_routing::ensure_metrics_registered_prometheus(®istry) { tracing::warn!("Failed to register global routing decision metrics: {}", e); } + if let Err(e) = + crate::global_routing_transport::ensure_metrics_registered_prometheus(®istry) + { + tracing::warn!("Failed to register global routing WAN metrics: {}", e); + } if let Err(e) = ensure_tokio_perf_metrics_registered_prometheus(®istry) { tracing::warn!("Failed to register tokio perf metrics: {}", e); } diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 98eb556d0201..77856ef1d667 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -5457,11 +5457,12 @@ impl .await?; attach_agent_context_from_context(&mut common_request, &context); + let supplied_prompt_token_digest = request + .nvext + .as_ref() + .and_then(|nvext| nvext.prompt_token_digest.as_deref()); crate::global_routing_envelope::enforce_prompt_token_digest( - request - .nvext - .as_ref() - .and_then(|nvext| nvext.prompt_token_digest.as_deref()), + supplied_prompt_token_digest, &common_request.token_ids, &request_id, )?; @@ -5472,18 +5473,24 @@ impl prompt_injected_reasoning, )?; - if let Some(decision) = crate::global_routing::decide( - &common_request.model, - &common_request.token_ids, - self.kv_cache_block_size, - &request_id, - common_request.multi_modal_data.is_some() || common_request.mm_routing_info.is_some(), - ) - .await - .map_err(|error| crate::http::service::error::HttpError { - code: 503, - message: error.to_string(), - })? { + // A valid supplied digest marks the request as already selected and + // authenticated by the global router. Execute it in this regional + // frontend instead of recursively making another WAN decision. + if supplied_prompt_token_digest.is_none() + && let Some(decision) = crate::global_routing::decide( + &common_request.model, + &common_request.token_ids, + self.kv_cache_block_size, + &request_id, + common_request.multi_modal_data.is_some() + || common_request.mm_routing_info.is_some(), + ) + .await + .map_err(|error| crate::http::service::error::HttpError { + code: 503, + message: error.to_string(), + })? + { let mut signed_request = request.clone(); signed_request.inner.stream = Some(original_stream_flag); let normalized_body = serde_json::to_value(&signed_request).map_err(|error| { @@ -5675,7 +5682,7 @@ impl let _stage_guard = StageGuard::new(STAGE_PREPROCESS, ""); // unpack the request - let (mut request, mut context) = request.into_parts(); + let (mut request, context) = request.into_parts(); let request_id = context.id().to_string(); // Preserve original streaming flag @@ -5739,68 +5746,6 @@ impl &request_id, )?; - if let Some(decision) = crate::global_routing::decide( - &common_request.model, - &common_request.token_ids, - self.kv_cache_block_size, - &request_id, - common_request.multi_modal_data.is_some() || common_request.mm_routing_info.is_some(), - ) - .await - .map_err(|error| crate::http::service::error::HttpError { - code: 503, - message: error.to_string(), - })? { - let mut signed_request = request.clone(); - signed_request.inner.stream = Some(original_stream_flag); - let normalized_body = serde_json::to_value(&signed_request).map_err(|error| { - crate::http::service::error::HttpError { - code: 500, - message: format!("failed to serialize normalized request: {error}"), - } - })?; - let signed = crate::global_routing_envelope::mint_signed_routing_decision( - decision, - context - .metadata() - .get(crate::global_routing_envelope::TRUSTED_AUTH_METADATA_KEY) - .map(String::as_str), - &request_id, - "/v1/completions", - &common_request.model, - normalized_body, - &common_request.token_ids, - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_secs(), - ) - .map_err(|error| crate::http::service::error::HttpError { - code: 503, - message: error.to_string(), - })?; - let wan_response = crate::global_routing_transport::dispatch_signed(&signed) - .await - .map_err(|error| crate::http::service::error::HttpError { - code: 502, - message: error.to_string(), - })?; - context.insert( - crate::global_routing_envelope::SIGNED_ENVELOPE_CONTEXT_KEY, - signed, - ); - context.insert( - crate::global_routing_transport::WAN_RESPONSE_CONTEXT_KEY, - crate::global_routing_transport::WanResponseHandle::new(wan_response), - ); - return Ok(ResponseStream::new( - Box::pin(stream::empty()), - Arc::new(dynamo_runtime::pipeline::context::StreamContext::from( - context, - )), - )); - } - let trace_state = crate::request_trace::build_request_end_trace_state( &common_request, &tracker, From 7b20d320c83c343f543fa85bfc963c3ff7714522 Mon Sep 17 00:00:00 2001 From: skeptrune Date: Wed, 26 Aug 2026 15:06:50 -0700 Subject: [PATCH 7/8] fix(global-routing): route partial-block prompts exactly --- components/global-ckf-consumer/src/api.rs | 20 +++++++--- components/global-ckf-consumer/src/policy.rs | 42 +++++++++++++------- 2 files changed, 42 insertions(+), 20 deletions(-) diff --git a/components/global-ckf-consumer/src/api.rs b/components/global-ckf-consumer/src/api.rs index 5889eb14e3c3..3d0126140766 100644 --- a/components/global-ckf-consumer/src/api.rs +++ b/components/global-ckf-consumer/src/api.rs @@ -411,6 +411,7 @@ async fn token_decision( &request.model, request.role, &hashes, + request.token_ids.len() as u64, request.block_size, request.local_dc, request.stable_tie_key, @@ -475,6 +476,7 @@ pub(crate) fn evaluate_decision( model: &str, role: QueryRole, hashes: &[u64], + request_token_count: u64, block_size: u32, local_dc: u64, stable_tie_key: u64, @@ -489,7 +491,7 @@ pub(crate) fn evaluate_decision( } let input = PolicyInput { local_dc: dynamo_kv_router::identity::DcId::new(local_dc), - query_block_count: hashes.len() as u64, + query_token_count: request_token_count, native_block_size_tokens: block_size as u64, stable_tie_key, }; @@ -600,11 +602,6 @@ fn native_hashes( } else { request.token_ids.len() / stride }; - if block_count == 0 && !request.token_ids.is_empty() { - return Err(ApiError::bad_request( - "token_ids does not contain one complete native block", - )); - } if block_count > max_query_blocks { return Err(ApiError::payload_too_large( "token_ids exceeds the configured block limit", @@ -887,6 +884,17 @@ mod tests { ); } + #[test] + fn token_contract_accepts_a_partial_first_block() { + let mut request = request(); + request.token_ids = vec![1, 2, 3]; + request.block_size = 256; + request.block_mm_infos = None; + request.is_eagle = None; + + assert!(native_hashes(&request, 8).unwrap().is_empty()); + } + #[test] fn token_contract_rejects_ambiguous_hash_input() { let value = serde_json::json!({ diff --git a/components/global-ckf-consumer/src/policy.rs b/components/global-ckf-consumer/src/policy.rs index 9f6cb96140f9..c7ae00d3e20c 100644 --- a/components/global-ckf-consumer/src/policy.rs +++ b/components/global-ckf-consumer/src/policy.rs @@ -51,7 +51,7 @@ pub struct PoolFacts { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct PolicyInput { pub local_dc: DcId, - pub query_block_count: u64, + pub query_token_count: u64, pub native_block_size_tokens: u64, pub stable_tie_key: u64, } @@ -158,7 +158,10 @@ fn evaluate_pool( { Some(IneligibleReason::OccupancyIncomplete) } - Some(_) if facts.matched_prefix_blocks > input.query_block_count => { + Some(_) + if facts.matched_prefix_blocks + > input.query_token_count / input.native_block_size_tokens => + { Some(IneligibleReason::PrefixExceedsQuery) } Some(_) => None, @@ -170,12 +173,13 @@ fn evaluate_pool( } let occupancy = facts.occupancy.expect("eligible pool has occupancy"); - let uncached_blocks = input.query_block_count - facts.matched_prefix_blocks; - let uncached_prefill_tokens = uncached_blocks + let cached_tokens = facts + .matched_prefix_blocks .checked_mul(input.native_block_size_tokens) .ok_or(PolicyError::TokenArithmeticOverflow { pool_id: facts.pool_id, })?; + let uncached_prefill_tokens = input.query_token_count - cached_tokens; Ok(Ok(EligiblePool { pool_id: facts.pool_id, matched_prefix_blocks: facts.matched_prefix_blocks, @@ -252,7 +256,7 @@ mod tests { fn input() -> PolicyInput { PolicyInput { local_dc: DcId::new(1), - query_block_count: 10, + query_token_count: 2_600, native_block_size_tokens: 256, stable_tie_key: 42, } @@ -328,6 +332,24 @@ mod tests { assert_eq!(selected([facts(2, 10, 50), facts(1, 10, 50)]), pool(1)); } + #[test] + fn partial_block_tokens_are_counted_exactly() { + let mut short = input(); + short.query_token_count = 40; + let selected = select_pool(short, [facts(1, 0, 0)]) + .unwrap() + .selected + .unwrap(); + assert_eq!(selected.matched_prefix_blocks, 0); + assert_eq!(selected.uncached_prefill_tokens, 40); + + let selected = select_pool(input(), [facts(1, 10, 0)]) + .unwrap() + .selected + .unwrap(); + assert_eq!(selected.uncached_prefill_tokens, 40); + } + #[test] fn stable_remote_tie_is_deterministic_and_input_order_independent() { let east = facts(2, 10, 0); @@ -366,20 +388,12 @@ mod tests { } #[test] - fn invalid_prefix_and_arithmetic_overflow_never_saturate_into_a_route() { + fn invalid_prefix_never_saturates_into_a_route() { let invalid = facts(1, 11, 0); let decision = select_pool(input(), [invalid]).unwrap(); assert_eq!( decision.ineligible[0].reason, IneligibleReason::PrefixExceedsQuery ); - - let mut huge = input(); - huge.query_block_count = u64::MAX; - huge.native_block_size_tokens = 2; - assert!(matches!( - select_pool(huge, [facts(1, 0, 0)]), - Err(PolicyError::TokenArithmeticOverflow { .. }) - )); } } From 153fb63b2923d13abd5042cfc3800e98a5f7ac12 Mon Sep 17 00:00:00 2001 From: bhaktatejas922 Date: Wed, 26 Aug 2026 16:33:04 -0700 Subject: [PATCH 8/8] fix(global-routing): preserve unary WAN contract --- lib/llm/src/preprocessor.rs | 54 +++++++++++++++++++++++++++++++++++-- 1 file changed, 52 insertions(+), 2 deletions(-) diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs index 77856ef1d667..6cf1c47f8a69 100644 --- a/lib/llm/src/preprocessor.rs +++ b/lib/llm/src/preprocessor.rs @@ -97,6 +97,24 @@ pub use crate::protocols::common::preprocessor::PreprocessedEmbeddingRequest; use crate::protocols::common::llm_backend::EmbeddingsEngineOutput; +fn prepare_signed_wan_request( + request: &NvCreateChatCompletionRequest, + original_stream_flag: bool, +) -> NvCreateChatCompletionRequest { + let mut signed_request = request.clone(); + signed_request.inner.stream = Some(original_stream_flag); + + // `enable_usage_for_nonstreaming` adds stream_options for Dynamo's internal + // streaming aggregation. That field is invalid on the regional HTTP hop once + // the client-facing stream flag is restored to false. The regional frontend + // enables unary usage accounting again after it validates the signed request. + if !original_stream_flag { + signed_request.inner.stream_options = None; + } + + signed_request +} + fn routing_priorities(hints: Option<&AgentHints>) -> (Option, Option, Option) { let priority_jump = hints.and_then(|h| { h.priority @@ -5491,8 +5509,7 @@ impl message: error.to_string(), })? { - let mut signed_request = request.clone(); - signed_request.inner.stream = Some(original_stream_flag); + let signed_request = prepare_signed_wan_request(&request, original_stream_flag); let normalized_body = serde_json::to_value(&signed_request).map_err(|error| { crate::http::service::error::HttpError { code: 500, @@ -5933,6 +5950,39 @@ mod tests { FinishReason, Role, }; + #[test] + fn test_prepare_signed_wan_request_restores_client_stream_contract() { + let mut unary: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({ + "model": "morph-dsv4flash", + "messages": [{"role": "user", "content": "hello"}], + "stream": false + })) + .unwrap(); + unary.enable_usage_for_nonstreaming(false); + unary.inner.stream = Some(true); + + let signed_unary = prepare_signed_wan_request(&unary, false); + assert_eq!(signed_unary.inner.stream, Some(false)); + assert!(signed_unary.inner.stream_options.is_none()); + + let streaming: NvCreateChatCompletionRequest = serde_json::from_value(serde_json::json!({ + "model": "morph-dsv4flash", + "messages": [{"role": "user", "content": "hello"}], + "stream": true, + "stream_options": {"include_usage": true} + })) + .unwrap(); + + let signed_streaming = prepare_signed_wan_request(&streaming, true); + assert_eq!(signed_streaming.inner.stream, Some(true)); + assert!( + signed_streaming + .inner + .stream_options + .is_some_and(|options| options.include_usage) + ); + } + fn chat_stream_chunk( index: u32, role: Option,