feat(policy): add event-driven cache-aware routing to CacheAwarePolicy

bindings/python/src/lib.rs

@@ -375,6 +375,7 @@ struct Router {
     balance_rel_threshold: f32,
     eviction_interval_secs: u64,
     max_tree_size: usize,
+    block_size: usize,
     max_idle_secs: u64,
     assignment_mode: String,
     max_payload_size: usize,
@@ -477,6 +478,7 @@ impl Router {
                     balance_rel_threshold: self.balance_rel_threshold,
                     eviction_interval_secs: self.eviction_interval_secs,
                     max_tree_size: self.max_tree_size,
+                    block_size: self.block_size,
                 },
                 PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo {
                     load_check_interval_secs: 5,
@@ -699,6 +701,7 @@ impl Router {
         balance_rel_threshold = 1.5,
         eviction_interval_secs = 120,
         max_tree_size = 2usize.pow(26),
+        block_size = 16,
         max_idle_secs = 14400,
         assignment_mode = String::from("random"),
         max_payload_size = 512 * 1024 * 1024,
@@ -792,6 +795,7 @@ impl Router {
         balance_rel_threshold: f32,
         eviction_interval_secs: u64,
         max_tree_size: usize,
+        block_size: usize,
         max_idle_secs: u64,
         assignment_mode: String,
         max_payload_size: usize,
@@ -894,6 +898,7 @@ impl Router {
             balance_rel_threshold,
             eviction_interval_secs,
             max_tree_size,
+            block_size,
             max_idle_secs,
             assignment_mode,
             max_payload_size,

bindings/python/src/smg/router_args.py

@@ -36,6 +36,7 @@ class RouterArgs:
     balance_rel_threshold: float = 1.5
     eviction_interval_secs: int = 60
     max_tree_size: int = 2**26
+    block_size: int = 16
     max_idle_secs: int = 4 * 3600
     assignment_mode: str = "random"  # Mode for manual policy new routing key assignment
     max_payload_size: int = 512 * 1024 * 1024  # 512MB default for large batches
@@ -332,6 +333,12 @@ def add_cli_args(
             default=RouterArgs.max_tree_size,
             help="Maximum size of the approximation tree for cache-aware routing",
         )
+        routing_group.add_argument(
+            f"--{prefix}block-size",
+            type=int,
+            default=RouterArgs.block_size,
+            help="KV cache block size for event-driven cache-aware routing (default: 16)",
+        )
         routing_group.add_argument(
             f"--{prefix}max-idle-secs",
             type=int,

model_gateway/src/app_context.rs

@@ -16,7 +16,10 @@ use tracing::debug;
 
 use crate::{
     config::RouterConfig,
-    core::{steps::WorkflowEngines, JobQueue, LoadMonitor, WorkerRegistry, WorkerService},
+    core::{
+        steps::WorkflowEngines, JobQueue, KvEventMonitor, LoadMonitor, WorkerRegistry,
+        WorkerService,
+    },
     middleware::TokenBucket,
     observability::inflight_tracker::InFlightRequestTracker,
     policies::PolicyRegistry,
@@ -59,6 +62,7 @@ pub struct AppContext {
     pub wasm_manager: Option<Arc<WasmModuleManager>>,
     pub worker_service: Arc<WorkerService>,
     pub inflight_tracker: Arc<InFlightRequestTracker>,
+    pub kv_event_monitor: Option<Arc<KvEventMonitor>>,
 }
 
 impl std::fmt::Debug for AppContext {
@@ -87,6 +91,7 @@ pub struct AppContextBuilder {
     workflow_engines: Option<Arc<OnceLock<WorkflowEngines>>>,
     mcp_orchestrator: Option<Arc<OnceLock<Arc<McpOrchestrator>>>>,
     wasm_manager: Option<Arc<WasmModuleManager>>,
+    kv_event_monitor: Option<Arc<KvEventMonitor>>,
 }
 
 impl AppContext {
@@ -132,6 +137,7 @@ impl AppContextBuilder {
             workflow_engines: None,
             mcp_orchestrator: None,
             wasm_manager: None,
+            kv_event_monitor: None,
         }
     }
 
@@ -232,6 +238,11 @@ impl AppContextBuilder {
         self
     }
 
+    pub fn kv_event_monitor(mut self, kv_event_monitor: Option<Arc<KvEventMonitor>>) -> Self {
+        self.kv_event_monitor = kv_event_monitor;
+        self
+    }
+
     pub fn build(self) -> Result<AppContext, AppContextBuildError> {
         let router_config = self
             .router_config
@@ -289,6 +300,7 @@ impl AppContextBuilder {
             wasm_manager: self.wasm_manager,
             worker_service,
             inflight_tracker: InFlightRequestTracker::new(),
+            kv_event_monitor: self.kv_event_monitor,
         })
     }
 
@@ -314,6 +326,7 @@ impl AppContextBuilder {
             .with_mcp_orchestrator(&router_config)
             .await?
             .with_wasm_manager(&router_config)
+            .with_kv_event_monitor(&router_config)
             .router_config(router_config))
     }
 
@@ -544,6 +557,33 @@ impl AppContextBuilder {
         Ok(self)
     }
 
+    /// Create KV event monitor for event-driven cache-aware routing.
+    ///
+    /// The monitor is created when the default policy is cache_aware, regardless
+    /// of connection mode. The monitor itself is cheap (empty DashMaps) and stays
+    /// dormant until workers are added. The UpdatePoliciesStep gates subscriptions
+    /// on `cache_aware && gRPC`, so HTTP workers are never subscribed.
+    fn with_kv_event_monitor(mut self, config: &RouterConfig) -> Self {
+        use crate::config::types::PolicyConfig;
+
+        let is_cache_aware = matches!(config.policy, PolicyConfig::CacheAware { .. });
+
+        if is_cache_aware {
+            let monitor = Arc::new(KvEventMonitor::new(None));
+            debug!("Created KV event monitor for event-driven cache-aware routing");
+
+            // Inject monitor into PolicyRegistry — propagates to default_policy
+            // and any other existing cache-aware policies.
+            if let Some(ref registry) = self.policy_registry {
+                registry.set_kv_event_monitor(Some(Arc::clone(&monitor)));
+            }
+
+            self.kv_event_monitor = Some(monitor);
+        }
+
+        self
+    }
+
     /// Create wasm manager if enabled in config
     fn with_wasm_manager(mut self, config: &RouterConfig) -> Self {
         self.wasm_manager = if config.enable_wasm {

model_gateway/src/config/builder.rs

@@ -122,6 +122,7 @@ impl RouterConfigBuilder {
             balance_rel_threshold,
             eviction_interval_secs,
             max_tree_size,
+            block_size: 16,
         };
         self
     }

model_gateway/src/config/types.rs

@@ -250,6 +250,8 @@ pub enum PolicyConfig {
         balance_rel_threshold: f32,
         eviction_interval_secs: u64,
         max_tree_size: usize,
+        #[serde(default = "default_block_size")]
+        block_size: usize,
     },
 
     #[serde(rename = "power_of_two")]
@@ -307,6 +309,10 @@ pub enum PolicyConfig {
     },
 }
 
+fn default_block_size() -> usize {
+    16
+}
+
 fn default_prefix_token_count() -> usize {
     256
 }
@@ -773,6 +779,7 @@ mod tests {
             balance_rel_threshold: 1.5,
             eviction_interval_secs: 300,
             max_tree_size: 1000,
+            block_size: 16,
         };
         assert_eq!(cache_aware.name(), "cache_aware");
 
@@ -794,6 +801,7 @@ mod tests {
             balance_rel_threshold: 1.5,
             eviction_interval_secs: 300,
             max_tree_size: 1000,
+            block_size: 16,
         };
         let json = serde_json::to_string(&cache_aware).unwrap();
         assert!(json.contains("\"type\":\"cache_aware\""));
@@ -816,6 +824,7 @@ mod tests {
             balance_rel_threshold: 2.0,
             eviction_interval_secs: 600,
             max_tree_size: 5000,
+            block_size: 16,
         };
 
         match cache_aware {
@@ -825,6 +834,7 @@ mod tests {
                 balance_rel_threshold,
                 eviction_interval_secs,
                 max_tree_size,
+                ..
             } => {
                 assert!((cache_threshold - 0.75).abs() < 0.0001);
                 assert_eq!(balance_abs_threshold, 20);
@@ -1218,6 +1228,7 @@ mod tests {
                 balance_rel_threshold: 1.1,
                 eviction_interval_secs: 60,
                 max_tree_size: 1000,
+                block_size: 16,
             }),
             decode_policy: Some(PolicyConfig::PowerOfTwo {
                 load_check_interval_secs: 60,
@@ -1248,6 +1259,7 @@ mod tests {
                 balance_rel_threshold: 1.1,
                 eviction_interval_secs: 60,
                 max_tree_size: 1000,
+                block_size: 16,
             }),
             decode_policy: None,
         };
@@ -1304,6 +1316,7 @@ mod tests {
             balance_rel_threshold: 1.5,
             eviction_interval_secs: 300,
             max_tree_size: 2000,
+            block_size: 16,
         };
 
         match pd.get_prefill_policy(&main_policy) {

model_gateway/src/config/validation.rs

@@ -173,7 +173,16 @@ impl ConfigValidator {
                 balance_rel_threshold,
                 eviction_interval_secs,
                 max_tree_size,
+                block_size,
             } => {
+                if *block_size == 0 {
+                    return Err(ConfigError::InvalidValue {
+                        field: "block_size".to_string(),
+                        value: block_size.to_string(),
+                        reason: "Must be > 0".to_string(),
+                    });
+                }
+
                 if !(0.0..=1.0).contains(cache_threshold) {
                     return Err(ConfigError::InvalidValue {
                         field: "cache_threshold".to_string(),
@@ -742,6 +751,7 @@ mod tests {
                 balance_rel_threshold: 1.1,
                 eviction_interval_secs: 60,
                 max_tree_size: 1000,
+                block_size: 16,
             },
         );
 
@@ -761,6 +771,7 @@ mod tests {
                 balance_rel_threshold: 1.1,
                 eviction_interval_secs: 60,
                 max_tree_size: 1000,
+                block_size: 16,
             },
         );
 
@@ -815,6 +826,7 @@ mod tests {
                 balance_rel_threshold: 1.1,
                 eviction_interval_secs: 60,
                 max_tree_size: 1000,
+                block_size: 16,
             },
         );
 
@@ -859,6 +871,7 @@ mod tests {
                     balance_rel_threshold: 1.1,
                     eviction_interval_secs: 60,
                     max_tree_size: 1000,
+                    block_size: 16,
                 }),
                 decode_policy: Some(PolicyConfig::PowerOfTwo {
                     load_check_interval_secs: 60,