From 875ef5d5cc3ea1cff22eca739b88e99b10afa374 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Fri, 28 Aug 2026 11:58:36 -0700 Subject: [PATCH] fix(router): buffer bodies when routing-key override is enabled Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- model_gateway/src/config/types.rs | 4 +- model_gateway/src/main.rs | 5 +- model_gateway/src/policies/registry.rs | 7 ++ .../src/routers/common/body_policy.rs | 27 +++++-- model_gateway/src/routers/http/router.rs | 75 +++++++------------ model_gateway/tests/common/mod.rs | 5 +- .../tests/routing/stream_request_body_test.rs | 67 +++++++++++++++++ 7 files changed, 129 insertions(+), 61 deletions(-) diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index 6aff2c95d7..556dcf410c 100755 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -517,8 +517,8 @@ pub enum ManualAssignmentMode { /// `_t` and per-retry `_r` suffixes stripped, so every turn of a /// conversation shares one key) wins over the routing-key headers; the first /// configured header carrying a valid value is the fallback when no rid is -/// present. Raw-streamed requests have no readable body and therefore derive -/// keys from the headers only. +/// present. An enabled override keeps automatic body forwarding buffered so +/// body `rid` precedence is preserved. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct RoutingKeyOverrideConfig { /// When false, policies are used unchanged. diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 6c34069c23..c851b7eb7a 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -428,8 +428,9 @@ struct CliArgs { /// Sticky sessions: route every request of a conversation to the same /// worker, on any policy. The key is derived from the request body's rid /// with per-turn/per-retry suffixes stripped (conv_t2_r1 -> conv), - /// falling back to the routing-key headers when no rid is present; - /// raw-streamed requests carry no readable rid and use the headers only. + /// falling back to the routing-key headers when no rid is present. + /// Enabling this keeps automatic body forwarding buffered so body rid + /// precedence is preserved. /// Reuses the manual eviction/idle/assignment knobs for the sticky map #[arg( long, diff --git a/model_gateway/src/policies/registry.rs b/model_gateway/src/policies/registry.rs index 2a1bb058ed..439b7c3bc8 100644 --- a/model_gateway/src/policies/registry.rs +++ b/model_gateway/src/policies/registry.rs @@ -208,6 +208,13 @@ impl PolicyRegistry { }) } + /// Whether sticky routing may derive its preferred key from the request + /// body's `rid`, requiring automatic body-path selection to keep the body + /// readable. + pub(crate) fn routing_key_override_enabled(&self) -> bool { + self.routing_key_sticky.is_some() + } + /// Resolve the effective sticky key: the rid-derived key wins, the /// configured routing-key headers are the fallback when no rid is /// present. Header keys get the same lineage stripping as rid keys, so a diff --git a/model_gateway/src/routers/common/body_policy.rs b/model_gateway/src/routers/common/body_policy.rs index bada4cb0c6..b14aea3c5c 100644 --- a/model_gateway/src/routers/common/body_policy.rs +++ b/model_gateway/src/routers/common/body_policy.rs @@ -22,6 +22,7 @@ pub(crate) enum BodyPath { pub(crate) const BODY_PATH_STREAMED: &str = "streamed"; pub(crate) const BODY_PATH_BUFFERED: &str = "buffered"; +pub(crate) const REASON_ROUTING_KEY_OVERRIDE: &str = "routing_key_override"; pub(crate) const REASON_POLICY_NEEDS_TEXT: &str = "policy_needs_text"; pub(crate) const REASON_MODEL_AMBIGUOUS: &str = "model_ambiguous"; pub(crate) const REASON_WORKER_MUTATES_BODY: &str = "worker_mutates_body"; @@ -35,6 +36,7 @@ pub(crate) const REASON_PURE_FORWARD: &str = "pure_forward"; /// Per-request inputs, gathered by the caller before the body arrives. pub(crate) struct BodyPathInputs { + pub routing_key_override: bool, pub policy_needs_text: bool, pub model_ambiguous: bool, pub worker_mutates_body: bool, @@ -45,14 +47,18 @@ pub(crate) struct BodyPathInputs { } /// Decide the request body path before the body arrives. Buffer when -/// something must read the body here: a text-routing policy without a valid -/// hint-header waiver, a registry serving more than one model (content-blind -/// selection could land the request on the wrong model's worker), a -/// body-mutating worker, a WASM request hook, or a missing/invalid -/// Content-Length. Otherwise, with router retries enabled, buffer up to +/// something must read the body here: routing-key override (body `rid` wins +/// over any header key), a text-routing policy without a valid hint-header +/// waiver, a registry serving more than one model (content-blind selection +/// could land the request on the wrong model's worker), a body-mutating +/// worker, a WASM request hook, or a missing/invalid Content-Length. +/// Otherwise, with router retries enabled, buffer up to /// `max_buffered_request_bytes` to keep the request replayable — a larger /// body streams and forfeits router retries. Otherwise stream, at any size. pub(crate) fn decide_body_path(inputs: &BodyPathInputs) -> BodyPath { + if inputs.routing_key_override { + return BodyPath::Buffer(REASON_ROUTING_KEY_OVERRIDE); + } if inputs.policy_needs_text { return BodyPath::Buffer(REASON_POLICY_NEEDS_TEXT); } @@ -83,6 +89,7 @@ mod tests { fn eligible(content_length: u64) -> BodyPathInputs { BodyPathInputs { + routing_key_override: false, policy_needs_text: false, model_ambiguous: false, worker_mutates_body: false, @@ -96,6 +103,7 @@ mod tests { #[test] fn hard_buffer_reasons_apply_in_first_match_order() { let all = BodyPathInputs { + routing_key_override: true, policy_needs_text: true, model_ambiguous: true, worker_mutates_body: true, @@ -105,10 +113,17 @@ mod tests { }; assert_eq!( decide_body_path(&all), - BodyPath::Buffer(REASON_POLICY_NEEDS_TEXT) + BodyPath::Buffer(REASON_ROUTING_KEY_OVERRIDE) ); for (inputs, reason) in [ + ( + BodyPathInputs { + routing_key_override: true, + ..eligible(64) + }, + REASON_ROUTING_KEY_OVERRIDE, + ), ( BodyPathInputs { policy_needs_text: true, diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index 26b9ae50b3..625617f215 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -1650,6 +1650,7 @@ impl Router { .worker_registry .get_routing_pool(crate::worker::UNKNOWN_MODEL_ID, RoutingPool::HttpRegular); decide_body_path(&BodyPathInputs { + routing_key_override: self.policy_registry.routing_key_override_enabled(), policy_needs_text: self .policy_registry .any_policy_needs_request_text(Some(headers)), @@ -1688,8 +1689,8 @@ impl Router { let model_id = crate::worker::UNKNOWN_MODEL_ID; // Buffered-path parity: a valid tokens hint is exactly what selection // would have received there (text is never extracted alongside it). - // Streamed requests have no readable body, hence no rid key; the - // sticky override keys them by the header alone. + // Streamed requests have no readable body, hence no rid key; + // routing-key override is excluded by the body-path gate above. let hinted_tokens = header_utils::parse_routing_tokens_hint(Some(req.headers())); let Some(worker) = self.select_worker_for_model( model_id, @@ -2167,7 +2168,7 @@ mod tests { routers::common::{ body_policy::{ REASON_NO_CONTENT_LENGTH, REASON_POLICY_NEEDS_TEXT, REASON_RETRYABLE, - REASON_RETRY_FORFEITED, REASON_WASM_REQUEST_HOOK, + REASON_RETRY_FORFEITED, REASON_ROUTING_KEY_OVERRIDE, REASON_WASM_REQUEST_HOOK, }, request_lease::test_probe::{spawn_release_gated_stub, DropProbeRequest}, }, @@ -3038,61 +3039,35 @@ mod tests { } } - #[tokio::test] - async fn key_hint_streams_under_cache_aware_only_with_override() { - // Without the sticky override nothing consumes the key content-blind. - let router = streaming_router( + #[test] + fn routing_key_override_always_decides_buffer() { + let cache_aware = streaming_router_with_key_override( cache_aware_policy(), 1024 * 1024, vec![plain_worker("http://worker1:8080")], ); - let req = with_header( - streamed_request(&[b"{\"text\":\"hello\"}"]), - "x-smg-routing-key", - "media-1", - ); - assert!(router - .route_streaming_request(req, "/generate", false) - .await - .is_err()); - let (url_a, _cap_a) = spawn_capture_stub("application/json", "{}").await; - let (url_b, _cap_b) = spawn_capture_stub("application/json", "{}").await; - let router = streaming_router_with_key_override( - cache_aware_policy(), + let no_hint = headers_with_content_length(Some("64")); + let mut key_hint = no_hint.clone(); + key_hint.insert("x-smg-routing-key", "request-unique-key".parse().unwrap()); + let mut tokens_hint = no_hint.clone(); + tokens_hint.insert("x-smg-routing-tokens", "1,2,3".parse().unwrap()); + for headers in [&no_hint, &key_hint, &tokens_hint] { + assert_eq!( + cache_aware.request_body_path(headers, false), + BodyPath::Buffer(REASON_ROUTING_KEY_OVERRIDE) + ); + } + + let text_free = streaming_router_with_key_override( + least_load_policy(), 1024 * 1024, - vec![plain_worker(&url_a), plain_worker(&url_b)], + vec![plain_worker("http://worker1:8080")], ); - let first = routed_worker_id( - &router, - with_header( - streamed_request(&[b"{\"text\":\"hello\"}"]), - "x-smg-routing-key", - "media-1", - ), - ) - .await; - let second = routed_worker_id( - &router, - with_header( - streamed_request(&[b"{\"text\":\"other\"}"]), - "x-smg-routing-key", - "media-1", - ), - ) - .await; - assert_eq!(first, second, "keyed requests must stick to one worker"); - - // Over-cap keys are ignored by the same extractor selection uses. - let req = with_header( - streamed_request(&[b"{\"text\":\"hello\"}"]), - "x-smg-routing-key", - &"k".repeat(129), + assert_eq!( + text_free.request_body_path(&no_hint, false), + BodyPath::Buffer(REASON_ROUTING_KEY_OVERRIDE) ); - assert!(router - .route_streaming_request(req, "/generate", false) - .await - .is_err()); } /// With retries disabled the parsed request must be freed at dispatch: diff --git a/model_gateway/tests/common/mod.rs b/model_gateway/tests/common/mod.rs index 2ebfbb56fe..6d444bb233 100644 --- a/model_gateway/tests/common/mod.rs +++ b/model_gateway/tests/common/mod.rs @@ -348,7 +348,10 @@ async fn build_test_app_context( // Initialize registries let worker_registry = Arc::new(WorkerRegistry::new()); - let policy_registry = Arc::new(PolicyRegistry::new(config.policy.clone())); + let policy_registry = Arc::new(PolicyRegistry::with_override( + config.policy.clone(), + config.routing_key_override.clone(), + )); // Initialize storage backends (Memory for tests). let response_storage = Arc::new(MemoryResponseStorage::new()); diff --git a/model_gateway/tests/routing/stream_request_body_test.rs b/model_gateway/tests/routing/stream_request_body_test.rs index 3e9062034f..11c8a6238c 100644 --- a/model_gateway/tests/routing/stream_request_body_test.rs +++ b/model_gateway/tests/routing/stream_request_body_test.rs @@ -273,6 +273,73 @@ async fn retries_disabled_streams_a_small_body() { ctx.shutdown().await; } +#[tokio::test] +async fn routing_key_override_buffers_and_prefers_body_rid() { + let (url_a, captured_a) = spawn_capture_worker().await; + let (url_b, captured_b) = spawn_capture_worker().await; + + let mut config = base_config(64 * 1024 * 1024); + config.disable_retries = true; + config.policy = PolicyConfig::RoundRobin; + config.routing_key_override.enabled = true; + config.health_check.disable_health_check = true; + let ctx = AppTestContext::new_with_config(config, vec![]).await; + for url in [&url_a, &url_b] { + let worker: Arc = Arc::new( + BasicWorkerBuilder::new(url) + .worker_type(WorkerType::Regular) + .models(vec![ModelCard::new("mock-model")]) + .health_config(openai_protocol::worker::HealthCheckConfig { + disable_health_check: true, + ..Default::default() + }) + .build(), + ); + ctx.app_context.worker_registry.register(worker); + } + let app = ctx.create_app(); + + for (rid, header_key) in [("conv_t1", "request-1"), ("conv_t2", "request-2")] { + let payload = serde_json::to_vec(&json!({ + "input_ids": [1, 2, 3], + "stream": false, + "rid": rid, + })) + .unwrap(); + let mut req = + chunked_json_request("/generate", payload.clone(), Some(payload.len() as u64)); + req.headers_mut() + .insert("x-smg-routing-key", header_key.parse().unwrap()); + req.headers_mut() + .insert("x-smg-routing-tokens", "1,2,3".parse().unwrap()); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } + + let captured_a = captured_a.lock().await; + let captured_b = captured_b.lock().await; + assert!( + matches!((captured_a.len(), captured_b.len()), (2, 0) | (0, 2)), + "body rid must override distinct per-request header keys" + ); + for ((headers, body), rid) in captured_a + .iter() + .chain(captured_b.iter()) + .zip(["conv_t1", "conv_t2"]) + { + assert!( + headers.get(CONTENT_LENGTH).is_some(), + "routing-key override must use the buffered typed path" + ); + let body: Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["rid"].as_str(), Some(rid)); + assert_eq!(body["input_ids"], json!([1, 2, 3])); + } + drop(captured_a); + drop(captured_b); + ctx.shutdown().await; +} + #[tokio::test] async fn default_config_buffers_small_and_streams_past_one_mib() { let (app, captured, ctx) =