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/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..3d0126140766 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.token_ids.len() as u64, + 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. +#[allow(clippy::too_many_arguments)] +pub(crate) fn evaluate_decision( + state: &AppState, + model: &str, + role: QueryRole, + hashes: &[u64], + request_token_count: 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), - query_block_count: hashes.len() as u64, - native_block_size_tokens: request.query.block_size as u64, - stable_tie_key: request.stable_tie_key, + local_dc: dynamo_kv_router::identity::DcId::new(local_dc), + query_token_count: request_token_count, + 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 { @@ -572,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", @@ -790,7 +815,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 { @@ -859,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/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/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 { .. }) - )); } } 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.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 new file mode 100644 index 000000000000..1af9215b3822 --- /dev/null +++ b/lib/llm/src/global_routing_envelope.rs @@ -0,0 +1,776 @@ +// 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 +//! +//! ```json +//! {"version":1,"key_id":"...","payload":"...","signature":"..."} +//! ``` +//! +//! 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.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, OnceLock}; + +use prometheus::{IntCounterVec, Opts, Registry}; +use serde::{Deserialize, Serialize}; + +use crate::protocols::TokenIdType; + +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"; + +static SIGNING_CONFIG: OnceLock> = OnceLock::new(); + +#[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, + 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, +} + +/// 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 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)] +#[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 + /// 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)] +#[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_generation: u64, +} + +#[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 key identifier is unknown")] + UnknownKey, + #[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_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( + 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 { + 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 { + 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 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 { + keys: std::collections::HashMap, +} + +impl EnvelopeVerifier { + /// Load a SPKI PEM Ed25519 public key (`openssl pkey -pubout`). + pub fn from_public_key_pem( + key_id: impl Into, + pem: &str, + ) -> Result { + use ed25519_dalek::pkcs8::DecodePublicKey; + Ok(Self { + keys: [( + key_id.into(), + ed25519_dalek::VerifyingKey::from_public_key_pem(pem)?, + )] + .into_iter() + .collect(), + }) + } + + /// 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: &SignedEnvelope, + local_region: &str, + now_unix_seconds: u64, + ) -> Result { + use base64::Engine; + + let engine = &base64::engine::general_purpose::URL_SAFE_NO_PAD; + 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(&wire.signature) + .ok() + .and_then(|bytes| ed25519_dalek::Signature::from_slice(&bytes).ok()) + .ok_or(EnvelopeError::Malformed)?; + 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 = + 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) + } +} + +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") +}); + +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(ENVELOPES.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(), + }) +} + +/// SHA256 digest of the exact native token sequence, as lowercase hex. +/// +/// 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 { + use sha2::Digest; + let mut hasher = sha2::Sha256::new(); + for token_id in token_ids { + hasher.update(token_id.to_be_bytes()); + } + hex_digest(hasher.finalize()) +} + +/// 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 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"); + } + hex +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 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]), + } + } + + 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(), + 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(), + 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(), + 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_generation: 42, + }, + } + } + + #[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(); + 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) + )); + + // 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(); + 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(), + }) + ); + + // Malformed encoding. + let malformed = SignedEnvelope { + version: 1, + key_id: "test-key-1".into(), + payload: "***".into(), + signature: "***".into(), + }; + assert_eq!( + 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!( + prompt_token_digest(&[1, 2, 3]), + "7b0b5ea3ff36958c8e32ccf24b71da9ac68e51d0881bf75e62b837ec9ea6f3a5" + ); + assert_eq!( + body_digest(&envelope.request.normalized_body).unwrap(), + "b26a7dbf98b6cda6b6efad40e14948127007a81c1dacc0e18c5c71604a5504d9" + ); + } + + #[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 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_eq!(prompt_token_digest(&[]).len(), 64); + } + + /// 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] + 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); + + 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!(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(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/global_routing_transport.rs b/lib/llm/src/global_routing_transport.rs new file mode 100644 index 000000000000..f4cf28a492a9 --- /dev/null +++ b/lib/llm/src/global_routing_transport.rs @@ -0,0 +1,855 @@ +// 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, LazyLock, Mutex, OnceLock}; + +use async_trait::async_trait; +use base64::Engine; +use bytes::Bytes; +use futures::Stream; +use prometheus::{IntCounterVec, Opts, Registry}; +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"; + +/// 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}")] + 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)?; + + WAN_ATTEMPTS.with_label_values(&[region]).inc(); + + // 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| { + WAN_OUTCOMES + .with_label_values(&[region, "before_response", "failure"]) + .inc(); + DispatchError::BeforeResponse(error.to_string()) + })?; + let status = response.status(); + let headers = response.headers().clone(); + 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, + body: Box::pin(body), + }) + } +} + +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, + 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 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"); + let dispatch_sites = preprocessor + .match_indices("global_routing_transport::dispatch_signed(&signed)") + .map(|(offset, _)| offset) + .collect::>(); + 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") + .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(), + 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 3156fea15e87..42be9a7725e1 100644 --- a/lib/llm/src/http/service/openai.rs +++ b/lib/llm/src/http/service/openai.rs @@ -98,6 +98,71 @@ 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( + 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; + 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 = @@ -2894,6 +2959,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/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index dee46c0e9139..c9563cdef387 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -1212,6 +1212,19 @@ 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) = 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/lib.rs b/lib/llm/src/lib.rs index 37f048a3be6a..5ef9d2596b7c 100644 --- a/lib/llm/src/lib.rs +++ b/lib/llm/src/lib.rs @@ -17,7 +17,9 @@ pub mod first_token; pub mod fpm_publisher; pub mod fpm_trace; pub mod frontend_config; -pub mod global_routing_shadow; +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 5f2719b2c078..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 @@ -5386,7 +5404,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(); @@ -5457,13 +5475,15 @@ impl .await?; attach_agent_context_from_context(&mut common_request, &context); - crate::global_routing_shadow::observe( - &common_request.model, + 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( + supplied_prompt_token_digest, &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, @@ -5471,6 +5491,77 @@ impl prompt_injected_reasoning, )?; + // 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 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, + 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(), + })?; + 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"); let trace_state = crate::request_trace::build_request_end_trace_state( &common_request, @@ -5663,13 +5754,14 @@ 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_shadow::observe( - &common_request.model, + 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, - self.kv_cache_block_size, &request_id, - common_request.multi_modal_data.is_some() || common_request.mm_routing_info.is_some(), - ); + )?; let trace_state = crate::request_trace::build_request_end_trace_state( &common_request, @@ -5858,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, 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() } } 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,