refactor(worker): extract health loop into WorkerManager

bindings/golang/src/policy.rs

@@ -6,6 +6,7 @@
 //! Rust gateway.
 
 use std::{
+    any::Any,
     ffi::{CStr, CString},
     os::raw::c_char,
     ptr,
@@ -31,7 +32,7 @@ use smg::{
         utils::{generate_tool_constraints, process_chat_messages},
     },
     worker::{
-        circuit_breaker::CircuitBreaker,
+        circuit_breaker::{CircuitBreaker, CircuitState},
         resilience::ResolvedResilience,
         worker::{RuntimeType, WorkerMetadata, WorkerRoutingKeyLoad},
         ConnectionMode, Worker, WorkerResult, WorkerType,
@@ -106,6 +107,10 @@ impl std::fmt::Debug for GrpcWorker {
 
 #[async_trait]
 impl Worker for GrpcWorker {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
     fn url(&self) -> &str {
         &self.endpoint
     }
@@ -130,11 +135,33 @@ impl Worker for GrpcWorker {
         self.status.store(status as u8, Ordering::Relaxed);
     }
 
+    fn revision(&self) -> u64 {
+        0
+    }
+
     async fn check_health_async(&self) -> WorkerResult<()> {
         // FFI workers don't do their own health checks
         Ok(())
     }
 
+    // FFI workers don't run the state machine — these counters are unused.
+    // The Go SDK manages worker health via direct set_healthy() calls.
+    fn consecutive_failures_increment(&self) -> usize {
+        0
+    }
+    fn consecutive_failures_reset(&self) {}
+    fn consecutive_successes_increment(&self) -> usize {
+        0
+    }
+    fn consecutive_successes_reset(&self) {}
+    fn total_pending_probes(&self) -> usize {
+        0
+    }
+    fn total_pending_probes_increment(&self) -> usize {
+        0
+    }
+    fn total_pending_probes_reset(&self) {}
+
     fn load(&self) -> usize {
         self.load.load(Ordering::Relaxed)
     }
@@ -147,8 +174,16 @@ impl Worker for GrpcWorker {
         self.load.fetch_sub(1, Ordering::Relaxed);
     }
 
-    fn worker_routing_key_load(&self) -> &WorkerRoutingKeyLoad {
-        &self.routing_key_load
+    fn routing_key_load(&self) -> usize {
+        self.routing_key_load.value()
+    }
+
+    fn increment_routing_key_load(&self, routing_key: &str) {
+        self.routing_key_load.increment(routing_key);
+    }
+
+    fn decrement_routing_key_load(&self, routing_key: &str) {
+        self.routing_key_load.decrement(routing_key);
     }
 
     fn processed_requests(&self) -> usize {
@@ -163,8 +198,16 @@ impl Worker for GrpcWorker {
         &self.metadata
     }
 
-    fn circuit_breaker(&self) -> &CircuitBreaker {
-        &self.circuit_breaker
+    fn circuit_breaker_state(&self) -> CircuitState {
+        self.circuit_breaker.state()
+    }
+
+    fn circuit_breaker_can_execute(&self) -> bool {
+        self.circuit_breaker.can_execute()
+    }
+
+    fn record_circuit_breaker_outcome(&self, success: bool) {
+        self.circuit_breaker.record_outcome(success);
     }
 
     fn resilience(&self) -> &ResolvedResilience {

model_gateway/src/observability/metrics_ws/collectors.rs

@@ -206,7 +206,7 @@ fn collect_workers(context: &AppContext) -> Value {
                 "is_healthy": w.is_healthy(),
                 "load": w.load(),
                 "processed_requests": w.processed_requests(),
-                "circuit_breaker": w.circuit_breaker().state().to_string(),
+                "circuit_breaker": w.circuit_breaker_state().to_string(),
             })
         })
         .collect();

model_gateway/src/policies/manual.rs

@@ -292,9 +292,7 @@ fn min_load_select(workers: &[Arc<dyn Worker>], healthy_indices: &[usize]) -> us
 }
 
 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()
-    })
+    select_min_by(healthy_indices, |idx| workers[idx].routing_key_load())
 }
 
 #[cfg(test)]
@@ -793,9 +791,7 @@ mod tests {
             assert_eq!(branch, ExecutionBranch::Vacant);
 
             let selected_idx = result.unwrap();
-            workers[selected_idx]
-                .worker_routing_key_load()
-                .increment(&routing_key);
+            workers[selected_idx].increment_routing_key_load(&routing_key);
         }
 
         let distribution: HashMap<_, usize> = policy
@@ -822,13 +818,13 @@ mod tests {
         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");
+        workers[0].increment_routing_key_load("existing-1");
+        workers[0].increment_routing_key_load("existing-2");
+        workers[1].increment_routing_key_load("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);
+        assert_eq!(workers[0].routing_key_load(), 2);
+        assert_eq!(workers[1].routing_key_load(), 1);
+        assert_eq!(workers[2].routing_key_load(), 0);
 
         let headers = headers_with_routing_key("new-key");
         let info = SelectWorkerInfo {
@@ -878,9 +874,9 @@ mod tests {
         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");
+        workers[0].increment_routing_key_load("key-0");
+        workers[1].increment_routing_key_load("key-1");
+        workers[1].increment_routing_key_load("key-2");
 
         let headers = headers_with_routing_key("new-key");
         let info = SelectWorkerInfo {
@@ -916,9 +912,9 @@ mod tests {
         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");
+        workers[0].increment_routing_key_load("key-1");
+        workers[0].increment_routing_key_load("key-2");
+        workers[0].increment_routing_key_load("key-3");
 
         let mut selected_worker_0 = false;
         for i in 0..50 {

model_gateway/src/policies/mod.rs

@@ -144,7 +144,7 @@ pub(crate) fn get_healthy_worker_indices(workers: &[Arc<dyn Worker>]) -> Vec<usi
     workers
         .iter()
         .enumerate()
-        .filter(|(_, w)| w.is_healthy() && w.circuit_breaker().can_execute())
+        .filter(|(_, w)| w.is_healthy() && w.circuit_breaker_can_execute())
         .map(|(idx, _)| idx)
         .collect()
 }

model_gateway/src/server.rs

@@ -66,7 +66,10 @@ use crate::{
     },
     service_discovery::{start_service_discovery, ServiceDiscoveryConfig},
     wasm::route::{add_wasm_module, list_wasm_modules, remove_wasm_module},
-    worker::{manager::WorkerManager, worker::WorkerType},
+    worker::{
+        manager::{WorkerManager, WorkerManagerConfig},
+        worker::WorkerType,
+    },
     workflow::{
         job_queue::{JobQueue, JobQueueConfig},
         Job, TokenizerConfigRequest, WorkflowEngines,
@@ -1075,23 +1078,26 @@ pub async fn startup(config: ServerConfig) -> Result<(), Box<dyn std::error::Err
     let router_manager = RouterManager::from_config(&config, &app_context).await?;
     let router: Arc<dyn RouterTrait> = router_manager.clone();
 
-    // Health checker handle must outlive the server to keep the background task alive.
-    // HealthChecker aborts its task on Drop, so binding it here keeps it alive until
-    // the server shuts down.
-    let _health_checker = if config.router_config.health_check.disable_health_check {
-        info!("Global health checks disabled via CLI/config; skipping health checker");
+    // WorkerManager owns the background health check loop. Its handle must
+    // outlive the server to keep the task alive — bind it here so its Drop
+    // (which aborts the task) runs at server shutdown.
+    let _worker_manager = if config.router_config.health_check.disable_health_check {
+        info!("Global health checks disabled via CLI/config; skipping WorkerManager");
         None
     } else {
-        let hc = app_context.worker_registry.start_health_checker(
-            config.router_config.health_check.check_interval_secs,
-            config.router_config.health_check.remove_unhealthy_workers,
+        let manager = WorkerManager::start(
+            app_context.worker_registry.clone(),
+            WorkerManagerConfig {
+                default_check_interval_secs: config.router_config.health_check.check_interval_secs,
+                remove_unhealthy: config.router_config.health_check.remove_unhealthy_workers,
+            },
             app_context.worker_job_queue.get().cloned(),
         );
         debug!(
-            "Started health checker for workers with {}s interval",
+            "Started WorkerManager health check loop with {}s default interval",
             config.router_config.health_check.check_interval_secs
         );
-        Some(hc)
+        Some(manager)
     };
 
     // LoadMonitor groups are started dynamically when workers are registered.

model_gateway/src/service_discovery.rs

@@ -640,6 +640,7 @@ async fn handle_pod_event(
             );
             let job = Job::RemoveWorker {
                 url: old_url.clone(),
+                expected_revision: None,
             };
             if let Some(job_queue) = app_context.worker_job_queue.get() {
                 if let Err(e) = job_queue.submit(job).await {
@@ -789,6 +790,7 @@ async fn handle_pod_deletion(
 
         let job = Job::RemoveWorker {
             url: worker_url.clone(),
+            expected_revision: None,
         };
 
         if let Some(job_queue) = app_context.worker_job_queue.get() {
@@ -930,6 +932,7 @@ async fn reconcile_pods(
         );
         let job = Job::RemoveWorker {
             url: worker_url.clone(),
+            expected_revision: None,
         };
         if let Some(job_queue) = app_context.worker_job_queue.get() {
             match job_queue.submit(job).await {

model_gateway/src/worker/builder.rs

@@ -10,7 +10,7 @@ use super::{
     circuit_breaker::{CircuitBreaker, CircuitBreakerConfig},
     resilience::ResolvedResilience,
     worker::{
-        BasicWorker, ConnectionMode, RuntimeType, WorkerMetadata, WorkerRoutingKeyLoad, WorkerType,
+        BasicWorker, ConnectionMode, RuntimeType, WorkerMetadata, WorkerRuntime, WorkerType,
         DEFAULT_WORKER_HTTP_TIMEOUT_SECS,
     },
 };
@@ -227,10 +227,7 @@ impl BasicWorkerBuilder {
 
     /// Build the BasicWorker instance
     pub fn build(mut self) -> BasicWorker {
-        use std::sync::{
-            atomic::{AtomicU8, AtomicUsize},
-            Arc,
-        };
+        use std::sync::Arc;
 
         use tokio::sync::OnceCell;
 
@@ -287,17 +284,11 @@ impl BasicWorkerBuilder {
         let resilience = self.resilience.unwrap_or_default();
 
         BasicWorker {
-            load_counter: Arc::new(AtomicUsize::new(0)),
-            worker_routing_key_load: Arc::new(WorkerRoutingKeyLoad::new(&metadata.spec.url)),
-            processed_counter: Arc::new(AtomicUsize::new(0)),
-            status: Arc::new(AtomicU8::new(initial_status as u8)),
-            consecutive_failures: Arc::new(AtomicUsize::new(0)),
-            consecutive_successes: Arc::new(AtomicUsize::new(0)),
-            total_pending_probes: Arc::new(AtomicUsize::new(0)),
-            circuit_breaker: CircuitBreaker::with_config_and_label(
+            runtime: ArcSwap::from_pointee(WorkerRuntime::new(&metadata.spec.url, initial_status)),
+            circuit_breaker: ArcSwap::from_pointee(CircuitBreaker::with_config_and_label(
                 self.circuit_breaker_config,
                 metadata.spec.url.clone(),
-            ),
+            )),
             metadata,
             grpc_client,
             models_override: Arc::new(ArcSwap::from_pointee(WorkerModels::Wildcard)),

model_gateway/src/worker/circuit_breaker.rs

@@ -8,7 +8,7 @@ use tracing::info;
 use crate::observability::metrics::Metrics;
 
 /// Circuit breaker configuration
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct CircuitBreakerConfig {
     /// Number of consecutive failures to open the circuit
     pub failure_threshold: u32,
@@ -143,6 +143,11 @@ impl CircuitBreaker {
         &self.metric_label
     }
 
+    /// Get the configuration used to create this circuit breaker.
+    pub fn config(&self) -> &CircuitBreakerConfig {
+        &self.config
+    }
+
     /// Check if a request can be executed (lock-free hot path)
     #[inline]
     pub fn can_execute(&self) -> bool {