diff --git a/experimental/sgl-router/src/server/error.rs b/experimental/sgl-router/src/server/error.rs index 253ba5e2c..b6445169a 100644 --- a/experimental/sgl-router/src/server/error.rs +++ b/experimental/sgl-router/src/server/error.rs @@ -17,6 +17,21 @@ pub enum ApiError { #[error("model not found: {0}")] ModelNotFound(String), + /// A request refused by the fleet-wide sampling contract + /// (`--override-sampling-params` under `--sampling-param-conflict + /// reject`). + /// + /// Distinct from [`Self::BadRequest`] on purpose: rolling a contract out + /// across a fleet turns previously-served client traffic into 400s, and + /// the operator's first question is how much and on which parameter. + /// Folded into `bad_request` that is unanswerable — the code would be the + /// same one malformed JSON and a missing `model` field already emit. + /// `param` is a `&'static str` from + /// [`crate::config::SamplingField::wire_name`], which keeps it usable as a + /// bounded metric label. + #[error("{param} violates this deployment's sampling contract: {detail}")] + SamplingContract { param: &'static str, detail: String }, + /// Could not reach the upstream worker (connect refused, DNS, TLS, request /// build error). `source` captures the full anyhow chain for server-side /// logging; clients see a generic message. @@ -114,6 +129,9 @@ impl ApiError { fn status_and_code(&self) -> (StatusCode, &'static str) { match self { ApiError::BadRequest(_) => (StatusCode::BAD_REQUEST, "bad_request"), + ApiError::SamplingContract { .. } => { + (StatusCode::BAD_REQUEST, "sampling_contract_violation") + } ApiError::ModelNotFound(_) => (StatusCode::NOT_FOUND, "model_not_found"), ApiError::UpstreamUnreachable { .. } => { (StatusCode::BAD_GATEWAY, "upstream_unreachable") @@ -242,7 +260,9 @@ impl IntoResponse for ApiError { ); "service unavailable".to_string() } - ApiError::BadRequest(_) | ApiError::ModelNotFound(_) => self.to_string(), + ApiError::BadRequest(_) + | ApiError::ModelNotFound(_) + | ApiError::SamplingContract { .. } => self.to_string(), }; let mut resp = ( status, @@ -442,4 +462,27 @@ mod tests { "ApiError::Internal must not leak anyhow chain to client; got: {body_str}" ); } + /// A sampling-contract rejection must not be filed under `bad_request`: + /// an operator rolling `--sampling-param-conflict reject` across a fleet + /// has to be able to alert on contract rejections without them being + /// indistinguishable from clients sending malformed JSON. + #[test] + fn sampling_contract_has_a_distinct_code_from_other_bad_requests() { + let err = ApiError::SamplingContract { + param: "temperature", + detail: "got 0.5, expected 1 (or omit the field)".into(), + }; + let (status, code_header, env) = parse_envelope(err.into_response()); + assert_eq!(status, StatusCode::BAD_REQUEST); + assert_eq!(code_header.as_deref(), Some("sampling_contract_violation")); + assert_eq!(env.error.code, "sampling_contract_violation"); + assert_ne!(env.error.code, "bad_request"); + // The parameter and both values reach the client: a 400 here is + // actionable without an operator explaining it. + assert!( + env.error.message.contains("temperature") && env.error.message.contains("0.5"), + "got: {}", + env.error.message + ); + } } diff --git a/experimental/sgl-router/src/server/metrics.rs b/experimental/sgl-router/src/server/metrics.rs index bcf6b0f35..68af70db3 100644 --- a/experimental/sgl-router/src/server/metrics.rs +++ b/experimental/sgl-router/src/server/metrics.rs @@ -42,6 +42,7 @@ //! | `sgl_router_cache_aware_decisions_total` | Counter | `model_id`, `decision` | //! | `sgl_router_diverted_overlap_blocks` | Histogram | `model_id` | //! | `sgl_router_ingress_tokenize_errors_total` | Counter | `model_id` | +//! | `sgl_router_sampling_contract_rejections_total` | Counter | `param` | //! //! `sgl_router_cache_aware_decisions_total` records exactly one decision per //! cache-aware prefill selection that resolves a worker, so the labels sum to @@ -346,6 +347,7 @@ pub struct MetricsRegistry { cache_aware_decisions_total: Mutex>>, diverted_overlap_blocks: Mutex>, ingress_tokenize_errors_total: Mutex>>, + sampling_contract_rejections_total: Mutex>>, } #[derive(Debug, Hash, Eq, PartialEq, Clone)] @@ -745,6 +747,24 @@ impl MetricsRegistry { counter.fetch_add(1, Ordering::Relaxed); } + /// Bump `sgl_router_sampling_contract_rejections_total{param}`. + /// + /// Recorded when the fleet-wide sampling contract refuses a request under + /// `--sampling-param-conflict reject`. This is the rollout gauge for the + /// flag: it answers "how much client traffic is the contract turning away, + /// and on which parameter" — which is otherwise unanswerable, because the + /// rejection reaches the client as a 400 like any other. `param` is a + /// wire name from a fixed enum, so the label set is bounded. + pub fn record_sampling_contract_rejection(&self, param: &'static str) { + let mut guard = self.sampling_contract_rejections_total.lock(); + let counter = guard + .entry(param) + .or_insert_with(|| Arc::new(AtomicU64::new(0))) + .clone(); + drop(guard); + counter.fetch_add(1, Ordering::Relaxed); + } + /// Render the registry as a Prometheus 0.0.4 exposition-format string /// with no live worker snapshot. The per-worker gauges emit only their /// HELP/TYPE headers and a zeroed pool-size series. Production scrapes @@ -1188,6 +1208,26 @@ impl MetricsRegistry { } drop(guard); + // sampling_contract_rejections_total + out.push_str( + "# HELP sgl_router_sampling_contract_rejections_total Requests refused by the fleet-wide sampling contract (--override-sampling-params under --sampling-param-conflict reject), by parameter.\n", + ); + out.push_str("# TYPE sgl_router_sampling_contract_rejections_total counter\n"); + let guard = self.sampling_contract_rejections_total.lock(); + let mut entries: Vec<(&str, u64)> = guard + .iter() + .map(|(k, v)| (*k, v.load(Ordering::Relaxed))) + .collect(); + entries.sort_by_key(|entry| entry.0); + for (param, value) in entries { + out.push_str(&format!( + "sgl_router_sampling_contract_rejections_total{{param=\"{}\"}} {}\n", + escape_label(param), + value, + )); + } + drop(guard); + out } } @@ -1739,4 +1779,29 @@ mod tests { "render did not escape backslash; got:\n{out}", ); } + /// The contract's rollout gauge: absent until a request is actually + /// refused, then keyed by the parameter that refused it. + #[test] + fn sampling_contract_rejections_are_keyed_by_param() { + let reg = MetricsRegistry::new(); + let out = reg.render(); + assert!(out.contains("# TYPE sgl_router_sampling_contract_rejections_total counter")); + assert!( + !out.contains("sgl_router_sampling_contract_rejections_total{"), + "must emit no series before the first rejection" + ); + + reg.record_sampling_contract_rejection("temperature"); + reg.record_sampling_contract_rejection("temperature"); + reg.record_sampling_contract_rejection("top_p"); + let out = reg.render(); + assert!( + out.contains(r#"sgl_router_sampling_contract_rejections_total{param="temperature"} 2"#), + "got:\n{out}" + ); + assert!( + out.contains(r#"sgl_router_sampling_contract_rejections_total{param="top_p"} 1"#), + "got:\n{out}" + ); + } } diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index f49564810..6b82fb108 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -1,7 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 -use crate::config::SessionAffinityMode; +use crate::config::{ + ConflictPolicy, ParamSpec, SamplingField, SamplingOverrides, SessionAffinityMode, +}; use crate::discovery::{ModelId, WorkerMode}; use crate::policies::kv_events::{compute_block_hashes, compute_block_hashes_bigram}; use crate::policies::registry::{PdPoolResolver, PdResolveError}; @@ -23,7 +25,6 @@ use axum::http::{HeaderMap, HeaderName, HeaderValue, Response}; use bytes::Bytes; use serde::de::IgnoredAny; use serde::Deserialize; -use std::collections::HashMap; use std::sync::Arc; /// Observability header carrying the final decode-pool URL for a @@ -54,35 +55,355 @@ const CHARS_PER_TOKEN_ESTIMATE: usize = 4; /// enforced by the `DefaultBodyLimit`, and returns 413 PAYLOAD_TOO_LARGE. pub const MAX_CHAT_BODY_BYTES: usize = 32 << 20; -/// Minimal probe over the request body — we only need the `stream` field -/// and the `model` field to decide between buffered vs SSE forwarding and -/// to select a worker. Deserializing into this struct (vs `serde_json::Value`) -/// does two things: +/// Minimal probe over the request body — the fields the ROUTER itself acts on, +/// plus the sampling parameters the fleet-wide contract governs. +/// Deserializing into this struct (vs `serde_json::Value`) does two things: /// -/// 1. Avoids the per-field heap allocation of `Value` for a multi-MiB body. +/// 1. Avoids building a `Value` tree over a multi-MiB body. NOTHING here +/// retains client-sized data: an unrecognized key is skipped through +/// `IgnoredAny`, and a sampling value that is not a number is drained the +/// same way (see [`ProbedValue`]). /// 2. Pins the contract: the body MUST be a JSON object. Degenerate /// shapes (`null`, `[]`, `"hi"`) fail at this step rather than being /// silently forwarded with `stream=false`. /// /// All other fields are ignored — the worker is authoritative for the /// full request schema. -#[derive(Debug, Deserialize)] +#[derive(Debug, Default)] struct RequestProbe { - #[serde(default)] stream: Option, - #[serde(default)] model: Option, /// Explicit output budget used by Decode Bucket routing. - #[serde(default)] max_tokens: Option, - #[serde(default)] max_completion_tokens: Option, + /// What the request said about each governed sampling parameter, indexed + /// by [`SamplingField::index`] so governing another needs no change here. + /// Read by [`apply_sampling_overrides`]; see [`ProbedValue`]. + sampling: [ProbedValue; SamplingField::ALL.len()], +} + +/// What the request said about one governed sampling parameter. +/// +/// WHY not `Option`: these keys carry client-controlled +/// JSON of arbitrary size and the probe runs on every request whether or not a +/// contract is configured, so holding a `Value` would let +/// `{"temperature": [1, 1, ...]}` allocate a tree proportional to a +/// [`MAX_CHAT_BODY_BYTES`] body for a field nothing then reads. Each variant is +/// resolved during deserialization; nothing client-sized is retained. +#[derive(Debug, Clone, Copy, PartialEq, Default)] +enum ProbedValue { + /// Omitted, or an explicit `null`. The OpenAI API types these parameters + /// as nullable with a documented default, so `null` asks for the default — + /// and on a governed fleet the configured value IS the default. Both mean + /// the configured value is injected. + #[default] + Absent, + /// A JSON number, a bool, or a string the engine's pydantic lax mode reads + /// as one — see [`parse_as_engine_number`]. + Number(f64), + /// Present, and NOT readable as a number by this probe. That is a + /// statement about the probe, not about the request: the engine's coercion + /// rules are laxer and undocumented, so a value landing here may still be + /// a number downstream. `reject` therefore refuses it rather than + /// forwarding it — see [`apply_sampling_overrides`]. + Unusable, +} + +/// Longest numeric string the probe will read. A sampling value is a short +/// literal, so the cap keeps a client-sized string off the parse path +/// entirely; an over-long one stays [`ProbedValue::Unusable`], which `reject` +/// refuses rather than forwards. +const MAX_SAMPLING_NUMERIC_LEN: usize = 64; + +/// Read a string the way the engine's pydantic lax mode reads it. +/// +/// Pydantic parses the TRIMMED string first; failing that it strips +/// underscores — refusing a leading one, a trailing one, or a doubled one — +/// and parses WITHOUT trimming. So `"1_0"` is 10 and `" 1.0 "` is 1, but +/// `" 1_0 "` is an error. That is not Python's own numeric-literal rule +/// either: pydantic takes `1._5`, `1e_5` and `-_1`, each a `SyntaxError` in +/// Python source. +/// +/// WHY this is written out rather than approximated: a contract that forwards +/// what it cannot parse is only as strong as this function's fidelity to a +/// transitive dependency's undocumented coercion table, across a fleet whose +/// engines need not even share a pydantic version. It is not, because +/// [`apply_sampling_overrides`] refuses the residue — this function only +/// decides how much of what the engine accepts is answered precisely instead +/// of with a 400. +fn parse_as_engine_number(s: &str) -> Option { + // Bounded BEFORE the first parse, not just before the underscore path: the + // string is client-controlled and can run to the body limit, and parsing + // one is linear in its length. Nothing that long is a sampling value. + let trimmed = s.trim(); + if trimmed.len() > MAX_SAMPLING_NUMERIC_LEN { + return None; + } + if let Ok(v) = trimmed.parse::() { + return Some(v); + } + if s.len() > MAX_SAMPLING_NUMERIC_LEN + || !s.contains('_') + || s.starts_with('_') + || s.ends_with('_') + || s.contains("__") + { + return None; + } + // Stripped into a fixed stack buffer: the value is client-controlled, and + // this type exists to keep client-sized allocations off the request path. + let mut buf = [0u8; MAX_SAMPLING_NUMERIC_LEN]; + let mut len = 0; + for &b in s.as_bytes() { + if b != b'_' { + buf[len] = b; + len += 1; + } + } + std::str::from_utf8(&buf[..len]).ok()?.parse().ok() +} + +impl<'de> Deserialize<'de> for ProbedValue { + fn deserialize>(d: D) -> Result { + struct ValueVisitor; + impl<'de> serde::de::Visitor<'de> for ValueVisitor { + type Value = ProbedValue; + + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.write_str("a sampling parameter value") + } + + fn visit_i64(self, v: i64) -> Result { + Ok(ProbedValue::Number(v as f64)) + } + + fn visit_u64(self, v: u64) -> Result { + Ok(ProbedValue::Number(v as f64)) + } + + fn visit_f64(self, v: f64) -> Result { + Ok(ProbedValue::Number(v)) + } + + /// Parsed, not retained: the engine reads a numeric string as a + /// number, so a `reject` contract must too or `"1.5"` slips past + /// a pin of 1. + fn visit_str(self, v: &str) -> Result { + Ok(parse_as_engine_number(v).map_or(ProbedValue::Unusable, ProbedValue::Number)) + } + + fn visit_unit(self) -> Result { + Ok(ProbedValue::Absent) + } + + /// The engine reads a JSON bool as a number, so a pin of 0 must + /// see `false` as the 0 the engine will sample with rather than as + /// something it cannot judge. + fn visit_bool(self, v: bool) -> Result { + Ok(ProbedValue::Number(if v { 1.0 } else { 0.0 })) + } + + /// Drained, never collected — the allocation this type exists to + /// avoid. + fn visit_seq>( + self, + mut seq: A, + ) -> Result { + while seq.next_element::()?.is_some() {} + Ok(ProbedValue::Unusable) + } + + /// Drained, never collected — as [`Self::visit_seq`]. + fn visit_map>( + self, + mut map: M, + ) -> Result { + while map.next_entry::()?.is_some() {} + Ok(ProbedValue::Unusable) + } + } + d.deserialize_any(ValueVisitor) + } +} + +/// A field the ROUTER acts on, as opposed to one it only forwards. A repeated +/// occurrence of one of these is a 400: a body that says two different things +/// about how to route itself is ambiguous at the edge, and the router must not +/// decide on one copy while the engine serves the other. +#[derive(Debug, Clone, Copy)] +enum RoutingKey { + Stream, + Model, + MaxTokens, + MaxCompletionTokens, +} + +impl RoutingKey { + /// Bit position in the visitor's occurrence mask. Unlike + /// [`SamplingField::index`] this indexes no array, so it needs no + /// agreement with a separate length constant. + const fn bit(self) -> u8 { + match self { + Self::Stream => 1 << 0, + Self::Model => 1 << 1, + Self::MaxTokens => 1 << 2, + Self::MaxCompletionTokens => 1 << 3, + } + } + + const fn wire_name(self) -> &'static str { + match self { + Self::Stream => "stream", + Self::Model => "model", + Self::MaxTokens => "max_tokens", + Self::MaxCompletionTokens => "max_completion_tokens", + } + } +} + +/// One key of the request object, resolved WITHOUT allocating: the visitor +/// matches a borrowed `&str` and keeps only a discriminant, so a body's +/// unrecognized majority costs no `String` per key. +enum ProbeKey { + Routing(RoutingKey), + /// A governed sampling parameter. A repeat LAST-WINS, matching the + /// engine's own `json.loads` — so the value the contract compares against + /// and the value the engine actually samples with are the same one. + Sampling(SamplingField), + Other, +} + +impl<'de> Deserialize<'de> for ProbeKey { + fn deserialize>(d: D) -> Result { + struct KeyVisitor; + impl serde::de::Visitor<'_> for KeyVisitor { + type Value = ProbeKey; + + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.write_str("a request field name") + } + + fn visit_str(self, v: &str) -> Result { + Ok(match v { + "stream" => ProbeKey::Routing(RoutingKey::Stream), + "model" => ProbeKey::Routing(RoutingKey::Model), + "max_tokens" => ProbeKey::Routing(RoutingKey::MaxTokens), + "max_completion_tokens" => ProbeKey::Routing(RoutingKey::MaxCompletionTokens), + other => match SamplingField::from_wire_name(other) { + Some(field) => ProbeKey::Sampling(field), + None => ProbeKey::Other, + }, + }) + } + } + d.deserialize_str(KeyVisitor) + } +} + +/// Reads the probed fields out of the request object. +/// +/// `read_sampling` is false only on [`probe_without_sampling_values`]'s +/// fallback pass, where a sampling value is scanned rather than converted. +struct ProbeVisitor { + read_sampling: bool, +} + +impl<'de> serde::de::Visitor<'de> for ProbeVisitor { + type Value = RequestProbe; + + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.write_str("a JSON object") + } + + fn visit_map>(self, mut map: M) -> Result { + let mut probe = RequestProbe::default(); + let mut seen = 0u8; + while let Some(key) = map.next_key::()? { + match key { + ProbeKey::Routing(r) => { + if seen & r.bit() != 0 { + return Err(serde::de::Error::custom(format_args!( + "duplicate field `{}`", + r.wire_name() + ))); + } + seen |= r.bit(); + match r { + RoutingKey::Stream => probe.stream = map.next_value()?, + RoutingKey::Model => probe.model = map.next_value()?, + RoutingKey::MaxTokens => probe.max_tokens = map.next_value()?, + RoutingKey::MaxCompletionTokens => { + probe.max_completion_tokens = map.next_value()? + } + } + } + ProbeKey::Sampling(field) => { + probe.sampling[field.index()] = if self.read_sampling { + map.next_value()? + } else { + // Scanned, not converted — the whole point of + // this pass. The key IS present, so it is + // unreadable rather than absent: a contract + // must not inject over a value the client sent. + map.next_value::()?; + ProbedValue::Unusable + }; + } + ProbeKey::Other => { + map.next_value::()?; + } + } + } + Ok(probe) + } +} + +impl<'de> Deserialize<'de> for RequestProbe { + fn deserialize>(d: D) -> Result { + // `deserialize_map` is what pins the body to a JSON object: `null`, + // `[]` and `"hi"` are all rejected here, so no separate shape-anchoring + // pass is needed. + d.deserialize_map(ProbeVisitor { + read_sampling: true, + }) + } +} + +/// Re-read a body without converting its sampling values. +/// +/// serde_json converts a number literal while resolving [`ProbedValue`], so a +/// governed key holding one outside `f64`'s range — `{"temperature": 1e400}` — +/// fails the WHOLE parse, where the `IgnoredAny` that covered these keys +/// before the contract existed scanned past it. Left alone that 400s a body +/// the router used to forward, on a request nothing has opted in to, with a +/// message ("body must be a JSON object") that is not even true of it. +/// +/// So the probe falls back to this pass, which reads the routing keys exactly +/// as before and marks each sampling key present-but-unreadable. A governed +/// `reject` then refuses it, which is the same answer it gives any other value +/// it cannot represent. +fn probe_without_sampling_values(body: &[u8]) -> Result { + let mut de = serde_json::Deserializer::from_slice(body); + let probe = serde::Deserializer::deserialize_map( + &mut de, + ProbeVisitor { + read_sampling: false, + }, + )?; + de.end()?; + Ok(probe) } impl RequestProbe { fn requested_max_output_tokens(&self) -> Option { self.max_completion_tokens.or(self.max_tokens) } + + /// Lets [`apply_sampling_overrides`] loop over whatever the operator + /// configured instead of repeating a per-field ladder. + fn sampling_field(&self, field: SamplingField) -> ProbedValue { + self.sampling[field.index()] + } } /// RAII guard that records `sgl_router_request_duration_seconds` when @@ -134,11 +455,15 @@ pub async fn chat_completions( body: Bytes, ) -> Result, ApiError> { let start = std::time::Instant::now(); - let probe = parse_probe(&body)?; + let mut probe = parse_probe(&body)?; let streaming = probe.stream.unwrap_or(false); let requested_max_output_tokens = probe.requested_max_output_tokens(); + // `take`n rather than borrowed: the sampling contract further down reads + // the rest of `probe`, but nothing reads `model` again, so the `String` + // moves out instead of being cloned on every request. let model_str = probe .model + .take() .ok_or_else(|| ApiError::BadRequest("missing `model` field".into()))?; let model_id = ModelId(model_str.clone()); @@ -167,6 +492,16 @@ pub async fn chat_completions( .get(&model_id) .ok_or_else(|| ApiError::ModelNotFound(model_str.clone()))?; + // Fleet-wide sampling contract (`--override-sampling-params`), applied + // once the model is known to be served (so a request naming an unknown + // model still gets that answer) and before anything is admitted: under + // `reject` a numeric value differing from the configured one is a 400 + // here, costing no queue slot and no engine round-trip. Either way the + // configured values for fields the request omitted come back as the + // inject-set for the forwarded body. + let inject_sampling = + apply_sampling_overrides(&ctx.config.model.sampling_overrides, &probe, &ctx.metrics)?; + // Tokenize once at ingress whenever it can pay off — decoupled from the // routing policy, because forwarding `input_ids` is a property of the // MODEL (does it have a chat encoder so the router can produce @@ -494,10 +829,15 @@ pub async fn chat_completions( let bootstrap_room = bootstrap.as_ref().map(|b| b.room); // Build the body forwarded to the engine(s) exactly once — injecting the - // `input_ids` and/or bootstrap fields, or forwarding the original bytes - // untouched when neither applies. - let outgoing_body = - build_outgoing_body(&body, request_value, forward_input_ids, bootstrap.as_ref())?; + // `input_ids`, bootstrap fields and sampling values, or forwarding the + // original bytes untouched when none applies. + let outgoing_body = build_outgoing_body( + &body, + request_value, + forward_input_ids, + bootstrap.as_ref(), + &inject_sampling, + )?; let result = if let Some(decode_worker) = decode_peer { // PD-disagg dispatch (Pattern B — spawn prefill, await decode). @@ -901,37 +1241,36 @@ struct BootstrapFields { } /// Build the body forwarded to the engine, injecting (when present) the -/// precomputed `input_ids` and/or the PD `bootstrap_*` fields into the -/// already-parsed request object and serializing once. When neither is -/// needed, returns the original bytes unchanged (no re-serialize). +/// precomputed `input_ids`, the PD `bootstrap_*` fields and the fleet-wide +/// sampling values into the already-parsed request object and serializing +/// once. When none is needed, returns the original bytes unchanged (no +/// re-serialize). /// -/// `input_ids`: the router-computed prompt tokens. When set, the engine skips -/// its own chat-template tokenization; `messages` are retained in the body so -/// the engine still derives stop tokens / tool-call constraint and the OpenAI -/// response shape. The caller sets this only when the tokens are -/// engine-equivalent and `input_ids_safe_to_forward` held. +/// `input_ids`: the router-computed prompt tokens. When set the engine skips +/// its own chat-template tokenization, so `messages` are retained for the stop +/// tokens, tool-call constraint and response shape it still derives from them. +/// Set only when `input_ids_safe_to_forward` held. /// -/// `value` is the already-parsed request body when one is on hand (the -/// cache-aware path parses once at ingress); it is consumed so the mutation -/// reuses that parse. It is `None` only for a load-only policy in PD mode — a -/// path that never parses at ingress — so the bootstrap injection re-parses -/// the bytes here (matching the pre-refactor behavior). The body shape was -/// validated by `parse_probe`; the non-object arm defends against a TOCTOU -/// regression rather than panicking. +/// `value` is the ingress parse when one is on hand (the cache-aware path +/// parses once at ingress); it is consumed so the mutation reuses that parse. +/// It is `None` for a load-only policy — a path that never parses at ingress — +/// so injection re-parses the bytes here. The body shape was validated by +/// `parse_probe`; the non-object arm defends against a TOCTOU regression +/// rather than panicking. fn build_outgoing_body( body: &Bytes, value: Option, input_ids: Option<&[u32]>, bootstrap: Option<&BootstrapFields>, + sampling: &[(SamplingField, serde_json::Number)], ) -> Result { - if input_ids.is_none() && bootstrap.is_none() { + if input_ids.is_none() && bootstrap.is_none() && sampling.is_empty() { // Nothing to inject — forward the original bytes (cheap Arc clone). return Ok(body.clone()); } let parsed = match value { Some(v) => v, - // Load-only + PD: the ingress skipped the parse, so re-parse for the - // bootstrap injection (input_ids is never set on this path). + // The ingress skipped the parse, so re-parse for the injection. None => serde_json::from_slice(body).map_err(|_| { ApiError::BadRequest("invalid request: body must be a JSON object".to_string()) })?, @@ -944,6 +1283,15 @@ fn build_outgoing_body( )); } }; + // The caller passes the inject-set from `apply_sampling_overrides` — + // configured values for fields the request omitted — so writing them here + // never masks a client value, in either conflict mode. + for (field, value) in sampling { + obj.insert( + field.wire_name().to_string(), + serde_json::Value::Number(value.clone()), + ); + } if let Some(ids) = input_ids { obj.insert( "input_ids".to_string(), @@ -1114,6 +1462,120 @@ fn request_is_multimodal(value: &serde_json::Value) -> bool { }) } +/// Apply the fleet-wide sampling contract (`--override-sampling-params` / +/// `--sampling-param-conflict`) to one request, before admission. +/// +/// Returns the inject-set: the configured value for every exact-valued +/// parameter the request OMITTED, which [`build_outgoing_body`] writes into +/// the forwarded body so the engine's own defaults can't drift from what the +/// operator declared. A band ([`ParamSpec::Range`]) names no single value, so +/// it never injects. +/// +/// For a parameter the request DID send: +/// * [`ConflictPolicy::Allow`] forwards the client value untouched, which +/// is why the mode check comes before any comparison; +/// * [`ConflictPolicy::Reject`] 400s a numeric value that differs from the +/// configured one (or falls outside the band) — never a silent rewrite, +/// which is the one behavior no client can detect; +/// * [`ConflictPolicy::Reject`] also 400s a value this probe cannot read as +/// a number ([`ProbedValue::Unusable`]). `reject` is a promise that +/// nothing but the configured value reaches the engine, and the engine's +/// coercion rules are laxer than [`parse_as_engine_number`] and +/// undocumented — a bool and an underscored numeric string were both once +/// numbers to the engine and unreadable here. Forwarding the residue +/// makes the contract only as strong as this probe's fidelity to a +/// transitive Python dependency, so the residue is refused instead. +/// Under `allow` it keeps flowing, because `allow` makes no promise to +/// break. +/// +/// A rejection is counted per parameter before it is returned, because a +/// contract rollout turns served traffic into 400s and the operator needs to +/// see how much and where. +fn apply_sampling_overrides( + overrides: &SamplingOverrides, + probe: &RequestProbe, + metrics: &MetricsRegistry, +) -> Result, ApiError> { + // Length is a startup constant, and `Vec::with_capacity(0)` does not + // allocate — so an unconfigured contract still costs nothing. + let mut inject = Vec::with_capacity(overrides.params.len()); + // Every configured parameter is judged even after one has failed. The + // counter is how an operator sizes a rollout's blast radius per parameter, + // and stopping at the first violation would report zero for every + // parameter that sorts after it — a fleet violating both `temperature` and + // `top_p` on every request would look like it violates only `temperature`. + // The client is still told about one parameter, so the 400 stays one + // sentence. + let mut first_violation: Option = None; + for (&field, spec) in &overrides.params { + let name = field.wire_name(); + let violation = |detail: String, first: &mut Option| { + metrics.record_sampling_contract_rejection(name); + let err = ApiError::SamplingContract { + param: name, + detail, + }; + if first.is_none() { + *first = Some(err); + } + }; + let got = match probe.sampling_field(field) { + // A band names no single value, so it never injects — a request + // that omits the parameter gets the engine's own default. + ProbedValue::Absent => { + if let ParamSpec::Exact(v) = spec { + inject.push((field, v.clone())); + } + continue; + } + // `allow` forwards a client value untouched, so everything below + // is `reject`-only. + _ if overrides.conflict == ConflictPolicy::Allow => continue, + // A value the router cannot read as a number is a value it cannot + // prove conforms. Under `reject` that is a refusal, not a pass. + ProbedValue::Unusable => { + violation( + match spec { + ParamSpec::Exact(want) => { + format!("expected {want} (or omit the field), got a non-numeric value") + } + &ParamSpec::Range { lo, hi } => { + format!( + "must be a number between {lo} and {hi}, got a non-numeric value" + ) + } + }, + &mut first_violation, + ); + continue; + } + ProbedValue::Number(got) => got, + }; + match spec { + ParamSpec::Exact(want) => { + if Some(got) != want.as_f64() { + violation( + format!("got {got}, expected {want} (or omit the field)"), + &mut first_violation, + ); + } + } + &ParamSpec::Range { lo, hi } => { + if !(lo..=hi).contains(&got) { + violation( + format!("must be between {lo} and {hi}, got {got}"), + &mut first_violation, + ); + } + } + } + } + match first_violation { + Some(err) => Err(err), + None => Ok(inject), + } +} + fn parse_probe(body: &Bytes) -> Result { // We deliberately do NOT echo the serde error into the client-visible // message — that risks leaking field-level detail and is also of little @@ -1121,24 +1583,27 @@ fn parse_probe(body: &Bytes) -> Result { // Server-side, the full error is logged with `tracing::debug!` for // operator triage. // - // Two-step deserialize: - // 1. `Map` *anchors* the shape to a JSON object. - // This rejects `null` / `[]` / `"hi"` (all valid JSON but not - // request shape) without walking the full value into a - // `serde_json::Value` per field. - // 2. `RequestProbe` (struct of `Option` + `Option`) - // lifts out only the fields we care about — `stream` and `model`. - // Other fields are ignored; the worker is authoritative for the - // rest of the schema. - let _: HashMap = serde_json::from_slice(body).map_err(|e| { - tracing::debug!(error = %e, "chat-completions body rejected as non-object JSON"); - ApiError::BadRequest("invalid request: body must be a JSON object".to_string()) - })?; - let probe: RequestProbe = serde_json::from_slice(body).map_err(|e| { - tracing::debug!(error = %e, "chat-completions request-probe deserialize failed"); - ApiError::BadRequest("invalid request: body must be a JSON object".to_string()) - })?; - Ok(probe) + // ONE deserialize. [`RequestProbe`]'s hand-written `visit_map` both + // anchors the shape (its `deserialize_map` rejects `null` / `[]` / `"hi"` + // — valid JSON, not request shape) and lifts out the probed fields, + // skipping the unknown majority through `IgnoredAny`. It never builds a + // `serde_json::Value` and allocates nothing per unrecognized key, so the + // shape check costs no separate pass over a multi-MiB body. + let err = match serde_json::from_slice::(body) { + Ok(probe) => return Ok(probe), + Err(e) => e, + }; + // The one failure this pass can invent that the fallback cannot is + // converting a sampling number literal; malformed JSON, a non-object body + // and a duplicated routing key all fail both, so retrying here cannot + // launder a body that is genuinely bad. + if let Ok(probe) = probe_without_sampling_values(body) { + return Ok(probe); + } + tracing::debug!(error = %err, "chat-completions request-probe deserialize failed"); + Err(ApiError::BadRequest( + "invalid request: body must be a JSON object".to_string(), + )) } #[cfg(test)] @@ -1213,7 +1678,8 @@ mod tests { port: None, room: 42, }; - let injected = build_outgoing_body(&body, Some(value), None, Some(&bootstrap)).unwrap(); + let injected = + build_outgoing_body(&body, Some(value), None, Some(&bootstrap), &[]).unwrap(); let parsed: serde_json::Value = serde_json::from_slice(&injected).unwrap(); assert_eq!(parsed.get("bootstrap_port"), Some(&serde_json::Value::Null)); assert_eq!( @@ -1234,7 +1700,7 @@ mod tests { Bytes::from_static(br#"{"model":"x","messages":[{"role":"user","content":"hi"}]}"#); let value: serde_json::Value = serde_json::from_slice(&body).unwrap(); let ids = [1u32, 2, 3]; - let out = build_outgoing_body(&body, Some(value), Some(&ids), None).unwrap(); + let out = build_outgoing_body(&body, Some(value), Some(&ids), None, &[]).unwrap(); let parsed: serde_json::Value = serde_json::from_slice(&out).unwrap(); assert_eq!(parsed.get("input_ids"), Some(&serde_json::json!([1, 2, 3]))); assert!( @@ -1249,7 +1715,7 @@ mod tests { fn build_outgoing_body_no_injection_returns_original_bytes() { let body = Bytes::from_static(br#"{"model":"x","messages":[]}"#); let value: serde_json::Value = serde_json::from_slice(&body).unwrap(); - let out = build_outgoing_body(&body, Some(value), None, None).unwrap(); + let out = build_outgoing_body(&body, Some(value), None, None, &[]).unwrap(); assert_eq!( out, body, "no injection must forward the original bytes unchanged" @@ -1269,7 +1735,8 @@ mod tests { port: Some(9), room: 5, }; - let out = build_outgoing_body(&body, Some(value), Some(&ids), Some(&bootstrap)).unwrap(); + let out = + build_outgoing_body(&body, Some(value), Some(&ids), Some(&bootstrap), &[]).unwrap(); let parsed: serde_json::Value = serde_json::from_slice(&out).unwrap(); assert_eq!(parsed.get("input_ids"), Some(&serde_json::json!([7, 8]))); assert_eq!( @@ -1364,7 +1831,7 @@ mod tests { port: Some(1), room: 2, }; - let out = build_outgoing_body(&body, None, None, Some(&bootstrap)).unwrap(); + let out = build_outgoing_body(&body, None, None, Some(&bootstrap), &[]).unwrap(); let parsed: serde_json::Value = serde_json::from_slice(&out).unwrap(); assert_eq!( parsed.get("bootstrap_room"), @@ -1555,4 +2022,662 @@ mod tests { other => panic!("expected BadRequest, got {other:?}"), } } + + fn probe_of(body: &str) -> RequestProbe { + parse_probe(&Bytes::copy_from_slice(body.as_bytes())).unwrap() + } + + /// Build a [`SamplingOverrides`] the only way production does — through + /// the flag parser. + /// + /// Hand-building the struct here would re-implement `canonical_number`'s + /// integral normalization in the test, so a regression in the parser would + /// leave these assertions green while the forwarded body changed. + fn overrides_of(conflict: ConflictPolicy, json: &str) -> SamplingOverrides { + crate::config::parse_sampling_overrides(json, conflict).expect("test config must parse") + } + + /// A throwaway registry, so the contract's rejection counter has somewhere + /// to land. + fn metrics() -> Arc { + MetricsRegistry::new() + } + + /// With nothing configured the sampling contract is inert: no request is + /// ever inspected, and nothing is injected. + #[test] + fn unconfigured_sampling_overrides_inject_nothing() { + let overrides = SamplingOverrides::default(); + let p = probe_of(r#"{"model":"x","temperature":0.7,"n":4}"#); + assert_eq!( + apply_sampling_overrides(&overrides, &p, &metrics()).unwrap(), + vec![] + ); + } + + /// The `reject` contract at the decision level: omitted -> inject the + /// configured value; equal to it -> pass untouched; any other numeric + /// value -> 400; non-numeric garbage -> forwarded for the engine's own + /// schema error. A band admits its range, 400s outside it, and injects + /// nothing. + #[test] + fn reject_mode_pins_configured_values_and_admits_a_band() { + let overrides = overrides_of( + ConflictPolicy::Reject, + r#"{"top_p": 0.95, "frequency_penalty": 0.0, "presence_penalty": 0.0, + "n": 1, "temperature": {"min": 0, "max": 1}}"#, + ); + + // Omitted params: accepted, exact values injected, band injects nothing. + let p = probe_of(r#"{"model":"x","messages":[]}"#); + let inject = apply_sampling_overrides(&overrides, &p, &metrics()).unwrap(); + assert_eq!( + inject + .iter() + .map(|(f, v)| (f.wire_name(), v.to_string())) + .collect::>(), + vec![ + ("top_p", "0.95".to_string()), + ("frequency_penalty", "0.0".to_string()), + ("presence_penalty", "0.0".to_string()), + ("n", "1".to_string()), + ] + ); + + for accepted in [ + r#"{"model":"x","temperature":0.0}"#, + r#"{"model":"x","temperature":0.6}"#, + r#"{"model":"x","temperature":1.0}"#, + r#"{"model":"x","top_p":0.95}"#, + r#"{"model":"x","presence_penalty":0}"#, + r#"{"model":"x","frequency_penalty":0}"#, + r#"{"model":"x","n":1}"#, + ] { + let p = probe_of(accepted); + apply_sampling_overrides(&overrides, &p, &metrics()) + .unwrap_or_else(|e| panic!("{accepted} must be accepted: {e:?}")); + } + + // Nothing is injected over a field the client already sent. + let p = probe_of(r#"{"model":"x","top_p":0.95}"#); + let inject = apply_sampling_overrides(&overrides, &p, &metrics()).unwrap(); + assert!(!inject.iter().any(|(f, _)| *f == SamplingField::TopP)); + + for rejected in [ + r#"{"model":"x","temperature":1.1}"#, + r#"{"model":"x","temperature":2.0}"#, + r#"{"model":"x","temperature":-0.1}"#, + r#"{"model":"x","top_p":0.8}"#, + r#"{"model":"x","presence_penalty":0.5}"#, + r#"{"model":"x","frequency_penalty":0.5}"#, + r#"{"model":"x","n":2}"#, + ] { + let p = probe_of(rejected); + let err = apply_sampling_overrides(&overrides, &p, &metrics()) + .expect_err(&format!("{rejected} must be rejected")); + assert!( + matches!(err, ApiError::SamplingContract { .. }), + "{rejected}: got {err:?}" + ); + } + + // A bool is a number to the engine, so it is judged like one: `n: true` + // IS the configured `n: 1` and passes. + let p = probe_of(r#"{"model":"x","n":true}"#); + let inject = apply_sampling_overrides(&overrides, &p, &metrics()).unwrap(); + assert!(!inject.iter().any(|(f, _)| *f == SamplingField::N)); + + // Garbage this probe cannot read as a number is refused under `reject` + // rather than forwarded: `reject` promises the engine sees nothing but + // the configured value, and the engine's coercion rules are laxer than + // ours, so "not a number here" does not mean "not a number there". + let p = probe_of(r#"{"model":"x","top_p":"hot"}"#); + let err = apply_sampling_overrides(&overrides, &p, &metrics()) + .expect_err("an unreadable value must not slip past a pin"); + assert!(matches!(err, ApiError::SamplingContract { .. }), "{err:?}"); + + // Under `allow` it keeps flowing — `allow` makes no promise to break — + // and is still never injected over. + let allow = overrides_of(ConflictPolicy::Allow, r#"{"top_p": 0.95}"#); + let inject = apply_sampling_overrides(&allow, &p, &metrics()).unwrap(); + assert!(inject.is_empty(), "a client value is never overwritten"); + + // Numeric strings coerce the way the engine's pydantic lax mode does: + // "0.95" equals the configured value, "0.8" differs and 400s here. + let p = probe_of(r#"{"model":"x","top_p":"0.95"}"#); + assert!(apply_sampling_overrides(&overrides, &p, &metrics()).is_ok()); + let p = probe_of(r#"{"model":"x","top_p":"0.8"}"#); + assert!(apply_sampling_overrides(&overrides, &p, &metrics()).is_err()); + + // An exact temperature (no band) rejects differing values and injects + // when absent, like every other parameter. + let exact = overrides_of(ConflictPolicy::Reject, r#"{"temperature": 1.0}"#); + let p = probe_of(r#"{"model":"x","temperature":0.6}"#); + assert!(apply_sampling_overrides(&exact, &p, &metrics()).is_err()); + let p = probe_of(r#"{"model":"x"}"#); + assert_eq!( + apply_sampling_overrides(&exact, &p, &metrics()).unwrap(), + vec![( + SamplingField::Temperature, + serde_json::Number::from_f64(1.0).unwrap() + )] + ); + } + + /// `allow` keeps the fill-when-absent half of the contract and drops the + /// rejection half: a client value — right, wrong or garbage — is forwarded + /// untouched, so the configured values are fleet-wide defaults. + #[test] + fn allow_mode_never_rejects_and_never_masks_a_client_value() { + let overrides = overrides_of( + ConflictPolicy::Allow, + r#"{"temperature": 1, "top_p": 0.95, "n": 1}"#, + ); + + // Omitted -> injected, exactly as under `reject`. + let p = probe_of(r#"{"model":"x"}"#); + assert_eq!( + apply_sampling_overrides(&overrides, &p, &metrics()) + .unwrap() + .iter() + .map(|(f, _)| f.wire_name()) + .collect::>(), + vec!["temperature", "top_p", "n"] + ); + + // Every value `reject` would 400 is accepted here, and nothing is + // injected over it — the client's value reaches the engine. + let p = probe_of(r#"{"model":"x","temperature":0.6,"top_p":0.8,"n":4}"#); + assert_eq!( + apply_sampling_overrides(&overrides, &p, &metrics()).unwrap(), + vec![] + ); + + // Partial overlap: the client set temperature, so only the untouched + // parameters are filled in. + let p = probe_of(r#"{"model":"x","temperature":0.6}"#); + assert_eq!( + apply_sampling_overrides(&overrides, &p, &metrics()) + .unwrap() + .iter() + .map(|(f, _)| f.wire_name()) + .collect::>(), + vec!["top_p", "n"] + ); + } + + /// An explicit `null` is absent for the engine, so it is absent here too: + /// the configured value is injected rather than the field being read as a + /// client-supplied conflict. + #[test] + fn explicit_null_sampling_value_counts_as_omitted() { + let overrides = overrides_of(ConflictPolicy::Reject, r#"{"temperature": 1}"#); + let p = probe_of(r#"{"model":"x","temperature":null}"#); + assert_eq!( + apply_sampling_overrides(&overrides, &p, &metrics()) + .unwrap() + .iter() + .map(|(f, _)| f.wire_name()) + .collect::>(), + vec!["temperature"] + ); + } + + /// The inject-set lands in the forwarded body, and rides the same + /// `build_outgoing_body` serialize as the forwarded `input_ids` — so + /// chat-encoder traffic gets the contract too, not just the raw-prompt + /// path. + #[test] + fn build_outgoing_body_injects_sampling_overrides_alongside_input_ids() { + let body = + Bytes::from_static(br#"{"model":"x","messages":[{"role":"user","content":"hi"}]}"#); + let value: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let ids = [1u32, 2, 3]; + let overrides = overrides_of( + ConflictPolicy::Reject, + r#"{"top_p": 0.95, "top_k": 1000, "frequency_penalty": 0.0, + "presence_penalty": 0.0, "n": 1}"#, + ); + let inject = apply_sampling_overrides( + &overrides, + &probe_of(r#"{"model":"x","messages":[{"role":"user","content":"hi"}]}"#), + &metrics(), + ) + .unwrap(); + let out = build_outgoing_body(&body, Some(value), Some(&ids), None, &inject).unwrap(); + let parsed: serde_json::Value = serde_json::from_slice(&out).unwrap(); + assert_eq!(parsed.get("top_p"), Some(&serde_json::json!(0.95))); + // `top_k` and `n` are engine-typed `int`: the injected literals must + // not be `1000.0` / `1.0`. + assert_eq!(parsed.get("top_k"), Some(&serde_json::json!(1000))); + assert_eq!(parsed.get("n"), Some(&serde_json::json!(1))); + assert_eq!( + parsed.get("frequency_penalty"), + Some(&serde_json::json!(0.0)) + ); + assert_eq!( + parsed.get("presence_penalty"), + Some(&serde_json::json!(0.0)) + ); + assert_eq!(parsed.get("input_ids"), Some(&serde_json::json!([1, 2, 3]))); + assert!(parsed.get("messages").is_some()); + } + /// A repeated sampling key LAST-WINS instead of 400ing, matching the + /// engine's own `json.loads`: the value the contract compares against must + /// be the value the engine will actually sample with. Probing these fields + /// must not turn a body the router used to forward into a rejection. + #[test] + fn duplicate_sampling_key_takes_the_last_value_like_the_engine() { + let overrides = overrides_of(ConflictPolicy::Reject, r#"{"temperature": 1}"#); + + // Last value matches the contract -> accepted. + let p = probe_of(r#"{"model":"x","temperature":0.5,"temperature":1}"#); + assert!(apply_sampling_overrides(&overrides, &p, &metrics()).is_ok()); + + // Last value differs -> rejected on THAT value, not the first one. + let p = probe_of(r#"{"model":"x","temperature":1,"temperature":0.5}"#); + let err = apply_sampling_overrides(&overrides, &p, &metrics()).unwrap_err(); + assert!( + format!("{err}").contains("got 0.5"), + "must judge the last value; got {err}" + ); + } + + /// The router-actionable fields keep the stricter pre-existing contract: + /// a body that says two different things about how to route itself is + /// ambiguous at the edge. See + /// `parse_probe_handles_duplicate_stream_keys`. + #[test] + fn duplicate_routing_key_still_rejects() { + for body in [ + r#"{"model":"a","model":"b"}"#, + r#"{"stream":true,"stream":false}"#, + r#"{"max_tokens":1,"max_tokens":2}"#, + r#"{"max_completion_tokens":1,"max_completion_tokens":2}"#, + // An explicit null still counts as an occurrence, so this is a + // duplicate even though the first value reads as `None`. + r#"{"stream":null,"stream":true}"#, + ] { + let b = Bytes::copy_from_slice(body.as_bytes()); + assert!( + parse_probe(&b).is_err(), + "{body} must be rejected as ambiguous" + ); + } + } + + /// A sampling key carrying a huge non-numeric value must cost nothing: it + /// is drained, not materialized, and it does not fail the probe. Probing + /// these keys must neither put a client-sized allocation on the request + /// path nor reject a body an ungoverned router forwards. + #[test] + fn oversized_non_numeric_sampling_value_is_drained_not_materialized() { + let big_array = format!("[{}]", "1,".repeat(50_000) + "1"); + let big_string = format!("\"{}\"", "x".repeat(200_000)); + let big_object = format!("{{{}\"k\":1}}", "\"j\":[[[1]]],".repeat(10_000)); + for value in [&big_array, &big_string, &big_object] { + let body = format!(r#"{{"model":"x","temperature":{value}}}"#); + let probe = probe_of(&body); + assert_eq!( + probe.sampling_field(SamplingField::Temperature), + ProbedValue::Unusable, + "a non-numeric value must collapse to Unusable" + ); + // ...and it is never injected over, in either mode: the client + // sent something, so there is no omission to fill. + let allow = overrides_of(ConflictPolicy::Allow, r#"{"temperature": 1}"#); + let inject = apply_sampling_overrides(&allow, &probe, &metrics()).unwrap(); + assert!(inject.is_empty(), "must not inject over a client value"); + let reject = overrides_of(ConflictPolicy::Reject, r#"{"temperature": 1}"#); + assert!( + apply_sampling_overrides(&reject, &probe, &metrics()).is_err(), + "reject must refuse a value it cannot read as a number" + ); + } + } + + /// A numeric string is read the way the engine's pydantic lax mode reads + /// it, but is not retained as a string. + #[test] + fn probed_sampling_values_normalize_to_numbers() { + let p = probe_of(r#"{"model":"x","temperature":" 0.7 ","top_k":40,"min_p":0.05}"#); + assert_eq!( + p.sampling_field(SamplingField::Temperature), + ProbedValue::Number(0.7) + ); + assert_eq!( + p.sampling_field(SamplingField::TopK), + ProbedValue::Number(40.0) + ); + assert_eq!( + p.sampling_field(SamplingField::MinP), + ProbedValue::Number(0.05) + ); + assert_eq!(p.sampling_field(SamplingField::N), ProbedValue::Absent); + } + + /// A contract rejection must be distinguishable from every other 400 — + /// malformed JSON, a missing `model`, a bad SLO header — and must name the + /// parameter, since that is what an operator rolling the flag out needs. + #[test] + fn contract_rejection_has_its_own_error_code_and_counter() { + let overrides = overrides_of(ConflictPolicy::Reject, r#"{"top_p": 0.95}"#); + let metrics = metrics(); + let p = probe_of(r#"{"model":"x","top_p":0.5}"#); + + let err = apply_sampling_overrides(&overrides, &p, &metrics).unwrap_err(); + let ApiError::SamplingContract { param, .. } = &err else { + panic!("expected SamplingContract, got {err:?}"); + }; + assert_eq!(*param, "top_p"); + // The parameter and the offending value both reach the client, so a + // 400 is self-explanatory without an operator in the loop. The + // distinct `x-router-error-code` is pinned in `server::error`. + let msg = format!("{err}"); + assert!(msg.contains("top_p") && msg.contains("0.5"), "got {msg}"); + + apply_sampling_overrides(&overrides, &p, &metrics).unwrap_err(); + assert!( + metrics + .render() + .contains(r#"sgl_router_sampling_contract_rejections_total{param="top_p"} 2"#), + "rejections must be counted per parameter:\n{}", + metrics.render() + ); + } + + /// Nothing configured -> the body is forwarded as the same `Bytes`, with + /// neither a parse nor a copy. + #[test] + fn build_outgoing_body_without_injection_forwards_the_same_allocation() { + let body = Bytes::from_static(br#"{"model":"x"}"#); + let out = build_outgoing_body(&body, None, None, None, &[]).unwrap(); + assert_eq!( + out.as_ptr(), + body.as_ptr(), + "must be an Arc clone, not a copy" + ); + } + + /// The engine reads each of these as a number, so a contract that waved + /// them through would pin nothing. `false` is 0 and `"0.5_0"` is 0.5 to + /// pydantic — a pin of 1 must refuse both, and a pin of the value they + /// coerce to must accept them, since the engine samples with exactly + /// that. + #[test] + fn values_the_engine_reads_as_numbers_are_judged_not_waved_through() { + for (body_value, engine_sees) in [("false", 0.0), ("true", 1.0), (r#""0.5_0""#, 0.5)] { + let probe = probe_of(&format!(r#"{{"model":"x","temperature":{body_value}}}"#)); + assert_eq!( + probe.sampling_field(SamplingField::Temperature), + ProbedValue::Number(engine_sees), + "{body_value} must be read as the number the engine will use" + ); + + let pinned_elsewhere = overrides_of(ConflictPolicy::Reject, r#"{"temperature": 2}"#); + assert!( + apply_sampling_overrides(&pinned_elsewhere, &probe, &metrics()).is_err(), + "{body_value} differs from the pin and must be rejected" + ); + + let pinned_here = overrides_of( + ConflictPolicy::Reject, + &format!(r#"{{"temperature": {engine_sees}}}"#), + ); + assert!( + apply_sampling_overrides(&pinned_here, &probe, &metrics()).is_ok(), + "{body_value} IS the pinned value to the engine, so it must pass" + ); + } + } + + /// `parse_as_engine_number` against the engine's actual answers. + /// + /// Every expectation here was produced by running the value through + /// pydantic 2.13.5 on sglang's own field declaration + /// (`temperature: Optional[float] = None`, no validator, no strict + /// config), not derived from a reading of the rules — deriving them is + /// what gets this wrong. Python's numeric-literal rule, the obvious + /// guess, disagrees with pydantic on `1._5`, `1_.5`, `1e_5`, `1_e5` and + /// `-_1`. + #[test] + fn numeric_strings_are_read_the_way_the_engine_reads_them() { + #[rustfmt::skip] + let cases: &[(&str, Option)] = &[ + ("1_0", Some(10.0_f64)), + ("1_000.5", Some(1000.5_f64)), + ("0.5_0", Some(0.5_f64)), + ("1_0.5_0", Some(10.5_f64)), + ("0.5e1_0", Some(5000000000.0_f64)), + ("1_2_3", Some(123.0_f64)), + ("1_000_000", Some(1000000.0_f64)), + ("0_1", Some(1.0_f64)), + ("1_0.0_1", Some(10.01_f64)), + ("-1_0", Some(-10.0_f64)), + ("+1_0", Some(10.0_f64)), + ("1_0e1_0", Some(100000000000.0_f64)), + ("1._5", Some(1.5_f64)), + ("1_.5", Some(1.5_f64)), + ("1e_5", Some(100000.0_f64)), + ("1_e5", Some(100000.0_f64)), + ("-_1", Some(-1.0_f64)), + ("+_1", Some(1.0_f64)), + ("._5", Some(0.5_f64)), + ("-_.5", Some(-0.5_f64)), + ("1_._5", Some(1.5_f64)), + ("+_.5", Some(0.5_f64)), + ("1_.", Some(1.0_f64)), + ("_1", None), + ("1_", None), + ("1__0", None), + ("_", None), + ("__", None), + ("._", None), + ("-_", None), + ("_.5", None), + ("1e5_", None), + ("_1.5", None), + ("1.5_", None), + ("0_x10", None), + ("1_0e_1_0", Some(100000000000.0_f64)), + (" 1_0 ", None), + ("_ 1", None), + ("1 _0", None), + (" __1 ", None), + ("\t1_0\n", None), + ("0.5", Some(0.5_f64)), + (" 1.5 ", Some(1.5_f64)), + ("1e-1", Some(0.1_f64)), + ("+1.5", Some(1.5_f64)), + (".5", Some(0.5_f64)), + ("1.", Some(1.0_f64)), + ("-0", Some(-0.0_f64)), + ("1E5", Some(100000.0_f64)), + ("inf", Some(f64::INFINITY)), + ("-inf", Some(f64::NEG_INFINITY)), + ("Infinity", Some(f64::INFINITY)), + ("nan", Some(f64::NAN)), + ("NaN", Some(f64::NAN)), + ("0x10", None), + ("0b101", None), + ("0o17", None), + ("1,5", None), + ("1.5f", None), + ("", None), + (" ", None), + ("abc", None), + ("1e400", Some(f64::INFINITY)), + ]; + for &(input, want) in cases { + let got = parse_as_engine_number(input); + match (got, want) { + (Some(g), Some(w)) if g.is_nan() && w.is_nan() => {} + _ => assert_eq!(got, want, "parse_as_engine_number({input:?})"), + } + } + } + + /// A string long enough to be worth an allocation is not normalized — + /// the underscore path runs in a fixed stack buffer. Under `reject` the + /// over-long value is refused, which is the safe direction: the contract + /// never silently forwards what it declined to read. + #[test] + fn overlong_numeric_string_is_not_normalized_and_is_refused() { + let long = format!("1{}", "_0".repeat(MAX_SAMPLING_NUMERIC_LEN)); + assert!(long.len() > MAX_SAMPLING_NUMERIC_LEN); + assert_eq!(parse_as_engine_number(&long), None); + + let probe = probe_of(&format!(r#"{{"model":"x","temperature":"{long}"}}"#)); + let overrides = overrides_of(ConflictPolicy::Reject, r#"{"temperature": 1}"#); + assert!(apply_sampling_overrides(&overrides, &probe, &metrics()).is_err()); + } + + /// `allow` promises nothing, so it has nothing to fail open on: a value + /// the probe cannot read keeps flowing, and is still never injected over. + #[test] + fn allow_never_rejects_an_unreadable_value() { + let probe = probe_of(r#"{"model":"x","temperature":"abc","top_p":[1]}"#); + let overrides = overrides_of(ConflictPolicy::Allow, r#"{"temperature": 1, "top_p": 0.9}"#); + let inject = apply_sampling_overrides(&overrides, &probe, &metrics()).unwrap(); + assert!(inject.is_empty(), "a client value is never overwritten"); + } + + /// A refused unreadable value is a contract violation like any other: same + /// error code, same per-parameter counter, and it names the parameter and + /// the expectation without echoing the client's value back. + #[test] + fn unreadable_value_rejection_is_counted_and_named() { + for (config, body, expected_detail) in [ + ( + r#"{"temperature": 1}"#, + r#"{"model":"x","temperature":"abc"}"#, + "expected 1 (or omit the field), got a non-numeric value", + ), + ( + r#"{"temperature": {"min": 0.5, "max": 1.5}}"#, + r#"{"model":"x","temperature":{"a":1}}"#, + "must be a number between 0.5 and 1.5, got a non-numeric value", + ), + ] { + let overrides = overrides_of(ConflictPolicy::Reject, config); + let metrics = metrics(); + let probe = probe_of(body); + + let err = apply_sampling_overrides(&overrides, &probe, &metrics).unwrap_err(); + match &err { + ApiError::SamplingContract { param, detail } => { + assert_eq!(*param, "temperature"); + assert_eq!(detail, expected_detail); + } + other => panic!("expected SamplingContract, got {other:?}"), + } + assert!( + metrics.render().contains( + r#"sgl_router_sampling_contract_rejections_total{param="temperature"} 1"# + ), + "a refusal must be visible to an operator rolling the flag out:\n{}", + metrics.render() + ); + } + } + + /// A number literal outside `f64`'s range must not fail the probe. + /// + /// serde_json converts a literal while resolving `ProbedValue`, and errors + /// rather than saturating — so probing the sampling keys turned + /// `{"temperature": 1e400}` into a 400 whose message ("body must be a JSON + /// object") was not true of it, on a router with no contract configured. + /// An unprobed key holding the same literal was unaffected, which is the + /// asymmetry that gives the bug away. + #[test] + fn out_of_range_number_literal_does_not_fail_the_probe() { + for body in [ + r#"{"model":"x","temperature":1e400}"#, + r#"{"model":"x","temperature":-1e309}"#, + r#"{"model":"x","top_k":1E1000,"stream":true}"#, + ] { + let probe = parse_probe(&Bytes::copy_from_slice(body.as_bytes())) + .unwrap_or_else(|e| panic!("{body} must still parse: {e:?}")); + assert_eq!(probe.model.as_deref(), Some("x"), "{body}"); + } + // Routing fields still come through on the fallback pass. + let probe = probe_of(r#"{"model":"x","temperature":1e400,"stream":true}"#); + assert_eq!(probe.stream, Some(true)); + + // The key was PRESENT, so it is unreadable rather than absent: a + // contract must not inject over a value the client sent... + let probe = probe_of(r#"{"model":"x","temperature":1e400}"#); + assert_eq!( + probe.sampling_field(SamplingField::Temperature), + ProbedValue::Unusable + ); + let allow = overrides_of(ConflictPolicy::Allow, r#"{"temperature": 1}"#); + assert!(apply_sampling_overrides(&allow, &probe, &metrics()) + .unwrap() + .is_empty()); + // ...and `reject` refuses it, as it does any value it cannot read. + let reject = overrides_of(ConflictPolicy::Reject, r#"{"temperature": 1}"#); + assert!(apply_sampling_overrides(&reject, &probe, &metrics()).is_err()); + + // A body that is genuinely malformed still 400s -- the fallback must + // not launder one. + for bad in [ + r#"{"model":"x","stream":true,"stream":false}"#, + "[]", + "null", + "{oops", + ] { + assert!( + parse_probe(&Bytes::copy_from_slice(bad.as_bytes())).is_err(), + "{bad} must still be rejected" + ); + } + } + + /// The counter answers "how much traffic is the contract turning away, and + /// on which parameter". Stopping at the first violation would report zero + /// for every parameter sorting after it, so a fleet violating two would + /// look like it violates one — and the operator would roll the second out + /// believing it had no blast radius. + #[test] + fn every_violated_parameter_is_counted_not_just_the_first() { + let overrides = overrides_of( + ConflictPolicy::Reject, + r#"{"temperature": 1, "top_p": 0.95, "n": 1}"#, + ); + let metrics = metrics(); + let probe = probe_of(r#"{"model":"x","temperature":0.7,"top_p":0.8,"n":2}"#); + + let err = apply_sampling_overrides(&overrides, &probe, &metrics).unwrap_err(); + // The client still hears about one parameter. + let ApiError::SamplingContract { param, .. } = &err else { + panic!("expected SamplingContract, got {err:?}"); + }; + assert_eq!(*param, "temperature"); + + let rendered = metrics.render(); + for name in ["temperature", "top_p", "n"] { + assert!( + rendered.contains(&format!( + r#"sgl_router_sampling_contract_rejections_total{{param="{name}"}} 1"# + )), + "{name} must be counted:\n{rendered}" + ); + } + } + + /// The cap applies to the plain parse too, not only the underscore path: + /// the string is client-controlled and parsing one is linear in its + /// length, so an unbounded digit run would be free CPU amplification on + /// every request that carries one. + #[test] + fn overlong_plain_numeric_string_is_not_parsed() { + let long = "1".repeat(MAX_SAMPLING_NUMERIC_LEN + 1); + assert_eq!(parse_as_engine_number(&long), None); + assert_eq!( + parse_as_engine_number(&"1".repeat(MAX_SAMPLING_NUMERIC_LEN)), + "1".repeat(MAX_SAMPLING_NUMERIC_LEN).parse::().ok(), + "a value at the cap is still read" + ); + } }