dynamo-llm 1.6.2

// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::{
    collections::{HashMap, HashSet},
    sync::{
        Arc, Mutex,
        atomic::{AtomicBool, AtomicUsize, Ordering},
    },
    time::{Duration, Instant},
};

use crate::kv_router::SelectionPolicySource;
use dynamo_kv_router::{
    config::KvRouterConfig,
    protocols::RoutingConstraints,
    scheduling::{
        ClassifierError, ClassifyEvent, ClassifyFuture, ClassifyRequest, RequestClassifier,
    },
};
use dynamo_runtime::{
    DistributedRuntime, Runtime,
    component::{Client, Instance},
    discovery::EventTransportKind,
    distributed::{DiscoveryBackend, DistributedConfig, RequestPlaneMode},
    error::{BackendError, ErrorType, match_error_chain},
    pipeline::{
        AddressedRequest, AsyncEngineContext, Context, ManyIn, Operator, PushRouter, RouterMode,
        ServerStreamingEngine, StreamingDispatch, context::Controller,
    },
    storage::kv::Selector,
    traits::DistributedRuntimeProvider,
};
use tokio::sync::{Notify, mpsc, watch};

use super::*;
use crate::{
    http::service::metrics::Metrics,
    kv_router::{RoutingLoadContext, routing_host::request_guard::RouteObservation},
    local_model::runtime_config::ModelRuntimeConfig,
    lora::{LoraReplicaConfig, LoraRoutingTable, LoraStateTracker},
    migration::Migration,
    protocols::common::{
        extensions::{AgentContextBuilder, SESSION_AFFINITY_CONTEXT_KEY, SessionAffinityId},
        preprocessor::{MmRoutingInfo, RoutingHints},
        timing::RequestTracker,
    },
};

fn request() -> PreprocessedRequest {
    PreprocessedRequest::builder()
        .model("test".to_string())
        .token_ids(vec![1])
        .stop_conditions(Default::default())
        .sampling_options(Default::default())
        .output_options(Default::default())
        .build()
        .unwrap()
}

async fn test_load_context(client: &Client) -> Arc<RoutingLoadContext> {
    RoutingLoadContext::start(
        client.clone(),
        crate::kv_router::RouterLoadSource::Aggregated,
        crate::discovery::LoadThresholdHandle::new(Default::default()),
        &client.endpoint.drt().child_token(),
        None,
    )
    .await
    .unwrap()
}

#[test]
fn classify_response_item_separates_terminal_failures_from_healthy_frames() {
    let outcome =
        |output: &LLMEngineOutput| classify_response_item(&Annotated::from_data(output.clone()));

    let mut output = LLMEngineOutput::default();
    assert!(matches!(outcome(&output), ResponseItemOutcome::Healthy));

    // Carries only a bare message, so migration has nothing to act on: drain for a typed error.
    output.finish_reason = Some(FinishReason::Error("decode failed".to_string()));
    assert!(matches!(
        outcome(&output),
        ResponseItemOutcome::DrainableTerminal
    ));

    output.finish_reason = Some(FinishReason::Cancelled);
    assert!(matches!(
        outcome(&output),
        ResponseItemOutcome::DrainableTerminal
    ));

    output.finish_reason = Some(FinishReason::Length);
    assert!(matches!(outcome(&output), ResponseItemOutcome::Healthy));
}

#[test]
fn selector_state_remains_owned_by_the_scheduler_actor() {
    fn assert_send_sync<T: Send + Sync>() {}

    assert_send_sync::<RoutingHost>();
}

#[test]
fn builtin_policies_declare_capabilities() {
    for mode in [RouterMode::RoundRobin, RouterMode::Random] {
        let selector = BuiltinWorkerSelector::new(mode).unwrap();
        assert_eq!(selector.required_worker_inputs(), WorkerInputs::NONE);
    }
    for mode in [RouterMode::PowerOfTwoChoices, RouterMode::LeastLoaded] {
        let selector = BuiltinWorkerSelector::new(mode).unwrap();
        assert_eq!(selector.required_worker_inputs(), WorkerInputs::OCCUPANCY);
    }
}

#[tokio::test]
async fn builtin_host_constructs_only_declared_capabilities() {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace("builtin-capabilities".to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    let load_context = test_load_context(&client).await;
    let inner = PushRouter::from_client(client.clone(), RouterMode::RoundRobin)
        .await
        .unwrap();
    let host = RoutingHost::new_builtin(inner, load_context.clone()).unwrap();

    assert_eq!(host.required_worker_inputs(), WorkerInputs::NONE);
    assert!(host.hosted_occupancy.is_none());

    drop(host);

    client.override_discovered_instances(vec![1, 2]);
    client.override_instance_avail(vec![1, 2]);
    let inner = PushRouter::from_client(client, RouterMode::PowerOfTwoChoices)
        .await
        .unwrap();
    let host = RoutingHost::new_builtin(inner, load_context).unwrap();
    let RoutingPolicy::Builtin(selector) = &host.policy else {
        unreachable!()
    };
    assert_eq!(host.required_worker_inputs(), WorkerInputs::OCCUPANCY);
    assert!(host.hosted_occupancy.is_some());
    let selection = host
        .hosted_occupancy
        .as_ref()
        .unwrap()
        .select_and_reserve(&host.inner, selector, Some(1))
        .unwrap();
    assert_eq!(selection.worker_id, 1);
    assert_eq!(selection.occupancy, 1);
    assert_eq!(host.inner.occupancy_for_test(1), 1);
    let mut guard: RequestGuard = RequestGuard::new_builtin(
        Arc::clone(&host.request_metrics),
        selection.worker_id,
        Some(selection.reservation),
        None,
        &request(),
    );
    assert_eq!(guard.retarget_worker(1), Some(1));
    guard.abort().await;
    assert_eq!(host.inner.occupancy_for_test(1), 0);

    drop(host);
    runtime.shutdown();
}

#[tokio::test]
async fn builtin_occupancy_selection_uses_all_selectable_workers() {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace("builtin-occupancy-workers".to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    let load_context = test_load_context(&client).await;
    let inner = PushRouter::from_client(client.clone(), RouterMode::LeastLoaded)
        .await
        .unwrap();
    client.override_discovered_instances(vec![1, 2]);
    client.override_instance_avail(vec![1, 2]);
    let host = RoutingHost::new_builtin(inner, load_context).unwrap();
    let RoutingPolicy::Builtin(selector) = &host.policy else {
        unreachable!()
    };
    let occupancy = host.hosted_occupancy.as_ref().unwrap();
    let first = occupancy
        .select_and_reserve(&host.inner, selector, Some(1))
        .unwrap();
    let second = occupancy
        .select_and_reserve(&host.inner, selector, None)
        .unwrap();

    assert_eq!(first.worker_id, 1);
    assert_eq!(second.worker_id, 2);
    assert_eq!(second.candidate_count, 2);
    assert_eq!(second.occupancy, 1);

    drop(second);
    drop(first);
    drop(host);
    runtime.shutdown();
}

#[tokio::test]
async fn builtin_direct_without_worker_is_invalid_argument() {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace("builtin-direct-invalid-argument".to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    let load_context = test_load_context(&client).await;
    let inner = PushRouter::from_client(client, RouterMode::Direct)
        .await
        .unwrap();
    let (host, _) = builtin_host_with_affinity(
        inner,
        load_context,
        crate::session_affinity::SessionAffinityMode::Hard,
    );

    let error = host
        .generate(affinity_request("direct-unbound", None))
        .await
        .unwrap_err();
    assert!(match_error_chain(
        error.as_ref(),
        &[ErrorType::InvalidArgument],
        &[]
    ));

    drop(host);
    runtime.shutdown();
}

#[tokio::test]
async fn builtin_direct_uses_bound_soft_affinity_as_exact_target() {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace("builtin-direct-soft-affinity".to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    endpoint.register_endpoint_instance().await.unwrap();
    let worker_id = client.wait_for_instances().await.unwrap()[0].id();
    let load_context = test_load_context(&client).await;
    let dispatch = Arc::new(CompletedBuiltinDispatch::default());
    let inner = PushRouter::from_client_with_dispatch(
        client,
        RouterMode::Direct,
        Arc::clone(&dispatch) as Arc<dyn StreamingDispatch<_, _>>,
    )
    .await
    .unwrap();
    let (host, _) = builtin_host_with_affinity(
        inner,
        load_context,
        crate::session_affinity::SessionAffinityMode::Soft,
    );
    let session_id = SessionAffinityId::new("direct-soft-bound");
    bind_affinity_target(&host, &session_id, AffinityTarget::worker(worker_id)).await;

    let mut stream = host
        .generate(affinity_request("direct-soft-bound", None))
        .await
        .unwrap();
    while stream.next().await.is_some() {}

    assert_eq!(dispatch.worker_ids.lock().unwrap().as_slice(), &[worker_id]);
    drop(host);
    runtime.shutdown();
}

#[derive(Default)]
struct CompletedBuiltinDispatch {
    worker_ids: Mutex<Vec<u64>>,
}

#[async_trait]
impl StreamingDispatch<PreprocessedRequest, Annotated<LLMEngineOutput>>
    for CompletedBuiltinDispatch
{
    async fn generate(
        &self,
        request: SingleIn<AddressedRequest<PreprocessedRequest>>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        let (addressed, context) = request.transfer(());
        let (_, _, instance) = addressed.into_parts();
        self.worker_ids
            .lock()
            .unwrap()
            .push(instance.expect("selected worker instance").id());
        let output = Annotated::from_data(LLMEngineOutput {
            finish_reason: Some(FinishReason::Stop),
            ..Default::default()
        });
        Ok(ResponseStream::new(
            Box::pin(stream::once(async move { output })),
            context.context(),
        ))
    }

    async fn generate_bidirectional(
        &self,
        _instance: Instance,
        _address: String,
        _input: ManyIn<PreprocessedRequest>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        unreachable!("the routing host dispatches unary requests")
    }
}

#[tokio::test]
#[serial_test::serial]
async fn builtin_hard_affinity_ignores_local_inhibition() {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace("builtin-direct-local-inhibition".to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    let load_context = test_load_context(&client).await;
    endpoint.register_endpoint_instance().await.unwrap();
    let worker_id = client.wait_for_instances().await.unwrap()[0].id();
    let dispatch = Arc::new(CompletedBuiltinDispatch::default());
    let inner = PushRouter::from_client_with_dispatch(
        client.clone(),
        RouterMode::Direct,
        Arc::clone(&dispatch) as Arc<dyn StreamingDispatch<_, _>>,
    )
    .await
    .unwrap();
    let (host, affinity) = builtin_host_with_affinity(
        inner,
        load_context,
        crate::session_affinity::SessionAffinityMode::Hard,
    );

    let session_id = SessionAffinityId::new("local-inhibition");
    bind_affinity_target(&host, &session_id, AffinityTarget::worker(worker_id)).await;

    client.report_instance_down(worker_id);
    assert!(client.instance_ids().contains(&worker_id));
    assert!(!client.instance_ids_avail().contains(&worker_id));

    let mut stream = host
        .generate(affinity_request("local-inhibition", None))
        .await
        .unwrap();
    while stream.next().await.is_some() {}
    assert_eq!(dispatch.worker_ids.lock().unwrap().as_slice(), &[worker_id]);
    assert_eq!(
        affinity.query_target(&session_id, None).unwrap(),
        Some(AffinityTarget::worker(worker_id))
    );

    drop(host);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn builtin_lora_keeps_separate_selection_and_cleanup() {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace("builtin-lora-capability".to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    let load_context = test_load_context(&client).await;
    endpoint.register_endpoint_instance().await.unwrap();
    let worker_id = client.wait_for_instances().await.unwrap()[0].id();
    let stale_worker = worker_id.wrapping_add(1);
    client.override_instance_avail(vec![stale_worker, worker_id]);
    let dispatch = Arc::new(CompletedBuiltinDispatch::default());
    let inner = PushRouter::from_client_with_dispatch(
        client,
        RouterMode::RoundRobin,
        Arc::clone(&dispatch) as Arc<dyn StreamingDispatch<_, _>>,
    )
    .await
    .unwrap();
    let routing_table = LoraRoutingTable::new();
    routing_table.update_allocation(
        "adapter".to_string(),
        LoraReplicaConfig {
            lora_name: "adapter".to_string(),
            replica_factor: 1,
            replica_set: vec![WorkerWithDpRank::new(worker_id, 0)],
            updated_at: Instant::now(),
            is_active: true,
        },
    );
    let filter = Arc::new(LoraFilter::new(routing_table, LoraStateTracker::new()));
    let estimator = Arc::new(LoadEstimator::new());
    let host = RoutingHost::new_builtin_with_capabilities(
        inner,
        load_context,
        None,
        Some((filter, Arc::clone(&estimator))),
    )
    .unwrap();

    let mut base_stream = host.generate(Context::new(request())).await.unwrap();
    while base_stream.next().await.is_some() {}

    let mut adapter_request = request();
    adapter_request.routing_mut().lora_name = Some("adapter".to_string());
    adapter_request.routing_mut().backend_instance_id = Some(stale_worker);
    let mut adapter_stream = host.generate(Context::new(adapter_request)).await.unwrap();
    assert_eq!(estimator.get_inflight_counts().get("adapter"), Some(&1));
    while adapter_stream.next().await.is_some() {}

    assert_eq!(dispatch.worker_ids.lock().unwrap().len(), 2);
    assert_eq!(dispatch.worker_ids.lock().unwrap()[1], worker_id);
    assert!(!estimator.get_inflight_counts().contains_key("adapter"));

    drop(host);
    runtime.shutdown();
}

fn affinity_request(
    session_id: &str,
    explicit_worker: Option<u64>,
) -> SingleIn<PreprocessedRequest> {
    let mut content = request();
    content.routing_mut().backend_instance_id = explicit_worker;
    let mut request = Context::new(content);
    request.insert(
        SESSION_AFFINITY_CONTEXT_KEY,
        SessionAffinityId::new(session_id),
    );
    request
}

#[tokio::test]
#[serial_test::serial]
async fn builtin_affinity_uses_common_host_for_every_policy() {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let component = distributed
        .namespace("builtin-affinity-lifecycle".to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap();

    for (index, mode) in [
        RouterMode::RoundRobin,
        RouterMode::Random,
        RouterMode::PowerOfTwoChoices,
        RouterMode::LeastLoaded,
        RouterMode::DeviceAwareWeighted,
        RouterMode::Direct,
    ]
    .into_iter()
    .enumerate()
    {
        let endpoint = component.endpoint(format!("mode-{index}"));
        let client = endpoint.client().await.unwrap();
        let load_context = test_load_context(&client).await;
        endpoint.register_endpoint_instance().await.unwrap();
        let worker_id = client.wait_for_instances().await.unwrap()[0].id();
        let dispatch = Arc::new(CompletedBuiltinDispatch::default());
        let inner = PushRouter::from_client_with_dispatch(
            client,
            mode,
            Arc::clone(&dispatch) as Arc<dyn StreamingDispatch<_, _>>,
        )
        .await
        .unwrap();
        let (host, affinity) = builtin_host_with_affinity(
            inner,
            load_context,
            crate::session_affinity::SessionAffinityMode::Hard,
        );
        let session_id = format!("session-{index}");
        let affinity_id = SessionAffinityId::new(session_id.clone());
        let explicit_worker = (mode == RouterMode::Direct).then_some(worker_id);

        let mut first = host
            .generate(affinity_request(&session_id, explicit_worker))
            .await
            .unwrap();
        while first.next().await.is_some() {}
        assert_eq!(
            affinity.query_target(&affinity_id, None).unwrap(),
            Some(AffinityTarget::worker(worker_id))
        );

        let mut second = host
            .generate(affinity_request(&session_id, explicit_worker))
            .await
            .unwrap();
        while second.next().await.is_some() {}
        assert_eq!(
            dispatch.worker_ids.lock().unwrap().as_slice(),
            &[worker_id; 2]
        );
    }

    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn builtin_hard_affinity_ignores_overload_while_soft_affinity_falls_back() {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace("builtin-affinity-modes".to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    endpoint.register_endpoint_instance().await.unwrap();
    let worker_id = client.wait_for_instances().await.unwrap()[0].id();
    let fallback_worker_id = worker_id + 1;
    client.override_instance_avail(vec![worker_id, fallback_worker_id]);
    client.set_overloaded_instances(&[worker_id]);
    let load_context = test_load_context(&client).await;
    let inner = PushRouter::from_client(client, RouterMode::RoundRobin)
        .await
        .unwrap();
    let host = RoutingHost::new_builtin(inner, load_context).unwrap();
    let request = Context::new(request());

    let hard = host
        .select_hosted_worker(&request, Some(AffinityTarget::worker(worker_id)), None)
        .unwrap();
    let soft = host
        .select_hosted_worker(&request, None, Some(AffinityTarget::worker(worker_id)))
        .unwrap();

    assert_eq!(hard.initial_worker, worker_id);
    assert_eq!(
        hard.target_constraint,
        Some(AffinityTarget::worker(worker_id))
    );
    assert_eq!(soft.initial_worker, fallback_worker_id);
    assert_eq!(soft.target_constraint, None);

    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn builtin_direct_distinguishes_unknown_requests_from_stale_affinity() {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace("builtin-direct-worker-loss".to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    let load_context = test_load_context(&client).await;
    endpoint.register_endpoint_instance().await.unwrap();
    let real_worker = client.wait_for_instances().await.unwrap()[0].id();
    let stale_worker = real_worker.wrapping_add(1);
    client.override_instance_avail(vec![stale_worker, real_worker]);
    let dispatch = Arc::new(CompletedBuiltinDispatch::default());
    let inner = PushRouter::from_client_with_dispatch(
        client,
        RouterMode::Direct,
        Arc::clone(&dispatch) as Arc<dyn StreamingDispatch<_, _>>,
    )
    .await
    .unwrap();
    let (host, affinity) = builtin_host_with_affinity(
        inner,
        load_context,
        crate::session_affinity::SessionAffinityMode::Hard,
    );

    let mut standalone = request();
    standalone.routing_mut().backend_instance_id = Some(stale_worker);
    let error = host.generate(Context::new(standalone)).await.unwrap_err();
    assert!(match_error_chain(
        error.as_ref(),
        &[ErrorType::InvalidArgument],
        &[]
    ));
    assert!(dispatch.worker_ids.lock().unwrap().is_empty());

    let session_id = SessionAffinityId::new("direct-affinity");
    bind_affinity_target(&host, &session_id, AffinityTarget::worker(stale_worker)).await;
    let error = host
        .generate(affinity_request("direct-affinity", None))
        .await
        .unwrap_err();
    assert!(!match_error_chain(
        error.as_ref(),
        &[ErrorType::InvalidArgument],
        &[]
    ));
    assert!(dispatch.worker_ids.lock().unwrap().is_empty());
    assert_eq!(affinity.query_target(&session_id, None).unwrap(), None);

    let mut valid = request();
    valid.routing_mut().backend_instance_id = Some(real_worker);
    let mut stream = host.generate(Context::new(valid)).await.unwrap();
    while stream.next().await.is_some() {}
    assert_eq!(
        dispatch.worker_ids.lock().unwrap().as_slice(),
        &[real_worker]
    );

    drop(host);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn terminal_item_does_not_skip_transport_eof() {
    let (router, runtime) = router(None).await;
    let inputs = router.required_worker_inputs();
    assert!(inputs.contains(WorkerInputs::CACHE));
    assert!(inputs.contains(WorkerInputs::LOAD));
    let context = Context::new(()).context();
    let drained = Arc::new(AtomicBool::new(false));
    let source_drained = Arc::clone(&drained);
    let source = ResponseStream::new(
        Box::pin(async_stream::stream! {
            yield Annotated::from_data(LLMEngineOutput {
                finish_reason: Some(FinishReason::Stop),
                ..Default::default()
            });
            source_drained.store(true, Ordering::Release);
        }),
        Arc::clone(&context),
    );
    let guard = RequestGuard::new_kv_with_cleanup(
        Arc::clone(&router.request_metrics),
        KvRequestCleanup::new(
            Arc::clone(router.kv_router()),
            "terminal-drain".to_string(),
            WorkerWithDpRank::from_worker_id(0),
            None,
        ),
        &request(),
        Some(RouteObservation {
            prompt_tokens: 1,
            best_router_tokens: 0,
            selected_router_tokens: 0,
        }),
        None,
    );
    let monitored = monitor_response_stream(source, context, guard);
    tokio::pin!(monitored);

    assert!(monitored.next().await.is_some());
    assert!(monitored.next().await.is_none());
    assert!(drained.load(Ordering::Acquire));

    drop(router);
    runtime.shutdown();
}

struct KvHitSnapshot {
    best: u64,
    selected: u64,
    reused: u64,
}

fn kv_hit_snapshot(metrics: &crate::kv_router::metrics::RouterRequestMetrics) -> KvHitSnapshot {
    // The fixture request carries no tracker, so the guard labels it `aggregated`.
    let phase = RequestPhase::Aggregated.as_str();
    let model = "test";
    let counter = |h: &prometheus::IntCounterVec| h.with_label_values(&[phase, model]).get();
    KvHitSnapshot {
        best: counter(&metrics.kv_best_eligible_cached_prefix_tokens),
        selected: counter(&metrics.kv_selected_cached_prefix_tokens),
        reused: counter(&metrics.kv_worker_reused_tokens),
    }
}

/// Drive one tracked attempt through the real guard and return the metric deltas.
async fn run_kv_hit_attempt(final_frame: LLMEngineOutput) -> (KvHitSnapshot, KvHitSnapshot) {
    let (router, runtime) = router(None).await;
    let metrics = Arc::clone(&router.request_metrics);
    let before = kv_hit_snapshot(&metrics);
    let context = Context::new(()).context();
    let source = ResponseStream::new(
        Box::pin(async_stream::stream! {
            yield Annotated::from_data(LLMEngineOutput {
                token_ids: vec![7],
                ..Default::default()
            });
            yield Annotated::from_data(final_frame);
        }),
        Arc::clone(&context),
    );
    let guard = RequestGuard::new_kv_with_cleanup(
        Arc::clone(&metrics),
        KvRequestCleanup::new(
            Arc::clone(router.kv_router()),
            "kv-hit-attempt".to_string(),
            WorkerWithDpRank::from_worker_id(0),
            None,
        ),
        &request(),
        Some(RouteObservation {
            prompt_tokens: 100,
            best_router_tokens: 75,
            selected_router_tokens: 60,
        }),
        None,
    );
    let monitored = monitor_response_stream(source, context, guard);
    tokio::pin!(monitored);
    while monitored.next().await.is_some() {}
    let after = kv_hit_snapshot(&metrics);
    drop(router);
    runtime.shutdown();
    (before, after)
}

#[tokio::test]
#[serial_test::serial]
async fn kv_cache_hit_complete_attempt_records_every_stage_once() {
    let (before, after) = run_kv_hit_attempt(LLMEngineOutput {
        finish_reason: Some(FinishReason::Stop),
        engine_data: Some(serde_json::json!({
            "kv_cache_hit": {
                "prompt_tokens": 100,
                "reused_tokens": 85,
            }
        })),
        ..Default::default()
    })
    .await;
    assert_eq!(after.best - before.best, 75);
    assert_eq!(after.selected - before.selected, 60);
    assert_eq!(after.reused - before.reused, 85);
}

#[tokio::test]
#[serial_test::serial]
async fn kv_cache_hit_attempt_without_worker_report_contributes_zero() {
    let (before, after) = run_kv_hit_attempt(LLMEngineOutput {
        finish_reason: Some(FinishReason::Stop),
        ..Default::default()
    })
    .await;
    assert_eq!(after.best - before.best, 75);
    assert_eq!(after.selected - before.selected, 60);
    assert_eq!(after.reused - before.reused, 0);
}

#[tokio::test]
#[serial_test::serial]
async fn kv_cache_hit_cancelled_attempt_keeps_reported_reuse() {
    let (before, after) = run_kv_hit_attempt(LLMEngineOutput {
        finish_reason: Some(FinishReason::Cancelled),
        engine_data: Some(serde_json::json!({
            "kv_cache_hit": {
                "prompt_tokens": 100,
                "reused_tokens": 85,
            }
        })),
        ..Default::default()
    })
    .await;
    assert_eq!(after.best - before.best, 75);
    assert_eq!(after.selected - before.selected, 60);
    assert_eq!(after.reused - before.reused, 85);
}

fn cancelled_frame() -> Annotated<LLMEngineOutput> {
    Annotated::from_data(LLMEngineOutput {
        finish_reason: Some(FinishReason::Cancelled),
        ..Default::default()
    })
}

#[rstest::rstest]
#[case(0)]
#[case(1)]
#[tokio::test]
#[serial_test::serial]
async fn kv_cache_hit_counts_immediately_once_in_selection_phase(#[case] reused: u64) {
    let (router, runtime) = router(None).await;
    let metrics = crate::kv_router::metrics::RouterRequestMetrics::for_test(
        &dynamo_runtime::MetricsRegistry::new(),
    );
    let tracker = Arc::new(RequestTracker::new());
    let permit = tracker.set_phase(RequestPhase::Prefill).await;
    let mut req = request();
    req.tracker = Some(tracker.clone());
    let mut guard = RequestGuard::new_kv_with_cleanup(
        metrics.clone(),
        KvRequestCleanup::new(
            Arc::clone(router.kv_router()),
            "phase-test".to_string(),
            WorkerWithDpRank::from_worker_id(0),
            None,
        ),
        &req,
        Some(RouteObservation {
            prompt_tokens: 1,
            best_router_tokens: 1,
            selected_router_tokens: 1,
        }),
        None,
    );
    drop(permit);
    let _permit = tracker.set_phase(RequestPhase::Decode).await;
    guard
        .on_item(&Annotated::from_data(LLMEngineOutput {
            engine_data: Some(serde_json::json!({"kv_cache_hit": {
                "prompt_tokens": 1, "reused_tokens": reused
            }})),
            ..Default::default()
        }))
        .await;
    assert_eq!(
        metrics
            .kv_worker_reused_tokens
            .with_label_values(&["prefill", "test"])
            .get(),
        reused
    );
    guard
        .on_item(&Annotated::from_data(LLMEngineOutput {
            engine_data: Some(serde_json::json!({"kv_cache_hit": {
                "prompt_tokens": 1, "reused_tokens": 9
            }})),
            ..Default::default()
        }))
        .await;
    guard.abort().await;
    drop(guard);
    assert_eq!(
        metrics
            .kv_worker_reused_tokens
            .with_label_values(&["prefill", "test"])
            .get(),
        reused
    );
    assert_eq!(
        metrics
            .kv_worker_reused_tokens
            .with_label_values(&["decode", "test"])
            .get(),
        0
    );
    drop(router);
    runtime.shutdown();
}

fn engine_shutdown_frame() -> Annotated<LLMEngineOutput> {
    Annotated {
        data: None,
        id: None,
        event: Some("error".to_string()),
        comment: None,
        error: Some(
            DynamoError::builder()
                .error_type(ErrorType::Backend(BackendError::EngineShutdown))
                .message("engine is shutting down")
                .build(),
        ),
    }
}

fn is_engine_shutdown(item: &Annotated<LLMEngineOutput>) -> bool {
    item.error.as_ref().is_some_and(|error| {
        match_error_chain(
            error,
            &[ErrorType::Backend(BackendError::EngineShutdown)],
            &[],
        )
    })
}

/// A trailing `EngineShutdown` error must reach migration, and the `Cancelled` frame that
/// preceded it must not be forwarded once it does.
#[tokio::test]
#[serial_test::serial]
async fn shutdown_cancellation_drains_trailing_engine_shutdown_error() {
    let (router, runtime) = router(None).await;
    let context = Context::new(()).context();
    // Reaching this at all is what the old code could not do.
    let polled_past_cancel = Arc::new(AtomicBool::new(false));
    let source_polled = Arc::clone(&polled_past_cancel);
    let source = ResponseStream::new(
        Box::pin(async_stream::stream! {
            yield cancelled_frame();
            source_polled.store(true, Ordering::Release);
            yield engine_shutdown_frame();
        }),
        Arc::clone(&context),
    );
    let guard = RequestGuard::new_kv_with_cleanup(
        Arc::clone(&router.request_metrics),
        KvRequestCleanup::new(
            Arc::clone(router.kv_router()),
            "shutdown-drain".to_string(),
            WorkerWithDpRank::from_worker_id(0),
            None,
        ),
        &request(),
        None,
        None,
    );
    let monitored = monitor_response_stream(source, context, guard);
    tokio::pin!(monitored);

    let shutdown = monitored
        .next()
        .await
        .expect("the trailing engine-shutdown error must not be swallowed");
    assert!(
        is_engine_shutdown(&shutdown),
        "migration keys off the typed error, got {:?}",
        shutdown.error
    );
    assert!(monitored.next().await.is_none());
    assert!(polled_past_cancel.load(Ordering::Acquire));

    drop(router);
    runtime.shutdown();
}

/// A client cancel arrives through the context, not as a frame, so it must still preempt.
#[tokio::test]
#[serial_test::serial]
async fn client_cancellation_still_ends_stream_without_draining() {
    let (router, runtime) = router(None).await;
    let controller = Controller::new("client-cancelled-drain".to_string());
    controller.stop();
    let cancelled_request = Context::with_controller((), controller);
    let context = cancelled_request.context();
    let source_polled = Arc::new(AtomicBool::new(false));
    let polled = Arc::clone(&source_polled);
    let source = ResponseStream::new(
        Box::pin(async_stream::stream! {
            polled.store(true, Ordering::Release);
            yield cancelled_frame();
            yield engine_shutdown_frame();
        }),
        Arc::clone(&context),
    );
    let guard = RequestGuard::new_kv_with_cleanup(
        Arc::clone(&router.request_metrics),
        KvRequestCleanup::new(
            Arc::clone(router.kv_router()),
            "client-cancelled-drain".to_string(),
            WorkerWithDpRank::from_worker_id(0),
            None,
        ),
        &request(),
        None,
        None,
    );
    let monitored = monitor_response_stream(source, context, guard);
    tokio::pin!(monitored);

    assert!(monitored.next().await.is_none());
    assert!(
        !source_polled.load(Ordering::Acquire),
        "a client cancel must not read further worker frames"
    );

    drop(router);
    runtime.shutdown();
}

/// A worker that sends a terminal frame and then stops talking without closing the transport
/// must not hold the request open; the drain is bounded.
#[tokio::test]
#[serial_test::serial]
async fn drain_without_trailing_error_gives_up_at_the_deadline() {
    let (router, runtime) = router(None).await;
    // Auto-advances the drain deadline once the task idles, so this costs no wall clock.
    tokio::time::pause();
    let context = Context::new(()).context();
    let source = ResponseStream::new(
        Box::pin(async_stream::stream! {
            yield cancelled_frame();
            // Never ends and never sends the error: the transport is wedged open.
            std::future::pending::<()>().await;
        }),
        Arc::clone(&context),
    );
    let guard = RequestGuard::new_kv_with_cleanup(
        Arc::clone(&router.request_metrics),
        KvRequestCleanup::new(
            Arc::clone(router.kv_router()),
            "shutdown-drain-deadline".to_string(),
            WorkerWithDpRank::from_worker_id(0),
            None,
        ),
        &request(),
        None,
        None,
    );
    let monitored = monitor_response_stream(source, context, guard);
    tokio::pin!(monitored);

    // Generous relative to DRAIN_TIMEOUT: the assertion is that the drain is bounded at all.
    let deadline = DRAIN_TIMEOUT * 4;
    let item = tokio::time::timeout(deadline, monitored.next())
        .await
        .expect("the drain must give up at DRAIN_TIMEOUT")
        .expect("the withheld frame must be released once the drain gives up");
    assert!(matches!(
        item.data
            .as_ref()
            .and_then(|data| data.finish_reason.as_ref()),
        Some(FinishReason::Cancelled)
    ));
    assert!(
        tokio::time::timeout(deadline, monitored.next())
            .await
            .expect("the stream must end after the drain gives up")
            .is_none()
    );

    drop(router);
    runtime.shutdown();
}

/// The drain window must actually be armed for DRAIN_TIMEOUT, not merely exist. A worker
/// that sends its typed error a scheduler tick after the `Cancelled` frame -- well inside
/// the window -- must still have that error reach migration, and the withheld frame must
/// not be published ahead of it. Removing the `drain_deadline` reset leaves the deadline at
/// its already-elapsed initial value, which ends the stream on the `Cancelled` frame and
/// fails this test.
#[tokio::test]
#[serial_test::serial]
async fn trailing_error_within_the_drain_window_still_reaches_migration() {
    let (router, runtime) = router(None).await;
    tokio::time::pause();
    let context = Context::new(()).context();
    let source = ResponseStream::new(
        Box::pin(async_stream::stream! {
            yield cancelled_frame();
            // Comfortably inside the window: the drain must still be open here.
            tokio::time::sleep(DRAIN_TIMEOUT / 2).await;
            yield engine_shutdown_frame();
        }),
        Arc::clone(&context),
    );
    let guard = RequestGuard::new_kv_with_cleanup(
        Arc::clone(&router.request_metrics),
        KvRequestCleanup::new(
            Arc::clone(router.kv_router()),
            "drain-window-armed".to_string(),
            WorkerWithDpRank::from_worker_id(0),
            None,
        ),
        &request(),
        None,
        None,
    );
    let monitored = monitor_response_stream(source, context, guard);
    tokio::pin!(monitored);

    let first = monitored
        .next()
        .await
        .expect("the trailing engine-shutdown error must survive the drain window");
    assert!(
        is_engine_shutdown(&first),
        "migration keys off the typed error, and the withheld cancel frame must not \
         precede it; got {first:?}"
    );
    assert!(
        monitored.next().await.is_none(),
        "the shutdown error is the whole response"
    );

    drop(router);
    runtime.shutdown();
}

/// `biased` polls the stream arm before the drain deadline, so a worker that stays
/// continuously ready with terminal frames must not be able to postpone the deadline
/// forever. The drain is armed once and checked against the clock on every frame.
#[tokio::test]
#[serial_test::serial]
async fn always_ready_terminals_cannot_starve_the_drain_deadline() {
    let (router, runtime) = router(None).await;
    tokio::time::pause();
    let context = Context::new(()).context();
    let source = ResponseStream::new(
        Box::pin(async_stream::stream! {
            loop {
                yield cancelled_frame();
            }
        }),
        Arc::clone(&context),
    );
    let guard = RequestGuard::new_kv_with_cleanup(
        Arc::clone(&router.request_metrics),
        KvRequestCleanup::new(
            Arc::clone(router.kv_router()),
            "starvation-guard".to_string(),
            WorkerWithDpRank::from_worker_id(0),
            None,
        ),
        &request(),
        None,
        None,
    );
    let monitored = monitor_response_stream(source, context, guard);
    tokio::pin!(monitored);

    // Paused time does not auto-advance while the source is always ready, so the consumer
    // moves the clock past the deadline itself.
    let mut frames = 0u32;
    while monitored.next().await.is_some() {
        frames += 1;
        if frames == 8 {
            tokio::time::advance(DRAIN_TIMEOUT + Duration::from_millis(1)).await;
        }
        assert!(
            frames < 1_000,
            "the drain deadline was starved by an always-ready terminal stream"
        );
    }

    drop(router);
    runtime.shutdown();
}

/// The selection future is nested inside every request future (deeper still
/// on the disaggregated path). A debug build once overflowed a worker stack
/// when this grew and `await_with_cleanup_policy` held it by value; keep it
/// small enough that nesting stays cheap.
#[tokio::test]
#[serial_test::serial]
async fn kv_selection_future_stays_small() {
    let (router, _runtime) = router_with_workers(None, &[1]).await;
    let request = Context::with_id_and_metadata(
        request(),
        "selection-future-size".to_string(),
        Default::default(),
    );
    let budget = CleanupBudget::default();
    let future = router.select_with_affinity(&request, RequestPhase::Aggregated, false, &budget);
    let size = std::mem::size_of_val(&future);
    drop(future);
    eprintln!("kv selection future size: {size} bytes");
    assert!(size < 16 * 1024, "kv selection future is {size} bytes");
}

/// Transport EOF ends the drain and releases the booking.
#[tokio::test]
#[serial_test::serial]
async fn shutdown_cancellation_without_trailing_error_still_aborts() {
    let (router, runtime) = router(None).await;
    let context_id = "shutdown-cancel-without-error".to_string();
    let cancelled_request =
        Context::with_id_and_metadata(request(), context_id.clone(), Default::default());
    let (mut selection, _) = router
        .select_with_affinity(
            &cancelled_request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    let cancelled_worker = selection.worker;
    let guard = router
        .track_selection(
            &cancelled_request,
            &mut selection,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    let source = ResponseStream::new(
        Box::pin(stream::once(async { cancelled_frame() })),
        cancelled_request.context().clone(),
    );
    let monitored = monitor_response_stream(source, cancelled_request.context().clone(), guard);
    tokio::pin!(monitored);

    // Below DRAIN_TIMEOUT so only transport EOF can satisfy these assertions: at 10s the
    // test would also pass if the drain deadline, not EOF, had ended the stream.
    let eof_bound = DRAIN_TIMEOUT / 2;
    let item = tokio::time::timeout(eof_bound, monitored.next())
        .await
        .expect("the drain must end at transport EOF, not block on a trailing error")
        .expect("the cancelled frame must be yielded once EOF proves it was the last");
    assert!(matches!(
        item.data
            .as_ref()
            .and_then(|data| data.finish_reason.as_ref()),
        Some(FinishReason::Cancelled)
    ));
    assert!(
        tokio::time::timeout(eof_bound, monitored.next())
            .await
            .expect("the drain must end at transport EOF, not wait for a trailing error")
            .is_none()
    );

    let retry_request =
        Context::with_id_and_metadata(request(), context_id.clone(), Default::default());
    let (mut retry_selection, _) = router
        .select_with_affinity(
            &retry_request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    assert_eq!(retry_selection.worker, cancelled_worker);
    let mut retry_guard = router
        .track_selection(
            &retry_request,
            &mut retry_selection,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .expect("the booking must be released when the drain reaches transport EOF");
    retry_guard.abort().await;

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn stream_failure_releases_booking_before_error_is_observable() {
    let (router, runtime) = router(None).await;
    let context_id = "stream-failure-cleanup".to_string();
    let failed_request =
        Context::with_id_and_metadata(request(), context_id.clone(), Default::default());
    let (mut failed_selection, _) = router
        .select_with_affinity(
            &failed_request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    let failed_worker = failed_selection.worker;
    let failed_guard = router
        .track_selection(
            &failed_request,
            &mut failed_selection,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    let failure = Annotated {
        data: None,
        id: None,
        event: Some("error".to_string()),
        comment: None,
        error: Some(
            DynamoError::builder()
                .error_type(ErrorType::WorkerOverloaded)
                .message("selected worker is overloaded")
                .build(),
        ),
    };
    let source = ResponseStream::new(
        Box::pin(stream::once(async move { failure })),
        failed_request.context().clone(),
    );
    let monitored = monitor_response_stream(source, failed_request.context().clone(), failed_guard);
    tokio::pin!(monitored);

    let item = monitored.next().await.expect("failed item must be yielded");
    assert!(item.error.is_some());

    // The monitored stream is still suspended at its yield point. Rebooking
    // the same id on the same worker proves cleanup completed before the
    // failure became visible, rather than relying on EOF or Drop cleanup.
    let retry_request =
        Context::with_id_and_metadata(request(), context_id.clone(), Default::default());
    let (mut retry_selection, _) = router
        .select_with_affinity(
            &retry_request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    assert_eq!(retry_selection.worker, failed_worker);
    let mut retry_guard = router
        .track_selection(
            &retry_request,
            &mut retry_selection,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .expect("same-worker booking must be released before yielding the error");
    retry_guard.abort().await;

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn output_block_accounting_tracks_grouped_chunks() {
    for (track_output_blocks, expected_output_tokens, expected_output_blocks) in
        [(true, None, 3), (true, Some(66), 3), (false, None, 0)]
    {
        let config = KvRouterConfig {
            skip_initial_worker_wait: true,
            use_kv_events: false,
            router_track_active_blocks: true,
            router_track_output_blocks: track_output_blocks,
            ..Default::default()
        };
        let (router, runtime) = router_with_config(
            None,
            HashMap::from([(7, ModelRuntimeConfig::default())]),
            config,
        )
        .await;
        let prompt_tokens = 16_usize;
        let chunks = [1_usize, 32, 15];
        let mut input = request();
        input.token_ids = (1..=prompt_tokens as u32).collect::<Vec<_>>().into();
        input.routing = Some(RoutingHints {
            expected_output_tokens,
            ..Default::default()
        });
        let input = Context::new(input);
        let (mut selection, _) = router
            .select_with_affinity(
                &input,
                RequestPhase::Aggregated,
                false,
                &CleanupBudget::default(),
            )
            .await
            .unwrap();
        let mut guard = router
            .track_selection(
                &input,
                &mut selection,
                RequestPhase::Aggregated,
                false,
                &CleanupBudget::default(),
            )
            .await
            .unwrap();
        let mut next_token = prompt_tokens as u32 + 1;
        let initial_loads = router
            .kv_router()
            .get_potential_loads(&[], None, None, None, None)
            .await
            .unwrap();
        let initial_blocks = initial_loads
            .iter()
            .find(|load| load.worker_id == 7 && load.dp_rank == 0)
            .unwrap()
            .potential_decode_blocks;
        for size in &chunks {
            let token_ids = (next_token..next_token + *size as u32).collect();
            next_token += *size as u32;
            guard
                .on_item(&Annotated::from_data(LLMEngineOutput {
                    token_ids,
                    index: Some(0),
                    ..Default::default()
                }))
                .await;
        }
        let loads = router
            .kv_router()
            .get_potential_loads(&[], None, None, None, None)
            .await
            .unwrap();
        let load = loads
            .iter()
            .find(|load| load.worker_id == 7 && load.dp_rank == 0)
            .unwrap();
        assert_eq!(load.active_requests, 1);
        // This single-request fixture has no shared prompt blocks, so decay applies
        // to both prompt and output blocks. Check accounting and OSL propagation here;
        // the sequence tests check the number of local load observations.
        // The last boundary is observed at output length 33; OSL 66 gives 0.5 decay.
        let decay = if expected_output_tokens.is_some() {
            0.5
        } else {
            1.0
        };
        let expected_blocks =
            ((initial_blocks + expected_output_blocks) as f64 * decay).round() as usize;
        println!(
            "track_output_blocks={track_output_blocks} prompt_tokens={prompt_tokens} output_tokens=48 block_size=16 chunks={chunks:?} initial_blocks={initial_blocks} observed_blocks={} expected_blocks={expected_blocks}",
            load.potential_decode_blocks
        );
        let observed_blocks = load.potential_decode_blocks;
        guard.finish().await;
        let loads = router
            .kv_router()
            .get_potential_loads(&[], None, None, None, None)
            .await
            .unwrap();
        let load = loads
            .iter()
            .find(|load| load.worker_id == 7 && load.dp_rank == 0)
            .unwrap();
        assert_eq!(load.active_requests, 0);
        assert_eq!(load.potential_decode_blocks, 0);
        drop(router);
        runtime.shutdown();
        assert_eq!(
            observed_blocks, expected_blocks,
            "incorrect grouped-output accounting with tracking={track_output_blocks}, osl={expected_output_tokens:?}"
        );
    }
}

async fn router(session_affinity_ttl: Option<Duration>) -> (RoutingHost, Runtime) {
    router_with_workers(session_affinity_ttl, &[7]).await
}

async fn router_with_workers(
    session_affinity_ttl: Option<Duration>,
    worker_ids: &[u64],
) -> (RoutingHost, Runtime) {
    let workers = worker_ids
        .iter()
        .copied()
        .map(|worker_id| (worker_id, ModelRuntimeConfig::default()))
        .collect();
    router_with_worker_configs(session_affinity_ttl, workers).await
}

async fn router_with_worker_configs(
    session_affinity_ttl: Option<Duration>,
    workers: HashMap<u64, ModelRuntimeConfig>,
) -> (RoutingHost, Runtime) {
    // Keep the large construction future off debug test stacks.
    Box::pin(router_with_worker_configs_and_classifier(
        session_affinity_ttl,
        workers,
        None::<RecordingClassifier>,
    ))
    .await
}

async fn router_with_classifier(
    classifier: impl RequestClassifier,
    session_affinity_ttl: Option<Duration>,
) -> (RoutingHost, Runtime) {
    router_with_worker_configs_and_classifier(
        session_affinity_ttl,
        HashMap::from([(7, ModelRuntimeConfig::default())]),
        Some(classifier),
    )
    .await
}

async fn router_with_worker_configs_and_classifier(
    session_affinity_ttl: Option<Duration>,
    workers: HashMap<u64, ModelRuntimeConfig>,
    classifier: Option<impl RequestClassifier>,
) -> (RoutingHost, Runtime) {
    let config = KvRouterConfig {
        skip_initial_worker_wait: true,
        use_kv_events: false,
        router_track_active_blocks: false,
        ..Default::default()
    };
    Box::pin(router_with_config_and_classifier(
        session_affinity_ttl,
        workers,
        config,
        classifier,
    ))
    .await
}

async fn router_with_config(
    session_affinity_ttl: Option<Duration>,
    workers: HashMap<u64, ModelRuntimeConfig>,
    config: KvRouterConfig,
) -> (RoutingHost, Runtime) {
    Box::pin(router_with_config_and_classifier(
        session_affinity_ttl,
        workers,
        config,
        None::<RecordingClassifier>,
    ))
    .await
}

async fn router_with_config_and_classifier(
    session_affinity_ttl: Option<Duration>,
    workers: HashMap<u64, ModelRuntimeConfig>,
    config: KvRouterConfig,
    classifier: Option<impl RequestClassifier>,
) -> (RoutingHost, Runtime) {
    let runtime = Runtime::from_current().unwrap();
    // Each runtime has its own in-memory discovery store, so fixture namespaces can repeat.
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let component = distributed
        .namespace("affinity-selection-cancellation".to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap();
    let endpoint = component.endpoint("generate");
    let client = endpoint.client().await.unwrap();
    let worker_ids = workers.keys().copied().collect::<Vec<_>>();
    let (_tx, workers) = watch::channel(workers);
    let mut chooser = KvRouter::new(
        endpoint,
        client.clone(),
        workers,
        None,
        16,
        SelectionPolicySource::Registry,
        Some(config),
        None,
        "decode",
        None,
        false,
        None,
        None,
    )
    .await
    .unwrap();
    if let Some(classifier) = classifier {
        chooser = chooser.with_request_classifier(classifier).unwrap();
    }
    let inner = PushRouter::from_client(client, RouterMode::KV)
        .await
        .unwrap();
    let mut router = RoutingHost::new(inner, Arc::new(chooser), session_affinity_ttl).unwrap();
    // Parallel tests must not share the process-global request counters.
    router.request_metrics = router.request_metrics.with_isolated_counters_for_test();
    router
        .inner
        .client
        .override_discovered_instances(worker_ids.clone());
    router.inner.client.override_instance_avail(worker_ids);
    (router, runtime)
}

/// A dispatch plane that yields before completing, unlike [`CompletedBuiltinDispatch`].
/// A recorded worker id therefore means the dispatch survived, not merely started.
#[derive(Default)]
struct PendingThenCompletedDispatch {
    worker_ids: Mutex<Vec<u64>>,
}

#[async_trait]
impl StreamingDispatch<PreprocessedRequest, Annotated<LLMEngineOutput>>
    for PendingThenCompletedDispatch
{
    async fn generate(
        &self,
        request: SingleIn<AddressedRequest<PreprocessedRequest>>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        tokio::task::yield_now().await;
        let (addressed, context) = request.transfer(());
        let (_, _, instance) = addressed.into_parts();
        self.worker_ids
            .lock()
            .unwrap()
            .push(instance.expect("selected worker instance").id());
        let output = Annotated::from_data(LLMEngineOutput {
            finish_reason: Some(FinishReason::Stop),
            ..Default::default()
        });
        Ok(ResponseStream::new(
            Box::pin(stream::once(async move { output })),
            context.context(),
        ))
    }

    async fn generate_bidirectional(
        &self,
        _instance: Instance,
        _address: String,
        _input: ManyIn<PreprocessedRequest>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        unreachable!("the routing host dispatches unary requests")
    }
}

/// Build a KV routing host whose dispatch plane is the recording mock, so a test
/// can observe whether the selected worker was actually asked to serve the request.
async fn router_with_recorded_dispatch(
    namespace: &str,
) -> (RoutingHost, Arc<PendingThenCompletedDispatch>, u64, Runtime) {
    router_with_recorded_dispatch_and_affinity(namespace, None).await
}

async fn router_with_recorded_dispatch_and_affinity(
    namespace: &str,
    session_affinity_ttl: Option<Duration>,
) -> (RoutingHost, Arc<PendingThenCompletedDispatch>, u64, Runtime) {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace(namespace.to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    endpoint.register_endpoint_instance().await.unwrap();
    let worker_id = client.wait_for_instances().await.unwrap()[0].id();
    let workers = HashMap::from([(worker_id, ModelRuntimeConfig::default())]);
    let (_workers_tx, workers) = watch::channel(workers);
    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,
        crate::kv_router::SelectionPolicySource::Registry,
        Some(config),
        None,
        "decode",
        None,
        false,
        None,
        None,
    )
    .await
    .unwrap();
    let dispatch = Arc::new(PendingThenCompletedDispatch::default());
    let inner = PushRouter::from_client_with_dispatch(
        client,
        RouterMode::KV,
        Arc::clone(&dispatch) as Arc<dyn StreamingDispatch<_, _>>,
    )
    .await
    .unwrap();
    let router = RoutingHost::new(inner, Arc::new(chooser), session_affinity_ttl).unwrap();
    (router, dispatch, worker_id, runtime)
}

/// Select and admit a request in the given phase, then stop its context, leaving
/// the request poised at the dispatch stage with an already-cancelled context.
async fn admitted_request_stopped_before_dispatch(
    router: &RoutingHost,
    context_id: &str,
    phase: RequestPhase,
) -> (SingleIn<PreprocessedRequest>, WorkerSelection, RequestGuard) {
    let tracker = Arc::new(RequestTracker::new());
    // The permit only serializes further phase changes; the phase itself sticks
    // once set, and `dispatch_selection` reads it back off the tracker.
    drop(tracker.set_phase(phase).await);
    let mut input = request();
    input.tracker = Some(Arc::clone(&tracker));
    let request = Context::with_id_and_metadata(input, context_id.to_string(), Default::default());

    let (mut selection, _) = router
        .select_with_affinity(&request, phase, false, &CleanupBudget::default())
        .await
        .unwrap();
    let guard = router
        .track_selection(
            &request,
            &mut selection,
            phase,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();

    // The client went away while the request was in flight upstream — for the
    // decode leg, while remote prefill was still running.
    request.context().stop();
    assert!(request.context().is_stopped());

    (request, selection, guard)
}

/// Covers the `CancelWhenStopped` arm inside `dispatch_selection` itself. The
/// `generate()`-level tests cancel earlier, at selection, so without this the
/// arm could be replaced by a bare `dispatch.await` and stay green.
#[tokio::test]
#[serial_test::serial]
async fn kv_cancellation_after_admission_still_stops_the_aggregated_dispatch() {
    let (router, dispatch, _worker_id, runtime) =
        router_with_recorded_dispatch("kv-aggregated-dispatch-after-stop").await;
    let (request, selection, guard) = admitted_request_stopped_before_dispatch(
        &router,
        "aggregated-pre-dispatch-cancellation",
        RequestPhase::Aggregated,
    )
    .await;

    let error = router
        .dispatch_selection(request, selection, guard, &CleanupBudget::default())
        .await
        .expect_err("aggregated dispatch must still be cancelled before it is issued");
    assert!(match_error_chain(
        error.as_ref(),
        &[ErrorType::Cancelled],
        &[]
    ));
    assert!(
        dispatch.worker_ids.lock().unwrap().is_empty(),
        "no worker should have been asked to serve a cancelled aggregated request"
    );

    drop(router);
    runtime.shutdown();
}

#[derive(Debug, PartialEq, Eq)]
enum ClassifierObservation {
    Completed(usize),
    /// The typed cause the plugin received, if the abort carried one.
    Aborted {
        cause: Option<ErrorType>,
    },
}

fn abort_cause_type(error: &(dyn std::error::Error + 'static)) -> Option<ErrorType> {
    let mut cause = Some(error);
    while let Some(current) = cause {
        if let Some(typed) = current.downcast_ref::<DynamoError>() {
            return Some(typed.error_type());
        }
        cause = current.source();
    }
    None
}

struct RecordingClassifier {
    calls: Arc<AtomicUsize>,
    observations: mpsc::UnboundedSender<ClassifierObservation>,
}

#[async_trait::async_trait]
impl RequestClassifier for RecordingClassifier {
    fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture {
        self.calls.fetch_add(1, Ordering::Relaxed);
        Box::pin(async move { Ok(request) })
    }

    async fn on_event(&mut self, event: ClassifyEvent) {
        let observation = match event {
            ClassifyEvent::Completed {
                context_tokens: Some(context_tokens),
                ..
            } => Some(ClassifierObservation::Completed(context_tokens)),
            ClassifyEvent::Aborted { error, .. } => Some(ClassifierObservation::Aborted {
                cause: error.as_deref().and_then(|error| abort_cause_type(error)),
            }),
            _ => None,
        };
        if let Some(observation) = observation {
            self.observations.send(observation).unwrap();
        }
    }
}

#[derive(Debug, thiserror::Error)]
#[error("classifier rejected request")]
struct ClassifierRejected;

struct RejectingClassifier {
    calls: Arc<AtomicUsize>,
    observations: mpsc::UnboundedSender<Option<(bool, String)>>,
}

#[async_trait::async_trait]
impl RequestClassifier for RejectingClassifier {
    fn classify(&mut self, _request: ClassifyRequest) -> ClassifyFuture {
        self.calls.fetch_add(1, Ordering::Relaxed);
        Box::pin(async { Err(Box::new(ClassifierRejected) as Box<ClassifierError>) })
    }

    async fn on_event(&mut self, event: ClassifyEvent) {
        if let ClassifyEvent::Aborted { error, .. } = event {
            self.observations
                .send(error.map(|error| {
                    (
                        error.downcast_ref::<ClassifierRejected>().is_some(),
                        error.to_string(),
                    )
                }))
                .unwrap();
        }
    }
}

#[tokio::test]
async fn classifier_failure_aborts_once_with_original_error() {
    let calls = Arc::new(AtomicUsize::new(0));
    let (observations_tx, mut observations_rx) = mpsc::unbounded_channel();
    let (router, runtime) = router_with_classifier(
        RejectingClassifier {
            calls: Arc::clone(&calls),
            observations: observations_tx,
        },
        Some(Duration::from_secs(10)),
    )
    .await;
    let session_id = SessionAffinityId::new("classifier-failure-stale-affinity");
    bind_affinity_target(&router, &session_id, AffinityTarget::new(7, Some(1))).await;
    let mut request = Context::new(request());
    request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id);

    let client_error = match router
        .select_with_affinity(
            &request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
    {
        Ok(_) => panic!("classifier rejection unexpectedly selected a worker"),
        Err(error) => error,
    };
    assert!(!format!("{client_error:#}").contains("classifier rejected request"));
    assert_eq!(calls.load(Ordering::Relaxed), 1);
    assert_eq!(
        tokio::time::timeout(Duration::from_secs(1), observations_rx.recv())
            .await
            .expect("classifier abort event timed out"),
        Some(Some((true, "classifier rejected request".to_string())))
    );
    assert!(observations_rx.try_recv().is_err());

    drop(router);
    runtime.shutdown();
}

struct TypedRejectingClassifier {
    error_type: ErrorType,
    calls: Arc<AtomicUsize>,
    observations: mpsc::UnboundedSender<ClassifierObservation>,
}

#[async_trait::async_trait]
impl RequestClassifier for TypedRejectingClassifier {
    fn classify(&mut self, _request: ClassifyRequest) -> ClassifyFuture {
        self.calls.fetch_add(1, Ordering::Relaxed);
        let error_type = self.error_type;
        Box::pin(async move {
            Err(Box::new(
                DynamoError::builder()
                    .error_type(error_type)
                    .message("classifier shed load")
                    .build(),
            ) as Box<ClassifierError>)
        })
    }

    async fn on_event(&mut self, event: ClassifyEvent) {
        if let ClassifyEvent::Aborted { error, .. } = event {
            self.observations
                .send(ClassifierObservation::Aborted {
                    cause: error.as_deref().and_then(|error| abort_cause_type(error)),
                })
                .unwrap();
        }
    }
}

#[tokio::test]
async fn typed_classifier_rejection_reaches_client_without_migration() {
    for error_type in [
        ErrorType::ResourceExhausted,
        ErrorType::WorkerOverloaded,
        ErrorType::CannotConnect,
    ] {
        let calls = Arc::new(AtomicUsize::new(0));
        let (observations_tx, mut observations_rx) = mpsc::unbounded_channel();
        let (router, runtime) = router_with_classifier(
            TypedRejectingClassifier {
                error_type,
                calls: Arc::clone(&calls),
                observations: observations_tx,
            },
            None,
        )
        .await;
        let engine: ServerStreamingEngine<PreprocessedRequest, Annotated<LLMEngineOutput>> =
            Arc::new(router);
        let migration = Migration::new(2, None, "test".to_string(), Arc::new(Metrics::new()));
        let client_error = match migration
            .generate(Context::new(request()), engine.clone())
            .await
        {
            Ok(_) => panic!("typed rejection unexpectedly selected a worker"),
            Err(error) => error,
        };

        assert_eq!(
            calls.load(Ordering::Relaxed),
            1,
            "classifier rejection {error_type:?} must not be retried"
        );
        assert!(match_error_chain(client_error.as_ref(), &[error_type], &[]));
        assert!(format!("{client_error:#}").contains("classifier shed load"));
        assert_eq!(
            tokio::time::timeout(Duration::from_secs(1), observations_rx.recv())
                .await
                .expect("classifier abort event timed out"),
            Some(ClassifierObservation::Aborted {
                cause: Some(error_type),
            })
        );
        assert!(observations_rx.try_recv().is_err());

        drop(engine);
        runtime.shutdown();
    }
}

#[tokio::test]
async fn stale_affinity_rebind_classifies_once_without_intermediate_abort() {
    let calls = Arc::new(AtomicUsize::new(0));
    let (observations_tx, mut observations_rx) = mpsc::unbounded_channel();
    let (router, runtime) = router_with_classifier(
        RecordingClassifier {
            calls: Arc::clone(&calls),
            observations: observations_tx,
        },
        Some(Duration::from_secs(10)),
    )
    .await;
    let session_id = SessionAffinityId::new("classifier-stale-affinity");
    bind_affinity_target(&router, &session_id, AffinityTarget::new(7, Some(1))).await;
    let mut request = Context::new(request());
    request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id);

    let (mut selection, operation) = router
        .select_with_affinity(
            &request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();

    assert_eq!(calls.load(Ordering::Relaxed), 1);
    assert!(observations_rx.try_recv().is_err());
    let mut lifecycle = selection.request_lifecycle.take().unwrap();
    lifecycle.observe_context_tokens(1);
    lifecycle.complete();
    assert_eq!(
        observations_rx.recv().await,
        Some(ClassifierObservation::Completed(1))
    );
    router.kv_router().free(request.id()).await.unwrap();

    drop(operation);
    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn query_only_selection_bypasses_classifier_and_request_lifecycle() {
    let calls = Arc::new(AtomicUsize::new(0));
    let (observations_tx, mut observations_rx) = mpsc::unbounded_channel();
    let (router, runtime) = router_with_classifier(
        RecordingClassifier {
            calls: Arc::clone(&calls),
            observations: observations_tx,
        },
        None,
    )
    .await;
    let request = Context::new(request());

    let (selection, operation) = router
        .select_with_affinity(
            &request,
            RequestPhase::Aggregated,
            true,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();

    assert_eq!(calls.load(Ordering::Relaxed), 0);
    assert!(selection.request_lifecycle.is_none());
    assert!(operation.is_none());
    drop(selection);
    assert!(!router.prefill_worker_busy(&request, 0.5).await.unwrap());
    assert_eq!(calls.load(Ordering::Relaxed), 0);
    assert!(observations_rx.try_recv().is_err());

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn planned_route_admission_classifies_and_carries_lifecycle() {
    let calls = Arc::new(AtomicUsize::new(0));
    let (observations_tx, mut observations_rx) = mpsc::unbounded_channel();
    let (router, runtime) = router_with_classifier(
        RecordingClassifier {
            calls: Arc::clone(&calls),
            observations: observations_tx,
        },
        None,
    )
    .await;
    let request = Context::new(request());

    let preview = router
        .preview_kv_route(&request, RequestPhase::Decode)
        .await
        .unwrap();
    assert_eq!(calls.load(Ordering::Relaxed), 0, "preview is query-only");
    let mut plan = router
        .plan_kv_route_from_preview(&request, preview)
        .await
        .unwrap();
    assert_eq!(calls.load(Ordering::Relaxed), 1);

    let mut lifecycle = plan.selection.request_lifecycle.take().unwrap();
    lifecycle.observe_context_tokens(1);
    lifecycle.complete();
    assert_eq!(
        observations_rx.recv().await,
        Some(ClassifierObservation::Completed(1))
    );
    router.kv_router().free(request.id()).await.unwrap();

    drop(plan);
    drop(router);
    runtime.shutdown();
}

struct PausingClassifier {
    entered: Arc<Notify>,
    resumed: Arc<Notify>,
}

impl RequestClassifier for PausingClassifier {
    fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture {
        let entered = Arc::clone(&self.entered);
        let resumed = Arc::clone(&self.resumed);
        Box::pin(async move {
            entered.notify_one();
            resumed.notified().await;
            Ok(request)
        })
    }
}

#[tokio::test]
async fn classifier_pause_defers_admission_until_resumed() {
    let entered = Arc::new(Notify::new());
    let resumed = Arc::new(Notify::new());
    let (router, runtime) = router_with_classifier(
        PausingClassifier {
            entered: Arc::clone(&entered),
            resumed: Arc::clone(&resumed),
        },
        None,
    )
    .await;
    let router = Arc::new(router);
    let request = Context::new(request());
    let request_id = request.id().to_string();
    let selection = {
        let router = Arc::clone(&router);
        tokio::spawn(async move {
            router
                .select_with_affinity(
                    &request,
                    RequestPhase::Aggregated,
                    false,
                    &CleanupBudget::default(),
                )
                .await
        })
    };

    entered.notified().await;
    assert!(!selection.is_finished());
    resumed.notify_one();
    let (selection, _) = selection.await.unwrap().unwrap();
    assert!(selection.request_lifecycle.is_some());
    router.kv_router().free(&request_id).await.unwrap();
    drop(selection);

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn tracked_admission_without_lifecycle_bypasses_classifier() {
    let calls = Arc::new(AtomicUsize::new(0));
    let (observations_tx, mut observations_rx) = mpsc::unbounded_channel();
    let (router, runtime) = router_with_classifier(
        RecordingClassifier {
            calls: Arc::clone(&calls),
            observations: observations_tx,
        },
        None,
    )
    .await;

    // `find_best_match_details_with_policy_class` with `update_states` is a tracked admission
    // that never calls `begin_request_lifecycle` — the same funnel as the
    // Python bindings `best_worker` and the standalone `RouterRequest::New`
    // path. The plugin would receive no lifecycle events for it, so admission
    // must use the default queue inputs instead of the classifier.
    let outcome = router
        .kv_router()
        .find_best_match_details_with_policy_class(
            Some("unregistered-tracked"),
            &[1, 2, 3, 4],
            None,
            None,
            true,
            false,
            None,
            None,
            0.0,
            0,
            None,
            None,
            None,
            None,
            None,
            RoutingConstraints::default(),
        )
        .await
        .unwrap();
    assert!(matches!(
        outcome,
        crate::kv_router::FindBestMatchOutcome::Routed { .. }
    ));
    assert_eq!(calls.load(Ordering::Relaxed), 0);
    router
        .kv_router()
        .free("unregistered-tracked")
        .await
        .unwrap();
    assert!(observations_rx.try_recv().is_err());

    drop(router);
    runtime.shutdown();
}

struct TokenContractClassifier {
    classified: mpsc::UnboundedSender<(usize, dynamo_kv_router::scheduling::RequestProgress)>,
    completed: mpsc::UnboundedSender<usize>,
}

#[async_trait::async_trait]
impl RequestClassifier for TokenContractClassifier {
    fn classify(&mut self, request: ClassifyRequest) -> ClassifyFuture {
        self.classified
            .send((request.input_tokens(), request.progress().clone()))
            .unwrap();
        Box::pin(async move { Ok(request) })
    }

    async fn on_event(&mut self, event: ClassifyEvent) {
        if let ClassifyEvent::Completed {
            context_tokens: Some(context_tokens),
            ..
        } = event
        {
            self.completed.send(context_tokens).unwrap();
        }
    }
}

#[tokio::test]
async fn classifier_uses_scheduler_token_basis() {
    let (classified_tx, mut classified_rx) = mpsc::unbounded_channel();
    let (completed_tx, mut completed_rx) = mpsc::unbounded_channel();
    let (router, runtime) = router_with_classifier(
        TokenContractClassifier {
            classified: classified_tx,
            completed: completed_tx,
        },
        None,
    )
    .await;
    let mut content = request();
    content.mm_routing_info = Some(MmRoutingInfo {
        routing_token_ids: vec![1, 2, 3, 4, 5, 0, 0, 0],
        block_mm_infos: vec![None],
        expanded_prompt_len: 5,
    });
    let request = Context::new(content);
    let (mut selection, _) = router
        .select_with_affinity(
            &request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    // Classification and the scheduler queue share one token basis: the
    // routing tokens (`isl_tokens`, 8 here), not the multimodal expanded
    // prompt length (5), so a pass-through classifier cannot shift queue
    // bucketing, limits, or DRR cost for multimodal requests.
    let (input_tokens, progress) = classified_rx.recv().await.unwrap();
    assert_eq!(input_tokens, 8);
    assert_eq!(progress.context_tokens(), 8);

    let mut guard = router
        .track_selection(
            &request,
            &mut selection,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    guard.mark_dispatched();
    let output = LLMEngineOutput {
        completion_usage: Some(dynamo_protocols::types::CompletionUsage {
            prompt_tokens: 5,
            completion_tokens: 6,
            total_tokens: 11,
            prompt_tokens_details: None,
            completion_tokens_details: None,
        }),
        ..Default::default()
    };
    guard.on_item(&Annotated::from_data(output)).await;
    assert_eq!(progress.context_tokens(), 11);
    assert!(completed_rx.try_recv().is_err());
    guard.finish().await;
    assert_eq!(completed_rx.recv().await, Some(11));

    drop(router);
    runtime.shutdown();
}

async fn track_request(
    router: &RoutingHost,
    is_query_only: bool,
) -> (SingleIn<PreprocessedRequest>, WorkerSelection, RequestGuard) {
    let request = Context::new(request());
    let (mut selection, _) = router
        .select_with_affinity(
            &request,
            RequestPhase::Aggregated,
            is_query_only,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    let guard = router
        .track_selection(
            &request,
            &mut selection,
            RequestPhase::Aggregated,
            is_query_only,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    (request, selection, guard)
}

/// Scheduler-reported loads with no prompt, the way the reservation tests read them.
async fn potential_loads(router: &RoutingHost) -> Vec<dynamo_kv_router::protocols::PotentialLoad> {
    router
        .kv_router()
        .get_potential_loads(&[], None, None, None, None)
        .await
        .unwrap()
}

fn active_requests_for(
    loads: &[dynamo_kv_router::protocols::PotentialLoad],
    worker_id: u64,
    dp_rank: u32,
) -> usize {
    loads
        .iter()
        .find(|load| load.worker_id == worker_id && load.dp_rank == dp_rank)
        .expect("selected worker must be reported")
        .active_requests
}

/// Preview and admit one decode route, the two stages the plan tests exercise.
async fn plan_decode_route(
    router: &RoutingHost,
    request: &Context<PreprocessedRequest>,
) -> RoutePlan {
    let preview = router
        .preview_kv_route(request, RequestPhase::Decode)
        .await
        .expect("decode preview should select one request");
    router
        .plan_kv_route_from_preview(request, preview)
        .await
        .expect("decode plan should admit one request")
}

#[tokio::test]
#[serial_test::serial]
async fn route_plan_from_preview_holds_and_releases_the_decode_reservation() {
    let (router, runtime) = router(None).await;
    let request = Context::new(request());

    let plan = plan_decode_route(&router, &request).await;
    assert_eq!(plan.signals.worker.worker_id, 7);
    assert_eq!(
        router.request_metrics.requests_started_total.get(),
        0,
        "a topology decision is not a started request"
    );
    let admitted_loads = potential_loads(&router).await;
    assert_eq!(active_requests_for(&admitted_loads, 7, 0), 1);

    plan.abort().await;
    let released_loads = potential_loads(&router).await;
    assert!(
        released_loads.iter().all(|load| load.active_requests == 0),
        "abandoned plans must release their scheduler reservation: {released_loads:?}"
    );
    assert_eq!(
        router.request_metrics.requests_started_total.get(),
        0,
        "an abandoned plan must not count as a started request"
    );

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn route_preview_does_not_admit_a_request() {
    let (router, runtime) = router(None).await;
    let request = Context::new(request());

    let preview = router
        .preview_kv_route(&request, RequestPhase::Decode)
        .await
        .expect("decode preview should select one request");
    assert_eq!(preview.signals.worker.worker_id, 7);
    let loads = potential_loads(&router).await;
    assert!(loads.iter().all(|load| load.active_requests == 0));
    assert_eq!(router.request_metrics.requests_started_total.get(), 0);

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn route_plan_from_preview_admits_the_previewed_worker() {
    let (router, runtime) = router_with_workers(None, &[7, 8]).await;
    let request = Context::new(request());
    let preview = router
        .preview_kv_route(&request, RequestPhase::Decode)
        .await
        .unwrap();
    let previewed_worker = preview.signals.worker;

    let plan = router
        .plan_kv_route_from_preview(&request, preview)
        .await
        .unwrap();
    assert_eq!(plan.signals.worker, previewed_worker);
    let loads = potential_loads(&router).await;
    assert_eq!(
        active_requests_for(&loads, previewed_worker.worker_id, previewed_worker.dp_rank),
        1
    );
    assert_eq!(
        loads.iter().filter(|load| load.active_requests > 0).count(),
        1
    );

    plan.abort().await;
    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn route_preview_does_not_acquire_session_affinity() {
    let (router, runtime) = router(Some(Duration::from_secs(10))).await;
    let session_id = SessionAffinityId::new("preview-only");
    let mut request = Context::new(request());
    request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id.clone());

    router
        .preview_kv_route(&request, RequestPhase::Decode)
        .await
        .unwrap();

    let affinity = router.affinity.as_ref().unwrap();
    let acquisition = tokio::time::timeout(
        Duration::from_millis(100),
        affinity.acquire(&session_id, None),
    )
    .await
    .expect("preview must not leave affinity initialization pending")
    .unwrap();
    assert!(matches!(acquisition, Hold::Initialize(_)));
    drop(acquisition);

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn planned_dispatch_transfers_the_reservation_to_request_cleanup() {
    let (router, runtime) = router(None).await;
    let request = Context::new(request());
    let plan = plan_decode_route(&router, &request).await;

    assert!(router.dispatch_kv_plan(request, plan).await.is_err());
    assert_eq!(router.request_metrics.requests_started_total.get(), 1);
    let loads = potential_loads(&router).await;
    assert!(loads.iter().all(|load| load.active_requests == 0));

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn prefill_busy_probe_does_not_admit_a_request() {
    let (router, runtime) = router(None).await;
    let request = Context::new(request());

    assert!(!router.prefill_worker_busy(&request, 0.5).await.unwrap());
    let loads = potential_loads(&router).await;
    assert!(loads.iter().all(|load| load.active_requests == 0));
    assert_eq!(router.request_metrics.requests_started_total.get(), 0);

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn aborted_route_plan_drops_pending_affinity_initialization() {
    let (router, runtime) = router(Some(Duration::from_secs(10))).await;
    let session_id = SessionAffinityId::new("abandoned-route-plan");
    let mut request = Context::new(request());
    request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id.clone());

    let preview = router
        .preview_kv_route(&request, RequestPhase::Decode)
        .await
        .unwrap();
    router
        .plan_kv_route_from_preview(&request, preview)
        .await
        .unwrap()
        .abort()
        .await;

    let affinity = router.affinity.as_ref().unwrap();
    let acquisition = tokio::time::timeout(
        Duration::from_millis(100),
        affinity.acquire(&session_id, None),
    )
    .await
    .expect("abandoned plan must not leave affinity initialization pending")
    .unwrap();
    assert!(matches!(acquisition, Hold::Initialize(_)));
    drop(acquisition);

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn session_affinity_disabled_does_not_create_coordinator() {
    let (router, runtime) = router(None).await;
    assert!(router.affinity.is_none());

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn router_request_counters_follow_admission_and_completion_lifecycle() {
    let (router, runtime) = router(None).await;
    let metrics = router.request_metrics.clone();

    let controller = Controller::new("pre-admission-cancellation".to_string());
    controller.stop();
    let cancelled_request = Context::with_controller(request(), controller);
    assert!(
        router
            .select_with_affinity(
                &cancelled_request,
                RequestPhase::Aggregated,
                false,
                &CleanupBudget::default()
            )
            .await
            .is_err()
    );
    assert_eq!(metrics.requests_started_total.get(), 0);

    let (_, _, mut query_guard) = track_request(&router, true).await;
    query_guard.abort().await;
    drop(query_guard);
    assert_eq!(metrics.requests_started_total.get(), 0);

    let (_, selection, mut cancelled_guard) = track_request(&router, false).await;

    assert_eq!(metrics.requests_started_total.get(), 1);
    assert_eq!(metrics.requests_total.get(), 0);
    // The booking is held from admission until the guard finishes.
    let loads = potential_loads(&router).await;
    assert_eq!(
        loads
            .iter()
            .find(|load| load.worker_id == selection.worker.worker_id)
            .expect("tracked worker must be reported")
            .active_requests,
        1,
        "a tracked request stays booked while its guard is live: {loads:?}"
    );

    // Admission remains counted even when the request aborts before dispatch.
    cancelled_guard.abort().await;
    drop(cancelled_guard);
    assert_eq!(metrics.requests_started_total.get(), 1);
    assert_eq!(metrics.requests_total.get(), 0);
    let loads = potential_loads(&router).await;
    assert!(
        loads.iter().all(|load| load.active_requests == 0),
        "an aborted guard frees its booking: {loads:?}"
    );

    let mut failed_input = request();
    failed_input.migration_state = Some(Default::default());
    let migration_state = failed_input.migration_state.clone().unwrap();
    let failed_request = Context::new(failed_input);
    let (mut failed_selection, _) = router
        .select_with_affinity(
            &failed_request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    let failed_worker = failed_selection.worker.worker_id;
    let failed_dispatch_guard = router
        .track_selection(
            &failed_request,
            &mut failed_selection,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    assert!(
        router
            .dispatch_selection(
                failed_request,
                failed_selection,
                failed_dispatch_guard,
                &CleanupBudget::default()
            )
            .await
            .is_err()
    );
    assert_eq!(migration_state.excluded_worker_ids(), vec![failed_worker]);
    assert_eq!(metrics.requests_started_total.get(), 2);
    assert_eq!(metrics.requests_total.get(), 0);

    let (_, _, mut completed_guard) = track_request(&router, false).await;
    completed_guard.start_dispatch("aggregated");
    completed_guard.mark_dispatched();
    completed_guard.finish().await;
    drop(completed_guard);
    assert_eq!(metrics.requests_started_total.get(), 3);
    assert_eq!(metrics.requests_total.get(), 1);

    let mut builtin_guard =
        RequestGuard::new_builtin(Arc::clone(&metrics), 7, None, None, &request());
    assert_eq!(metrics.requests_started_total.get(), 4);
    builtin_guard.abort().await;
    drop(builtin_guard);
    assert_eq!(metrics.requests_total.get(), 1);

    let (other_router, other_runtime) = router_with_workers(None, &[8]).await;
    let (_, _, mut other_guard) = track_request(&other_router, false).await;
    other_guard.start_dispatch("aggregated");
    other_guard.mark_dispatched();
    other_guard.finish().await;
    drop(other_guard);
    assert_eq!(metrics.requests_started_total.get(), 4);
    assert_eq!(metrics.requests_total.get(), 1);
    drop(other_router);
    other_runtime.shutdown();

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn session_affinity_post_selection_failures_preserve_binding() {
    let (router, runtime) = router(Some(Duration::from_secs(10))).await;
    let affinity = router.affinity.as_ref().unwrap();
    let session_id = SessionAffinityId::new("cancelled-after-selection");
    let original_target = AffinityTarget {
        worker_id: 7,
        dp_rank: Some(0),
    };
    let Hold::Initialize(initializer) = affinity.acquire(&session_id, None).await.unwrap() else {
        panic!("first request must initialize");
    };
    drop(
        initializer
            .commit(crate::session_affinity::to_table(original_target))
            .unwrap(),
    );

    let operation = Some(affinity.acquire(&session_id, None).await.unwrap());
    drop(operation);
    assert_eq!(
        affinity.query_target(&session_id, None).unwrap(),
        Some(original_target)
    );

    let operation = Some(affinity.acquire(&session_id, None).await.unwrap());
    drop(operation);
    assert_eq!(
        affinity.query_target(&session_id, None).unwrap(),
        Some(original_target)
    );

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn session_affinity_existing_selection_cancellation_preserves_binding_without_retry() {
    let (router, runtime) = router(Some(Duration::from_secs(10))).await;
    let session_id = SessionAffinityId::new("cancelled-selection");
    let original_target = AffinityTarget {
        worker_id: 7,
        dp_rank: Some(0),
    };
    let Hold::Initialize(initializer) = router
        .affinity
        .as_ref()
        .unwrap()
        .acquire(&session_id, None)
        .await
        .unwrap()
    else {
        panic!("first request must initialize");
    };
    drop(
        initializer
            .commit(crate::session_affinity::to_table(original_target))
            .unwrap(),
    );

    let controller = Controller::new("cancelled-selection-request".to_string());
    controller.stop();
    let mut request = Context::with_controller(request(), controller);
    request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id.clone());

    let Err(error) = router
        .select_with_affinity(
            &request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
    else {
        panic!("stopped request must return cancellation");
    };
    assert!(match_error_chain(
        error.as_ref(),
        &[ErrorType::Cancelled],
        &[]
    ));
    assert_eq!(
        router
            .affinity
            .as_ref()
            .unwrap()
            .query_target(&session_id, None)
            .unwrap(),
        Some(original_target)
    );

    let Hold::Bound { target, lease } = router
        .affinity
        .as_ref()
        .unwrap()
        .acquire(&session_id, None)
        .await
        .unwrap()
    else {
        panic!("cancellation must preserve the existing binding");
    };
    assert_eq!(target, crate::session_affinity::to_table(original_target));
    drop(lease);

    drop(router);
    runtime.shutdown();
}

/// A builtin host over `inner` whose session-affinity coordinator has a 10s
/// TTL and binds in `mode`.
fn builtin_host_with_affinity(
    inner: PushRouter<PreprocessedRequest, Annotated<LLMEngineOutput>>,
    load_context: Arc<RoutingLoadContext>,
    mode: crate::session_affinity::SessionAffinityMode,
) -> (RoutingHost, AffinityCoordinator) {
    let affinity = AffinityCoordinator::new(Duration::from_secs(10), mode).unwrap();
    let host =
        RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(affinity.clone()))
            .unwrap();
    (host, affinity)
}

async fn bind_affinity_target(
    router: &RoutingHost,
    session_id: &SessionAffinityId,
    target: AffinityTarget,
) {
    let Hold::Initialize(initializer) = router
        .affinity
        .as_ref()
        .unwrap()
        .acquire(session_id, None)
        .await
        .unwrap()
    else {
        panic!("first request must initialize");
    };
    drop(
        initializer
            .commit(crate::session_affinity::to_table(target))
            .unwrap(),
    );
}

#[tokio::test]
async fn request_constraints_preserve_worker_only_affinity() {
    let mut request_worker = ModelRuntimeConfig::default();
    request_worker.taints.insert("request-pool".to_string());
    let workers = HashMap::from([(7, ModelRuntimeConfig::default()), (8, request_worker)]);
    let (router, runtime) =
        router_with_worker_configs(Some(Duration::from_secs(10)), workers).await;
    let target = AffinityTarget::worker(7);

    let allowlist_session = SessionAffinityId::new("affinity-allowlist-conflict");
    bind_affinity_target(&router, &allowlist_session, target).await;
    let mut allowlist_input = request();
    allowlist_input.routing_mut().allowed_worker_ids = Some(HashSet::from([8]));
    let mut allowlist_request = Context::new(allowlist_input);
    allowlist_request.insert(SESSION_AFFINITY_CONTEXT_KEY, allowlist_session.clone());
    assert!(
        router
            .select_with_affinity(
                &allowlist_request,
                RequestPhase::Aggregated,
                false,
                &CleanupBudget::default()
            )
            .await
            .is_err()
    );
    assert_eq!(
        router
            .affinity
            .as_ref()
            .unwrap()
            .query_target(&allowlist_session, None)
            .unwrap(),
        Some(target)
    );

    let taint_session = SessionAffinityId::new("affinity-taint-conflict");
    bind_affinity_target(&router, &taint_session, target).await;
    let mut taint_input = request();
    taint_input.routing_mut().routing_constraints = Some(RoutingConstraints {
        required_taints: HashSet::from(["request-pool".to_string()]),
        ..Default::default()
    });
    let mut taint_request = Context::new(taint_input);
    taint_request.insert(SESSION_AFFINITY_CONTEXT_KEY, taint_session.clone());
    assert!(
        router
            .select_with_affinity(
                &taint_request,
                RequestPhase::Aggregated,
                false,
                &CleanupBudget::default()
            )
            .await
            .is_err()
    );
    assert_eq!(
        router
            .affinity
            .as_ref()
            .unwrap()
            .query_target(&taint_session, None)
            .unwrap(),
        Some(target)
    );

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn stale_affinity_rank_recovers_within_request() {
    let (router, runtime) = router(Some(Duration::from_secs(10))).await;
    let session_id = SessionAffinityId::new("stale-affinity-rank");
    bind_affinity_target(&router, &session_id, AffinityTarget::new(7, Some(1))).await;
    let mut request = Context::new(request());
    request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id);

    let (selection, operation) = router
        .select_with_affinity(
            &request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    assert_eq!(selection.worker, WorkerWithDpRank::new(7, 0));
    assert!(matches!(operation, Some(Hold::Initialize(_))));
    router.kv_router().free(request.id()).await.unwrap();

    drop(operation);
    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn config_watch_gap_preserves_hard_affinity() {
    let (router, runtime) = router(Some(Duration::from_secs(10))).await;
    let session_id = SessionAffinityId::new("config-watch-gap");
    let target = AffinityTarget::new(8, Some(0));
    router
        .inner
        .client
        .override_discovered_instances(vec![7, target.worker_id]);
    router
        .inner
        .client
        .override_instance_avail(vec![7, target.worker_id]);
    bind_affinity_target(&router, &session_id, target).await;
    assert!(router.affinity_target_is_valid(target));

    let mut request = Context::new(request());
    request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id.clone());
    assert!(
        router
            .select_with_affinity(
                &request,
                RequestPhase::Aggregated,
                false,
                &CleanupBudget::default()
            )
            .await
            .is_err()
    );
    assert_eq!(
        router
            .affinity
            .as_ref()
            .unwrap()
            .query_target(&session_id, None)
            .unwrap(),
        Some(target)
    );

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
async fn migration_exclusion_preserves_hard_affinity_without_widening_or_escaping_hard_pins() {
    let mut constrained_worker = ModelRuntimeConfig::default();
    constrained_worker.taints.insert("retry-pool".to_string());
    let workers = HashMap::from([
        (7, constrained_worker),
        (8, ModelRuntimeConfig::default()),
        (9, ModelRuntimeConfig::default()),
    ]);
    let (router, runtime) =
        router_with_worker_configs(Some(Duration::from_secs(10)), workers).await;
    let session_id = SessionAffinityId::new("migration-exclusion");
    let original_target = AffinityTarget {
        worker_id: 7,
        dp_rank: Some(0),
    };
    let Hold::Initialize(initializer) = router
        .affinity
        .as_ref()
        .unwrap()
        .acquire(&session_id, None)
        .await
        .unwrap()
    else {
        panic!("first request must initialize");
    };
    drop(
        initializer
            .commit(crate::session_affinity::to_table(original_target))
            .unwrap(),
    );

    let mut retry_input = request();
    retry_input.routing_mut().allowed_worker_ids = Some(HashSet::from([7, 8]));
    retry_input.migration_state = Some(Default::default());
    retry_input
        .migration_state
        .as_ref()
        .unwrap()
        .record_failure(
            7,
            Some(
                DynamoError::builder()
                    .error_type(ErrorType::WorkerOverloaded)
                    .message("worker 7 overloaded")
                    .build(),
            ),
        );
    let mut retry_request = Context::new(retry_input);
    retry_request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id.clone());

    let Err(error) = router
        .select_with_affinity(
            &retry_request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
    else {
        panic!("migration exclusions must not rebind hard affinity");
    };
    assert!(error.to_string().contains("worker 7"));
    assert_eq!(
        router
            .affinity
            .as_ref()
            .unwrap()
            .query_target(&session_id, None)
            .unwrap(),
        Some(original_target)
    );

    let mut exhausted_input = request();
    exhausted_input.routing_mut().allowed_worker_ids = Some(HashSet::from([7, 10]));
    exhausted_input.migration_state = Some(Default::default());
    exhausted_input
        .migration_state
        .as_ref()
        .unwrap()
        .record_failure(
            7,
            Some(
                DynamoError::builder()
                    .error_type(ErrorType::WorkerOverloaded)
                    .message("worker 7 overloaded")
                    .build(),
            ),
        );
    let exhausted_request = Context::new(exhausted_input);
    let Err(error) = router
        .select_with_affinity(
            &exhausted_request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
    else {
        panic!("exhausting the constrained worker set must reject the retry");
    };
    assert!(match_error_chain(
        error.as_ref(),
        &[ErrorType::ResourceExhausted],
        &[]
    ));

    let mut constrained_input = request();
    constrained_input.routing_mut().routing_constraints = Some(RoutingConstraints {
        required_taints: HashSet::from(["retry-pool".to_string()]),
        ..Default::default()
    });
    constrained_input.migration_state = Some(Default::default());
    constrained_input
        .migration_state
        .as_ref()
        .unwrap()
        .record_failure(
            7,
            Some(
                DynamoError::builder()
                    .error_type(ErrorType::WorkerOverloaded)
                    .message("worker 7 overloaded")
                    .build(),
            ),
        );
    let constrained_request = Context::new(constrained_input);
    let Err(error) = router
        .select_with_affinity(
            &constrained_request,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
        )
        .await
    else {
        panic!("routing constraints must not be widened during retry");
    };
    assert!(match_error_chain(
        error.as_ref(),
        &[ErrorType::ResourceExhausted],
        &[]
    ));

    let mut pinned_input = request();
    let routing = pinned_input.routing_mut();
    routing.backend_instance_id = Some(7);
    routing.dp_rank = Some(0);
    pinned_input.migration_state = Some(Default::default());
    pinned_input
        .migration_state
        .as_ref()
        .unwrap()
        .record_failure(7, None);
    let pinned_request = Context::new(pinned_input);
    let (selection, _) = router
        .select_with_affinity(
            &pinned_request,
            RequestPhase::Aggregated,
            true,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    assert_eq!(selection.worker.worker_id, 7);

    drop(router);
    runtime.shutdown();
}

#[derive(Default)]
struct RejectFirstDispatch {
    attempts: Mutex<Vec<(u64, Vec<u64>)>>,
}

#[async_trait]
impl StreamingDispatch<PreprocessedRequest, Annotated<LLMEngineOutput>> for RejectFirstDispatch {
    async fn generate(
        &self,
        request: SingleIn<AddressedRequest<PreprocessedRequest>>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        let (addressed, context) = request.transfer(());
        let (request, _, instance) = addressed.into_parts();
        let worker_id = instance.expect("selected worker instance").id();
        let excluded_worker_ids = request
            .migration_state
            .as_ref()
            .map(|state| state.excluded_worker_ids())
            .unwrap_or_default();
        let attempt = {
            let mut attempts = self.attempts.lock().unwrap();
            attempts.push((worker_id, excluded_worker_ids));
            attempts.len()
        };

        if attempt == 1 {
            let output = Annotated {
                data: None,
                id: None,
                event: Some("error".to_string()),
                comment: None,
                error: Some(
                    DynamoError::builder()
                        .error_type(ErrorType::WorkerOverloaded)
                        .message("selected worker is overloaded")
                        .build(),
                ),
            };
            return Ok(ResponseStream::new(
                Box::pin(stream::once(async move { output })),
                context.context(),
            ));
        }

        let output = Annotated::from_data(LLMEngineOutput {
            token_ids: vec![2],
            finish_reason: Some(FinishReason::Stop),
            ..Default::default()
        });
        Ok(ResponseStream::new(
            Box::pin(stream::once(async move { output })),
            context.context(),
        ))
    }

    async fn generate_bidirectional(
        &self,
        _instance: Instance,
        _address: String,
        _input: ManyIn<PreprocessedRequest>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        unreachable!("the KV router dispatches unary requests")
    }
}

/// A two-worker routing host wired to a caller-supplied dispatch, for driving migration
/// end to end by choosing what the first worker answers with.
struct MigrationHarness {
    runtime: Runtime,
    chooser: Arc<KvRouter>,
    engine: ServerStreamingEngine<PreprocessedRequest, Annotated<LLMEngineOutput>>,
    registered_ids: HashSet<u64>,
    // Kept alive: dropping these tears down discovery for the workers the router must see.
    _store: tempfile::TempDir,
    _drts: Vec<DistributedRuntime>,
}

async fn two_worker_migration_harness(
    namespace: &str,
    dispatch: Arc<dyn StreamingDispatch<PreprocessedRequest, Annotated<LLMEngineOutput>>>,
    classifier: Option<RecordingClassifier>,
) -> MigrationHarness {
    async fn shared_drt(runtime: Runtime, store_path: &std::path::Path) -> DistributedRuntime {
        DistributedRuntime::new(
            runtime,
            DistributedConfig {
                discovery_backend: DiscoveryBackend::KvStore(Selector::File(
                    store_path.to_path_buf(),
                )),
                nats_config: None,
                request_plane: RequestPlaneMode::Tcp,
                response_plane: None,
                event_transport_kind: EventTransportKind::Zmq,
            },
        )
        .await
        .unwrap()
    }

    let runtime = Runtime::from_current().unwrap();
    let store = tempfile::tempdir().unwrap();
    let router_drt = shared_drt(runtime.clone(), store.path()).await;
    let first_worker_drt = shared_drt(runtime.clone(), store.path()).await;
    let second_worker_drt = shared_drt(runtime.clone(), store.path()).await;
    let endpoint_for = |drt: &DistributedRuntime| {
        drt.namespace(namespace.to_string())
            .unwrap()
            .component("workers".to_string())
            .unwrap()
            .endpoint("generate")
    };
    let first_worker_endpoint = endpoint_for(&first_worker_drt);
    let second_worker_endpoint = endpoint_for(&second_worker_drt);
    first_worker_endpoint
        .register_endpoint_instance()
        .await
        .unwrap();
    second_worker_endpoint
        .register_endpoint_instance()
        .await
        .unwrap();

    let endpoint = endpoint_for(&router_drt);
    let client = endpoint.client().await.unwrap();
    let load_context = test_load_context(&client).await;
    let instances = tokio::time::timeout(Duration::from_secs(5), async {
        let mut source = client.instance_source.as_ref().clone();
        loop {
            let instances = source.borrow_and_update().clone();
            if instances.len() == 2 {
                return instances;
            }
            source.changed().await.unwrap();
        }
    })
    .await
    .expect("both workers must be discovered");
    let registered_ids = instances
        .into_iter()
        .map(|instance| instance.id())
        .collect::<HashSet<_>>();
    assert_eq!(registered_ids.len(), 2);

    let workers = registered_ids
        .iter()
        .copied()
        .map(|worker_id| (worker_id, ModelRuntimeConfig::default()))
        .collect::<HashMap<_, _>>();
    let (_workers_tx, workers) = watch::channel(workers);
    let config = KvRouterConfig {
        skip_initial_worker_wait: true,
        use_kv_events: false,
        router_track_active_blocks: false,
        ..Default::default()
    };
    let mut chooser = KvRouter::new_with_worker_role_and_scheduler_load(
        endpoint,
        client.clone(),
        workers,
        None,
        16,
        SelectionPolicySource::Registry,
        Some(config),
        None,
        None,
        "decode",
        None,
        false,
        None,
        None,
        load_context.scheduler_load_sender(),
        load_context.cancellation_token(),
    )
    .await
    .unwrap();
    if let Some(classifier) = classifier {
        chooser = chooser.with_request_classifier(classifier).unwrap();
    }
    let push_router =
        PushRouter::from_client_with_dispatch(client.clone(), RouterMode::KV, dispatch)
            .await
            .unwrap();
    let chooser = Arc::new(chooser);
    let engine: ServerStreamingEngine<PreprocessedRequest, Annotated<LLMEngineOutput>> = Arc::new(
        RoutingHost::new_with_load_context(
            push_router,
            chooser.clone(),
            load_context,
            None,
            crate::session_affinity::SessionAffinityMode::Hard,
        )
        .unwrap(),
    );

    MigrationHarness {
        runtime,
        chooser,
        engine,
        registered_ids,
        _store: store,
        _drts: vec![router_drt, first_worker_drt, second_worker_drt],
    }
}

#[tokio::test]
#[serial_test::serial]
async fn worker_overload_stream_migration_releases_and_reselects() {
    let dispatch = Arc::new(RejectFirstDispatch::default());
    let harness =
        two_worker_migration_harness("worker-overload-migration", dispatch.clone(), None).await;
    let migration = Migration::new(1, None, "test".to_string(), Arc::new(Metrics::new()));

    let responses: Vec<_> = migration
        .generate(Context::new(request()), harness.engine.clone())
        .await
        .unwrap()
        .collect()
        .await;

    assert_eq!(responses.len(), 1);
    assert!(responses[0].error.is_none());
    assert_eq!(responses[0].data.as_ref().unwrap().token_ids, vec![2]);
    let attempts = {
        let attempts = dispatch.attempts.lock().unwrap();
        attempts.clone()
    };
    assert_eq!(attempts.len(), 2);
    let failed_worker = attempts[0].0;
    let retried_worker = attempts[1].0;
    assert_ne!(failed_worker, retried_worker);
    assert!(harness.registered_ids.contains(&failed_worker));
    assert!(harness.registered_ids.contains(&retried_worker));
    assert!(attempts[0].1.is_empty());
    assert_eq!(attempts[1].1, vec![failed_worker]);
    let loads = harness
        .chooser
        .get_potential_loads(&[], None, None, None, None)
        .await
        .unwrap();
    assert!(
        loads.iter().all(|load| load.active_requests == 0),
        "all scheduler bookings must be released after migration: {loads:?}"
    );
    harness.runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn stream_migration_retry_continues_one_classifier_lifecycle() {
    let calls = Arc::new(AtomicUsize::new(0));
    let (observations_tx, mut observations_rx) = mpsc::unbounded_channel();
    let dispatch = Arc::new(RejectFirstDispatch::default());
    let harness = two_worker_migration_harness(
        "stream-migration-classifier-lifecycle",
        dispatch.clone(),
        Some(RecordingClassifier {
            calls: Arc::clone(&calls),
            observations: observations_tx,
        }),
    )
    .await;
    let migration = Migration::new(1, None, "test".to_string(), Arc::new(Metrics::new()));

    let responses: Vec<_> = migration
        .generate(Context::new(request()), harness.engine.clone())
        .await
        .unwrap()
        .collect()
        .await;

    assert_eq!(responses.len(), 1);
    assert!(responses[0].error.is_none());
    let attempts = {
        let attempts = dispatch.attempts.lock().unwrap();
        attempts.clone()
    };
    assert_eq!(attempts.len(), 2, "the failed stream must be retried once");
    assert_ne!(attempts[0].0, attempts[1].0);

    // The retry continues the failed attempt's lifecycle, so the plugin must
    // see exactly one terminal event for the logical request: the retry's
    // Completed, with no Aborted from the failed stream before it.
    let observation = tokio::time::timeout(Duration::from_secs(1), observations_rx.recv())
        .await
        .expect("classifier terminal event timed out")
        .expect("classifier event channel closed");
    assert!(
        matches!(observation, ClassifierObservation::Completed(_)),
        "the failed stream must not abort the lifecycle Migration retries: {observation:?}"
    );
    assert!(
        observations_rx.try_recv().is_err(),
        "the logical request must emit exactly one terminal lifecycle event"
    );
    assert_eq!(
        calls.load(Ordering::Relaxed),
        1,
        "the retry must reuse the cached classification"
    );
    let loads = harness
        .chooser
        .get_potential_loads(&[], None, None, None, None)
        .await
        .unwrap();
    assert!(
        loads.iter().all(|load| load.active_requests == 0),
        "both attempt bookings must be released after classifier completion: {loads:?}"
    );
    harness.runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn prefill_retry_preserves_parked_decode_classifier_lifecycle() {
    let calls = Arc::new(AtomicUsize::new(0));
    let (observations_tx, mut observations_rx) = mpsc::unbounded_channel();
    let (decode, decode_runtime) = Box::pin(router_with_classifier(
        RecordingClassifier {
            calls: Arc::clone(&calls),
            observations: observations_tx,
        },
        None,
    ))
    .await;
    let (prefill, prefill_runtime) = Box::pin(router(None)).await;
    let mut content = request();
    content.migration_state = Some(Default::default());
    let request = Context::new(content);

    let (mut selected, _) = decode
        .select_with_affinity(
            &request,
            RequestPhase::Decode,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    let mut guard = decode
        .track_selection(
            &request,
            &mut selected,
            RequestPhase::Decode,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    guard.mark_dispatched();
    assert!(guard.release_for_retry().await);
    assert!(
        potential_loads(&decode)
            .await
            .iter()
            .all(|load| load.active_requests == 0)
    );

    let (mut selected, _) = prefill
        .select_with_affinity(
            &request,
            RequestPhase::Prefill,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    assert!(
        selected.request_lifecycle.is_none(),
        "prefill must not claim the decode lifecycle"
    );
    let mut guard = prefill
        .track_selection(
            &request,
            &mut selected,
            RequestPhase::Prefill,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    guard.mark_dispatched();
    guard.finish().await;

    let (mut selected, _) = decode
        .select_with_affinity(
            &request,
            RequestPhase::Decode,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    let mut guard = decode
        .track_selection(
            &request,
            &mut selected,
            RequestPhase::Decode,
            false,
            &CleanupBudget::default(),
        )
        .await
        .unwrap();
    guard.mark_dispatched();
    guard.finish().await;
    let observation = tokio::time::timeout(Duration::from_secs(1), observations_rx.recv())
        .await
        .unwrap()
        .unwrap();
    assert!(matches!(observation, ClassifierObservation::Completed(_)));
    assert!(observations_rx.try_recv().is_err());
    assert_eq!(
        calls.load(Ordering::Relaxed),
        1,
        "decode retry reuses its classification"
    );
    drop(decode);
    drop(prefill);
    decode_runtime.shutdown();
    prefill_runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn pinned_request_stream_failure_aborts_parked_lifecycle_with_cause() {
    let calls = Arc::new(AtomicUsize::new(0));
    let (observations_tx, mut observations_rx) = mpsc::unbounded_channel();
    let dispatch = Arc::new(RejectFirstDispatch::default());
    let harness = two_worker_migration_harness(
        "pinned-stream-failure-classifier-lifecycle",
        dispatch.clone(),
        Some(RecordingClassifier {
            calls: Arc::clone(&calls),
            observations: observations_tx,
        }),
    )
    .await;
    let migration = Migration::new(1, None, "test".to_string(), Arc::new(Metrics::new()));
    let pinned = *harness.registered_ids.iter().next().unwrap();
    let mut content = request();
    content.routing = Some(RoutingHints {
        backend_instance_id: Some(pinned),
        ..Default::default()
    });

    let responses: Vec<_> = migration
        .generate(Context::new(content), harness.engine.clone())
        .await
        .unwrap()
        .collect()
        .await;

    // The failure is type-migratable so the stream host parks the lifecycle,
    // but the explicit pin vetoes the retry. The plugin must still receive
    // exactly one terminal event carrying the failure, not a bare Drop abort.
    assert_eq!(responses.len(), 1);
    assert!(responses[0].error.is_some());
    let attempts = dispatch.attempts.lock().unwrap().len();
    assert_eq!(
        attempts, 1,
        "an explicit pin must block the migration retry"
    );
    assert_eq!(
        tokio::time::timeout(Duration::from_secs(1), observations_rx.recv())
            .await
            .expect("classifier abort event timed out"),
        Some(ClassifierObservation::Aborted {
            cause: Some(ErrorType::WorkerOverloaded),
        })
    );
    assert!(observations_rx.try_recv().is_err());
    harness.runtime.shutdown();
}

/// Every attempt fails in-stream with a migratable worker error.
#[derive(Default)]
struct AlwaysOverloadedDispatch {
    attempts: Mutex<Vec<u64>>,
}

#[async_trait]
impl StreamingDispatch<PreprocessedRequest, Annotated<LLMEngineOutput>>
    for AlwaysOverloadedDispatch
{
    async fn generate(
        &self,
        request: SingleIn<AddressedRequest<PreprocessedRequest>>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        let (addressed, context) = request.transfer(());
        let (_, _, instance) = addressed.into_parts();
        self.attempts
            .lock()
            .unwrap()
            .push(instance.expect("selected worker instance").id());
        let output = Annotated {
            data: None,
            id: None,
            event: Some("error".to_string()),
            comment: None,
            error: Some(
                DynamoError::builder()
                    .error_type(ErrorType::WorkerOverloaded)
                    .message("selected worker is overloaded")
                    .build(),
            ),
        };
        Ok(ResponseStream::new(
            Box::pin(stream::once(async move { output })),
            context.context(),
        ))
    }

    async fn generate_bidirectional(
        &self,
        _instance: Instance,
        _address: String,
        _input: ManyIn<PreprocessedRequest>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        unreachable!("the routing host dispatches unary requests")
    }
}

#[tokio::test]
#[serial_test::serial]
async fn exhausted_migration_aborts_parked_lifecycle_with_worker_failure() {
    let calls = Arc::new(AtomicUsize::new(0));
    let (observations_tx, mut observations_rx) = mpsc::unbounded_channel();
    let dispatch = Arc::new(AlwaysOverloadedDispatch::default());
    let harness = two_worker_migration_harness(
        "exhausted-migration-classifier-lifecycle",
        dispatch.clone(),
        Some(RecordingClassifier {
            calls: Arc::clone(&calls),
            observations: observations_tx,
        }),
    )
    .await;
    let migration = Migration::new(1, None, "test".to_string(), Arc::new(Metrics::new()));

    let responses: Vec<_> = migration
        .generate(Context::new(request()), harness.engine.clone())
        .await
        .unwrap()
        .collect()
        .await;

    // Both attempts fail in-stream and the retry budget is spent. The client
    // receives the final attempt's worker error, so the classifier's single
    // terminal event must carry that same failure, not the synthetic
    // "migration limit exhausted" that retry accounting produces.
    assert_eq!(
        dispatch.attempts.lock().unwrap().len(),
        2,
        "one migration retry must run before the budget is exhausted"
    );
    let error = responses
        .last()
        .and_then(|response| response.error.as_ref())
        .expect("the final worker error reaches the client");
    assert!(match_error_chain(
        error,
        &[ErrorType::WorkerOverloaded],
        &[]
    ));
    assert_eq!(
        tokio::time::timeout(Duration::from_secs(1), observations_rx.recv())
            .await
            .expect("classifier abort event timed out"),
        Some(ClassifierObservation::Aborted {
            cause: Some(ErrorType::WorkerOverloaded),
        })
    );
    assert!(observations_rx.try_recv().is_err());
    harness.runtime.shutdown();
}

/// A gracefully shutting-down worker: `Cancelled` data frame, then the reason.
#[derive(Default)]
struct ShutdownAfterCancelDispatch {
    attempts: Mutex<Vec<(u64, Vec<u64>)>>,
}

#[async_trait]
impl StreamingDispatch<PreprocessedRequest, Annotated<LLMEngineOutput>>
    for ShutdownAfterCancelDispatch
{
    async fn generate(
        &self,
        request: SingleIn<AddressedRequest<PreprocessedRequest>>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        let (addressed, context) = request.transfer(());
        let (request, _, instance) = addressed.into_parts();
        let worker_id = instance.expect("selected worker instance").id();
        let excluded_worker_ids = request
            .migration_state
            .as_ref()
            .map(|state| state.excluded_worker_ids())
            .unwrap_or_default();
        let attempt = {
            let mut attempts = self.attempts.lock().unwrap();
            attempts.push((worker_id, excluded_worker_ids));
            attempts.len()
        };

        if attempt == 1 {
            return Ok(ResponseStream::new(
                Box::pin(stream::iter(vec![
                    cancelled_frame(),
                    engine_shutdown_frame(),
                ])),
                context.context(),
            ));
        }

        let output = Annotated::from_data(LLMEngineOutput {
            token_ids: vec![2],
            finish_reason: Some(FinishReason::Stop),
            ..Default::default()
        });
        Ok(ResponseStream::new(
            Box::pin(stream::once(async move { output })),
            context.context(),
        ))
    }

    async fn generate_bidirectional(
        &self,
        _instance: Instance,
        _address: String,
        _input: ManyIn<PreprocessedRequest>,
    ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
        unreachable!("the KV router dispatches unary requests")
    }
}

/// The shutdown error must reach migration, which is the only layer that can move the request.
#[tokio::test]
#[serial_test::serial]
async fn engine_shutdown_after_cancel_frame_migrates_and_reselects() {
    let dispatch = Arc::new(ShutdownAfterCancelDispatch::default());
    let harness =
        two_worker_migration_harness("engine-shutdown-migration", dispatch.clone(), None).await;
    let migration = Migration::new(1, None, "test".to_string(), Arc::new(Metrics::new()));

    let responses: Vec<_> = migration
        .generate(Context::new(request()), harness.engine.clone())
        .await
        .unwrap()
        .collect()
        .await;

    let attempts = {
        let attempts = dispatch.attempts.lock().unwrap();
        attempts.clone()
    };
    assert_eq!(attempts.len(), 2, "the shutdown must trigger a migration");
    let failed_worker = attempts[0].0;
    let retried_worker = attempts[1].0;
    assert_ne!(failed_worker, retried_worker);
    assert!(harness.registered_ids.contains(&failed_worker));
    assert!(harness.registered_ids.contains(&retried_worker));
    assert!(attempts[0].1.is_empty());
    assert_eq!(attempts[1].1, vec![failed_worker]);
    assert_eq!(
        responses.len(),
        1,
        "a migrated request must read as one clean stream; the aborted attempt's \
         terminal frame must not surface ahead of the retry: {responses:?}"
    );
    let last = responses.last().expect("the retry must produce output");
    assert!(last.error.is_none());
    assert_eq!(last.data.as_ref().unwrap().token_ids, vec![2]);
    assert_eq!(
        last.data.as_ref().unwrap().finish_reason,
        Some(FinishReason::Stop)
    );
    let loads = harness
        .chooser
        .get_potential_loads(&[], None, None, None, None)
        .await
        .unwrap();
    assert!(
        loads.iter().all(|load| load.active_requests == 0),
        "all scheduler bookings must be released after migration: {loads:?}"
    );
    harness.runtime.shutdown();
}

// A client that disconnects while remote prefill is still running. These enter
// through `RoutingHost::generate` with the context stopped on entry, which is the
// production sequence (see `prefill_router/mod.rs`, NVBugs 5969206).
//
// The dispatch mock must yield before recording. One that is ready on its first
// poll wins the biased race and stays green even on an unfixed build.

/// Build a builtin-plane routing host (round-robin, random, occupancy, direct)
/// whose dispatch plane records the workers it was asked to serve.
async fn builtin_host_with_recorded_dispatch(
    namespace: &str,
    mode: RouterMode,
) -> (RoutingHost, Arc<PendingThenCompletedDispatch>, u64, Runtime) {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace(namespace.to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    endpoint.register_endpoint_instance().await.unwrap();
    let worker_id = client.wait_for_instances().await.unwrap()[0].id();
    let load_context = test_load_context(&client).await;
    let dispatch = Arc::new(PendingThenCompletedDispatch::default());
    let inner = PushRouter::from_client_with_dispatch(
        client,
        mode,
        Arc::clone(&dispatch) as Arc<dyn StreamingDispatch<_, _>>,
    )
    .await
    .unwrap();
    let host = RoutingHost::new_builtin(inner, load_context).unwrap();
    (host, dispatch, worker_id, runtime)
}

/// A request in `phase` whose client has already gone away, ready to be handed
/// to `generate()` exactly as the prefill router hands over the decode leg.
///
/// `staged_kv` mirrors what the prefill router sets: true on the leg that
/// follows a completed remote prefill, false everywhere else — including the
/// conditional-disaggregation bypass, which reaches decode with nothing staged.
async fn stopped_request_in_phase(
    phase: RequestPhase,
    context_id: &str,
    target: Option<u64>,
    staged_kv: bool,
) -> SingleIn<PreprocessedRequest> {
    let tracker = Arc::new(RequestTracker::new());
    drop(tracker.set_phase(phase).await);
    let mut input = request();
    input.tracker = Some(Arc::clone(&tracker));
    input.staged_kv_cleanup = staged_kv;
    if let Some(worker_id) = target {
        input.routing_mut().backend_instance_id = Some(worker_id);
    }
    let request = Context::with_id_and_metadata(input, context_id.to_string(), Default::default());
    request.context().stop();
    assert!(request.context().is_stopped());
    request
}

/// The builtin plane, which is what the failing vLLM cancellation test uses
/// (`tests/fault_tolerance/cancellation/utils.py` sets `router_mode="round-robin"`).
/// Covers all four dispatch arms of `select_and_dispatch_builtin`.
#[tokio::test]
#[serial_test::serial]
async fn builtin_stopped_decode_request_reaches_the_worker_on_every_dispatch_path() {
    // (label, mode, pin the worker explicitly?)
    let cases: [(&str, RouterMode, bool); 4] = [
        ("round-robin", RouterMode::RoundRobin, false),
        ("least-loaded", RouterMode::LeastLoaded, false),
        ("explicit-target", RouterMode::RoundRobin, true),
        ("direct", RouterMode::Direct, true),
    ];

    for (index, (label, mode, pin_worker)) in cases.into_iter().enumerate() {
        let (host, dispatch, worker_id, runtime) = builtin_host_with_recorded_dispatch(
            &format!("builtin-decode-after-stop-{index}"),
            mode,
        )
        .await;
        let request = stopped_request_in_phase(
            RequestPhase::Decode,
            &format!("post-prefill-decode-cleanup-{label}"),
            pin_worker.then_some(worker_id),
            true,
        )
        .await;

        let mut stream = host.generate(request).await.unwrap_or_else(|error| {
            panic!("{label}: decode dispatch must survive a stopped context: {error:#}")
        });
        while stream.next().await.is_some() {}

        assert_eq!(
            dispatch.worker_ids.lock().unwrap().as_slice(),
            &[worker_id],
            "{label}: the decode worker must receive the request exactly once so it can run its KV-transfer cleanup"
        );

        drop(host);
        runtime.shutdown();
    }
}

/// The carve-out must stay narrow. Nothing without staged KV blocks may reach a
/// worker after the client has gone — including a decode request, which is what
/// the conditional-disaggregation bypass produces: `RequestPhase::Decode`
/// without a remote prefill behind it, so nothing to release.
#[tokio::test]
#[serial_test::serial]
async fn builtin_stopped_requests_without_staged_kv_never_reach_a_worker() {
    let cases = [
        ("prefill", RequestPhase::Prefill),
        ("aggregated", RequestPhase::Aggregated),
        ("conditional-disagg-bypass-decode", RequestPhase::Decode),
    ];

    for (index, (label, phase)) in cases.into_iter().enumerate() {
        let (host, dispatch, _worker_id, runtime) = builtin_host_with_recorded_dispatch(
            &format!("builtin-unstaged-after-stop-{index}"),
            RouterMode::RoundRobin,
        )
        .await;
        let request =
            stopped_request_in_phase(phase, &format!("builtin-cancelled-{label}"), None, false)
                .await;

        let error = host
            .generate(request)
            .await
            .expect_err("a stopped request with nothing staged must be cancelled before dispatch");
        assert!(
            match_error_chain(error.as_ref(), &[ErrorType::Cancelled], &[]),
            "{label}: expected a cancellation error, got {error:#}"
        );
        assert!(
            dispatch.worker_ids.lock().unwrap().is_empty(),
            "{label}: no worker should have been asked to serve a cancelled request"
        );

        drop(host);
        runtime.shutdown();
    }
}

/// The same guarantee on the KV plane, entered through `generate()`. This walks
/// worker selection, routing-decision recording, and dispatch in turn — the
/// stages a stopped decode context has to survive in production.
#[tokio::test]
#[serial_test::serial]
async fn kv_stopped_decode_request_reaches_the_worker_through_generate() {
    let (router, dispatch, worker_id, runtime) =
        router_with_recorded_dispatch("kv-decode-generate-after-stop").await;
    let request = stopped_request_in_phase(
        RequestPhase::Decode,
        "kv-post-prefill-decode-cleanup",
        None,
        true,
    )
    .await;

    let mut stream = router
        .generate(request)
        .await
        .expect("kv decode dispatch must survive an already-stopped context");
    while stream.next().await.is_some() {}

    assert_eq!(
        dispatch.worker_ids.lock().unwrap().as_slice(),
        &[worker_id],
        "the decode worker must receive the request so it can run its KV-transfer cleanup"
    );

    drop(router);
    runtime.shutdown();
}

/// One budget for the whole conditional route.
///
/// `preview_kv_route`, `plan_kv_route_from_preview` and `dispatch_kv_plan` run
/// in sequence on the bypass path. Each used to open its own budget, so a
/// stopped request could draw close to three full timeouts. Driving the real
/// chain here — rather than reusing one budget object by hand — is what makes a
/// regression visible: on the old code the plan reports a full budget again.
#[tokio::test]
#[serial_test::serial]
async fn conditional_route_stages_share_one_cleanup_budget() {
    // Real time rather than a paused clock: the router harness needs live
    // instance discovery, which never resolves under `start_paused`.
    let (router, runtime) = router(None).await;
    let request = Context::new(request());

    let preview = router
        .preview_kv_route(&request, RequestPhase::Decode)
        .await
        .unwrap();

    // Start the clock, then let measurable time pass before the next stage.
    let opening_budget = preview.cleanup_budget_remaining();
    let waited = Duration::from_millis(50);
    tokio::time::sleep(waited).await;

    let plan = router
        .plan_kv_route_from_preview(&request, preview)
        .await
        .unwrap();
    let planned_budget = plan.cleanup_budget_remaining();

    // A restarted budget reports the full constant again, so this is what
    // separates one shared budget from three fresh ones.
    assert!(
        planned_budget < opening_budget,
        "the plan restarted the budget: preview left {opening_budget:?}, plan reports {planned_budget:?}"
    );
    assert!(
        opening_budget - planned_budget >= waited,
        "time spent between stages must come out of the shared budget"
    );

    plan.abort().await;
    drop(router);
    runtime.shutdown();
}

/// The session-affinity wait runs upstream of every stage the cleanup policy
/// wraps, and it cancels on a stopped context of its own accord. A decode leg
/// with staged KV therefore used to die there — before the carve-out could
/// apply — whenever another request held the same session in `Initializing`.
#[tokio::test]
#[serial_test::serial]
async fn kv_stopped_decode_request_survives_a_contended_session_affinity_wait() {
    let (router, dispatch, worker_id, runtime) = router_with_recorded_dispatch_and_affinity(
        "kv-decode-affinity-contention",
        Some(Duration::from_secs(10)),
    )
    .await;
    let session_id = SessionAffinityId::new("contended-decode-session");

    // Another request is mid-initialization for the same session, so the decode
    // leg has to wait rather than take the slot immediately.
    let router = Arc::new(router);
    let coordinator = router.affinity.as_ref().unwrap().clone();
    let Hold::Initialize(holder) = coordinator.acquire(&session_id, None).await.unwrap() else {
        panic!("the first acquisition must initialize the session");
    };

    let mut input = request();
    let tracker = Arc::new(RequestTracker::new());
    drop(tracker.set_phase(RequestPhase::Decode).await);
    input.tracker = Some(Arc::clone(&tracker));
    input.staged_kv_cleanup = true;
    let mut decode_request = Context::with_id_and_metadata(
        input,
        "affinity-contended-decode".to_string(),
        Default::default(),
    );
    decode_request.insert(SESSION_AFFINITY_CONTEXT_KEY, session_id.clone());
    decode_request.context().stop();

    let generate_router = Arc::clone(&router);
    let generate = tokio::spawn(async move { generate_router.generate(decode_request).await });

    // Only release the session once the decode leg is provably parked on the
    // wait; otherwise the test could pass without ever exercising it.
    coordinator.wait_for_initializing_waiter().await;
    drop(holder);

    let mut stream = generate
        .await
        .unwrap()
        .expect("a staged-KV decode leg must survive the affinity wait");
    while stream.next().await.is_some() {}

    assert_eq!(
        dispatch.worker_ids.lock().unwrap().as_slice(),
        &[worker_id],
        "the decode worker must receive the request so it can run its KV-transfer cleanup"
    );

    drop(router);
    runtime.shutdown();
}

/// The same narrowing on the KV plane. A bypass decode request holds nothing a
/// worker needs to release, so a disconnect is abandoned immediately.
#[tokio::test]
#[serial_test::serial]
async fn kv_stopped_decode_request_without_staged_kv_never_reaches_a_worker() {
    let (router, dispatch, _worker_id, runtime) =
        router_with_recorded_dispatch("kv-bypass-decode-after-stop").await;
    let request =
        stopped_request_in_phase(RequestPhase::Decode, "kv-bypass-decode", None, false).await;

    let error = router
        .generate(request)
        .await
        .expect_err("a bypass decode request with nothing staged must be cancelled");
    assert!(match_error_chain(
        error.as_ref(),
        &[ErrorType::Cancelled],
        &[]
    ));
    assert!(
        dispatch.worker_ids.lock().unwrap().is_empty(),
        "no worker should have been asked to serve a cancelled bypass request"
    );

    drop(router);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn unknown_explicit_workers_are_rejected_before_builtin_dispatch() {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let component = distributed
        .namespace("unknown-explicit-workers")
        .unwrap()
        .component("workers")
        .unwrap();
    let endpoint = component.endpoint("generate");
    let client = endpoint.client().await.unwrap();
    let load_context = test_load_context(&client).await;
    endpoint.register_endpoint_instance().await.unwrap();
    let live_worker = client.wait_for_instances().await.unwrap()[0].id();
    let dispatch = Arc::new(CompletedBuiltinDispatch::default());
    let inner = PushRouter::from_client_with_dispatch(
        client,
        RouterMode::RoundRobin,
        Arc::clone(&dispatch) as Arc<dyn StreamingDispatch<_, _>>,
    )
    .await
    .unwrap();
    let host = RoutingHost::new_builtin(inner, load_context).unwrap();
    let worker_id = live_worker.wrapping_add(1);
    for (field, phase) in [
        ("backend_instance_id", RequestPhase::Aggregated),
        ("decode_worker_id", RequestPhase::Decode),
        ("prefill_worker_id", RequestPhase::Prefill),
    ] {
        let mut content = request();
        content.routing =
            Some(serde_json::from_value(serde_json::json!({field: worker_id})).unwrap());
        let error = host
            .select_and_dispatch_builtin(Context::new(content), phase, |_, _| Ok(()))
            .await
            .unwrap_err();
        assert!(match_error_chain(
            error.as_ref(),
            &[ErrorType::InvalidArgument],
            &[]
        ));
        assert_eq!(
            error.downcast_ref::<DynamoError>().unwrap().message(),
            format!("nvext.{field}={worker_id} does not identify a known worker")
        );
    }
    assert!(dispatch.worker_ids.lock().unwrap().is_empty());
    let mut valid = request();
    valid.routing_mut().decode_worker_id = Some(live_worker);
    let mut stream = host.generate(Context::new(valid)).await.unwrap();
    while stream.next().await.is_some() {}
    assert_eq!(
        dispatch.worker_ids.lock().unwrap().as_slice(),
        &[live_worker]
    );
    drop(host);
    runtime.shutdown();
}

#[tokio::test]
async fn unknown_explicit_workers_are_rejected_before_kv_admission() {
    let (host, dispatch, worker_id, runtime) =
        router_with_recorded_dispatch("unknown-kv-worker").await;
    let unknown = worker_id.wrapping_add(1);
    assert!(!host.inner.client.is_instance_live(unknown));
    for (field, phase) in [
        ("backend_instance_id", RequestPhase::Aggregated),
        ("decode_worker_id", RequestPhase::Decode),
        ("prefill_worker_id", RequestPhase::Prefill),
    ] {
        let mut content = request();
        content.routing =
            Some(serde_json::from_value(serde_json::json!({field: unknown})).unwrap());
        let request = Context::new(content);
        let error = host
            .select_with_affinity(&request, phase, false, &CleanupBudget::default())
            .await
            .err()
            .expect("unknown worker must fail before admission");
        assert!(match_error_chain(
            error.as_ref(),
            &[ErrorType::InvalidArgument],
            &[]
        ));
        assert_eq!(
            error.downcast_ref::<DynamoError>().unwrap().message(),
            format!("nvext.{field}={unknown} does not identify a known worker")
        );
        let error = host
            .preview_kv_route(&request, phase)
            .await
            .err()
            .expect("unknown preview target");
        assert!(match_error_chain(
            error.as_ref(),
            &[ErrorType::InvalidArgument],
            &[]
        ));
    }
    assert!(dispatch.worker_ids.lock().unwrap().is_empty());
    assert!(
        potential_loads(&host)
            .await
            .iter()
            .all(|load| load.active_requests == 0)
    );
    drop(host);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn explicit_worker_validation_preserves_precedence_and_discovery_membership() {
    let (host, runtime) = router_with_workers(None, &[u64::MAX]).await;
    let cases = [
        (serde_json::json!({}), RequestPhase::Aggregated),
        (
            serde_json::json!({"backend_instance_id": null, "decode_worker_id": null,
            "prefill_worker_id": null, "dp_rank": null}),
            RequestPhase::Aggregated,
        ),
        (
            serde_json::json!({"backend_instance_id": u64::MAX}),
            RequestPhase::Prefill,
        ),
        (
            serde_json::json!({"backend_instance_id": u64::MAX, "decode_worker_id": null}),
            RequestPhase::Decode,
        ),
        (
            serde_json::json!({"backend_instance_id": 42, "decode_worker_id": u64::MAX,
            "prefill_worker_id": 42}),
            RequestPhase::Aggregated,
        ),
        (
            serde_json::json!({"backend_instance_id": 42, "decode_worker_id": 42,
            "prefill_worker_id": u64::MAX}),
            RequestPhase::Prefill,
        ),
    ];
    host.inner.client.override_instance_avail(vec![]);
    for (value, phase) in cases {
        let _: crate::protocols::common::extensions::NvExt =
            serde_json::from_value(value.clone()).unwrap();
        let mut content = request();
        content.routing = Some(serde_json::from_value(value).unwrap());
        host.validate_explicit_worker(&content, phase).unwrap();
    }
    for (field, phase) in [
        ("decode_worker_id", RequestPhase::Aggregated),
        ("prefill_worker_id", RequestPhase::Prefill),
    ] {
        let mut content = request();
        content.routing = Some(
            serde_json::from_value(serde_json::json!({
                "backend_instance_id": u64::MAX, field: 42
            }))
            .unwrap(),
        );
        let error = host.validate_explicit_worker(&content, phase).unwrap_err();
        assert_eq!(
            error.downcast_ref::<DynamoError>().unwrap().message(),
            format!("nvext.{field}=42 does not identify a known worker")
        );
    }
    drop(host);
    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn explicit_worker_disappearing_after_preview_is_not_revalidated() {
    let (host, runtime) = router(None).await;
    let mut content = request();
    content.routing_mut().decode_worker_id = Some(7);
    let request = Context::new(content);
    let preview = host
        .preview_kv_route(&request, RequestPhase::Decode)
        .await
        .unwrap();
    host.inner.client.override_discovered_instances(vec![]);
    // A failure after successful selection must keep its service-error classification.
    let error = match host.plan_kv_route_from_preview(&request, preview).await {
        Ok(plan) => host.dispatch_kv_plan(request, plan).await.unwrap_err(),
        Err(error) => error,
    };
    assert!(!match_error_chain(
        error.as_ref(),
        &[ErrorType::InvalidArgument],
        &[]
    ));
    drop(host);
    runtime.shutdown();
}

fn subagent_request(
    session_id: &str,
    parent_session_id: Option<&str>,
) -> SingleIn<PreprocessedRequest> {
    let mut agent_context = AgentContextBuilder::default();
    agent_context.session_id(session_id.to_string());
    if let Some(parent_session_id) = parent_session_id {
        agent_context.parent_session_id(parent_session_id.to_string());
    }
    let mut content = request();
    content.agent_context = Some(agent_context.build().unwrap());
    let mut request = Context::new(content);
    request.insert(
        SESSION_AFFINITY_CONTEXT_KEY,
        SessionAffinityId::new(session_id),
    );
    request
}

fn parent_group_id(parent_session_id: &str) -> SessionAffinityId {
    SessionAffinityId::new(crate::session_affinity::subagent_group_affinity_id(
        parent_session_id,
    ))
}

async fn affinity_mode_host(namespace: &str) -> (Runtime, RoutingHost) {
    let runtime = Runtime::from_current().unwrap();
    let distributed = DistributedRuntime::new(runtime.clone(), DistributedConfig::process_local())
        .await
        .unwrap();
    let endpoint = distributed
        .namespace(namespace.to_string())
        .unwrap()
        .component("workers".to_string())
        .unwrap()
        .endpoint("generate".to_string());
    let client = endpoint.client().await.unwrap();
    endpoint.register_endpoint_instance().await.unwrap();
    client.wait_for_instances().await.unwrap();
    let load_context = test_load_context(&client).await;
    let inner = PushRouter::from_client(client, RouterMode::RoundRobin)
        .await
        .unwrap();
    let coordinator = AffinityCoordinator::new(
        Duration::from_secs(60),
        crate::session_affinity::SessionAffinityMode::Hard,
    )
    .unwrap();
    let host =
        RoutingHost::new_builtin_with_coordinator(inner, load_context, Some(coordinator)).unwrap();
    (runtime, host)
}

#[tokio::test]
#[serial_test::serial]
async fn parent_group_binding_resolves_siblings_to_one_key() {
    let (runtime, host) = affinity_mode_host("subagent-group-binding").await;

    let key = |request: &SingleIn<PreprocessedRequest>| {
        host.group_binding_id(request, None)
            .unwrap_or_else(|| {
                crate::session_affinity::affinity_id(request)
                    .unwrap()
                    .unwrap()
            })
            .as_str()
            .to_string()
    };

    let sibling_a = key(&subagent_request("child-1", Some("parent-1")));
    let sibling_b = key(&subagent_request("child-2", Some("parent-1")));
    let other_parent = key(&subagent_request("child-3", Some("parent-2")));
    let main_agent = key(&subagent_request("parent-1", None));

    assert_eq!(sibling_a, sibling_b);
    assert_ne!(sibling_a, other_parent);
    assert_ne!(sibling_a, main_agent);
    assert_eq!(main_agent, "parent-1");

    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn parent_group_binding_offers_siblings_the_committed_worker() {
    let workers = [7u64, 9]
        .into_iter()
        .map(|worker_id| (worker_id, ModelRuntimeConfig::default()))
        .collect();
    let (router, runtime) =
        router_with_worker_configs(Some(Duration::from_secs(60)), workers).await;

    let offered = async |request: &SingleIn<PreprocessedRequest>| {
        router
            .select_with_session_affinity(
                request,
                RequestPhase::Aggregated,
                false,
                &CleanupBudget::default(),
                |target| std::future::ready(Ok(target)),
            )
            .await
            .unwrap()
    };

    let first = subagent_request("child-1", Some("parent-1"));
    let (offered_first, hold) = offered(&first).await;
    assert_eq!(offered_first, None, "a new group starts unbound");
    let mut stream = router
        .bind_affinity(
            hold,
            AffinityTarget::new(7, Some(0)),
            ResponseStream::new(
                Box::pin(stream::iter([Annotated::from_data(
                    LLMEngineOutput::default(),
                )])),
                first.context(),
            ),
        )
        .unwrap();
    while stream.next().await.is_some() {}

    let sibling = subagent_request("child-2", Some("parent-1"));
    let (offered_sibling, _) = offered(&sibling).await;
    assert_eq!(offered_sibling, Some(AffinityTarget::new(7, Some(0))));

    let other_parent = subagent_request("child-3", Some("parent-2"));
    let (offered_other, _) = offered(&other_parent).await;
    assert_eq!(
        offered_other, None,
        "a different parent is a different group"
    );

    runtime.shutdown();
}

#[tokio::test]
#[serial_test::serial]
async fn hard_parent_group_recovers_when_the_bound_worker_leaves() {
    let workers = [7u64, 9]
        .into_iter()
        .map(|worker_id| (worker_id, ModelRuntimeConfig::default()))
        .collect();
    let (router, runtime) =
        router_with_worker_configs(Some(Duration::from_secs(60)), workers).await;
    let affinity = router.affinity.as_ref().unwrap();
    let group = parent_group_id("parent-1");

    let Hold::Initialize(initializer) = affinity.acquire(&group, None).await.unwrap() else {
        panic!("the first subagent must initialize the group");
    };
    drop(
        initializer
            .commit(crate::session_affinity::to_table(AffinityTarget::new(
                7,
                Some(0),
            )))
            .unwrap(),
    );
    assert_eq!(
        affinity.query_target(&group, None).unwrap(),
        Some(AffinityTarget::new(7, Some(0)))
    );

    // Worker 7 leaves the pool. Hard mode must invalidate the dead pin and retry unbound rather
    // than fail every sibling until the TTL expires.
    router.inner.client.override_discovered_instances(vec![9]);
    router.inner.client.override_instance_avail(vec![9]);
    let sibling = subagent_request("child-2", Some("parent-1"));
    let (selected, operation) = router
        .select_with_session_affinity(
            &sibling,
            RequestPhase::Aggregated,
            false,
            &CleanupBudget::default(),
            |target: Option<AffinityTarget>| {
                // The dead pin is rejected at dispatch time; the unbound retry is accepted.
                std::future::ready(match target {
                    Some(target) if target.worker_id == 7 => Err(anyhow::anyhow!("worker gone")),
                    other => Ok(other),
                })
            },
        )
        .await
        .unwrap();
    assert_eq!(selected, None, "the retry runs unbound");
    assert!(matches!(operation, Some(Hold::Initialize(_))));
    assert_eq!(affinity.query_target(&group, None).unwrap(), None);

    runtime.shutdown();
}