fix(router): backport unknown worker target validation to 1.5 (#14858)

lib/llm/src/kv_router/prefill_router/mod.rs

@@ -20,6 +20,7 @@ use dynamo_kv_router::{
     selector::{DefaultWorkerSelector, WorkerSelector},
 };
 use dynamo_runtime::{
+    error::{ErrorType, match_error_chain},
     pipeline::{
         AsyncEngineContextProvider, Context, ManyOut, Operator, ResponseStream, RouterMode,
         ServerStreamingEngine, SingleIn, async_trait, propagate_first_response_guard,
@@ -400,7 +401,13 @@ where
                 Ok(None) => {
                     (req, context) = conditional_request.into_parts();
                 }
-                Err(error) if crate::kv_router::routing_host::is_cancelled(&error) => {
+                Err(error)
+                    if match_error_chain(
+                        error.as_ref(),
+                        &[ErrorType::Cancelled, ErrorType::InvalidArgument],
+                        &[],
+                    ) =>
+                {
                     return Err(error);
                 }
                 Err(error) => {
@@ -846,6 +853,76 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn conditional_disagg_rejects_unknown_decode_target_before_fallback() {
+        use crate::kv_router::KvRouter;
+        use dynamo_runtime::{
+            DistributedRuntime, Runtime, distributed::DistributedConfig, pipeline::PushRouter,
+        };
+
+        let runtime = Runtime::from_current().unwrap();
+        let distributed =
+            DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
+                .await
+                .unwrap();
+        let endpoint = distributed
+            .namespace("conditional-invalid-worker")
+            .unwrap()
+            .component("decode")
+            .unwrap()
+            .endpoint("generate");
+        let client = endpoint.client().await.unwrap();
+        let (_workers_tx, workers) = watch::channel(HashMap::<u64, ModelRuntimeConfig>::new());
+        let config = KvRouterConfig {
+            skip_initial_worker_wait: true,
+            use_kv_events: false,
+            router_track_active_blocks: false,
+            ..Default::default()
+        };
+        let chooser = KvRouter::new(
+            endpoint,
+            client.clone(),
+            workers,
+            None,
+            16,
+            DefaultWorkerSelector::new(Some(config.clone()), "decode"),
+            Some(config),
+            None,
+            "decode",
+            None,
+            false,
+            None,
+            None,
+        )
+        .await
+        .unwrap();
+        let inner = PushRouter::from_client(client, RouterMode::KV)
+            .await
+            .unwrap();
+        let host = Arc::new(RoutingHost::new(inner, Arc::new(chooser), None).unwrap());
+        let router = active_conditional_router();
+        assert!(router.decode_routing_host.set(host).is_ok());
+        let downstream = Arc::new(QueryOnlyDecodeHost::default());
+        let mut request = query_only_request();
+        request.annotations.clear();
+        request.routing_mut().decode_worker_id = Some(u64::MAX);
+
+        // No prefill binding: the old fallback would continue to downstream instead
+        // of returning the preview's client error.
+        let error = router
+            .generate(SingleIn::new(request), downstream.clone())
+            .await
+            .unwrap_err();
+        assert!(match_error_chain(
+            error.as_ref(),
+            &[ErrorType::InvalidArgument],
+            &[]
+        ));
+        assert_eq!(downstream.requests.load(Ordering::Relaxed), 0);
+        drop(router);
+        runtime.shutdown();
+    }
+
     #[tokio::test]
     async fn conditional_disagg_query_only_forwards_to_decode_host() {
         let router = active_conditional_router();

lib/llm/src/kv_router/routing_host.rs

@@ -576,6 +576,40 @@ where
         }
     }
 
+    /// Check request-supplied targets only at initial selection, not after a route preview
+    /// or when dispatching: a worker disappearing later is a service failure.
+    fn validate_explicit_worker(
+        &self,
+        request: &PreprocessedRequest,
+        phase: RequestPhase,
+    ) -> Result<(), Error> {
+        let Some(routing) = request.routing.as_ref() else {
+            return Ok(());
+        };
+        let target = match phase {
+            RequestPhase::Prefill => routing
+                .prefill_worker_id
+                .map(|id| (id, "prefill_worker_id")),
+            RequestPhase::Decode | RequestPhase::Aggregated => {
+                routing.decode_worker_id.map(|id| (id, "decode_worker_id"))
+            }
+        }
+        .or_else(|| {
+            routing
+                .backend_instance_id
+                .map(|id| (id, "backend_instance_id"))
+        });
+        // Validate discovery membership, not health: unavailable workers are still known.
+        if let Some((worker_id, field)) = target
+            && !self.inner.client.is_instance_discovered(worker_id)
+        {
+            return Err(invalid_argument(format!(
+                "nvext.{field}={worker_id} does not identify a known worker"
+            )));
+        }
+        Ok(())
+    }
+
     fn affinity_target_is_valid(&self, target: AffinityTarget) -> bool {
         if !self.inner.client.is_instance_discovered(target.worker_id) {
             return false;

lib/llm/src/kv_router/routing_host/builtin.rs

@@ -323,6 +323,7 @@ where
                 None,
             )
         } else {
+            self.validate_explicit_worker(request.content(), phase)?;
             self.select_with_session_affinity(&request, phase, is_query_only, &budget, |target| {
                 let pinned_target = explicit.or(match self.session_affinity_mode {
                     SessionAffinityMode::Hard => target,

lib/llm/src/kv_router/routing_host/kv.rs

@@ -93,6 +93,7 @@ where
         is_query_only: bool,
         budget: &CleanupBudget,
     ) -> Result<(WorkerSelection, Option<AffinityAcquire>), Error> {
+        self.validate_explicit_worker(request.content(), phase)?;
         self.select_with_session_affinity(request, phase, is_query_only, budget, |target| {
             self.select_request(request, phase, is_query_only, target, budget)
         })
@@ -132,6 +133,7 @@ where
         if self.kv_router_if_enabled().is_none() {
             return Err(anyhow::anyhow!("KV route previews require KV routing"));
         }
+        self.validate_explicit_worker(request.content(), phase)?;
 
         let phase_label = phase.to_string();
         let route_guard = StageGuard::new(STAGE_ROUTE, &phase_label);