aisimulate-core 0.12.0

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

//! vLLM G1 manager over a minimal physical block-pool model.
//!
//! Each request lease owns its physical-copy IDs and visibility state. The
//! manager owns KV-event metadata, while the pool owns occupancy, duplicate
//! copies, prefix pins, and LRU eviction.

use uuid::Uuid;

pub(crate) use crate::engine::cache::vllm_block_pool::SourceReuseDependency;
use crate::engine::cache::vllm_block_pool::{
    BlockCopyId, BlockReservation, ReserveOutcome, VllmBlockPool,
};
use crate::engine::common::hashing::{BlockHash, SequenceHash};
use crate::engine::common::kv_cache_trace;
use crate::engine::common::protocols::{KvEventPublishers, PrefillCost};
use crate::engine::common::sequence::{BlockIdentity, RequestSequence};
use crate::engine::{KvBlock, KvEvent, KvEventData, StoredBlocks};

struct PendingStore {
    parent_hash: Option<SequenceHash>,
    local_hash: Option<BlockHash>,
    token_ids: Option<Vec<u32>>,
}

struct LeaseCommitOptions<'a> {
    cache_fresh: bool,
    sequence: Option<&'a RequestSequence>,
    write_dependencies: &'a [SourceReuseDependency],
}

#[derive(Debug)]
struct PendingCapacityWrites {
    block_indices: Vec<usize>,
    dependencies: Vec<SourceReuseDependency>,
}

#[derive(Debug)]
struct BlockLeaseEntry {
    identity: BlockIdentity,
    copy: Option<BlockCopyId>,
    /// Whether a freshly allocated full block still needs to become cache-visible.
    pending_cache: bool,
    /// Fresh capacity can be owned while an earlier transfer still reads it.
    /// The scheduler must authorize its next write after installing a fence.
    capacity_write_pending: bool,
}

/// Move-only native-G1 ownership token attached to one scheduler request.
#[derive(Debug)]
#[must_use = "a native block lease must be finished, aborted, retracted, or moved into a hold"]
pub(crate) struct BlockRequestLease {
    owner: Uuid,
    entries: Vec<BlockLeaseEntry>,
    allocated_tokens: usize,
    /// Present only while newly acquired capacity awaits a scheduler-installed
    /// write fence. The default path does not allocate this state.
    pending_capacity_writes: Option<Box<PendingCapacityWrites>>,
}

impl BlockRequestLease {
    pub(crate) fn new(owner: Uuid, identities: Vec<BlockIdentity>) -> Self {
        let mut entries = Vec::with_capacity(identities.capacity());
        entries.extend(identities.into_iter().map(|identity| BlockLeaseEntry {
            identity,
            copy: None,
            pending_cache: false,
            capacity_write_pending: false,
        }));
        Self {
            owner,
            entries,
            allocated_tokens: 0,
            pending_capacity_writes: None,
        }
    }

    pub(crate) fn owner(&self) -> Uuid {
        self.owner
    }

    pub(crate) fn allocated_tokens(&self) -> usize {
        self.allocated_tokens
    }

    pub(crate) fn resident_block_count(&self) -> usize {
        self.entries
            .iter()
            .filter(|entry| entry.copy.is_some())
            .count()
    }

    /// Project one framework-native identity without exposing physical G1
    /// copy identity.
    pub(crate) fn sequence_hash(&self, block_index: usize) -> Option<SequenceHash> {
        self.entries
            .get(block_index)
            .and_then(|entry| entry.identity.sequence_hash)
    }

    #[cfg(test)]
    pub(crate) fn entry_capacity(&self) -> usize {
        self.entries.capacity()
    }

    pub(crate) fn append_partial(&mut self) {
        // One scheduler decision can materialize more than one token (for
        // example speculative decoding). In that case the previous partial
        // block may already be logically complete but still await
        // `finalize_lease_computed_prefix`, so its identity is intentionally
        // unresolved while the next partial entry is opened.
        self.entries.push(BlockLeaseEntry {
            identity: BlockIdentity::partial(),
            copy: None,
            pending_cache: false,
            capacity_write_pending: false,
        });
    }

    fn debug_assert_owner(&self, owner: Uuid) {
        debug_assert_eq!(self.owner, owner, "native lease owner mismatch");
    }
}

struct StoredBlock {
    hash: SequenceHash,
    metadata: PendingStore,
}

struct StoreGroup {
    parent_hash: Option<SequenceHash>,
    blocks: Vec<SequenceHash>,
    local_hashes: Option<Vec<BlockHash>>,
    token_ids: Option<Vec<Vec<u32>>>,
}

impl StoreGroup {
    fn from_block(block: StoredBlock) -> Self {
        let PendingStore {
            parent_hash,
            local_hash,
            token_ids,
        } = block.metadata;
        Self {
            parent_hash,
            blocks: vec![block.hash],
            local_hashes: local_hash.map(|hash| vec![hash]),
            token_ids: token_ids.map(|ids| vec![ids]),
        }
    }

    fn can_append(&self, block: &StoredBlock) -> bool {
        self.local_hashes.is_some() == block.metadata.local_hash.is_some()
            && self.token_ids.is_some() == block.metadata.token_ids.is_some()
    }

    fn push(&mut self, block: StoredBlock) {
        self.blocks.push(block.hash);
        if let (Some(hashes), Some(hash)) = (&mut self.local_hashes, block.metadata.local_hash) {
            hashes.push(hash);
        }
        if let (Some(token_ids), Some(ids)) = (&mut self.token_ids, block.metadata.token_ids) {
            token_ids.push(ids);
        }
    }
}

pub(crate) struct DecodeBlockReservation {
    pool: BlockReservation,
}

#[must_use = "a destination reservation must be activated or explicitly cancelled"]
pub(crate) struct DestinationReservation {
    request_id: Uuid,
    block_count: usize,
    pool: BlockReservation,
}

impl DestinationReservation {
    pub(crate) fn transferable_prompt_tokens(&self, block_size: usize) -> usize {
        self.pool.fresh_len().saturating_mul(block_size)
    }

    #[cfg(test)]
    pub(crate) fn len(&self) -> usize {
        self.pool.len()
    }
}

/// Short-lived, unpinned source view used by the synchronous
/// snapshot/prepare/attach transition.
#[must_use = "a native store source snapshot must be attached or discarded synchronously"]
pub(crate) struct StoreSourceSnapshot<'a> {
    owner: Uuid,
    lease: &'a BlockRequestLease,
    block_indices: &'a [usize],
    copies: Vec<BlockCopyId>,
}

impl StoreSourceSnapshot<'_> {
    pub(crate) fn len(&self) -> usize {
        self.copies.len()
    }

    pub(crate) fn sequence_hashes(&self) -> impl ExactSizeIterator<Item = SequenceHash> + '_ {
        self.block_indices.iter().map(|&block_index| {
            self.lease.entries[block_index]
                .identity
                .sequence_hash
                .expect("validated store source lost its sequence hash")
        })
    }
}

/// Result of ordinary request-owned G1 allocation.
///
/// Capacity acquisition is complete in `Ready`, but the scheduler must fence
/// and authorize any returned dependencies before the write-producing pass.
#[must_use = "allocation dependencies must be fenced before compute"]
pub(crate) enum NativeAllocation<T> {
    Ready {
        value: T,
        dependencies: Vec<SourceReuseDependency>,
    },
    CapacityExhausted,
}

pub(super) enum VllmAcquire<T> {
    Ready(T),
    CapacityExhausted,
}

pub(crate) struct VllmKvManager {
    pool: VllmBlockPool,
    block_size: usize,
    enable_prefix_caching: bool,
    kv_event_publishers: KvEventPublishers,
    dp_rank: u32,
    next_event_id: u64,
}

impl VllmKvManager {
    pub(crate) fn new_with_event_sink(
        max_capacity: usize,
        block_size: usize,
        enable_prefix_caching: bool,
        kv_event_publishers: KvEventPublishers,
        dp_rank: u32,
    ) -> Self {
        assert!(block_size > 0, "block_size must be > 0");
        if !kv_event_publishers.is_empty() {
            tracing::info!(dp_rank, block_size, "VllmKvManager initialized");
        }
        Self {
            pool: VllmBlockPool::new(max_capacity),
            block_size,
            enable_prefix_caching,
            kv_event_publishers,
            dp_rank,
            next_event_id: 0,
        }
    }

    /// Atomically allocate the native lease through `cumulative_tokens`.
    ///
    /// Capacity is reserved before either physical residency or the allocation
    /// watermark changes, so exhaustion leaves the lease unchanged.
    pub(crate) fn allocate_lease(
        &mut self,
        owner: Uuid,
        lease: &mut BlockRequestLease,
        cumulative_tokens: usize,
        reusable_prefix_blocks: usize,
    ) -> NativeAllocation<usize> {
        lease.debug_assert_owner(owner);
        let previous_blocks = lease
            .allocated_tokens
            .div_ceil(self.block_size)
            .min(lease.entries.len());
        let target_blocks = cumulative_tokens
            .div_ceil(self.block_size)
            .min(lease.entries.len());
        if target_blocks <= previous_blocks {
            // A dependency-bearing earlier attempt may already own a larger
            // suffix than this pass can compute. Keep physical ownership
            // monotonic and report the same pending write dependency again.
            lease.allocated_tokens = lease.allocated_tokens.max(cumulative_tokens);
            return NativeAllocation::Ready {
                value: 0,
                dependencies: self.pending_capacity_write_dependencies(lease),
            };
        }
        assert!(
            lease.pending_capacity_writes.is_none(),
            "native lease cannot grow before prior write dependencies are authorized"
        );
        let newly_reusable_prefix_blocks = reusable_prefix_blocks.saturating_sub(previous_blocks);
        assert!(
            newly_reusable_prefix_blocks <= target_blocks - previous_blocks,
            "reusable prefix exceeds the newly allocated block range"
        );
        assert!(self.enable_prefix_caching || reusable_prefix_blocks == 0);

        let count = target_blocks - previous_blocks;
        let prefix = lease.entries[previous_blocks..previous_blocks + newly_reusable_prefix_blocks]
            .iter()
            .map(|entry| {
                entry
                    .identity
                    .sequence_hash
                    .expect("reusable prefix must contain only complete blocks")
            });
        let Some(ReserveOutcome {
            mut reservation,
            removed,
        }) = self.pool.reserve_exact_prefix(prefix, count)
        else {
            return NativeAllocation::CapacityExhausted;
        };
        assert_eq!(
            reservation.len() - reservation.fresh_len(),
            newly_reusable_prefix_blocks,
            "exact native prefix reservation returned the wrong hit count"
        );
        let fresh_blocks = count - newly_reusable_prefix_blocks;
        let dependencies = self
            .pool
            .reservation_next_pending_dependencies(&reservation, fresh_blocks);
        self.publish_removed(removed);
        self.commit_lease_range(
            lease,
            previous_blocks,
            target_blocks,
            &mut reservation,
            LeaseCommitOptions {
                cache_fresh: false,
                sequence: None,
                write_dependencies: &dependencies,
            },
        );
        assert_eq!(reservation.len(), 0, "native reservation was not consumed");
        self.pool.cancel(reservation);
        lease.allocated_tokens = cumulative_tokens;
        NativeAllocation::Ready {
            value: count,
            dependencies,
        }
    }

    /// Record that the scheduler ordered every pending write behind its source
    /// dependencies. This authorizes writes without terminalizing the reads.
    pub(crate) fn authorize_lease_writes_after_dependencies(
        &mut self,
        owner: Uuid,
        lease: &mut BlockRequestLease,
        dependencies: &[SourceReuseDependency],
    ) {
        lease.debug_assert_owner(owner);
        assert!(
            !dependencies.is_empty(),
            "native compute authorization requires at least one dependency"
        );
        let pending = lease
            .pending_capacity_writes
            .as_ref()
            .expect("native compute authorization has no pending capacity writes");
        assert_eq!(
            pending.dependencies.as_slice(),
            dependencies,
            "native compute fence does not cover acquired capacity dependencies"
        );
        #[cfg(debug_assertions)]
        {
            let actual = self
                .pool
                .copies_pending_dependencies(pending.block_indices.iter().map(|&block_index| {
                    lease.entries[block_index]
                        .copy
                        .expect("pending capacity write lost its request-owned copy")
                }));
            debug_assert_eq!(actual.as_slice(), dependencies);
        }
        self.pool.authorize_source_reuse_writes(
            pending.block_indices.iter().map(|&block_index| {
                lease.entries[block_index]
                    .copy
                    .expect("pending capacity write lost its request-owned copy")
            }),
            dependencies,
        );
        let pending = lease
            .pending_capacity_writes
            .take()
            .expect("checked pending write state disappeared");
        for block_index in pending.block_indices {
            lease.entries[block_index].capacity_write_pending = false;
        }
    }

    pub(crate) fn allocate_lease_from_decode_reservation(
        &mut self,
        owner: Uuid,
        lease: &mut BlockRequestLease,
        cumulative_tokens: usize,
        reservation: &mut DecodeBlockReservation,
    ) -> NativeAllocation<usize> {
        lease.debug_assert_owner(owner);
        let previous_blocks = lease
            .allocated_tokens
            .div_ceil(self.block_size)
            .min(lease.entries.len());
        let target_blocks = cumulative_tokens
            .div_ceil(self.block_size)
            .min(lease.entries.len());
        if target_blocks <= previous_blocks {
            lease.allocated_tokens = lease.allocated_tokens.max(cumulative_tokens);
            return NativeAllocation::Ready {
                value: 0,
                dependencies: self.pending_capacity_write_dependencies(lease),
            };
        }
        assert!(
            lease.pending_capacity_writes.is_none(),
            "native lease cannot grow before prior write dependencies are authorized"
        );
        let count = target_blocks - previous_blocks;
        assert!(
            reservation.pool.fresh_len() >= count,
            "decode reservation does not cover the native lease growth"
        );
        let dependencies = self
            .pool
            .reservation_next_pending_dependencies(&reservation.pool, count);
        self.commit_lease_range(
            lease,
            previous_blocks,
            target_blocks,
            &mut reservation.pool,
            LeaseCommitOptions {
                cache_fresh: false,
                sequence: None,
                write_dependencies: &dependencies,
            },
        );
        lease.allocated_tokens = cumulative_tokens;
        NativeAllocation::Ready {
            value: count,
            dependencies,
        }
    }

    pub(crate) fn finalize_lease_computed_prefix(
        &mut self,
        owner: Uuid,
        sequence: &mut RequestSequence,
        lease: &mut BlockRequestLease,
        computed_before: usize,
        computed_after: usize,
    ) {
        lease.debug_assert_owner(owner);
        assert!(
            computed_before <= computed_after,
            "computed token count cannot move backwards during one scheduling decision"
        );
        let first_new_block = computed_before / self.block_size;
        let completed_blocks = (computed_after / self.block_size).min(lease.entries.len());
        if first_new_block >= completed_blocks {
            return;
        }

        let materialize_store_events =
            self.enable_prefix_caching && self.materialize_store_events();
        let mut stores = materialize_store_events
            .then(|| Vec::with_capacity(completed_blocks - first_new_block));
        for position in first_new_block..completed_blocks {
            let parent_hash = position
                .checked_sub(1)
                .and_then(|parent| lease.entries[parent].identity.sequence_hash);
            if lease.entries[position].identity.sequence_hash.is_none() {
                lease.entries[position].identity =
                    sequence.complete_block_identity(position, parent_hash);
                lease.entries[position].pending_cache = self.enable_prefix_caching;
            }

            let entry = &mut lease.entries[position];
            assert!(
                !entry.capacity_write_pending,
                "cannot finalize dependency-bearing capacity before write authorization"
            );
            if !entry.pending_cache {
                if let Some(stores) = &mut stores {
                    stores.push(None);
                }
                sequence.discard_completed_block(position);
                continue;
            }
            entry.pending_cache = false;
            let copy = entry
                .copy
                .expect("computed native block must retain physical residency");
            let hash = entry
                .identity
                .sequence_hash
                .expect("computed native block must have a sequence hash");
            let became_visible = self.pool.cache_private(copy, hash);
            if let Some(stores) = &mut stores {
                stores.push(became_visible.then(|| StoredBlock {
                    hash,
                    metadata: PendingStore {
                        parent_hash,
                        local_hash: entry.identity.local_hash,
                        token_ids: sequence.block_token_ids(position),
                    },
                }));
            }
            sequence.discard_completed_block(position);
        }
        if let Some(stores) = stores {
            self.publish_store_sequence(stores);
        }
        #[cfg(debug_assertions)]
        sequence.debug_assert_finalized_range(
            lease.entries.len(),
            lease.entries[first_new_block..completed_blocks]
                .iter()
                .map(|entry| entry.identity),
            lease.entries.last().map(|entry| entry.identity),
        );
    }

    pub(crate) fn reserve_destination_lease(
        &mut self,
        owner: Uuid,
        sequence: &RequestSequence,
        lease: &BlockRequestLease,
        _eviction_now_ms: Option<f64>,
    ) -> VllmAcquire<DestinationReservation> {
        lease.debug_assert_owner(owner);
        assert_eq!(
            lease.resident_block_count(),
            0,
            "destination request already owns physical blocks"
        );
        let prompt_blocks = sequence
            .num_input_tokens()
            .div_ceil(self.block_size)
            .min(lease.entries.len());
        let prefix_candidates = lease.entries[..prompt_blocks]
            .iter()
            .map_while(|entry| entry.identity.sequence_hash);
        let Some(outcome) = self
            .pool
            .reserve_resident_prefix(prefix_candidates, prompt_blocks)
        else {
            return VllmAcquire::CapacityExhausted;
        };
        self.publish_removed(outcome.removed);
        VllmAcquire::Ready(DestinationReservation {
            request_id: owner,
            block_count: prompt_blocks,
            pool: outcome.reservation,
        })
    }

    /// Reserve an exact already-authorized G1 prefix plus only the contiguous
    /// suffix selected by an external logical-cache lookup.
    pub(crate) fn reserve_external_prefix_lease(
        &mut self,
        owner: Uuid,
        lease: &BlockRequestLease,
        g1_prefix_blocks: usize,
        transferred_suffix: &[SequenceHash],
    ) -> VllmAcquire<DestinationReservation> {
        lease.debug_assert_owner(owner);
        assert_eq!(
            lease.resident_block_count(),
            0,
            "external destination request already owns physical blocks"
        );
        assert!(
            !transferred_suffix.is_empty(),
            "external destination suffix must not be empty"
        );
        let block_count = g1_prefix_blocks
            .checked_add(transferred_suffix.len())
            .expect("external destination block count overflow");
        assert!(
            block_count <= lease.entries.len(),
            "external destination exceeds the request block table"
        );
        for (offset, &hash) in transferred_suffix.iter().enumerate() {
            assert_eq!(
                lease.entries[g1_prefix_blocks + offset]
                    .identity
                    .sequence_hash,
                Some(hash),
                "external suffix does not match the request sequence"
            );
        }
        let prefix = lease.entries[..g1_prefix_blocks].iter().map(|entry| {
            entry
                .identity
                .sequence_hash
                .expect("G1 destination prefix must contain complete blocks")
        });
        let Some(outcome) = self.pool.reserve_exact_prefix(prefix, block_count) else {
            return VllmAcquire::CapacityExhausted;
        };
        assert_eq!(
            outcome.reservation.fresh_len(),
            transferred_suffix.len(),
            "external reservation changed the exact tier boundary"
        );
        self.publish_removed(outcome.removed);
        VllmAcquire::Ready(DestinationReservation {
            request_id: owner,
            block_count,
            pool: outcome.reservation,
        })
    }

    pub(crate) fn destination_pending_dependencies(
        &self,
        reservation: &DestinationReservation,
    ) -> Vec<SourceReuseDependency> {
        self.pool
            .reservation_pending_dependencies(&reservation.pool)
    }

    /// Capture a pure, unpinned view of completed request-owned source blocks.
    ///
    /// The returned snapshot also proves that dependency attachment is valid
    /// at this point. The caller may synchronously admit a store using the
    /// projected hashes, then consume the snapshot through
    /// [`Self::attach_store_source_dependency`] without yielding or mutating G1.
    pub(crate) fn snapshot_store_sources<'a>(
        &self,
        owner: Uuid,
        lease: &'a BlockRequestLease,
        block_indices: &'a [usize],
    ) -> Option<StoreSourceSnapshot<'a>> {
        lease.debug_assert_owner(owner);
        assert!(
            !block_indices.is_empty(),
            "native store source cohort must not be empty"
        );
        let mut copies = Vec::with_capacity(block_indices.len());
        let mut previous = None;
        for &block_index in block_indices {
            if previous.is_some_and(|previous| previous >= block_index) {
                return None;
            }
            previous = Some(block_index);
            let entry = lease.entries.get(block_index)?;
            entry.identity.sequence_hash?;
            let copy = entry.copy?;
            if entry.pending_cache || entry.capacity_write_pending {
                return None;
            }
            copies.push(copy);
        }
        if !self.pool.can_attach_source_reuse_dependency(&copies) {
            return None;
        }
        Some(StoreSourceSnapshot {
            owner,
            lease,
            block_indices,
            copies,
        })
    }

    /// Infallibly attach one prepared transfer identity to a validated source
    /// snapshot. Callers must keep this in the same non-yielding transition as
    /// snapshot validation and external store admission.
    pub(crate) fn attach_store_source_dependency(
        &mut self,
        owner: Uuid,
        snapshot: StoreSourceSnapshot<'_>,
        dependency: SourceReuseDependency,
    ) {
        snapshot.lease.debug_assert_owner(owner);
        assert_eq!(snapshot.owner, owner, "native store source owner mismatch");
        self.pool
            .attach_source_reuse_dependency(&snapshot.copies, dependency);
    }

    pub(crate) fn satisfy_source_reuse_dependency(
        &mut self,
        dependency: SourceReuseDependency,
    ) -> bool {
        self.pool.satisfy_source_reuse_dependency(dependency)
    }

    pub(crate) fn is_source_reuse_dependency_pending(
        &self,
        dependency: SourceReuseDependency,
    ) -> bool {
        self.pool.is_source_reuse_dependency_pending(dependency)
    }

    pub(crate) fn activate_destination_lease(
        &mut self,
        owner: Uuid,
        sequence: &RequestSequence,
        lease: &mut BlockRequestLease,
        mut reservation: DestinationReservation,
    ) {
        lease.debug_assert_owner(owner);
        debug_assert_eq!(
            lease.resident_block_count(),
            0,
            "destination request already owns physical blocks"
        );
        assert_eq!(reservation.request_id, owner, "destination owner mismatch");
        let prompt_blocks = reservation.block_count;
        assert!(
            prompt_blocks
                <= sequence
                    .num_input_tokens()
                    .div_ceil(self.block_size)
                    .min(lease.entries.len()),
            "destination reservation exceeds the request prompt"
        );
        assert!(
            self.pool
                .reservation_pending_dependencies(&reservation.pool)
                .is_empty(),
            "cannot activate a destination before source reuse dependencies are terminal"
        );
        self.commit_lease_range(
            lease,
            0,
            prompt_blocks,
            &mut reservation.pool,
            LeaseCommitOptions {
                cache_fresh: self.enable_prefix_caching,
                sequence: Some(sequence),
                write_dependencies: &[],
            },
        );
        lease.allocated_tokens = (prompt_blocks * self.block_size).min(sequence.num_input_tokens());
        assert_eq!(
            reservation.pool.len(),
            0,
            "destination reservation was not consumed"
        );
        self.pool.cancel(reservation.pool);
    }

    pub(crate) fn preempt_lease(&mut self, owner: Uuid, lease: &mut BlockRequestLease) {
        lease.debug_assert_owner(owner);
        self.release_lease_entries(lease);
        lease.allocated_tokens = 0;
    }

    pub(crate) fn finish_lease(&mut self, owner: Uuid, mut lease: BlockRequestLease) {
        lease.debug_assert_owner(owner);
        self.release_lease_entries(&mut lease);
        lease.allocated_tokens = 0;
    }

    fn release_lease_entries(&mut self, lease: &mut BlockRequestLease) {
        lease.pending_capacity_writes = None;
        for entry in lease.entries.iter_mut().rev() {
            if let Some(copy) = entry.copy.take() {
                self.pool.release(copy);
            }
            entry.pending_cache = false;
            entry.capacity_write_pending = false;
        }
    }

    fn pending_capacity_write_dependencies(
        &self,
        lease: &mut BlockRequestLease,
    ) -> Vec<SourceReuseDependency> {
        let Some(pending) = lease.pending_capacity_writes.as_mut() else {
            return Vec::new();
        };
        pending
            .dependencies
            .retain(|dependency| self.pool.is_source_reuse_dependency_pending(*dependency));
        if pending.dependencies.is_empty() {
            let pending = lease
                .pending_capacity_writes
                .take()
                .expect("checked pending write state disappeared");
            for block_index in pending.block_indices {
                lease.entries[block_index].capacity_write_pending = false;
            }
            Vec::new()
        } else {
            pending.dependencies.clone()
        }
    }

    fn commit_lease_range(
        &mut self,
        lease: &mut BlockRequestLease,
        start: usize,
        end: usize,
        reservation: &mut BlockReservation,
        options: LeaseCommitOptions<'_>,
    ) {
        let LeaseCommitOptions {
            cache_fresh,
            sequence,
            write_dependencies,
        } = options;
        let fresh_write_pending = !write_dependencies.is_empty();
        assert!(start <= end && end <= lease.entries.len());
        let prefix_len = reservation.len() - reservation.fresh_len();
        let mut prefix_copies = self.pool.activate_prefix(reservation);
        assert_eq!(prefix_copies.len(), prefix_len);
        let mut pending_indices = fresh_write_pending
            .then(|| Vec::with_capacity((end - start).saturating_sub(prefix_len)));
        let materialize_store_events = self.materialize_store_events();
        let mut stores =
            (cache_fresh && materialize_store_events).then(|| Vec::with_capacity(end - start));

        for (offset, position) in (start..end).enumerate() {
            let parent_hash = position
                .checked_sub(1)
                .and_then(|parent| lease.entries[parent].identity.sequence_hash);
            let entry = &mut lease.entries[position];
            assert!(
                entry.copy.is_none(),
                "native lease entry is already resident"
            );
            if offset < prefix_len {
                let (hash, copy) = prefix_copies
                    .next()
                    .expect("prefix reservation returned too few copies");
                assert_eq!(
                    entry.identity.sequence_hash,
                    Some(hash),
                    "reserved prefix hash changed before activation"
                );
                entry.copy = Some(copy);
                entry.pending_cache = false;
                entry.capacity_write_pending = false;
                if let Some(stores) = &mut stores {
                    stores.push(None);
                }
                continue;
            }

            let Some(hash) = entry.identity.sequence_hash else {
                entry.copy = Some(self.pool.allocate_private(reservation));
                entry.pending_cache = false;
                entry.capacity_write_pending = fresh_write_pending;
                if let Some(indices) = &mut pending_indices {
                    indices.push(position);
                }
                if let Some(stores) = &mut stores {
                    stores.push(None);
                }
                continue;
            };
            if cache_fresh && self.enable_prefix_caching {
                let (copy, became_visible) = self.pool.allocate_cached(reservation, hash);
                entry.copy = Some(copy);
                entry.pending_cache = false;
                entry.capacity_write_pending = fresh_write_pending;
                if let Some(indices) = &mut pending_indices {
                    indices.push(position);
                }
                if let Some(stores) = &mut stores {
                    stores.push(became_visible.then(|| StoredBlock {
                        hash,
                        metadata: PendingStore {
                            parent_hash,
                            local_hash: entry.identity.local_hash,
                            token_ids: sequence.and_then(|seq| seq.block_token_ids(position)),
                        },
                    }));
                }
            } else {
                entry.copy = Some(self.pool.allocate_private(reservation));
                entry.pending_cache = self.enable_prefix_caching;
                entry.capacity_write_pending = fresh_write_pending;
                if let Some(indices) = &mut pending_indices {
                    indices.push(position);
                }
                if let Some(stores) = &mut stores {
                    stores.push(None);
                }
            }
        }
        assert!(prefix_copies.next().is_none());
        if let Some(block_indices) = pending_indices {
            assert!(
                lease.pending_capacity_writes.is_none(),
                "new capacity was acquired before prior write dependencies were authorized"
            );
            lease.pending_capacity_writes = Some(Box::new(PendingCapacityWrites {
                block_indices,
                dependencies: write_dependencies.to_vec(),
            }));
        }
        if let Some(stores) = stores {
            self.publish_store_sequence(stores);
        }
    }

    pub(crate) fn get_lease_prefill_cost(
        &self,
        sequence: &RequestSequence,
        lease: &BlockRequestLease,
    ) -> PrefillCost {
        let (overlap_blocks, active_overlap_blocks) =
            if self.enable_prefix_caching && sequence.enable_prefix_caching() {
                let mut overlap = 0;
                let mut active = 0;
                for entry in &lease.entries {
                    let Some(hash) = entry.identity.sequence_hash else {
                        break;
                    };
                    let Some(hit) = self.pool.prefix_hit(hash) else {
                        break;
                    };
                    overlap += 1;
                    active += usize::from(hit.is_active);
                }
                (overlap, active)
            } else {
                (0, 0)
            };
        let new_blocks = lease.entries.len() - overlap_blocks;
        let cached_tokens = (overlap_blocks * self.block_size).min(sequence.len());
        let active_cached_tokens = (active_overlap_blocks * self.block_size).min(sequence.len());
        PrefillCost {
            new_blocks,
            new_tokens: sequence.len() - cached_tokens,
            cached_tokens,
            active_cached_tokens,
        }
    }

    pub(crate) fn reserve_decode_blocks(
        &mut self,
        count: usize,
    ) -> VllmAcquire<DecodeBlockReservation> {
        let Some(outcome) = self.pool.reserve(&[], count) else {
            return VllmAcquire::CapacityExhausted;
        };
        self.publish_removed(outcome.removed);
        VllmAcquire::Ready(DecodeBlockReservation {
            pool: outcome.reservation,
        })
    }

    pub(crate) fn release_decode_reservation(&mut self, reservation: DecodeBlockReservation) {
        self.pool.cancel(reservation.pool);
    }

    pub(crate) fn cancel_destination(&mut self, reservation: DestinationReservation) {
        self.pool.cancel(reservation.pool);
    }

    fn materialize_store_events(&self) -> bool {
        !self.kv_event_publishers.is_empty() || *kv_cache_trace::KV_CACHE_TRACE_ENABLED
    }

    fn publish_store_sequence(&mut self, stores: Vec<Option<StoredBlock>>) {
        let mut group: Option<StoreGroup> = None;
        for store in stores {
            let Some(store) = store else {
                self.flush_store_group(&mut group);
                continue;
            };
            if group
                .as_ref()
                .is_some_and(|current| !current.can_append(&store))
            {
                self.flush_store_group(&mut group);
            }
            match &mut group {
                Some(current) => current.push(store),
                None => group = Some(StoreGroup::from_block(store)),
            }
        }
        self.flush_store_group(&mut group);
    }

    fn flush_store_group(&mut self, group: &mut Option<StoreGroup>) {
        let Some(group) = group.take() else {
            return;
        };
        self.publish_kv_event(
            group.blocks,
            group.local_hashes.as_deref().unwrap_or(&[]),
            group.parent_hash,
            true,
            group.token_ids,
        );
    }

    fn publish_removed(&mut self, hashes: Vec<SequenceHash>) {
        if !hashes.is_empty() {
            self.publish_kv_event(hashes, &[], None, false, None);
        }
    }

    fn publish_kv_event(
        &mut self,
        full_blocks: Vec<SequenceHash>,
        local_hashes: &[BlockHash],
        parent_hash: Option<SequenceHash>,
        is_store: bool,
        token_ids: Option<Vec<Vec<u32>>>,
    ) {
        if !self.enable_prefix_caching || full_blocks.is_empty() {
            return;
        }
        if *kv_cache_trace::KV_CACHE_TRACE_ENABLED {
            kv_cache_trace::log_vllm_trace(
                if is_store { "allocation" } else { "eviction" },
                self.dp_rank,
                self.block_size,
                self.num_active_blocks(),
                self.num_inactive_blocks(),
                self.max_capacity(),
            );
        }
        if self.kv_event_publishers.is_empty() {
            return;
        }
        assert!(local_hashes.is_empty() || local_hashes.len() == full_blocks.len());
        assert!(
            token_ids
                .as_ref()
                .is_none_or(|ids| ids.len() == full_blocks.len())
        );

        let data = if is_store {
            KvEventData::Stored(StoredBlocks {
                parent_hash,
                start_position: None,
                blocks: full_blocks
                    .into_iter()
                    .enumerate()
                    .map(|(index, hash)| KvBlock {
                        block_hash: hash,
                        tokens_hash: local_hashes.get(index).copied().unwrap_or_default(),
                        token_ids: token_ids.as_ref().and_then(|ids| ids.get(index).cloned()),
                    })
                    .collect(),
            })
        } else {
            KvEventData::Removed {
                block_hashes: full_blocks,
            }
        };
        let event = KvEvent {
            event_id: self.next_event_id,
            data,
            dp_rank: self.dp_rank,
        };
        self.next_event_id = self
            .next_event_id
            .checked_add(1)
            .unwrap_or_else(|| panic!("KV event ID overflow"));
        if let Err(error) = self
            .kv_event_publishers
            .publish(event, token_ids.as_deref())
        {
            tracing::warn!(error = %error, "failed to publish native G1 KV event");
        }
    }

    pub(crate) fn num_active_blocks(&self) -> usize {
        self.pool.num_active()
    }

    pub(crate) fn num_inactive_blocks(&self) -> usize {
        self.pool.num_inactive()
    }

    pub(crate) fn max_capacity(&self) -> usize {
        self.pool.capacity()
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};

    use super::*;
    use crate::engine::common::protocols::KvCacheEventSink;

    #[derive(Default)]
    struct CapturingNativeSink {
        events: Mutex<Vec<KvEvent>>,
    }

    impl CapturingNativeSink {
        fn take(&self) -> Vec<KvEvent> {
            std::mem::take(&mut *self.events.lock().unwrap())
        }
    }

    impl KvCacheEventSink for CapturingNativeSink {
        fn publish(&self, event: KvEvent) -> anyhow::Result<()> {
            self.events.lock().unwrap().push(event);
            Ok(())
        }
    }

    fn request(
        owner: Uuid,
        hashes: &[u64],
        emit_token_ids: bool,
    ) -> (RequestSequence, BlockRequestLease) {
        let tokens = (0..hashes.len() * 4).map(|token| token as u32).collect();
        let (sequence, _) = RequestSequence::new(tokens, 0, 0, 4, true, true, emit_token_ids, None);
        let identities = hashes
            .iter()
            .copied()
            .map(|hash| BlockIdentity {
                sequence_hash: Some(hash),
                local_hash: Some(hash + 100),
            })
            .collect();
        (sequence, BlockRequestLease::new(owner, identities))
    }

    trait TestReady<T> {
        fn ready(self) -> T;
    }

    impl<T> TestReady<T> for VllmAcquire<T> {
        fn ready(self) -> T {
            match self {
                VllmAcquire::Ready(value) => value,
                VllmAcquire::CapacityExhausted => panic!("unexpected allocation failure"),
            }
        }
    }

    impl<T> TestReady<T> for NativeAllocation<T> {
        fn ready(self) -> T {
            match self {
                NativeAllocation::Ready {
                    value,
                    dependencies,
                } => {
                    assert!(dependencies.is_empty());
                    value
                }
                NativeAllocation::CapacityExhausted => {
                    panic!("unexpected allocation failure")
                }
            }
        }
    }

    fn ready<T>(outcome: impl TestReady<T>) -> T {
        outcome.ready()
    }

    fn finish_source_with_dependency(
        manager: &mut VllmKvManager,
        owner: Uuid,
        hashes: &[u64],
        block_indices: &[usize],
        dependency: SourceReuseDependency,
    ) {
        let (mut sequence, mut lease) = request(owner, hashes, false);
        ready(manager.allocate_lease(owner, &mut lease, hashes.len() * 4, 0));
        manager.finalize_lease_computed_prefix(
            owner,
            &mut sequence,
            &mut lease,
            0,
            hashes.len() * 4,
        );
        let snapshot = manager
            .snapshot_store_sources(owner, &lease, block_indices)
            .expect("completed source should be snapshotable");
        assert!(snapshot.sequence_hashes().eq(hashes.iter().copied()));
        manager.attach_store_source_dependency(owner, snapshot, dependency);
        manager.finish_lease(owner, lease);
        assert_eq!(manager.num_inactive_blocks(), hashes.len());
    }

    #[test]
    fn duplicate_full_hashes_consume_physical_capacity() {
        let mut manager =
            VllmKvManager::new_with_event_sink(2, 4, true, KvEventPublishers::default(), 0);
        let mut leases = Vec::new();
        for owner in [Uuid::from_u128(1), Uuid::from_u128(2)] {
            let (mut sequence, mut lease) = request(owner, &[7], false);
            ready(manager.allocate_lease(owner, &mut lease, 4, 0));
            manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 0, 4);
            leases.push(lease);
        }
        assert_eq!(manager.num_active_blocks(), 2);
        let third = Uuid::from_u128(3);
        let (_, mut third_lease) = request(third, &[8], false);
        assert!(matches!(
            manager.allocate_lease(third, &mut third_lease, 4, 0),
            NativeAllocation::CapacityExhausted
        ));
        assert_eq!(third_lease.allocated_tokens(), 0);
        assert_eq!(third_lease.resident_block_count(), 0);
    }

    #[test]
    fn authorized_prefix_reuses_one_physical_copy() {
        let mut manager =
            VllmKvManager::new_with_event_sink(2, 4, true, KvEventPublishers::default(), 0);
        let first = Uuid::from_u128(1);
        let (mut first_sequence, mut first_lease) = request(first, &[7], false);
        ready(manager.allocate_lease(first, &mut first_lease, 4, 0));
        manager.finalize_lease_computed_prefix(first, &mut first_sequence, &mut first_lease, 0, 4);
        manager.finish_lease(first, first_lease);

        let second = Uuid::from_u128(2);
        let (_, mut second_lease) = request(second, &[7], false);
        ready(manager.allocate_lease(second, &mut second_lease, 4, 1));
        assert_eq!(manager.num_active_blocks(), 1);
        assert_eq!(manager.num_inactive_blocks(), 0);
    }

    #[test]
    fn later_native_allocation_does_not_reacquire_owned_prefix() {
        let mut manager =
            VllmKvManager::new_with_event_sink(2, 4, true, KvEventPublishers::default(), 0);

        let seed = Uuid::from_u128(1);
        let (mut seed_sequence, mut seed_lease) = request(seed, &[7], false);
        ready(manager.allocate_lease(seed, &mut seed_lease, 4, 0));
        manager.finalize_lease_computed_prefix(seed, &mut seed_sequence, &mut seed_lease, 0, 4);
        let prefix_copy = seed_lease.entries[0].copy;
        manager.finish_lease(seed, seed_lease);

        let owner = Uuid::from_u128(3);
        let (_, mut lease) = request(owner, &[7, 8], false);
        assert_eq!(ready(manager.allocate_lease(owner, &mut lease, 4, 1)), 1);
        assert_eq!(lease.entries[0].copy, prefix_copy);

        assert_eq!(ready(manager.allocate_lease(owner, &mut lease, 8, 1)), 1);
        assert_eq!(lease.entries[0].copy, prefix_copy);
        assert_eq!(lease.allocated_tokens(), 8);
        assert_eq!(lease.resident_block_count(), 2);
        assert_eq!(manager.num_active_blocks(), 2);
    }

    #[test]
    #[should_panic(expected = "reusable prefix exceeds the newly allocated block range")]
    fn native_allocation_rejects_excessive_reusable_prefix() {
        let mut manager =
            VllmKvManager::new_with_event_sink(2, 4, true, KvEventPublishers::default(), 0);
        let owner = Uuid::from_u128(4);
        let (_, mut lease) = request(owner, &[7, 8], false);

        let _ = manager.allocate_lease(owner, &mut lease, 4, 2);
    }

    #[test]
    fn full_block_is_hidden_until_computed() {
        let mut manager =
            VllmKvManager::new_with_event_sink(2, 4, true, KvEventPublishers::default(), 0);
        let owner = Uuid::from_u128(1);
        let (mut sequence, mut lease) = request(owner, &[7], false);
        ready(manager.allocate_lease(owner, &mut lease, 4, 0));
        assert!(manager.pool.prefix_hit(7).is_none());
        manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 0, 4);
        assert!(manager.pool.prefix_hit(7).is_some());
    }

    #[test]
    fn finalization_only_visits_blocks_completed_by_this_decision() {
        let mut manager =
            VllmKvManager::new_with_event_sink(2, 4, true, KvEventPublishers::default(), 0);
        let owner = Uuid::from_u128(1);
        let (mut sequence, mut lease) = request(owner, &[7, 8], false);
        ready(manager.allocate_lease(owner, &mut lease, 8, 0));

        manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 0, 4);
        assert!(manager.pool.prefix_hit(7).is_some());
        assert!(manager.pool.prefix_hit(8).is_none());

        manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 4, 8);
        assert!(manager.pool.prefix_hit(8).is_some());
    }

    #[test]
    fn finalization_handles_unaligned_decision_boundaries() {
        let mut manager =
            VllmKvManager::new_with_event_sink(3, 4, true, KvEventPublishers::default(), 0);
        let owner = Uuid::from_u128(1);
        let (mut sequence, mut lease) = request(owner, &[7, 8, 9], false);
        ready(manager.allocate_lease(owner, &mut lease, 12, 0));

        manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 3, 9);
        assert!(manager.pool.prefix_hit(7).is_some());
        assert!(manager.pool.prefix_hit(8).is_some());
        assert!(manager.pool.prefix_hit(9).is_none());
    }

    #[test]
    fn cached_prefix_watermark_finalizes_only_the_fresh_suffix() {
        let mut manager =
            VllmKvManager::new_with_event_sink(2, 4, true, KvEventPublishers::default(), 0);
        let seed = Uuid::from_u128(1);
        let (mut seed_sequence, mut seed_lease) = request(seed, &[7], false);
        ready(manager.allocate_lease(seed, &mut seed_lease, 4, 0));
        manager.finalize_lease_computed_prefix(seed, &mut seed_sequence, &mut seed_lease, 0, 4);
        manager.finish_lease(seed, seed_lease);

        let owner = Uuid::from_u128(2);
        let (mut sequence, mut lease) = request(owner, &[7, 8], false);
        ready(manager.allocate_lease(owner, &mut lease, 8, 1));
        manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 4, 8);
        assert!(manager.pool.prefix_hit(7).is_some());
        assert!(manager.pool.prefix_hit(8).is_some());
    }

    #[test]
    fn request_release_evicts_leaf_before_parent() {
        let mut manager =
            VllmKvManager::new_with_event_sink(2, 4, true, KvEventPublishers::default(), 0);
        let owner = Uuid::from_u128(1);
        let (mut sequence, mut lease) = request(owner, &[7, 8], false);
        ready(manager.allocate_lease(owner, &mut lease, 8, 0));
        manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 0, 8);
        manager.finish_lease(owner, lease);

        let next = Uuid::from_u128(2);
        let (_, mut next_lease) = request(next, &[9], false);
        ready(manager.allocate_lease(next, &mut next_lease, 4, 0));

        assert!(
            manager.pool.prefix_hit(7).is_some(),
            "parent should remain resident"
        );
        assert!(
            manager.pool.prefix_hit(8).is_none(),
            "leaf should be evicted first"
        );
    }

    #[test]
    fn vllm_prefix_cache_example_covers_the_full_block_lifecycle() {
        let sink = Arc::new(CapturingNativeSink::default());
        let publishers = KvEventPublishers::new(Some(sink.clone()));
        let mut manager = VllmKvManager::new_with_event_sink(10, 4, true, publishers, 0);

        // Time 1 in vLLM's prefix-caching example: three complete prompt
        // blocks plus one partial tail occupy four physical blocks.
        let first = Uuid::from_u128(1);
        let first_tokens = (0..15).collect::<Vec<u32>>();
        let (mut first_sequence, first_identities) =
            RequestSequence::new(first_tokens, 2, 2, 4, true, true, false, Some(vec![15, 16]));
        let mut first_lease = BlockRequestLease::new(first, first_identities);
        ready(manager.allocate_lease(first, &mut first_lease, 15, 0));
        manager.finalize_lease_computed_prefix(first, &mut first_sequence, &mut first_lease, 0, 15);
        assert_eq!(manager.num_active_blocks(), 4);
        assert_eq!(manager.num_inactive_blocks(), 0);

        // Time 2: output completes the partial block and opens a new partial
        // tail. Only the completed block becomes prefix-cache visible.
        let (_, opened_partial) = first_sequence.generate_token();
        assert!(!opened_partial);
        ready(manager.allocate_lease(first, &mut first_lease, 16, 0));
        manager.finalize_lease_computed_prefix(
            first,
            &mut first_sequence,
            &mut first_lease,
            15,
            16,
        );
        let (_, opened_partial) = first_sequence.generate_token();
        assert!(opened_partial);
        first_lease.append_partial();
        ready(manager.allocate_lease(first, &mut first_lease, 17, 0));
        manager.finalize_lease_computed_prefix(
            first,
            &mut first_sequence,
            &mut first_lease,
            16,
            17,
        );
        assert_eq!(manager.num_active_blocks(), 5);
        assert_eq!(manager.num_inactive_blocks(), 0);

        let first_hashes = first_lease.entries[..4]
            .iter()
            .map(|entry| entry.identity.sequence_hash.unwrap())
            .collect::<Vec<_>>();

        // Time 3: the second request shares ten prompt tokens with the first.
        // Only its first two complete blocks hit; its divergent complete block
        // and partial tail consume two additional physical blocks.
        let second = Uuid::from_u128(2);
        let mut second_tokens = (0..10).collect::<Vec<u32>>();
        second_tokens.extend([100, 101, 102, 103]);
        let (mut second_sequence, second_identities) =
            RequestSequence::new(second_tokens, 0, 0, 4, true, true, false, None);
        let mut second_lease = BlockRequestLease::new(second, second_identities);
        let prefill = manager.get_lease_prefill_cost(&second_sequence, &second_lease);
        assert_eq!(prefill.cached_tokens, 8);
        assert_eq!(prefill.active_cached_tokens, 8);
        assert_eq!(prefill.new_blocks, 2);
        assert_eq!(prefill.new_tokens, 6);
        ready(manager.allocate_lease(second, &mut second_lease, 14, 2));
        manager.finalize_lease_computed_prefix(
            second,
            &mut second_sequence,
            &mut second_lease,
            0,
            14,
        );
        let second_divergent_hash = second_lease.entries[2].identity.sequence_hash.unwrap();
        assert_eq!(manager.num_active_blocks(), 7);
        assert_eq!(manager.num_inactive_blocks(), 0);

        // Time 4: finishing request 0 frees its private tail and makes its
        // unique cached suffix inactive. Shared prefix blocks remain active.
        manager.finish_lease(first, first_lease);
        assert_eq!(manager.num_active_blocks(), 4);
        assert_eq!(manager.num_inactive_blocks(), 2);
        assert!(manager.pool.prefix_hit(first_hashes[0]).unwrap().is_active);
        assert!(manager.pool.prefix_hit(first_hashes[1]).unwrap().is_active);
        assert!(!manager.pool.prefix_hit(first_hashes[2]).unwrap().is_active);
        assert!(!manager.pool.prefix_hit(first_hashes[3]).unwrap().is_active);

        // Time 5: finishing request 1 frees its private tail and leaves all
        // five complete blocks cached but inactive.
        manager.finish_lease(second, second_lease);
        assert_eq!(manager.num_active_blocks(), 0);
        assert_eq!(manager.num_inactive_blocks(), 5);
        sink.take();

        // Five unused slots satisfy the next request first. Its four
        // remaining blocks evict the oldest release batch first, with each
        // request's suffixes preceding its prefixes.
        let pressure_owner = Uuid::from_u128(3);
        let pressure_hashes = (30..39).collect::<Vec<_>>();
        let (mut pressure_sequence, mut pressure_lease) =
            request(pressure_owner, &pressure_hashes, false);
        ready(manager.allocate_lease(pressure_owner, &mut pressure_lease, 36, 0));
        let events = sink.take();
        assert_eq!(events.len(), 1);
        let KvEventData::Removed { block_hashes } = &events[0].data else {
            panic!("capacity pressure must emit one removal event")
        };
        assert_eq!(
            block_hashes,
            &vec![
                first_hashes[3],
                first_hashes[2],
                second_divergent_hash,
                first_hashes[1],
            ]
        );
        assert_eq!(manager.num_active_blocks(), 9);
        assert_eq!(manager.num_inactive_blocks(), 1);

        manager.finalize_lease_computed_prefix(
            pressure_owner,
            &mut pressure_sequence,
            &mut pressure_lease,
            0,
            36,
        );
        manager.finish_lease(pressure_owner, pressure_lease);
        assert_eq!(manager.num_active_blocks(), 0);
        assert_eq!(manager.num_inactive_blocks(), 10);
    }

    #[test]
    fn external_destination_reserves_exact_prefix_and_suffix_then_becomes_real_hit() {
        let mut manager =
            VllmKvManager::new_with_event_sink(3, 4, true, KvEventPublishers::default(), 0);
        let seed = Uuid::from_u128(1);
        let (mut seed_sequence, mut seed_lease) = request(seed, &[7], false);
        ready(manager.allocate_lease(seed, &mut seed_lease, 4, 0));
        manager.finalize_lease_computed_prefix(seed, &mut seed_sequence, &mut seed_lease, 0, 4);
        let prefix_copy = seed_lease.entries[0].copy.unwrap();
        manager.finish_lease(seed, seed_lease);

        let owner = Uuid::from_u128(2);
        let (sequence, mut lease) = request(owner, &[7, 8, 9], false);
        let reservation = ready(manager.reserve_external_prefix_lease(owner, &lease, 1, &[8]));
        assert_eq!(reservation.block_count, 2);
        assert_eq!(reservation.pool.fresh_len(), 1);
        assert_eq!(manager.num_active_blocks(), 2);

        manager.activate_destination_lease(owner, &sequence, &mut lease, reservation);
        assert_eq!(lease.resident_block_count(), 2);
        assert_eq!(lease.allocated_tokens(), 8);
        assert_eq!(lease.entries[0].copy, Some(prefix_copy));
        assert!(manager.pool.prefix_hit(8).is_some());
        assert!(manager.pool.prefix_hit(9).is_none());
        assert_eq!(
            manager
                .get_lease_prefill_cost(&sequence, &lease)
                .cached_tokens,
            8
        );
    }

    #[test]
    fn source_dependency_propagates_but_write_authorization_is_not_terminal() {
        let mut manager =
            VllmKvManager::new_with_event_sink(1, 4, true, KvEventPublishers::default(), 0);
        let source = Uuid::from_u128(1);
        let dependency = SourceReuseDependency::from_adapter_id(23);
        finish_source_with_dependency(&mut manager, source, &[7], &[0], dependency);

        let owner = Uuid::from_u128(2);
        let (mut sequence, mut lease) = request(owner, &[8], false);
        let dependencies = match manager.allocate_lease(owner, &mut lease, 4, 0) {
            NativeAllocation::Ready {
                value: 1,
                dependencies,
            } => dependencies,
            _ => panic!("dependency-bearing capacity must still be acquired"),
        };
        assert_eq!(dependencies, vec![dependency]);
        manager.authorize_lease_writes_after_dependencies(owner, &mut lease, &dependencies);
        manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 0, 4);
        assert!(manager.pool.prefix_hit(8).is_some());
        assert!(manager.is_source_reuse_dependency_pending(dependency));

        assert!(
            manager
                .snapshot_store_sources(owner, &lease, &[0])
                .is_none(),
            "new stores must wait until the old source reader is terminal"
        );
        assert!(manager.satisfy_source_reuse_dependency(dependency));
        assert!(
            manager
                .snapshot_store_sources(owner, &lease, &[0])
                .is_some()
        );
    }

    #[test]
    fn fenced_multiblock_allocation_keeps_owned_suffix_across_smaller_target() {
        let mut manager =
            VllmKvManager::new_with_event_sink(2, 4, true, KvEventPublishers::default(), 0);
        let source = Uuid::from_u128(1);
        let dependency = SourceReuseDependency::from_adapter_id(29);
        finish_source_with_dependency(&mut manager, source, &[7, 8], &[0, 1], dependency);

        let owner = Uuid::from_u128(2);
        let (mut sequence, mut lease) = request(owner, &[9, 10], false);
        let dependencies = match manager.allocate_lease(owner, &mut lease, 8, 0) {
            NativeAllocation::Ready {
                value: 2,
                dependencies,
            } => dependencies,
            _ => panic!("dependency-bearing suffix must still be acquired"),
        };
        assert_eq!(dependencies, vec![dependency]);
        manager.authorize_lease_writes_after_dependencies(owner, &mut lease, &dependencies);
        assert!(matches!(
            manager.allocate_lease(owner, &mut lease, 4, 0),
            NativeAllocation::Ready {
                value: 0,
                ref dependencies,
            } if dependencies.is_empty()
        ));
        assert_eq!(lease.allocated_tokens(), 8);
        assert_eq!(lease.resident_block_count(), 2);

        manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 0, 4);
        manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 4, 8);
        assert!(manager.satisfy_source_reuse_dependency(dependency));
        assert!(manager.pool.prefix_hit(9).is_some());
        assert!(manager.pool.prefix_hit(10).is_some());
    }

    #[test]
    fn event_enabled_finalization_preserves_store_payload() {
        let sink = Arc::new(CapturingNativeSink::default());
        let publishers = KvEventPublishers::new(Some(sink.clone()));
        let mut manager = VllmKvManager::new_with_event_sink(2, 4, true, publishers, 3);
        let owner = Uuid::from_u128(1);
        let token_ids = [vec![4, 5, 6, 7]];
        let (mut sequence, mut lease) = request(owner, &[6, 7], true);
        ready(manager.allocate_lease(owner, &mut lease, 8, 0));
        manager.finalize_lease_computed_prefix(owner, &mut sequence, &mut lease, 4, 8);

        let mut events = sink.take();
        assert_eq!(events.len(), 1);
        let event = events.pop().unwrap();
        assert_eq!(event.event_id, 0);
        assert_eq!(event.dp_rank, 3);
        let KvEventData::Stored(stored) = event.data else {
            panic!("expected Stored event")
        };
        assert_eq!(stored.parent_hash, Some(6));
        assert_eq!(stored.blocks.len(), 1);
        assert_eq!(stored.blocks[0].block_hash, 7);
        assert_eq!(stored.blocks[0].tokens_hash, 107);
        assert_eq!(stored.blocks[0].token_ids, Some(token_ids[0].clone()));
    }
}