[sgl-router] refactor - session-aware policy (#40379)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

experimental/sgl-router/POLICY_DESIGN.md

@@ -172,7 +172,7 @@ bucket, both groups share this one request-length decision. Policy fallback on
 a cache/affinity miss stays within that group's candidates. There is no second
 pass with relaxed admission and no post-policy substitution.
 
-SLO ordering, cache lookup, and session/sticky policies are follow-ups. Their
+SLO ordering, cache lookup, global session modes, and sticky policies are follow-ups. Their
 integration must preserve bucket-first selection and the same-bucket PD rule.
 Cross-bucket affinity probing is not part of this interface. Session/routing
 keys still pass through `PickRequest` for policies operating inside the selected
@@ -294,6 +294,18 @@ separate mechanism.
 | `StickyPolicy` | Reuse an admitted routing-key binding; use the configured fallback for new or missing keys |
 | `CacheAwarePolicy` | Prefer a usable prefix under cache and pressure rules; use a load-based fallback on a miss |
 
+Session assignments are scoped by model, bucket ID, stage, and session key.
+`SessionAwarePolicy::new(store, engine_load)` receives shared state; the caller
+owns the store's idle timeout and eviction task. Missing or empty session keys
+use power-of-two without creating assignments. A new or out-of-group binding
+uses power-of-two with `AllowAll`, then the session policy checks its selected
+engine before binding. A concurrent live assignment wins, but is checked before
+returning it; rejection ends that attempt without rewriting the binding or
+retrying another engine. Existing bindings are reused regardless of pressure
+when admitted. Session policies can be attached independently to each role.
+Programmatic reorg callers configure `model.affinity.session_id_header` for HTTP
+header extraction; this does not enable legacy global modes or backup escape.
+
 Session and sticky policies do not create assignments for missing keys. A
 binding outside the candidates cannot win. A missing binding may invoke policy
 fallback within the group; hard admission rejection remains an error.
@@ -557,8 +569,15 @@ Implemented here:
 - `AppContext::chat_routing` configures legacy versus reorg routing on the same
   endpoint and carries the reorg model-resolver map.
 
-Follow-up order: concrete admission (#40271), then bucket SLO ordering
-in a separate PR, followed by remaining policies and production configuration.
+- `SessionAwarePolicy` reuses admitted model/bucket/role-scoped bindings from a
+  shared `AffinityStore`, falling back to power-of-two for new or keyless sessions.
+  Assignments follow admission; concurrent binding winners are rechecked.
+  Rejection preserves existing bindings and advances to the next bucket.
+  The caller owns expiry and sweeper lifecycle. A binding may remain after a
+  later PD group fails, because it records placement rather than dispatch.
+
+Follow-up order: cache-aware selection (#40366), concrete admission (#40271),
+then bucket SLO ordering, remaining policies, and production configuration.
 
 Not yet implemented in the reorg path:
 
@@ -566,7 +585,8 @@ Not yet implemented in the reorg path:
 - SLO estimates, targets, and bucket preference ordering.
 - CLI/configuration parsing, validation, and model-specific construction.
   The YAML above is illustrative; reorg resolvers are installed in code.
-- Session modes, prefix memoization, and cache-aware selection.
+- Global session modes and sticky routing-key affinity.
+- Prefix memoization and cache-aware selection.
 - Shared load interpretation, dispatch correction, and policy-specific
   dispatch-timestamp requirements.
 - PD compatibility filtering, retry integration, and legacy-route switchover.

experimental/sgl-router/src/policies_reorg/mod.rs

@@ -6,6 +6,7 @@
 
 pub mod admission;
 pub mod power_of_two;
+pub mod session_aware;
 
 use std::fmt::Debug;
 use std::sync::Arc;

experimental/sgl-router/src/policies_reorg/session_aware.rs

@@ -0,0 +1,124 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+//! Bucket-scoped session placement. Admission rejection preserves the binding
+//! and returns to the bucket loop; it never selects a backup inside the group.
+
+use std::sync::Arc;
+use std::time::Instant;
+
+use futures::future::BoxFuture;
+
+use crate::state::load_monitor::engine_reported_load::{
+    EngineReportedLoadSnapshot, EngineReportedLoadTable,
+};
+use crate::state::AffinityStore;
+use crate::workers::Worker;
+
+use super::admission::{AllowAll, Decision, EngineAdmission};
+use super::power_of_two::PowerOfTwoPolicy;
+use super::{Pick, PickError, PickRequest, Policy, Rejection};
+
+#[derive(Debug)]
+pub struct SessionAwarePolicy {
+    store: Arc<AffinityStore>,
+    engine_load: Arc<EngineReportedLoadTable>,
+    fallback: PowerOfTwoPolicy,
+    pub admission: Arc<dyn EngineAdmission>,
+}
+
+impl SessionAwarePolicy {
+    /// The caller owns the shared store's idle timeout and sweeper lifecycle.
+    /// This policy implements bucket-scoped affinity, without legacy global modes
+    /// or primary/backup pressure escape.
+    pub fn new(store: Arc<AffinityStore>, engine_load: Arc<EngineReportedLoadTable>) -> Self {
+        Self {
+            store,
+            fallback: PowerOfTwoPolicy::new(Arc::clone(&engine_load)),
+            engine_load,
+            admission: Arc::new(AllowAll),
+        }
+    }
+
+    fn assignment_key(request: &PickRequest<'_>) -> Option<String> {
+        let session = request.session_key.filter(|key| !key.is_empty())?;
+        // Length prefixes keep arbitrary model, bucket and session strings
+        // unambiguous, including embedded delimiters. Roles never share bindings.
+        Some(format!(
+            "session:{:?}:{}:{}{}:{}{}",
+            request.stage,
+            request.model.0.len(),
+            request.model.0,
+            request.bucket.len(),
+            request.bucket,
+            session
+        ))
+    }
+
+    fn check(
+        &self,
+        engine: &Worker,
+        request: &PickRequest<'_>,
+        load: &EngineReportedLoadSnapshot,
+    ) -> Result<(), PickError> {
+        match self
+            .admission
+            .check(engine, request, load.fresh_load_for_url(&engine.url))?
+        {
+            Decision::Allow => Ok(()),
+            Decision::Reject(reason) => Err(PickError::AdmissionRejected(Rejection {
+                engine: engine.id.clone(),
+                reason,
+            })),
+        }
+    }
+}
+
+impl Policy for SessionAwarePolicy {
+    fn pick<'a>(
+        &'a self,
+        engines: &'a [Arc<Worker>],
+        request: &'a PickRequest<'a>,
+    ) -> BoxFuture<'a, Result<Pick, PickError>> {
+        Box::pin(async move {
+            if engines.is_empty() {
+                return Err(PickError::NoCandidates);
+            }
+            let key = Self::assignment_key(request);
+            if let Some(bound) = key.as_ref().and_then(|key| self.store.bound(key, engines)) {
+                let load = self.engine_load.capture_snapshot(Instant::now());
+                self.check(bound, request, &load)?;
+                return Ok(Pick {
+                    engine: Arc::clone(bound),
+                    reason: "session_primary",
+                });
+            }
+
+            // The nested power-of-two policy uses AllowAll. The session owner
+            // checks its chosen engine before creating or replacing a binding.
+            let mut pick = self.pick_fallback(engines, request).await?;
+            let load = self.engine_load.capture_snapshot(Instant::now());
+            self.check(&pick.engine, request, &load)?;
+            let Some(key) = key else {
+                pick.reason = "no_session";
+                return Ok(pick);
+            };
+
+            let effective = self.store.bind(key, &pick.engine, engines);
+            if !Arc::ptr_eq(effective, &pick.engine) {
+                // A racing first assignment wins. Check it once, without
+                // rewriting a rejected binding or retrying another engine.
+                self.check(effective, request, &load)?;
+                pick.reason = "session_primary";
+            } else {
+                pick.reason = "assigned";
+            }
+            pick.engine = Arc::clone(effective);
+            Ok(pick)
+        })
+    }
+
+    fn fallback(&self) -> Option<&dyn Policy> {
+        Some(&self.fallback)
+    }
+}

experimental/sgl-router/tests/component/main.rs

@@ -13,5 +13,6 @@ mod policies;
 mod policies_reorg;
 mod policies_reorg_load;
 mod policies_reorg_power_of_two;
+mod policies_reorg_session_aware;
 mod tokenizer;
 mod workers;

experimental/sgl-router/tests/component/policies_reorg_session_aware.rs

@@ -0,0 +1,410 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec};
+use sgl_router::policies_reorg::admission::{Decision, EngineAdmission};
+use sgl_router::policies_reorg::session_aware::SessionAwarePolicy;
+use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage};
+use sgl_router::state::load_monitor::engine_reported_load::{
+    EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat,
+};
+use sgl_router::state::load_monitor::router_inflight_load::MockClock;
+use sgl_router::state::AffinityStore;
+use sgl_router::workers::Worker;
+
+fn engine(id: &str, active: usize) -> Arc<Worker> {
+    let engine = Arc::new(Worker::new(WorkerSpec {
+        id: WorkerId(id.into()),
+        url: format!("http://{id}"),
+        mode: Stage::Plain,
+        model_ids: vec![ModelId("m".into())],
+        bootstrap_port: None,
+    }));
+    engine.active_requests.store(active, Ordering::Relaxed);
+    engine
+}
+
+fn request(model: &ModelId) -> PickRequest<'_> {
+    PickRequest {
+        bucket: "bucket",
+        session_key: Some("session"),
+        ..PickRequest::new(model, Stage::Plain, 8)
+    }
+}
+
+fn policy() -> (SessionAwarePolicy, Arc<AffinityStore>) {
+    let store = AffinityStore::new(Duration::from_secs(60));
+    (
+        SessionAwarePolicy::new(store.clone(), EngineReportedLoadTable::new()),
+        store,
+    )
+}
+
+#[derive(Debug, Default)]
+struct Admission {
+    reject: AtomicBool,
+    invalid: AtomicBool,
+    calls: Mutex<Vec<(String, Option<u64>)>>,
+}
+
+impl EngineAdmission for Admission {
+    fn check(
+        &self,
+        engine: &Worker,
+        _: &PickRequest<'_>,
+        load: Option<&EngineReportedWorkerLoad>,
+    ) -> Result<Decision, PickError> {
+        self.calls
+            .lock()
+            .unwrap()
+            .push((engine.id.0.clone(), load.map(|load| load.num_waiting_reqs)));
+        if self.invalid.load(Ordering::Relaxed) {
+            return Err(PickError::InvalidSignal("invalid admission input".into()));
+        }
+        Ok(if self.reject.load(Ordering::Relaxed) {
+            Decision::Reject("full".into())
+        } else {
+            Decision::Allow
+        })
+    }
+}
+
+#[tokio::test]
+async fn new_sessions_choose_lower_pressure_then_reuse_binding() {
+    let (policy, store) = policy();
+    let engines = [engine("a", 0), engine("b", 9)];
+    let model = ModelId("m".into());
+    let request = request(&model);
+    let first = policy.pick(&engines, &request).await.unwrap();
+    assert_eq!(first.engine.id.0, "a");
+    assert_eq!(first.reason, "assigned");
+    engines[0].active_requests.store(100, Ordering::Relaxed);
+    for _ in 0..5 {
+        let pick = policy.pick(&engines, &request).await.unwrap();
+        assert!(Arc::ptr_eq(&pick.engine, &engines[0]));
+        assert_eq!(pick.reason, "session_primary");
+    }
+    assert_eq!(store.len(), 1);
+}
+
+#[tokio::test]
+async fn missing_and_empty_keys_use_admitted_power_of_two_without_binding() {
+    let (mut policy, store) = policy();
+    let admission = Arc::new(Admission::default());
+    policy.admission = admission.clone();
+    let model = ModelId("m".into());
+    let engines = [engine("a", 9), engine("b", 0)];
+    for key in [None, Some("")] {
+        let request = PickRequest {
+            session_key: key,
+            ..request(&model)
+        };
+        let pick = policy.pick(&engines, &request).await.unwrap();
+        assert_eq!(pick.engine.id.0, "b");
+        assert_eq!(pick.reason, "no_session");
+        admission.reject.store(true, Ordering::Relaxed);
+        assert!(matches!(
+            policy.pick(&engines, &request).await,
+            Err(PickError::AdmissionRejected(_))
+        ));
+        admission.reject.store(false, Ordering::Relaxed);
+    }
+    assert!(store.is_empty());
+    assert_eq!(admission.calls.lock().unwrap().len(), 4);
+}
+
+#[tokio::test]
+async fn rejected_new_and_existing_sessions_never_rebind_or_try_another_engine() {
+    let (mut policy, store) = policy();
+    let admission = Arc::new(Admission::default());
+    policy.admission = admission.clone();
+    let engines = [engine("a", 0), engine("b", 9)];
+    let model = ModelId("m".into());
+    let request = request(&model);
+    admission.reject.store(true, Ordering::Relaxed);
+    assert!(matches!(
+        policy.pick(&engines, &request).await,
+        Err(PickError::AdmissionRejected(_))
+    ));
+    assert!(store.is_empty());
+    admission.reject.store(false, Ordering::Relaxed);
+    policy.pick(&engines, &request).await.unwrap();
+    engines[0].active_requests.store(100, Ordering::Relaxed);
+    admission.reject.store(true, Ordering::Relaxed);
+    assert!(matches!(
+        policy.pick(&engines, &request).await,
+        Err(PickError::AdmissionRejected(_))
+    ));
+    admission.reject.store(false, Ordering::Relaxed);
+    assert_eq!(
+        policy.pick(&engines, &request).await.unwrap().engine.id.0,
+        "a"
+    );
+    assert_eq!(store.len(), 1);
+    assert_eq!(admission.calls.lock().unwrap().len(), 4);
+}
+
+#[tokio::test]
+async fn removed_binding_is_replaced_only_after_admission() {
+    let (mut policy, store) = policy();
+    let admission = Arc::new(Admission::default());
+    policy.admission = admission.clone();
+    let model = ModelId("m".into());
+    let request = request(&model);
+    let original = [engine("a", 0)];
+    let replacement = [engine("b", 0)];
+    policy.pick(&original, &request).await.unwrap();
+    admission.reject.store(true, Ordering::Relaxed);
+    assert!(matches!(
+        policy.pick(&replacement, &request).await,
+        Err(PickError::AdmissionRejected(_))
+    ));
+    admission.reject.store(false, Ordering::Relaxed);
+    assert_eq!(
+        policy.pick(&original, &request).await.unwrap().reason,
+        "session_primary"
+    );
+    assert_eq!(
+        policy.pick(&replacement, &request).await.unwrap().reason,
+        "assigned"
+    );
+    assert_eq!(store.len(), 1);
+}
+
+#[tokio::test]
+async fn bindings_are_scoped_by_model_bucket_role_and_session() {
+    let (policy, store) = policy();
+    let models = [ModelId("m".into()), ModelId("other".into())];
+    let original = [engine("a", 0)];
+    let fleet = [original[0].clone(), engine("b", 0)];
+    policy.pick(&original, &request(&models[0])).await.unwrap();
+    original[0].active_requests.store(100, Ordering::Relaxed);
+    let requests = [
+        request(&models[1]),
+        PickRequest {
+            bucket: "other",
+            ..request(&models[0])
+        },
+        PickRequest {
+            stage: Stage::Prefill,
+            ..request(&models[0])
+        },
+        PickRequest {
+            stage: Stage::Decode,
+            ..request(&models[0])
+        },
+        PickRequest {
+            session_key: Some("other"),
+            ..request(&models[0])
+        },
+    ];
+    for request in requests {
+        let pick = policy.pick(&fleet, &request).await.unwrap();
+        assert_eq!(pick.engine.id.0, "b");
+        assert_eq!(pick.reason, "assigned");
+    }
+    assert_eq!(
+        policy
+            .pick(&fleet, &request(&models[0]))
+            .await
+            .unwrap()
+            .engine
+            .id
+            .0,
+        "a"
+    );
+    assert_eq!(store.len(), 6);
+}
+
+#[tokio::test]
+async fn embedded_delimiters_do_not_alias_scopes() {
+    let (policy, store) = policy();
+    let models = [ModelId("a\0b".into()), ModelId("a".into())];
+    let first = PickRequest {
+        bucket: "c",
+        session_key: Some("d\0e"),
+        ..request(&models[0])
+    };
+    let second = PickRequest {
+        bucket: "b\0c",
+        session_key: Some("d\0e"),
+        ..request(&models[1])
+    };
+    let third = PickRequest {
+        bucket: "c\0d",
+        session_key: Some("e"),
+        ..request(&models[0])
+    };
+    let fleet = [engine("a", 0), engine("b", 9)];
+    policy.pick(&fleet, &first).await.unwrap();
+    fleet[0].active_requests.store(100, Ordering::Relaxed);
+    for request in [second, third] {
+        assert_eq!(
+            policy.pick(&fleet, &request).await.unwrap().engine.id.0,
+            "b"
+        );
+    }
+    assert_eq!(store.len(), 3);
+}
+
+#[tokio::test]
+async fn same_id_replacement_returns_the_live_candidate_instance() {
+    let (policy, _) = policy();
+    let model = ModelId("m".into());
+    let request = request(&model);
+    policy.pick(&[engine("a", 0)], &request).await.unwrap();
+    let replacement = [engine("a", 0)];
+    let pick = policy.pick(&replacement, &request).await.unwrap();
+    assert!(Arc::ptr_eq(&pick.engine, &replacement[0]));
+    assert_eq!(pick.reason, "session_primary");
+}
+
+#[tokio::test]
+async fn empty_candidates_skip_admission_and_invalid_signals_never_bind() {
+    let (mut policy, store) = policy();
+    let admission = Arc::new(Admission::default());
+    policy.admission = admission.clone();
+    let model = ModelId("m".into());
+    let request = request(&model);
+    assert!(matches!(
+        policy.pick(&[], &request).await,
+        Err(PickError::NoCandidates)
+    ));
+    assert!(admission.calls.lock().unwrap().is_empty());
+    admission.invalid.store(true, Ordering::Relaxed);
+    assert!(matches!(
+        policy.pick(&[engine("a", 0)], &request).await,
+        Err(PickError::InvalidSignal(_))
+    ));
+    assert!(store.is_empty());
+}
+
+#[tokio::test]
+async fn admission_receives_fresh_load_on_assignment_and_reuse() {
+    let store = AffinityStore::new(Duration::from_secs(60));
+    let table = EngineReportedLoadTable::new();
+    let mut policy = SessionAwarePolicy::new(store, table.clone());
+    let admission = Arc::new(Admission::default());
+    policy.admission = admission.clone();
+    let engines = [engine("a", 0)];
+    let model = ModelId("m".into());
+    let request = request(&model);
+    for waiting in [3, 7] {
+        table.set(
+            &engines[0].url,
+            0,
+            LoadStat {
+                num_waiting_reqs: waiting,
+                num_running_reqs: 1,
+                num_tokens: 10,
+                max_total_num_tokens: 100,
+                native_cache: None,
+            },
+            Instant::now(),
+        );
+        policy.pick(&engines, &request).await.unwrap();
+    }
+    assert_eq!(
+        *admission.calls.lock().unwrap(),
+        vec![("a".into(), Some(3)), ("a".into(), Some(7))]
+    );
+}
+
+#[tokio::test]
+async fn shared_store_refreshes_active_sessions_and_expires_idle_ones() {
+    let clock = Arc::new(MockClock::new(Instant::now()));
+    let store = AffinityStore::with_clock(Duration::from_secs(10), clock.clone());
+    let policy = SessionAwarePolicy::new(store.clone(), EngineReportedLoadTable::new());
+    let model = ModelId("m".into());
+    let fleet = [engine("a", 0), engine("b", 9)];
+    let hot = request(&model);
+    let cold = PickRequest {
+        session_key: Some("cold"),
+        ..hot
+    };
+    policy.pick(&fleet, &hot).await.unwrap();
+    policy.pick(&fleet, &cold).await.unwrap();
+    clock.advance(Duration::from_secs(8));
+    fleet[0].active_requests.store(100, Ordering::Relaxed);
+    policy.pick(&fleet, &hot).await.unwrap();
+    clock.advance(Duration::from_secs(8));
+    assert_eq!(store.sweep_expired(), 1);
+    assert_eq!(policy.pick(&fleet, &hot).await.unwrap().engine.id.0, "a");
+    assert_eq!(policy.pick(&fleet, &cold).await.unwrap().engine.id.0, "b");
+}
+
+#[derive(Debug)]
+struct RacingAdmission {
+    competitor: SessionAwarePolicy,
+    winner: Arc<Worker>,
+    reject_winner: bool,
+    calls: Mutex<Vec<String>>,
+}
+
+impl EngineAdmission for RacingAdmission {
+    fn check(
+        &self,
+        engine: &Worker,
+        request: &PickRequest<'_>,
+        _: Option<&EngineReportedWorkerLoad>,
+    ) -> Result<Decision, PickError> {
+        self.calls.lock().unwrap().push(engine.id.0.clone());
+        if engine.id.0 == "a" {
+            // Complete a competing first request after this request selected a,
+            // but before it can commit. The effective binding is now b.
+            futures::executor::block_on(
+                self.competitor
+                    .pick(std::slice::from_ref(&self.winner), request),
+            )?;
+        }
+        Ok(if self.reject_winner && engine.id == self.winner.id {
+            Decision::Reject("racing winner full".into())
+        } else {
+            Decision::Allow
+        })
+    }
+}
+
+#[tokio::test]
+async fn concurrent_assignment_winner_is_checked_and_preserved_on_rejection() {
+    for reject_winner in [false, true] {
+        let (mut policy, store) = policy();
+        let engines = [engine("a", 0), engine("b", 9)];
+        let admission = Arc::new(RacingAdmission {
+            competitor: SessionAwarePolicy::new(store.clone(), EngineReportedLoadTable::new()),
+            winner: engines[1].clone(),
+            reject_winner,
+            calls: Mutex::new(Vec::new()),
+        });
+        policy.admission = admission.clone();
+        let model = ModelId("m".into());
+        let request = request(&model);
+        let result = policy.pick(&engines, &request).await;
+        if reject_winner {
+            assert!(
+                matches!(result, Err(PickError::AdmissionRejected(rejection)) if rejection.engine.0 == "b")
+            );
+        } else {
+            let pick = result.unwrap();
+            assert_eq!(pick.engine.id.0, "b");
+            assert_eq!(pick.reason, "session_primary");
+        }
+        assert_eq!(*admission.calls.lock().unwrap(), vec!["a", "b"]);
+        assert_eq!(
+            admission
+                .competitor
+                .pick(&engines, &request)
+                .await
+                .unwrap()
+                .engine
+                .id
+                .0,
+            "b"
+        );
+        assert_eq!(store.len(), 1);
+    }
+}

experimental/sgl-router/tests/proxy/chat_routing/reorg.rs

@@ -12,6 +12,8 @@ use sgl_router::server::app_context::ChatRouting;
 use sgl_router::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad;
 use std::sync::Mutex;
 
+mod session_aware;
+
 type PickCall = (String, Stage, u64, Option<u64>);
 
 #[derive(Debug)]

experimental/sgl-router/tests/proxy/chat_routing/reorg/session_aware.rs

@@ -0,0 +1,119 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors
+// SPDX-License-Identifier: Apache-2.0
+
+use super::*;
+
+#[tokio::test]
+async fn session_aware_reuses_custom_header_binding_after_load_changes() {
+    use sgl_router::config::AffinityConfig;
+    use sgl_router::policies_reorg::session_aware::SessionAwarePolicy;
+    use sgl_router::state::load_monitor::engine_reported_load::EngineReportedLoadTable;
+    use sgl_router::state::AffinityStore;
+    use std::sync::atomic::Ordering;
+    use std::time::Duration;
+
+    let primary = MockWorker::start(vec![]).await;
+    let other = MockWorker::start(vec![]).await;
+    let store = AffinityStore::new(Duration::from_secs(60));
+    let policy = Arc::new(SessionAwarePolicy::new(
+        store.clone(),
+        EngineReportedLoadTable::new(),
+    ));
+    let mut ctx = context(
+        &[
+            ("primary", Stage::Plain, &primary),
+            ("other", Stage::Plain, &other),
+        ],
+        vec![Bucket::new(
+            "session",
+            BucketGroups::Plain(EngineGroup::new(policy)),
+        )],
+    );
+    Arc::get_mut(&mut ctx).unwrap().config.model.affinity = Some(AffinityConfig {
+        session_id_header: "x-test-session".into(),
+        ..Default::default()
+    });
+    let primary_worker = ctx.registry.get(&WorkerId("primary".into())).unwrap();
+    let other_worker = ctx.registry.get(&WorkerId("other".into())).unwrap();
+    other_worker.active_requests.store(10, Ordering::Relaxed);
+    let app = build_router(ctx);
+    for _ in 0..2 {
+        let mut req = request(body("hello"));
+        req.headers_mut()
+            .insert("x-test-session", "same-session".parse().unwrap());
+        let response = app.clone().oneshot(req).await.unwrap();
+        assert_eq!(response.status(), StatusCode::OK);
+        let _ = response.into_body().collect().await.unwrap();
+        assert!(primary.captured.lock().unwrap().last_body.is_some());
+        assert!(other.captured.lock().unwrap().last_body.is_none());
+        primary_worker.active_requests.store(100, Ordering::Relaxed);
+        other_worker.active_requests.store(0, Ordering::Relaxed);
+    }
+    assert_eq!(store.len(), 1);
+}
+
+#[tokio::test]
+async fn rejected_session_binding_advances_buckets_without_reassignment_or_dispatch() {
+    use sgl_router::config::AffinityConfig;
+    use sgl_router::policies_reorg::session_aware::SessionAwarePolicy;
+    use sgl_router::state::load_monitor::engine_reported_load::EngineReportedLoadTable;
+    use sgl_router::state::AffinityStore;
+    use std::time::Duration;
+
+    let primary = MockWorker::start(vec![]).await;
+    let backup = MockWorker::start(vec![]).await;
+    let store = AffinityStore::new(Duration::from_secs(60));
+    let table = EngineReportedLoadTable::new();
+    let mut rejected = SessionAwarePolicy::new(store.clone(), table.clone());
+    rejected.admission = Arc::new(RejectAll);
+    let accepted = SessionAwarePolicy::new(store.clone(), table.clone());
+    let mut ctx = context(
+        &[
+            ("primary", Stage::Plain, &primary),
+            ("backup", Stage::Plain, &backup),
+        ],
+        vec![
+            Bucket::new(
+                "a-primary",
+                BucketGroups::Plain(EngineGroup {
+                    worker_ids: Some([WorkerId("primary".into())].into_iter().collect()),
+                    policy: Arc::new(rejected),
+                }),
+            ),
+            Bucket::new(
+                "b-backup",
+                BucketGroups::Plain(EngineGroup {
+                    worker_ids: Some([WorkerId("backup".into())].into_iter().collect()),
+                    policy: Arc::new(accepted),
+                }),
+            ),
+        ],
+    );
+    Arc::get_mut(&mut ctx).unwrap().config.model.affinity = Some(AffinityConfig::default());
+    let seed = SessionAwarePolicy::new(store.clone(), table);
+    let model = ModelId("tiny".into());
+    let pick_request = PickRequest {
+        bucket: "a-primary",
+        session_key: Some("same-session"),
+        ..PickRequest::new(&model, Stage::Plain, 1)
+    };
+    let engine = ctx.registry.get(&WorkerId("primary".into())).unwrap();
+    seed.pick(std::slice::from_ref(&engine), &pick_request)
+        .await
+        .unwrap();
+    let mut req = request(body("hello"));
+    req.headers_mut()
+        .insert("x-session-id", "same-session".parse().unwrap());
+    let response = build_router(ctx).oneshot(req).await.unwrap();
+    assert_eq!(response.status(), StatusCode::OK);
+    let _ = response.into_body().collect().await.unwrap();
+    assert!(primary.captured.lock().unwrap().last_body.is_none());
+    assert!(backup.captured.lock().unwrap().last_body.is_some());
+    assert_eq!(store.len(), 2);
+    let retained = seed
+        .pick(std::slice::from_ref(&engine), &pick_request)
+        .await
+        .unwrap();
+    assert_eq!(retained.reason, "session_primary");
+    assert_eq!(retained.engine.id.0, "primary");
+}