fix(router): buffer bodies when routing-key override is enabled

model_gateway/src/config/types.rs

@@ -517,8 +517,8 @@ pub enum ManualAssignmentMode {
 /// `_t<n>` and per-retry `_r<n>` 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.

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,

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

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,

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:

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());

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<dyn Worker> = 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) =