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

//! Thin neutral contract adapter over the mechanically moved scheduler cores.

use std::collections::{HashMap, VecDeque};
use std::sync::Arc;

use anyhow::{Result, anyhow, ensure};
use uuid::Uuid;

use crate::engine::common::perf_model::PerfModel;
use crate::engine::common::protocols::{
    DirectRequest, EngineType, KvTransferTimingMode, MockEngineArgs,
    PreemptionMode as CorePreemptionMode, SglangArgs, WorkerType as CoreWorkerType,
};
use crate::engine::generalized::{CommandContext, RankEngine, RankIdentity, RankPass};
use crate::engine::host_offload::{G2Binding, HostTierHandle};
use crate::engine::{
    Admission, Backend, Command, CommandEffects, CommandResult, EngineConfig, ForwardPassMetrics,
    HandoffId, HostOffloadObserver, KvHintSource, LifecycleEvent, Metrics, Output,
    PassCompletionEffects, PassStartEffects, PendingPass, PreemptionMode, Request, TimingModel,
    TransferTimingMode, WorkerType,
};

use super::{
    EngineCore, EnginePassResult, KvEventVisibility, MockerMetrics,
    SchedulerCommand as CoreCommand, SchedulerCommandEffects as CoreCommandEffects,
    SchedulerCommandResult as CoreCommandResult, SchedulerLifecycleEvent as CoreLifecycle,
    SglangCore, VllmCore,
};

pub fn engine_seed_offset(identity: RankIdentity) -> Result<u64> {
    identity
        .worker_id
        .checked_mul(u64::from(identity.dp_size.get()))
        .and_then(|base| base.checked_add(u64::from(identity.dp_rank)))
        .ok_or_else(|| anyhow!("native mock-engine seed offset overflow"))
}

/// One preserved vLLM/SGLang scheduler rank behind the neutral contract.
pub struct SchedulerRank {
    core: EngineCore,
    handoff_requests: HashMap<HandoffId, Uuid>,
    /// vLLM `launch`: outputs the engine loop has not returned yet.
    held_outputs: Option<HeldOutputs>,
}

/// vLLM's engine loop returns a pass's outputs only after it launches the
/// next pass (`launch`), which may be after the pass ends.
#[derive(Default)]
struct HeldOutputs {
    /// The latest pass's end and outputs, until the next pass is formed.
    latest: Option<(f64, Vec<Output>)>,
    /// Outputs by the time the engine loop returns them, in pass order.
    returned: VecDeque<(f64, Vec<Output>)>,
}

impl HeldOutputs {
    /// Drop a cancelled request's outputs; true when any was held.
    fn discard(&mut self, request_id: Uuid) -> bool {
        let mut dropped = false;
        for (_, outputs) in self.latest.iter_mut().chain(self.returned.iter_mut()) {
            let before = outputs.len();
            outputs.retain(|output| output.request_id != request_id);
            dropped |= outputs.len() != before;
        }
        self.latest.take_if(|(_, outputs)| outputs.is_empty());
        self.returned.retain(|(_, outputs)| !outputs.is_empty());
        dropped
    }

    fn is_empty(&self) -> bool {
        self.latest.is_none() && self.returned.is_empty()
    }
}

impl SchedulerRank {
    pub(crate) fn set_belady_oracle(&mut self, oracle: crate::engine::belady::BeladyOracle) {
        match &mut self.core {
            EngineCore::Vllm(core) => core.set_belady_oracle(oracle),
            EngineCore::Sglang(core) => core.set_belady_oracle(oracle),
        }
    }

    pub(crate) fn set_g3_offload(
        &mut self,
        registry: crate::engine::g3_offload::SharedG3Tier,
        node: usize,
    ) {
        self.core.set_g3_offload(registry, node);
    }

    /// Pull hinted prefixes from peer ranks' G2. Returns this rank's own G2,
    /// or `None` for an engine without KVCR.
    pub(crate) fn set_kvcr(
        &mut self,
        registry: crate::engine::kvcr::SharedKvcrRegistry,
        rank: KvHintSource,
    ) -> Option<HostTierHandle> {
        self.core.set_kvcr(registry, rank)
    }

    pub(crate) fn set_host_offload_observer(&mut self, observer: Arc<dyn HostOffloadObserver>) {
        self.core.set_host_offload_observer(observer);
    }

    /// Per-rank G2 only. Shared G2 is observed through its deployment registry.
    pub(crate) fn per_rank_host_cache_metrics(&self) -> Option<crate::engine::HostCacheMetrics> {
        match &self.core {
            EngineCore::Vllm(core) => core.per_rank_host_cache_metrics(),
            EngineCore::Sglang(_) => None,
        }
    }

    pub fn new_with_timing_model(
        identity: RankIdentity,
        config: &EngineConfig,
        timing: Arc<dyn TimingModel>,
        seed_offset: u64,
    ) -> Result<Self> {
        Self::new_with_g2(identity, config, timing, seed_offset, None)
    }

    /// `g2` binds a cluster-shared G2 configuration to its deployment pool.
    pub(crate) fn new_with_g2(
        identity: RankIdentity,
        config: &EngineConfig,
        timing: Arc<dyn TimingModel>,
        seed_offset: u64,
        g2: Option<&G2Binding>,
    ) -> Result<Self> {
        config.validate()?;
        ensure!(
            config.g3_offload.is_none(),
            "g3_offload is Replay-owned; construct it through ReplaySpec"
        );
        ensure!(
            config.kvcr.is_none(),
            "kvcr is Replay-owned; construct it through ReplaySpec"
        );
        let mut args = core_args(config, timing);
        if config.backend == Backend::Sglang {
            normalize_sglang_attention_dp(&mut args, identity.dp_size.get())?;
        }
        let capture_kv_events = config.emit_kv_events;
        let core = match config.backend {
            Backend::Vllm | Backend::Trtllm => EngineCore::Vllm(VllmCore::new_with_worker_rank(
                args,
                identity.worker_id,
                identity.dp_rank,
                seed_offset,
                capture_kv_events,
                g2,
            )?),
            Backend::Sglang => EngineCore::Sglang(SglangCore::new_with_worker_rank(
                args,
                identity.worker_id,
                identity.dp_rank,
                seed_offset,
                capture_kv_events,
            )),
        };
        Ok(Self {
            core,
            handoff_requests: HashMap::new(),
            held_outputs: config.launch.map(|_| HeldOutputs::default()),
        })
    }

    fn core_command(command: Command) -> CoreCommand {
        match command {
            Command::Submit(request) => CoreCommand::Submit(core_request(request)),
            Command::CancelRequest { request_id, .. } => CoreCommand::CancelRequest { request_id },
            Command::SubmitHandoffPrefill {
                handoff_id,
                request,
            } => CoreCommand::SubmitHandoffPrefill {
                handoff_id,
                request: core_request(request),
            },
            Command::ReserveDestination {
                handoff_id,
                request,
            } => CoreCommand::ReserveDestination {
                handoff_id,
                request: core_request(request),
            },
            Command::ActivateDestination { handoff_id } => {
                CoreCommand::ActivateDestination { handoff_id }
            }
            Command::ReleaseSource { handoff_id } => CoreCommand::ReleaseSource { handoff_id },
            Command::CancelSource { handoff_id } => CoreCommand::CancelSource { handoff_id },
            Command::CancelDestination { handoff_id } => {
                CoreCommand::CancelDestination { handoff_id }
            }
        }
    }

    fn metrics(&self) -> Metrics {
        let metrics = match &self.core {
            EngineCore::Vllm(core) => core.mocker_metrics(),
            EngineCore::Sglang(core) => core.mocker_metrics(),
        };
        map_metrics(metrics)
    }
}

impl RankEngine for SchedulerRank {
    type Config = EngineConfig;
    type Command = Command;
    type CommandEffects = CommandEffects;
    type PassStartEffects = PassStartEffects;
    type PendingPass = PendingPass;
    type PassCompletionEffects = PassCompletionEffects;
    type InternalEffects = PassStartEffects;

    fn new(identity: RankIdentity, config: &Self::Config) -> Result<Self> {
        let timing = config.built_in_timing_model()?;
        let seed_offset = engine_seed_offset(identity)?;
        Self::new_with_timing_model(identity, config, timing, seed_offset)
    }

    fn apply_command_effects(
        &mut self,
        command: Self::Command,
        context: CommandContext,
        pending_pass: Option<&mut Self::PendingPass>,
    ) -> Result<Self::CommandEffects> {
        let pending_suppression = pending_output_suppression(&command, &self.handoff_requests);
        let handoff_update = handoff_tracking_update(&command);
        let core_command = Self::core_command(command);
        let mut effects = self.core.apply_command_effects_at(
            core_command,
            context.allow_immediate_admission(),
            context.now_ms,
        )?;
        // Preserve the scheduler's command boundary when no model step is in
        // flight: native G1 mutations produced by the command belong to its
        // returned effects. Mid-pass mutations remain buffered so neither a
        // command nor a due physical transfer can expose KV state before the
        // shared completion boundary.
        if !context.pass_in_flight {
            effects.kv_events.extend(self.core.drain_kv_events());
        }
        let suppressed_pending_output = if let (Some((request_id, discard_on_noop)), Some(pending)) =
            (pending_suppression, pending_pass)
            && (effects.result != CoreCommandResult::Noop || discard_on_noop)
        {
            let before = pending.effects.outputs.len();
            pending
                .effects
                .outputs
                .retain(|output| output.request_id != request_id);
            pending
                .effects
                .lifecycle_events
                .retain(|event| match *event {
                    LifecycleEvent::SourceHeld { request_id: id, .. }
                    | LifecycleEvent::DestinationReserved { request_id: id, .. }
                    | LifecycleEvent::TtftMilestone { request_id: id, .. } => id != request_id,
                });
            before != pending.effects.outputs.len()
        } else {
            false
        };
        let suppressed_pending_output = suppressed_pending_output
            | match (pending_suppression, self.held_outputs.as_mut()) {
                (Some((request_id, discard_on_noop)), Some(held))
                    if effects.result != CoreCommandResult::Noop || discard_on_noop =>
                {
                    held.discard(request_id)
                }
                _ => false,
            };
        if suppressed_pending_output
            && let Some((request_id, _)) = pending_suppression
            && !effects.retired_requests.contains(&request_id)
        {
            // A final output can be suppressed after the native scheduler has
            // already retired its request. Replay still owns its accounting
            // until the pass completion is observed, so publish the same
            // retirement delta that completion would have carried.
            effects.retired_requests.push(request_id);
        }
        if effects.result != CoreCommandResult::Noop || suppressed_pending_output {
            self.apply_handoff_tracking_update(handoff_update);
        }
        for request_id in &effects.retired_requests {
            self.handoff_requests
                .retain(|_, tracked_request| tracked_request != request_id);
        }
        map_command_effects(effects, self.metrics(), suppressed_pending_output)
    }

    fn is_ready(&self) -> bool {
        self.core.is_ready()
    }

    fn waiting_for_external_command(&self) -> bool {
        self.core.waiting_for_external_command()
    }

    fn prepare_group_pass(&mut self, wave_step: u64, dp_size: std::num::NonZeroU32) {
        self.core.prepare_group_pass(wave_step, dp_size.get());
    }

    fn prefill_in_pass(&self) -> bool {
        self.core.prefill_in_pass()
    }

    fn model_work_in_pass(&self) -> bool {
        self.core.model_work_in_pass()
    }

    fn finish_group_pass(&mut self, any_rank_prefilled: bool, any_rank_ran_model: bool) {
        self.core
            .finish_group_pass(any_rank_prefilled, any_rank_ran_model);
    }

    fn execute_pass(
        &mut self,
        now_ms: f64,
    ) -> Result<RankPass<Self::PassStartEffects, Self::PendingPass>> {
        let pass = self.core.try_execute_hidden_pass(now_ms)?;
        let end_ms = pass.end_ms;
        let previous_returned_ms = pass.previous_returned_ms;
        let (same_timestamp_retry, start_effects, completion_effects) = split_pass(pass)?;
        if let Some(held) = self.held_outputs.as_mut()
            && let Some((previous_end_ms, outputs)) = held.latest.take()
        {
            held.returned
                .push_back((previous_returned_ms.unwrap_or(previous_end_ms), outputs));
        }
        Ok(RankPass {
            end_ms,
            same_timestamp_retry,
            start_effects,
            pending: PendingPass {
                started_at_ms: now_ms,
                effects: completion_effects,
            },
        })
    }

    fn complete_pass(
        &mut self,
        mut pending: Self::PendingPass,
        end_ms: f64,
    ) -> Result<Self::PassCompletionEffects> {
        self.core.complete_engine_boundary(end_ms);
        // The preserved scheduler retries deferred destination reservations
        // when a forward pass releases capacity. Keep that wakeup at the
        // pass-completion boundary: command-time retry is suppressed while a
        // pass is in flight, and waiting until an unrelated later pass can
        // leave disaggregated replay permanently asleep.
        pending.effects.lifecycle_events.extend(
            self.core
                .retry_pending_destinations()
                .into_iter()
                .map(map_lifecycle),
        );
        let completion_kv_events = self.core.drain_kv_events();
        pending.effects.kv_events.extend(completion_kv_events);
        // Occupancy is authoritative at the shared completion boundary, but
        // SGLang cache hit/total are transient observations from this pass.
        // Preserve those fields while refreshing the rest of the snapshot;
        // a later live adapter latches the last non-empty observation.
        let sglang_cache_hit_tokens = pending.effects.metrics.sglang_cache_hit_tokens;
        let sglang_cache_total_tokens = pending.effects.metrics.sglang_cache_total_tokens;
        pending.effects.metrics = self.metrics();
        pending.effects.metrics.sglang_cache_hit_tokens = sglang_cache_hit_tokens;
        pending.effects.metrics.sglang_cache_total_tokens = sglang_cache_total_tokens;
        // Under the SGLang host loop a pass ends when the scheduler thread returns
        // to its loop, not when its forward does; FPM telemetry wants the forward.
        pending.effects.forward_pass_metrics.duration_ms = self
            .core
            .last_forward_ms()
            .unwrap_or((end_ms - pending.started_at_ms).max(0.0));
        for output in &pending.effects.outputs {
            if output.completed {
                self.handoff_requests
                    .retain(|_, request_id| *request_id != output.request_id);
            }
        }
        // vLLM `launch`: the engine loop returns this pass's outputs once it
        // has launched the next pass, which it forms now if it has work.
        if let Some(held) = self.held_outputs.as_mut()
            && !pending.effects.outputs.is_empty()
            && self.core.is_ready()
            && !self.core.waiting_for_external_command()
        {
            held.latest = Some((end_ms, std::mem::take(&mut pending.effects.outputs)));
        }
        Ok(pending.effects)
    }

    fn complete_idle_group_pass(
        &mut self,
        started_at_ms: f64,
        end_ms: f64,
    ) -> Result<Option<Self::PassCompletionEffects>> {
        self.core.complete_engine_boundary(end_ms);
        let lifecycle_events = self
            .core
            .retry_pending_destinations()
            .into_iter()
            .map(map_lifecycle)
            .collect::<Vec<_>>();
        let kv_events = self.core.drain_kv_events();
        Ok(Some(PassCompletionEffects {
            lifecycle_events,
            kv_events,
            metrics: self.metrics(),
            forward_pass_metrics: ForwardPassMetrics {
                duration_ms: (end_ms - started_at_ms).max(0.0),
                ..Default::default()
            },
            ..PassCompletionEffects::default()
        }))
    }

    fn next_internal_deadline_ms(&self) -> Option<f64> {
        self.core.next_internal_deadline_ms()
    }

    fn process_internal_work(
        &mut self,
        now_ms: f64,
        pass_in_flight: bool,
    ) -> Result<Self::InternalEffects> {
        if pass_in_flight {
            return Ok(PassStartEffects::default());
        }
        self.core.process_internal_work(now_ms);
        Ok(PassStartEffects {
            kv_events: self.core.drain_kv_events(),
            ..PassStartEffects::default()
        })
    }

    fn is_drained(&self) -> bool {
        self.core.is_drained() && self.held_outputs.as_ref().is_none_or(HeldOutputs::is_empty)
    }

    fn next_emission_ms(&self) -> Option<f64> {
        let held = self.held_outputs.as_ref()?;
        held.returned.front().map(|(at_ms, _)| *at_ms)
    }

    fn take_due_emissions(&mut self, now_ms: f64) -> Result<Option<Self::PassCompletionEffects>> {
        let Some(held) = self.held_outputs.as_mut() else {
            return Ok(None);
        };
        let mut outputs = Vec::new();
        while held
            .returned
            .front()
            .is_some_and(|(at_ms, _)| *at_ms <= now_ms)
        {
            let (_, due) = held.returned.pop_front().expect("peeked returned outputs");
            outputs.extend(due);
        }
        if outputs.is_empty() {
            return Ok(None);
        }
        Ok(Some(PassCompletionEffects {
            outputs,
            metrics: self.metrics(),
            ..PassCompletionEffects::default()
        }))
    }
}

#[derive(Clone, Copy)]
enum HandoffTrackingUpdate {
    None,
    Insert(HandoffId, Uuid),
    RemoveHandoff(HandoffId),
    RemoveRequest(Uuid),
}

impl SchedulerRank {
    fn apply_handoff_tracking_update(&mut self, update: HandoffTrackingUpdate) {
        match update {
            HandoffTrackingUpdate::None => {}
            HandoffTrackingUpdate::Insert(handoff_id, request_id) => {
                self.handoff_requests.insert(handoff_id, request_id);
            }
            HandoffTrackingUpdate::RemoveHandoff(handoff_id) => {
                self.handoff_requests.remove(&handoff_id);
            }
            HandoffTrackingUpdate::RemoveRequest(request_id) => self
                .handoff_requests
                .retain(|_, tracked_request| *tracked_request != request_id),
        }
    }
}

fn handoff_tracking_update(command: &Command) -> HandoffTrackingUpdate {
    match command {
        Command::SubmitHandoffPrefill {
            handoff_id,
            request,
        }
        | Command::ReserveDestination {
            handoff_id,
            request,
        } => HandoffTrackingUpdate::Insert(*handoff_id, request.request_id),
        Command::ReleaseSource { handoff_id }
        | Command::CancelSource { handoff_id }
        | Command::CancelDestination { handoff_id } => {
            HandoffTrackingUpdate::RemoveHandoff(*handoff_id)
        }
        Command::CancelRequest { request_id, .. } => {
            HandoffTrackingUpdate::RemoveRequest(*request_id)
        }
        Command::Submit(_) | Command::ActivateDestination { .. } => HandoffTrackingUpdate::None,
    }
}

fn core_args(config: &EngineConfig, timing: Arc<dyn TimingModel>) -> MockEngineArgs {
    MockEngineArgs {
        engine_type: match config.backend {
            Backend::Vllm => EngineType::Vllm,
            Backend::Sglang => EngineType::Sglang,
            Backend::Trtllm => EngineType::Trtllm,
        },
        num_gpu_blocks: config.num_gpu_blocks,
        block_size: config.block_size,
        max_model_len: config.max_model_len,
        max_num_seqs: Some(config.max_num_seqs),
        max_num_batched_tokens: Some(config.max_num_batched_tokens),
        prefill_schedule_interval: config.prefill_schedule_interval,
        prefill_decode_interval: config.prefill_decode_interval,
        max_concurrent_batches: config.max_concurrent_batches(),
        launch: config.launch,
        enable_prefix_caching: config.enable_prefix_caching,
        enable_chunked_prefill: config.enable_chunked_prefill,
        speedup_ratio: config.speedup_ratio,
        decode_speedup_ratio: config.decode_speedup_ratio,
        worker_type: match config.worker_type {
            WorkerType::Aggregated => CoreWorkerType::Aggregated,
            WorkerType::Prefill => CoreWorkerType::Prefill,
            WorkerType::Decode => CoreWorkerType::Decode,
        },
        perf_model: Arc::new(PerfModel::External { timing }),
        aic_nextn: config.aic_nextn,
        aic_nextn_accept_rates: config.aic_nextn_accept_rates.clone(),
        aic_mtp_seed: config.aic_mtp_seed,
        kv_transfer_bytes_per_token: config.kv_transfer_bytes_per_token,
        kv_cache_bytes_per_token: config.kv_cache_bytes_per_token,
        kv_cache_groups: config.kv_cache_groups.clone(),
        kv_cache_capacity_bytes: config.kv_cache_capacity_bytes,
        native_host_offload: config.native_host_offload.clone(),
        state_cache: config.state_cache,
        prefix_match_unit: config.prefix_match_unit,
        kv_transfer_bandwidth: config.kv_transfer_bandwidth,
        kv_transfer_timing_mode: match config.kv_transfer_timing_mode {
            TransferTimingMode::FullPrompt => KvTransferTimingMode::FullPrompt,
            TransferTimingMode::DestinationMissing => KvTransferTimingMode::DestinationMissing,
        },
        preemption_mode: match config.preemption_mode {
            PreemptionMode::Lifo => CorePreemptionMode::Lifo,
            PreemptionMode::Fifo => CorePreemptionMode::Fifo,
        },
        sglang: Some(SglangArgs {
            schedule_policy: Some(
                match config.sglang.schedule_policy {
                    crate::engine::SglangSchedulePolicy::Fifo => "fifo",
                    crate::engine::SglangSchedulePolicy::Lpm => "lpm",
                }
                .to_string(),
            ),
            page_size: Some(config.block_size),
            max_prefill_tokens: Some(config.sglang.max_prefill_tokens),
            chunked_prefill_size: Some(config.sglang.chunked_prefill_size),
            clip_max_new_tokens: Some(config.sglang.clip_max_new_tokens),
            schedule_conservativeness: Some(config.sglang.schedule_conservativeness),
            vlm_cache_bytes: Some(config.sglang.vlm_cache_bytes),
            host_loop: config.sglang.host_loop,
            frontend: config.frontend.clone(),
        }),
        emit_kv_events: config.emit_kv_events,
        emit_kv_token_ids: config.emit_kv_token_ids,
    }
}

/// Mirror SGLang's `handle_data_parallelism` launch-time DP-attention normalization before
/// cloning the rank-local scheduler configuration. Replay's `dp_size` is specifically attention
/// DP, so every SGLang rank receives the same normalized per-rank controls exactly once during
/// construction.
fn normalize_sglang_attention_dp(args: &mut MockEngineArgs, dp_size: u32) -> Result<()> {
    if dp_size == 1 {
        return Ok(());
    }
    let sglang = args
        .sglang
        .as_mut()
        .expect("materialized SGLang engine arguments must include scheduler controls");
    let dp_size = usize::try_from(dp_size).expect("u32 always fits usize on supported platforms");
    let chunked_prefill_size = sglang
        .chunked_prefill_size
        .expect("materialized SGLang chunked-prefill size");
    let per_rank_chunked_prefill_size = chunked_prefill_size / dp_size;
    ensure!(
        per_rank_chunked_prefill_size > 0,
        "SGLang attention DP size {dp_size} reduces chunked_prefill_size={chunked_prefill_size} to zero"
    );
    sglang.chunked_prefill_size = Some(per_rank_chunked_prefill_size);
    sglang.schedule_conservativeness = Some(
        sglang
            .schedule_conservativeness
            .expect("materialized SGLang schedule conservativeness")
            * 0.3,
    );
    Ok(())
}

fn core_request(request: Request) -> DirectRequest {
    DirectRequest {
        tokens: request.tokens,
        max_output_tokens: request.max_output_tokens,
        output_token_ids: request.output_token_ids,
        uuid: Some(request.request_id),
        arrival_timestamp_ms: None,
        session_id: request.session_id,
        images: request.images,
        kv_hint: request.kv_hint,
    }
}

fn pending_output_suppression(
    command: &Command,
    handoffs: &HashMap<HandoffId, Uuid>,
) -> Option<(Uuid, bool)> {
    match *command {
        Command::CancelRequest {
            request_id,
            discard_pending_output,
        } => Some((request_id, discard_pending_output)),
        Command::CancelSource { handoff_id } | Command::CancelDestination { handoff_id } => {
            handoffs
                .get(&handoff_id)
                .copied()
                .map(|request_id| (request_id, false))
        }
        _ => None,
    }
}

fn map_command_effects(
    effects: CoreCommandEffects,
    metrics: Metrics,
    suppressed_pending_output: bool,
) -> Result<CommandEffects> {
    let result = match effects.result {
        CoreCommandResult::Submitted(id) => CommandResult::Submitted(id),
        CoreCommandResult::DestinationAccepted { request_id } => {
            CommandResult::DestinationAccepted { request_id }
        }
        CoreCommandResult::Applied => CommandResult::Applied,
        CoreCommandResult::Noop => CommandResult::Noop,
    };
    Ok(CommandEffects {
        result,
        lifecycle_events: effects
            .lifecycle_events
            .into_iter()
            .map(map_lifecycle)
            .collect(),
        kv_events: effects.kv_events,
        retired_requests: effects.retired_requests,
        metrics,
        suppressed_pending_output,
    })
}

fn map_lifecycle(event: CoreLifecycle) -> LifecycleEvent {
    match event {
        CoreLifecycle::SourceHeld {
            handoff_id,
            request_id,
            transfer_timing,
        } => LifecycleEvent::SourceHeld {
            handoff_id,
            request_id,
            transfer_timing,
        },
        CoreLifecycle::DestinationReserved {
            handoff_id,
            request_id,
            transferable_prompt_tokens,
        } => LifecycleEvent::DestinationReserved {
            handoff_id,
            request_id,
            transferable_prompt_tokens,
        },
        CoreLifecycle::TtftMilestone {
            request_id,
            stage,
            at_ms,
        } => LifecycleEvent::TtftMilestone {
            request_id,
            stage,
            at_ms,
        },
    }
}

fn map_metrics(metrics: MockerMetrics) -> Metrics {
    Metrics {
        dp_rank: metrics.dp_rank,
        active_blocks: metrics.active_decode_blocks,
        inactive_blocks: metrics.inactive_decode_blocks,
        total_blocks: metrics.total_blocks,
        kv_cache_used_bytes: metrics.kv_cache_used_bytes,
        kv_cache_capacity_bytes: metrics.kv_cache_capacity_bytes,
        cache_usage: metrics.gpu_cache_usage_perc,
        physical_cache_usage: metrics.physical_gpu_cache_usage_perc,
        running_requests: metrics.running_requests,
        waiting_requests: metrics.waiting_requests,
        preemptions_total: metrics.vllm_preemptions_total,
        sglang_cache_hit_tokens: metrics.sglang_cache_hit_tokens,
        sglang_cache_total_tokens: metrics.sglang_cache_total_tokens,
    }
}

fn split_pass(
    pass: EnginePassResult,
) -> Result<(
    crate::engine::generalized::SameTimestampRetry,
    PassStartEffects,
    PassCompletionEffects,
)> {
    let EnginePassResult {
        committed_requests,
        same_timestamp_retry,
        output_signals,
        admissions,
        g1_prefix_cache_lookup,
        pressure_events,
        lifecycle_events,
        mocker_metrics,
        kv_event_visibility,
        kv_events,
        fpm,
        decode_acceptance,
        ..
    } = pass;
    let (start_kv, completion_kv) = match kv_event_visibility {
        KvEventVisibility::PassEnd => (Vec::new(), kv_events),
    };
    let start = PassStartEffects {
        committed_requests,
        admissions: admissions
            .into_iter()
            .map(|admission| Admission {
                request_id: admission.uuid,
                reused_input_tokens: admission.reused_input_tokens,
                cache_tier_attribution: admission.cache_tier_attribution,
            })
            .collect(),
        pressure_events,
        kv_events: start_kv,
        g1_prefix_cache_lookup,
    };
    let completion = PassCompletionEffects {
        outputs: output_signals
            .into_iter()
            .map(|output| Output {
                request_id: output.uuid,
                token_id: output.token_id,
                completed: output.completed,
                rejected: output.rejected,
                cached_tokens: output.cached_tokens,
            })
            .collect(),
        lifecycle_events: lifecycle_events.into_iter().map(map_lifecycle).collect(),
        kv_events: completion_kv,
        metrics: map_metrics(mocker_metrics),
        forward_pass_metrics: fpm.map(map_fpm).unwrap_or_default(),
        decode_acceptance,
    };
    Ok((same_timestamp_retry, start, completion))
}

fn map_fpm(fpm: crate::engine::common::protocols::ForwardPassSnapshot) -> ForwardPassMetrics {
    ForwardPassMetrics {
        num_prefill_requests: fpm.num_prefill_requests,
        sum_prefill_tokens: fpm.sum_prefill_tokens,
        var_prefill_length: fpm.var_prefill_length,
        sum_prefill_kv_tokens: fpm.sum_prefill_kv_tokens,
        num_decode_requests: fpm.num_decode_requests,
        sum_decode_kv_tokens: fpm.sum_decode_kv_tokens,
        var_decode_kv_tokens: fpm.var_decode_kv_tokens,
        num_queued_prefill: fpm.num_queued_prefill,
        sum_queued_prefill_tokens: fpm.sum_queued_prefill_tokens,
        var_queued_prefill_length: fpm.var_queued_prefill_length,
        num_queued_decode: fpm.num_queued_decode,
        sum_queued_decode_kv_tokens: fpm.sum_queued_decode_kv_tokens,
        var_queued_decode_kv_tokens: fpm.var_queued_decode_kv_tokens,
        duration_ms: fpm.wall_time_secs * 1_000.0,
    }
}

#[cfg(test)]
mod tests {
    use std::num::NonZeroU32;
    use std::sync::Mutex;

    use super::*;
    use crate::engine::{
        EncoderShape, HostOffloadObservation, HostOffloadObservationData, ImageSpec,
        NativeHostOffloadConfig, PressureKind, TimingModelConfig,
    };

    #[rstest::rstest]
    fn unsupported_native_images_reject_without_consuming_request_or_handoff_ids(
        #[values(Backend::Vllm, Backend::Trtllm)] backend: Backend,
        #[values(WorkerType::Aggregated, WorkerType::Prefill, WorkerType::Decode)]
        worker_type: WorkerType,
    ) {
        let config = EngineConfig {
            backend,
            worker_type,
            num_gpu_blocks: 8,
            block_size: 4,
            timing_model: TimingModelConfig::Fixed {
                prefill_ms: 1.0,
                decode_ms: 1.0,
            },
            ..EngineConfig::default()
        };
        let mut rank = SchedulerRank::new(
            RankIdentity {
                worker_id: 0,
                dp_rank: 0,
                dp_size: NonZeroU32::MIN,
            },
            &config,
        )
        .unwrap();
        let request_id = Uuid::from_u128(1);
        let handoff_id = HandoffId::from(Uuid::from_u128(2));
        let mut request = Request {
            request_id,
            tokens: (0..8).collect(),
            max_output_tokens: 1,
            output_token_ids: None,
            session_id: None,
            images: vec![ImageSpec {
                identity: 1,
                token_start: 4,
                token_end: 8,
                encoder: EncoderShape {
                    sequences: 1,
                    patch_tokens: 16,
                    transformer_tokens: 16,
                    output_tokens: 4,
                },
                embedding_bytes: 128,
            }],
            kv_hint: None,
        };
        let command = |request| match worker_type {
            WorkerType::Aggregated => Command::Submit(request),
            WorkerType::Prefill => Command::SubmitHandoffPrefill {
                handoff_id,
                request,
            },
            WorkerType::Decode => Command::ReserveDestination {
                handoff_id,
                request,
            },
        };
        let context = CommandContext {
            now_ms: 0.0,
            pass_in_flight: false,
        };
        let before = rank.metrics();
        let error = rank
            .apply_command_effects(command(request.clone()), context, None)
            .unwrap_err();
        assert!(
            error
                .to_string()
                .contains("images require the SGLang backend"),
            "{error:#}"
        );
        assert_eq!(rank.metrics(), before);
        assert!(rank.is_drained());
        assert!(rank.handoff_requests.is_empty());

        // The same request and handoff IDs remain usable for a text request.
        request.images.clear();
        let effects = rank
            .apply_command_effects(command(request), context, None)
            .unwrap();
        let expected = if worker_type == WorkerType::Decode {
            CommandResult::DestinationAccepted { request_id }
        } else {
            CommandResult::Submitted(request_id)
        };
        assert_eq!(effects.result, expected);
        if worker_type == WorkerType::Decode {
            rank.apply_command_effects(Command::ActivateDestination { handoff_id }, context, None)
                .unwrap();
        }
        let pass = rank.execute_pass(0.0).unwrap();
        let completion = rank.complete_pass(pass.pending, pass.end_ms).unwrap();
        if worker_type == WorkerType::Prefill {
            assert!(completion.lifecycle_events.iter().any(|event| matches!(
                event,
                LifecycleEvent::SourceHeld { handoff_id: id, .. } if *id == handoff_id
            )));
            rank.apply_command_effects(
                Command::ReleaseSource { handoff_id },
                CommandContext {
                    now_ms: pass.end_ms,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
        } else {
            assert!(
                completion
                    .outputs
                    .iter()
                    .any(|output| output.request_id == request_id && output.completed)
            );
        }
        assert!(rank.is_drained());
    }

    #[test]
    fn state_cache_parameters_reach_rank_local_args() {
        let config: EngineConfig = serde_json::from_value(serde_json::json!({
            "num_gpu_blocks":8,"block_size":64,"kv_cache_bytes_per_token":16,
            "state_cache": {"bytes_per_request":1500}
        }))
        .unwrap();
        let args = core_args(&config, config.built_in_timing_model().unwrap());
        assert_eq!(args.state_cache, config.state_cache);
        assert_eq!(args.num_gpu_blocks, 8);
        assert_eq!(args.block_size, 64);

        let legacy = EngineConfig::default();
        let args = core_args(&legacy, legacy.built_in_timing_model().unwrap());
        assert!(args.state_cache.is_none());
        assert_eq!(args.num_gpu_blocks, legacy.num_gpu_blocks);
        assert_eq!(args.block_size, legacy.block_size);
    }

    #[test]
    fn kvcr_is_replay_owned() {
        let config = EngineConfig {
            kv_cache_bytes_per_token: Some(250_000),
            native_host_offload: Some(NativeHostOffloadConfig::new(4)),
            kvcr: Some(crate::engine::KvcrConfig::default()),
            ..EngineConfig::default()
        };
        let identity = RankIdentity {
            worker_id: 1,
            dp_rank: 0,
            dp_size: NonZeroU32::MIN,
        };
        let timing = config.built_in_timing_model().unwrap();
        let error = SchedulerRank::new_with_g2(identity, &config, timing, 0, None)
            .err()
            .expect("a bare rank must reject kvcr");
        assert_eq!(
            error.to_string(),
            "kvcr is Replay-owned; construct it through ReplaySpec"
        );
    }

    #[test]
    fn sglang_attention_dp_normalizes_per_rank_scheduler_controls_once() {
        let args = MockEngineArgs::builder()
            .engine_type(EngineType::Sglang)
            .sglang(Some(SglangArgs {
                chunked_prefill_size: Some(16),
                schedule_conservativeness: Some(2.0),
                ..Default::default()
            }))
            .build()
            .unwrap();

        let mut dp1 = args.clone();
        normalize_sglang_attention_dp(&mut dp1, 1).unwrap();
        let dp1 = dp1.sglang.unwrap();
        assert_eq!(dp1.chunked_prefill_size, Some(16));
        assert_eq!(dp1.schedule_conservativeness, Some(2.0));

        let mut dp4 = args;
        normalize_sglang_attention_dp(&mut dp4, 4).unwrap();
        let dp4 = dp4.sglang.unwrap();
        assert_eq!(dp4.chunked_prefill_size, Some(4));
        assert_eq!(dp4.schedule_conservativeness, Some(0.6));
    }

    #[test]
    fn sglang_attention_dp_normalizes_memory_pressure_admission() {
        let second_request_admitted = |dp_size| {
            let config = EngineConfig {
                backend: Backend::Sglang,
                num_gpu_blocks: 8,
                block_size: 4,
                max_num_seqs: 4,
                max_num_batched_tokens: 64,
                sglang: crate::engine::SglangConfig {
                    chunked_prefill_size: 64,
                    schedule_conservativeness: 1.0,
                    ..Default::default()
                },
                timing_model: TimingModelConfig::Fixed {
                    prefill_ms: 1.0,
                    decode_ms: 1.0,
                },
                ..EngineConfig::default()
            };
            let mut rank = SchedulerRank::new(
                RankIdentity {
                    worker_id: 1,
                    dp_rank: 0,
                    dp_size: NonZeroU32::new(dp_size).unwrap(),
                },
                &config,
            )
            .unwrap();
            let first = Uuid::from_u128(94_001);
            rank.apply_command_effects(
                Command::Submit(Request {
                    request_id: first,
                    tokens: vec![1, 2, 3, 4],
                    max_output_tokens: 20,
                    output_token_ids: None,
                    session_id: None,
                    images: Vec::new(),
                    kv_hint: None,
                }),
                CommandContext {
                    now_ms: 0.0,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
            let first_pass = rank.execute_pass(0.0).unwrap();
            let first_end_ms = first_pass.end_ms;
            rank.complete_pass(first_pass.pending, first_end_ms)
                .unwrap();

            let second = Uuid::from_u128(94_002);
            rank.apply_command_effects(
                Command::Submit(Request {
                    request_id: second,
                    tokens: (10..22).collect(),
                    max_output_tokens: 1,
                    output_token_ids: None,
                    session_id: None,
                    images: Vec::new(),
                    kv_hint: None,
                }),
                CommandContext {
                    now_ms: first_end_ms,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
            rank.execute_pass(first_end_ms)
                .unwrap()
                .start_effects
                .admissions
                .iter()
                .any(|admission| admission.request_id == second)
        };

        assert!(!second_request_admitted(1));
        assert!(second_request_admitted(4));
    }

    #[derive(Debug, Clone, Copy, PartialEq)]
    struct CapturedHostEvent {
        request_id: Uuid,
        kind: &'static str,
        at_ms: f64,
    }

    #[derive(Default)]
    struct HostEventCapture(Mutex<Vec<CapturedHostEvent>>);

    impl HostEventCapture {
        fn snapshot(&self) -> Vec<CapturedHostEvent> {
            self.0.lock().expect("host event capture poisoned").clone()
        }
    }

    impl HostOffloadObserver for HostEventCapture {
        fn record(&self, observation: HostOffloadObservation<'_>) {
            let (kind, at_ms) = match observation.event {
                HostOffloadObservationData::LoadQueued { at_ms, .. } => ("load_queued", at_ms),
                HostOffloadObservationData::LoadCompleted { at_ms, .. } => {
                    ("load_completed", at_ms)
                }
                HostOffloadObservationData::LoadCancelled { at_ms, .. } => {
                    ("load_cancelled", at_ms)
                }
                _ => return,
            };
            self.0
                .lock()
                .expect("host event capture poisoned")
                .push(CapturedHostEvent {
                    request_id: observation.request_id,
                    kind,
                    at_ms,
                });
        }
    }

    fn rank() -> SchedulerRank {
        rank_for_worker(WorkerType::Aggregated)
    }

    fn rank_config(worker_type: WorkerType) -> EngineConfig {
        EngineConfig {
            worker_type,
            num_gpu_blocks: 8,
            block_size: 4,
            max_num_seqs: 2,
            max_num_batched_tokens: 16,
            speedup_ratio: 0.0,
            timing_model: TimingModelConfig::Fixed {
                prefill_ms: 10.0,
                decode_ms: 10.0,
            },
            ..EngineConfig::default()
        }
    }

    fn rank_for_worker(worker_type: WorkerType) -> SchedulerRank {
        let config = rank_config(worker_type);
        SchedulerRank::new(
            RankIdentity {
                worker_id: 1,
                dp_rank: 0,
                dp_size: NonZeroU32::MIN,
            },
            &config,
        )
        .unwrap()
    }

    fn host_rank(observer: Arc<HostEventCapture>) -> SchedulerRank {
        host_rank_with_capacity(observer, 2, 4)
    }

    fn host_rank_with_capacity(
        observer: Arc<HostEventCapture>,
        g1_blocks: usize,
        host_blocks: usize,
    ) -> SchedulerRank {
        let config = EngineConfig {
            num_gpu_blocks: g1_blocks,
            block_size: 4,
            max_num_seqs: 4,
            max_num_batched_tokens: 16,
            kv_cache_bytes_per_token: Some(250_000),
            native_host_offload: Some(
                NativeHostOffloadConfig::new(host_blocks).with_bandwidths(1.0, 1.0),
            ),
            speedup_ratio: 0.0,
            timing_model: TimingModelConfig::Fixed {
                prefill_ms: 10.0,
                decode_ms: 10.0,
            },
            ..EngineConfig::default()
        };
        let mut rank = SchedulerRank::new(
            RankIdentity {
                worker_id: 2,
                dp_rank: 0,
                dp_size: NonZeroU32::MIN,
            },
            &config,
        )
        .unwrap();
        rank.set_host_offload_observer(observer);
        rank
    }

    fn submit_completed_prompt(
        rank: &mut SchedulerRank,
        request_id: Uuid,
        tokens: Vec<u32>,
        now_ms: f64,
    ) -> f64 {
        let effects = rank
            .apply_command_effects(
                Command::Submit(Request {
                    request_id,
                    tokens,
                    max_output_tokens: 0,
                    output_token_ids: None,
                    session_id: None,
                    images: Vec::new(),
                    kv_hint: None,
                }),
                CommandContext {
                    now_ms,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
        assert_eq!(effects.result, CommandResult::Submitted(request_id));
        let pass = rank.execute_pass(now_ms).unwrap();
        let end_ms = pass.end_ms;
        rank.complete_pass(pass.pending, end_ms).unwrap();
        end_ms
    }

    /// Seed G2 with one prompt, evict it from G1, then queue an H2D owned by a
    /// still-pending source handoff. Returns `(handoff_id, request_id, due_ms)`.
    fn queue_source_h2d(rank: &mut SchedulerRank) -> (HandoffId, Uuid, f64) {
        let seed_id = Uuid::from_u128(93_001);
        let seed_end = submit_completed_prompt(rank, seed_id, vec![1, 2, 3, 4], 0.0);
        let store_due = rank.next_internal_deadline_ms().unwrap();
        assert!(store_due >= seed_end);
        rank.process_internal_work(store_due, false).unwrap();

        let evict_id = Uuid::from_u128(93_002);
        let evict_end =
            submit_completed_prompt(rank, evict_id, vec![5, 6, 7, 8, 9, 10, 11, 12], store_due);
        let mut restore_at_ms = evict_end;
        while let Some(deadline) = rank.next_internal_deadline_ms() {
            restore_at_ms = restore_at_ms.max(deadline);
            rank.process_internal_work(restore_at_ms, false).unwrap();
        }

        let handoff_id = HandoffId::from(Uuid::from_u128(93_003));
        let restore_id = Uuid::from_u128(93_004);
        let effects = rank
            .apply_command_effects(
                Command::SubmitHandoffPrefill {
                    handoff_id,
                    request: Request {
                        request_id: restore_id,
                        tokens: vec![1, 2, 3, 4],
                        max_output_tokens: 0,
                        output_token_ids: None,
                        session_id: None,
                        images: Vec::new(),
                        kv_hint: None,
                    },
                },
                CommandContext {
                    now_ms: restore_at_ms,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
        assert_eq!(effects.result, CommandResult::Submitted(restore_id));
        let pass = rank.execute_pass(restore_at_ms).unwrap();
        rank.complete_pass(pass.pending, pass.end_ms).unwrap();
        let due_ms = rank.next_internal_deadline_ms().unwrap();
        (handoff_id, restore_id, due_ms)
    }

    fn start_request_pass(
        rank: &mut SchedulerRank,
        request_id: Uuid,
        output_token_ids: Vec<u32>,
    ) -> PendingPass {
        let effects = rank
            .apply_command_effects(
                Command::Submit(Request {
                    request_id,
                    tokens: vec![1, 2, 3, 4],
                    max_output_tokens: output_token_ids.len(),
                    output_token_ids: Some(output_token_ids),
                    session_id: None,
                    images: Vec::new(),
                    kv_hint: None,
                }),
                CommandContext {
                    now_ms: 0.0,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
        assert_eq!(effects.result, CommandResult::Submitted(request_id));
        let pass = rank.execute_pass(0.0).unwrap();
        assert!(
            pass.pending
                .effects
                .outputs
                .iter()
                .any(|output| output.request_id == request_id)
        );
        pass.pending
    }

    #[test]
    fn ordinary_cancel_suppresses_pending_output_when_scheduler_state_is_removed() {
        let request_id = Uuid::from_u128(90_001);
        let mut rank = rank();
        let mut pending = start_request_pass(&mut rank, request_id, vec![5, 6]);

        let effects = rank
            .apply_command_effects(
                Command::CancelRequest {
                    request_id,
                    discard_pending_output: false,
                },
                CommandContext {
                    now_ms: 1.0,
                    pass_in_flight: true,
                },
                Some(&mut pending),
            )
            .unwrap();

        assert_eq!(effects.result, CommandResult::Applied);
        assert!(effects.suppressed_pending_output);
        assert_eq!(effects.retired_requests, vec![request_id]);
        assert!(pending.effects.outputs.is_empty());
    }

    #[test]
    fn a_cancel_drops_outputs_vllm_has_not_returned() {
        // With `launch`, the first pass's output waits in the rank for the
        // next pass's launch; cancelling the request drops it.
        let request_id = Uuid::from_u128(90_002);
        let config = EngineConfig {
            launch: Some(crate::engine::PassLaunchConfig {
                decode_ms: 1.0,
                piecewise_ms: 1.0,
                eager_ms: 1.0,
                max_cudagraph_capture_size: 16,
                eager_queue_fraction: 0.0,
            }),
            ..rank_config(WorkerType::Aggregated)
        };
        let mut rank = SchedulerRank::new(
            RankIdentity {
                worker_id: 1,
                dp_rank: 0,
                dp_size: NonZeroU32::MIN,
            },
            &config,
        )
        .unwrap();
        let pending = start_request_pass(&mut rank, request_id, vec![5, 6]);
        assert!(rank.complete_pass(pending, 0.0).unwrap().outputs.is_empty());
        let effects = rank
            .apply_command_effects(
                Command::CancelRequest {
                    request_id,
                    discard_pending_output: false,
                },
                CommandContext {
                    now_ms: 0.0,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
        assert!(effects.suppressed_pending_output);
        assert!(rank.is_drained());
    }

    #[test]
    fn ordinary_noop_cancel_preserves_pending_output() {
        let request_id = Uuid::from_u128(90_002);
        let mut rank = rank();
        let mut pending = start_request_pass(&mut rank, request_id, vec![5]);

        let effects = rank
            .apply_command_effects(
                Command::CancelRequest {
                    request_id,
                    discard_pending_output: false,
                },
                CommandContext {
                    now_ms: 1.0,
                    pass_in_flight: true,
                },
                Some(&mut pending),
            )
            .unwrap();

        assert_eq!(effects.result, CommandResult::Noop);
        assert!(!effects.suppressed_pending_output);
        assert_eq!(pending.effects.outputs.len(), 1);
    }

    #[test]
    fn explicit_discard_suppresses_pending_output_after_noop_cancellation() {
        let request_id = Uuid::from_u128(90_003);
        let mut rank = rank();
        let mut pending = start_request_pass(&mut rank, request_id, vec![5]);

        let effects = rank
            .apply_command_effects(
                Command::CancelRequest {
                    request_id,
                    discard_pending_output: true,
                },
                CommandContext {
                    now_ms: 1.0,
                    pass_in_flight: true,
                },
                Some(&mut pending),
            )
            .unwrap();

        assert_eq!(effects.result, CommandResult::Noop);
        assert!(effects.suppressed_pending_output);
        assert_eq!(effects.retired_requests, vec![request_id]);
        assert!(pending.effects.outputs.is_empty());
    }

    #[test]
    fn mid_pass_command_and_internal_call_keep_due_h2d_hidden() {
        let observer = Arc::new(HostEventCapture::default());
        let mut rank = host_rank(Arc::clone(&observer));
        let (_handoff_id, restore_id, h2d_due_ms) = queue_source_h2d(&mut rank);
        assert!(
            observer
                .snapshot()
                .iter()
                .any(|event| { event.request_id == restore_id && event.kind == "load_queued" })
        );

        let busy_id = Uuid::from_u128(93_005);
        let submission = rank
            .apply_command_effects(
                Command::Submit(Request {
                    request_id: busy_id,
                    tokens: vec![21, 22, 23, 24],
                    max_output_tokens: 0,
                    output_token_ids: None,
                    session_id: None,
                    images: Vec::new(),
                    kv_hint: None,
                }),
                CommandContext {
                    now_ms: h2d_due_ms - 0.5,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
        assert_eq!(submission.result, CommandResult::Submitted(busy_id));
        let pass = rank.execute_pass(h2d_due_ms - 0.5).unwrap();
        assert!(pass.end_ms > h2d_due_ms);
        let mut pending = pass.pending;

        let arrival_id = Uuid::from_u128(93_006);
        let arrival = rank
            .apply_command_effects(
                Command::Submit(Request {
                    request_id: arrival_id,
                    tokens: vec![31, 32, 33, 34],
                    max_output_tokens: 0,
                    output_token_ids: None,
                    session_id: None,
                    images: Vec::new(),
                    kv_hint: None,
                }),
                CommandContext {
                    now_ms: h2d_due_ms + 0.5,
                    pass_in_flight: true,
                },
                Some(&mut pending),
            )
            .unwrap();
        assert_eq!(arrival.result, CommandResult::Submitted(arrival_id));
        assert!(arrival.kv_events.is_empty());
        assert!(
            !observer
                .snapshot()
                .iter()
                .any(|event| { event.request_id == restore_id && event.kind == "load_completed" })
        );

        let internal = rank.process_internal_work(h2d_due_ms + 0.5, true).unwrap();
        assert_eq!(internal, PassStartEffects::default());
        assert!(
            !observer
                .snapshot()
                .iter()
                .any(|event| { event.request_id == restore_id && event.kind == "load_completed" })
        );

        rank.complete_pass(pending, pass.end_ms).unwrap();
        let events = observer.snapshot();
        assert!(
            events.iter().any(|event| {
                event.request_id == restore_id
                    && event.kind == "load_completed"
                    // The observation retains the physical completion timestamp,
                    // even though it is emitted only at the later model boundary.
                    && event.at_ms == h2d_due_ms
            }),
            "events: {events:?}, pass end: {}",
            pass.end_ms
        );
    }

    #[test]
    fn rejected_submit_at_due_deadline_does_not_settle_host_work() {
        let observer = Arc::new(HostEventCapture::default());
        let mut rank = host_rank(Arc::clone(&observer));
        let (_handoff_id, restore_id, h2d_due_ms) = queue_source_h2d(&mut rank);
        let events_before = observer.snapshot();
        let metrics_before = rank.metrics();
        let deadline_before = rank.next_internal_deadline_ms();

        let error = rank
            .apply_command_effects(
                Command::Submit(Request {
                    request_id: restore_id,
                    tokens: vec![41, 42, 43, 44],
                    max_output_tokens: 0,
                    output_token_ids: None,
                    session_id: None,
                    images: Vec::new(),
                    kv_hint: None,
                }),
                CommandContext {
                    now_ms: h2d_due_ms + 1.0,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap_err();
        assert!(format!("{error:#}").contains("already active"));
        assert_eq!(rank.metrics(), metrics_before);
        assert_eq!(rank.next_internal_deadline_ms(), deadline_before);
        assert_eq!(observer.snapshot(), events_before);
        assert!(rank.core.drain_kv_events().is_empty());
    }

    #[test]
    fn same_pass_queues_multiple_host_loads_with_checked_headroom_updates() {
        let observer = Arc::new(HostEventCapture::default());
        let mut rank = host_rank_with_capacity(Arc::clone(&observer), 4, 8);
        let mut now_ms = 0.0;
        for (request_id, tokens) in [
            (Uuid::from_u128(95_001), vec![1, 2, 3, 4]),
            (Uuid::from_u128(95_002), vec![11, 12, 13, 14]),
        ] {
            now_ms = submit_completed_prompt(&mut rank, request_id, tokens, now_ms);
            while let Some(deadline) = rank.next_internal_deadline_ms() {
                now_ms = now_ms.max(deadline);
                rank.process_internal_work(now_ms, false).unwrap();
            }
        }
        now_ms = submit_completed_prompt(
            &mut rank,
            Uuid::from_u128(95_003),
            (100..116).collect(),
            now_ms,
        );
        while let Some(deadline) = rank.next_internal_deadline_ms() {
            now_ms = now_ms.max(deadline);
            rank.process_internal_work(now_ms, false).unwrap();
        }

        let loads = [
            (Uuid::from_u128(95_004), vec![1, 2, 3, 4, 21, 22, 23, 24]),
            (
                Uuid::from_u128(95_005),
                vec![11, 12, 13, 14, 31, 32, 33, 34],
            ),
        ];
        for (request_id, tokens) in &loads {
            let effects = rank
                .apply_command_effects(
                    Command::Submit(Request {
                        request_id: *request_id,
                        tokens: tokens.clone(),
                        max_output_tokens: 0,
                        output_token_ids: None,
                        session_id: None,
                        images: Vec::new(),
                        kv_hint: None,
                    }),
                    CommandContext {
                        now_ms,
                        pass_in_flight: false,
                    },
                    None,
                )
                .unwrap();
            assert_eq!(effects.result, CommandResult::Submitted(*request_id));
        }

        let pass = rank.execute_pass(now_ms).unwrap();
        let events = observer.snapshot();
        for (request_id, _) in loads {
            assert!(events.iter().any(|event| {
                event.request_id == request_id
                    && event.kind == "load_queued"
                    && event.at_ms == now_ms
            }));
        }
        rank.complete_pass(pass.pending, pass.end_ms).unwrap();
    }

    #[test]
    fn mid_pass_request_and_source_cancel_observe_command_time_without_completing_h2d() {
        for cancel_source in [false, true] {
            let observer = Arc::new(HostEventCapture::default());
            let mut rank = host_rank(Arc::clone(&observer));
            let (handoff_id, restore_id, h2d_due_ms) = queue_source_h2d(&mut rank);

            let busy_id = Uuid::from_u128(94_000 + u128::from(cancel_source));
            rank.apply_command_effects(
                Command::Submit(Request {
                    request_id: busy_id,
                    tokens: vec![21, 22, 23, 24],
                    max_output_tokens: 0,
                    output_token_ids: None,
                    session_id: None,
                    images: Vec::new(),
                    kv_hint: None,
                }),
                CommandContext {
                    now_ms: h2d_due_ms - 0.5,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
            let pass = rank.execute_pass(h2d_due_ms - 0.5).unwrap();
            assert!(pass.end_ms > h2d_due_ms);
            let mut pending = pass.pending;
            let cancel_at_ms = h2d_due_ms + 0.5;
            let command = if cancel_source {
                Command::CancelSource { handoff_id }
            } else {
                Command::CancelRequest {
                    request_id: restore_id,
                    discard_pending_output: false,
                }
            };
            let effects = rank
                .apply_command_effects(
                    command,
                    CommandContext {
                        now_ms: cancel_at_ms,
                        pass_in_flight: true,
                    },
                    Some(&mut pending),
                )
                .unwrap();
            assert_eq!(effects.result, CommandResult::Applied);
            let events = observer.snapshot();
            assert!(events.iter().any(|event| {
                event.request_id == restore_id
                    && event.kind == "load_cancelled"
                    && event.at_ms == cancel_at_ms
            }));
            assert!(
                !events.iter().any(|event| {
                    event.request_id == restore_id && event.kind == "load_completed"
                })
            );

            rank.complete_pass(pending, pass.end_ms).unwrap();
            assert!(
                !observer.snapshot().iter().any(|event| {
                    event.request_id == restore_id && event.kind == "load_completed"
                })
            );
        }
    }

    #[test]
    fn cancel_source_uses_command_time_outside_a_pass() {
        let observer = Arc::new(HostEventCapture::default());
        let mut rank = host_rank(Arc::clone(&observer));
        let (handoff_id, restore_id, h2d_due_ms) = queue_source_h2d(&mut rank);
        let cancel_at_ms = h2d_due_ms - 0.25;

        let effects = rank
            .apply_command_effects(
                Command::CancelSource { handoff_id },
                CommandContext {
                    now_ms: cancel_at_ms,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
        assert_eq!(effects.result, CommandResult::Applied);
        assert!(observer.snapshot().iter().any(|event| {
            event.request_id == restore_id
                && event.kind == "load_cancelled"
                && event.at_ms == cancel_at_ms
        }));
    }

    #[test]
    fn handoff_tracking_is_inserted_on_success_and_cleared_by_cancel() {
        let mut rank = rank_for_worker(WorkerType::Decode);
        let handoff_id = HandoffId::from(Uuid::from_u128(91_001));
        let request_id = Uuid::from_u128(91_002);
        let reservation = rank
            .apply_command_effects(
                Command::ReserveDestination {
                    handoff_id,
                    request: Request {
                        request_id,
                        tokens: vec![1, 2, 3, 4],
                        max_output_tokens: 1,
                        output_token_ids: Some(vec![5]),
                        session_id: None,
                        images: Vec::new(),
                        kv_hint: None,
                    },
                },
                CommandContext {
                    now_ms: 0.0,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
        assert!(matches!(
            reservation.result,
            CommandResult::DestinationAccepted { .. }
        ));
        assert_eq!(rank.handoff_requests.get(&handoff_id), Some(&request_id));

        let cancellation = rank
            .apply_command_effects(
                Command::CancelDestination { handoff_id },
                CommandContext {
                    now_ms: 0.0,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
        assert_eq!(cancellation.result, CommandResult::Applied);
        assert!(!rank.handoff_requests.contains_key(&handoff_id));
    }

    #[test]
    fn pass_start_exposes_vllm_preemption_pressure_event() {
        let config = EngineConfig {
            num_gpu_blocks: 6,
            block_size: 4,
            max_num_seqs: 2,
            max_num_batched_tokens: 16,
            enable_prefix_caching: false,
            enable_chunked_prefill: true,
            speedup_ratio: 0.0,
            preemption_mode: PreemptionMode::Lifo,
            timing_model: TimingModelConfig::Fixed {
                prefill_ms: 10.0,
                decode_ms: 10.0,
            },
            ..EngineConfig::default()
        };
        let mut rank = SchedulerRank::new(
            RankIdentity {
                worker_id: 7,
                dp_rank: 0,
                dp_size: NonZeroU32::MIN,
            },
            &config,
        )
        .unwrap();
        let first = Uuid::from_u128(92_001);
        let second = Uuid::from_u128(92_002);
        for (request_id, tokens) in [
            (first, (0..8).collect::<Vec<_>>()),
            (second, (100..108).collect::<Vec<_>>()),
        ] {
            let effects = rank
                .apply_command_effects(
                    Command::Submit(Request {
                        request_id,
                        tokens,
                        max_output_tokens: 8,
                        output_token_ids: None,
                        session_id: None,
                        images: Vec::new(),
                        kv_hint: None,
                    }),
                    CommandContext {
                        now_ms: 0.0,
                        pass_in_flight: false,
                    },
                    None,
                )
                .unwrap();
            assert_eq!(effects.result, CommandResult::Submitted(request_id));
        }

        let mut now_ms = 0.0;
        let mut observed = None;
        for _ in 0..16 {
            let pass = rank.execute_pass(now_ms).unwrap();
            if let Some(event) = pass.start_effects.pressure_events.first() {
                assert_eq!(pass.start_effects.pressure_events.len(), 1);
                observed = Some(event.clone());
            }
            let end_ms = pass.end_ms;
            rank.complete_pass(pass.pending, end_ms).unwrap();
            if observed.is_some() {
                break;
            }
            now_ms = end_ms.max(now_ms + 1.0);
        }

        let event = observed.expect("tight native G1 capacity should preempt one vLLM request");
        assert_eq!(event.at_ms, now_ms);
        assert_eq!(event.kind, PressureKind::VllmPreemption);
        assert_eq!(event.request_id, second);
        assert_eq!(event.state_before.running_requests, 2);
        assert_eq!(event.state_before.waiting_requests, Some(0));
        assert_eq!(event.state_after.running_requests, 1);
        assert_eq!(event.state_after.waiting_requests, Some(1));
        assert!(event.request_active_blocks_before > 0);
        assert!(event.state_after.active_blocks < event.state_before.active_blocks);
        assert_eq!(event.logical_available_blocks_before, None);
        assert_eq!(event.required_blocks_before, None);
    }

    #[test]
    fn terminal_handoff_output_clears_tracking() {
        let mut rank = rank_for_worker(WorkerType::Prefill);
        let handoff_id = HandoffId::from(Uuid::from_u128(92_001));
        let request_id = Uuid::from_u128(92_002);
        let submission = rank
            .apply_command_effects(
                Command::SubmitHandoffPrefill {
                    handoff_id,
                    request: Request {
                        request_id,
                        tokens: vec![1, 2, 3, 4],
                        max_output_tokens: 1,
                        output_token_ids: Some(vec![5]),
                        session_id: None,
                        images: Vec::new(),
                        kv_hint: None,
                    },
                },
                CommandContext {
                    now_ms: 0.0,
                    pass_in_flight: false,
                },
                None,
            )
            .unwrap();
        assert_eq!(submission.result, CommandResult::Submitted(request_id));
        assert_eq!(rank.handoff_requests.get(&handoff_id), Some(&request_id));

        let pass = rank.execute_pass(0.0).unwrap();
        let completion = rank.complete_pass(pass.pending, pass.end_ms).unwrap();
        assert!(
            completion
                .outputs
                .iter()
                .any(|output| output.request_id == request_id && output.completed)
        );
        assert!(!rank.handoff_requests.contains_key(&handoff_id));
    }

    #[test]
    fn cold_prefill_restores_decode_owned_g2_without_a_handoff_to_the_consumer() {
        use crate::engine::host_offload::G2Binding;

        for (shared, h2d_bandwidth_gbps) in [(false, 1.0), (true, 1.0), (true, 0.01)] {
            let binding = G2Binding {
                registry: Arc::default(),
                tensor_parallel_size: 1,
            };
            let build = |worker_type, worker_id| {
                let host = NativeHostOffloadConfig::new(8).with_bandwidths(1.0, h2d_bandwidth_gbps);
                let config = EngineConfig {
                    worker_type,
                    num_gpu_blocks: 4,
                    block_size: 4,
                    max_num_seqs: 1,
                    max_num_batched_tokens: 16,
                    kv_cache_bytes_per_token: Some(250_000),
                    native_host_offload: Some(if shared {
                        host.cluster_shared("cross-role-tp1")
                    } else {
                        host
                    }),
                    timing_model: TimingModelConfig::Fixed {
                        prefill_ms: 2.0,
                        decode_ms: 1.0,
                    },
                    ..EngineConfig::default()
                };
                SchedulerRank::new_with_g2(
                    RankIdentity {
                        worker_id,
                        dp_rank: 0,
                        dp_size: NonZeroU32::MIN,
                    },
                    &config,
                    config.built_in_timing_model().unwrap(),
                    0,
                    Some(&binding),
                )
                .unwrap()
            };
            // Exactly one P rank and one D rank. Seed only D through the typed
            // native destination handshake, representing historical incoming
            // prompt KV. P has never computed or received any of these tokens.
            let mut prefill = build(WorkerType::Prefill, 0);
            let mut decode = build(WorkerType::Decode, 1);
            let history = Uuid::from_u128(95_001);
            let historical_handoff = HandoffId::from(Uuid::from_u128(95_002));
            let command = |rank: &mut SchedulerRank, command, now_ms| {
                rank.apply_command_effects(
                    command,
                    CommandContext {
                        now_ms,
                        pass_in_flight: false,
                    },
                    None,
                )
                .unwrap()
            };
            command(
                &mut decode,
                Command::ReserveDestination {
                    handoff_id: historical_handoff,
                    request: Request {
                        images: Vec::new(),
                        request_id: history,
                        tokens: (1..=9).collect(),
                        max_output_tokens: 2,
                        output_token_ids: None,
                        session_id: None,
                        kv_hint: None,
                    },
                },
                0.0,
            );
            command(
                &mut decode,
                Command::ActivateDestination {
                    handoff_id: historical_handoff,
                },
                0.0,
            );
            let mut now_ms = 0.0;
            let mut completed = false;
            for _ in 0..8 {
                let pass = decode.execute_pass(now_ms).unwrap();
                now_ms = pass.end_ms;
                let effects = decode.complete_pass(pass.pending, now_ms).unwrap();
                completed |= effects
                    .outputs
                    .iter()
                    .any(|output| output.request_id == history && output.completed);
                if completed {
                    break;
                }
            }
            assert!(completed, "historical decode must finish");
            while let Some(deadline) = decode.next_internal_deadline_ms() {
                now_ms = now_ms.max(deadline);
                decode.process_internal_work(now_ms, false).unwrap();
            }
            assert!(prefill.core.is_drained());
            if shared {
                assert_eq!(binding.registry.occupancy(), Some((8, 2, 2)));
            }

            // The consumer enters P as a new source request, not a destination
            // activation. Its only possible reuse is D's completed host copy.
            let consumer = Uuid::from_u128(95_003);
            let consumer_handoff = HandoffId::from(Uuid::from_u128(95_004));
            command(
                &mut prefill,
                Command::SubmitHandoffPrefill {
                    handoff_id: consumer_handoff,
                    request: Request {
                        images: Vec::new(),
                        request_id: consumer,
                        tokens: (1..=9).collect(),
                        max_output_tokens: 1,
                        output_token_ids: None,
                        session_id: None,
                        kv_hint: None,
                    },
                },
                now_ms,
            );
            let submitted_at_ms = now_ms;
            let mut admission = None;
            let mut admitted_at_ms = None;
            for _ in 0..8 {
                let pass = prefill.execute_pass(now_ms).unwrap();
                admission = pass
                    .start_effects
                    .admissions
                    .iter()
                    .find(|admission| admission.request_id == consumer)
                    .copied()
                    .or(admission);
                if admission.is_some() {
                    admitted_at_ms = Some(now_ms);
                }
                now_ms = pass.end_ms;
                prefill.complete_pass(pass.pending, now_ms).unwrap();
                if admission.is_some() {
                    break;
                }
                if let Some(deadline) = prefill.next_internal_deadline_ms() {
                    now_ms = now_ms.max(deadline);
                    prefill.process_internal_work(now_ms, false).unwrap();
                }
            }
            let attribution = admission.unwrap().cache_tier_attribution.unwrap();
            assert_eq!(attribution.g1_reused_input_tokens, 0);
            assert_eq!(
                attribution.host_reused_input_tokens,
                if shared { 8 } else { 0 }
            );
            // Two 1 MB blocks must finish H2D before admission. Observe the
            // start of the admitted pass, excluding its fixed 2 ms computation.
            // The slow-link case rules out an unrelated scheduler wait.
            let expected_wait_ms = if shared {
                2.0 / h2d_bandwidth_gbps
            } else {
                0.0
            };
            let actual_wait_ms = admitted_at_ms.unwrap() - submitted_at_ms;
            assert!(
                (actual_wait_ms - expected_wait_ms).abs() < 1e-9,
                "shared={shared}, H2D={h2d_bandwidth_gbps} GB/s: admission wait {actual_wait_ms} ms, expected {expected_wait_ms} ms"
            );
            command(
                &mut prefill,
                Command::ReleaseSource {
                    handoff_id: consumer_handoff,
                },
                now_ms,
            );
        }
    }
}