chore: vendor sglang v0.5.10 snapshot

third_party/sglang/sgl-model-gateway/src/policies/cache_aware.rs

@@ -0,0 +1,887 @@
+/*
+    Cache-Aware Load Balancing Router
+
+    This router combines two strategies to optimize both cache utilization and request distribution:
+
+    1. Cache-Aware Routing (Approximate Tree)
+    2. Load Balancing (Shortest Queue with Balance Thresholds)
+
+    The router dynamically switches between these strategies based on load conditions:
+    - Uses load balancing when the system is imbalanced
+    - Uses cache-aware routing when the system is balanced
+
+    A system is considered imbalanced if both conditions are met:
+    1. (max - min) > abs_threshold
+    2. max > rel_threshold * min
+
+    Strategy Details:
+
+    1. Cache-Aware Routing (Approximate Tree)
+    -------------------------------------------
+    This strategy maintains an approximate radix tree for each worker based on request history,
+    eliminating the need for direct cache state queries. The tree stores raw text characters
+    instead of token IDs to avoid tokenization overhead.
+
+    Process:
+    a. For each request, find the worker with the highest prefix match
+    b. If match rate > cache_threshold:
+    Route to the worker with highest match (likely has relevant data cached)
+    c. If match rate ≤ cache_threshold:
+    Route to the worker with smallest tree size (most available cache capacity)
+    d. Background maintenance:
+    Periodically evict least recently used leaf nodes to prevent memory overflow
+
+    2. Load Balancing (Shortest Queue)
+    -------------------------------------------
+    This strategy tracks pending request counts per worker and routes new requests
+    to the least busy worker when the system is detected to be imbalanced.
+
+    Configuration Parameters:
+    ------------------------
+    1. cache_threshold: (float, 0.0 to 1.0)
+    Minimum prefix match ratio to use highest-match routing.
+    Below this threshold, routes to worker with most available cache space.
+
+    2. balance_abs_threshold: (integer)
+    Absolute difference threshold for load imbalance detection.
+    System is potentially imbalanced if (max_load - min_load) > abs_threshold
+
+    3. balance_rel_threshold: (float)
+    Relative ratio threshold for load imbalance detection.
+    System is potentially imbalanced if max_load > min_load * rel_threshold
+    Used in conjunction with abs_threshold to determine final imbalance state.
+
+    4. eviction_interval_secs: (integer)
+    Interval between LRU eviction cycles for the approximate trees.
+
+    5. max_tree_size: (integer)
+    Maximum nodes per tree. When exceeded, LRU leaf nodes are evicted
+    during the next eviction cycle.
+*/
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use dashmap::DashMap;
+use rand::Rng;
+use smg_mesh::{tree_ops::TreeOperation, OptionalMeshSyncManager};
+use tracing::{debug, warn};
+
+use super::{
+    get_healthy_worker_indices, normalize_model_key, tree::Tree, utils::PeriodicTask,
+    CacheAwareConfig, LoadBalancingPolicy, SelectWorkerInfo,
+};
+use crate::core::{Worker, UNKNOWN_MODEL_ID};
+
+/// Cache-aware routing policy
+///
+/// Routes requests based on cache affinity when load is balanced,
+/// switches to shortest-queue routing when load is imbalanced.
+/// Maintains separate trees per model for multi-model support.
+/// Supports mesh synchronization of tree operations across cluster nodes.
+/// When mesh is not enabled, the policy works independently without synchronization.
+#[derive(Debug)]
+pub struct CacheAwarePolicy {
+    config: CacheAwareConfig,
+    trees: Arc<DashMap<String, Arc<Tree>>>,
+    mesh_sync: OptionalMeshSyncManager,
+    _eviction_task: Option<PeriodicTask>,
+}
+
+impl CacheAwarePolicy {
+    pub fn new() -> Self {
+        Self::with_config(CacheAwareConfig::default())
+    }
+
+    pub fn with_config(config: CacheAwareConfig) -> Self {
+        let trees = Arc::new(DashMap::<String, Arc<Tree>>::new());
+
+        // Start background eviction thread if configured
+        let eviction_task = if config.eviction_interval_secs > 0 {
+            let trees_clone = Arc::clone(&trees);
+            let max_tree_size = config.max_tree_size;
+
+            Some(PeriodicTask::spawn(
+                config.eviction_interval_secs,
+                "Eviction",
+                move || {
+                    for tree_ref in trees_clone.iter() {
+                        let model_id = tree_ref.key();
+                        let tree = tree_ref.value();
+                        tree.evict_tenant_by_size(max_tree_size);
+
+                        debug!(
+                            "Cache eviction completed for model {}, max_size: {}",
+                            model_id, max_tree_size
+                        );
+                    }
+                },
+            ))
+        } else {
+            None
+        };
+
+        Self {
+            config,
+            trees,
+            mesh_sync: None,
+            _eviction_task: eviction_task,
+        }
+    }
+
+    /// Set mesh sync manager (can be called after construction)
+    pub fn set_mesh_sync(&mut self, mesh_sync: OptionalMeshSyncManager) {
+        self.mesh_sync = mesh_sync.clone();
+        if mesh_sync.is_some() {
+            self.restore_tree_state_from_mesh();
+        }
+    }
+
+    /// Initialize the tree with worker URLs (used only during initial setup)
+    pub fn init_workers(&self, workers: &[Arc<dyn Worker>]) {
+        // Group workers by model
+        let mut model_workers: std::collections::HashMap<String, Vec<&Arc<dyn Worker>>> =
+            std::collections::HashMap::new();
+        for worker in workers {
+            let tree_key = normalize_model_key(worker.model_id());
+            model_workers
+                .entry(tree_key.to_string())
+                .or_default()
+                .push(worker);
+        }
+
+        // Initialize tree for each model
+        for (tree_key, model_workers) in model_workers {
+            let tree = self
+                .trees
+                .entry(tree_key)
+                .or_insert_with(|| Arc::new(Tree::new()));
+            for worker in model_workers {
+                tree.insert("", worker.url());
+            }
+        }
+    }
+
+    /// Add a single worker to the tree (incremental update)
+    pub fn add_worker(&self, worker: &dyn Worker) {
+        let tree_key = normalize_model_key(worker.model_id());
+        let tree = self
+            .trees
+            .entry(tree_key.to_string())
+            .or_insert_with(|| Arc::new(Tree::new()));
+        tree.insert("", worker.url());
+    }
+
+    /// Add a worker by URL and model (for backward compatibility)
+    pub fn add_worker_by_url(&self, url: &str, model_id: &str) {
+        let tree = self
+            .trees
+            .entry(model_id.to_string())
+            .or_insert_with(|| Arc::new(Tree::new()));
+        tree.insert("", url);
+    }
+
+    /// Remove a worker from the tree
+    pub fn remove_worker(&self, worker: &dyn Worker) {
+        let tree_key = normalize_model_key(worker.model_id());
+        if let Some(tree) = self.trees.get(tree_key) {
+            tree.remove_tenant(worker.url());
+        }
+    }
+
+    /// Remove a worker by URL (removes from all model trees for backward compatibility)
+    pub fn remove_worker_by_url(&self, url: &str) {
+        // Remove from all trees since we don't know which model it belongs to
+        for tree_ref in self.trees.iter() {
+            tree_ref.value().remove_tenant(url);
+        }
+    }
+
+    /// Restore tree state from mesh store
+    /// This is called during initialization to rebuild trees from synchronized state
+    fn restore_tree_state_from_mesh(&self) {
+        if let Some(ref mesh_sync) = self.mesh_sync {
+            // Get all tree states from mesh
+            // We need to iterate through all models that have tree states
+            // For now, we'll restore trees for models that are already in our trees map
+            // In a full implementation, we might want to query mesh for all tree states
+
+            for tree_ref in self.trees.iter() {
+                let model_id = tree_ref.key();
+                if let Some(tree_state) = mesh_sync.get_tree_state(model_id) {
+                    debug!(
+                        "Restoring tree state for model {} with {} operations",
+                        model_id,
+                        tree_state.operations.len()
+                    );
+
+                    let tree = tree_ref.value();
+                    // Apply all operations to rebuild the tree
+                    for operation in &tree_state.operations {
+                        match operation {
+                            TreeOperation::Insert(insert_op) => {
+                                tree.insert(&insert_op.text, &insert_op.tenant);
+                            }
+                            TreeOperation::Remove(remove_op) => {
+                                tree.remove_tenant(&remove_op.tenant);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    /// Normalize model_id for mesh synchronization
+    /// Converts empty model_id to UNKNOWN_MODEL_ID for consistency
+    fn normalize_mesh_model_id(model_id: &str) -> &str {
+        if model_id.is_empty() {
+            UNKNOWN_MODEL_ID
+        } else {
+            model_id
+        }
+    }
+
+    /// Apply remote tree operation from mesh
+    /// This is called when receiving tree state updates from other nodes
+    pub fn apply_remote_tree_operation(&self, model_id: &str, operation: &TreeOperation) {
+        let tree_key = Self::normalize_mesh_model_id(model_id);
+
+        let tree = self
+            .trees
+            .entry(tree_key.to_string())
+            .or_insert_with(|| Arc::new(Tree::new()));
+
+        match operation {
+            TreeOperation::Insert(insert_op) => {
+                tree.insert(&insert_op.text, &insert_op.tenant);
+                debug!(
+                    "Applied remote tree insert: model={}, text={}, tenant={}",
+                    model_id, insert_op.text, insert_op.tenant
+                );
+            }
+            TreeOperation::Remove(remove_op) => {
+                tree.remove_tenant(&remove_op.tenant);
+                debug!(
+                    "Applied remote tree remove: model={}, tenant={}",
+                    model_id, remove_op.tenant
+                );
+            }
+        }
+    }
+
+    /// Run cache eviction to prevent unbounded growth
+    pub fn evict_cache(&self, max_size: usize) {
+        for tree_ref in self.trees.iter() {
+            let model_id = tree_ref.key();
+            let tree = tree_ref.value();
+            tree.evict_tenant_by_size(max_size);
+            debug!(
+                "Cache eviction for model {}, max_size: {}",
+                model_id, max_size
+            );
+        }
+    }
+
+    fn select_worker_min_load(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        request_text: &Option<&str>,
+        healthy_indices: &[usize],
+        model_id: &str,
+        max_load: usize,
+        min_load: usize,
+    ) -> Option<usize> {
+        // Log load balancing trigger (only compute worker loads if debug enabled)
+        if tracing::enabled!(tracing::Level::DEBUG) {
+            let worker_loads: Vec<(&str, usize)> =
+                workers.iter().map(|w| (w.url(), w.load())).collect();
+            debug!(
+                "Load balancing triggered | max: {} | min: {} | workers: {:?}",
+                max_load, min_load, worker_loads
+            );
+        }
+
+        // Use shortest queue when imbalanced
+        let min_load_idx = healthy_indices
+            .iter()
+            .min_by_key(|&&idx| workers[idx].load())
+            .copied()?;
+
+        // Even in imbalanced mode, update the tree to maintain cache state
+        if let Some(text) = request_text {
+            // Get the tree reference without locking the entire HashMap
+            // DashMap only locks the specific shard containing this key
+            let tree = self.trees.get(model_id).map(|entry| entry.value().clone());
+
+            if let Some(tree) = tree {
+                let worker_url = workers[min_load_idx].url();
+                // Now we can work with the tree without holding the HashMap lock
+                tree.insert(text, worker_url);
+
+                // Sync insert operation to mesh if enabled (no-op if mesh is not enabled)
+                if let Some(ref mesh_sync) = self.mesh_sync {
+                    use smg_mesh::tree_ops::TreeInsertOp;
+                    let op = TreeOperation::Insert(TreeInsertOp {
+                        text: text.to_string(),
+                        tenant: worker_url.to_string(),
+                    });
+                    let mesh_model_id = Self::normalize_mesh_model_id(model_id);
+                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) {
+                        warn!("Failed to sync tree insert operation to mesh: {}", e);
+                    }
+                }
+            } else {
+                debug!(
+                    "Warning: No tree found for model '{}', skipping cache update",
+                    model_id
+                );
+            }
+        }
+
+        // Increment processed counter
+        workers[min_load_idx].increment_processed();
+
+        Some(min_load_idx)
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for CacheAwarePolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let request_text = info.request_text;
+        let healthy_indices = get_healthy_worker_indices(workers);
+
+        if healthy_indices.is_empty() {
+            return None;
+        }
+
+        // Determine the model for this set of workers (router pre-filters by model)
+        // All workers should be from the same model
+        let model_id = normalize_model_key(workers[healthy_indices[0]].model_id());
+
+        // Get current load statistics - compute min/max in single pass without allocation
+        let (min_load, max_load) = workers.iter().fold((usize::MAX, 0usize), |(min, max), w| {
+            let load = w.load();
+            (min.min(load), max.max(load))
+        });
+        let min_load = if min_load == usize::MAX { 0 } else { min_load };
+
+        // Check if load is imbalanced
+        let is_imbalanced = max_load.saturating_sub(min_load) > self.config.balance_abs_threshold
+            && (max_load as f32) > (min_load as f32 * self.config.balance_rel_threshold);
+
+        if is_imbalanced {
+            return self.select_worker_min_load(
+                workers,
+                &request_text,
+                &healthy_indices,
+                model_id,
+                max_load,
+                min_load,
+            );
+        }
+
+        // Use cache-aware routing when balanced
+        let text = request_text.unwrap_or("");
+
+        // Get the tree reference without locking the entire HashMap
+        // DashMap only locks the specific shard containing this key
+        let tree = self.trees.get(model_id).map(|entry| entry.value().clone());
+
+        if let Some(tree) = tree {
+            // 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
+            };
+
+            // Select worker without String allocation
+            let selected_idx = if match_rate > self.config.cache_threshold {
+                // Cache hit path: find worker by URL (compare &str directly, no allocation)
+                let tenant_url: &str = &result.tenant;
+                workers
+                    .iter()
+                    .position(|w| w.url() == tenant_url)
+                    .filter(|&idx| workers[idx].is_healthy())
+            } else {
+                // Low cache match: use worker with minimum load
+                healthy_indices
+                    .iter()
+                    .min_by_key(|&&idx| workers[idx].load())
+                    .copied()
+            };
+
+            if let Some(idx) = selected_idx {
+                // Update the tree with this request (use worker URL directly, no allocation)
+                tree.insert(text, workers[idx].url());
+
+                // Sync insert operation to mesh if enabled (no-op if mesh is not enabled)
+                if let Some(ref mesh_sync) = self.mesh_sync {
+                    use smg_mesh::tree_ops::TreeInsertOp;
+                    let op = TreeOperation::Insert(TreeInsertOp {
+                        text: text.to_string(),
+                        tenant: workers[idx].url().to_string(),
+                    });
+                    let mesh_model_id = Self::normalize_mesh_model_id(model_id);
+                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) {
+                        warn!("Failed to sync tree insert operation to mesh: {}", e);
+                    }
+                }
+
+                // Increment processed counter
+                workers[idx].increment_processed();
+
+                return Some(idx);
+            }
+
+            // Selected worker no longer exists or unhealthy, remove stale tenant from tree
+            if match_rate > self.config.cache_threshold {
+                let tenant_url: &str = &result.tenant;
+                tree.remove_tenant(tenant_url);
+                debug!("Removed stale worker {} from cache tree", tenant_url);
+
+                // Sync removal to mesh if enabled (no-op if mesh is not enabled)
+                if let Some(ref mesh_sync) = self.mesh_sync {
+                    use smg_mesh::tree_ops::TreeRemoveOp;
+                    let op = TreeOperation::Remove(TreeRemoveOp {
+                        tenant: tenant_url.to_string(),
+                    });
+                    let mesh_model_id = Self::normalize_mesh_model_id(model_id);
+                    if let Err(e) = mesh_sync.sync_tree_operation(mesh_model_id.to_string(), op) {
+                        warn!("Failed to sync tree remove operation to mesh: {}", e);
+                    }
+                }
+            }
+
+            // Fallback to first healthy worker
+            healthy_indices.first().copied()
+        } else {
+            // No tree for this model, log warning and use random selection
+            debug!(
+                "Warning: No tree found for model '{}', using random worker selection",
+                model_id
+            );
+            // Return a random healthy worker
+            let mut rng = rand::rng();
+            let random_idx = rng.random_range(0..healthy_indices.len());
+            Some(healthy_indices[random_idx])
+        }
+    }
+
+    fn on_request_complete(&self, worker_url: &str, success: bool) {
+        // Could track success rates per worker for more intelligent routing
+        if !success {
+            // Optionally reduce affinity for failed requests
+            tracing::debug!(
+                "Request to {} completed with success={}",
+                worker_url,
+                success
+            );
+        }
+    }
+
+    fn name(&self) -> &'static str {
+        "cache_aware"
+    }
+
+    fn needs_request_text(&self) -> bool {
+        true // Cache-aware policy needs request text for cache affinity
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+impl Default for CacheAwarePolicy {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    #[tokio::test]
+    async fn test_cache_aware_with_balanced_load() {
+        // Create policy without eviction thread for testing
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0, // Disable eviction thread
+            ..Default::default()
+        };
+        let policy = CacheAwarePolicy::with_config(config);
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .api_key("test_api_key")
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .api_key("test_api_key")
+                    .build(),
+            ),
+        ];
+
+        // Initialize the policy with workers
+        policy.init_workers(&workers);
+
+        // First request should be distributed
+        let idx1 = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("hello world"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+
+        // Same request should go to same worker (cache hit)
+        let idx2 = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("hello world"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+        assert_eq!(idx1, idx2);
+
+        // Similar request should also go to same worker
+        let idx3 = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("hello"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+        assert_eq!(idx1, idx3);
+    }
+
+    #[tokio::test]
+    async fn test_cache_aware_with_imbalanced_load() {
+        let policy = CacheAwarePolicy::with_config(CacheAwareConfig {
+            cache_threshold: 0.5,
+            balance_abs_threshold: 5,
+            balance_rel_threshold: 2.0,
+            eviction_interval_secs: 0, // Disable eviction thread
+            max_tree_size: 10000,
+        });
+
+        let worker1 = BasicWorkerBuilder::new("http://w1:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+        let worker2 = BasicWorkerBuilder::new("http://w2:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+
+        // Create significant load imbalance
+        for _ in 0..20 {
+            worker1.increment_load();
+        }
+        // worker2 has load 0
+
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(worker1), Arc::new(worker2)];
+        policy.init_workers(&workers);
+
+        // Should select worker2 (lower load) despite cache affinity
+        let info = SelectWorkerInfo {
+            request_text: Some("test"),
+            ..Default::default()
+        };
+        for _ in 0..5 {
+            let idx = policy.select_worker(&workers, &info).await.unwrap();
+            assert_eq!(idx, 1); // Should always pick worker2
+        }
+    }
+
+    #[tokio::test]
+    async fn test_cache_aware_worker_removal() {
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0, // Disable eviction thread
+            ..Default::default()
+        };
+        let policy = CacheAwarePolicy::with_config(config);
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        policy.init_workers(&workers);
+
+        // Route some requests
+        policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("test1"),
+                    ..Default::default()
+                },
+            )
+            .await;
+        policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("test2"),
+                    ..Default::default()
+                },
+            )
+            .await;
+
+        // Remove a worker
+        policy.remove_worker_by_url("http://w1:8000");
+        workers[0].set_healthy(false);
+
+        // All requests should now go to worker2
+        let idx = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("test1"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+        assert_eq!(idx, 1);
+    }
+
+    #[tokio::test]
+    async fn test_cache_aware_sync_tree_operation_to_mesh() {
+        use std::sync::Arc;
+
+        use smg_mesh::{stores::StateStores, sync::MeshSyncManager};
+
+        let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
+        let mesh_sync = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
+
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let mut policy = CacheAwarePolicy::with_config(config);
+        policy.set_mesh_sync(Some(mesh_sync.clone()));
+
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
+            BasicWorkerBuilder::new("http://w1:8000")
+                .worker_type(WorkerType::Regular)
+                .api_key("test_api_key")
+                .build(),
+        )];
+
+        policy.init_workers(&workers);
+
+        // Select worker with a request - should sync to mesh
+        let _idx = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("test request"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+
+        // Verify tree operation was synced to mesh (under UNKNOWN_MODEL_ID since no model was specified)
+        let tree_state = mesh_sync.get_tree_state(UNKNOWN_MODEL_ID);
+        assert!(tree_state.is_some());
+        let tree = tree_state.unwrap();
+        assert!(!tree.operations.is_empty());
+    }
+
+    #[test]
+    fn test_cache_aware_restore_tree_state_from_mesh() {
+        use std::sync::Arc;
+
+        use smg_mesh::{
+            stores::StateStores,
+            sync::MeshSyncManager,
+            tree_ops::{TreeInsertOp, TreeOperation},
+        };
+
+        let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
+        let mesh_sync = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
+
+        // Pre-populate mesh with tree state
+        let op1 = TreeOperation::Insert(TreeInsertOp {
+            text: "test_text_1".to_string(),
+            tenant: "http://w1:8000".to_string(),
+        });
+        mesh_sync
+            .sync_tree_operation("model1".to_string(), op1)
+            .unwrap();
+
+        let op2 = TreeOperation::Insert(TreeInsertOp {
+            text: "test_text_2".to_string(),
+            tenant: "http://w2:8000".to_string(),
+        });
+        mesh_sync
+            .sync_tree_operation("model1".to_string(), op2)
+            .unwrap();
+
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let mut policy = CacheAwarePolicy::with_config(config);
+        policy.set_mesh_sync(Some(mesh_sync.clone()));
+
+        // Initialize with a model to trigger restore
+        let _workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
+            BasicWorkerBuilder::new("http://w1:8000")
+                .worker_type(WorkerType::Regular)
+                .api_key("test_api_key")
+                .build(),
+        )];
+
+        // Create a tree entry for model1 to trigger restore
+        let _tree = policy
+            .trees
+            .entry("model1".to_string())
+            .or_insert_with(|| Arc::new(Tree::new()));
+
+        // Manually trigger restore (normally done in constructor)
+        // For testing, we'll verify the tree state exists in mesh
+        let tree_state = mesh_sync.get_tree_state("model1");
+        assert!(tree_state.is_some());
+        let state = tree_state.unwrap();
+        assert_eq!(state.operations.len(), 2);
+    }
+
+    #[test]
+    fn test_cache_aware_apply_remote_tree_operation() {
+        use std::sync::Arc;
+
+        use smg_mesh::{
+            stores::StateStores,
+            sync::MeshSyncManager,
+            tree_ops::{TreeInsertOp, TreeOperation},
+        };
+
+        let stores = Arc::new(StateStores::with_self_name("node1".to_string()));
+        let mesh_sync = Arc::new(MeshSyncManager::new(stores, "node1".to_string()));
+
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let mut policy = CacheAwarePolicy::with_config(config);
+        policy.set_mesh_sync(Some(mesh_sync.clone()));
+
+        // Apply remote tree operation
+        let remote_op = TreeOperation::Insert(TreeInsertOp {
+            text: "remote_text".to_string(),
+            tenant: "http://remote:8000".to_string(),
+        });
+
+        policy.apply_remote_tree_operation("model1", &remote_op);
+
+        // Verify the tree was updated
+        let tree = policy.trees.get("model1");
+        assert!(tree.is_some());
+    }
+
+    #[test]
+    fn test_cache_aware_multi_node_consistency() {
+        use std::sync::Arc;
+
+        use smg_mesh::{
+            stores::StateStores,
+            sync::MeshSyncManager,
+            tree_ops::{TreeInsertOp, TreeOperation},
+        };
+
+        // Simulate two nodes
+        let stores1 = Arc::new(StateStores::with_self_name("node1".to_string()));
+        let mesh_sync1 = Arc::new(MeshSyncManager::new(stores1.clone(), "node1".to_string()));
+
+        let stores2 = Arc::new(StateStores::with_self_name("node2".to_string()));
+        let mesh_sync2 = Arc::new(MeshSyncManager::new(stores2.clone(), "node2".to_string()));
+
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+
+        let mut _policy1 = CacheAwarePolicy::with_config(config.clone());
+        _policy1.set_mesh_sync(Some(mesh_sync1.clone()));
+        let mut _policy2 = CacheAwarePolicy::with_config(config);
+        _policy2.set_mesh_sync(Some(mesh_sync2.clone()));
+
+        // Node1 syncs a tree operation
+        let op = TreeOperation::Insert(TreeInsertOp {
+            text: "shared_text".to_string(),
+            tenant: "http://shared:8000".to_string(),
+        });
+        mesh_sync1
+            .sync_tree_operation("model1".to_string(), op.clone())
+            .unwrap();
+
+        // Node2 should be able to get the tree state
+        let tree_state = mesh_sync2.get_tree_state("model1");
+        // Note: In a real scenario, this would be synced via gossip protocol
+        // For unit test, we verify the sync mechanism works
+        // Tree state may or may not exist depending on sync timing
+        let _ = tree_state;
+    }
+
+    #[tokio::test]
+    async fn test_cache_aware_without_mesh() {
+        let config = CacheAwareConfig {
+            eviction_interval_secs: 0,
+            ..Default::default()
+        };
+        let policy = CacheAwarePolicy::with_config(config);
+
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
+            BasicWorkerBuilder::new("http://w1:8000")
+                .worker_type(WorkerType::Regular)
+                .api_key("test_api_key")
+                .build(),
+        )];
+
+        policy.init_workers(&workers);
+
+        // Should work without mesh
+        let idx = policy
+            .select_worker(
+                &workers,
+                &SelectWorkerInfo {
+                    request_text: Some("test request"),
+                    ..Default::default()
+                },
+            )
+            .await
+            .unwrap();
+        assert_eq!(idx, 0);
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/consistent_hashing.rs

@@ -0,0 +1,521 @@
+//! Consistent hashing routing policy with header-based routing support
+//!
+//! Supports two routing mechanisms via HTTP headers:
+//! - `X-SMG-Target-Worker`: Direct routing by worker index (0-based), returns None if unavailable
+//! - `X-SMG-Routing-Key`: Consistent hash routing for session affinity
+//!
+//! ## Consistent Hashing
+//!
+//! Uses a pre-computed hash ring from WorkerRegistry where:
+//! 1. Each worker is placed at a fixed position based on hash(worker_url)
+//! 2. Keys are hashed to the ring, then walk clockwise to find first healthy worker
+//! 3. When workers scale up/down, only keys in the affected range redistribute (~1/N keys move)
+//!
+//! The ring is built once when workers are added/removed, not per-request.
+//! This ensures O(log n) lookup performance.
+//!
+//! Complexity: O(log n) binary search + O(k) walk where k = consecutive unhealthy workers.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use rand::Rng as _;
+
+use super::{LoadBalancingPolicy, SelectWorkerInfo};
+use crate::{
+    core::Worker,
+    observability::metrics::Metrics,
+    routers::header_utils::{extract_routing_key, extract_target_worker},
+};
+
+/// Execution branch for metrics
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Branch {
+    NoHealthyWorkers,
+    TargetWorkerHit,
+    TargetWorkerMiss,
+    RoutingKeyHit,
+    RandomFallback,
+}
+
+impl Branch {
+    #[inline]
+    const fn as_str(&self) -> &'static str {
+        match self {
+            Self::NoHealthyWorkers => "no_healthy_workers",
+            Self::TargetWorkerHit => "target_worker_hit",
+            Self::TargetWorkerMiss => "target_worker_miss",
+            Self::RoutingKeyHit => "routing_key_hit",
+            Self::RandomFallback => "random_fallback",
+        }
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct ConsistentHashingPolicy;
+
+impl ConsistentHashingPolicy {
+    pub fn new() -> Self {
+        Self
+    }
+
+    /// Use consistent hashing to find a worker for the given key.
+    /// Uses pre-computed ring from SelectWorkerInfo if available.
+    ///
+    /// The ring returns a worker URL, which we then map to an index in the workers array.
+    /// This correctly handles filtered worker arrays since we match by URL, not by index.
+    ///
+    /// Complexity: O(n) to build healthy URL map + O(log n) ring lookup + O(k) walk
+    fn find_by_consistent_hash(
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo,
+        key: &str,
+    ) -> Option<usize> {
+        // Build URL→index map for healthy workers: O(n) once, O(1) lookups
+        let healthy_url_to_idx: std::collections::HashMap<&str, usize> = workers
+            .iter()
+            .enumerate()
+            .filter(|(_, w)| w.is_healthy())
+            .map(|(i, w)| (w.url(), i))
+            .collect();
+
+        if healthy_url_to_idx.is_empty() {
+            return None;
+        }
+
+        // Use pre-computed ring if available
+        if let Some(ref ring) = info.hash_ring {
+            // O(1) lookup per URL checked instead of O(n)
+            let url = ring.find_healthy_url(key, |url| healthy_url_to_idx.contains_key(url))?;
+            return healthy_url_to_idx.get(url).copied();
+        }
+
+        // Fallback: no ring provided, use simple modulo (less optimal but functional)
+        // This shouldn't happen in normal operation as WorkerSelectionStage provides the ring
+        let mut healthy_indices: Vec<usize> = healthy_url_to_idx.values().copied().collect();
+        healthy_indices.sort_unstable(); // Ensure deterministic order
+
+        // Use blake3 for consistent hashing in fallback too
+        let hash = blake3::hash(key.as_bytes());
+        let hash_val = u64::from_le_bytes(hash.as_bytes()[..8].try_into().unwrap());
+        let idx = (hash_val as usize) % healthy_indices.len();
+        Some(healthy_indices[idx])
+    }
+
+    fn select_worker_impl(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo,
+    ) -> (Option<usize>, Branch) {
+        if workers.is_empty() {
+            return (None, Branch::NoHealthyWorkers);
+        }
+
+        let target_worker = extract_target_worker(info.headers);
+        let routing_key = extract_routing_key(info.headers);
+
+        // Priority 1: X-SMG-Target-Worker - direct routing by worker index
+        // O(1) parse + O(1) bounds check + O(1) health check
+        if let Some(idx_str) = target_worker {
+            if let Ok(idx) = idx_str.parse::<usize>() {
+                if idx < workers.len() && workers[idx].is_healthy() {
+                    return (Some(idx), Branch::TargetWorkerHit);
+                }
+            }
+            return (None, Branch::TargetWorkerMiss);
+        }
+
+        // Priority 2: X-SMG-Routing-Key - consistent hash routing (O(log n))
+        if let Some(key) = routing_key {
+            return match Self::find_by_consistent_hash(workers, info, key) {
+                Some(idx) => (Some(idx), Branch::RoutingKeyHit),
+                None => (None, Branch::NoHealthyWorkers),
+            };
+        }
+
+        // Priority 3: Implicit routing key from stable headers (session affinity)
+        let implicit_key = info.headers.and_then(|h| {
+            h.get("authorization")
+                .or_else(|| h.get("x-forwarded-for"))
+                .or_else(|| h.get("cookie"))
+                .and_then(|v| v.to_str().ok())
+                .filter(|s| !s.is_empty())
+        });
+
+        if let Some(key) = implicit_key {
+            return match Self::find_by_consistent_hash(workers, info, key) {
+                Some(idx) => (Some(idx), Branch::RoutingKeyHit),
+                None => (None, Branch::NoHealthyWorkers),
+            };
+        }
+
+        // Fallback: random selection (truly anonymous client)
+        let healthy_count = workers.iter().filter(|w| w.is_healthy()).count();
+        if healthy_count == 0 {
+            return (None, Branch::NoHealthyWorkers);
+        }
+
+        let random_healthy_idx = rand::rng().random_range(0..healthy_count);
+        let idx = workers
+            .iter()
+            .enumerate()
+            .filter(|(_, w)| w.is_healthy())
+            .nth(random_healthy_idx)
+            .map(|(i, _)| i)
+            .unwrap();
+
+        (Some(idx), Branch::RandomFallback)
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for ConsistentHashingPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let (result, branch) = self.select_worker_impl(workers, info);
+        Metrics::record_worker_consistent_hashing_policy_branch(branch.as_str());
+        result
+    }
+
+    fn name(&self) -> &'static str {
+        "consistent_hashing"
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, HashRing, WorkerType};
+
+    fn headers_with_routing_key(key: &str) -> http::HeaderMap {
+        let mut headers = http::HeaderMap::new();
+        headers.insert("x-smg-routing-key", key.parse().unwrap());
+        headers
+    }
+
+    fn headers_with_target_worker(idx: usize) -> http::HeaderMap {
+        let mut headers = http::HeaderMap::new();
+        headers.insert("x-smg-target-worker", idx.to_string().parse().unwrap());
+        headers
+    }
+
+    fn create_workers(urls: &[&str]) -> Vec<Arc<dyn Worker>> {
+        urls.iter()
+            .map(|url| {
+                Arc::new(
+                    BasicWorkerBuilder::new(*url)
+                        .worker_type(WorkerType::Regular)
+                        .build(),
+                ) as Arc<dyn Worker>
+            })
+            .collect()
+    }
+
+    #[tokio::test]
+    async fn test_consistent_routing() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let headers = headers_with_routing_key("user-123");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, _) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+
+        // Same key should always route to same worker
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(result, Some(first_idx));
+            assert_eq!(branch, Branch::RoutingKeyHit);
+        }
+    }
+
+    #[tokio::test]
+    async fn test_different_keys_distribute() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let mut distribution = HashMap::new();
+        for i in 0..100 {
+            let headers = headers_with_routing_key(&format!("user-{}", i));
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                ..Default::default()
+            };
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            *distribution.entry(result.unwrap()).or_insert(0) += 1;
+        }
+
+        assert!(distribution.len() > 1, "Should distribute across workers");
+    }
+
+    #[tokio::test]
+    async fn test_target_worker_hit() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_target_worker(1);
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, Some(1));
+        assert_eq!(branch, Branch::TargetWorkerHit);
+    }
+
+    #[tokio::test]
+    async fn test_target_worker_miss_out_of_bounds() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_target_worker(5); // Out of bounds
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::TargetWorkerMiss);
+    }
+
+    #[tokio::test]
+    async fn test_target_worker_miss_unhealthy() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+        workers[1].set_healthy(false);
+
+        let headers = headers_with_target_worker(1);
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::TargetWorkerMiss);
+    }
+
+    #[tokio::test]
+    async fn test_target_worker_priority_over_routing_key() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let mut headers = http::HeaderMap::new();
+        headers.insert("x-smg-target-worker", "1".parse().unwrap());
+        headers.insert("x-smg-routing-key", "some-key".parse().unwrap());
+
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, Some(1));
+        assert_eq!(branch, Branch::TargetWorkerHit);
+    }
+
+    #[tokio::test]
+    async fn test_fallback_random_distribution() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        // Without routing headers, should distribute randomly across workers
+        let mut distribution = HashMap::new();
+        for _ in 0..100 {
+            let info = SelectWorkerInfo::default();
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert!(result.is_some());
+            assert_eq!(branch, Branch::RandomFallback);
+            *distribution.entry(result.unwrap()).or_insert(0) += 1;
+        }
+
+        // Should distribute across multiple workers (not always same one)
+        assert!(
+            distribution.len() > 1,
+            "Random fallback should distribute across workers"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_no_healthy_workers() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000"]);
+        workers[0].set_healthy(false);
+
+        let headers = headers_with_routing_key("test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::NoHealthyWorkers);
+    }
+
+    #[tokio::test]
+    async fn test_empty_workers() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![];
+
+        let info = SelectWorkerInfo::default();
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::NoHealthyWorkers);
+    }
+
+    #[tokio::test]
+    async fn test_consistent_hash_minimal_redistribution() {
+        // Test that consistent hashing moves fewer keys than random redistribution
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&[
+            "http://w0:8000",
+            "http://w1:8000",
+            "http://w2:8000",
+            "http://w3:8000",
+        ]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        // Record which worker each key routes to with all workers healthy
+        let mut key_to_worker_before: HashMap<String, usize> = HashMap::new();
+        for i in 0..100 {
+            let key = format!("user-{}", i);
+            let headers = headers_with_routing_key(&key);
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                hash_ring: Some(ring.clone()),
+                ..Default::default()
+            };
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            key_to_worker_before.insert(key, result.unwrap());
+        }
+
+        // Mark worker 1 as unhealthy
+        workers[1].set_healthy(false);
+
+        // Record new routing and count how many keys moved
+        let mut moved_count = 0;
+        for i in 0..100 {
+            let key = format!("user-{}", i);
+            let headers = headers_with_routing_key(&key);
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                hash_ring: Some(ring.clone()),
+                ..Default::default()
+            };
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            let new_worker = result.unwrap();
+            let old_worker = key_to_worker_before[&key];
+
+            if new_worker != old_worker {
+                moved_count += 1;
+            }
+        }
+
+        // With consistent hashing, approximately 1/N keys should move (N = worker count)
+        // Random redistribution would move approximately (N-1)/N = 75% of keys
+        // Verify we're significantly better than random (< 50% moved)
+        let keys_on_failed_worker = key_to_worker_before.values().filter(|&&w| w == 1).count();
+        assert!(
+            moved_count <= keys_on_failed_worker + 5,
+            "Consistent hashing should only move keys from failed worker (+small variance). \
+             Expected ~{}, got {}",
+            keys_on_failed_worker,
+            moved_count
+        );
+        assert!(
+            moved_count < 50,
+            "Consistent hashing should move fewer than 50% of keys (random would move ~75%), got {}%",
+            moved_count
+        );
+    }
+
+    #[tokio::test]
+    async fn test_routing_key_failover_and_recovery() {
+        // Test that when a worker fails, keys move to another worker,
+        // and when it recovers, keys return to the original worker
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w0:8000", "http://w1:8000", "http://w2:8000"]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        // Find which worker a key routes to when all are healthy
+        let test_key = "session-abc-123";
+        let headers = headers_with_routing_key(test_key);
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            hash_ring: Some(ring.clone()),
+            ..Default::default()
+        };
+        let (result, _) = policy.select_worker_impl(&workers, &info);
+        let original_idx = result.unwrap();
+
+        // Mark that worker unhealthy
+        workers[original_idx].set_healthy(false);
+
+        // Key should now route to a different healthy worker
+        let (failover_result, _) = policy.select_worker_impl(&workers, &info);
+        let failover_idx = failover_result.unwrap();
+        assert_ne!(
+            failover_idx, original_idx,
+            "Should failover to different worker"
+        );
+        assert!(
+            workers[failover_idx].is_healthy(),
+            "Failover target should be healthy"
+        );
+
+        // Failover should be consistent
+        for _ in 0..5 {
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(result, Some(failover_idx), "Failover should be consistent");
+        }
+
+        // Recover the original worker
+        workers[original_idx].set_healthy(true);
+
+        // Key should route back to original worker
+        let (recovered_result, _) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(
+            recovered_result,
+            Some(original_idx),
+            "Should return to original worker after recovery"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_empty_routing_key_uses_fallback() {
+        let policy = ConsistentHashingPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_routing_key("");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert!(result.is_some());
+        assert_eq!(branch, Branch::RandomFallback);
+    }
+
+    #[tokio::test]
+    async fn test_policy_name() {
+        let policy = ConsistentHashingPolicy::new();
+        assert_eq!(policy.name(), "consistent_hashing");
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/factory.rs

@@ -0,0 +1,156 @@
+//! Factory for creating load balancing policies
+
+use std::sync::Arc;
+
+use super::{
+    BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy,
+    LoadBalancingPolicy, ManualConfig, ManualPolicy, PowerOfTwoPolicy, PrefixHashConfig,
+    PrefixHashPolicy, RandomPolicy, RoundRobinPolicy,
+};
+use crate::config::PolicyConfig;
+
+/// Factory for creating policy instances
+pub struct PolicyFactory;
+
+impl PolicyFactory {
+    /// Create a policy from configuration
+    pub fn create_from_config(config: &PolicyConfig) -> Arc<dyn LoadBalancingPolicy> {
+        match config {
+            PolicyConfig::Random => Arc::new(RandomPolicy::new()),
+            PolicyConfig::RoundRobin => Arc::new(RoundRobinPolicy::new()),
+            PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()),
+            PolicyConfig::CacheAware {
+                cache_threshold,
+                balance_abs_threshold,
+                balance_rel_threshold,
+                eviction_interval_secs,
+                max_tree_size,
+            } => {
+                let config = CacheAwareConfig {
+                    cache_threshold: *cache_threshold,
+                    balance_abs_threshold: *balance_abs_threshold,
+                    balance_rel_threshold: *balance_rel_threshold,
+                    eviction_interval_secs: *eviction_interval_secs,
+                    max_tree_size: *max_tree_size,
+                };
+                Arc::new(CacheAwarePolicy::with_config(config))
+            }
+            PolicyConfig::Bucket {
+                balance_abs_threshold,
+                balance_rel_threshold,
+                bucket_adjust_interval_secs,
+            } => {
+                let config = BucketConfig {
+                    balance_abs_threshold: *balance_abs_threshold,
+                    balance_rel_threshold: *balance_rel_threshold,
+                    bucket_adjust_interval_secs: *bucket_adjust_interval_secs,
+                };
+                Arc::new(BucketPolicy::with_config(config))
+            }
+            PolicyConfig::Manual {
+                eviction_interval_secs,
+                max_idle_secs,
+                assignment_mode,
+            } => {
+                let config = ManualConfig {
+                    eviction_interval_secs: *eviction_interval_secs,
+                    max_idle_secs: *max_idle_secs,
+                    assignment_mode: *assignment_mode,
+                };
+                Arc::new(ManualPolicy::with_config(config))
+            }
+            PolicyConfig::ConsistentHashing => Arc::new(ConsistentHashingPolicy::new()),
+            PolicyConfig::PrefixHash {
+                prefix_token_count,
+                load_factor,
+            } => {
+                let config = PrefixHashConfig {
+                    prefix_token_count: *prefix_token_count,
+                    load_factor: *load_factor,
+                };
+                Arc::new(PrefixHashPolicy::new(config))
+            }
+        }
+    }
+
+    /// Create a policy by name (for dynamic loading)
+    pub fn create_by_name(name: &str) -> Option<Arc<dyn LoadBalancingPolicy>> {
+        match name.to_lowercase().as_str() {
+            "random" => Some(Arc::new(RandomPolicy::new())),
+            "round_robin" | "roundrobin" => Some(Arc::new(RoundRobinPolicy::new())),
+            "power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())),
+            "cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())),
+            "bucket" => Some(Arc::new(BucketPolicy::new())),
+            "manual" => Some(Arc::new(ManualPolicy::new())),
+            "consistent_hashing" | "consistenthashing" => {
+                Some(Arc::new(ConsistentHashingPolicy::new()))
+            }
+            "prefix_hash" | "prefixhash" => Some(Arc::new(PrefixHashPolicy::with_defaults())),
+            _ => None,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_create_from_config() {
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::Random);
+        assert_eq!(policy.name(), "random");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::RoundRobin);
+        assert_eq!(policy.name(), "round_robin");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::PowerOfTwo {
+            load_check_interval_secs: 60,
+        });
+        assert_eq!(policy.name(), "power_of_two");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::CacheAware {
+            cache_threshold: 0.7,
+            balance_abs_threshold: 10,
+            balance_rel_threshold: 1.5,
+            eviction_interval_secs: 30,
+            max_tree_size: 1000,
+        });
+        assert_eq!(policy.name(), "cache_aware");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::Bucket {
+            balance_abs_threshold: 10,
+            balance_rel_threshold: 1.5,
+            bucket_adjust_interval_secs: 5,
+        });
+        assert_eq!(policy.name(), "bucket");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::Manual {
+            eviction_interval_secs: 60,
+            max_idle_secs: 4 * 3600,
+            assignment_mode: Default::default(),
+        });
+        assert_eq!(policy.name(), "manual");
+
+        let policy = PolicyFactory::create_from_config(&PolicyConfig::ConsistentHashing);
+        assert_eq!(policy.name(), "consistent_hashing");
+    }
+
+    #[tokio::test]
+    async fn test_create_by_name() {
+        assert!(PolicyFactory::create_by_name("random").is_some());
+        assert!(PolicyFactory::create_by_name("RANDOM").is_some());
+        assert!(PolicyFactory::create_by_name("round_robin").is_some());
+        assert!(PolicyFactory::create_by_name("RoundRobin").is_some());
+        assert!(PolicyFactory::create_by_name("power_of_two").is_some());
+        assert!(PolicyFactory::create_by_name("PowerOfTwo").is_some());
+        assert!(PolicyFactory::create_by_name("cache_aware").is_some());
+        assert!(PolicyFactory::create_by_name("CacheAware").is_some());
+        assert!(PolicyFactory::create_by_name("bucket").is_some());
+        assert!(PolicyFactory::create_by_name("Bucket").is_some());
+        assert!(PolicyFactory::create_by_name("manual").is_some());
+        assert!(PolicyFactory::create_by_name("Manual").is_some());
+        assert!(PolicyFactory::create_by_name("consistent_hashing").is_some());
+        assert!(PolicyFactory::create_by_name("ConsistentHashing").is_some());
+        assert!(PolicyFactory::create_by_name("unknown").is_none());
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/manual.rs

@@ -0,0 +1,981 @@
+//! Manual routing policy based on routing key header
+//!
+//! This policy provides sticky session routing where each unique routing key
+//! is consistently mapped to the same worker. Unlike consistent hashing,
+//! this policy:
+//! - Does NOT redistribute any sessions when workers are added
+//! - Only remaps sessions when their assigned worker becomes unhealthy
+//! - Maintains up to 2 candidate workers per routing key for fast failover
+//!
+//! Use this when you need stronger stickiness guarantees than consistent hashing,
+//! for example with stateful chat sessions where context is stored on the worker.
+//!
+//! ## Header
+//! - `X-SMG-Routing-Key`: The routing key for sticky session routing
+
+use std::{sync::Arc, time::Instant};
+
+use async_trait::async_trait;
+use dashmap::{mapref::entry::Entry, DashMap};
+use rand::Rng;
+use tracing::info;
+
+use super::{
+    get_healthy_worker_indices, utils::PeriodicTask, LoadBalancingPolicy, SelectWorkerInfo,
+};
+use crate::{
+    config::ManualAssignmentMode, core::Worker, observability::metrics::Metrics,
+    routers::header_utils::extract_routing_key,
+};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum ExecutionBranch {
+    NoHealthyWorkers,
+    OccupiedHit,
+    OccupiedMiss,
+    Vacant,
+    NoRoutingId,
+}
+
+impl ExecutionBranch {
+    fn as_str(&self) -> &'static str {
+        match self {
+            Self::NoHealthyWorkers => "no_healthy_workers",
+            Self::OccupiedHit => "occupied_hit",
+            Self::OccupiedMiss => "occupied_miss",
+            Self::Vacant => "vacant",
+            Self::NoRoutingId => "no_routing_id",
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+struct RoutingId(String);
+
+impl RoutingId {
+    fn new(id: impl Into<String>) -> Self {
+        Self(id.into())
+    }
+}
+
+const MAX_CANDIDATE_WORKERS: usize = 2;
+
+#[derive(Debug, Clone)]
+pub struct ManualConfig {
+    pub eviction_interval_secs: u64,
+    pub max_idle_secs: u64,
+    pub assignment_mode: ManualAssignmentMode,
+}
+
+impl Default for ManualConfig {
+    fn default() -> Self {
+        Self {
+            eviction_interval_secs: 60,
+            max_idle_secs: 4 * 3600,
+            assignment_mode: ManualAssignmentMode::Random,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+struct Node {
+    candi_worker_urls: Vec<String>,
+    last_access: Instant,
+}
+
+impl Node {
+    fn push_bounded(&mut self, url: String) {
+        while self.candi_worker_urls.len() >= MAX_CANDIDATE_WORKERS {
+            self.candi_worker_urls.remove(0);
+        }
+        self.candi_worker_urls.push(url);
+    }
+}
+
+#[derive(Debug)]
+pub struct ManualPolicy {
+    routing_map: Arc<DashMap<RoutingId, Node>>,
+    assignment_mode: ManualAssignmentMode,
+    _eviction_task: Option<PeriodicTask>,
+}
+
+impl Default for ManualPolicy {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl ManualPolicy {
+    pub fn new() -> Self {
+        Self::with_config(ManualConfig::default())
+    }
+
+    pub fn with_config(config: ManualConfig) -> Self {
+        use std::time::Duration;
+
+        let routing_map = Arc::new(DashMap::<RoutingId, Node>::new());
+
+        let eviction_task = if config.eviction_interval_secs > 0 && config.max_idle_secs > 0 {
+            let routing_map_clone = Arc::clone(&routing_map);
+            let max_idle = Duration::from_secs(config.max_idle_secs);
+
+            Some(PeriodicTask::spawn(
+                config.eviction_interval_secs,
+                "ManualPolicyEviction",
+                move || {
+                    let now = Instant::now();
+                    let before_size = routing_map_clone.len();
+
+                    routing_map_clone
+                        .retain(|_, node| now.duration_since(node.last_access) < max_idle);
+
+                    let evicted_count = before_size - routing_map_clone.len();
+                    if evicted_count > 0 {
+                        info!(
+                            "ManualPolicy TTL eviction: evicted {} entries, remaining {} (max_idle: {}s)",
+                            evicted_count,
+                            routing_map_clone.len(),
+                            max_idle.as_secs()
+                        );
+                    }
+                },
+            ))
+        } else {
+            None
+        };
+
+        Self {
+            routing_map,
+            assignment_mode: config.assignment_mode,
+            _eviction_task: eviction_task,
+        }
+    }
+
+    fn select_new_worker(&self, workers: &[Arc<dyn Worker>], healthy_indices: &[usize]) -> usize {
+        match self.assignment_mode {
+            ManualAssignmentMode::Random => random_select(healthy_indices),
+            ManualAssignmentMode::MinLoad => min_load_select(workers, healthy_indices),
+            ManualAssignmentMode::MinGroup => min_group_select(workers, healthy_indices),
+        }
+    }
+
+    fn select_by_routing_id(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        routing_id: &str,
+        healthy_indices: &[usize],
+    ) -> (usize, ExecutionBranch) {
+        let routing_id = RoutingId::new(routing_id);
+
+        match self.routing_map.entry(routing_id) {
+            Entry::Occupied(mut entry) => {
+                let node = entry.get_mut();
+                node.last_access = Instant::now();
+                if let Some(idx) =
+                    find_healthy_worker(&node.candi_worker_urls, workers, healthy_indices)
+                {
+                    (idx, ExecutionBranch::OccupiedHit)
+                } else {
+                    let selected_idx = self.select_new_worker(workers, healthy_indices);
+                    node.push_bounded(workers[selected_idx].url().to_string());
+                    (selected_idx, ExecutionBranch::OccupiedMiss)
+                }
+            }
+            Entry::Vacant(entry) => {
+                let selected_idx = self.select_new_worker(workers, healthy_indices);
+                entry.insert(Node {
+                    candi_worker_urls: vec![workers[selected_idx].url().to_string()],
+                    last_access: Instant::now(),
+                });
+                (selected_idx, ExecutionBranch::Vacant)
+            }
+        }
+    }
+
+    fn select_worker_impl(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> (Option<usize>, ExecutionBranch) {
+        let healthy_indices = get_healthy_worker_indices(workers);
+        if healthy_indices.is_empty() {
+            return (None, ExecutionBranch::NoHealthyWorkers);
+        }
+
+        if let Some(routing_id) = extract_routing_key(info.headers) {
+            let (idx, branch) = self.select_by_routing_id(workers, routing_id, &healthy_indices);
+            return (Some(idx), branch);
+        }
+
+        (
+            Some(random_select(&healthy_indices)),
+            ExecutionBranch::NoRoutingId,
+        )
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for ManualPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let (result, branch) = self.select_worker_impl(workers, info);
+        Metrics::record_worker_manual_policy_branch(branch.as_str());
+        Metrics::set_manual_policy_cache_entries(self.routing_map.len());
+        result
+    }
+
+    fn name(&self) -> &'static str {
+        "manual"
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+fn find_healthy_worker(
+    urls: &[String],
+    workers: &[Arc<dyn Worker>],
+    healthy_indices: &[usize],
+) -> Option<usize> {
+    for url in urls {
+        if let Some(idx) = find_worker_index_by_url(workers, url) {
+            if healthy_indices.contains(&idx) {
+                return Some(idx);
+            }
+        }
+    }
+    None
+}
+
+fn find_worker_index_by_url(workers: &[Arc<dyn Worker>], url: &str) -> Option<usize> {
+    workers.iter().position(|w| w.url() == url)
+}
+
+fn random_select(healthy_indices: &[usize]) -> usize {
+    let mut rng = rand::rng();
+    let random_idx = rng.random_range(0..healthy_indices.len());
+    healthy_indices[random_idx]
+}
+
+fn select_min_by<K, V, F>(indices: &[K], get_value: F) -> K
+where
+    K: Copy,
+    V: Ord,
+    F: Fn(K) -> V,
+{
+    let mut min_val: Option<V> = None;
+    let mut candidates = Vec::new();
+
+    for &idx in indices {
+        let val = get_value(idx);
+        match min_val.as_ref().map(|m| val.cmp(m)) {
+            None | Some(std::cmp::Ordering::Less) => {
+                min_val = Some(val);
+                candidates.clear();
+                candidates.push(idx);
+            }
+            Some(std::cmp::Ordering::Equal) => {
+                candidates.push(idx);
+            }
+            Some(std::cmp::Ordering::Greater) => {}
+        }
+    }
+
+    if candidates.len() == 1 {
+        candidates[0]
+    } else {
+        let mut rng = rand::rng();
+        candidates[rng.random_range(0..candidates.len())]
+    }
+}
+
+fn min_load_select(workers: &[Arc<dyn Worker>], healthy_indices: &[usize]) -> usize {
+    select_min_by(healthy_indices, |idx| workers[idx].load())
+}
+
+fn min_group_select(workers: &[Arc<dyn Worker>], healthy_indices: &[usize]) -> usize {
+    select_min_by(healthy_indices, |idx| {
+        workers[idx].worker_routing_key_load().value()
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    fn create_workers(urls: &[&str]) -> Vec<Arc<dyn Worker>> {
+        urls.iter()
+            .map(|url| {
+                Arc::new(
+                    BasicWorkerBuilder::new(*url)
+                        .worker_type(WorkerType::Regular)
+                        .build(),
+                ) as Arc<dyn Worker>
+            })
+            .collect()
+    }
+
+    fn headers_with_routing_key(key: &str) -> http::HeaderMap {
+        let mut headers = http::HeaderMap::new();
+        headers.insert("x-smg-routing-key", key.parse().unwrap());
+        headers
+    }
+
+    #[test]
+    fn test_manual_consistent_routing() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let headers = headers_with_routing_key("user-123");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, branch) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(
+                result,
+                Some(first_idx),
+                "Same routing_id should route to same worker"
+            );
+            assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        }
+    }
+
+    #[test]
+    fn test_manual_different_routing_ids() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let mut distribution = HashMap::new();
+        for i in 0..100 {
+            let headers = headers_with_routing_key(&format!("user-{}", i));
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                ..Default::default()
+            };
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(branch, ExecutionBranch::Vacant);
+            *distribution.entry(result.unwrap()).or_insert(0) += 1;
+        }
+
+        assert!(
+            distribution.len() > 1,
+            "Should distribute across multiple workers"
+        );
+    }
+
+    #[test]
+    fn test_manual_fallback_random() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let mut counts = HashMap::new();
+        for _ in 0..100 {
+            let info = SelectWorkerInfo::default();
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(branch, ExecutionBranch::NoRoutingId);
+            if let Some(idx) = result {
+                *counts.entry(idx).or_insert(0) += 1;
+            }
+        }
+
+        assert_eq!(counts.len(), 2, "Random fallback should use all workers");
+    }
+
+    #[test]
+    fn test_manual_with_unhealthy_workers() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        workers[0].set_healthy(false);
+
+        let headers = headers_with_routing_key("test-routing-id");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, Some(1), "Should only select healthy worker");
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(result, Some(1), "Should only select healthy worker");
+            assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        }
+    }
+
+    #[test]
+    fn test_manual_no_healthy_workers() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000"]);
+
+        workers[0].set_healthy(false);
+        let headers = headers_with_routing_key("test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, ExecutionBranch::NoHealthyWorkers);
+    }
+
+    #[test]
+    fn test_manual_empty_routing_id() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let mut counts = HashMap::new();
+        for _ in 0..100 {
+            let headers = headers_with_routing_key("");
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                ..Default::default()
+            };
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(branch, ExecutionBranch::NoRoutingId);
+            if let Some(idx) = result {
+                *counts.entry(idx).or_insert(0) += 1;
+            }
+        }
+
+        assert_eq!(
+            counts.len(),
+            2,
+            "Empty routing_id should use random fallback"
+        );
+    }
+
+    #[test]
+    fn test_manual_remaps_when_worker_becomes_unhealthy() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_routing_key("sticky-user");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, branch) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        workers[first_idx].set_healthy(false);
+
+        let (new_result, branch) = policy.select_worker_impl(&workers, &info);
+        let new_idx = new_result.unwrap();
+        assert_ne!(new_idx, first_idx, "Should remap to healthy worker");
+        assert_eq!(branch, ExecutionBranch::OccupiedMiss);
+
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(
+                result,
+                Some(new_idx),
+                "Should consistently route to new worker"
+            );
+            assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        }
+    }
+
+    #[test]
+    fn test_manual_empty_workers() {
+        let policy = ManualPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![];
+        let headers = headers_with_routing_key("test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, ExecutionBranch::NoHealthyWorkers);
+    }
+
+    #[test]
+    fn test_manual_single_worker() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000"]);
+
+        let headers = headers_with_routing_key("single-test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, Some(0));
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(result, Some(0));
+            assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        }
+    }
+
+    #[test]
+    fn test_manual_worker_recovery() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_routing_key("recovery-test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, branch) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        workers[first_idx].set_healthy(false);
+
+        let (second_result, branch) = policy.select_worker_impl(&workers, &info);
+        let second_idx = second_result.unwrap();
+        assert_ne!(second_idx, first_idx);
+        assert_eq!(branch, ExecutionBranch::OccupiedMiss);
+
+        workers[first_idx].set_healthy(true);
+
+        let (after_recovery, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(
+            after_recovery,
+            Some(first_idx),
+            "Should return to original worker after recovery since it's first in candidate list"
+        );
+        assert_eq!(branch, ExecutionBranch::OccupiedHit);
+    }
+
+    #[test]
+    fn test_manual_max_candidate_workers_eviction() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let headers = headers_with_routing_key("eviction-test");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, branch) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+        assert_eq!(branch, ExecutionBranch::Vacant);
+
+        workers[first_idx].set_healthy(false);
+
+        let (second_result, branch) = policy.select_worker_impl(&workers, &info);
+        let second_idx = second_result.unwrap();
+        assert_ne!(second_idx, first_idx);
+        assert_eq!(branch, ExecutionBranch::OccupiedMiss);
+
+        workers[second_idx].set_healthy(false);
+
+        let remaining_idx = (0..3).find(|&i| i != first_idx && i != second_idx).unwrap();
+        let (third_result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(
+            third_result,
+            Some(remaining_idx),
+            "Should select the only remaining healthy worker"
+        );
+        assert_eq!(branch, ExecutionBranch::OccupiedMiss);
+
+        workers[first_idx].set_healthy(true);
+
+        let (idx_after_restore, branch) = policy.select_worker_impl(&workers, &info);
+        assert_ne!(
+            idx_after_restore,
+            Some(first_idx),
+            "First worker should be evicted from candidates due to MAX_CANDIDATE_WORKERS=2"
+        );
+        assert_eq!(branch, ExecutionBranch::OccupiedHit);
+    }
+
+    #[test]
+    fn test_manual_policy_name() {
+        let policy = ManualPolicy::new();
+        assert_eq!(policy.name(), "manual");
+    }
+
+    #[test]
+    fn test_manual_routing_info_push_bounded() {
+        let mut info = Node {
+            candi_worker_urls: vec!["http://w1:8000".to_string()],
+            last_access: Instant::now(),
+        };
+
+        info.push_bounded("http://w2:8000".to_string());
+        assert_eq!(info.candi_worker_urls.len(), 2);
+        assert_eq!(info.candi_worker_urls[0], "http://w1:8000");
+        assert_eq!(info.candi_worker_urls[1], "http://w2:8000");
+
+        info.push_bounded("http://w3:8000".to_string());
+        assert_eq!(info.candi_worker_urls.len(), 2);
+        assert_eq!(
+            info.candi_worker_urls[0], "http://w2:8000",
+            "Oldest entry should be removed"
+        );
+        assert_eq!(info.candi_worker_urls[1], "http://w3:8000");
+    }
+
+    #[test]
+    fn test_manual_find_healthy_worker_priority() {
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        let urls = vec![
+            "http://w1:8000".to_string(),
+            "http://w2:8000".to_string(),
+            "http://w3:8000".to_string(),
+        ];
+        let healthy_indices = vec![0, 1, 2];
+
+        let result = find_healthy_worker(&urls, &workers, &healthy_indices);
+        assert_eq!(
+            result,
+            Some(0),
+            "Should return first healthy worker in urls"
+        );
+
+        workers[0].set_healthy(false);
+        let healthy_indices = vec![1, 2];
+        let result = find_healthy_worker(&urls, &workers, &healthy_indices);
+        assert_eq!(result, Some(1), "Should skip unhealthy and return next");
+
+        workers[1].set_healthy(false);
+        let healthy_indices = vec![2];
+        let result = find_healthy_worker(&urls, &workers, &healthy_indices);
+        assert_eq!(result, Some(2), "Should return last healthy worker");
+
+        workers[2].set_healthy(false);
+        let healthy_indices: Vec<usize> = vec![];
+        let result = find_healthy_worker(&urls, &workers, &healthy_indices);
+        assert_eq!(result, None, "Should return None when no healthy workers");
+    }
+
+    #[test]
+    fn test_manual_find_worker_index_by_url() {
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        assert_eq!(
+            find_worker_index_by_url(&workers, "http://w1:8000"),
+            Some(0)
+        );
+        assert_eq!(
+            find_worker_index_by_url(&workers, "http://w2:8000"),
+            Some(1)
+        );
+        assert_eq!(
+            find_worker_index_by_url(&workers, "http://w3:8000"),
+            None,
+            "Should return None for unknown URL"
+        );
+    }
+
+    #[test]
+    fn test_manual_config_default() {
+        let config = ManualConfig::default();
+        assert_eq!(config.eviction_interval_secs, 60);
+        assert_eq!(config.max_idle_secs, 4 * 3600);
+    }
+
+    #[test]
+    fn test_manual_with_disabled_eviction() {
+        let config = ManualConfig {
+            eviction_interval_secs: 0,
+            max_idle_secs: 3600,
+            assignment_mode: ManualAssignmentMode::Random,
+        };
+        let policy = ManualPolicy::with_config(config);
+        assert!(policy._eviction_task.is_none());
+    }
+
+    #[test]
+    fn test_manual_last_access_updates() {
+        let policy = ManualPolicy::new();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+        let headers = headers_with_routing_key("test-key");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        let routing_id = RoutingId::new("test-key");
+
+        // Vacant: first access
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(branch, ExecutionBranch::Vacant);
+        let first_idx = result.unwrap();
+        let access_after_vacant = policy.routing_map.get(&routing_id).unwrap().last_access;
+        assert!(access_after_vacant.elapsed().as_millis() < 100);
+
+        std::thread::sleep(std::time::Duration::from_millis(10));
+
+        // OccupiedHit: same worker still healthy
+        let (_, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        let access_after_hit = policy.routing_map.get(&routing_id).unwrap().last_access;
+        assert!(access_after_hit > access_after_vacant);
+
+        std::thread::sleep(std::time::Duration::from_millis(10));
+
+        // OccupiedMiss: worker becomes unhealthy
+        workers[first_idx].set_healthy(false);
+        let (_, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(branch, ExecutionBranch::OccupiedMiss);
+        let access_after_miss = policy.routing_map.get(&routing_id).unwrap().last_access;
+        assert!(access_after_miss > access_after_hit);
+    }
+
+    #[test]
+    fn test_manual_ttl_eviction_logic() {
+        use std::time::Duration;
+
+        let config = ManualConfig {
+            eviction_interval_secs: 2,
+            max_idle_secs: 2,
+            assignment_mode: ManualAssignmentMode::Random,
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        let headers = headers_with_routing_key("key-0");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        policy.select_worker_impl(&workers, &info);
+
+        assert_eq!(policy.routing_map.len(), 1);
+
+        std::thread::sleep(Duration::from_secs(4));
+
+        assert_eq!(policy.routing_map.len(), 0);
+    }
+
+    #[test]
+    fn test_min_group_select_distributes_evenly() {
+        let config = ManualConfig {
+            assignment_mode: ManualAssignmentMode::MinGroup,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        for i in 0..9 {
+            let routing_key = format!("key-{}", i);
+            let headers = headers_with_routing_key(&routing_key);
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                ..Default::default()
+            };
+
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert!(result.is_some());
+            assert_eq!(branch, ExecutionBranch::Vacant);
+
+            let selected_idx = result.unwrap();
+            workers[selected_idx]
+                .worker_routing_key_load()
+                .increment(&routing_key);
+        }
+
+        let distribution: HashMap<_, usize> = policy
+            .routing_map
+            .iter()
+            .map(|e| e.candi_worker_urls.first().unwrap().clone())
+            .fold(HashMap::new(), |mut acc, url| {
+                *acc.entry(url).or_default() += 1;
+                acc
+            });
+
+        assert_eq!(distribution.len(), 3, "Should use all 3 workers");
+        for count in distribution.values() {
+            assert_eq!(*count, 3, "Each worker should have exactly 3 routing keys");
+        }
+    }
+
+    #[test]
+    fn test_min_group_select_prefers_worker_with_fewer_routing_keys() {
+        let config = ManualConfig {
+            assignment_mode: ManualAssignmentMode::MinGroup,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        workers[0].worker_routing_key_load().increment("existing-1");
+        workers[0].worker_routing_key_load().increment("existing-2");
+        workers[1].worker_routing_key_load().increment("existing-3");
+
+        assert_eq!(workers[0].worker_routing_key_load().value(), 2);
+        assert_eq!(workers[1].worker_routing_key_load().value(), 1);
+        assert_eq!(workers[2].worker_routing_key_load().value(), 0);
+
+        let headers = headers_with_routing_key("new-key");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        let (result, _) = policy.select_worker_impl(&workers, &info);
+        let selected_idx = result.unwrap();
+
+        assert_eq!(selected_idx, 2, "Should select worker with 0 routing keys");
+    }
+
+    #[test]
+    fn test_min_load_select_prefers_worker_with_fewer_requests() {
+        let config = ManualConfig {
+            assignment_mode: ManualAssignmentMode::MinLoad,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+
+        workers[0].increment_load();
+        workers[0].increment_load();
+        workers[1].increment_load();
+
+        assert_eq!(workers[0].load(), 2);
+        assert_eq!(workers[1].load(), 1);
+        assert_eq!(workers[2].load(), 0);
+
+        let headers = headers_with_routing_key("new-key");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+        let (result, _) = policy.select_worker_impl(&workers, &info);
+        let selected_idx = result.unwrap();
+
+        assert_eq!(selected_idx, 2, "Should select worker with 0 load");
+    }
+
+    #[test]
+    fn test_min_group_sticky_after_assignment() {
+        let config = ManualConfig {
+            assignment_mode: ManualAssignmentMode::MinGroup,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        workers[0].worker_routing_key_load().increment("key-0");
+        workers[1].worker_routing_key_load().increment("key-1");
+        workers[1].worker_routing_key_load().increment("key-2");
+
+        let headers = headers_with_routing_key("new-key");
+        let info = SelectWorkerInfo {
+            headers: Some(&headers),
+            ..Default::default()
+        };
+
+        let (first_result, branch) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+        assert_eq!(branch, ExecutionBranch::Vacant);
+        assert_eq!(
+            first_idx, 0,
+            "Should select worker 0 (has 1 routing key vs 2)"
+        );
+
+        for _ in 0..10 {
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(
+                result,
+                Some(first_idx),
+                "Same routing key should route to same worker"
+            );
+            assert_eq!(branch, ExecutionBranch::OccupiedHit);
+        }
+    }
+
+    #[test]
+    fn test_random_mode_does_not_consider_load() {
+        let config = ManualConfig {
+            assignment_mode: ManualAssignmentMode::Random,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+
+        workers[0].worker_routing_key_load().increment("key-1");
+        workers[0].worker_routing_key_load().increment("key-2");
+        workers[0].worker_routing_key_load().increment("key-3");
+
+        let mut selected_worker_0 = false;
+        for i in 0..50 {
+            let headers = headers_with_routing_key(&format!("test-{}", i));
+            let info = SelectWorkerInfo {
+                headers: Some(&headers),
+                ..Default::default()
+            };
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            if result == Some(0) {
+                selected_worker_0 = true;
+                break;
+            }
+        }
+        assert!(
+            selected_worker_0,
+            "Random mode should sometimes select worker 0 despite higher load"
+        );
+    }
+
+    fn assert_no_routing_key_uses_random(
+        assignment_mode: ManualAssignmentMode,
+        setup_load: impl Fn(&[Arc<dyn Worker>]),
+    ) {
+        let config = ManualConfig {
+            assignment_mode,
+            ..Default::default()
+        };
+        let policy = ManualPolicy::with_config(config);
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000"]);
+        setup_load(&workers);
+
+        let mut selected_worker_0 = false;
+        for _ in 0..50 {
+            let info = SelectWorkerInfo::default();
+            let (result, branch) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(branch, ExecutionBranch::NoRoutingId);
+            if result == Some(0) {
+                selected_worker_0 = true;
+                break;
+            }
+        }
+        assert!(
+            selected_worker_0,
+            "Should randomly select worker 0 despite higher load"
+        );
+    }
+
+    #[test]
+    fn test_no_routing_key_uses_random_even_with_min_load_mode() {
+        assert_no_routing_key_uses_random(ManualAssignmentMode::MinLoad, |workers| {
+            workers[0].increment_load();
+            workers[0].increment_load();
+        });
+    }
+
+    #[test]
+    fn test_no_routing_key_uses_random_even_with_min_group_mode() {
+        assert_no_routing_key_uses_random(ManualAssignmentMode::MinGroup, |workers| {
+            workers[0].worker_routing_key_load().increment("k1");
+            workers[0].worker_routing_key_load().increment("k2");
+        });
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/mod.rs

@@ -0,0 +1,213 @@
+//! Load balancing policies for SGLang router
+//!
+//! This module provides a unified abstraction for routing policies that work
+//! across both regular and prefill-decode (PD) routing modes.
+
+use std::{fmt::Debug, sync::Arc};
+
+use async_trait::async_trait;
+use smg_mesh::OptionalMeshSyncManager;
+
+use crate::core::{HashRing, Worker};
+
+mod bucket;
+mod cache_aware;
+mod consistent_hashing;
+mod factory;
+mod manual;
+mod power_of_two;
+mod prefix_hash;
+mod random;
+mod registry;
+mod round_robin;
+pub mod tree;
+pub(crate) mod utils;
+pub use bucket::BucketPolicy;
+pub use cache_aware::CacheAwarePolicy;
+pub use consistent_hashing::ConsistentHashingPolicy;
+pub use factory::PolicyFactory;
+pub use manual::{ManualConfig, ManualPolicy};
+pub use power_of_two::PowerOfTwoPolicy;
+pub use prefix_hash::{PrefixHashConfig, PrefixHashPolicy};
+pub use random::RandomPolicy;
+pub use registry::PolicyRegistry;
+pub use round_robin::RoundRobinPolicy;
+pub use tree::PrefixMatchResult;
+
+/// Core trait for load balancing policies
+///
+/// This trait provides a unified interface for implementing routing algorithms
+/// that can work with both regular single-worker selection and PD dual-worker selection.
+#[async_trait]
+pub trait LoadBalancingPolicy: Send + Sync + Debug {
+    /// Select a single worker from the available workers
+    ///
+    /// This is used for regular routing mode where requests go to a single worker.
+    /// Now uses Arc<dyn Worker> for better performance and to avoid unnecessary cloning.
+    ///
+    /// # Arguments
+    /// * `workers` - Available workers to select from
+    /// * `info` - Additional information for routing decisions
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize>;
+
+    /// Update policy state after request completion
+    ///
+    /// This is called when a request completes (successfully or not) to allow
+    /// policies to update their internal state.
+    fn on_request_complete(&self, _worker_url: &str, _success: bool) {
+        // Default: no-op for stateless policies
+    }
+
+    /// Get policy name for metrics and debugging
+    fn name(&self) -> &'static str;
+
+    /// Check if this policy needs request text for routing decisions
+    fn needs_request_text(&self) -> bool {
+        false // Default: most policies don't need request text
+    }
+
+    /// Update worker load information
+    ///
+    /// This is called periodically with current load information for load-aware policies.
+    fn update_loads(&self, _loads: &std::collections::HashMap<String, isize>) {
+        // Default: no-op for policies that don't use load information
+    }
+
+    /// Set mesh sync manager
+    fn set_mesh_sync(&mut self, _mesh_sync: OptionalMeshSyncManager) {
+        // Default: no-op for policies that don't use mesh sync
+    }
+
+    /// Reset any internal state
+    ///
+    /// This is useful for policies that maintain state (e.g., round-robin counters).
+    fn reset(&self) {
+        // Default: no-op for stateless policies
+    }
+
+    /// Get as Any for downcasting
+    fn as_any(&self) -> &dyn std::any::Any;
+}
+
+/// Configuration for cache-aware policy
+#[derive(Debug, Clone)]
+pub struct CacheAwareConfig {
+    pub cache_threshold: f32,
+    pub balance_abs_threshold: usize,
+    pub balance_rel_threshold: f32,
+    pub eviction_interval_secs: u64,
+    pub max_tree_size: usize,
+}
+
+impl Default for CacheAwareConfig {
+    fn default() -> Self {
+        Self {
+            cache_threshold: 0.5,
+            balance_abs_threshold: 32,
+            balance_rel_threshold: 1.1,
+            eviction_interval_secs: 30,
+            max_tree_size: 10000,
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct BucketConfig {
+    pub balance_abs_threshold: usize,
+    pub balance_rel_threshold: f32,
+    pub bucket_adjust_interval_secs: usize,
+}
+
+impl Default for BucketConfig {
+    fn default() -> Self {
+        Self {
+            balance_abs_threshold: 32,
+            balance_rel_threshold: 1.0001,
+            bucket_adjust_interval_secs: 5,
+        }
+    }
+}
+
+/// Helper function to filter healthy workers and return their indices
+pub(crate) fn get_healthy_worker_indices(workers: &[Arc<dyn Worker>]) -> Vec<usize> {
+    workers
+        .iter()
+        .enumerate()
+        .filter(|(_, w)| w.is_healthy() && w.circuit_breaker().can_execute())
+        .map(|(idx, _)| idx)
+        .collect()
+}
+
+/// Helper function to normalize model_id to a key for policy lookups.
+///
+/// Returns UNKNOWN_MODEL_ID for empty model_ids to ensure consistent behavior
+/// across single-model and multi-model deployments.
+#[inline]
+pub(crate) fn normalize_model_key(model_id: &str) -> &str {
+    if model_id.is_empty() {
+        crate::core::UNKNOWN_MODEL_ID
+    } else {
+        model_id
+    }
+}
+
+/// Information passed to policy for worker selection
+#[derive(Debug, Clone, Default)]
+pub struct SelectWorkerInfo<'a> {
+    /// Request text for cache-aware routing
+    pub request_text: Option<&'a str>,
+    /// Tokenized request for prefix-hash routing
+    /// Used by PrefixHashPolicy for token-based prefix hashing
+    pub tokens: Option<&'a [u32]>,
+    /// HTTP headers for header-based routing policies
+    /// Policies can extract routing information from headers like:
+    /// - X-SMG-Target-Worker: Direct routing to a specific worker by index
+    /// - X-SMG-Routing-Key: Consistent hash routing for session affinity
+    pub headers: Option<&'a http::HeaderMap>,
+    /// Pre-computed hash ring for O(log n) consistent hashing
+    /// Built and cached by WorkerRegistry, passed through to avoid per-request rebuilds
+    pub hash_ring: Option<Arc<HashRing>>,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    #[tokio::test]
+    async fn test_get_healthy_worker_indices() {
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .api_key("test_api_key")
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .api_key("test_api_key2")
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w3:8000")
+                    .worker_type(WorkerType::Regular)
+                    .api_key("test_api_key")
+                    .build(),
+            ),
+        ];
+
+        // All healthy initially
+        let indices = get_healthy_worker_indices(&workers);
+        assert_eq!(indices, vec![0, 1, 2]);
+
+        // Mark one unhealthy
+        workers[1].set_healthy(false);
+        let indices = get_healthy_worker_indices(&workers);
+        assert_eq!(indices, vec![0, 2]);
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/power_of_two.rs

@@ -0,0 +1,389 @@
+//! Power-of-two choices load balancing policy
+
+use std::{
+    collections::HashMap,
+    sync::{Arc, RwLock},
+};
+
+use async_trait::async_trait;
+use rand::Rng;
+use tracing::debug;
+
+use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
+use crate::core::Worker;
+
+/// Power-of-two choices policy
+///
+/// Randomly selects two workers and routes to the one with lower load.
+/// This provides good load distribution with minimal coordination overhead.
+#[derive(Debug)]
+pub struct PowerOfTwoPolicy {
+    /// Cached load information from external monitoring
+    cached_loads: RwLock<HashMap<String, isize>>,
+}
+
+impl PowerOfTwoPolicy {
+    pub fn new() -> Self {
+        Self {
+            cached_loads: RwLock::new(HashMap::new()),
+        }
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for PowerOfTwoPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        _info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let healthy_indices = get_healthy_worker_indices(workers);
+
+        if healthy_indices.is_empty() {
+            return None;
+        }
+
+        if healthy_indices.len() == 1 {
+            return Some(healthy_indices[0]);
+        }
+
+        // Select two random workers - use offset to guarantee different selection in O(1)
+        let mut rng = rand::rng();
+        let idx1 = rng.random_range(0..healthy_indices.len());
+        // Pick idx2 from remaining indices: offset by 1 + random from (len-1) to guarantee different
+        let idx2 =
+            (idx1 + 1 + rng.random_range(0..healthy_indices.len() - 1)) % healthy_indices.len();
+
+        let worker_idx1 = healthy_indices[idx1];
+        let worker_idx2 = healthy_indices[idx2];
+        let worker1 = &workers[worker_idx1];
+        let worker2 = &workers[worker_idx2];
+
+        // Access cached loads safely
+        let loads_guard = self.cached_loads.read().ok();
+
+        // Try to get high-fidelity token loads for BOTH workers
+        let load1_tokens = loads_guard
+            .as_ref()
+            .and_then(|m| m.get(worker1.url()).copied());
+        let load2_tokens = loads_guard
+            .as_ref()
+            .and_then(|m| m.get(worker2.url()).copied());
+
+        // If either worker is missing token data (e.g. monitor failure),
+        // we must degrade BOTH to request counts to ensure fairness.
+        let (load1, load2) = match (load1_tokens, load2_tokens) {
+            (Some(t1), Some(t2)) => {
+                // Both have token data. Compare Tokens.
+                (t1, t2)
+            }
+            _ => {
+                // If One or both are missing token data.
+                // Fallback to local request counts for BOTH.
+                (worker1.load() as isize, worker2.load() as isize)
+            }
+        };
+
+        // Select worker with lower load
+        let selected_idx = if load1 <= load2 {
+            worker_idx1
+        } else {
+            worker_idx2
+        };
+
+        debug!(
+            "Power-of-two selection: {}={} vs {}={} -> selected {}",
+            worker1.url(),
+            load1,
+            worker2.url(),
+            load2,
+            workers[selected_idx].url()
+        );
+
+        // Increment processed counter
+        workers[selected_idx].increment_processed();
+
+        Some(selected_idx)
+    }
+
+    fn name(&self) -> &'static str {
+        "power_of_two"
+    }
+
+    fn update_loads(&self, loads: &HashMap<String, isize>) {
+        if let Ok(mut cached) = self.cached_loads.write() {
+            *cached = loads.clone();
+        }
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+impl Default for PowerOfTwoPolicy {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    #[tokio::test]
+    async fn test_power_of_two_selection() {
+        let policy = PowerOfTwoPolicy::new();
+        let worker1 = BasicWorkerBuilder::new("http://w1:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+        let worker2 = BasicWorkerBuilder::new("http://w2:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+        let worker3 = BasicWorkerBuilder::new("http://w3:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+
+        // Set different loads
+        for _ in 0..10 {
+            worker1.increment_load();
+        }
+        for _ in 0..5 {
+            worker2.increment_load();
+        }
+        // worker3 has load 0
+
+        let workers: Vec<Arc<dyn Worker>> =
+            vec![Arc::new(worker1), Arc::new(worker2), Arc::new(worker3)];
+
+        // Run multiple selections
+        let mut selected_counts = [0; 3];
+        let info = SelectWorkerInfo::default();
+        for _ in 0..100 {
+            if let Some(idx) = policy.select_worker(&workers, &info).await {
+                selected_counts[idx] += 1;
+            }
+        }
+
+        // Worker with lowest load (worker3) should be selected most often
+        assert!(selected_counts[2] > selected_counts[1]);
+        assert!(selected_counts[1] > selected_counts[0]);
+    }
+
+    #[tokio::test]
+    async fn test_power_of_two_with_cached_loads() {
+        let policy = PowerOfTwoPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        // Update cached loads
+        let mut loads = HashMap::new();
+        loads.insert("http://w1:8000".to_string(), 100);
+        loads.insert("http://w2:8000".to_string(), 10);
+        policy.update_loads(&loads);
+
+        // Should prefer worker2 with lower cached load
+        let mut w2_selected = 0;
+        let info = SelectWorkerInfo::default();
+        for _ in 0..50 {
+            if let Some(idx) = policy.select_worker(&workers, &info).await {
+                if idx == 1 {
+                    w2_selected += 1;
+                }
+            }
+        }
+
+        // Worker2 should be selected significantly more often
+        assert!(w2_selected > 35); // Should win most of the time
+    }
+
+    #[tokio::test]
+    async fn test_power_of_two_single_worker() {
+        let policy = PowerOfTwoPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
+            BasicWorkerBuilder::new("http://w1:8000")
+                .worker_type(WorkerType::Regular)
+                .build(),
+        )];
+
+        // With single worker, should always select it
+        assert_eq!(
+            policy
+                .select_worker(&workers, &SelectWorkerInfo::default())
+                .await,
+            Some(0)
+        );
+    }
+
+    #[tokio::test]
+    async fn test_reproduce_incompatible_metric_bug() {
+        use std::{collections::HashMap, sync::Arc};
+
+        use crate::core::{BasicWorkerBuilder, WorkerType};
+
+        // 1. Setup the policy
+        let policy = PowerOfTwoPolicy::new();
+
+        // 2. Create Worker A: Idle (0 reqs), but has high token usage in cache
+        let worker_a = BasicWorkerBuilder::new("http://worker_a:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+
+        // 3. Create Worker B: Busy (5 reqs), but missing from cache
+        let worker_b = BasicWorkerBuilder::new("http://worker_b:8000")
+            .worker_type(WorkerType::Regular)
+            .build();
+
+        // Manually increment load on Worker B to simulate active requests
+        for _ in 0..5 {
+            worker_b.increment_load();
+        }
+
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(worker_a), Arc::new(worker_b)];
+
+        // 4. Simulate LoadMonitor update:
+        // Only Worker A gets a token report. Worker B is missing (e.g. monitor failure).
+        let mut loads = HashMap::new();
+        loads.insert("http://worker_a:8000".to_string(), 50_000); // 50k tokens load
+        policy.update_loads(&loads);
+
+        // 5. Run selection
+        let selected_idx = policy
+            .select_worker(&workers, &SelectWorkerInfo::default())
+            .await
+            .expect("Should select a worker");
+
+        // 6. Verify the Fix
+        // Logic:
+        // - Worker A has token load (50k) but Worker B has NO token load.
+        // - Policy should fallback to request counts for BOTH.
+        // - A has 0 requests, B has 5 requests.
+        // - 0 <= 5, so A should be selected.
+
+        if selected_idx == 0 {
+            println!("Bug Fixed: System correctly fell back to request counts and selected idle Worker A.");
+        } else {
+            println!(
+                "Bug PERSISTS: Selected Worker B (Load: 5 reqs) over Worker A (Load: 50k tokens)"
+            );
+        }
+
+        // Assert that the CORRECT worker (A, index 0) is selected
+        assert_eq!(
+            selected_idx, 0,
+            "The policy failed to handle incompatible metrics. Should select idle Worker A."
+        );
+    }
+    #[tokio::test]
+    async fn test_power_of_two_edge_cases() {
+        use std::{collections::HashMap, sync::Arc};
+
+        use crate::core::{BasicWorkerBuilder, WorkerType};
+
+        let policy = PowerOfTwoPolicy::new();
+
+        // Helper to create a worker with specific request load
+        let create_worker = |url: &str, reqs: usize| {
+            let w = BasicWorkerBuilder::new(url)
+                .worker_type(WorkerType::Regular)
+                .build();
+            for _ in 0..reqs {
+                w.increment_load();
+            }
+            Arc::new(w)
+        };
+
+        //  Scenario 1: Happy Path (Both have Token Data)
+        // Worker A: 10 requests, but only 1,000 tokens (Light usage) -> Should be CHOSEN
+        // Worker B:  2 requests, but 100,000 tokens (Heavy usage) -> Should be AVOIDED
+        // This proves we use high-fidelity metrics when available, ignoring request counts.
+        let w_a = create_worker("http://a:8000", 10);
+        let w_b = create_worker("http://b:8000", 2);
+        let workers_1: Vec<Arc<dyn Worker>> = vec![w_a.clone(), w_b.clone()];
+
+        let mut loads_1 = HashMap::new();
+        loads_1.insert("http://a:8000".to_string(), 1_000);
+        loads_1.insert("http://b:8000".to_string(), 100_000);
+        policy.update_loads(&loads_1);
+
+        let idx_1 = policy
+            .select_worker(&workers_1, &SelectWorkerInfo::default())
+            .await
+            .unwrap();
+        assert_eq!(
+            idx_1, 0,
+            "Happy Path Failed: Should select Worker A (fewer tokens) despite higher request count"
+        );
+
+        // Scenario 2: Partial Failure (Worker A has tokens, Worker B is missing)
+        // Worker A: 10 requests, 1,000 tokens (Cached)
+        // Worker B:  2 requests, MISSING cache
+        // Logic: Fallback to requests -> Compare 10 (A) vs 2 (B) -> Select B
+        let w_c = create_worker("http://c:8000", 10);
+        let w_d = create_worker("http://d:8000", 2);
+        let workers_2: Vec<Arc<dyn Worker>> = vec![w_c.clone(), w_d.clone()];
+
+        let mut loads_2 = HashMap::new();
+        loads_2.insert("http://c:8000".to_string(), 1_000);
+        // http://d:8000 is MISSING
+        policy.update_loads(&loads_2);
+
+        let idx_2 = policy
+            .select_worker(&workers_2, &SelectWorkerInfo::default())
+            .await
+            .unwrap();
+        assert_eq!(idx_2, 1, "Partial Fail 1 Failed: Should fallback to requests and select Worker B (fewer requests)");
+
+        // Scenario 3: Partial Failure (Worker A is missing, Worker B has tokens)
+        // Worker A:  2 requests, MISSING cache
+        // Worker B: 10 requests, 1,000 tokens (Cached)
+        // Logic: Fallback to requests -> Compare 2 (A) vs 10 (B) -> Select A
+        let w_e = create_worker("http://e:8000", 2);
+        let w_f = create_worker("http://f:8000", 10);
+        let workers_3: Vec<Arc<dyn Worker>> = vec![w_e.clone(), w_f.clone()];
+
+        let mut loads_3 = HashMap::new();
+        // http://e:8000 is MISSING
+        loads_3.insert("http://f:8000".to_string(), 1_000);
+        policy.update_loads(&loads_3);
+
+        let idx_3 = policy
+            .select_worker(&workers_3, &SelectWorkerInfo::default())
+            .await
+            .unwrap();
+        assert_eq!(idx_3, 0, "Partial Fail 2 Failed: Should fallback to requests and select Worker A (fewer requests)");
+
+        // Scenario 4: Total Failure (Both missing)
+        // Worker A: 5 requests
+        // Worker B: 3 requests
+        // Logic: Requests vs Requests -> Select B
+        let w_g = create_worker("http://g:8000", 5);
+        let w_h = create_worker("http://h:8000", 3);
+        let workers_4: Vec<Arc<dyn Worker>> = vec![w_g.clone(), w_h.clone()];
+
+        let loads_4 = HashMap::new();
+        policy.update_loads(&loads_4);
+
+        let idx_4 = policy
+            .select_worker(&workers_4, &SelectWorkerInfo::default())
+            .await
+            .unwrap();
+        assert_eq!(
+            idx_4, 1,
+            "Total Fail Failed: Should select Worker B based on request count"
+        );
+
+        println!("All edge case tests passed successfully.");
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/prefix_hash.rs

@@ -0,0 +1,414 @@
+//! Prefix Hash routing policy for KV cache-aware load balancing
+//!
+//! A lightweight alternative to the full radix tree cache_aware policy.
+//! Routes requests based on a hash of their prefix tokens to maximize
+//! KV cache hits across workers.
+//!
+//! ## Algorithm
+//!
+//! 1. Extract first N tokens from the request (configurable prefix length)
+//! 2. Hash the token sequence using xxhash for fast, stable hashing
+//! 3. Use consistent hash ring to find the target worker
+//! 4. If worker is overloaded (load > avg * load_factor), find least loaded
+//! 5. Return least loaded worker that passes load check, or initial if all overloaded
+//!
+//! ## Complexity
+//!
+//! - Hash computation: O(prefix_length)
+//! - Ring lookup: O(log n) binary search
+//! - Load balance fallback: O(n) scan for least loaded
+//!
+//! ## Comparison with cache_aware
+//!
+//! | Aspect          | prefix_hash       | cache_aware (radix) |
+//! |-----------------|-------------------|---------------------|
+//! | Lookup          | O(log n)          | O(prefix_len)       |
+//! | Memory          | O(workers × vn)   | O(total_tokens)     |
+//! | Update          | O(1)              | O(prefix_len)       |
+//! | Precision       | Prefix grouping   | Exact matching      |
+//!
+//! prefix_hash trades optimal cache utilization for predictable O(log n) performance.
+
+use std::sync::Arc;
+
+use super::{LoadBalancingPolicy, SelectWorkerInfo};
+use crate::{core::Worker, observability::metrics::Metrics};
+
+/// Configuration for the PrefixHash load balancing policy
+#[derive(Debug, Clone)]
+pub struct PrefixHashConfig {
+    /// Number of prefix tokens to use for hashing.
+    /// Longer prefixes = more precise routing but less grouping.
+    /// Shorter prefixes = more requests grouped together.
+    /// Default: 256 tokens (~1 paragraph of text)
+    pub prefix_token_count: usize,
+
+    /// Load factor threshold for walking the ring.
+    /// If a worker's load > (total_load / num_workers) * load_factor,
+    /// walk clockwise to the next worker.
+    /// Default: 1.25 (125% of average load)
+    pub load_factor: f64,
+}
+
+impl Default for PrefixHashConfig {
+    fn default() -> Self {
+        Self {
+            prefix_token_count: 256,
+            load_factor: 1.25,
+        }
+    }
+}
+
+/// Execution branch for metrics
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum Branch {
+    NoHealthyWorkers,
+    NoTokens,
+    RingHit,
+    LoadBalanceWalk,
+    FallbackLeastLoad,
+}
+
+impl Branch {
+    #[inline]
+    const fn as_str(&self) -> &'static str {
+        match self {
+            Self::NoHealthyWorkers => "no_healthy_workers",
+            Self::NoTokens => "no_tokens",
+            Self::RingHit => "ring_hit",
+            Self::LoadBalanceWalk => "load_balance_walk",
+            Self::FallbackLeastLoad => "fallback_least_load",
+        }
+    }
+}
+
+/// Prefix Hash load balancing policy
+///
+/// Routes requests based on prefix token hash for KV cache locality.
+/// Uses consistent hashing with bounded load balancing.
+#[derive(Debug)]
+pub struct PrefixHashPolicy {
+    config: PrefixHashConfig,
+}
+
+impl PrefixHashPolicy {
+    /// Create a new PrefixHashPolicy with the given configuration
+    pub fn new(config: PrefixHashConfig) -> Self {
+        Self { config }
+    }
+
+    /// Create a new PrefixHashPolicy with default configuration
+    pub fn with_defaults() -> Self {
+        Self::new(PrefixHashConfig::default())
+    }
+
+    /// Compute hash of prefix tokens using xxhash
+    #[inline]
+    fn compute_prefix_hash(&self, tokens: &[u32]) -> u64 {
+        let prefix_len = tokens.len().min(self.config.prefix_token_count);
+        let prefix = &tokens[..prefix_len];
+
+        let bytes: &[u8] = bytemuck::cast_slice(prefix);
+        xxhash_rust::xxh3::xxh3_64(bytes)
+    }
+
+    /// Check if a worker's load is acceptable
+    #[inline]
+    fn load_ok(&self, worker_load: usize, total_load: usize, num_workers: usize) -> bool {
+        if total_load == 0 || num_workers == 0 {
+            return true;
+        }
+
+        // Average load per worker (with +1 to simulate incoming request)
+        let avg_load = (total_load + 1) as f64 / num_workers as f64;
+        let threshold = avg_load * self.config.load_factor;
+
+        (worker_load as f64) <= threshold
+    }
+
+    /// Find worker using consistent hash ring with load balancing
+    fn find_worker_with_load_balance(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo,
+        prefix_hash: u64,
+    ) -> (Option<usize>, Branch) {
+        // Build healthy worker URL to index map
+        let healthy_workers: Vec<(usize, &Arc<dyn Worker>)> = workers
+            .iter()
+            .enumerate()
+            .filter(|(_, w)| w.is_healthy())
+            .collect();
+
+        if healthy_workers.is_empty() {
+            return (None, Branch::NoHealthyWorkers);
+        }
+
+        // Calculate total load for load balancing
+        let total_load: usize = healthy_workers.iter().map(|(_, w)| w.load()).sum();
+        let num_workers = healthy_workers.len();
+
+        // Use pre-computed ring if available
+        if let Some(ref ring) = info.hash_ring {
+            // Convert prefix hash to a ring key string for lookup
+            let key = format!("{:016x}", prefix_hash);
+
+            // Build URL to (index, worker) map for healthy workers
+            let healthy_url_map: std::collections::HashMap<&str, (usize, &Arc<dyn Worker>)> =
+                healthy_workers
+                    .iter()
+                    .map(|(idx, w)| (w.url(), (*idx, *w)))
+                    .collect();
+
+            // Find initial worker from ring
+            if let Some(initial_url) =
+                ring.find_healthy_url(&key, |url| healthy_url_map.contains_key(url))
+            {
+                if let Some(&(idx, worker)) = healthy_url_map.get(initial_url) {
+                    let worker_load = worker.load();
+
+                    // Check if initial worker has acceptable load
+                    if self.load_ok(worker_load, total_load, num_workers) {
+                        return (Some(idx), Branch::RingHit);
+                    }
+
+                    // Initial worker overloaded, find least loaded healthy worker
+                    // This is a simpler approach than walking the ring
+                    let least_loaded = healthy_workers
+                        .iter()
+                        .filter(|(_, w)| self.load_ok(w.load(), total_load, num_workers))
+                        .min_by_key(|(_, w)| w.load());
+
+                    if let Some(&(idx, _)) = least_loaded {
+                        return (Some(idx), Branch::LoadBalanceWalk);
+                    }
+
+                    // All workers overloaded, use initial worker anyway
+                    return (Some(idx), Branch::LoadBalanceWalk);
+                }
+            }
+        }
+
+        // Fallback: no ring or ring lookup failed, use least loaded worker
+        let least_loaded = healthy_workers
+            .iter()
+            .min_by_key(|(_, w)| w.load())
+            .map(|(idx, _)| *idx);
+
+        (least_loaded, Branch::FallbackLeastLoad)
+    }
+
+    fn select_worker_impl(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo,
+    ) -> (Option<usize>, Branch) {
+        if workers.is_empty() {
+            return (None, Branch::NoHealthyWorkers);
+        }
+
+        // Get tokens from SelectWorkerInfo
+        let tokens = match info.tokens {
+            Some(t) if !t.is_empty() => t,
+            _ => return (None, Branch::NoTokens),
+        };
+
+        // Compute prefix hash
+        let prefix_hash = self.compute_prefix_hash(tokens);
+
+        // Find worker using ring with load balancing
+        self.find_worker_with_load_balance(workers, info, prefix_hash)
+    }
+}
+
+#[async_trait::async_trait]
+impl LoadBalancingPolicy for PrefixHashPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let (result, branch) = self.select_worker_impl(workers, info);
+        Metrics::record_worker_prefix_hash_policy_branch(branch.as_str());
+        result
+    }
+
+    fn name(&self) -> &'static str {
+        "prefix_hash"
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, HashRing, WorkerType};
+
+    fn create_workers(urls: &[&str]) -> Vec<Arc<dyn Worker>> {
+        urls.iter()
+            .map(|url| {
+                Arc::new(
+                    BasicWorkerBuilder::new(*url)
+                        .worker_type(WorkerType::Regular)
+                        .build(),
+                ) as Arc<dyn Worker>
+            })
+            .collect()
+    }
+
+    #[test]
+    fn test_prefix_hash_consistent_routing() {
+        let policy = PrefixHashPolicy::with_defaults();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        // Same tokens should always route to same worker
+        let tokens: Vec<u32> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+        let info = SelectWorkerInfo {
+            tokens: Some(&tokens),
+            hash_ring: Some(ring.clone()),
+            ..Default::default()
+        };
+
+        let (first_result, _) = policy.select_worker_impl(&workers, &info);
+        let first_idx = first_result.unwrap();
+
+        // Verify consistency
+        for _ in 0..10 {
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            assert_eq!(result, Some(first_idx));
+        }
+    }
+
+    #[test]
+    fn test_different_prefixes_distribute() {
+        let policy = PrefixHashPolicy::with_defaults();
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        let mut distribution = std::collections::HashMap::new();
+
+        // Different token sequences should distribute across workers
+        for i in 0..100 {
+            let tokens: Vec<u32> = vec![i, i + 1, i + 2, i + 3];
+            let info = SelectWorkerInfo {
+                tokens: Some(&tokens),
+                hash_ring: Some(ring.clone()),
+                ..Default::default()
+            };
+
+            let (result, _) = policy.select_worker_impl(&workers, &info);
+            *distribution.entry(result.unwrap()).or_insert(0) += 1;
+        }
+
+        assert!(
+            distribution.len() > 1,
+            "Should distribute across workers, got {:?}",
+            distribution
+        );
+    }
+
+    #[test]
+    fn test_shared_prefix_routes_same() {
+        let policy = PrefixHashPolicy::new(PrefixHashConfig {
+            prefix_token_count: 5, // Only look at first 5 tokens
+            ..Default::default()
+        });
+        let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        // Two sequences with same first 5 tokens should route to same worker
+        let tokens1: Vec<u32> = vec![1, 2, 3, 4, 5, 100, 200, 300];
+        let tokens2: Vec<u32> = vec![1, 2, 3, 4, 5, 999, 888, 777];
+
+        let info1 = SelectWorkerInfo {
+            tokens: Some(&tokens1),
+            hash_ring: Some(ring.clone()),
+            ..Default::default()
+        };
+        let info2 = SelectWorkerInfo {
+            tokens: Some(&tokens2),
+            hash_ring: Some(ring.clone()),
+            ..Default::default()
+        };
+
+        let (result1, _) = policy.select_worker_impl(&workers, &info1);
+        let (result2, _) = policy.select_worker_impl(&workers, &info2);
+
+        assert_eq!(result1, result2, "Same prefix should route to same worker");
+    }
+
+    #[test]
+    fn test_no_tokens_returns_none() {
+        let policy = PrefixHashPolicy::with_defaults();
+        let workers = create_workers(&["http://w1:8000"]);
+        let ring = Arc::new(HashRing::new(&workers));
+
+        // Empty tokens
+        let tokens: Vec<u32> = vec![];
+        let info = SelectWorkerInfo {
+            tokens: Some(&tokens),
+            hash_ring: Some(ring.clone()),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::NoTokens);
+
+        // No tokens field
+        let info_no_tokens = SelectWorkerInfo {
+            tokens: None,
+            hash_ring: Some(ring),
+            ..Default::default()
+        };
+
+        let (result2, branch2) = policy.select_worker_impl(&workers, &info_no_tokens);
+        assert_eq!(result2, None);
+        assert_eq!(branch2, Branch::NoTokens);
+    }
+
+    #[test]
+    fn test_no_healthy_workers() {
+        let policy = PrefixHashPolicy::with_defaults();
+        let workers = create_workers(&["http://w1:8000"]);
+        workers[0].set_healthy(false);
+
+        let ring = Arc::new(HashRing::new(&workers));
+        let tokens: Vec<u32> = vec![1, 2, 3];
+        let info = SelectWorkerInfo {
+            tokens: Some(&tokens),
+            hash_ring: Some(ring),
+            ..Default::default()
+        };
+
+        let (result, branch) = policy.select_worker_impl(&workers, &info);
+        assert_eq!(result, None);
+        assert_eq!(branch, Branch::NoHealthyWorkers);
+    }
+
+    #[test]
+    fn test_load_ok_calculation() {
+        let policy = PrefixHashPolicy::new(PrefixHashConfig {
+            load_factor: 1.25,
+            ..Default::default()
+        });
+
+        // Total load 100, 4 workers -> avg 25, threshold 31.25
+        assert!(policy.load_ok(30, 100, 4)); // 30 <= 31.25
+        assert!(!policy.load_ok(35, 100, 4)); // 35 > 31.25
+
+        // Edge cases
+        assert!(policy.load_ok(0, 0, 4)); // No load = OK
+        assert!(policy.load_ok(100, 0, 0)); // No workers = OK (shouldn't happen)
+    }
+
+    #[test]
+    fn test_policy_name() {
+        let policy = PrefixHashPolicy::with_defaults();
+        assert_eq!(policy.name(), "prefix_hash");
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/random.rs

@@ -0,0 +1,138 @@
+//! Random load balancing policy
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use rand::Rng;
+
+use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
+use crate::core::Worker;
+
+/// Random selection policy
+///
+/// Selects workers randomly with uniform distribution among healthy workers.
+#[derive(Debug, Default)]
+pub struct RandomPolicy;
+
+impl RandomPolicy {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for RandomPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        _info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let healthy_indices = get_healthy_worker_indices(workers);
+
+        if healthy_indices.is_empty() {
+            return None;
+        }
+
+        let mut rng = rand::rng();
+        let random_idx = rng.random_range(0..healthy_indices.len());
+
+        Some(healthy_indices[random_idx])
+    }
+
+    fn name(&self) -> &'static str {
+        "random"
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    #[tokio::test]
+    async fn test_random_selection() {
+        let policy = RandomPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w3:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        let mut counts = HashMap::new();
+        for _ in 0..100 {
+            if let Some(idx) = policy
+                .select_worker(&workers, &SelectWorkerInfo::default())
+                .await
+            {
+                *counts.entry(idx).or_insert(0) += 1;
+            }
+        }
+
+        assert_eq!(counts.len(), 3);
+        assert!(counts.values().all(|&count| count > 0));
+    }
+
+    #[tokio::test]
+    async fn test_random_with_unhealthy_workers() {
+        let policy = RandomPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        workers[0].set_healthy(false);
+
+        for _ in 0..10 {
+            assert_eq!(
+                policy
+                    .select_worker(&workers, &SelectWorkerInfo::default())
+                    .await,
+                Some(1)
+            );
+        }
+    }
+
+    #[tokio::test]
+    async fn test_random_no_healthy_workers() {
+        let policy = RandomPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![Arc::new(
+            BasicWorkerBuilder::new("http://w1:8000")
+                .worker_type(WorkerType::Regular)
+                .build(),
+        )];
+
+        workers[0].set_healthy(false);
+        assert_eq!(
+            policy
+                .select_worker(&workers, &SelectWorkerInfo::default())
+                .await,
+            None
+        );
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/registry.rs

@@ -0,0 +1,512 @@
+use std::sync::{Arc, OnceLock, RwLock};
+
+use dashmap::DashMap;
+use serde_json;
+use smg_mesh::OptionalMeshSyncManager;
+use tracing::{debug, info, warn};
+
+/// Policy Registry for managing model-to-policy mappings
+///
+/// This registry manages the dynamic assignment of load balancing policies to models.
+/// When the first worker of a new model is added, it determines the policy for that model.
+/// All subsequent workers of the same model use the established policy.
+/// When the last worker of a model is removed, the policy mapping is cleaned up.
+use super::{BucketPolicy, CacheAwarePolicy, LoadBalancingPolicy, PolicyFactory};
+use crate::{config::types::PolicyConfig, core::Worker};
+
+/// Registry for managing model-to-policy mappings
+#[derive(Clone)]
+pub struct PolicyRegistry {
+    /// Model ID -> Policy instance mapping (lock-free reads via DashMap)
+    model_policies: Arc<DashMap<String, Arc<dyn LoadBalancingPolicy>>>,
+
+    /// Model ID -> Worker count for cleanup tracking (lock-free reads via DashMap)
+    model_worker_counts: Arc<DashMap<String, usize>>,
+
+    /// Default policy instance (cached, immutable after creation)
+    default_policy: Arc<dyn LoadBalancingPolicy>,
+
+    /// Prefill policy for PD mode (set once at startup, lock-free reads via OnceLock)
+    prefill_policy: Arc<OnceLock<Arc<dyn LoadBalancingPolicy>>>,
+
+    /// Decode policy for PD mode (set once at startup, lock-free reads via OnceLock)
+    decode_policy: Arc<OnceLock<Arc<dyn LoadBalancingPolicy>>>,
+
+    /// Optional mesh sync manager for state synchronization
+    /// When None, the registry works independently without mesh synchronization
+    /// Uses RwLock for thread-safe access when setting mesh_sync after initialization
+    mesh_sync: Arc<RwLock<OptionalMeshSyncManager>>,
+}
+
+impl PolicyRegistry {
+    /// Create a new PolicyRegistry with a default policy
+    pub fn new(default_policy_config: PolicyConfig) -> Self {
+        let default_policy = Self::create_policy_from_config(&default_policy_config);
+
+        Self {
+            model_policies: Arc::new(DashMap::new()),
+            model_worker_counts: Arc::new(DashMap::new()),
+            default_policy,
+            prefill_policy: Arc::new(OnceLock::new()),
+            decode_policy: Arc::new(OnceLock::new()),
+            mesh_sync: Arc::new(RwLock::new(None)),
+        }
+    }
+
+    /// Set mesh sync manager (thread-safe, can be called after initialization)
+    pub fn set_mesh_sync(&self, mesh_sync: OptionalMeshSyncManager) {
+        *self.mesh_sync.write().unwrap() = mesh_sync;
+    }
+
+    /// Called when a worker is added
+    /// Returns the policy that should be used for this worker's model
+    pub fn on_worker_added(
+        &self,
+        model_id: &str,
+        policy_hint: Option<&str>,
+    ) -> Arc<dyn LoadBalancingPolicy> {
+        // Increment worker count using DashMap entry API
+        let count = self
+            .model_worker_counts
+            .entry(model_id.to_string())
+            .and_modify(|c| *c += 1)
+            .or_insert(1);
+        debug!("Worker added for model {}, count: {}", model_id, *count);
+        drop(count); // Release the entry lock
+
+        // Check if model already has a policy (lock-free read via DashMap)
+        if let Some(existing_policy) = self.model_policies.get(model_id) {
+            debug!(
+                "Model {} already has policy: {}",
+                model_id,
+                existing_policy.name()
+            );
+            return Arc::clone(&existing_policy);
+        }
+
+        // New model - determine policy
+        let policy = self.determine_policy_for_model(model_id, policy_hint);
+
+        info!(
+            "Assigning policy {} to new model {}",
+            policy.name(),
+            model_id
+        );
+
+        // Store policy for this model (DashMap handles concurrent inserts)
+        self.model_policies
+            .insert(model_id.to_string(), Arc::clone(&policy));
+
+        // Sync to mesh if enabled (no-op if mesh is not enabled)
+        if let Some(ref mesh_sync) = *self.mesh_sync.read().unwrap() {
+            // Serialize policy config (simplified - just store policy name for now)
+            let config = serde_json::to_vec(&policy.name()).unwrap_or_default();
+            mesh_sync.sync_policy_state(model_id.to_string(), policy.name().to_string(), config);
+        }
+
+        policy
+    }
+
+    /// Called when a worker is removed
+    pub fn on_worker_removed(&self, model_id: &str) {
+        // Decrement worker count and check if cleanup needed
+        let should_cleanup = if let Some(mut count_ref) = self.model_worker_counts.get_mut(model_id)
+        {
+            *count_ref = count_ref.saturating_sub(1);
+            debug!(
+                "Worker removed for model {}, count: {}",
+                model_id, *count_ref
+            );
+            if *count_ref == 0 {
+                drop(count_ref); // Release before remove
+                self.model_worker_counts.remove(model_id);
+                true
+            } else {
+                false
+            }
+        } else {
+            warn!(
+                "Attempted to remove worker for model {} with no registered workers",
+                model_id
+            );
+            false
+        };
+
+        // Clean up policy if this was the last worker
+        if should_cleanup {
+            if let Some((_, policy)) = self.model_policies.remove(model_id) {
+                info!(
+                    "Removed policy {} for model {} (last worker removed)",
+                    policy.name(),
+                    model_id
+                );
+            }
+
+            // Sync removal to mesh if enabled (no-op if mesh is not enabled)
+            if let Some(ref mesh_sync) = *self.mesh_sync.read().unwrap() {
+                mesh_sync.remove_policy_state(model_id);
+            }
+        }
+    }
+
+    /// Get the policy for a model (lock-free via DashMap)
+    pub fn get_policy(&self, model_id: &str) -> Option<Arc<dyn LoadBalancingPolicy>> {
+        self.model_policies.get(model_id).map(|r| Arc::clone(&r))
+    }
+
+    /// Get the default policy
+    pub fn get_default_policy(&self) -> Arc<dyn LoadBalancingPolicy> {
+        Arc::clone(&self.default_policy)
+    }
+
+    /// Get policy for a model, or default if not found
+    pub fn get_policy_or_default(&self, model_id: &str) -> Arc<dyn LoadBalancingPolicy> {
+        self.get_policy(model_id)
+            .unwrap_or_else(|| self.get_default_policy())
+    }
+
+    /// Determine policy for a new model
+    fn determine_policy_for_model(
+        &self,
+        model_id: &str,
+        policy_hint: Option<&str>,
+    ) -> Arc<dyn LoadBalancingPolicy> {
+        // 1. Check policy hint from worker
+        if let Some(policy_type) = policy_hint {
+            debug!("Using policy hint '{}' for model {}", policy_type, model_id);
+            return self.create_policy_from_type(policy_type);
+        }
+
+        // 2. Use default policy
+        debug!("Using default policy for model {}", model_id);
+        Arc::clone(&self.default_policy)
+    }
+
+    /// Create a policy from a type string (delegates to PolicyFactory)
+    fn create_policy_from_type(&self, policy_type: &str) -> Arc<dyn LoadBalancingPolicy> {
+        if policy_type == "cache_aware" {
+            let mut cache_aware = CacheAwarePolicy::new();
+            let mesh_sync = &*self.mesh_sync.read().unwrap();
+            cache_aware.set_mesh_sync(mesh_sync.clone());
+            Arc::new(cache_aware)
+        } else {
+            PolicyFactory::create_by_name(policy_type).unwrap_or_else(|| {
+                warn!("Unknown policy type '{}', using default", policy_type);
+                Arc::clone(&self.default_policy)
+            })
+        }
+    }
+
+    /// Create a policy from a PolicyConfig (delegates to PolicyFactory)
+    fn create_policy_from_config(config: &PolicyConfig) -> Arc<dyn LoadBalancingPolicy> {
+        PolicyFactory::create_from_config(config)
+    }
+
+    /// Get current model->policy mappings (for debugging/monitoring)
+    pub fn get_all_mappings(&self) -> std::collections::HashMap<String, String> {
+        self.model_policies
+            .iter()
+            .map(|entry| (entry.key().clone(), entry.value().name().to_string()))
+            .collect()
+    }
+
+    /// Get worker counts per model
+    pub fn get_worker_counts(&self) -> std::collections::HashMap<String, usize> {
+        self.model_worker_counts
+            .iter()
+            .map(|entry| (entry.key().clone(), *entry.value()))
+            .collect()
+    }
+
+    /// Clear all policies (useful for testing)
+    pub fn clear(&self) {
+        self.model_policies.clear();
+        self.model_worker_counts.clear();
+    }
+
+    /// Set the prefill policy for PD mode (lock-free, set once at startup)
+    pub fn set_prefill_policy(&self, policy: Arc<dyn LoadBalancingPolicy>) {
+        // OnceLock::set returns Err if already set, which we ignore since
+        // the policy should only be set once at startup
+        let _ = self.prefill_policy.set(policy);
+    }
+
+    /// Set the decode policy for PD mode (lock-free, set once at startup)
+    pub fn set_decode_policy(&self, policy: Arc<dyn LoadBalancingPolicy>) {
+        // OnceLock::set returns Err if already set, which we ignore since
+        // the policy should only be set once at startup
+        let _ = self.decode_policy.set(policy);
+    }
+
+    /// Get the prefill policy for PD mode, or default if not set (lock-free)
+    pub fn get_prefill_policy(&self) -> Arc<dyn LoadBalancingPolicy> {
+        self.prefill_policy
+            .get()
+            .map(Arc::clone)
+            .unwrap_or_else(|| self.get_default_policy())
+    }
+
+    /// Get the decode policy for PD mode, or default if not set (lock-free)
+    pub fn get_decode_policy(&self) -> Arc<dyn LoadBalancingPolicy> {
+        self.decode_policy
+            .get()
+            .map(Arc::clone)
+            .unwrap_or_else(|| self.get_default_policy())
+    }
+
+    /// Get all PowerOfTwo policies that need load updates (lock-free)
+    pub fn get_all_power_of_two_policies(&self) -> Vec<Arc<dyn LoadBalancingPolicy>> {
+        let mut power_of_two_policies = Vec::new();
+
+        if self.default_policy.name() == "power_of_two" {
+            power_of_two_policies.push(Arc::clone(&self.default_policy));
+        }
+
+        // Get prefill and decode policies (lock-free via OnceLock::get)
+        let prefill_policy_opt = self.prefill_policy.get();
+        let decode_policy_opt = self.decode_policy.get();
+
+        if let Some(policy) = prefill_policy_opt {
+            if policy.name() == "power_of_two" && !Arc::ptr_eq(policy, &self.default_policy) {
+                power_of_two_policies.push(Arc::clone(policy));
+            }
+        }
+
+        if let Some(policy) = decode_policy_opt {
+            if policy.name() == "power_of_two"
+                && !Arc::ptr_eq(policy, &self.default_policy)
+                && !prefill_policy_opt.is_some_and(|p| Arc::ptr_eq(p, policy))
+            {
+                power_of_two_policies.push(Arc::clone(policy));
+            }
+        }
+
+        for entry in self.model_policies.iter() {
+            let policy = entry.value();
+            if policy.name() == "power_of_two" {
+                let already_added = power_of_two_policies.iter().any(|p| Arc::ptr_eq(p, policy));
+                if !already_added {
+                    power_of_two_policies.push(Arc::clone(policy));
+                }
+            }
+        }
+
+        power_of_two_policies
+    }
+
+    /// Initialize cache-aware policy with workers if applicable
+    /// This should be called after workers are registered for a model
+    pub fn init_cache_aware_policy(&self, model_id: &str, workers: &[Arc<dyn Worker>]) {
+        // Get the policy for this model
+        if let Some(policy) = self.get_policy(model_id) {
+            if policy.name() == "cache_aware" {
+                if let Some(cache_aware) = policy.as_any().downcast_ref::<CacheAwarePolicy>() {
+                    debug!(
+                        "Initializing cache-aware policy with {} workers for model {}",
+                        workers.len(),
+                        model_id
+                    );
+                    cache_aware.init_workers(workers);
+                }
+            }
+        }
+    }
+
+    /// Remove a worker from cache-aware policy if applicable
+    /// This should be called when a worker is being removed
+    pub fn remove_worker_from_cache_aware(&self, model_id: &str, worker_url: &str) {
+        // Get the policy for this model
+        if let Some(policy) = self.get_policy(model_id) {
+            if policy.name() == "cache_aware" {
+                if let Some(cache_aware) = policy.as_any().downcast_ref::<CacheAwarePolicy>() {
+                    cache_aware.remove_worker_by_url(worker_url);
+                    debug!(
+                        "Removed worker {} from cache-aware policy for model {}",
+                        worker_url, model_id
+                    );
+                }
+            }
+        }
+    }
+
+    /// Initialize cache-aware policies for PD mode (prefill and decode) - lock-free
+    pub fn init_pd_cache_aware_policies(
+        &self,
+        prefill_workers: &[Arc<dyn Worker>],
+        decode_workers: &[Arc<dyn Worker>],
+    ) {
+        // Initialize prefill policy if it's cache-aware (lock-free via OnceLock::get)
+        if let Some(prefill_policy) = self.prefill_policy.get() {
+            if prefill_policy.name() == "cache_aware" {
+                if let Some(cache_aware) =
+                    prefill_policy.as_any().downcast_ref::<CacheAwarePolicy>()
+                {
+                    if !prefill_workers.is_empty() {
+                        debug!(
+                            "Initializing prefill cache-aware policy with {} workers",
+                            prefill_workers.len()
+                        );
+                        cache_aware.init_workers(prefill_workers);
+                    }
+                }
+            }
+        }
+
+        // Initialize decode policy if it's cache-aware (lock-free via OnceLock::get)
+        if let Some(decode_policy) = self.decode_policy.get() {
+            if decode_policy.name() == "cache_aware" {
+                if let Some(cache_aware) = decode_policy.as_any().downcast_ref::<CacheAwarePolicy>()
+                {
+                    if !decode_workers.is_empty() {
+                        debug!(
+                            "Initializing decode cache-aware policy with {} workers",
+                            decode_workers.len()
+                        );
+                        cache_aware.init_workers(decode_workers);
+                    }
+                }
+            }
+        }
+    }
+
+    /// Initialize bucket policies for PD mode - lock-free
+    pub fn init_pd_bucket_policies(&self, prefill_workers: &[Arc<dyn Worker>]) {
+        // Initialize prefill policy if it's bucket (lock-free via OnceLock::get)
+        if let Some(prefill_policy) = self.prefill_policy.get() {
+            if prefill_policy.name() == "bucket" {
+                if let Some(bucket) = prefill_policy.as_any().downcast_ref::<BucketPolicy>() {
+                    if !prefill_workers.is_empty() {
+                        debug!(
+                            "Initializing prefill bucket policy with {} workers",
+                            prefill_workers.len()
+                        );
+                        bucket.init_prefill_worker_urls(prefill_workers);
+                    }
+                }
+            }
+        }
+    }
+
+    /// Apply remote tree operation to cache-aware policy for a model
+    /// This is called when receiving tree state updates from mesh
+    pub fn apply_remote_tree_operation(
+        &self,
+        model_id: &str,
+        operation: &smg_mesh::tree_ops::TreeOperation,
+    ) {
+        // Try to find the policy for this model
+        if let Some(policy) = self.get_policy(model_id) {
+            if policy.name() == "cache_aware" {
+                if let Some(cache_aware) = policy.as_any().downcast_ref::<CacheAwarePolicy>() {
+                    cache_aware.apply_remote_tree_operation(model_id, operation);
+                }
+            }
+        }
+
+        // Also check default policy if it's cache-aware
+        if self.default_policy.name() == "cache_aware" {
+            if let Some(cache_aware) = self
+                .default_policy
+                .as_any()
+                .downcast_ref::<CacheAwarePolicy>()
+            {
+                cache_aware.apply_remote_tree_operation(model_id, operation);
+            }
+        }
+
+        // Check prefill and decode policies for PD mode
+        if let Some(prefill_policy) = self.prefill_policy.get() {
+            if prefill_policy.name() == "cache_aware" {
+                if let Some(cache_aware) =
+                    prefill_policy.as_any().downcast_ref::<CacheAwarePolicy>()
+                {
+                    cache_aware.apply_remote_tree_operation(model_id, operation);
+                }
+            }
+        }
+
+        if let Some(decode_policy) = self.decode_policy.get() {
+            if decode_policy.name() == "cache_aware" {
+                if let Some(cache_aware) = decode_policy.as_any().downcast_ref::<CacheAwarePolicy>()
+                {
+                    cache_aware.apply_remote_tree_operation(model_id, operation);
+                }
+            }
+        }
+    }
+}
+
+impl std::fmt::Debug for PolicyRegistry {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("PolicyRegistry")
+            .field("model_policies", &self.model_policies)
+            .field("model_worker_counts", &self.model_worker_counts)
+            .field("default_policy", &self.default_policy.name())
+            .finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_policy_registry_basic() {
+        let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
+
+        // First worker of a model sets the policy
+        let policy1 = registry.on_worker_added("llama-3", Some("cache_aware"));
+        assert_eq!(policy1.name(), "cache_aware");
+
+        // Second worker of same model uses existing policy
+        let policy2 = registry.on_worker_added("llama-3", Some("round_robin"));
+        assert_eq!(policy2.name(), "cache_aware"); // Ignores hint, uses existing
+
+        // Different model can have different policy
+        let policy3 = registry.on_worker_added("gpt-4", Some("random"));
+        assert_eq!(policy3.name(), "random");
+
+        // Check mappings
+        let mappings = registry.get_all_mappings();
+        assert_eq!(mappings.get("llama-3").unwrap(), "cache_aware");
+        assert_eq!(mappings.get("gpt-4").unwrap(), "random");
+
+        // Check worker counts
+        let counts = registry.get_worker_counts();
+        assert_eq!(*counts.get("llama-3").unwrap(), 2);
+        assert_eq!(*counts.get("gpt-4").unwrap(), 1);
+    }
+
+    #[tokio::test]
+    async fn test_policy_registry_cleanup() {
+        let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
+
+        // Add workers
+        registry.on_worker_added("llama-3", Some("cache_aware"));
+        registry.on_worker_added("llama-3", None);
+        assert_eq!(registry.get_worker_counts().get("llama-3"), Some(&2));
+
+        // Remove one worker - policy should remain
+        registry.on_worker_removed("llama-3");
+        assert!(registry.get_policy("llama-3").is_some());
+        assert_eq!(registry.get_worker_counts().get("llama-3"), Some(&1));
+
+        // Remove last worker - policy should be cleaned up
+        registry.on_worker_removed("llama-3");
+        assert!(registry.get_policy("llama-3").is_none());
+        assert_eq!(registry.get_worker_counts().get("llama-3"), None);
+    }
+
+    #[tokio::test]
+    async fn test_default_policy() {
+        let registry = PolicyRegistry::new(PolicyConfig::RoundRobin);
+
+        // No hint, no template - uses default
+        let policy = registry.on_worker_added("unknown-model", None);
+        assert_eq!(policy.name(), "round_robin");
+
+        // Get default directly
+        let default = registry.get_default_policy();
+        assert_eq!(default.name(), "round_robin");
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/round_robin.rs

@@ -0,0 +1,149 @@
+//! Round-robin load balancing policy
+
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Arc,
+};
+
+use async_trait::async_trait;
+
+use super::{get_healthy_worker_indices, LoadBalancingPolicy, SelectWorkerInfo};
+use crate::core::Worker;
+
+/// Round-robin selection policy
+///
+/// Selects workers in sequential order, cycling through all healthy workers.
+#[derive(Debug, Default)]
+pub struct RoundRobinPolicy {
+    counter: AtomicUsize,
+}
+
+impl RoundRobinPolicy {
+    pub fn new() -> Self {
+        Self {
+            counter: AtomicUsize::new(0),
+        }
+    }
+}
+
+#[async_trait]
+impl LoadBalancingPolicy for RoundRobinPolicy {
+    async fn select_worker(
+        &self,
+        workers: &[Arc<dyn Worker>],
+        _info: &SelectWorkerInfo<'_>,
+    ) -> Option<usize> {
+        let healthy_indices = get_healthy_worker_indices(workers);
+
+        if healthy_indices.is_empty() {
+            return None;
+        }
+
+        // Get and increment counter atomically
+        let count = self.counter.fetch_add(1, Ordering::Relaxed);
+        let selected_idx = count % healthy_indices.len();
+
+        Some(healthy_indices[selected_idx])
+    }
+
+    fn name(&self) -> &'static str {
+        "round_robin"
+    }
+
+    fn reset(&self) {
+        self.counter.store(0, Ordering::Relaxed);
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::{BasicWorkerBuilder, WorkerType};
+
+    #[tokio::test]
+    async fn test_round_robin_selection() {
+        let policy = RoundRobinPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w3:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        let info = SelectWorkerInfo::default();
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(1));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(2));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(1));
+    }
+
+    #[tokio::test]
+    async fn test_round_robin_with_unhealthy_workers() {
+        let policy = RoundRobinPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w3:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        workers[1].set_healthy(false);
+
+        let info = SelectWorkerInfo::default();
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(2));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(2));
+    }
+
+    #[tokio::test]
+    async fn test_round_robin_reset() {
+        let policy = RoundRobinPolicy::new();
+        let workers: Vec<Arc<dyn Worker>> = vec![
+            Arc::new(
+                BasicWorkerBuilder::new("http://w1:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+            Arc::new(
+                BasicWorkerBuilder::new("http://w2:8000")
+                    .worker_type(WorkerType::Regular)
+                    .build(),
+            ),
+        ];
+
+        let info = SelectWorkerInfo::default();
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(1));
+
+        policy.reset();
+        assert_eq!(policy.select_worker(&workers, &info).await, Some(0));
+    }
+}

third_party/sglang/sgl-model-gateway/src/policies/utils.rs

@@ -0,0 +1,110 @@
+use std::{
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+    thread::{self, JoinHandle},
+    time::Duration,
+};
+
+use tracing::debug;
+
+#[derive(Debug)]
+pub struct PeriodicTask {
+    debug_name: &'static str,
+    shutdown_flag: Arc<AtomicBool>,
+    handle: Option<JoinHandle<()>>,
+}
+
+impl PeriodicTask {
+    /// Spawn a background thread that periodically executes a task.
+    pub fn spawn<F>(interval_secs: u64, debug_name: &'static str, task: F) -> Self
+    where
+        F: Fn() + Send + 'static,
+    {
+        let shutdown_flag = Arc::new(AtomicBool::new(false));
+        let shutdown_clone = Arc::clone(&shutdown_flag);
+
+        let handle = thread::spawn(move || {
+            let check_interval_ms = 100u64;
+            let total_sleep_ms = interval_secs * 1000;
+
+            loop {
+                // Sleep in small increments, checking shutdown flag periodically
+                let mut slept_ms = 0u64;
+                while slept_ms < total_sleep_ms {
+                    if shutdown_clone.load(Ordering::Relaxed) {
+                        debug!("{} thread received shutdown signal", debug_name);
+                        return;
+                    }
+                    thread::sleep(Duration::from_millis(check_interval_ms));
+                    slept_ms += check_interval_ms;
+                }
+
+                // Check shutdown before starting task
+                if shutdown_clone.load(Ordering::Relaxed) {
+                    debug!("{} thread received shutdown signal", debug_name);
+                    return;
+                }
+
+                task();
+            }
+        });
+
+        Self {
+            debug_name,
+            shutdown_flag,
+            handle: Some(handle),
+        }
+    }
+}
+
+impl Drop for PeriodicTask {
+    fn drop(&mut self) {
+        self.shutdown_flag.store(true, Ordering::Relaxed);
+
+        if let Some(handle) = self.handle.take() {
+            match handle.join() {
+                Ok(()) => debug!("{} thread shut down cleanly", self.debug_name),
+                Err(_) => debug!("{} thread panicked during shutdown", self.debug_name),
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::{sync::atomic::AtomicUsize, time::Instant};
+
+    use super::*;
+
+    #[test]
+    fn test_periodic_task_executes() {
+        let counter = Arc::new(AtomicUsize::new(0));
+        let counter_clone = Arc::clone(&counter);
+
+        let _task = PeriodicTask::spawn(1, "test", move || {
+            counter_clone.fetch_add(1, Ordering::SeqCst);
+        });
+
+        // Wait for at least one execution
+        thread::sleep(Duration::from_millis(1200));
+        assert!(counter.load(Ordering::SeqCst) >= 1);
+
+        // Task will be stopped on drop
+    }
+
+    #[test]
+    fn test_periodic_task_responds_to_shutdown() {
+        let task = PeriodicTask::spawn(60, "test", || {
+            // Long interval task
+        });
+
+        let start = Instant::now();
+        drop(task);
+        let elapsed = start.elapsed();
+
+        // Should shutdown within ~200ms (2 check intervals), not 60 seconds
+        assert!(elapsed < Duration::from_millis(500));
+    }
+}