[Bugfix] Content-aware routing should work with token ids

src/policies/cache_aware.rs

@@ -63,6 +63,7 @@ use super::{get_healthy_worker_indices, CacheAwareConfig, LoadBalancingPolicy, R
 use crate::core::Worker;
 use crate::metrics::RouterMetrics;
 use crate::policies::normalize_model_key;
+use crate::protocols::spec::{is_token_id_key, TOKEN_ID_SEPARATOR};
 use crate::tree::Tree;
 use dashmap::DashMap;
 use rand::Rng;
@@ -72,6 +73,26 @@ use std::thread;
 use std::time::Duration;
 use tracing::{debug, info};
 
+/// Cache-match ratio, computed per token for token-id routing keys and otherwise for chars.
+fn token_aware_match_rate(text: &str, matched_char_count: usize, input_char_count: usize) -> f32 {
+    if input_char_count == 0 {
+        return 0.0;
+    }
+    if is_token_id_key(text) {
+        let total_tokens = text.matches(TOKEN_ID_SEPARATOR).count();
+        if total_tokens == 0 {
+            return 0.0;
+        }
+        let matched_tokens = text
+            .chars()
+            .take(matched_char_count)
+            .filter(|c| *c == TOKEN_ID_SEPARATOR)
+            .count();
+        return matched_tokens as f32 / total_tokens as f32;
+    }
+    matched_char_count as f32 / input_char_count as f32
+}
+
 /// Cache-aware routing policy
 ///
 /// Routes requests based on cache affinity when load is balanced,
@@ -300,11 +321,8 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
         // Now we work with the tree without holding the HashMap lock
         // Use prefix_match_with_counts to avoid redundant chars().count() calls
         let result = tree.prefix_match_with_counts(text);
-        let match_rate = if result.input_char_count == 0 {
-            0.0
-        } else {
-            result.matched_char_count as f32 / result.input_char_count as f32
-        };
+        let match_rate =
+            token_aware_match_rate(text, result.matched_char_count, result.input_char_count);
 
         debug!(
             "Cache match for model '{}': matched_chars={}, input_chars={}, match_rate={:.2}",
@@ -479,6 +497,42 @@ mod tests {
     use super::*;
     use crate::core::{BasicWorker, WorkerType};
 
+    #[test]
+    fn test_token_aware_match_rate_ignores_partial_token_ids() {
+        // "\u{1e}123457\u{1f}8\u{1f}" vs cached "\u{1e}123456\u{1f}7\u{1f}" share 6 chars
+        // but zero complete tokens -> rate 0.
+        let incoming = "\u{1e}123457\u{1f}8\u{1f}";
+        let input = incoming.chars().count();
+        assert_eq!(token_aware_match_rate(incoming, 6, input), 0.0);
+    }
+
+    #[test]
+    fn test_token_aware_match_rate_counts_whole_tokens() {
+        // Matched through two separators plus a partial third token -> 2 of 3.
+        let incoming = "\u{1e}10\u{1f}20\u{1f}30\u{1f}";
+        let input = incoming.chars().count();
+        assert_eq!(token_aware_match_rate(incoming, 8, input), 2.0 / 3.0);
+        // Full match -> 1.0.
+        assert_eq!(token_aware_match_rate(incoming, input, input), 1.0);
+    }
+
+    #[test]
+    fn test_token_aware_match_rate_text_keys_use_chars() {
+        // Plain-text keys are scored by character ratio, unchanged.
+        assert_eq!(token_aware_match_rate("hello world", 6, 11), 6.0 / 11.0);
+    }
+
+    #[test]
+    fn test_token_aware_match_rate_untagged_separator_text_uses_chars() {
+        // A prompt containing a raw separator is still text, not a token-id key.
+        let incoming = "hello\u{1f}world";
+        let input = incoming.chars().count();
+        assert_eq!(
+            token_aware_match_rate(incoming, 6, input),
+            6.0 / input as f32
+        );
+    }
+
     #[test]
     fn test_cache_aware_with_balanced_load() {
         // Create policy without eviction thread for testing

src/protocols/spec.rs

@@ -1,6 +1,7 @@
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::collections::HashMap;
+use std::fmt::Write;
 
 // # Protocol Specifications
 //
@@ -2001,11 +2002,7 @@ impl GenerationRequest for InferenceGenerateRequest {
     }
 
     fn extract_text_for_routing(&self) -> String {
-        self.token_ids
-            .iter()
-            .map(|id| id.to_string())
-            .collect::<Vec<String>>()
-            .join(" ")
+        encode_token_ids(&self.token_ids)
     }
 }
 
@@ -2336,6 +2333,52 @@ pub enum PromptInput {
     String(String),
 }
 
+/// Record separator that opens every encoded token-id routing key, and separates the
+/// sequences of a batch. Routing keys that do not start with it are user text.
+pub(crate) const TOKEN_ID_TAG: char = '\u{1e}';
+
+/// Unit separator that terminates each encoded token id, so cache-aware routing
+/// (a character radix) credits only complete token IDs rather than shared
+/// leading digits of two different tokens.
+pub(crate) const TOKEN_ID_SEPARATOR: char = '\u{1f}';
+
+/// Upper bound on the characters one encoded token id occupies (`-2147483648` plus separator).
+const TOKEN_ID_ENCODED_LEN: usize = 12;
+
+fn write_token_ids(out: &mut String, ids: &[i32]) {
+    for id in ids {
+        let _ = write!(out, "{id}{TOKEN_ID_SEPARATOR}");
+    }
+}
+
+pub(crate) fn encode_token_ids(ids: &[i32]) -> String {
+    let mut out = String::with_capacity(1 + ids.len() * TOKEN_ID_ENCODED_LEN);
+    out.push(TOKEN_ID_TAG);
+    write_token_ids(&mut out, ids);
+    out
+}
+
+pub(crate) fn encode_token_id_batches(batches: &[Vec<i32>]) -> String {
+    let total_ids: usize = batches.iter().map(|b| b.len()).sum();
+    let mut out = String::with_capacity(batches.len() + total_ids * TOKEN_ID_ENCODED_LEN);
+    for batch in batches {
+        out.push(TOKEN_ID_TAG);
+        write_token_ids(&mut out, batch);
+    }
+    out
+}
+
+/// Whether a routing key is an encoded token-id sequence rather than user text.
+///
+/// Only keys that carry the tag and consist entirely of encoded-token characters
+/// qualify, so arbitrary prompt text can never be scored as if it were token ids.
+pub(crate) fn is_token_id_key(key: &str) -> bool {
+    let mut chars = key.chars();
+    chars.next() == Some(TOKEN_ID_TAG)
+        && chars
+            .all(|c| c.is_ascii_digit() || c == '-' || c == TOKEN_ID_SEPARATOR || c == TOKEN_ID_TAG)
+}
+
 impl PromptInput {
     /// Get the number of items in the PromptInput
     pub fn len(&self) -> usize {
@@ -2363,16 +2406,8 @@ impl PromptInput {
         match self {
             PromptInput::String(s) => s.clone(),
             PromptInput::StringArray(arr) => arr.join(" "),
-            PromptInput::IntArray(ids) => {
-                // Convert token IDs to string representation for routing
-                // Format: "token_ids:<count>" to indicate this is a token-based prompt
-                format!("token_ids:{}", ids.len())
-            }
-            PromptInput::IntBatch(batches) => {
-                // For batches, use total token count
-                let total_tokens: usize = batches.iter().map(|b| b.len()).sum();
-                format!("token_ids_batch:{}:{}", batches.len(), total_tokens)
-            }
+            PromptInput::IntArray(ids) => encode_token_ids(ids),
+            PromptInput::IntBatch(batches) => encode_token_id_batches(batches),
         }
     }
 
@@ -2422,7 +2457,35 @@ mod tests {
             "sampling_params": {"max_tokens": 4096, "seed": 42, "logprobs": 1}
         });
         let req: InferenceGenerateRequest = serde_json::from_value(body).unwrap();
-        assert_eq!(req.extract_text_for_routing(), "151644 8948 198 2610");
+        assert_eq!(
+            req.extract_text_for_routing(),
+            "\u{1e}151644\u{1f}8948\u{1f}198\u{1f}2610\u{1f}"
+        );
+    }
+
+    #[test]
+    fn test_prompt_input_token_ids_routing_key() {
+        let single = PromptInput::IntArray(vec![151644, 8948, 198, 2610]);
+        assert_eq!(
+            single.extract_text_for_routing(),
+            "\u{1e}151644\u{1f}8948\u{1f}198\u{1f}2610\u{1f}"
+        );
+        let batch = PromptInput::IntBatch(vec![vec![1, 2], vec![3, 4]]);
+        assert_eq!(
+            batch.extract_text_for_routing(),
+            "\u{1e}1\u{1f}2\u{1f}\u{1e}3\u{1f}4\u{1f}"
+        );
+    }
+
+    #[test]
+    fn test_is_token_id_key_rejects_user_text() {
+        assert!(is_token_id_key(
+            &PromptInput::IntArray(vec![1, 2]).extract_text_for_routing()
+        ));
+        assert!(!is_token_id_key("hello\u{1f}world"));
+        assert!(!is_token_id_key("\u{1e}hello\u{1f}world"));
+        assert!(!is_token_id_key("plain text"));
+        assert!(!is_token_id_key(""));
     }
 
     #[test]

tests/test_prompt_input.rs

@@ -67,7 +67,10 @@ fn test_prompt_input_single_int_array() {
     assert_eq!(req.prompt.len(), 1);
     assert!(!req.prompt.is_empty());
     assert!(req.prompt.is_token_based());
-    assert_eq!(req.prompt.extract_text_for_routing(), "token_ids:5");
+    assert_eq!(
+        req.prompt.extract_text_for_routing(),
+        "\u{1e}128000\u{1f}9906\u{1f}11\u{1f}1917\u{1f}0\u{1f}"
+    );
     assert_eq!(req.prompt.estimated_token_count(), 5);
 }
 
@@ -98,7 +101,8 @@ fn test_prompt_input_int_batch() {
     assert!(req.prompt.is_token_based());
     assert_eq!(
         req.prompt.extract_text_for_routing(),
-        "token_ids_batch:2:10"
+        "\u{1e}128000\u{1f}9906\u{1f}11\u{1f}1917\u{1f}0\u{1f}\
+         \u{1e}128001\u{1f}9906\u{1f}11\u{1f}1917\u{1f}1\u{1f}"
     );
     assert_eq!(req.prompt.estimated_token_count(), 10);
 }