feat(llm): resolve KV-routed sessions through the selection core

lib/kv-router/src/services/selection/affinity/README.md

@@ -31,10 +31,14 @@ The resolver takes a `TargetLiveness` from its host: whether a bound target can
 |---|---|---|---|---|---|---|
 | Selection service (`core/run.rs`, `core/reservations.rs`) | `session_context.session_id` or legacy `session_id` on the wire request | partition catalog (`CatalogLiveness`) | `Book`: resolve, book, commit | after the booking lands, before the reservation is installed | the reservation row | core shutdown token |
 | Standalone EPP (`deploy/inference-gateway/ext-proc/src/selector.rs`) | request session header | same as the selection service it embeds | `Book` | same | same | same |
-| Frontend KV (`lib/llm/src/kv_router/routing_host`) | `x-dynamo-session-affinity-id` header, subagent group key | discovery plus DP-rank range | `Lease`; the host drives the shared table directly (`AffinityCoordinator`, `lib/llm/src/session_affinity`) and passes a `Hard` binding to the core as an exact pin, a `Soft` one as a target | after dispatch returns a response stream | the response stream | client disconnect (or the request's cleanup budget for a decode leg with staged KV) |
-| Frontend builtin modes and prefill (`routing_host/builtin.rs`, `prefill_router`) | same as frontend KV | discovery plus DP-rank range | the host drives the shared table directly (`AffinityCoordinator`) | after dispatch | the response stream | client disconnect |
+| Frontend KV (`lib/llm/src/kv_router/routing_host`) | `x-dynamo-session-affinity-id` header, subagent group key | discovery plus DP-rank range (`DiscoveryLiveness`, installed on the partition resolver) | `Lease`: core resolves and returns the hold in `Selected`; host runs `check_dispatch` before dispatching and commits after | after dispatch returns a response stream | the response stream | client disconnect drops the selection future (or the request's cleanup budget for a decode leg with staged KV) |
+| Frontend builtin modes (`routing_host/builtin.rs`) | same as frontend KV | discovery plus DP-rank range | the host drives the shared table directly (`AffinityCoordinator`) | after dispatch | the response stream | client disconnect |
 | Runtime EPP (`ext-proc/src/epp.rs`) | none | n/a | n/a | n/a | n/a | no decode session affinity |
 
+The selection service, the standalone EPP, and the frontend KV plane (decode and prefill) use the resolver through the core. The frontend builtin modes drive the same table directly through `AffinityCoordinator` (`lib/llm/src/session_affinity`).
+
+Where the hosts differ on a live `Hard` mismatch (the bound worker is live but the request's own constraints steered selection elsewhere): the core rejects at commit and the table drops the binding, so the session re-binds on its next request; the frontend rejects before dispatch with `check_dispatch` and releases its hold without touching the binding, since nothing ran and the request, not the session, caused the mismatch.
+
 ## Errors
 
 Hosts map `AffinityError` onto their own error types. The core maps `InvalidArgument` (the request contradicts the binding or exceeds the session-id limit) to `BadRequest`, `ResourceExhausted` to `NotReady`, `Cancelled` to scheduler shutdown, and `Dropped` to `Internal`. A full table never fails `resolve`: the request routes without affinity and `full_table_fallbacks` counts it (a plain counter on the resolver; not exported as a metric).

lib/kv-router/src/services/selection/affinity/resolver.rs

@@ -224,6 +224,31 @@ impl AffinityResolver {
         Ok(None)
     }
 
+    /// Whether `commit` would reject `dispatched` for a live `Hard` binding.
+    ///
+    /// For a host that commits after dispatch, this is the check to run
+    /// before dispatching: a `Hard` binding whose worker is live but was not
+    /// selected (the request's own constraints excluded it) would be rejected
+    /// at commit, after the worker has already started the request. The hold
+    /// is left as it is; the host decides whether to release or invalidate it.
+    pub fn check_dispatch(
+        &self,
+        hold: &Hold,
+        dispatched: AffinityTarget,
+    ) -> Result<(), AffinityError> {
+        if self.table.mode() != SessionAffinityMode::Hard {
+            return Ok(());
+        }
+        let Some(target) = hold.target() else {
+            return Ok(());
+        };
+        if !self.liveness.is_schedulable(target) {
+            // `commit` fails this binding over to `dispatched` instead.
+            return Ok(());
+        }
+        validate_dispatch_target(hold.session_id(), target, dispatched)
+    }
+
     /// Give up a hold because the selection policy filtered out every
     /// candidate. In `Hard` mode that includes the bound worker, which the
     /// scheduler had limited selection to: drop the binding so the session
@@ -624,6 +649,36 @@ mod tests {
         ));
     }
 
+    #[tokio::test]
+    async fn check_dispatch_rejects_only_a_live_hard_mismatch() {
+        let (hard, liveness) = make_resolver(SessionAffinityMode::Hard, &[1, 2]);
+        bind(&hard, "s", target(1, Some(0))).await;
+        let held = hold(&hard, "s").await;
+        assert!(hard.check_dispatch(&held, target(1, Some(0))).is_ok());
+        assert!(matches!(
+            hard.check_dispatch(&held, target(2, Some(0))),
+            Err(AffinityError::InvalidArgument(_))
+        ));
+        assert_eq!(
+            bound(&hard, "s"),
+            Some(target(1, Some(0))),
+            "the check changes nothing"
+        );
+        liveness.set(&[2]);
+        assert!(
+            hard.check_dispatch(&held, target(2, Some(0))).is_ok(),
+            "a departed binding is left for commit to fail over"
+        );
+        drop(held);
+
+        let (soft, _) = make_resolver(SessionAffinityMode::Soft, &[1, 2]);
+        bind(&soft, "s", target(1, Some(0))).await;
+        let held = hold(&soft, "s").await;
+        assert!(soft.check_dispatch(&held, target(2, Some(0))).is_ok());
+        let fresh = hold(&hard, "fresh").await;
+        assert!(hard.check_dispatch(&fresh, target(2, Some(0))).is_ok());
+    }
+
     #[tokio::test]
     async fn filtered_hard_binding_is_released_and_soft_binding_kept() {
         let (hard, _) = make_resolver(SessionAffinityMode::Hard, &[1]);

lib/kv-router/src/services/selection/affinity/table.rs

@@ -209,6 +209,13 @@ impl Hold {
         }
     }
 
+    pub fn session_id(&self) -> &str {
+        match self {
+            Self::Initialize(initialization) => &initialization.session_id,
+            Self::Bound { lease, .. } => lease.session_id(),
+        }
+    }
+
     /// The session id, shared with the hold so a caller that consumes the
     /// hold can still name the session without copying it.
     pub(crate) fn shared_session_id(&self) -> Arc<str> {

lib/kv-router/src/services/selection/core/mod.rs

@@ -111,15 +111,32 @@ impl SelectionPartition {
         &self.0.indexer
     }
 
-    /// Return this partition's affinity table, initialized with `config`.
+    /// Return this partition's affinity table, initialized with `config` and
+    /// the catalog as its liveness source.
     pub fn session_affinity(
         &self,
         config: SessionAffinityConfig,
     ) -> Result<SessionAffinity, SelectionError> {
         self.0
-            .session_affinity(config)
+            .session_affinity_with(config, None)
             .map(|resolver| resolver.table().clone())
     }
+
+    /// Return this partition's session resolver, initialized with `config`
+    /// and the host's own liveness source. A resolver already installed with
+    /// the same configuration is returned as is, whatever its liveness.
+    pub fn session_affinity_with_liveness(
+        &self,
+        config: SessionAffinityConfig,
+        liveness: Arc<dyn TargetLiveness>,
+    ) -> Result<&AffinityResolver, SelectionError> {
+        self.0.session_affinity_with(config, Some(liveness))
+    }
+
+    /// This partition's session resolver, if session affinity is configured.
+    pub fn session_resolver(&self) -> Option<&AffinityResolver> {
+        self.0.affinity.get()
+    }
 }
 
 struct SelectionEntry {
@@ -150,9 +167,12 @@ impl TargetLiveness for CatalogLiveness {
 }
 
 impl SelectionEntry {
-    fn session_affinity(
+    /// The partition's resolver, created on first use with `liveness` or,
+    /// when the host supplies none, the catalog.
+    fn session_affinity_with(
         &self,
         config: SessionAffinityConfig,
+        liveness: Option<Arc<dyn TargetLiveness>>,
     ) -> Result<&AffinityResolver, SelectionError> {
         let resolver = self
             .affinity
@@ -163,13 +183,13 @@ impl SelectionEntry {
                 {
                     table.enable_replication(config.process_id(), sink);
                 }
-                Ok(AffinityResolver::new(
-                    table,
+                let liveness = liveness.unwrap_or_else(|| {
                     Arc::new(CatalogLiveness {
                         catalog: Arc::clone(&self.catalog),
                         key: self.key.clone(),
-                    }),
-                ))
+                    })
+                });
+                Ok(AffinityResolver::new(table, liveness))
             })?;
         let table = resolver.table();
         if table.ttl() != config.ttl || table.mode() != config.mode {

lib/kv-router/src/services/selection/core/operation.rs

@@ -4,6 +4,7 @@
 //! The one selection operation every host runs. Wire handlers and embedding
 //! hosts build a [`SelectionOperation`] and consume a [`SelectionOutcome`].
 
+use std::borrow::Cow;
 use std::collections::HashSet;
 use std::time::Duration;
 
@@ -19,6 +20,7 @@ use crate::scheduling::config::RouterConfigOverride;
 use crate::scheduling::queue::BookingHandle;
 use crate::scheduling::{AdvisoryWorkerLoad, QueueRejection, SchedulingResponse, SessionContext};
 
+use super::super::affinity::Hold;
 use super::super::error::SelectionError;
 use super::super::input::PromptView;
 
@@ -34,7 +36,7 @@ pub struct SelectionOperation<'a> {
     pub strict_priority: u32,
     pub policy_class: Option<String>,
     pub session_context: Option<SessionContext>,
-    pub session: SessionBinding,
+    pub session: SessionBinding<'a>,
     pub affinity_target: Option<WorkerAffinityTarget>,
     pub pinned_worker: Option<WorkerWithDpRank>,
     pub allowed_worker_ids: Option<HashSet<WorkerId>>,
@@ -80,15 +82,22 @@ impl SelectionAdmission {
 }
 
 /// What the core does with the partition's session table for this request.
-pub enum SessionBinding {
+///
+/// `requested_target` is an explicit target the request carries; a bound
+/// session must agree with it, and a new session binds only to it.
+pub enum SessionBinding<'a> {
     None,
-    /// Hold the session, steer to its worker, and bind it to the worker booked.
+    /// Hold the session and steer to its worker. Under `Book` the core binds
+    /// the worker booked; under `Lease` it returns the hold in
+    /// [`Selected::affinity_hold`] for the host to commit after dispatch.
     Managed {
-        session_id: String,
+        session_id: Cow<'a, str>,
+        requested_target: Option<WorkerAffinityTarget>,
     },
     /// Steer to the session's worker without holding or binding it.
     Query {
-        session_id: String,
+        session_id: Cow<'a, str>,
+        requested_target: Option<WorkerAffinityTarget>,
     },
 }
 
@@ -144,4 +153,8 @@ pub struct Selected {
     /// The booking's handle; `Lease` admission only. Dropping it frees the
     /// booking, `commit` hands it to the caller's own cleanup.
     pub booking: Option<BookingHandle>,
+    /// The session held for this request; `Lease` admission with a managed
+    /// session only. The host commits it through the partition's resolver
+    /// once it knows where the request went; dropping it releases the hold.
+    pub affinity_hold: Option<Hold>,
 }

lib/kv-router/src/services/selection/core/run.rs

@@ -185,6 +185,7 @@ impl SelectionCore {
             routing_hashes: _,
             shared_cache_hits: _,
             booking: _,
+            affinity_hold: _,
         } = selected;
         let booked = sequence_hashes.is_some();
         let potential_decode_blocks = response.potential_decode_blocks as u64;
@@ -276,24 +277,36 @@ impl SelectionCore {
         // scheduler's pin rule; a new session is bound to the worker booked.
         // An explicit target from the request is a preference the policy reads.
         let resolver = entry.affinity.get();
-        if matches!(session, SessionBinding::Managed { .. }) && claim.is_none() {
+        if matches!(session, SessionBinding::Managed { .. }) && !book {
             return Err(SelectionError::Internal(
-                "a managed session binding requires Book admission".to_string(),
+                "a managed session binding requires a booking admission".to_string(),
             ));
         }
         let mut affinity_hold = None;
         let affinity = match (&session, resolver) {
-            (SessionBinding::Managed { session_id }, Some(resolver)) => {
+            (
+                SessionBinding::Managed {
+                    session_id,
+                    requested_target,
+                },
+                Some(resolver),
+            ) => {
                 let resolution = resolver
-                    .resolve(session_id, None, self.cancel_token.cancelled())
+                    .resolve(session_id, *requested_target, self.cancel_token.cancelled())
                     .await
                     .map_err(affinity_error)?;
                 affinity_hold = resolution.hold;
                 resolution.affinity
             }
-            (SessionBinding::Query { session_id }, Some(resolver)) => {
-                resolver.query(session_id, None).map_err(affinity_error)?
-            }
+            (
+                SessionBinding::Query {
+                    session_id,
+                    requested_target,
+                },
+                Some(resolver),
+            ) => resolver
+                .query(session_id, *requested_target)
+                .map_err(affinity_error)?,
             _ => affinity_target.map(AffinityRequirement::soft),
         };
         // Router hints are attached to bookings only, and only when a worker in
@@ -373,7 +386,7 @@ impl SelectionCore {
                 prompt.lora_name.map(str::to_string),
                 track_prefill_tokens,
                 match &session {
-                    SessionBinding::Query { session_id } => Some(session_id.clone()),
+                    SessionBinding::Query { session_id, .. } => Some(session_id.to_string()),
                     _ => None,
                 },
             )
@@ -505,7 +518,7 @@ impl SelectionCore {
         // The routing hashes go to exactly one of: the reservation recorded
         // now, or the replay cache a later reservation records from.
         let mut routing_hashes = routing_hashes;
-        let booking = if let Some(claim) = claim {
+        let (booking, affinity_hold) = if let Some(claim) = claim {
             let Some(booking) = booking else {
                 return Err(SelectionError::Internal(
                     "booked selection has no booking handle".to_string(),
@@ -524,9 +537,10 @@ impl SelectionCore {
                     .await;
             }
             claim.install(booking, affinity_lease)?;
-            None
+            (None, None)
         } else {
-            booking
+            // A `Lease` host commits the session itself, after dispatch.
+            (booking, affinity_hold)
         };
 
         if let Some((cache_id, sequence_hashes, lora_name, track_prefill_tokens, session_id)) =
@@ -564,6 +578,7 @@ impl SelectionCore {
             routing_hashes: returned_routing_hashes,
             shared_cache_hits: host_shared_cache_hits,
             booking,
+            affinity_hold,
         }))
     }
 
@@ -701,13 +716,15 @@ fn session_binding(
     context: Option<&SessionContext>,
     is_steerable: bool,
     is_booking: bool,
-) -> SessionBinding {
+) -> SessionBinding<'static> {
     match context.filter(|_| is_steerable) {
         Some(context) if is_booking => SessionBinding::Managed {
-            session_id: context.session_id().to_string(),
+            session_id: context.session_id().to_string().into(),
+            requested_target: None,
         },
         Some(context) => SessionBinding::Query {
-            session_id: context.session_id().to_string(),
+            session_id: context.session_id().to_string().into(),
+            requested_target: None,
         },
         None => SessionBinding::None,
     }