feat(router): truncate routing tokens at configured media boundaries

Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>

Author: slin1237

bindings/python/src/lib.rs

@@ -516,6 +516,7 @@ struct Router {
     least_load_max_waiting_requests: u32,
     stream_request_bodies_over: u64,
     stream_body_stall_timeout_secs: u64,
+    routing_token_boundaries: Vec<u32>,
 }
 
 impl Router {
@@ -867,6 +868,7 @@ impl Router {
             .upstream_pool_idle_timeout_secs(self.upstream_pool_idle_timeout_secs)
             .stream_request_bodies_over(self.stream_request_bodies_over)
             .stream_body_stall_timeout_secs(self.stream_body_stall_timeout_secs)
+            .routing_token_boundaries(self.routing_token_boundaries.clone())
             .multimodal_tensor_transport(multimodal_tensor_transport)
             .multimodal_shm_min_bytes(self.multimodal_shm_min_bytes)
             .routing_key_override(config::RoutingKeyOverrideConfig {
@@ -1030,6 +1032,7 @@ impl Router {
         least_load_max_waiting_requests = 0,
         stream_request_bodies_over = 0,
         stream_body_stall_timeout_secs = 300,
+        routing_token_boundaries = vec![],
     ))]
     #[expect(clippy::too_many_arguments)]
     #[expect(
@@ -1170,6 +1173,7 @@ impl Router {
         least_load_max_waiting_requests: u32,
         stream_request_bodies_over: u64,
         stream_body_stall_timeout_secs: u64,
+        routing_token_boundaries: Vec<u32>,
     ) -> PyResult<Self> {
         let mut all_urls = worker_urls.clone();
 
@@ -1324,6 +1328,7 @@ impl Router {
             least_load_max_waiting_requests,
             stream_request_bodies_over,
             stream_body_stall_timeout_secs,
+            routing_token_boundaries,
         })
     }
 

bindings/python/src/smg/router_args.py

@@ -216,6 +216,7 @@ class RouterArgs:
     least_load_max_waiting_requests: int = 0
     stream_request_bodies_over: int = 0
     stream_body_stall_timeout_secs: int = 300
+    routing_token_boundaries: list[int] = dataclasses.field(default_factory=list)
 
     @staticmethod
     def add_cli_args(
@@ -613,6 +614,21 @@ def add_cli_args(
                 " via --stream-request-bodies-over. 0 disables"
             ),
         )
+        routing_group.add_argument(
+            f"--{prefix}routing-token-boundaries",
+            type=int,
+            nargs="*",
+            action="extend",
+            # None (not []) so an unset prefixed variant falls back to the
+            # unprefixed backend value in from_cli_args.
+            default=None,
+            help=(
+                "Token ids that end the routing-relevant prefix. Routing"
+                " tokens are truncated at the first occurrence of any listed"
+                " id before worker selection (e.g. multimodal placeholder"
+                " ids). Empty disables truncation"
+            ),
+        )
         routing_group.add_argument(
             f"--{prefix}dp-aware",
             action="store_true",

bindings/python/tests/test_arg_parser.py

@@ -833,6 +833,26 @@ def test_prefixed_repeated_selector_accumulates(self):
 
         assert router_args.selector == {"component": "engine", "env": "prod"}
 
+    def test_routing_token_boundaries_fall_back_to_backend_value(self):
+        """An unset prefixed list must not shadow the backend's value."""
+        parser = argparse.ArgumentParser()
+        RouterArgs.add_cli_args(parser, use_router_prefix=True)
+        namespace = parser.parse_args([])
+        namespace.routing_token_boundaries = [200091, 200038]
+
+        router_args = RouterArgs.from_cli_args(namespace, use_router_prefix=True)
+
+        assert router_args.routing_token_boundaries == [200091, 200038]
+
+    def test_routing_token_boundaries_prefixed_value_wins(self):
+        parser = argparse.ArgumentParser()
+        RouterArgs.add_cli_args(parser, use_router_prefix=True)
+        namespace = parser.parse_args(["--router-routing-token-boundaries", "200091", "200038"])
+
+        router_args = RouterArgs.from_cli_args(namespace, use_router_prefix=True)
+
+        assert router_args.routing_token_boundaries == [200091, 200038]
+
     def test_parse_repeated_model_alias_args(self):
         router_args = parse_router_args(
             [
@@ -1225,6 +1245,7 @@ class TestRouterArgsFieldOrder:
         "least_load_max_waiting_requests",
         "stream_request_bodies_over",
         "stream_body_stall_timeout_secs",
+        "routing_token_boundaries",
     ]
 
     def test_complete_field_sequence_is_frozen(self):
@@ -1249,6 +1270,7 @@ def test_new_fields_appended_after_positional_reserve(self):
             "least_load_max_waiting_requests",
             "stream_request_bodies_over",
             "stream_body_stall_timeout_secs",
+            "routing_token_boundaries",
         ):
             assert names.index(appended) > marker, (
                 f"{appended} must be appended after worker_startup_delay to "

model_gateway/src/config/builder.rs

@@ -215,6 +215,11 @@ impl RouterConfigBuilder {
         self
     }
 
+    pub fn routing_token_boundaries(mut self, ids: Vec<u32>) -> Self {
+        self.config.routing_token_boundaries = ids;
+        self
+    }
+
     pub fn upstream_pool_idle_timeout_secs(mut self, secs: u64) -> Self {
         self.config.upstream_pool_idle_timeout_secs = secs;
         self

model_gateway/src/config/types.rs

@@ -37,6 +37,13 @@ pub struct RouterConfig {
     /// Per-request sticky-routing override (honors `X-SMG-Routing-Key`).
     #[serde(default)]
     pub routing_key_override: RoutingKeyOverrideConfig,
+    /// Token ids that end the routing-relevant prefix. Routing tokens are
+    /// truncated at the first occurrence of any listed id before worker
+    /// selection: content past a media placeholder is never shareable across
+    /// conversations, and match ratios over the full sequence shrink as
+    /// conversations grow. Empty disables truncation.
+    #[serde(default)]
+    pub routing_token_boundaries: Vec<u32>,
     pub host: String,
     pub port: u16,
     /// Dedicated port for the isolated Kubernetes liveness/readiness/health
@@ -912,6 +919,7 @@ impl Default for RouterConfig {
             },
             policy: PolicyConfig::Random,
             routing_key_override: RoutingKeyOverrideConfig::default(),
+            routing_token_boundaries: Vec::new(),
             host: "0.0.0.0".to_string(),
             port: 3001,
             health_check_port: None,
@@ -1206,6 +1214,26 @@ mod tests {
         assert_eq!(with.stream_body_stall_timeout_secs, 0);
     }
 
+    #[test]
+    fn test_routing_token_boundaries_serde_default_and_roundtrip() {
+        // Config files predating the field deserialize to no boundaries.
+        let mut json: serde_json::Value = serde_json::to_value(RouterConfig::default()).unwrap();
+        json.as_object_mut()
+            .unwrap()
+            .remove("routing_token_boundaries")
+            .unwrap();
+        let without: RouterConfig = serde_json::from_value(json).unwrap();
+        assert!(without.routing_token_boundaries.is_empty());
+
+        let config = RouterConfig::builder()
+            .regular_mode(vec![])
+            .routing_token_boundaries(vec![200091, 200038])
+            .build_unchecked();
+        let json = serde_json::to_string(&config).unwrap();
+        let with: RouterConfig = serde_json::from_str(&json).unwrap();
+        assert_eq!(with.routing_token_boundaries, vec![200091, 200038]);
+    }
+
     #[test]
     fn test_routing_mode_is_pd_mode() {
         let regular = RoutingMode::Regular {

model_gateway/src/main.rs

@@ -312,6 +312,12 @@ struct CliArgs {
     #[arg(long, default_value_t = false, help_heading = "Routing Policy")]
     routing_key_override: bool,
 
+    /// Token ids that end the routing-relevant prefix (e.g. multimodal
+    /// placeholder ids); routing tokens are truncated at the first occurrence
+    /// before worker selection. Empty disables
+    #[arg(long, num_args = 0.., help_heading = "Routing Policy")]
+    routing_token_boundaries: Vec<u32>,
+
     /// Enable IGW (Inference Gateway) mode for multi-model support
     #[arg(long, default_value_t = false, help_heading = "Routing Policy")]
     enable_igw: bool,
@@ -1663,6 +1669,7 @@ impl CliArgs {
                 max_idle_secs: self.max_idle_secs,
                 assignment_mode: Self::parse_assignment_mode(&self.assignment_mode),
             })
+            .routing_token_boundaries(self.routing_token_boundaries.clone())
             .retries(!self.disable_retries)
             .upstream_http2(self.upstream_http2)
             .circuit_breaker(!self.disable_circuit_breaker)
@@ -1960,6 +1967,18 @@ mod tests {
         assert_eq!(defaults.stream_body_stall_timeout_secs, 300);
     }
 
+    /// Routing-token boundaries are a router-only setting and must flow into
+    /// `RouterConfig`.
+    #[test]
+    fn routing_token_boundaries_flag_flows_into_router_config() {
+        let cli = cli_args_from(&["--routing-token-boundaries", "200091", "200038"]);
+        let router_config = cli.to_router_config(vec![], vec![]).unwrap();
+        assert_eq!(router_config.routing_token_boundaries, vec![200091, 200038]);
+
+        let defaults = cli_args_from(&[]).to_router_config(vec![], vec![]).unwrap();
+        assert!(defaults.routing_token_boundaries.is_empty());
+    }
+
     /// `--health-check-port` must flow into BOTH conversion paths
     /// (`to_router_config` and `to_server_config`), mirroring the main
     /// listener `--port` field exactly. This is the two-path config-plumbing

model_gateway/src/observability/metrics.rs

@@ -1289,6 +1289,15 @@ impl Metrics {
         .set(count as f64);
     }
 
+    /// Record a routing-token truncation at a configured boundary id
+    pub fn record_routing_tokens_truncated(router_type: &'static str) {
+        counter!(
+            "smg_routing_tokens_truncated_total",
+            "router_type" => router_type
+        )
+        .increment(1);
+    }
+
     /// Record health check result
     pub fn record_worker_health_check(worker_type: &'static str, result: &'static str) {
         counter!(

model_gateway/src/routers/common/mod.rs

@@ -28,5 +28,6 @@ pub mod openai_bridge;
 pub mod persistence_utils;
 pub mod realtime;
 pub mod retry;
+pub(crate) mod routing_tokens;
 pub mod sse;
 pub mod worker_selection;

model_gateway/src/routers/common/routing_tokens.rs

@@ -0,0 +1,68 @@
+//! Routing-token boundary truncation shared by the HTTP, PD and gRPC
+//! selection paths.
+
+use crate::observability::metrics::Metrics;
+
+/// Slice `tokens` at the first configured boundary id. Content past a
+/// boundary is never shareable across conversations, and match ratios over
+/// the full sequence shrink as conversations grow.
+pub(crate) fn truncate_slice<'a>(
+    tokens: &'a [u32],
+    boundaries: &[u32],
+    router_type: &'static str,
+) -> &'a [u32] {
+    if boundaries.is_empty() {
+        return tokens;
+    }
+    match tokens.iter().position(|id| boundaries.contains(id)) {
+        Some(cut) => {
+            Metrics::record_routing_tokens_truncated(router_type);
+            &tokens[..cut]
+        }
+        None => tokens,
+    }
+}
+
+/// Owned variant: truncates in place, no reallocation.
+pub(crate) fn truncate_owned(
+    mut tokens: Vec<u32>,
+    boundaries: &[u32],
+    router_type: &'static str,
+) -> Vec<u32> {
+    let keep = truncate_slice(&tokens, boundaries, router_type).len();
+    tokens.truncate(keep);
+    tokens
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn cuts_at_first_of_several_boundaries() {
+        assert_eq!(
+            truncate_slice(&[1, 2, 900, 3, 901], &[900, 901], "http"),
+            &[1, 2]
+        );
+        assert_eq!(
+            truncate_owned(vec![1, 2, 901, 3, 900], &[900, 901], "http"),
+            vec![1, 2]
+        );
+    }
+
+    #[test]
+    fn boundary_first_yields_empty() {
+        assert_eq!(truncate_slice(&[900, 1], &[900], "http"), &[] as &[u32]);
+        assert_eq!(
+            truncate_owned(vec![900, 1], &[900], "http"),
+            Vec::<u32>::new()
+        );
+    }
+
+    #[test]
+    fn no_boundaries_or_no_match_is_identity() {
+        assert_eq!(truncate_slice(&[1, 900, 2], &[], "http"), &[1, 900, 2]);
+        assert_eq!(truncate_slice(&[1, 2, 3], &[900], "http"), &[1, 2, 3]);
+        assert_eq!(truncate_owned(vec![1, 2], &[900], "http"), vec![1, 2]);
+    }
+}