fix(zmq): correct transport-identity leaks in metrics, admin, and PD paths

crates/protocols/src/worker.rs

@@ -1136,6 +1136,10 @@ pub struct FlushCacheResult {
     pub http_workers: usize,
     #[serde(default)]
     pub grpc_workers: usize,
+    /// Workers skipped because their transport has no cache-flush RPC
+    /// (direct-ZMQ engines). Keeps `total = http + grpc + zmq` exact.
+    #[serde(default)]
+    pub zmq_workers: usize,
     pub message: String,
 }
 
@@ -1325,6 +1329,7 @@ impl IntoResponse for FlushCacheResult {
             "workers_flushed": self.successful.len(),
             "total_http_workers": self.http_workers,
             "total_grpc_workers": self.grpc_workers,
+            "total_zmq_workers_skipped": self.zmq_workers,
             "total_workers": self.total_workers
         });
 

model_gateway/src/routers/grpc/common/stages/encode.rs

@@ -320,6 +320,15 @@ fn prepare_items(
             prefill: BackendClient::Grpc(GrpcClient::TokenSpeed(_)),
             ..
         } => prepare_tokenspeed_items(intermediate, workers),
+        // TokenSpeed supports EPD encode, but only over gRPC — name the
+        // transport, not the engine, so the error points at the real limit.
+        ClientSelection::Disaggregated {
+            prefill: prefill @ BackendClient::Zmq(_),
+            ..
+        } if prefill.runtime_type() == RuntimeType::TokenSpeed => Err(anyhow!(
+            "EPD encode requires a gRPC TokenSpeed prefill worker; the direct-ZMQ backend has \
+             no encode dispatch"
+        )),
         ClientSelection::Disaggregated { prefill, .. } => Err(anyhow!(
             "EPD encode is not implemented for {} backend",
             backend_name(prefill)

model_gateway/src/routers/grpc/common/stages/helpers.rs

@@ -452,7 +452,13 @@ fn inject_sglang_bootstrap_metadata(
         bootstrap_room: room_id,
     };
 
-    let sglang_request = request.as_sglang_mut();
+    // Guarded by the caller's runtime check, but match defensively: a non-SGLang
+    // proto here (e.g. a ZMQ backend reporting an unexpected runtime) must not
+    // take down the request task via the panicking accessor.
+    let ProtoGenerateRequest::Sglang(sglang_request) = request else {
+        warn!("PD bootstrap metadata requested for a non-SGLang request; skipping injection");
+        return;
+    };
     sglang_request.disaggregated_params = Some(disagg_params);
 
     debug!(
@@ -768,6 +774,39 @@ mod stop_resolution_tests {
         );
     }
 
+    #[test]
+    fn pd_bootstrap_injection_skips_non_sglang_requests() {
+        use super::{RuntimeType, Worker, WorkerSelection};
+        use crate::worker::{BasicWorkerBuilder, WorkerType};
+
+        // An SGLang-runtime worker selection paired with a non-SGLang proto
+        // (e.g. a misreporting backend) must skip injection, not panic.
+        let worker: Arc<dyn Worker> = Arc::new(
+            BasicWorkerBuilder::new("grpc://prefill:30000")
+                .worker_type(WorkerType::Prefill)
+                .build(),
+        );
+        let selection = WorkerSelection::Disaggregated {
+            encode_assignments: None,
+            prefill: worker.clone(),
+            decode: worker,
+            runtime_type: RuntimeType::Sglang,
+        };
+
+        let mut req = vllm_request(vec!["."], vec![7]);
+        let before = match &req {
+            ProtoGenerateRequest::Vllm(inner) => (**inner).clone(),
+            _ => panic!("vllm_request builds a Vllm variant"),
+        };
+        super::maybe_inject_pd_metadata(&mut req, &selection);
+        match &req {
+            ProtoGenerateRequest::Vllm(inner) => {
+                assert_eq!(**inner, before, "request must be untouched");
+            }
+            _ => panic!("variant must be unchanged"),
+        }
+    }
+
     #[test]
     fn tokenspeed_zmq_multi_token_relies_on_router_decoder() {
         // "Hello world" => [1, 2]: not a flat stop id, so it must not forward.

model_gateway/src/routers/grpc/common/stages/request_execution.rs

@@ -25,11 +25,29 @@ use crate::{
             utils::tonic_ext::{TonicResultExt, TonicStatusExt},
         },
     },
-    worker::{RuntimeType, DEFAULT_BOOTSTRAP_PORT, MOONCAKE_CONNECTOR, NIXL_CONNECTOR},
+    worker::{
+        ConnectionModeExt, RuntimeType, DEFAULT_BOOTSTRAP_PORT, MOONCAKE_CONNECTOR, NIXL_CONNECTOR,
+    },
 };
 
 type StreamResult = Result<ProtoStream, tonic::Status>;
 
+/// Metric connection labels for the PD legs (a leg can be gRPC or ZMQ).
+fn pd_leg_labels(workers: &WorkerSelection) -> (&'static str, &'static str) {
+    match workers {
+        WorkerSelection::Disaggregated {
+            prefill, decode, ..
+        } => (
+            prefill.connection_mode().as_metric_label(),
+            decode.connection_mode().as_metric_label(),
+        ),
+        WorkerSelection::Single { worker } => {
+            let label = worker.connection_mode().as_metric_label();
+            (label, label)
+        }
+    }
+}
+
 /// KV-transfer params tagged onto the NIXL prefill leg so the engine pins its
 /// KV blocks and returns the handoff params for the decode worker.
 const NIXL_PREFILL_KV_PARAMS: &str = r#"{"do_remote_decode":true,"do_remote_prefill":false}"#;
@@ -489,11 +507,13 @@ impl RequestExecutionStage {
             decode_result.cb_status_code(),
         );
 
+        let (prefill_label, decode_label) = pd_leg_labels(workers);
+
         // Handle prefill result
         let prefill_stream = prefill_result.map_err(|e| {
             Metrics::record_worker_error(
                 metrics_labels::WORKER_PREFILL,
-                metrics_labels::CONNECTION_GRPC,
+                prefill_label,
                 metrics_labels::ERROR_BACKEND,
             );
             error!(function = "execute_parallel_pd", error = %e, "Prefill worker failed to start");
@@ -507,7 +527,7 @@ impl RequestExecutionStage {
         let decode_stream = decode_result.map_err(|e| {
             Metrics::record_worker_error(
                 metrics_labels::WORKER_DECODE,
-                metrics_labels::CONNECTION_GRPC,
+                decode_label,
                 metrics_labels::ERROR_BACKEND,
             );
             error!(function = "execute_parallel_pd", error = %e, "Decode worker failed to start");
@@ -648,6 +668,7 @@ impl RequestExecutionStage {
         );
 
         // Send to prefill, wait for completion
+        let (prefill_label, decode_label) = pd_leg_labels(workers);
         let prefill_start = Instant::now();
         let mut prefill_stream = prefill_client
             .generate(prefill_request)
@@ -656,7 +677,7 @@ impl RequestExecutionStage {
                 workers.record_outcome_prefill(e.http_status().as_u16());
                 Metrics::record_worker_error(
                     metrics_labels::WORKER_PREFILL,
-                    metrics_labels::CONNECTION_GRPC,
+                    prefill_label,
                     metrics_labels::ERROR_BACKEND,
                 );
                 error!(function = "execute_sequential_pd", error = %e, "Prefill worker failed to start");
@@ -678,7 +699,7 @@ impl RequestExecutionStage {
                     workers.record_outcome_prefill(e.http_status().as_u16());
                     Metrics::record_worker_error(
                         metrics_labels::WORKER_PREFILL,
-                        metrics_labels::CONNECTION_GRPC,
+                        prefill_label,
                         metrics_labels::ERROR_BACKEND,
                     );
                     error!(function = "execute_sequential_pd", error = %e, "Prefill stream error");
@@ -769,7 +790,7 @@ impl RequestExecutionStage {
             workers.record_outcome_decode(e.http_status().as_u16());
             Metrics::record_worker_error(
                 metrics_labels::WORKER_DECODE,
-                metrics_labels::CONNECTION_GRPC,
+                decode_label,
                 metrics_labels::ERROR_BACKEND,
             );
             error!(function = "execute_sequential_pd", error = %e, "Decode worker failed to start");
@@ -803,9 +824,42 @@ impl RequestExecutionStage {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+
     use smg_grpc_client::vllm_proto as vllm;
 
     use super::*;
+    use crate::worker::{BasicWorkerBuilder, ConnectionMode, Worker, WorkerType};
+
+    #[test]
+    fn pd_leg_labels_reflect_each_legs_transport() {
+        let prefill: Arc<dyn Worker> = Arc::new(
+            BasicWorkerBuilder::new("ipc:///tmp/smg-test-prefill")
+                .worker_type(WorkerType::Prefill)
+                .connection_mode(ConnectionMode::Zmq)
+                .build(),
+        );
+        let decode: Arc<dyn Worker> = Arc::new(
+            BasicWorkerBuilder::new("grpc://decode:30000")
+                .worker_type(WorkerType::Decode)
+                .connection_mode(ConnectionMode::Grpc)
+                .build(),
+        );
+        let selection = WorkerSelection::Disaggregated {
+            encode_assignments: None,
+            prefill,
+            decode,
+            runtime_type: RuntimeType::TokenSpeed,
+        };
+        assert_eq!(
+            pd_leg_labels(&selection),
+            (
+                metrics_labels::CONNECTION_ZMQ,
+                metrics_labels::CONNECTION_GRPC
+            ),
+            "each PD leg must carry its own transport label"
+        );
+    }
 
     #[test]
     fn kv_connector_mode_mooncake_uses_bootstrap_metadata() {

model_gateway/src/routers/grpc/common/stages/worker_selection.rs

@@ -21,8 +21,8 @@ use crate::{
         },
     },
     worker::{
-        ConnectionModeExt, HashRing, RuntimeType, Worker, WorkerRegistry, WorkerType,
-        UNKNOWN_MODEL_ID,
+        ConnectionMode, ConnectionModeExt, HashRing, RuntimeType, Worker, WorkerRegistry,
+        WorkerType, UNKNOWN_MODEL_ID,
     },
 };
 
@@ -399,13 +399,17 @@ impl WorkerSelectionStage {
         // Record worker selection metrics for both prefill and decode
         Metrics::record_worker_selection(
             metrics_labels::WORKER_PREFILL,
-            metrics_labels::CONNECTION_GRPC,
+            available_prefill[prefill_idx]
+                .connection_mode()
+                .as_metric_label(),
             model,
             prefill_policy.name(),
         );
         Metrics::record_worker_selection(
             metrics_labels::WORKER_DECODE,
-            metrics_labels::CONNECTION_GRPC,
+            available_decode[decode_idx]
+                .connection_mode()
+                .as_metric_label(),
             model,
             decode_policy.name(),
         );
@@ -452,7 +456,14 @@ impl WorkerSelectionStage {
             .fold((Vec::new(), Vec::new(), Vec::new()), |mut acc, w| {
                 if w.connection_mode().uses_grpc_pipeline() && w.is_available() {
                     match w.metadata().spec.worker_type {
-                        WorkerType::Encode => acc.0.push(w),
+                        // Encode dispatch is a gRPC encoder RPC sent to the
+                        // worker's URL; a direct-ZMQ worker has no encode
+                        // path, so only gRPC workers qualify for this pool.
+                        WorkerType::Encode => {
+                            if *w.connection_mode() == ConnectionMode::Grpc {
+                                acc.0.push(w);
+                            }
+                        }
                         WorkerType::Prefill => acc.1.push(w),
                         WorkerType::Decode => acc.2.push(w),
                         WorkerType::Regular => {}
@@ -577,13 +588,17 @@ impl WorkerSelectionStage {
         // recorded in assign_encode_workers.
         Metrics::record_worker_selection(
             metrics_labels::WORKER_PREFILL,
-            metrics_labels::CONNECTION_GRPC,
+            available_prefill[prefill_idx]
+                .connection_mode()
+                .as_metric_label(),
             model_id,
             prefill_policy.name(),
         );
         Metrics::record_worker_selection(
             metrics_labels::WORKER_DECODE,
-            metrics_labels::CONNECTION_GRPC,
+            available_decode[decode_idx]
+                .connection_mode()
+                .as_metric_label(),
             model_id,
             decode_policy.name(),
         );

model_gateway/src/routers/grpc/proto_wrapper.rs

@@ -1971,6 +1971,12 @@ impl ProtoStream {
                 .next()
                 .await
                 .map(|result| result.map(|r| ProtoGenerateResponse::TokenSpeed(Box::new(r)))),
+            // Every ZMQ engine (including TokenSpeed) emits vllm-shaped
+            // responses: the adapter translates wire output into
+            // `vllm::GenerateResponse`, so variant checks like
+            // `is_tokenspeed()` on a response are unreliable for ZMQ-backed
+            // streams. Key response-side engine logic on the worker's
+            // `runtime_type()`, never on the response variant.
             Self::Zmq(stream) => stream
                 .next()
                 .await