[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::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 text.contains(TOKEN_ID_SEPARATOR) {
+        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,31 @@ mod tests {
     use super::*;
     use crate::core::{BasicWorker, WorkerType};
 
+    #[test]
+    fn test_token_aware_match_rate_ignores_partial_token_ids() {
+        // "123457\u{1f}8\u{1f}" vs cached "123456\u{1f}7\u{1f}" share 5 chars but
+        // zero complete tokens -> rate 0.
+        let incoming = "123457\u{1f}8\u{1f}";
+        let input = incoming.chars().count();
+        assert_eq!(token_aware_match_rate(incoming, 5, 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 = "10\u{1f}20\u{1f}30\u{1f}";
+        let input = incoming.chars().count();
+        assert_eq!(token_aware_match_rate(incoming, 7, 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 have no separator -> character ratio, unchanged.
+        assert_eq!(token_aware_match_rate("hello world", 6, 11), 6.0 / 11.0);
+    }
+
     #[test]
     fn test_cache_aware_with_balanced_load() {
         // Create policy without eviction thread for testing

src/protocols/spec.rs

@@ -2336,6 +2336,17 @@ pub enum PromptInput {
     String(String),
 }
 
+/// 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. Never present in normal text.
+pub(crate) const TOKEN_ID_SEPARATOR: char = '\u{1f}';
+
+fn encode_token_ids(ids: &[i32]) -> String {
+    ids.iter()
+        .map(|id| format!("{id}{TOKEN_ID_SEPARATOR}"))
+        .collect()
+}
+
 impl PromptInput {
     /// Get the number of items in the PromptInput
     pub fn len(&self) -> usize {
@@ -2363,16 +2374,12 @@ 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) => batches
+                .iter()
+                .map(|b| encode_token_ids(b))
+                .collect::<Vec<_>>()
+                .join(";"),
         }
     }
 
@@ -2425,6 +2432,20 @@ mod tests {
         assert_eq!(req.extract_text_for_routing(), "151644 8948 198 2610");
     }
 
+    #[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(),
+            "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(),
+            "1\u{1f}2\u{1f};3\u{1f}4\u{1f}"
+        );
+    }
+
     #[test]
     fn test_inference_generate_lossless_passthrough() {
         let body = serde_json::json!({

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(),
+        "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"
+        "128000\u{1f}9906\u{1f}11\u{1f}1917\u{1f}0\u{1f}\
+         ;128001\u{1f}9906\u{1f}11\u{1f}1917\u{1f}1\u{1f}"
     );
     assert_eq!(req.prompt.estimated_token_count(), 10);
 }