feat: adding request and scheduled sequence (#6933)

Signed-off-by: Ryan Olson <rolson@nvidia.com>

Cargo.lock

@@ -3985,14 +3985,18 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "async-stream",
+ "bincode 2.0.1",
  "bytes",
  "derive_builder",
  "dynamo-tokens",
  "futures",
+ "indexmap 2.13.0",
  "lru 0.16.3",
  "parking_lot",
  "prometheus",
  "proptest",
+ "rand 0.9.2",
+ "rand_chacha 0.9.0",
  "rmp-serde",
  "rstest 0.26.1",
  "serde",
@@ -4001,6 +4005,7 @@ dependencies = [
  "tokio",
  "tokio-stream",
  "tracing",
+ "tracing-subscriber",
  "xxhash-rust",
 ]
 

Cargo.toml

@@ -93,6 +93,7 @@ modelexpress-client = { version = "0.2.0" }
 modelexpress-common = { version = "0.2.0" }
 
 humantime = { version = "2.2.0" }
+indexmap = { version = "2" }
 libc = { version = "0.2" }
 oneshot = { version = "0.1.13", features = ["std", "async"] }
 ordered-float = "4"

lib/kvbm-logical/Cargo.toml

@@ -15,8 +15,10 @@ dynamo-tokens = { workspace = true }
 anyhow = { workspace = true }
 async-stream = "0.3"
 bytes = "1.10"
+bincode = { version = "2.0.1", features = ["serde", "derive"] }
 derive_builder = { workspace = true }
 futures = { workspace = true }
+indexmap = { workspace = true }
 lru = "0.16"
 parking_lot = { workspace = true }
 prometheus = { workspace = true }
@@ -34,6 +36,9 @@ testing = []
 
 [dev-dependencies]
 proptest = "1.5.0"
+rand = { workspace = true }
+rand_chacha = "0.9"
 rstest = "0.26"
 serde_json = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
+tracing-subscriber = { workspace = true }

lib/kvbm-logical/src/integrations/mod.rs

+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Composition layer wiring together [`BlockSequence`](crate::BlockSequence),
+//! [`LogicalBlockAssignments`](crate::LogicalBlockAssignments), and
+//! [`BlockManager`](crate::BlockManager) into higher-level request lifecycle types.
+
+mod request;
+mod scheduled;
+
+pub use request::RequestSequence;
+pub use scheduled::{
+    ApplyError, DecodeOutcome, NoopDelegate, SchedulableSequence, SchedulableSequenceBuilder,
+    ScheduleError, SequenceDelegate, SequenceEvent, SequenceState,
+};

lib/kvbm-logical/src/integrations/request.rs

+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! # Request Sequence
+//!
+//! Low-level request primitive with direct RAII block lifecycle management.
+//!
+//! [`RequestSequence`] composes [`BlockSequence`](crate::BlockSequence),
+//! [`LogicalBlockAssignments`](crate::LogicalBlockAssignments), and
+//! [`BlockManager`](crate::BlockManager) into a single type that exposes
+//! individual block lifecycle operations without opinionation about
+//! scheduling policy.
+//!
+//! For a structured two-phase schedule/apply layer built on top of this,
+//! see [`SchedulableSequence`](crate::SchedulableSequence).
+//!
+//! ## When to use
+//!
+//! Use `RequestSequence` directly when you need full control over block
+//! allocation, staging, and registration timing. Use `SchedulableSequence`
+//! when you want state-machine enforcement of the prefill/decode protocol.
+//!
+//! ## Block lifecycle
+//!
+//! Blocks flow through three phases:
+//!
+//! 1. **Unassigned** -- freshly allocated `MutableBlock`s waiting to be paired
+//!    with token data
+//! 2. **Staged** -- `CompleteBlock`s paired with token data but not yet
+//!    committed to the registry
+//! 3. **Assigned** -- `ImmutableBlock`s registered in the block manager,
+//!    visible for prefix matching
+//!
+//! ## Basic usage
+//!
+//! ```ignore
+//! use kvbm_logical::{RequestSequence, BlockManager};
+//!
+//! // 1. Construct with tokens only (no manager interaction)
+//! let tokens: Vec<u32> = (0..8).collect();
+//! let mut seq = RequestSequence::<MyMeta>::new(tokens, 10, 4);
+//! // total_tokens=8, num_blocks=2, nothing allocated yet
+//!
+//! // 2. Prefix match against the cache
+//! let matched_count = seq.match_and_add_prefix(&manager).unwrap();
+//!
+//! // 3. Allocate blocks for the rest
+//! let remaining = seq.num_blocks() - matched_count;
+//! seq.allocate_blocks(remaining, &manager);
+//!
+//! // 4. Stage and register
+//! seq.complete_and_register_pending(&manager);
+//! // Now: assigned_blocks() == num_blocks()
+//! ```
+//!
+//! ## Generation loop
+//!
+//! After initial setup, generate tokens one at a time. Each `append_token`
+//! call returns `Some(block_index)` when a block boundary is crossed,
+//! signaling that `complete_and_register_pending` should be called and a
+//! new generation block allocated.
+//!
+//! ```ignore
+//! while !seq.is_complete() {
+//!     let token = model.forward(&seq);
+//!     let crossed = seq.append_token(token);
+//!     if crossed.is_some() {
+//!         seq.complete_and_register_pending(&manager);
+//!         seq.allocate_blocks(1, &manager);
+//!     }
+//! }
+//! ```
+//!
+//! ## Preemption and reacquire
+//!
+//! Release all blocks (RAII returns them to pools), then re-acquire later.
+//! Prefix-matched blocks may come from cache, saving re-computation.
+//!
+//! ```ignore
+//! // Preempt
+//! seq.release();
+//! assert_eq!(seq.assigned_blocks(), 0);
+//!
+//! // Later: reacquire
+//! let success = seq.reacquire(&manager);
+//! // Prefix cache hits are reflected in prefix_matched_blocks()
+//! ```
+//!
+//! ## Key accessors
+//!
+//! | Method                    | Description                                    |
+//! |---------------------------|------------------------------------------------|
+//! | `total_tokens()`          | Input + generated token count                  |
+//! | `num_input_tokens()`      | Original input token count                     |
+//! | `generated_tokens()`      | Tokens appended via `append_token`             |
+//! | `num_blocks()`            | Complete token blocks in the sequence          |
+//! | `assigned_blocks()`       | Registered/cache-matched blocks                |
+//! | `staged_blocks()`         | Completed but not yet registered               |
+//! | `unassigned_blocks()`     | Allocated but not yet paired with token data   |
+//! | `prefix_matched_blocks()` | Blocks matched from cache                      |
+//! | `is_complete()`           | `generated_tokens >= max_output_tokens`         |
+//! | `new_tokens_for_prefill()`| Tokens not covered by cache hits               |
+
+use crate::KvbmSequenceHashProvider;
+use crate::blocks::{BlockMetadata, ImmutableBlock};
+use crate::manager::BlockManager;
+use crate::sequence::{BlockSequence, LogicalBlockAssignmentError, LogicalBlockAssignments};
+
+use dynamo_tokens::Token;
+
+/// Manages a request's block lifecycle through direct RAII integration with
+/// [`BlockManager`], bypassing the `MoveBlock` signal protocol.
+///
+/// Composes [`BlockSequence`] (token data and hashing) and
+/// [`LogicalBlockAssignments`] (RAII block guards) into a single unit that
+/// handles:
+///
+/// - **Construction**: prefix matching, allocation, and registration of input blocks
+/// - **Generation**: token-by-token extension with automatic block promotion
+/// - **Preemption**: release all blocks, re-acquire later with potential cache hits
+///
+/// # Block lifecycle
+///
+/// The generation block (always at most one) lives as the single unassigned
+/// entry in [`LogicalBlockAssignments`]. When a block boundary is crossed:
+///
+/// 1. `stage()` — the generation `MutableBlock` becomes a `CompleteBlock`
+/// 2. `register()` — the `CompleteBlock` becomes an `ImmutableBlock`
+/// 3. `allocate_blocks(1)` → `extend_blocks()` — a new generation block
+pub struct RequestSequence<T: BlockMetadata> {
+    sequence: BlockSequence,
+    assignments: LogicalBlockAssignments<T>,
+    generated_tokens: usize,
+    max_output_tokens: usize,
+    num_input_tokens: usize,
+    prefix_matched_blocks: usize,
+}
+
+impl<T: BlockMetadata> RequestSequence<T> {
+    // =====================================================================
+    // Minimal constructor (no manager interaction)
+    // =====================================================================
+
+    /// Creates a `RequestSequence` with token data only. No blocks are
+    /// allocated and no manager interaction occurs.
+    ///
+    /// The caller must use [`match_and_add_prefix`], [`allocate_blocks`],
+    /// and [`complete_and_register_pending`] to search, allocate, and
+    /// register blocks.
+    ///
+    /// [`match_and_add_prefix`]: Self::match_and_add_prefix
+    /// [`allocate_blocks`]: Self::allocate_blocks
+    /// [`complete_and_register_pending`]: Self::complete_and_register_pending
+    pub fn new(tokens: Vec<Token>, max_output_tokens: usize, block_size: u32) -> Self {
+        let num_input_tokens = tokens.len();
+        let sequence = BlockSequence::new(tokens, block_size, None);
+        let assignments = LogicalBlockAssignments::new();
+
+        Self {
+            sequence,
+            assignments,
+            generated_tokens: 0,
+            max_output_tokens,
+            num_input_tokens,
+            prefix_matched_blocks: 0,
+        }
+    }
+
+    // =====================================================================
+    // Individual block operations
+    // =====================================================================
+
+    /// Search for prefix cache hits and add matched blocks in one step.
+    ///
+    /// This is the standard entry point for prefix matching on a fresh
+    /// sequence. Combines [`match_prefix`](Self::match_prefix) and
+    /// [`add_matched_blocks`](Self::add_matched_blocks).
+    ///
+    /// # Panics
+    ///
+    /// Panics if the sequence already has assigned blocks (i.e. this is
+    /// not a fresh sequence).
+    pub fn match_and_add_prefix(
+        &mut self,
+        manager: &BlockManager<T>,
+    ) -> Result<usize, LogicalBlockAssignmentError<T>> {
+        assert!(
+            self.assignments.is_empty(),
+            "match_and_add_prefix called on sequence with existing assignments"
+        );
+        let matched = self.match_prefix(manager);
+        if matched.is_empty() {
+            return Ok(0);
+        }
+        self.add_matched_blocks(matched)
+    }
+
+    /// Search for prefix cache hits against the manager's pools.
+    ///
+    /// Returns matched [`ImmutableBlock`]s in sequence order. Pass the result
+    /// to [`add_matched_blocks`](Self::add_matched_blocks).
+    fn match_prefix(&self, manager: &BlockManager<T>) -> Vec<ImmutableBlock<T>> {
+        let hashes = self.sequence.all_sequence_hashes();
+        manager.match_blocks(&hashes)
+    }
+
+    /// Add prefix-matched immutable blocks as assigned.
+    ///
+    /// Accumulates the internal `prefix_matched_blocks` counter so this
+    /// method can be called more than once (e.g. partial prefix matches
+    /// applied in separate batches).
+    ///
+    /// Returns the number of blocks added.
+    fn add_matched_blocks(
+        &mut self,
+        blocks: Vec<ImmutableBlock<T>>,
+    ) -> Result<usize, LogicalBlockAssignmentError<T>> {
+        let count = blocks.len();
+        let start = self.assignments.assigned_count();
+        let end = start + count;
+        let sequence_blocks = self.sequence.blocks();
+
+        assert!(
+            end <= sequence_blocks.len(),
+            "matched blocks exceed completed sequence blocks"
+        );
+
+        for (i, (block, seq_block)) in blocks.iter().zip(&sequence_blocks[start..end]).enumerate() {
+            let expected = seq_block.kvbm_sequence_hash();
+            let actual = block.sequence_hash();
+            if expected != actual {
+                return Err(LogicalBlockAssignmentError::SequenceHashMismatch {
+                    position: start + i,
+                    expected,
+                    actual,
+                    blocks,
+                });
+            }
+        }
+
+        self.assignments.extend_assigned(blocks)?;
+        self.prefix_matched_blocks += count;
+        Ok(count)
+    }
+
+    /// Allocate mutable blocks from the manager and store as unassigned.
+    ///
+    /// Returns `false` if allocation fails (insufficient blocks).
+    pub fn allocate_blocks(&mut self, count: usize, manager: &BlockManager<T>) -> bool {
+        if count == 0 {
+            return true;
+        }
+        let Some(new_blocks) = manager.allocate_blocks(count) else {
+            return false;
+        };
+        self.assignments.extend_blocks(new_blocks).is_ok()
+    }
+
+    /// Stage all unassigned blocks that have corresponding completed token
+    /// data in the sequence, without registering them.
+    ///
+    /// Moves unassigned `MutableBlock`s → `CompleteBlock`s (staged). Staged
+    /// blocks are not yet visible to `match_blocks()` — call
+    /// [`register_staged`](Self::register_staged) after the GPU has computed
+    /// their KV data.
+    pub fn stage_pending(&mut self) {
+        let start = self.assignments.assigned_count() + self.assignments.staged_count();
+        let completed = self.sequence.blocks().len();
+        if start < completed {
+            let blocks_slice = &self.sequence.blocks()[start..completed];
+            self.assignments
+                .stage(blocks_slice)
+                .expect("staging should not fail (block sizes and counts match)");
+        }
+    }
+
+    /// Register all staged blocks with the block manager.
+    ///
+    /// Moves staged `CompleteBlock`s → `ImmutableBlock`s (assigned). After
+    /// registration, blocks become visible to `match_blocks()` for prefix
+    /// reuse by future requests.
+    ///
+    /// Returns the number of blocks registered.
+    pub fn register_staged(&mut self, manager: &BlockManager<T>) -> usize {
+        self.assignments.register(manager)
+    }
+
+    /// Stage and register all unassigned blocks that have corresponding
+    /// completed token data in the sequence.
+    ///
+    /// Computes the offset from `assigned_count + staged_count` and stages
+    /// all completed token blocks beyond that offset. This converts
+    /// unassigned `MutableBlock`s → `CompleteBlock`s → `ImmutableBlock`s.
+    ///
+    /// A generation block (not yet filled) remains unassigned because it has
+    /// no corresponding entry in `sequence.blocks()`.
+    pub fn complete_and_register_pending(&mut self, manager: &BlockManager<T>) {
+        self.stage_pending();
+        self.register_staged(manager);
+    }
+
+    // =====================================================================
+    // Token-only append (no block lifecycle)
+    // =====================================================================
+
+    /// Append a generated token to the sequence. Increments `generated_tokens`.
+    ///
+    /// Returns `Some(block_index)` if a block boundary was crossed (the block
+    /// at that index is now complete), `None` otherwise.
+    ///
+    /// Does **not** stage, register, or allocate — the caller handles block
+    /// lifecycle via [`complete_and_register_pending`] and [`allocate_blocks`].
+    ///
+    /// [`complete_and_register_pending`]: Self::complete_and_register_pending
+    /// [`allocate_blocks`]: Self::allocate_blocks
+    ///
+    /// # Panics
+    ///
+    /// Panics if `generated_tokens >= max_output_tokens`.
+    pub fn append_token(&mut self, token: Token) -> Option<usize> {
+        assert!(
+            self.generated_tokens < self.max_output_tokens,
+            "Cannot generate more tokens: reached max_output_tokens limit"
+        );
+
+        let completed_block = self
+            .sequence
+            .append_token(token)
+            .expect("Token append failed");
+
+        self.generated_tokens += 1;
+        completed_block
+    }
+
+    /// Whether `generated_tokens >= max_output_tokens`.
+    pub fn is_complete(&self) -> bool {
+        self.generated_tokens >= self.max_output_tokens
+    }
+
+    // =====================================================================
+    // Release / reacquire
+    // =====================================================================
+
+    /// Releases all block assignments (RAII returns them to pools).
+    pub fn release(&mut self) {
+        self.assignments.clear();
+    }
+
+    /// Re-acquires blocks from the manager after a release/preemption.
+    ///
+    /// Uses the sequence's current token state (input + generated) to
+    /// match prefix blocks from pools and allocate the remainder.
+    ///
+    /// Returns `true` if all blocks were successfully acquired.
+    pub fn reacquire(&mut self, manager: &BlockManager<T>) -> bool {
+        assert!(
+            self.assignments.is_empty(),
+            "reacquire called with existing assignments"
+        );
+
+        let completed_blocks = self.sequence.blocks().len();
+
+        // Step 1: Prefix match
+        let hashes = self.sequence.all_sequence_hashes();
+        let matched = manager.match_blocks(&hashes);
+        let matched_count = matched.len();
+
+        if !matched.is_empty() && self.assignments.extend_assigned(matched).is_err() {
+            self.assignments.clear();
+            return false;
+        }
+
+        // Step 2: Allocate remaining complete blocks (gen block allocated by schedule_decode)
+        let remaining_complete = completed_blocks - matched_count;
+        let total = remaining_complete;
+
+        if total > 0 {
+            let Some(new_blocks) = manager.allocate_blocks(total) else {
+                self.assignments.clear();
+                return false;
+            };
+            if self.assignments.extend_blocks(new_blocks).is_err() {
+                self.assignments.clear();
+                return false;
+            }
+        }
+
+        // Step 3: Stage and register remaining complete blocks
+        if remaining_complete > 0 {
+            let blocks_slice = &self.sequence.blocks()[matched_count..completed_blocks];
+            if self.assignments.stage(blocks_slice).is_err() {
+                self.assignments.clear();
+                return false;
+            }
+            self.assignments.register(manager);
+        }
+
+        self.prefix_matched_blocks = matched_count;
+        true
+    }
+
+    // =====================================================================
+    // Accessors
+    // =====================================================================
+
+    /// Number of tokens generated so far.
+    pub fn generated_tokens(&self) -> usize {
+        self.generated_tokens
+    }
+
+    /// Maximum number of output tokens.
+    pub fn max_output_tokens(&self) -> usize {
+        self.max_output_tokens
+    }
+
+    /// Number of input tokens the sequence was created with.
+    pub fn num_input_tokens(&self) -> usize {
+        self.num_input_tokens
+    }
+
+    /// Total token count (input + generated).
+    pub fn total_tokens(&self) -> usize {
+        self.sequence.total_tokens()
+    }
+
+    /// Number of tokens remaining to be generated.
+    pub fn remaining_tokens(&self) -> usize {
+        self.max_output_tokens.saturating_sub(self.generated_tokens)
+    }
+
+    /// Number of completed blocks in the token sequence.
+    pub fn num_blocks(&self) -> usize {
+        self.sequence.blocks().len()
+    }
+
+    /// Number of complete token blocks in the sequence.
+    /// Alias for `num_blocks()` with a more descriptive name for scheduling.
+    pub fn complete_sequence_blocks(&self) -> usize {
+        self.sequence.blocks().len()
+    }
+
+    /// Number of blocks currently assigned (registered or cache-matched).
+    pub fn assigned_blocks(&self) -> usize {
+        self.assignments.assigned_count()
+    }
+
+    /// Number of blocks currently staged (completed but not registered).
+    pub fn staged_blocks(&self) -> usize {
+        self.assignments.staged_count()
+    }
+
+    /// Number of unassigned blocks (the generation block, if any).
+    pub fn unassigned_blocks(&self) -> usize {
+        self.assignments.unassigned_count()
+    }
+
+    /// Number of blocks that were prefix-matched during construction or reacquire.
+    pub fn prefix_matched_blocks(&self) -> usize {
+        self.prefix_matched_blocks
+    }
+
+    /// Number of new (non-cached) tokens for prefill cost calculation.
+    pub fn new_tokens_for_prefill(&self) -> usize {
+        let cached_tokens = self.prefix_matched_blocks * self.sequence.block_size();
+        self.total_tokens().saturating_sub(cached_tokens)
+    }
+
+    /// Reference to the underlying `BlockSequence`.
+    pub fn sequence(&self) -> &BlockSequence {
+        &self.sequence
+    }
+
+    /// Reference to the underlying `LogicalBlockAssignments`.
+    pub fn assignments(&self) -> &LogicalBlockAssignments<T> {
+        &self.assignments
+    }
+
+    /// Block size used by this sequence.
+    pub fn block_size(&self) -> usize {
+        self.sequence.block_size()
+    }
+
+    /// All block IDs in order: assigned ++ staged ++ unassigned.
+    ///
+    /// Block IDs identity-map to page indices in the GPU page pool.
+    pub fn page_indices(&self) -> Vec<u32> {
+        self.assignments
+            .all_block_ids()
+            .map(|&id| id as u32)
+            .collect()
+    }
+
+    /// Drop excess unassigned blocks beyond `keep` count.
+    /// Returns the number of blocks dropped (RAII returns them to reset pool).
+    pub fn drop_excess_unassigned(&mut self, keep: usize) -> usize {
+        let mut dropped = 0;
+        while self.assignments.unassigned_count() > keep {
+            if self.assignments.pop_last_unassigned().is_some() {
+                dropped += 1;
+            } else {
+                break;
+            }
+        }
+        dropped
+    }
+
+    // =====================================================================
+    // Crate-internal mutation accessors
+    // =====================================================================
+
+    /// Mutable access to assignments for higher-level wrappers.
+    pub(crate) fn assignments_mut(&mut self) -> &mut LogicalBlockAssignments<T> {
+        &mut self.assignments
+    }
+
+    /// Mutable access to the underlying `BlockSequence`.
+    #[allow(dead_code)]
+    pub(crate) fn sequence_mut(&mut self) -> &mut BlockSequence {
+        &mut self.sequence
+    }
+
+    /// Bulk-increment the generated token counter (for speculative decode).
+    #[allow(dead_code)]
+    pub(crate) fn add_generated_tokens(&mut self, count: usize) {
+        self.generated_tokens += count;
+    }
+}
+
+impl<T: BlockMetadata> std::fmt::Debug for RequestSequence<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("RequestSequence")
+            .field("total_tokens", &self.sequence.total_tokens())
+            .field("num_blocks", &self.sequence.blocks().len())
+            .field("generated_tokens", &self.generated_tokens)
+            .field("max_output_tokens", &self.max_output_tokens)
+            .field("num_input_tokens", &self.num_input_tokens)
+            .field("prefix_matched_blocks", &self.prefix_matched_blocks)
+            .field("assigned", &self.assignments.assigned_count())
+            .field("staged", &self.assignments.staged_count())
+            .field("unassigned", &self.assignments.unassigned_count())
+            .finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::testing::{TestMeta, create_test_manager};
+
+    const BLOCK_SIZE: u32 = 4;
+
+    fn make_tokens(n: usize) -> Vec<Token> {
+        (0..n as u32).collect()
+    }
+
+    /// Builds a prefilled `RequestSequence` with blocks allocated, staged,
+    /// and registered — the same end state as the removed `with_manager()`.
+    fn build_prefilled(
+        tokens: Vec<Token>,
+        max_output_tokens: usize,
+        block_size: u32,
+        manager: &BlockManager<TestMeta>,
+    ) -> Option<RequestSequence<TestMeta>> {
+        let mut seq = RequestSequence::new(tokens, max_output_tokens, block_size);
+        let completed_blocks = seq.num_blocks();
+        let matched_count = seq.match_and_add_prefix(manager).ok()?;
+        let remaining_complete = completed_blocks - matched_count;
+        let needs_generation = max_output_tokens > 0;
+        let total_to_allocate = remaining_complete + usize::from(needs_generation);
+        if !seq.allocate_blocks(total_to_allocate, manager) {
+            return None;
+        }
+        seq.complete_and_register_pending(manager);
+        Some(seq)
+    }
+
+    // =========================================================================
+    // Construction
+    // =========================================================================
+
+    #[test]
+    fn test_new_minimal_constructor() {
+        let tokens = make_tokens(8);
+        let seq = RequestSequence::<TestMeta>::new(tokens, 10, BLOCK_SIZE);
+
+        assert_eq!(seq.num_input_tokens(), 8);
+        assert_eq!(seq.total_tokens(), 8);
+        assert_eq!(seq.num_blocks(), 2);
+        assert_eq!(seq.generated_tokens(), 0);
+        assert_eq!(seq.max_output_tokens(), 10);
+        assert_eq!(seq.block_size(), BLOCK_SIZE as usize);
+
+        assert_eq!(seq.assigned_blocks(), 0);
+        assert_eq!(seq.staged_blocks(), 0);
+        assert_eq!(seq.unassigned_blocks(), 0);
+        assert_eq!(seq.prefix_matched_blocks(), 0);
+    }
+
+    #[test]
+    fn test_build_prefilled_basic() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let tokens = make_tokens(8); // 2 complete blocks
+        let seq = build_prefilled(tokens, 10, BLOCK_SIZE, &manager).unwrap();
+
+        assert_eq!(seq.num_input_tokens(), 8);
+        assert_eq!(seq.total_tokens(), 8);
+        assert_eq!(seq.num_blocks(), 2);
+        assert_eq!(seq.generated_tokens(), 0);
+        assert_eq!(seq.max_output_tokens(), 10);
+
+        assert_eq!(seq.assigned_blocks(), 2);
+        assert_eq!(seq.staged_blocks(), 0);
+        assert_eq!(seq.unassigned_blocks(), 1);
+    }
+
+    #[test]
+    fn test_build_prefilled_partial_tokens() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let seq = build_prefilled(make_tokens(6), 10, BLOCK_SIZE, &manager).unwrap();
+
+        assert_eq!(seq.num_blocks(), 1);
+        assert_eq!(seq.assigned_blocks(), 1);
+        assert_eq!(seq.unassigned_blocks(), 1);
+    }
+
+    #[test]
+    fn test_build_prefilled_empty_tokens() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let seq = build_prefilled(vec![], 10, BLOCK_SIZE, &manager).unwrap();
+
+        assert_eq!(seq.num_blocks(), 0);
+        assert_eq!(seq.assigned_blocks(), 0);
+        assert_eq!(seq.unassigned_blocks(), 1);
+    }
+
+    #[test]
+    fn test_build_prefilled_zero_max() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let seq = build_prefilled(make_tokens(8), 0, BLOCK_SIZE, &manager).unwrap();
+
+        assert_eq!(seq.assigned_blocks(), 2);
+        assert_eq!(seq.unassigned_blocks(), 0);
+    }
+
+    #[test]
+    fn test_build_prefilled_allocation_failure() {
+        let manager = create_test_manager::<TestMeta>(2);
+        let result = build_prefilled(make_tokens(12), 10, BLOCK_SIZE, &manager);
+        assert!(result.is_none());
+        assert_eq!(manager.available_blocks(), 2);
+    }
+
+    // =========================================================================
+    // Prefix matching
+    // =========================================================================
+
+    #[test]
+    fn test_prefix_cache_hit() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let tokens = make_tokens(8);
+
+        // Populate the manager with blocks for these tokens
+        let seq_for_populate = BlockSequence::new(tokens.clone(), BLOCK_SIZE, None);
+        let mutables = manager.allocate_blocks(2).unwrap();
+        let registered: Vec<_> = mutables
+            .into_iter()
+            .zip(seq_for_populate.blocks().iter())
+            .map(|(m, tb)| manager.register_block(m.complete(tb).unwrap()))
+            .collect();
+        drop(registered);
+
+        let seq = build_prefilled(tokens, 10, BLOCK_SIZE, &manager).unwrap();
+        assert_eq!(seq.prefix_matched_blocks(), 2);
+    }
+
+    #[test]
+    fn test_partial_prefix_cache_hit() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let tokens = make_tokens(12);
+
+        let seq_for_populate = BlockSequence::new(tokens[..4].to_vec(), BLOCK_SIZE, None);
+        let mutables = manager.allocate_blocks(1).unwrap();
+        let registered: Vec<_> = mutables
+            .into_iter()
+            .zip(seq_for_populate.blocks().iter())
+            .map(|(m, tb)| manager.register_block(m.complete(tb).unwrap()))
+            .collect();
+        drop(registered);
+
+        let seq = build_prefilled(tokens, 10, BLOCK_SIZE, &manager).unwrap();
+        assert_eq!(seq.prefix_matched_blocks(), 1);
+        assert_eq!(seq.assigned_blocks(), 3);
+        assert_eq!(seq.unassigned_blocks(), 1);
+    }
+
+    #[test]
+    #[should_panic(expected = "matched blocks exceed completed sequence blocks")]
+    fn test_add_matched_blocks_panics_when_matched_exceeds_completed() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = RequestSequence::<TestMeta>::new(make_tokens(4), 10, BLOCK_SIZE);
+
+        let source = BlockSequence::new(make_tokens(8), BLOCK_SIZE, None);
+        let mutables = manager.allocate_blocks(2).unwrap();
+        let matched: Vec<_> = mutables
+            .into_iter()
+            .zip(source.blocks().iter())
+            .map(|(m, tb)| manager.register_block(m.complete(tb).unwrap()))
+            .collect();
+
+        let _ = seq.add_matched_blocks(matched);
+    }
+
+    #[test]
+    fn test_add_matched_blocks_returns_error_on_hash_mismatch() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = RequestSequence::<TestMeta>::new(make_tokens(4), 10, BLOCK_SIZE);
+
+        let source = BlockSequence::new(vec![100, 101, 102, 103], BLOCK_SIZE, None);
+        let mutable = manager
+            .allocate_blocks(1)
+            .unwrap()
+            .into_iter()
+            .next()
+            .unwrap();
+        let mismatched = manager.register_block(mutable.complete(&source.blocks()[0]).unwrap());
+
+        let result = seq.add_matched_blocks(vec![mismatched]);
+        assert!(result.is_err());
+        match result.unwrap_err() {
+            LogicalBlockAssignmentError::SequenceHashMismatch {
+                position, blocks, ..
+            } => {
+                assert_eq!(position, 0);
+                assert_eq!(blocks.len(), 1);
+            }
+            other => panic!("expected SequenceHashMismatch, got: {other:?}"),
+        }
+    }
+
+    #[test]
+    fn test_new_tokens_for_prefill() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let tokens = make_tokens(12);
+
+        let seq_for_populate = BlockSequence::new(tokens[..4].to_vec(), BLOCK_SIZE, None);
+        let mutables = manager.allocate_blocks(1).unwrap();
+        let registered: Vec<_> = mutables
+            .into_iter()
+            .zip(seq_for_populate.blocks().iter())
+            .map(|(m, tb)| manager.register_block(m.complete(tb).unwrap()))
+            .collect();
+        drop(registered);
+
+        let seq = build_prefilled(tokens, 10, BLOCK_SIZE, &manager).unwrap();
+        assert_eq!(seq.prefix_matched_blocks(), 1);
+        assert_eq!(seq.new_tokens_for_prefill(), 8);
+    }
+
+    // =========================================================================
+    // Token append (individual ops)
+    // =========================================================================
+
+    #[test]
+    fn test_append_token_no_boundary() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = build_prefilled(make_tokens(5), 10, BLOCK_SIZE, &manager).unwrap();
+
+        // 5 tokens = 1 complete + 1 partial (1 token). Append 1 → 2 partial tokens, no boundary.
+        assert!(seq.append_token(100).is_none());
+        assert_eq!(seq.generated_tokens(), 1);
+        assert_eq!(seq.total_tokens(), 6);
+    }
+
+    #[test]
+    fn test_append_token_crosses_boundary() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = build_prefilled(make_tokens(7), 10, BLOCK_SIZE, &manager).unwrap();
+
+        // 7 tokens = 1 complete + 3 partial. 1 more completes block 1.
+        let block_idx = seq.append_token(100);
+        assert!(block_idx.is_some());
+        assert_eq!(seq.num_blocks(), 2);
+    }
+
+    #[test]
+    #[should_panic(expected = "Cannot generate more tokens")]
+    fn test_append_token_panics_after_max() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = build_prefilled(make_tokens(4), 1, BLOCK_SIZE, &manager).unwrap();
+
+        seq.append_token(100); // generated_tokens == 1 == max
+        seq.append_token(101); // panics
+    }
+
+    #[test]
+    fn test_is_complete_transitions() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = build_prefilled(make_tokens(4), 3, BLOCK_SIZE, &manager).unwrap();
+
+        assert!(!seq.is_complete());
+        seq.append_token(100);
+        assert!(!seq.is_complete());
+        seq.append_token(101);
+        assert!(!seq.is_complete());
+        seq.append_token(102);
+        assert!(seq.is_complete());
+    }
+
+    #[test]
+    fn test_is_complete_zero_max() {
+        let seq = RequestSequence::<TestMeta>::new(make_tokens(4), 0, BLOCK_SIZE);
+        assert!(seq.is_complete());
+    }
+
+    // =========================================================================
+    // Modular decode loop
+    // =========================================================================
+
+    #[test]
+    fn test_modular_decode_loop() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = build_prefilled(make_tokens(4), 8, BLOCK_SIZE, &manager).unwrap();
+
+        // Generate 4 tokens → complete a block
+        for i in 0..3 {
+            assert!(seq.append_token(100 + i).is_none());
+        }
+        let block_idx = seq.append_token(103);
+        assert!(block_idx.is_some());
+        assert!(!seq.is_complete());
+
+        seq.complete_and_register_pending(&manager);
+        assert_eq!(seq.assigned_blocks(), 2);
+        assert_eq!(seq.unassigned_blocks(), 0);
+        assert!(seq.allocate_blocks(1, &manager));
+        assert_eq!(seq.unassigned_blocks(), 1);
+
+        // Generate 4 more → complete + max
+        for i in 0..3 {
+            assert!(seq.append_token(200 + i).is_none());
+        }
+        assert!(seq.append_token(203).is_some());
+        assert!(seq.is_complete());
+
+        seq.complete_and_register_pending(&manager);
+        assert_eq!(seq.assigned_blocks(), 3);
+        seq.release();
+        assert_eq!(seq.assigned_blocks(), 0);
+    }
+
+    #[test]
+    fn test_modular_prefill_with_cache() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let tokens = make_tokens(8);
+
+        // Populate cache
+        let seq_for_populate = BlockSequence::new(tokens[..4].to_vec(), BLOCK_SIZE, None);
+        let mutables = manager.allocate_blocks(1).unwrap();
+        let registered: Vec<_> = mutables
+            .into_iter()
+            .zip(seq_for_populate.blocks().iter())
+            .map(|(m, tb)| manager.register_block(m.complete(tb).unwrap()))
+            .collect();
+        drop(registered);
+
+        let mut seq = RequestSequence::<TestMeta>::new(tokens, 10, BLOCK_SIZE);
+        let matched_count = seq.match_and_add_prefix(&manager).unwrap();
+        assert_eq!(matched_count, 1);
+        assert_eq!(seq.prefix_matched_blocks(), 1);
+
+        let remaining = seq.num_blocks() - matched_count;
+        assert!(seq.allocate_blocks(remaining + 1, &manager));
+        seq.complete_and_register_pending(&manager);
+
+        assert_eq!(seq.assigned_blocks(), 2);
+        assert_eq!(seq.unassigned_blocks(), 1);
+    }
+
+    // =========================================================================
+    // Allocate blocks
+    // =========================================================================
+
+    #[test]
+    fn test_allocate_blocks_failure() {
+        let manager = create_test_manager::<TestMeta>(2);
+        let mut seq = RequestSequence::<TestMeta>::new(make_tokens(4), 10, BLOCK_SIZE);
+
+        assert!(!seq.allocate_blocks(3, &manager));
+        assert!(seq.allocate_blocks(2, &manager));
+        assert_eq!(seq.unassigned_blocks(), 2);
+    }
+
+    #[test]
+    fn test_allocate_blocks_zero() {
+        let manager = create_test_manager::<TestMeta>(2);
+        let mut seq = RequestSequence::<TestMeta>::new(make_tokens(4), 10, BLOCK_SIZE);
+
+        assert!(seq.allocate_blocks(0, &manager));
+        assert_eq!(seq.unassigned_blocks(), 0);
+    }
+
+    // =========================================================================
+    // Release / reacquire
+    // =========================================================================
+
+    #[test]
+    fn test_release() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = build_prefilled(make_tokens(8), 10, BLOCK_SIZE, &manager).unwrap();
+
+        assert_eq!(seq.assigned_blocks(), 2);
+        assert_eq!(seq.unassigned_blocks(), 1);
+        let available_before = manager.available_blocks();
+
+        seq.release();
+
+        assert_eq!(seq.assigned_blocks(), 0);
+        assert_eq!(seq.staged_blocks(), 0);
+        assert_eq!(seq.unassigned_blocks(), 0);
+        assert!(manager.available_blocks() > available_before);
+    }
+
+    #[test]
+    fn test_release_idempotent() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = build_prefilled(make_tokens(8), 10, BLOCK_SIZE, &manager).unwrap();
+
+        seq.release();
+        seq.release();
+        assert_eq!(seq.assigned_blocks(), 0);
+    }
+
+    #[test]
+    fn test_release_returns_blocks_to_pools() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = build_prefilled(make_tokens(8), 10, BLOCK_SIZE, &manager).unwrap();
+
+        assert_eq!(manager.available_blocks(), 17);
+        seq.release();
+        assert_eq!(manager.available_blocks(), 20);
+    }
+
+    #[test]
+    fn test_reacquire_basic() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = build_prefilled(make_tokens(8), 10, BLOCK_SIZE, &manager).unwrap();
+
+        seq.append_token(100);
+        seq.append_token(101);
+
+        seq.release();
+        assert!(seq.reacquire(&manager));
+        assert_eq!(seq.assigned_blocks(), 2);
+        assert_eq!(seq.unassigned_blocks(), 0); // no gen block from reacquire
+        assert_eq!(seq.generated_tokens(), 2);
+    }
+
+    #[test]
+    fn test_reacquire_with_cache_hits() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = build_prefilled(make_tokens(8), 10, BLOCK_SIZE, &manager).unwrap();
+
+        seq.release();
+        assert!(seq.reacquire(&manager));
+        assert_eq!(seq.prefix_matched_blocks(), 2);
+    }
+
+    #[test]
+    fn test_reacquire_cleans_up_on_failure() {
+        let manager = create_test_manager::<TestMeta>(4);
+        let mut seq = build_prefilled(make_tokens(4), 10, BLOCK_SIZE, &manager).unwrap();
+
+        seq.release();
+        let _all = manager.allocate_blocks(4).unwrap();
+
+        assert!(!seq.reacquire(&manager));
+        assert_eq!(seq.assigned_blocks(), 0);
+        assert_eq!(seq.unassigned_blocks(), 0);
+    }
+
+    // =========================================================================
+    // RAII
+    // =========================================================================
+
+    #[test]
+    fn test_blocks_returned_on_drop() {
+        let manager = create_test_manager::<TestMeta>(20);
+        {
+            let _seq = build_prefilled(make_tokens(8), 10, BLOCK_SIZE, &manager).unwrap();
+            assert_eq!(manager.available_blocks(), 17);
+        }
+        assert_eq!(manager.available_blocks(), 20);
+    }
+
+    // =========================================================================
+    // Debug
+    // =========================================================================
+
+    #[test]
+    fn test_debug_impl() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let seq = build_prefilled(make_tokens(8), 10, BLOCK_SIZE, &manager).unwrap();
+
+        let debug_str = format!("{seq:?}");
+        assert!(debug_str.contains("RequestSequence"));
+        assert!(debug_str.contains("total_tokens"));
+    }
+}

lib/kvbm-logical/src/integrations/scheduled.rs

+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! # Schedulable Sequence
+//!
+//! Two-phase schedule/apply layer for LLM inference on top of
+//! [`RequestSequence`](crate::RequestSequence).
+//!
+//! [`SchedulableSequence`] enforces a state-machine protocol for prefill,
+//! decode, and speculative decode operations, tracks KV position, and
+//! maintains an append-only event history for observability.
+//!
+//! ## State machine
+//!
+//! ```text
+//!          schedule_prefill        apply_prefill
+//!  Idle ──────────────────► PrefillScheduled ─────► Idle
+//!   │                                                 │
+//!   │    schedule_decode          apply_decode         │
+//!   ├──────────────────► DecodeScheduled ─────────► Idle
+//!   │                                                 │
+//!   │   schedule_speculative    apply_speculative      │
+//!   └──────────────────► SpeculativeScheduled ────► Idle
+//!                   │
+//!                   │   revert_schedule
+//!                   └──────────────────────────► Idle
+//! ```
+//!
+//! Every `schedule_*` call validates preconditions and pre-allocates blocks.
+//! Every `apply_*` call commits the operation (appends tokens, registers
+//! blocks). `revert_schedule` undoes a schedule without applying,
+//! LIFO-releasing pre-allocated blocks.
+//!
+//! ## Dangling token tracking
+//!
+//! `SchedulableSequence` tracks which tokens have had their KV computed via
+//! `kv_position`. The difference `total_tokens - kv_position` gives the
+//! **tail token count** -- tokens whose KV hasn't been computed yet.
+//!
+//! After prefill with a generated token, `tail_tokens() == 1` (the first
+//! generated token is "dangling"). After each decode or speculative step,
+//! the count remains 1 (the newest token replaces the old dangling one).
+//!
+//! `schedule_decode` and `schedule_speculative` enforce exactly 1 tail
+//! token as a precondition.
+//!
+//! ## Typical lifecycle
+//!
+//! ```ignore
+//! use kvbm_logical::SchedulableSequence;
+//!
+//! let tokens: Vec<u32> = (0..8).collect();
+//! let mut seq = SchedulableSequence::<MyMeta>::builder()
+//!     .tokens(tokens)
+//!     .max_output_tokens(10)
+//!     .block_size(4)
+//!     .delegate(my_delegate)  // optional — defaults to NoopDelegate
+//!     .build()?;
+//!
+//! // 1. Optional prefix matching
+//! let matched = seq.match_and_add_prefix(&manager)?;
+//!
+//! // 2. Prefill (single chunk)
+//! seq.schedule_prefill(8 - matched * 4, &manager)?;
+//! seq.apply_prefill(Some(first_generated_token), &manager)?;
+//! // kv_position = 8, tail_tokens = 1
+//!
+//! // 3. Decode loop
+//! while !seq.is_complete() {
+//!     seq.schedule_decode(&manager)?;
+//!     let token = model.forward(&seq);
+//!     let outcome = seq.apply_decode(token, &manager)?;
+//!     // outcome: Continue | BlockCompleted | MaxLength | BlockCompletedAndMaxLength
+//! }
+//!
+//! // 4. Release
+//! seq.release()?;
+//! ```
+//!
+//! ## Chunked prefill
+//!
+//! Split prefill across multiple chunks. Only the **final** chunk (the one
+//! that reaches `num_input_tokens`) must provide a generated token.
+//!
+//! ```ignore
+//! // Chunk 1 (non-final): no token
+//! seq.schedule_prefill(4, &manager)?;
+//! seq.apply_prefill(None, &manager)?;
+//!
+//! // Chunk 2 (final): must provide first generated token
+//! seq.schedule_prefill(4, &manager)?;
+//! seq.apply_prefill(Some(first_token), &manager)?;
+//! ```
+//!
+//! ## Speculative decode
+//!
+//! Schedule a batch of draft tokens, then accept a prefix of them.
+//! Excess pre-allocated blocks are automatically released.
+//!
+//! ```ignore
+//! seq.schedule_speculative(5, &manager)?;
+//! // Model verifies draft tokens, accepts first 3
+//! let outcome = seq.apply_speculative(&[tok1, tok2, tok3], &manager)?;
+//! // Excess blocks LIFO-dropped, tail_tokens still 1
+//! ```
+//!
+//! ## Preemption
+//!
+//! Release and later reacquire blocks. Prefix cache hits reduce
+//! re-computation cost.
+//!
+//! ```ignore
+//! seq.release()?;
+//! // ... later ...
+//! let success = seq.reacquire(&manager)?;
+//! // Reacquire does not allocate a generation block;
+//! // the next schedule_decode handles that.
+//! seq.schedule_decode(&manager)?;
+//! ```
+//!
+//! ## Error handling
+//!
+//! | Error                     | When                                              |
+//! |---------------------------|---------------------------------------------------|
+//! | `ScheduleError::NotIdle`  | `schedule_*` called while already scheduled        |
+//! | `ScheduleError::PrefillNotComplete` | Decode/speculative before prefill done  |
+//! | `ScheduleError::PrefillComplete` | `schedule_prefill` after all input processed |
+//! | `ScheduleError::PrefillOverrun` | Chunk would exceed input token count        |
+//! | `ScheduleError::AllocationFailed` | Not enough blocks in the manager          |
+//! | `ScheduleError::GenerationComplete` | Already hit `max_output_tokens`         |
+//! | `ScheduleError::WrongDanglingCount` | Tail tokens != 1 for decode/speculative |
+//! | `ApplyError::WrongState`  | `apply_*` called in wrong state                   |
+//! | `ApplyError::TokenOnNonFinalChunk` | Token provided on non-final prefill chunk |
+//! | `ApplyError::MissingTokenOnFinalChunk` | Final prefill chunk missing token    |
+//! | `ApplyError::AcceptedExceedsScheduled` | More accepted than draft tokens     |
+//! | `ApplyError::AppendExceedsRemaining` | `append_tokens` exceeds output budget |
+//!
+//! ## Event delegate
+//!
+//! Every lifecycle transition is dispatched to a caller-provided
+//! [`SequenceDelegate`] via `on_event`. Events include `Created`,
+//! `PrefillScheduled`, `PrefillApplied`, `DecodeScheduled`,
+//! `DecodeApplied`, `SpeculativeScheduled`, `SpeculativeApplied`,
+//! `ScheduleReverted`, `UnassignedDropped`, `Released`, and `Reacquired`.
+//!
+//! When no delegate is provided (via the builder or `new(None)`), a
+//! [`NoopDelegate`] is used that silently discards all events.
+
+use std::sync::Arc;
+
+use derive_builder::Builder;
+
+use crate::blocks::BlockMetadata;
+use crate::manager::BlockManager;
+
+use super::request::RequestSequence;
+use dynamo_tokens::Token;
+
+// =============================================================================
+// State types
+// =============================================================================
+
+/// Current scheduling state of a [`SchedulableSequence`].
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SequenceState {
+    Idle,
+    PrefillScheduled {
+        num_tokens: usize,
+        blocks_allocated: usize,
+    },
+    DecodeScheduled {
+        blocks_allocated: usize,
+    },
+    SpeculativeScheduled {
+        num_tokens: usize,
+        blocks_allocated: usize,
+    },
+}
+
+/// Outcome of an [`apply_decode`](SchedulableSequence::apply_decode) or
+/// [`apply_speculative`](SchedulableSequence::apply_speculative) call.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum DecodeOutcome {
+    Continue,
+    BlockCompleted,
+    MaxLength,
+    BlockCompletedAndMaxLength,
+}
+
+/// Append-only event recording a lifecycle transition.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum SequenceEvent {
+    Created {
+        num_input_tokens: usize,
+        max_output_tokens: usize,
+        block_size: usize,
+    },
+    PrefixMatched {
+        blocks_matched: usize,
+    },
+    PrefillScheduled {
+        num_tokens: usize,
+        blocks_allocated: usize,
+    },
+    PrefillApplied {
+        num_tokens: usize,
+        blocks_registered: usize,
+        token_emitted: bool,
+    },
+    DecodeScheduled {
+        blocks_allocated: usize,
+    },
+    DecodeApplied {
+        token: Token,
+        block_completed: bool,
+    },
+    SpeculativeScheduled {
+        num_tokens: usize,
+        blocks_allocated: usize,
+    },
+    SpeculativeApplied {
+        accepted: usize,
+        scheduled: usize,
+        blocks_released: usize,
+    },
+    ScheduleReverted {
+        blocks_released: usize,
+    },
+    UnassignedDropped {
+        count: usize,
+    },
+    Released,
+    Reacquired {
+        prefix_matched: usize,
+        success: bool,
+    },
+}
+
+// =============================================================================
+// Delegate
+// =============================================================================
+
+/// Callback interface for [`SchedulableSequence`] lifecycle events.
+///
+/// Implementations receive every [`SequenceEvent`] as it occurs, enabling
+/// real-time metrics, logging, or external state updates without coupling
+/// event storage to the sequence itself.
+pub trait SequenceDelegate: Send + Sync {
+    fn on_event(&self, event: &SequenceEvent);
+}
+
+/// No-op delegate that silently discards all events.
+///
+/// Used as the default when no delegate is provided to
+/// [`SchedulableSequenceBuilder`].
+pub struct NoopDelegate;
+
+impl SequenceDelegate for NoopDelegate {
+    fn on_event(&self, _event: &SequenceEvent) {}
+}
+
+// =============================================================================
+// Builder
+// =============================================================================
+
+#[doc(hidden)]
+#[derive(Builder)]
+#[builder(
+    name = "SchedulableSequenceBuilder",
+    pattern = "owned",
+    build_fn(private, name = "build_params", error = "anyhow::Error")
+)]
+pub struct SchedulableSequenceParams {
+    tokens: Vec<Token>,
+    max_output_tokens: usize,
+    block_size: u32,
+    #[builder(default, setter(custom))]
+    delegate: Option<Arc<dyn SequenceDelegate>>,
+}
+
+impl SchedulableSequenceBuilder {
+    pub fn delegate(mut self, delegate: Arc<dyn SequenceDelegate>) -> Self {
+        self.delegate = Some(Some(delegate));
+        self
+    }
+
+    pub fn build<T: BlockMetadata>(self) -> anyhow::Result<SchedulableSequence<T>> {
+        let params = self.build_params()?;
+        Ok(SchedulableSequence::new(
+            params.tokens,
+            params.max_output_tokens,
+            params.block_size,
+            params.delegate,
+        ))
+    }
+}
+
+// =============================================================================
+// Errors
+// =============================================================================
+
+/// Error returned by `schedule_*` methods.
+#[derive(Debug, thiserror::Error)]
+pub enum ScheduleError {
+    #[error("expected Idle state, got {state:?}")]
+    NotIdle { state: SequenceState },
+    #[error("prefill overrun: position {position} + {num_tokens} > {num_input_tokens}")]
+    PrefillOverrun {
+        position: usize,
+        num_tokens: usize,
+        num_input_tokens: usize,
+    },
+    #[error("prefill already complete")]
+    PrefillComplete,
+    #[error("prefill not yet complete (position {position} < {num_input_tokens})")]
+    PrefillNotComplete {
+        position: usize,
+        num_input_tokens: usize,
+    },
+    #[error("allocation failed: needed {needed} blocks")]
+    AllocationFailed { needed: usize },
+    #[error("generation complete: {generated} >= {max_output}")]
+    GenerationComplete { generated: usize, max_output: usize },
+    #[error("expected {expected} dangling token(s), got {actual}")]
+    WrongDanglingCount { expected: usize, actual: usize },
+}
+
+/// Error returned by `apply_*` and `revert_schedule` methods.
+#[derive(Debug, thiserror::Error)]
+pub enum ApplyError {
+    #[error("expected {expected}, got {actual:?}")]
+    WrongState {
+        expected: &'static str,
+        actual: SequenceState,
+    },
+    #[error("token provided but prefill not completing this chunk")]
+    TokenOnNonFinalChunk,
+    #[error("accepted {accepted} tokens exceeds scheduled {scheduled}")]
+    AcceptedExceedsScheduled { accepted: usize, scheduled: usize },
+    #[error("final prefill chunk requires a generated token")]
+    MissingTokenOnFinalChunk,
+    #[error("append requested {requested} tokens but only {remaining} remain")]
+    AppendExceedsRemaining { requested: usize, remaining: usize },
+}
+
+// =============================================================================
+// SchedulableSequence
+// =============================================================================
+
+/// Generates simple `&self` forwarding methods to `self.inner`.
+macro_rules! delegate_to_inner {
+    ( $( $(#[$meta:meta])* $vis:vis fn $name:ident(&self) -> $ret:ty; )* ) => {
+        $( $(#[$meta])* $vis fn $name(&self) -> $ret { self.inner.$name() } )*
+    };
+}
+
+/// Two-phase schedule/apply wrapper over [`RequestSequence`].
+///
+/// Enforces a state machine protocol:
+/// - `Idle` → `schedule_*` → `Scheduled` → `apply_*` or `revert_schedule` → `Idle`
+///
+/// Dispatches lifecycle events to a caller-provided [`SequenceDelegate`].
+pub struct SchedulableSequence<T: BlockMetadata> {
+    inner: RequestSequence<T>,
+    state: SequenceState,
+    prefill_position: usize,
+    kv_position: usize,
+    delegate: Arc<dyn SequenceDelegate>,
+}
+
+impl<T: BlockMetadata> SchedulableSequence<T> {
+    // =====================================================================
+    // Construction
+    // =====================================================================
+
+    /// Returns a builder for configuring a `SchedulableSequence`.
+    pub fn builder() -> SchedulableSequenceBuilder {
+        SchedulableSequenceBuilder::default()
+    }
+
+    /// Creates a new `SchedulableSequence` wrapping a fresh `RequestSequence`.
+    ///
+    /// If `delegate` is `None`, a [`NoopDelegate`] is used.
+    pub fn new(
+        tokens: Vec<Token>,
+        max_output_tokens: usize,
+        block_size: u32,
+        delegate: Option<Arc<dyn SequenceDelegate>>,
+    ) -> Self {
+        let inner = RequestSequence::new(tokens, max_output_tokens, block_size);
+        let delegate = delegate.unwrap_or_else(|| Arc::new(NoopDelegate));
+        delegate.on_event(&SequenceEvent::Created {
+            num_input_tokens: inner.num_input_tokens(),
+            max_output_tokens,
+            block_size: block_size as usize,
+        });
+        Self {
+            inner,
+            state: SequenceState::Idle,
+            prefill_position: 0,
+            kv_position: 0,
+            delegate,
+        }
+    }
+
+    // =====================================================================
+    // Prefix matching (Idle only)
+    // =====================================================================
+
+    /// Match and add prefix blocks from the manager's pools.
+    ///
+    /// Advances `prefill_position` by `matched_blocks * block_size`.
+    pub fn match_and_add_prefix(
+        &mut self,
+        manager: &BlockManager<T>,
+    ) -> Result<usize, ScheduleError> {
+        self.require_idle()?;
+
+        let count = self
+            .inner
+            .match_and_add_prefix(manager)
+            .unwrap_or_else(|_| panic!("prefix match should not produce duplicates"));
+
+        if count > 0 {
+            self.prefill_position += count * self.inner.block_size();
+            self.kv_position = self.prefill_position;
+        }
+
+        self.delegate.on_event(&SequenceEvent::PrefixMatched {
+            blocks_matched: count,
+        });
+
+        Ok(count)
+    }
+
+    // =====================================================================
+    // Two-phase prefill
+    // =====================================================================
+
+    /// Schedule a prefill chunk of `num_tokens` tokens.
+    ///
+    /// Allocates the blocks needed to cover `prefill_position + num_tokens`
+    /// tokens, minus already-assigned/staged/unassigned blocks.
+    pub fn schedule_prefill(
+        &mut self,
+        num_tokens: usize,
+        manager: &BlockManager<T>,
+    ) -> Result<(), ScheduleError> {
+        self.require_idle()?;
+
+        if self.is_prefill_complete() {
+            return Err(ScheduleError::PrefillComplete);
+        }
+
+        let num_input = self.inner.num_input_tokens();
+        let new_position = self.prefill_position + num_tokens;
+        if new_position > num_input {
+            return Err(ScheduleError::PrefillOverrun {
+                position: self.prefill_position,
+                num_tokens,
+                num_input_tokens: num_input,
+            });
+        }
+
+        // How many total blocks are needed to cover tokens up to new_position?
+        let bs = self.inner.block_size();
+        let total_blocks_needed = new_position.div_ceil(bs);
+
+        // How many do we already have?
+        let already_have = self.inner.assigned_blocks()
+            + self.inner.staged_blocks()
+            + self.inner.unassigned_blocks();
+
+        let to_allocate = total_blocks_needed.saturating_sub(already_have);
+
+        if to_allocate > 0 && !self.inner.allocate_blocks(to_allocate, manager) {
+            return Err(ScheduleError::AllocationFailed {
+                needed: to_allocate,
+            });
+        }
+
+        self.state = SequenceState::PrefillScheduled {
+            num_tokens,
+            blocks_allocated: to_allocate,
+        };
+
+        self.delegate.on_event(&SequenceEvent::PrefillScheduled {
+            num_tokens,
+            blocks_allocated: to_allocate,
+        });
+
+        Ok(())
+    }
+
+    /// Apply a previously scheduled prefill chunk.
+    ///
+    /// If `token` is `Some`, it is the first generation token emitted on the
+    /// final prefill chunk. Providing a token on a non-final chunk returns an
+    /// error.
+    pub fn apply_prefill(
+        &mut self,
+        token: Option<Token>,
+        manager: &BlockManager<T>,
+    ) -> Result<(), ApplyError> {
+        let (num_tokens, _blocks_allocated) = match self.state {
+            SequenceState::PrefillScheduled {
+                num_tokens,
+                blocks_allocated,
+            } => (num_tokens, blocks_allocated),
+            other => {
+                return Err(ApplyError::WrongState {
+                    expected: "PrefillScheduled",
+                    actual: other,
+                });
+            }
+        };
+
+        let new_position = self.prefill_position + num_tokens;
+        let is_final = new_position >= self.inner.num_input_tokens();
+
+        if token.is_some() && !is_final {
+            return Err(ApplyError::TokenOnNonFinalChunk);
+        }
+        if is_final && token.is_none() && self.inner.max_output_tokens() > 0 {
+            return Err(ApplyError::MissingTokenOnFinalChunk);
+        }
+
+        let blocks_registered_before = self.inner.assigned_blocks();
+
+        // Stage and register the prefill blocks
+        self.inner.complete_and_register_pending(manager);
+        self.prefill_position = new_position;
+        self.kv_position = self.prefill_position;
+
+        // If a token was provided on the final chunk, append it.
+        // The token is "dangling" — its KV hasn't been computed yet.
+        // The block it may complete is NOT registered here; it will be
+        // staged during the next apply_decode after the model forward pass.
+        let token_emitted = token.is_some();
+        if let Some(tok) = token {
+            self.inner.append_token(tok);
+        }
+
+        let blocks_registered =
+            self.inner.assigned_blocks() - blocks_registered_before + self.inner.staged_blocks();
+
+        self.state = SequenceState::Idle;
+        self.delegate.on_event(&SequenceEvent::PrefillApplied {
+            num_tokens,
+            blocks_registered,
+            token_emitted,
+        });
+
+        Ok(())
+    }
+
+    // =====================================================================
+    // Two-phase decode
+    // =====================================================================
+
+    /// Schedule a single decode step.
+    ///
+    /// Allocates blocks for both pending completions (blocks completed at
+    /// the sequence level but not yet staged, e.g. from prefill's dangling
+    /// token crossing a boundary) and the generation block.
+    pub fn schedule_decode(&mut self, manager: &BlockManager<T>) -> Result<(), ScheduleError> {
+        self.require_idle()?;
+        self.require_prefill_complete()?;
+        self.require_not_complete()?;
+        self.require_one_dangling()?;
+
+        // Pending completions: complete blocks in sequence not yet assigned/staged
+        let complete_in_seq = self.inner.complete_sequence_blocks();
+        let registered = self.inner.assigned_blocks() + self.inner.staged_blocks();
+        let pending = complete_in_seq.saturating_sub(registered);
+
+        // Need: pending (for staging after KV computed) + 1 (gen block)
+        let need = pending + 1;
+        let have = self.inner.unassigned_blocks();
+        let to_allocate = need.saturating_sub(have);
+
+        if to_allocate > 0 && !self.inner.allocate_blocks(to_allocate, manager) {
+            return Err(ScheduleError::AllocationFailed {
+                needed: to_allocate,
+            });
+        }
+
+        self.state = SequenceState::DecodeScheduled {
+            blocks_allocated: to_allocate,
+        };
+        self.delegate.on_event(&SequenceEvent::DecodeScheduled {
+            blocks_allocated: to_allocate,
+        });
+
+        Ok(())
+    }
+
+    /// Apply a previously scheduled decode step.
+    pub fn apply_decode(
+        &mut self,
+        token: Token,
+        manager: &BlockManager<T>,
+    ) -> Result<DecodeOutcome, ApplyError> {
+        let _blocks_allocated = match self.state {
+            SequenceState::DecodeScheduled { blocks_allocated } => blocks_allocated,
+            other => {
+                return Err(ApplyError::WrongState {
+                    expected: "DecodeScheduled",
+                    actual: other,
+                });
+            }
+        };
+
+        let crossed = self.inner.append_token(token);
+        let block_completed = crossed.is_some();
+
+        // Always stage pending completions — handles both:
+        // 1. Blocks completed during prefill's token append (deferred staging)
+        // 2. Block just completed by this decode token
+        self.inner.complete_and_register_pending(manager);
+
+        self.kv_position += 1;
+        self.state = SequenceState::Idle;
+        self.delegate.on_event(&SequenceEvent::DecodeApplied {
+            token,
+            block_completed,
+        });
+
+        let is_complete = self.inner.is_complete();
+        Ok(match (block_completed, is_complete) {
+            (false, false) => DecodeOutcome::Continue,
+            (true, false) => DecodeOutcome::BlockCompleted,
+            (false, true) => DecodeOutcome::MaxLength,
+            (true, true) => DecodeOutcome::BlockCompletedAndMaxLength,
+        })
+    }
+
+    // =====================================================================
+    // Speculative decode
+    // =====================================================================
+
+    /// Schedule a speculative decode of `num_draft_tokens` tokens.
+    ///
+    /// Allocates enough blocks to accommodate the draft tokens, accounting
+    /// for the current partial block state and all already-held blocks
+    /// (assigned + staged + unassigned).
+    pub fn schedule_speculative(
+        &mut self,
+        num_draft_tokens: usize,
+        manager: &BlockManager<T>,
+    ) -> Result<(), ScheduleError> {
+        self.require_idle()?;
+        self.require_prefill_complete()?;
+        self.require_not_complete()?;
+        self.require_one_dangling()?;
+
+        // Clamp to remaining output budget to prevent append_token panics.
+        let num_draft_tokens = num_draft_tokens.min(self.inner.remaining_tokens());
+
+        let bs = self.inner.block_size();
+        let future_total = self.inner.total_tokens() + num_draft_tokens;
+        let future_blocks = future_total.div_ceil(bs);
+        let have = self.inner.assigned_blocks()
+            + self.inner.staged_blocks()
+            + self.inner.unassigned_blocks();
+        let to_allocate = future_blocks.saturating_sub(have);
+
+        if to_allocate > 0 && !self.inner.allocate_blocks(to_allocate, manager) {
+            return Err(ScheduleError::AllocationFailed {
+                needed: to_allocate,
+            });
+        }
+
+        self.state = SequenceState::SpeculativeScheduled {
+            num_tokens: num_draft_tokens,
+            blocks_allocated: to_allocate,
+        };
+        self.delegate
+            .on_event(&SequenceEvent::SpeculativeScheduled {
+                num_tokens: num_draft_tokens,
+                blocks_allocated: to_allocate,
+            });
+
+        Ok(())
+    }
+
+    /// Apply a speculative decode with `accepted` tokens (a prefix of the
+    /// scheduled draft).
+    ///
+    /// Excess unassigned blocks (allocated for rejected draft tokens) are
+    /// LIFO-dropped, returning them to the pool via RAII.
+    pub fn apply_speculative(
+        &mut self,
+        accepted: &[Token],
+        manager: &BlockManager<T>,
+    ) -> Result<DecodeOutcome, ApplyError> {
+        let (scheduled_tokens, _blocks_allocated) = match self.state {
+            SequenceState::SpeculativeScheduled {
+                num_tokens,
+                blocks_allocated,
+            } => (num_tokens, blocks_allocated),
+            other => {
+                return Err(ApplyError::WrongState {
+                    expected: "SpeculativeScheduled",
+                    actual: other,
+                });
+            }
+        };
+
+        if accepted.len() > scheduled_tokens {
+            return Err(ApplyError::AcceptedExceedsScheduled {
+                accepted: accepted.len(),
+                scheduled: scheduled_tokens,
+            });
+        }
+
+        // Append accepted tokens one at a time, tracking boundary crossings
+        let mut block_completed = false;
+        for &token in accepted {
+            let crossed = self.inner.append_token(token);
+            if crossed.is_some() {
+                block_completed = true;
+            }
+        }
+
+        // Stage and register all pending completions (including any from
+        // prefill's deferred staging and blocks just completed above)
+        self.inner.complete_and_register_pending(manager);
+
+        self.kv_position += accepted.len();
+
+        // LIFO-drop excess unassigned blocks.
+        // After appending accepted tokens, the generation block (if any) is the
+        // remaining unassigned. If we over-allocated for the draft, drop excess.
+        let excess = self.lifo_drop_excess_unassigned();
+
+        self.state = SequenceState::Idle;
+        self.delegate.on_event(&SequenceEvent::SpeculativeApplied {
+            accepted: accepted.len(),
+            scheduled: scheduled_tokens,
+            blocks_released: excess,
+        });
+
+        let is_complete = self.inner.is_complete();
+        Ok(match (block_completed, is_complete) {
+            (false, false) => DecodeOutcome::Continue,
+            (true, false) => DecodeOutcome::BlockCompleted,
+            (false, true) => DecodeOutcome::MaxLength,
+            (true, true) => DecodeOutcome::BlockCompletedAndMaxLength,
+        })
+    }
+
+    // =====================================================================
+    // Revert
+    // =====================================================================
+
+    /// Revert a scheduled (but not yet applied) operation.
+    ///
+    /// LIFO-pops the `blocks_allocated` unassigned blocks that were
+    /// pre-allocated during the schedule phase. The dropped RAII guards
+    /// return blocks to the manager's pools.
+    pub fn revert_schedule(&mut self) -> Result<(), ApplyError> {
+        let blocks_to_release = match self.state {
+            SequenceState::PrefillScheduled {
+                blocks_allocated, ..
+            } => blocks_allocated,
+            SequenceState::DecodeScheduled { blocks_allocated } => blocks_allocated,
+            SequenceState::SpeculativeScheduled {
+                blocks_allocated, ..
+            } => blocks_allocated,
+            other => {
+                return Err(ApplyError::WrongState {
+                    expected: "any Scheduled state",
+                    actual: other,
+                });
+            }
+        };
+
+        self.lifo_pop_unassigned(blocks_to_release);
+
+        self.state = SequenceState::Idle;
+        self.delegate.on_event(&SequenceEvent::ScheduleReverted {
+            blocks_released: blocks_to_release,
+        });
+
+        Ok(())
+    }
+
+    // =====================================================================
+    // Explicit LIFO drop of unassigned blocks
+    // =====================================================================
+
+    /// LIFO-drop up to `count` unassigned blocks. Returns the actual number
+    /// dropped. Valid only in Idle state.
+    pub fn drop_unassigned(&mut self, count: usize) -> usize {
+        assert!(
+            self.state == SequenceState::Idle,
+            "drop_unassigned called in non-Idle state: {:?}",
+            self.state
+        );
+        let dropped = self.lifo_pop_unassigned(count);
+        if dropped > 0 {
+            self.delegate
+                .on_event(&SequenceEvent::UnassignedDropped { count: dropped });
+        }
+        dropped
+    }
+
+    // =====================================================================
+    // Lifecycle
+    // =====================================================================
+
+    /// Release all block assignments (RAII returns them to pools).
+    pub fn release(&mut self) -> Result<(), ApplyError> {
+        self.require_idle_for_apply()?;
+        self.inner.release();
+        self.delegate.on_event(&SequenceEvent::Released);
+        Ok(())
+    }
+
+    /// Re-acquire blocks from the manager after a release/preemption.
+    pub fn reacquire(&mut self, manager: &BlockManager<T>) -> Result<bool, ApplyError> {
+        self.require_idle_for_apply()?;
+        let success = self.inner.reacquire(manager);
+        let prefix_matched = self.inner.prefix_matched_blocks();
+        self.delegate.on_event(&SequenceEvent::Reacquired {
+            prefix_matched,
+            success,
+        });
+        Ok(success)
+    }
+
+    // =====================================================================
+    // Token append (no KV advancement)
+    // =====================================================================
+
+    /// Append tokens to the sequence without advancing `kv_position`.
+    /// Each appended token increases the dangling count.
+    /// Requires Idle state.
+    pub fn append_tokens(&mut self, tokens: &[Token]) -> Result<(), ApplyError> {
+        self.require_idle_for_apply()?;
+        let remaining = self.inner.remaining_tokens();
+        if tokens.len() > remaining {
+            return Err(ApplyError::AppendExceedsRemaining {
+                requested: tokens.len(),
+                remaining,
+            });
+        }
+        for &token in tokens {
+            self.inner.append_token(token);
+        }
+        Ok(())
+    }
+
+    // =====================================================================
+    // Accessors
+    // =====================================================================
+
+    /// Current scheduling state.
+    pub fn state(&self) -> SequenceState {
+        self.state
+    }
+
+    /// How many input tokens have been processed so far.
+    pub fn prefill_position(&self) -> usize {
+        self.prefill_position
+    }
+
+    /// Whether all input tokens have been processed.
+    pub fn is_prefill_complete(&self) -> bool {
+        self.prefill_position >= self.inner.num_input_tokens()
+    }
+
+    /// Number of tokens whose KV has been computed.
+    pub fn kv_position(&self) -> usize {
+        self.kv_position
+    }
+
+    /// Number of tokens whose KV hasn't been computed yet.
+    /// After prefill: 1 (the first generated token). After decode: 1 (the new token).
+    pub fn tail_tokens(&self) -> usize {
+        self.inner.total_tokens().saturating_sub(self.kv_position)
+    }
+
+    /// Reference to the delegate.
+    pub fn delegate(&self) -> &Arc<dyn SequenceDelegate> {
+        &self.delegate
+    }
+
+    // Forwarded from RequestSequence
+
+    delegate_to_inner! {
+        pub fn generated_tokens(&self) -> usize;
+        pub fn max_output_tokens(&self) -> usize;
+        pub fn num_input_tokens(&self) -> usize;
+        pub fn total_tokens(&self) -> usize;
+        pub fn remaining_tokens(&self) -> usize;
+        pub fn num_blocks(&self) -> usize;
+        pub fn assigned_blocks(&self) -> usize;
+        pub fn staged_blocks(&self) -> usize;
+        pub fn unassigned_blocks(&self) -> usize;
+        pub fn prefix_matched_blocks(&self) -> usize;
+        pub fn block_size(&self) -> usize;
+        pub fn is_complete(&self) -> bool;
+    }
+
+    /// Reference to the underlying `RequestSequence`.
+    pub fn inner(&self) -> &RequestSequence<T> {
+        &self.inner
+    }
+
+    /// Mutable reference to the underlying `RequestSequence`.
+    #[allow(dead_code)]
+    pub(crate) fn inner_mut(&mut self) -> &mut RequestSequence<T> {
+        &mut self.inner
+    }
+
+    // =====================================================================
+    // Private helpers
+    // =====================================================================
+
+    fn require_idle(&self) -> Result<(), ScheduleError> {
+        if self.state != SequenceState::Idle {
+            return Err(ScheduleError::NotIdle { state: self.state });
+        }
+        Ok(())
+    }
+
+    fn require_idle_for_apply(&self) -> Result<(), ApplyError> {
+        if self.state != SequenceState::Idle {
+            return Err(ApplyError::WrongState {
+                expected: "Idle",
+                actual: self.state,
+            });
+        }
+        Ok(())
+    }
+
+    fn require_prefill_complete(&self) -> Result<(), ScheduleError> {
+        if !self.is_prefill_complete() {
+            return Err(ScheduleError::PrefillNotComplete {
+                position: self.prefill_position,
+                num_input_tokens: self.inner.num_input_tokens(),
+            });
+        }
+        Ok(())
+    }
+
+    fn require_not_complete(&self) -> Result<(), ScheduleError> {
+        if self.inner.is_complete() {
+            return Err(ScheduleError::GenerationComplete {
+                generated: self.inner.generated_tokens(),
+                max_output: self.inner.max_output_tokens(),
+            });
+        }
+        Ok(())
+    }
+
+    fn require_one_dangling(&self) -> Result<(), ScheduleError> {
+        let dangling = self.tail_tokens();
+        if dangling != 1 {
+            return Err(ScheduleError::WrongDanglingCount {
+                expected: 1,
+                actual: dangling,
+            });
+        }
+        Ok(())
+    }
+
+    /// LIFO-pop up to `count` unassigned blocks. Returns the actual count dropped.
+    fn lifo_pop_unassigned(&mut self, count: usize) -> usize {
+        let assignments = self.inner.assignments_mut();
+        let mut dropped = 0;
+        for _ in 0..count {
+            if assignments.pop_last_unassigned().is_some() {
+                dropped += 1;
+            } else {
+                break;
+            }
+        }
+        dropped
+    }
+
+    /// After speculative apply: drop any excess unassigned blocks beyond
+    /// what's needed for the current partial block (at most 1 gen block).
+    fn lifo_drop_excess_unassigned(&mut self) -> usize {
+        let bs = self.inner.block_size();
+        let total = self.inner.total_tokens();
+        // We need at most 1 unassigned (gen) block if there's a partial block in progress
+        // AND we haven't hit max output tokens.
+        let need_gen = if self.inner.is_complete() {
+            0
+        } else if !total.is_multiple_of(bs) {
+            // Partial block in progress — already have an unassigned block covering it
+            1
+        } else {
+            // On exact boundary — the last block was just completed & registered.
+            // We still keep 1 gen block for future decode unless complete.
+            1
+        };
+
+        let current = self.inner.unassigned_blocks();
+        let excess = current.saturating_sub(need_gen);
+        self.lifo_pop_unassigned(excess)
+    }
+}
+
+impl<T: BlockMetadata> std::fmt::Debug for SchedulableSequence<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("SchedulableSequence")
+            .field("state", &self.state)
+            .field("prefill_position", &self.prefill_position)
+            .field("kv_position", &self.kv_position)
+            .field("inner", &self.inner)
+            .finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::testing::{TestMeta, create_test_manager};
+    use std::sync::Mutex;
+
+    const BLOCK_SIZE: u32 = 4;
+
+    // =========================================================================
+    // Test delegate
+    // =========================================================================
+
+    struct CollectingDelegate {
+        events: Mutex<Vec<SequenceEvent>>,
+    }
+
+    impl CollectingDelegate {
+        fn new() -> Self {
+            Self {
+                events: Mutex::new(Vec::new()),
+            }
+        }
+
+        fn events(&self) -> Vec<SequenceEvent> {
+            self.events.lock().unwrap().clone()
+        }
+    }
+
+    impl SequenceDelegate for CollectingDelegate {
+        fn on_event(&self, event: &SequenceEvent) {
+            self.events.lock().unwrap().push(event.clone());
+        }
+    }
+
+    fn noop_delegate() -> Option<Arc<dyn SequenceDelegate>> {
+        None
+    }
+
+    fn make_tokens(n: usize) -> Vec<Token> {
+        (0..n as u32).collect()
+    }
+
+    // =========================================================================
+    // Construction
+    // =========================================================================
+
+    #[test]
+    fn test_new_starts_idle() {
+        let delegate = Arc::new(CollectingDelegate::new());
+        let seq = SchedulableSequence::<TestMeta>::new(
+            make_tokens(8),
+            10,
+            BLOCK_SIZE,
+            Some(delegate.clone()),
+        );
+
+        assert_eq!(seq.state(), SequenceState::Idle);
+        assert_eq!(seq.prefill_position(), 0);
+        assert_eq!(seq.kv_position(), 0);
+        assert_eq!(seq.tail_tokens(), 8);
+        assert_eq!(seq.num_input_tokens(), 8);
+        assert_eq!(seq.max_output_tokens(), 10);
+        assert_eq!(seq.block_size(), BLOCK_SIZE as usize);
+        assert!(!seq.is_prefill_complete());
+
+        let events = delegate.events();
+        assert_eq!(events.len(), 1);
+        assert_eq!(
+            events[0],
+            SequenceEvent::Created {
+                num_input_tokens: 8,
+                max_output_tokens: 10,
+                block_size: BLOCK_SIZE as usize,
+            }
+        );
+    }
+
+    // =========================================================================
+    // State machine enforcement
+    // =========================================================================
+
+    #[test]
+    fn test_schedule_prefill_requires_idle() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        seq.schedule_prefill(4, &manager).unwrap();
+        let err = seq.schedule_prefill(4, &manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::NotIdle { .. }));
+    }
+
+    #[test]
+    fn test_schedule_decode_requires_idle() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(4), 10, BLOCK_SIZE, noop_delegate());
+
+        // Prefill first
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+
+        // Schedule decode
+        seq.schedule_decode(&manager).unwrap();
+        let err = seq.schedule_decode(&manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::NotIdle { .. }));
+    }
+
+    #[test]
+    fn test_apply_prefill_requires_scheduled() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        let err = seq.apply_prefill(None, &manager).unwrap_err();
+        assert!(matches!(err, ApplyError::WrongState { .. }));
+    }
+
+    #[test]
+    fn test_apply_decode_requires_scheduled() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(4), 10, BLOCK_SIZE, noop_delegate());
+
+        let err = seq.apply_decode(100, &manager).unwrap_err();
+        assert!(matches!(err, ApplyError::WrongState { .. }));
+    }
+
+    #[test]
+    fn test_decode_requires_prefill_complete() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        let err = seq.schedule_decode(&manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::PrefillNotComplete { .. }));
+    }
+
+    #[test]
+    fn test_speculative_requires_prefill_complete() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        let err = seq.schedule_speculative(3, &manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::PrefillNotComplete { .. }));
+    }
+
+    // =========================================================================
+    // Prefill
+    // =========================================================================
+
+    #[test]
+    fn test_prefill_single_chunk() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        // Schedule and apply the full prefill (no gen block allocated)
+        seq.schedule_prefill(8, &manager).unwrap();
+        assert_eq!(
+            seq.state(),
+            SequenceState::PrefillScheduled {
+                num_tokens: 8,
+                blocks_allocated: 2, // 2 input blocks only
+            }
+        );
+
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+        assert_eq!(seq.state(), SequenceState::Idle);
+        assert_eq!(seq.prefill_position(), 8);
+        assert_eq!(seq.kv_position(), 8);
+        assert!(seq.is_prefill_complete());
+        assert_eq!(seq.assigned_blocks(), 2);
+        assert_eq!(seq.unassigned_blocks(), 0); // no gen block
+        assert_eq!(seq.tail_tokens(), 1); // token 1000
+    }
+
+    #[test]
+    fn test_prefill_chunked() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        // Chunk 1: first 4 tokens (1 block, non-final)
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(None, &manager).unwrap();
+        assert_eq!(seq.prefill_position(), 4);
+        assert_eq!(seq.kv_position(), 4);
+        assert!(!seq.is_prefill_complete());
+        assert_eq!(seq.assigned_blocks(), 1);
+
+        // Chunk 2: next 4 tokens (1 block, final — must provide token)
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+        assert_eq!(seq.prefill_position(), 8);
+        assert_eq!(seq.kv_position(), 8);
+        assert!(seq.is_prefill_complete());
+        assert_eq!(seq.assigned_blocks(), 2);
+        assert_eq!(seq.unassigned_blocks(), 0); // no gen block
+        assert_eq!(seq.tail_tokens(), 1);
+    }
+
+    #[test]
+    fn test_prefill_final_with_first_token() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(4), 10, BLOCK_SIZE, noop_delegate());
+
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(Some(100), &manager).unwrap();
+
+        assert!(seq.is_prefill_complete());
+        assert_eq!(seq.generated_tokens(), 1);
+        assert_eq!(seq.total_tokens(), 5);
+        assert_eq!(seq.kv_position(), 4);
+        assert_eq!(seq.tail_tokens(), 1);
+    }
+
+    #[test]
+    fn test_prefill_token_on_non_final_rejected() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        seq.schedule_prefill(4, &manager).unwrap();
+        let err = seq.apply_prefill(Some(100), &manager).unwrap_err();
+        assert!(matches!(err, ApplyError::TokenOnNonFinalChunk));
+    }
+
+    #[test]
+    fn test_prefill_overrun_rejected() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        let err = seq.schedule_prefill(9, &manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::PrefillOverrun { .. }));
+    }
+
+    #[test]
+    fn test_prefill_allocation_failure() {
+        let manager = create_test_manager::<TestMeta>(1); // only 1 block
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        let err = seq.schedule_prefill(8, &manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::AllocationFailed { .. }));
+        // State remains idle
+        assert_eq!(seq.state(), SequenceState::Idle);
+    }
+
+    #[test]
+    fn test_schedule_prefill_after_complete_rejected() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(4), 10, BLOCK_SIZE, noop_delegate());
+
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+
+        let err = seq.schedule_prefill(1, &manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::PrefillComplete));
+    }
+
+    #[test]
+    fn test_apply_prefill_none_on_final_rejected() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(4), 10, BLOCK_SIZE, noop_delegate());
+
+        seq.schedule_prefill(4, &manager).unwrap();
+        let err = seq.apply_prefill(None, &manager).unwrap_err();
+        assert!(matches!(err, ApplyError::MissingTokenOnFinalChunk));
+    }
+
+    // =========================================================================
+    // Decode
+    // =========================================================================
+
+    /// Helper: create a sequence with prefill done and first generated token.
+    ///
+    /// After this: `total_tokens = num_input + 1`, `generated_tokens = 1`,
+    /// `kv_position = num_input`, `tail_tokens = 1`, `unassigned_blocks = 0`.
+    fn prefilled_seq(
+        num_input: usize,
+        max_output: usize,
+        manager: &BlockManager<TestMeta>,
+    ) -> SchedulableSequence<TestMeta> {
+        let mut seq = SchedulableSequence::new(
+            make_tokens(num_input),
+            max_output,
+            BLOCK_SIZE,
+            noop_delegate(),
+        );
+        if num_input > 0 {
+            seq.schedule_prefill(num_input, manager).unwrap();
+            seq.apply_prefill(Some(1000), manager).unwrap();
+        }
+        seq
+    }
+
+    #[test]
+    fn test_decode_continue() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(5, 10, &manager);
+        // After prefill: total=6, gen=1, kv=5, assigned=1, unassigned=0
+
+        seq.schedule_decode(&manager).unwrap();
+        let outcome = seq.apply_decode(100, &manager).unwrap();
+        assert_eq!(outcome, DecodeOutcome::Continue);
+        assert_eq!(seq.generated_tokens(), 2); // 1 from prefill + 1 decode
+        assert_eq!(seq.state(), SequenceState::Idle);
+    }
+
+    #[test]
+    fn test_decode_block_completed() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(4, 10, &manager);
+        // After prefill: total=5, kv=4, assigned=1, unassigned=0
+
+        // 3 decodes to reach block boundary at total=8
+        for _ in 0..2 {
+            seq.schedule_decode(&manager).unwrap();
+            let outcome = seq.apply_decode(100, &manager).unwrap();
+            assert_eq!(outcome, DecodeOutcome::Continue);
+        }
+        seq.schedule_decode(&manager).unwrap();
+        let outcome = seq.apply_decode(100, &manager).unwrap();
+        assert_eq!(outcome, DecodeOutcome::BlockCompleted);
+        assert_eq!(seq.assigned_blocks(), 2);
+    }
+
+    #[test]
+    fn test_decode_max_length() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(5, 2, &manager);
+        // After prefill: gen=1, max=2. One more decode reaches max.
+
+        seq.schedule_decode(&manager).unwrap();
+        let outcome = seq.apply_decode(100, &manager).unwrap();
+        assert_eq!(outcome, DecodeOutcome::MaxLength);
+        assert!(seq.is_complete());
+    }
+
+    #[test]
+    fn test_decode_block_and_max() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(4, 4, &manager);
+        // After prefill: total=5, gen=1, max=4. Need 3 more decodes.
+        // At decode 3: total=8 (boundary) AND gen=4=max.
+
+        for _ in 0..2 {
+            seq.schedule_decode(&manager).unwrap();
+            seq.apply_decode(100, &manager).unwrap();
+        }
+        seq.schedule_decode(&manager).unwrap();
+        let outcome = seq.apply_decode(100, &manager).unwrap();
+        assert_eq!(outcome, DecodeOutcome::BlockCompletedAndMaxLength);
+    }
+
+    #[test]
+    fn test_decode_allocates_gen_block() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(4, 10, &manager);
+        // After prefill: 1 assigned, 0 unassigned
+
+        // First schedule_decode always allocates 1 (gen block)
+        seq.schedule_decode(&manager).unwrap();
+        assert_eq!(
+            seq.state(),
+            SequenceState::DecodeScheduled {
+                blocks_allocated: 1
+            }
+        );
+        assert_eq!(seq.unassigned_blocks(), 1);
+
+        // Decode until block boundary: 3 decodes total reach total=8
+        seq.apply_decode(100, &manager).unwrap(); // total=6
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(101, &manager).unwrap(); // total=7
+        seq.schedule_decode(&manager).unwrap();
+        let outcome = seq.apply_decode(102, &manager).unwrap(); // total=8, crosses boundary
+        assert_eq!(outcome, DecodeOutcome::BlockCompleted);
+        assert_eq!(seq.unassigned_blocks(), 0);
+
+        // Next schedule_decode should allocate 1 gen block
+        seq.schedule_decode(&manager).unwrap();
+        assert_eq!(
+            seq.state(),
+            SequenceState::DecodeScheduled {
+                blocks_allocated: 1
+            }
+        );
+        assert_eq!(seq.unassigned_blocks(), 1);
+    }
+
+    #[test]
+    fn test_decode_generation_complete_rejected() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(5, 2, &manager);
+        // After prefill: gen=1, max=2
+
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(100, &manager).unwrap();
+        assert!(seq.is_complete()); // gen=2=max
+
+        let err = seq.schedule_decode(&manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::GenerationComplete { .. }));
+    }
+
+    // =========================================================================
+    // Speculative decode
+    // =========================================================================
+
+    #[test]
+    fn test_speculative_basic() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(8, 10, &manager);
+        // After prefill: total=9, gen=1, kv=8, assigned=2, unassigned=0
+
+        // Schedule 2 draft tokens (stay within block)
+        seq.schedule_speculative(2, &manager).unwrap();
+        assert!(matches!(
+            seq.state(),
+            SequenceState::SpeculativeScheduled { num_tokens: 2, .. }
+        ));
+
+        // Accept both
+        let outcome = seq.apply_speculative(&[100, 101], &manager).unwrap();
+        assert_eq!(outcome, DecodeOutcome::Continue);
+        assert_eq!(seq.generated_tokens(), 3); // 1 from prefill + 2 speculative
+        assert_eq!(seq.state(), SequenceState::Idle);
+    }
+
+    #[test]
+    fn test_speculative_partial_accept() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let delegate = Arc::new(CollectingDelegate::new());
+        let mut seq =
+            SchedulableSequence::new(make_tokens(4), 10, BLOCK_SIZE, Some(delegate.clone()));
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+        // After prefill: total=5, assigned=1, unassigned=0
+
+        let avail_before = manager.available_blocks();
+
+        // Schedule 4 draft tokens
+        seq.schedule_speculative(4, &manager).unwrap();
+
+        // Accept only 2 → excess blocks should be released
+        let outcome = seq.apply_speculative(&[100, 101], &manager).unwrap();
+        assert_eq!(outcome, DecodeOutcome::Continue);
+        assert_eq!(seq.generated_tokens(), 3); // 1 from prefill + 2
+        assert_eq!(seq.unassigned_blocks(), 1); // keep 1 gen block
+
+        // Check delegate records the release
+        let events = delegate.events();
+        let last = events.last().unwrap();
+        if let SequenceEvent::SpeculativeApplied {
+            accepted,
+            scheduled,
+            ..
+        } = last
+        {
+            assert_eq!(*accepted, 2);
+            assert_eq!(*scheduled, 4);
+        } else {
+            panic!("expected SpeculativeApplied");
+        }
+
+        let _ = avail_before;
+    }
+
+    #[test]
+    fn test_speculative_single_accept() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(4, 10, &manager);
+
+        seq.schedule_speculative(5, &manager).unwrap();
+        let outcome = seq.apply_speculative(&[100], &manager).unwrap();
+        assert_eq!(outcome, DecodeOutcome::Continue);
+        assert_eq!(seq.generated_tokens(), 2); // 1 from prefill + 1
+    }
+
+    #[test]
+    fn test_speculative_zero_accept() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(4, 10, &manager);
+
+        let avail_before = manager.available_blocks();
+        seq.schedule_speculative(3, &manager).unwrap();
+        let avail_after_schedule = manager.available_blocks();
+
+        let outcome = seq.apply_speculative(&[], &manager).unwrap();
+        assert_eq!(outcome, DecodeOutcome::Continue);
+        assert_eq!(seq.generated_tokens(), 1); // 1 from prefill, 0 speculative
+
+        assert_eq!(seq.unassigned_blocks(), 1); // keep 1 gen block
+        assert!(manager.available_blocks() >= avail_after_schedule);
+        let _ = avail_before;
+    }
+
+    #[test]
+    fn test_speculative_exceeds_scheduled_rejected() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(4, 10, &manager);
+
+        seq.schedule_speculative(2, &manager).unwrap();
+        let err = seq
+            .apply_speculative(&[100, 101, 102], &manager)
+            .unwrap_err();
+        assert!(matches!(
+            err,
+            ApplyError::AcceptedExceedsScheduled {
+                accepted: 3,
+                scheduled: 2,
+            }
+        ));
+    }
+
+    #[test]
+    fn test_speculative_block_boundaries() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(7, 20, &manager);
+        // After prefill: total=8, kv=7, assigned=1, unassigned=0
+        // Block 1 complete at sequence level (4-6 + token 1000) but NOT registered
+
+        // 5 speculative tokens: registers pending block 1, crosses another boundary
+        seq.schedule_speculative(5, &manager).unwrap();
+        let outcome = seq
+            .apply_speculative(&[100, 101, 102, 103, 104], &manager)
+            .unwrap();
+        assert_eq!(outcome, DecodeOutcome::BlockCompleted);
+        assert_eq!(seq.generated_tokens(), 6); // 1 from prefill + 5
+        assert_eq!(seq.assigned_blocks(), 3); // block 0 + blocks 1,2 registered
+    }
+
+    // =========================================================================
+    // Revert
+    // =========================================================================
+
+    #[test]
+    fn test_revert_prefill() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        let avail_before = manager.available_blocks();
+        seq.schedule_prefill(4, &manager).unwrap();
+        assert!(manager.available_blocks() < avail_before);
+
+        seq.revert_schedule().unwrap();
+        assert_eq!(seq.state(), SequenceState::Idle);
+        assert_eq!(manager.available_blocks(), avail_before);
+    }
+
+    #[test]
+    fn test_revert_decode() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(4, 10, &manager);
+        // After prefill: total=5, assigned=1, unassigned=0
+
+        // Decode until boundary at total=8 → unassigned drops to 0
+        for _ in 0..3 {
+            seq.schedule_decode(&manager).unwrap();
+            seq.apply_decode(100, &manager).unwrap();
+        }
+        assert_eq!(seq.unassigned_blocks(), 0);
+
+        let avail_before = manager.available_blocks();
+        seq.schedule_decode(&manager).unwrap();
+        assert_eq!(manager.available_blocks(), avail_before - 1);
+
+        seq.revert_schedule().unwrap();
+        assert_eq!(seq.state(), SequenceState::Idle);
+        assert_eq!(manager.available_blocks(), avail_before);
+    }
+
+    #[test]
+    fn test_revert_speculative() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(4, 10, &manager);
+
+        let avail_before = manager.available_blocks();
+        seq.schedule_speculative(4, &manager).unwrap();
+        let allocated = avail_before - manager.available_blocks();
+
+        seq.revert_schedule().unwrap();
+        assert_eq!(seq.state(), SequenceState::Idle);
+        // Blocks allocated during schedule should be returned
+        assert_eq!(manager.available_blocks(), avail_before);
+        assert!(allocated > 0 || seq.unassigned_blocks() > 0);
+    }
+
+    #[test]
+    fn test_revert_returns_blocks_to_manager() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        let avail_before = manager.available_blocks();
+        seq.schedule_prefill(8, &manager).unwrap();
+        let avail_scheduled = manager.available_blocks();
+        assert!(avail_scheduled < avail_before);
+
+        seq.revert_schedule().unwrap();
+        assert_eq!(manager.available_blocks(), avail_before);
+    }
+
+    // =========================================================================
+    // Drop unassigned
+    // =========================================================================
+
+    #[test]
+    fn test_drop_unassigned_lifo() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(4, 10, &manager);
+        // After prefill: 1 assigned, 0 unassigned
+
+        // Do one decode to get 1 unassigned (gen block)
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(100, &manager).unwrap();
+        assert_eq!(seq.unassigned_blocks(), 1);
+
+        let dropped = seq.drop_unassigned(1);
+        assert_eq!(dropped, 1);
+        assert_eq!(seq.unassigned_blocks(), 0);
+    }
+
+    #[test]
+    fn test_drop_unassigned_partial() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(8, 10, &manager);
+        // After prefill: 2 assigned, 0 unassigned
+
+        // Do one decode to get 1 unassigned
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(100, &manager).unwrap();
+        assert_eq!(seq.unassigned_blocks(), 1);
+
+        // Try to drop 5, but only 1 available
+        let dropped = seq.drop_unassigned(5);
+        assert_eq!(dropped, 1);
+        assert_eq!(seq.unassigned_blocks(), 0);
+    }
+
+    #[test]
+    fn test_drop_unassigned_zero() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = prefilled_seq(4, 10, &manager);
+
+        // Do one decode to get 1 unassigned
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(100, &manager).unwrap();
+
+        let dropped = seq.drop_unassigned(0);
+        assert_eq!(dropped, 0);
+        assert_eq!(seq.unassigned_blocks(), 1); // unchanged
+    }
+
+    // =========================================================================
+    // Delegate event collection
+    // =========================================================================
+
+    #[test]
+    fn test_delegate_full_lifecycle() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let delegate = Arc::new(CollectingDelegate::new());
+        // max_output=3: 1 from prefill + 2 decodes
+        let mut seq = SchedulableSequence::<TestMeta>::new(
+            make_tokens(4),
+            3,
+            BLOCK_SIZE,
+            Some(delegate.clone()),
+        );
+
+        // Prefill
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+
+        // Decode token 1
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(100, &manager).unwrap();
+
+        // Decode token 2
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(101, &manager).unwrap();
+
+        // Release
+        seq.release().unwrap();
+
+        let h = delegate.events();
+        assert_eq!(h.len(), 8);
+
+        assert!(matches!(h[0], SequenceEvent::Created { .. }));
+        assert!(matches!(h[1], SequenceEvent::PrefillScheduled { .. }));
+        assert!(matches!(h[2], SequenceEvent::PrefillApplied { .. }));
+        assert!(matches!(h[3], SequenceEvent::DecodeScheduled { .. }));
+        assert!(matches!(h[4], SequenceEvent::DecodeApplied { .. }));
+        assert!(matches!(h[5], SequenceEvent::DecodeScheduled { .. }));
+        assert!(matches!(h[6], SequenceEvent::DecodeApplied { .. }));
+        assert!(matches!(h[7], SequenceEvent::Released));
+    }
+
+    // =========================================================================
+    // Integration: full lifecycle
+    // =========================================================================
+
+    #[test]
+    fn test_full_lifecycle_prefill_decode_release() {
+        let manager = create_test_manager::<TestMeta>(20);
+        // max_output=7: 1 from prefill + 6 decodes
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(6), 7, BLOCK_SIZE, noop_delegate());
+
+        // Prefill 6 tokens → 1 complete block + 2 partial
+        seq.schedule_prefill(6, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+
+        assert!(seq.is_prefill_complete());
+        assert_eq!(seq.assigned_blocks(), 1);
+        assert_eq!(seq.unassigned_blocks(), 1); // partial-tail block from div_ceil
+
+        // Decode 6 tokens
+        for i in 0..6u32 {
+            seq.schedule_decode(&manager).unwrap();
+            let outcome = seq.apply_decode(100 + i, &manager).unwrap();
+            if i < 5 {
+                match outcome {
+                    DecodeOutcome::Continue | DecodeOutcome::BlockCompleted => {}
+                    other => panic!("unexpected outcome at token {i}: {other:?}"),
+                }
+            } else {
+                // Last token (gen=7=max)
+                assert!(
+                    outcome == DecodeOutcome::MaxLength
+                        || outcome == DecodeOutcome::BlockCompletedAndMaxLength,
+                    "last token should hit max length, got: {outcome:?}"
+                );
+            }
+        }
+
+        assert!(seq.is_complete());
+        assert_eq!(seq.generated_tokens(), 7);
+        assert_eq!(seq.total_tokens(), 13);
+
+        seq.release().unwrap();
+        assert_eq!(seq.assigned_blocks(), 0);
+    }
+
+    #[test]
+    fn test_preempt_and_reacquire() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        // Prefill
+        seq.schedule_prefill(8, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+
+        // Decode 2 tokens
+        for _ in 0..2 {
+            seq.schedule_decode(&manager).unwrap();
+            seq.apply_decode(100, &manager).unwrap();
+        }
+        assert_eq!(seq.generated_tokens(), 3); // 1 from prefill + 2
+
+        // Preempt
+        seq.release().unwrap();
+        assert_eq!(seq.assigned_blocks(), 0);
+
+        // Reacquire (no gen block — allocated later by schedule_decode)
+        let success = seq.reacquire(&manager).unwrap();
+        assert!(success);
+        assert_eq!(seq.assigned_blocks(), 2);
+        assert_eq!(seq.unassigned_blocks(), 0); // no gen block from reacquire
+        assert_eq!(seq.generated_tokens(), 3);
+
+        // Continue decoding
+        seq.schedule_decode(&manager).unwrap();
+        let outcome = seq.apply_decode(200, &manager).unwrap();
+        assert_eq!(seq.generated_tokens(), 4);
+        let _ = outcome;
+    }
+
+    // =========================================================================
+    // Prefix matching
+    // =========================================================================
+
+    #[test]
+    fn test_match_and_add_prefix() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let tokens = make_tokens(8);
+
+        // Populate cache
+        let seq_for_populate = crate::BlockSequence::new(tokens[..4].to_vec(), BLOCK_SIZE, None);
+        let mutables = manager.allocate_blocks(1).unwrap();
+        let registered: Vec<_> = mutables
+            .into_iter()
+            .zip(seq_for_populate.blocks().iter())
+            .map(|(m, tb)| manager.register_block(m.complete(tb).unwrap()))
+            .collect();
+        drop(registered);
+
+        let mut seq = SchedulableSequence::<TestMeta>::new(tokens, 10, BLOCK_SIZE, noop_delegate());
+        let matched = seq.match_and_add_prefix(&manager).unwrap();
+        assert_eq!(matched, 1);
+        assert_eq!(seq.prefill_position(), 4); // 1 block * 4 tokens
+        assert_eq!(seq.kv_position(), 4); // cache hits have KV computed
+        assert_eq!(seq.assigned_blocks(), 1);
+
+        // Should only need 1 more input block (no gen block from prefill)
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+        assert_eq!(seq.assigned_blocks(), 2);
+        assert_eq!(seq.unassigned_blocks(), 0); // no gen block
+        assert_eq!(seq.tail_tokens(), 1);
+    }
+
+    #[test]
+    fn test_match_and_add_prefix_no_hits() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        let matched = seq.match_and_add_prefix(&manager).unwrap();
+        assert_eq!(matched, 0);
+        assert_eq!(seq.prefill_position(), 0);
+    }
+
+    // =========================================================================
+    // Edge cases
+    // =========================================================================
+
+    #[test]
+    fn test_empty_tokens_prefill() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = SchedulableSequence::<TestMeta>::new(vec![], 10, BLOCK_SIZE, noop_delegate());
+
+        assert!(seq.is_prefill_complete());
+
+        // Cannot schedule prefill (already complete)
+        let err = seq.schedule_prefill(0, &manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::PrefillComplete));
+
+        // Cannot schedule decode with 0 dangling tokens
+        let err = seq.schedule_decode(&manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::WrongDanglingCount { .. }));
+
+        // Append initial token to create dangling, then decode
+        seq.append_tokens(&[100]).unwrap();
+        assert_eq!(seq.tail_tokens(), 1);
+
+        seq.schedule_decode(&manager).unwrap();
+        let outcome = seq.apply_decode(101, &manager).unwrap();
+        assert_eq!(outcome, DecodeOutcome::Continue);
+    }
+
+    #[test]
+    fn test_zero_max_output_no_gen_block() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(4), 0, BLOCK_SIZE, noop_delegate());
+
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(None, &manager).unwrap();
+
+        assert_eq!(seq.assigned_blocks(), 1);
+        assert_eq!(seq.unassigned_blocks(), 0); // no gen block
+
+        // Can't decode since is_complete
+        let err = seq.schedule_decode(&manager).unwrap_err();
+        assert!(matches!(err, ScheduleError::GenerationComplete { .. }));
+    }
+
+    #[test]
+    fn test_debug_impl() {
+        let seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+        let debug_str = format!("{seq:?}");
+        assert!(debug_str.contains("SchedulableSequence"));
+        assert!(debug_str.contains("Idle"));
+    }
+
+    #[test]
+    fn test_revert_idle_rejected() {
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+        let err = seq.revert_schedule().unwrap_err();
+        assert!(matches!(err, ApplyError::WrongState { .. }));
+    }
+
+    #[test]
+    fn test_release_while_scheduled_rejected() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        seq.schedule_prefill(4, &manager).unwrap();
+        let err = seq.release().unwrap_err();
+        assert!(matches!(err, ApplyError::WrongState { .. }));
+    }
+
+    #[test]
+    fn test_reacquire_while_scheduled_rejected() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 10, BLOCK_SIZE, noop_delegate());
+
+        seq.schedule_prefill(4, &manager).unwrap();
+        let err = seq.reacquire(&manager).unwrap_err();
+        assert!(matches!(err, ApplyError::WrongState { .. }));
+    }
+
+    // =========================================================================
+    // Dangling token tracking
+    // =========================================================================
+
+    #[test]
+    fn test_dangling_tokens_tracking() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 20, BLOCK_SIZE, noop_delegate());
+
+        // Before prefill: all tokens are "not yet computed"
+        assert_eq!(seq.kv_position(), 0);
+        assert_eq!(seq.tail_tokens(), 8);
+
+        // After prefill with token: 1 dangling
+        seq.schedule_prefill(8, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+        assert_eq!(seq.kv_position(), 8);
+        assert_eq!(seq.tail_tokens(), 1);
+
+        // After decode: still 1 dangling
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(100, &manager).unwrap();
+        assert_eq!(seq.kv_position(), 9);
+        assert_eq!(seq.tail_tokens(), 1);
+
+        // After speculative (accept 2): still 1 dangling
+        seq.schedule_speculative(3, &manager).unwrap();
+        seq.apply_speculative(&[200, 201], &manager).unwrap();
+        assert_eq!(seq.kv_position(), 11);
+        assert_eq!(seq.tail_tokens(), 1);
+    }
+
+    #[test]
+    fn test_decode_requires_one_dangling() {
+        let manager = create_test_manager::<TestMeta>(20);
+
+        // 0 dangling: empty sequence with no tokens
+        let mut seq = SchedulableSequence::<TestMeta>::new(vec![], 10, BLOCK_SIZE, noop_delegate());
+        let err = seq.schedule_decode(&manager).unwrap_err();
+        assert!(matches!(
+            err,
+            ScheduleError::WrongDanglingCount {
+                expected: 1,
+                actual: 0,
+            }
+        ));
+
+        // 2 dangling: append 2 tokens
+        seq.append_tokens(&[100, 101]).unwrap();
+        assert_eq!(seq.tail_tokens(), 2);
+        let err = seq.schedule_decode(&manager).unwrap_err();
+        assert!(matches!(
+            err,
+            ScheduleError::WrongDanglingCount {
+                expected: 1,
+                actual: 2,
+            }
+        ));
+    }
+
+    #[test]
+    fn test_append_tokens_creates_dangling() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = SchedulableSequence::<TestMeta>::new(vec![], 10, BLOCK_SIZE, noop_delegate());
+
+        assert_eq!(seq.tail_tokens(), 0);
+
+        seq.append_tokens(&[100]).unwrap();
+        assert_eq!(seq.tail_tokens(), 1);
+        assert_eq!(seq.total_tokens(), 1);
+        assert_eq!(seq.kv_position(), 0);
+
+        // Now can schedule decode
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(101, &manager).unwrap();
+        assert_eq!(seq.tail_tokens(), 1);
+        assert_eq!(seq.kv_position(), 1);
+    }
+
+    #[test]
+    fn test_append_tokens_exceeding_remaining_returns_error_without_mutation() {
+        let mut seq = SchedulableSequence::<TestMeta>::new(vec![], 1, BLOCK_SIZE, noop_delegate());
+
+        let err = seq.append_tokens(&[100, 101]).unwrap_err();
+
+        assert!(matches!(
+            err,
+            ApplyError::AppendExceedsRemaining {
+                requested: 2,
+                remaining: 1,
+            }
+        ));
+        assert_eq!(seq.generated_tokens(), 0);
+        assert_eq!(seq.remaining_tokens(), 1);
+        assert_eq!(seq.total_tokens(), 0);
+        assert_eq!(seq.tail_tokens(), 0);
+        assert_eq!(seq.kv_position(), 0);
+    }
+
+    #[test]
+    fn test_kv_position_through_lifecycle() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(8), 20, BLOCK_SIZE, noop_delegate());
+
+        assert_eq!(seq.kv_position(), 0);
+
+        // Chunked prefill
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(None, &manager).unwrap();
+        assert_eq!(seq.kv_position(), 4);
+
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+        assert_eq!(seq.kv_position(), 8);
+
+        // Decode
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(100, &manager).unwrap();
+        assert_eq!(seq.kv_position(), 9);
+
+        // Speculative
+        seq.schedule_speculative(3, &manager).unwrap();
+        seq.apply_speculative(&[200, 201, 202], &manager).unwrap();
+        assert_eq!(seq.kv_position(), 12);
+    }
+
+    #[test]
+    fn test_pending_completion_staged_during_decode() {
+        let manager = create_test_manager::<TestMeta>(20);
+        // 7 input tokens: block 0 complete (0-3), block 1 partial (4-6)
+        let mut seq =
+            SchedulableSequence::<TestMeta>::new(make_tokens(7), 10, BLOCK_SIZE, noop_delegate());
+
+        seq.schedule_prefill(7, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+
+        // After prefill: token 1000 completed block 1 (4,5,6,1000) but NOT registered
+        assert_eq!(seq.assigned_blocks(), 1); // only block 0
+        assert_eq!(seq.unassigned_blocks(), 1); // partial-tail block from div_ceil
+        assert_eq!(seq.kv_position(), 7);
+        assert_eq!(seq.tail_tokens(), 1); // token 1000
+
+        // schedule_decode detects the pending completion
+        seq.schedule_decode(&manager).unwrap();
+        assert_eq!(
+            seq.state(),
+            SequenceState::DecodeScheduled {
+                blocks_allocated: 1
+            }
+        ); // 1 pending + 1 gen, but already had 1 unassigned
+
+        // apply_decode stages the pending block
+        let outcome = seq.apply_decode(100, &manager).unwrap();
+        assert_eq!(outcome, DecodeOutcome::Continue); // no boundary from THIS token
+        assert_eq!(seq.assigned_blocks(), 2); // block 1 now registered
+        assert_eq!(seq.unassigned_blocks(), 1); // gen block
+        assert_eq!(seq.kv_position(), 8);
+        assert_eq!(seq.tail_tokens(), 1);
+    }
+
+    // =========================================================================
+    // Builder
+    // =========================================================================
+
+    #[test]
+    fn test_builder_basic() {
+        let seq = SchedulableSequence::<TestMeta>::builder()
+            .tokens(make_tokens(8))
+            .max_output_tokens(10)
+            .block_size(BLOCK_SIZE)
+            .build::<TestMeta>()
+            .unwrap();
+
+        assert_eq!(seq.state(), SequenceState::Idle);
+        assert_eq!(seq.num_input_tokens(), 8);
+        assert_eq!(seq.max_output_tokens(), 10);
+        assert_eq!(seq.block_size(), BLOCK_SIZE as usize);
+    }
+
+    #[test]
+    fn test_builder_with_delegate() {
+        let delegate = Arc::new(CollectingDelegate::new());
+        let seq = SchedulableSequence::<TestMeta>::builder()
+            .tokens(make_tokens(4))
+            .max_output_tokens(5)
+            .block_size(BLOCK_SIZE)
+            .delegate(delegate.clone())
+            .build::<TestMeta>()
+            .unwrap();
+
+        assert_eq!(seq.num_input_tokens(), 4);
+
+        let events = delegate.events();
+        assert_eq!(events.len(), 1);
+        assert!(matches!(events[0], SequenceEvent::Created { .. }));
+    }
+
+    #[test]
+    fn test_builder_missing_required_field() {
+        let result = SchedulableSequence::<TestMeta>::builder()
+            .tokens(make_tokens(4))
+            // missing max_output_tokens and block_size
+            .build::<TestMeta>();
+
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_builder_default_noop_delegate() {
+        let manager = create_test_manager::<TestMeta>(20);
+        let mut seq = SchedulableSequence::<TestMeta>::builder()
+            .tokens(make_tokens(4))
+            .max_output_tokens(10)
+            .block_size(BLOCK_SIZE)
+            .build::<TestMeta>()
+            .unwrap();
+
+        // Verify the noop delegate doesn't panic — exercise the full lifecycle
+        seq.schedule_prefill(4, &manager).unwrap();
+        seq.apply_prefill(Some(1000), &manager).unwrap();
+        seq.schedule_decode(&manager).unwrap();
+        seq.apply_decode(100, &manager).unwrap();
+        seq.release().unwrap();
+    }
+}

lib/kvbm-logical/src/lib.rs

 // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-#![doc = include_str!("../README.md")]
+//! Logical block lifecycle management for KVBM.
+//!
+//! This crate provides the core block lifecycle system:
+//! - Type-safe state transitions (Reset -> Complete -> Registered)
+//! - Block registry with deduplication and attachments
+//! - Active/inactive/reset pool management
+//! - Event pipeline for distributed coordination
+//! - Block manager orchestration
 
 pub mod blocks;
 pub mod events;
+pub mod integrations;
 pub mod manager;
 pub mod metrics;
 pub mod pools;
 pub mod pubsub;
 pub mod registry;
+pub mod sequence;
 pub mod tinylfu;
 
 #[cfg(any(test, feature = "testing"))]
 pub mod testing;
 
+use bincode::{Decode, Encode};
+use serde::{Deserialize, Serialize};
+
 // Re-export common types and traits
 pub use blocks::{
     BlockError, BlockMetadata, CompleteBlock, ImmutableBlock, MutableBlock, WeakBlock,
 };
+pub use integrations::{
+    ApplyError, DecodeOutcome, NoopDelegate, RequestSequence, SchedulableSequence,
+    SchedulableSequenceBuilder, ScheduleError, SequenceDelegate, SequenceEvent, SequenceState,
+};
 pub use manager::BlockManager;
 pub use registry::BlockRegistry;
+pub use sequence::{
+    BlockSequence, BlockSequenceError, ExternalBlockAssignments, LogicalBlockAssignmentError,
+    LogicalBlockAssignments, zip_assigned, zip_assigned_pending,
+};
 
 pub type BlockId = usize;
 pub type SequenceHash = dynamo_tokens::PositionalLineageHash;
@@ -34,3 +54,19 @@ impl KvbmSequenceHashProvider for dynamo_tokens::TokenBlock {
         self.positional_lineage_hash()
     }
 }
+
+/// Logical layout handle type encoding the layout ID.
+///
+/// KVBM manages G1, G2 and G3 layouts directly. G4 is managed by an external service.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Encode, Decode, Serialize, Deserialize)]
+pub enum LogicalLayoutHandle {
+    /// Representation of GPU / Device Memory
+    G1,
+    /// Representation of CPU / Host Memory
+    G2,
+    /// Representation of Disk Storage
+    G3,
+    /// Representation of Blocks held in an external service
+    /// outside the control of the KVBM system.
+    G4,
+}

lib/kvbm-logical/src/sequence/assignments/external.rs

+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::ops::Range;
+
+use crate::{BlockId, KvbmSequenceHashProvider, SequenceHash};
+use dynamo_tokens::TokenBlock;
+
+use super::super::store::BlockStore;
+use crate::sequence::BlockSequenceError;
+
+/// Per-tier block_id tracking with an offset into the sequence.
+///
+/// Maintains an ordered mapping of `BlockId` → `SequenceHash` for assigned blocks,
+/// a staging area for blocks whose hashes have been computed but not yet committed,
+/// plus a FIFO queue of block_ids waiting for assignment. Index `i` in the assigned
+/// map corresponds to sequence position `offset + i`.
+///
+/// The three-phase lifecycle is:
+/// - **Unassigned** — block_ids queued for assignment (no hash yet).
+/// - **Staged** — block_ids paired with their `SequenceHash` but not yet committed.
+/// - **Assigned** — committed `BlockId → SequenceHash` pairs in positional order.
+///
+/// Multiple `ExternalBlockAssignments` instances can operate on the same `&[TokenBlock]` at
+/// different offsets (multi-tier).
+pub struct ExternalBlockAssignments {
+    store: BlockStore<(), SequenceHash, SequenceHash>,
+
+    /// Starting position in the sequence. Assignments begin at this position.
+    offset: usize,
+}
+
+impl std::fmt::Debug for ExternalBlockAssignments {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ExternalBlockAssignments")
+            .field("assigned_count", &self.store.assigned_count())
+            .field("staged_count", &self.store.staged_count())
+            .field("unassigned_count", &self.store.unassigned_count())
+            .field("offset", &self.offset)
+            .finish()
+    }
+}
+
+impl ExternalBlockAssignments {
+    /// Creates a new `ExternalBlockAssignments` starting at the given offset.
+    pub fn new(offset: usize) -> Self {
+        Self {
+            store: BlockStore::new(),
+            offset,
+        }
+    }
+
+    /// Returns the starting position in the sequence.
+    pub fn offset(&self) -> usize {
+        self.offset
+    }
+
+    /// Checks whether a block_id is known (assigned, staged, or unassigned).
+    pub fn contains(&self, block_id: &BlockId) -> bool {
+        self.store.contains(block_id)
+    }
+
+    /// Positional access: returns `(BlockId, SequenceHash)` at the given index
+    /// (relative to offset) in the assigned collection.
+    pub fn get_assigned(&self, index: usize) -> Option<(BlockId, SequenceHash)> {
+        self.store
+            .get_assigned(index)
+            .map(|(&id, &hash)| (id, hash))
+    }
+
+    /// Returns the number of assigned blocks.
+    pub fn assigned_count(&self) -> usize {
+        self.store.assigned_count()
+    }
+
+    /// Returns the number of staged blocks.
+    pub fn staged_count(&self) -> usize {
+        self.store.staged_count()
+    }
+
+    /// Returns the number of unassigned (pending) block_ids.
+    pub fn unassigned_count(&self) -> usize {
+        self.store.unassigned_count()
+    }
+
+    /// Positional access: returns `(BlockId, SequenceHash)` at the given index
+    /// (relative to the start of staged) in the staged collection.
+    pub fn get_staged(&self, index: usize) -> Option<(BlockId, SequenceHash)> {
+        self.store.get_staged(index).map(|(&id, &hash)| (id, hash))
+    }
+
+    /// Iterates over assigned blocks in positional order, yielding `(BlockId, SequenceHash)`.
+    pub fn assigned_iter(&self) -> impl Iterator<Item = (BlockId, SequenceHash)> + '_ {
+        self.store.assigned_iter().map(|(&id, &hash)| (id, hash))
+    }
+
+    /// Iterates over staged blocks in staging order, yielding `(BlockId, SequenceHash)`.
+    pub fn staged_iter(&self) -> impl Iterator<Item = (BlockId, SequenceHash)> + '_ {
+        self.store.staged_iter().map(|(&id, &hash)| (id, hash))
+    }
+
+    /// Iterates over unassigned block_ids in FIFO order.
+    pub fn unassigned_iter(&self) -> impl Iterator<Item = BlockId> + '_ {
+        self.store.unassigned_iter().map(|(&id, _)| id)
+    }
+
+    /// Clears all assigned, staged, and unassigned blocks, preserving the offset.
+    pub fn clear(&mut self) {
+        self.store.clear();
+    }
+
+    /// Takes all staged blocks, returning them as a `Vec`.
+    pub fn take_staged(&mut self) -> Vec<(BlockId, SequenceHash)> {
+        self.store.take_staged()
+    }
+
+    /// Returns the next sequence position to be assigned:
+    /// `offset + assigned_count + staged_count`.
+    pub fn next_position(&self) -> usize {
+        self.offset + self.store.assigned_count() + self.store.staged_count()
+    }
+
+    /// Absolute position range of assigned blocks: `offset..offset + assigned_count`.
+    pub fn assigned_positions(&self) -> Range<usize> {
+        self.offset..self.offset + self.store.assigned_count()
+    }
+
+    /// Absolute position range of staged blocks:
+    /// `offset + assigned_count .. offset + assigned_count + staged_count`.
+    pub fn staged_positions(&self) -> Range<usize> {
+        let start = self.offset + self.store.assigned_count();
+        start..start + self.store.staged_count()
+    }
+
+    /// Get the assigned `(BlockId, SequenceHash)` at an absolute sequence position.
+    ///
+    /// Returns `None` if `abs_pos` is outside [`assigned_positions()`](Self::assigned_positions).
+    pub fn get_at_position(&self, abs_pos: usize) -> Option<(BlockId, SequenceHash)> {
+        let relative = abs_pos.checked_sub(self.offset)?;
+        self.get_assigned(relative)
+    }
+
+    /// Absolute position range that pending (unassigned) blocks will occupy once
+    /// flushed: `next_position()..next_position() + unassigned_count()`.
+    pub fn pending_positions(&self) -> Range<usize> {
+        let start = self.next_position();
+        start..start + self.store.unassigned_count()
+    }
+
+    /// Get the pending `BlockId` at an absolute sequence position (FIFO order).
+    ///
+    /// Position `next_position()` maps to the first unassigned block,
+    /// `next_position() + 1` to the second, etc.
+    /// Returns `None` if `abs_pos` is outside [`pending_positions()`](Self::pending_positions).
+    pub fn get_pending_at_position(&self, abs_pos: usize) -> Option<BlockId> {
+        let start = self.next_position();
+        let relative = abs_pos.checked_sub(start)?;
+        self.store.get_unassigned(relative).map(|(&id, _)| id)
+    }
+
+    /// Add new block_ids to the unassigned queue.
+    ///
+    /// `block_ids` is the **full, ordered** list of block IDs allocated to this
+    /// assignment set. Known IDs (already in assigned, staged, or unassigned) must
+    /// form a contiguous prefix and are silently skipped. New IDs are appended to
+    /// the unassigned FIFO queue.
+    ///
+    /// This method does **not** assign blocks — call
+    /// [`assign_pending`](Self::assign_pending) to pair unassigned IDs with
+    /// available sequence blocks.
+    ///
+    /// # Block ID rules
+    ///
+    /// The list is partitioned into a **known prefix** and a **new suffix**:
+    ///
+    /// - **Known prefix** — IDs already present in `assigned`, `staged`, or
+    ///   `unassigned`. These are silently skipped. They must appear contiguously
+    ///   at the front of the list; interleaving a known ID after an unknown one
+    ///   is an [`OrderingViolation`](BlockSequenceError::OrderingViolation).
+    /// - **New suffix** — IDs not yet seen. These are appended (in order) to
+    ///   the unassigned FIFO queue.
+    ///
+    /// # Algorithm (two-phase, atomic)
+    ///
+    /// 1. **Validate & collect** — iterate `block_ids`. Known IDs must form a
+    ///    contiguous prefix (skip them). Unknown IDs are collected into a temp
+    ///    buffer. If a known ID appears after an unknown one →
+    ///    `OrderingViolation` error. No state is mutated until validation passes.
+    /// 2. **Commit** — push all new IDs to the unassigned queue.
+    pub fn extend_block_ids(
+        &mut self,
+        block_ids: impl IntoIterator<Item = BlockId>,
+    ) -> Result<(), BlockSequenceError> {
+        // Phase 1: Validate & collect
+        let mut new_ids = Vec::new();
+        let mut new_id_set = indexmap::IndexSet::new();
+        let mut first_new_index: Option<usize> = None;
+
+        for (i, id) in block_ids.into_iter().enumerate() {
+            if self.contains(&id) {
+                // Known ID — must come before any new IDs
+                if let Some(first_new) = first_new_index {
+                    return Err(BlockSequenceError::OrderingViolation {
+                        known_id: id,
+                        new_id: new_ids[0],
+                        known_index: i,
+                        first_new_index: first_new,
+                    });
+                }
+                // Skip — already known
+            } else {
+                // Unknown ID — collect, rejecting internal duplicates
+                if !new_id_set.insert(id) {
+                    return Err(BlockSequenceError::DuplicateBlockId { block_id: id });
+                }
+                if first_new_index.is_none() {
+                    first_new_index = Some(i);
+                }
+                new_ids.push(id);
+            }
+        }
+
+        // Phase 2: Commit — no errors from here on
+        for id in new_ids {
+            self.store.insert_unassigned(id, ());
+        }
+
+        Ok(())
+    }
+
+    /// Inserts pre-matched `(BlockId, SequenceHash)` pairs directly into the
+    /// assigned collection.
+    ///
+    /// This is the entry point for blocks whose hashes are already known (e.g.
+    /// cache hits). Two-phase atomic: collects all items, validates no duplicate
+    /// BlockIds across all three collections, then commits to assigned.
+    pub fn extend_assigned(
+        &mut self,
+        items: impl IntoIterator<Item = (BlockId, SequenceHash)>,
+    ) -> Result<usize, BlockSequenceError> {
+        let items: Vec<(BlockId, SequenceHash)> = items.into_iter().collect();
+
+        if let Err(block_id) = self
+            .store
+            .validate_no_duplicates(items.iter().map(|(id, _)| *id), items.len())
+        {
+            return Err(BlockSequenceError::DuplicateBlockId { block_id });
+        }
+
+        let count = items.len();
+        for (id, hash) in items {
+            self.store.insert_assigned(id, hash);
+        }
+
+        Ok(count)
+    }
+
+    /// FIFO drain from unassigned into staged, pairing each block_id with the
+    /// sequence hash from the corresponding `TokenBlock`.
+    ///
+    /// Staging starts at `sequence_blocks[self.next_position()]` and proceeds
+    /// forward, consuming one unassigned ID per available block. The loop stops
+    /// when either the unassigned queue is empty or there are no more sequence
+    /// blocks.
+    ///
+    /// Returns the range of newly staged indices (relative to the start of the
+    /// staged collection before this call).
+    ///
+    /// Each staged pair is validated: the position embedded in the block's
+    /// `kvbm_sequence_hash()` must equal the expected sequence index.
+    /// A mismatch returns [`BlockSequenceError::PositionMismatch`].
+    pub fn stage_pending(
+        &mut self,
+        sequence_blocks: &[TokenBlock],
+    ) -> Result<Range<usize>, BlockSequenceError> {
+        let staged_start_idx = self.store.staged_count();
+        let start_pos = self.next_position();
+
+        // How many sequence blocks are available starting from our next position?
+        let available_blocks = sequence_blocks.len().saturating_sub(start_pos);
+
+        // How many can we stage? Min of available blocks and unassigned count.
+        let to_stage = available_blocks.min(self.store.unassigned_count());
+
+        // Phase 1: Validate all positions before mutating
+        for i in 0..to_stage {
+            let seq_pos = start_pos + i;
+            let block = &sequence_blocks[seq_pos];
+            let hash = block.kvbm_sequence_hash();
+
+            let actual_pos = hash.position();
+            if actual_pos != seq_pos as u64 {
+                let block_id = self.store.get_unassigned(i).map(|(&id, _)| id).unwrap();
+                return Err(BlockSequenceError::PositionMismatch {
+                    expected: seq_pos,
+                    actual: actual_pos,
+                    block_id,
+                });
+            }
+        }
+
+        // Phase 2: Commit — no errors from here on
+        for i in 0..to_stage {
+            let seq_pos = start_pos + i;
+            let hash = sequence_blocks[seq_pos].kvbm_sequence_hash();
+            let (block_id, _) = self.store.shift_unassigned().unwrap();
+            self.store.insert_staged(block_id, hash);
+        }
+
+        let staged_end_idx = self.store.staged_count();
+        Ok(staged_start_idx..staged_end_idx)
+    }
+
+    /// Moves all staged blocks into assigned (infallible).
+    ///
+    /// Returns the range of newly assigned indices (relative to the start of
+    /// the assigned collection before this call).
+    pub fn commit_staged(&mut self) -> Range<usize> {
+        let start_idx = self.store.assigned_count();
+
+        while let Some((block_id, hash)) = self.store.shift_staged() {
+            self.store.insert_assigned(block_id, hash);
+        }
+
+        let end_idx = self.store.assigned_count();
+        start_idx..end_idx
+    }
+
+    /// Drain the unassigned FIFO queue into assigned, pairing each block_id
+    /// with the next available `TokenBlock` starting at `next_position()`.
+    ///
+    /// This is a convenience method equivalent to calling
+    /// [`stage_pending`](Self::stage_pending) followed by
+    /// [`commit_staged`](Self::commit_staged).
+    ///
+    /// Returns the range of newly assigned indices (relative to offset).
+    /// An empty range means no new assignments were made.
+    pub fn assign_pending(
+        &mut self,
+        sequence_blocks: &[TokenBlock],
+    ) -> Result<Range<usize>, BlockSequenceError> {
+        self.stage_pending(sequence_blocks)?;
+        Ok(self.commit_staged())
+    }
+}
+
+/// Zip two [`ExternalBlockAssignments`] over their overlapping assigned positions.
+///
+/// For each absolute position where **both** `a` and `b` have assigned blocks,
+/// yields `(position, a_block_id, b_block_id)`.
+///
+/// Iteration order: ascending position.
+pub fn zip_assigned(
+    a: &ExternalBlockAssignments,
+    b: &ExternalBlockAssignments,
+) -> Vec<(usize, BlockId, BlockId)> {
+    let a_range = a.assigned_positions();
+    let b_range = b.assigned_positions();
+    let start = a_range.start.max(b_range.start);
+    let end = a_range.end.min(b_range.end);
+
+    let mut result = Vec::new();
+    for pos in start..end {
+        // Both lookups are guaranteed to succeed within the intersection range.
+        let (a_id, _) = a.get_at_position(pos).unwrap();
+        let (b_id, _) = b.get_at_position(pos).unwrap();
+        result.push((pos, a_id, b_id));
+    }
+    result
+}
+
+/// Zip `src` assigned positions with `dst` pending positions.
+///
+/// For each absolute position where `src` has an assigned block and `dst`
+/// has a pending (unassigned) block, yields `(position, src_block_id, dst_block_id)`.
+///
+/// This is the onboard/offload planning primitive: the result tells you
+/// which source blocks to transfer into which destination blocks.
+pub fn zip_assigned_pending(
+    src: &ExternalBlockAssignments,
+    dst: &ExternalBlockAssignments,
+) -> Vec<(usize, BlockId, BlockId)> {
+    let src_range = src.assigned_positions();
+    let dst_range = dst.pending_positions();
+    let start = src_range.start.max(dst_range.start);
+    let end = src_range.end.min(dst_range.end);
+
+    let mut result = Vec::new();
+    for pos in start..end {
+        // Both lookups are guaranteed to succeed within the intersection range.
+        let (src_id, _) = src.get_at_position(pos).unwrap();
+        let dst_id = dst.get_pending_at_position(pos).unwrap();
+        result.push((pos, src_id, dst_id));
+    }
+    result
+}

lib/kvbm-logical/src/sequence/assignments/logical.rs

+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Manages physical block guards through three ordered collections,
+//! mirroring the block lifecycle state machine:
+//! `MutableBlock<T>` → `CompleteBlock<T>` → `ImmutableBlock<T>`.
+
+use crate::blocks::{BlockError, BlockMetadata, CompleteBlock, ImmutableBlock, MutableBlock};
+use crate::manager::BlockManager;
+use crate::{BlockId, SequenceHash};
+use dynamo_tokens::TokenBlock;
+
+use super::super::store::BlockStore;
+
+/// Error type for [`LogicalBlockAssignments`] operations.
+#[derive(Debug, thiserror::Error)]
+pub enum LogicalBlockAssignmentError<T: BlockMetadata> {
+    /// A mutable block_id in the input already exists in one of the three collections.
+    #[error("duplicate block_id {block_id} already present")]
+    DuplicateBlockId {
+        /// The first duplicate block_id detected.
+        block_id: BlockId,
+        /// All input blocks returned for recovery (no leaks).
+        blocks: Vec<MutableBlock<T>>,
+    },
+
+    /// An immutable block_id in the input already exists in one of the three collections.
+    #[error("duplicate block_id {block_id} already present")]
+    DuplicateAssignedBlockId {
+        /// The first duplicate block_id detected.
+        block_id: BlockId,
+        /// All input blocks returned for recovery (no leaks).
+        blocks: Vec<ImmutableBlock<T>>,
+    },
+
+    /// A matched block's sequence hash does not match the expected sequence hash.
+    #[error("sequence hash mismatch at position {position}: expected {expected}, got {actual}")]
+    SequenceHashMismatch {
+        /// The position in the sequence where the mismatch was detected.
+        position: usize,
+        /// The expected hash from the token sequence.
+        expected: SequenceHash,
+        /// The actual hash from the matched block.
+        actual: SequenceHash,
+        /// All input blocks returned for recovery (no leaks).
+        blocks: Vec<ImmutableBlock<T>>,
+    },
+}
+
+/// Manages the physical block guards (RAII types) through three ordered collections,
+/// mirroring the block lifecycle state machine:
+/// `MutableBlock<T>` → `CompleteBlock<T>` → `ImmutableBlock<T>`.
+///
+/// Provides the same ordered-collection semantics as
+/// [`ExternalBlockAssignments`](super::ExternalBlockAssignments) but at the guard level rather
+/// than the identity level.
+pub struct LogicalBlockAssignments<T: BlockMetadata> {
+    store: BlockStore<MutableBlock<T>, CompleteBlock<T>, ImmutableBlock<T>>,
+}
+
+impl<T: BlockMetadata> LogicalBlockAssignments<T> {
+    /// Creates an empty `LogicalBlockAssignments`.
+    pub fn new() -> Self {
+        Self {
+            store: BlockStore::new(),
+        }
+    }
+
+    // -- Counts & Queries --------------------------------------------------
+
+    /// Returns the number of assigned (registered/immutable) blocks.
+    pub fn assigned_count(&self) -> usize {
+        self.store.assigned_count()
+    }
+
+    /// Returns the number of staged (complete) blocks.
+    pub fn staged_count(&self) -> usize {
+        self.store.staged_count()
+    }
+
+    /// Returns the number of unassigned (mutable) blocks.
+    pub fn unassigned_count(&self) -> usize {
+        self.store.unassigned_count()
+    }
+
+    /// Returns `true` if all three collections are empty.
+    pub fn is_empty(&self) -> bool {
+        self.store.is_empty()
+    }
+
+    /// Checks whether a `BlockId` is present in any of the three collections.
+    pub fn contains(&self, block_id: &BlockId) -> bool {
+        self.store.contains(block_id)
+    }
+
+    // -- Positional Access -------------------------------------------------
+
+    /// Returns the assigned block at the given index (insertion order).
+    pub fn get_assigned(&self, index: usize) -> Option<(&BlockId, &ImmutableBlock<T>)> {
+        self.store.get_assigned(index)
+    }
+
+    /// Returns the staged block at the given index (staging order).
+    pub fn get_staged(&self, index: usize) -> Option<(&BlockId, &CompleteBlock<T>)> {
+        self.store.get_staged(index)
+    }
+
+    /// Returns the unassigned block at the given index (FIFO order).
+    pub fn get_unassigned(&self, index: usize) -> Option<(&BlockId, &MutableBlock<T>)> {
+        self.store.get_unassigned(index)
+    }
+
+    // -- Iteration ---------------------------------------------------------
+
+    /// Iterates over assigned blocks in positional order.
+    pub fn assigned_iter(&self) -> impl Iterator<Item = (&BlockId, &ImmutableBlock<T>)> {
+        self.store.assigned_iter()
+    }
+
+    /// Iterates over staged blocks in staging order.
+    pub fn staged_iter(&self) -> impl Iterator<Item = (&BlockId, &CompleteBlock<T>)> {
+        self.store.staged_iter()
+    }
+
+    /// Iterates over unassigned blocks in FIFO order.
+    pub fn unassigned_iter(&self) -> impl Iterator<Item = (&BlockId, &MutableBlock<T>)> {
+        self.store.unassigned_iter()
+    }
+
+    /// Iterates over all block IDs across all three collections in lifecycle
+    /// order: assigned → staged → unassigned.
+    pub fn all_block_ids(&self) -> impl Iterator<Item = &BlockId> {
+        self.store.all_block_ids()
+    }
+
+    // -- Mutation -----------------------------------------------------------
+
+    /// Adds mutable blocks to the unassigned queue.
+    ///
+    /// Two-phase atomic: collects all blocks, validates no duplicate BlockIds
+    /// across all three collections, then commits to unassigned. On error,
+    /// all blocks are returned in the error variant (no leaks).
+    pub fn extend_blocks(
+        &mut self,
+        blocks: impl IntoIterator<Item = MutableBlock<T>>,
+    ) -> Result<usize, LogicalBlockAssignmentError<T>> {
+        let blocks: Vec<MutableBlock<T>> = blocks.into_iter().collect();
+
+        if let Err(block_id) = self
+            .store
+            .validate_no_duplicates(blocks.iter().map(|b| b.block_id()), blocks.len())
+        {
+            return Err(LogicalBlockAssignmentError::DuplicateBlockId { block_id, blocks });
+        }
+
+        let count = blocks.len();
+        for block in blocks {
+            let id = block.block_id();
+            self.store.insert_unassigned(id, block);
+        }
+
+        Ok(count)
+    }
+
+    /// Inserts pre-matched immutable blocks directly into the assigned collection.
+    ///
+    /// This is the entry point for blocks retrieved via
+    /// [`BlockManager::match_blocks`] — blocks that already exist in the
+    /// manager's pools and can skip the unassigned → staged → assigned pipeline.
+    ///
+    /// Two-phase atomic: collects all blocks, validates no duplicate BlockIds
+    /// across all three collections, then commits to assigned. On error,
+    /// all blocks are returned in the error variant (no leaks).
+    pub fn extend_assigned(
+        &mut self,
+        blocks: impl IntoIterator<Item = ImmutableBlock<T>>,
+    ) -> Result<usize, LogicalBlockAssignmentError<T>> {
+        let blocks: Vec<ImmutableBlock<T>> = blocks.into_iter().collect();
+
+        if let Err(block_id) = self
+            .store
+            .validate_no_duplicates(blocks.iter().map(|b| b.block_id()), blocks.len())
+        {
+            return Err(LogicalBlockAssignmentError::DuplicateAssignedBlockId { block_id, blocks });
+        }
+
+        let count = blocks.len();
+        for block in blocks {
+            let id = block.block_id();
+            self.store.insert_assigned(id, block);
+        }
+
+        Ok(count)
+    }
+
+    /// FIFO drain from unassigned, completing each block with the corresponding
+    /// token block.
+    ///
+    /// Stages `min(sequence_blocks.len(), unassigned.len())` blocks. On
+    /// [`BlockError`], already-staged blocks remain in staged; the failed block
+    /// is returned in the error.
+    pub fn stage(
+        &mut self,
+        sequence_blocks: &[TokenBlock],
+    ) -> Result<usize, BlockError<MutableBlock<T>>> {
+        let to_stage = sequence_blocks.len().min(self.store.unassigned_count());
+
+        #[allow(clippy::needless_range_loop)]
+        for i in 0..to_stage {
+            let (block_id, mutable) = self.store.shift_unassigned().unwrap();
+            match mutable.complete(&sequence_blocks[i]) {
+                Ok(complete) => {
+                    self.store.insert_staged(block_id, complete);
+                }
+                Err(err) => {
+                    return Err(err);
+                }
+            }
+        }
+
+        Ok(to_stage)
+    }
+
+    /// Takes all staged blocks (FIFO order), registering each with the block
+    /// manager and moving them to assigned.
+    ///
+    /// Returns the number of blocks registered.
+    pub fn register(&mut self, manager: &BlockManager<T>) -> usize {
+        let count = self.store.staged_count();
+
+        while let Some((block_id, complete)) = self.store.shift_staged() {
+            let immutable = manager.register_block(complete);
+            self.store.insert_assigned(block_id, immutable);
+        }
+
+        count
+    }
+
+    /// LIFO-removes the last unassigned block.
+    pub fn pop_last_unassigned(&mut self) -> Option<(BlockId, MutableBlock<T>)> {
+        self.store.pop_unassigned()
+    }
+
+    /// Drops all guards across all three collections (RAII returns blocks to pools).
+    pub fn clear(&mut self) {
+        self.store.clear();
+    }
+
+    /// Takes all assigned blocks, returning them as a `Vec`.
+    pub fn take_assigned(&mut self) -> Vec<(BlockId, ImmutableBlock<T>)> {
+        self.store.take_assigned()
+    }
+
+    /// Takes all staged blocks, returning them as a `Vec`.
+    pub fn take_staged(&mut self) -> Vec<(BlockId, CompleteBlock<T>)> {
+        self.store.take_staged()
+    }
+
+    /// Takes all unassigned blocks, returning them as a `Vec`.
+    pub fn take_unassigned(&mut self) -> Vec<(BlockId, MutableBlock<T>)> {
+        self.store.take_unassigned()
+    }
+}
+
+impl<T: BlockMetadata> Default for LogicalBlockAssignments<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<T: BlockMetadata> std::fmt::Debug for LogicalBlockAssignments<T> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("LogicalBlockAssignments")
+            .field("assigned_count", &self.store.assigned_count())
+            .field("staged_count", &self.store.staged_count())
+            .field("unassigned_count", &self.store.unassigned_count())
+            .finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::sequence::BlockSequence;
+    use crate::testing::{TestMeta, create_test_manager};
+
+    const BLOCK_SIZE: u32 = 4;
+
+    fn create_sequence(num_blocks: usize) -> BlockSequence {
+        let total_tokens = num_blocks * BLOCK_SIZE as usize;
+        let tokens: Vec<u32> = (0..total_tokens as u32).collect();
+        BlockSequence::new(tokens, BLOCK_SIZE, None)
+    }
+
+    // =========================================================================
+    // Empty construction
+    // =========================================================================
+
+    #[test]
+    fn test_empty_construction() {
+        let la = LogicalBlockAssignments::<TestMeta>::new();
+        assert!(la.is_empty());
+        assert_eq!(la.assigned_count(), 0);
+        assert_eq!(la.staged_count(), 0);
+        assert_eq!(la.unassigned_count(), 0);
+        assert!(!la.contains(&0));
+    }
+
+    #[test]
+    fn test_default_is_empty() {
+        let la = LogicalBlockAssignments::<TestMeta>::default();
+        assert!(la.is_empty());
+    }
+
+    // =========================================================================
+    // extend_blocks
+    // =========================================================================
+
+    #[test]
+    fn test_extend_blocks_basic() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(5).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+
+        let mut la = LogicalBlockAssignments::new();
+        let count = la.extend_blocks(blocks).unwrap();
+
+        assert_eq!(count, 5);
+        assert_eq!(la.unassigned_count(), 5);
+        assert!(!la.is_empty());
+
+        // Verify FIFO ordering
+        for (i, expected_id) in ids.iter().enumerate() {
+            let (id, _) = la.get_unassigned(i).unwrap();
+            assert_eq!(id, expected_id);
+        }
+    }
+
+    // =========================================================================
+    // extend_assigned
+    // =========================================================================
+
+    /// Helper: allocate, complete, and register blocks through the manager,
+    /// returning `ImmutableBlock`s suitable for `extend_assigned`.
+    fn register_blocks_directly(
+        manager: &BlockManager<TestMeta>,
+        seq: &BlockSequence,
+        count: usize,
+    ) -> Vec<ImmutableBlock<TestMeta>> {
+        let mutables = manager.allocate_blocks(count).unwrap();
+        mutables
+            .into_iter()
+            .zip(seq.blocks().iter())
+            .map(|(m, tb)| manager.register_block(m.complete(tb).unwrap()))
+            .collect()
+    }
+
+    #[test]
+    fn test_extend_assigned_basic() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let seq = create_sequence(3);
+        let immutables = register_blocks_directly(&manager, &seq, 3);
+        let ids: Vec<BlockId> = immutables.iter().map(|b| b.block_id()).collect();
+        let hashes = seq.all_sequence_hashes();
+
+        let mut la = LogicalBlockAssignments::new();
+        let count = la.extend_assigned(immutables).unwrap();
+
+        assert_eq!(count, 3);
+        assert_eq!(la.assigned_count(), 3);
+        assert_eq!(la.staged_count(), 0);
+        assert_eq!(la.unassigned_count(), 0);
+
+        for (i, expected_id) in ids.iter().enumerate() {
+            let (id, imm) = la.get_assigned(i).unwrap();
+            assert_eq!(id, expected_id);
+            assert_eq!(imm.block_id(), *expected_id);
+            assert_eq!(imm.sequence_hash(), hashes[i]);
+        }
+    }
+
+    #[test]
+    fn test_extend_assigned_ordering_preserved() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let seq = create_sequence(5);
+        let immutables = register_blocks_directly(&manager, &seq, 5);
+        let ids: Vec<BlockId> = immutables.iter().map(|b| b.block_id()).collect();
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_assigned(immutables).unwrap();
+
+        // Verify insertion order matches input order
+        for (i, expected_id) in ids.iter().enumerate() {
+            let (id, _) = la.get_assigned(i).unwrap();
+            assert_eq!(id, expected_id);
+        }
+    }
+
+    #[test]
+    fn test_extend_assigned_contains() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let seq = create_sequence(3);
+        let immutables = register_blocks_directly(&manager, &seq, 3);
+        let ids: Vec<BlockId> = immutables.iter().map(|b| b.block_id()).collect();
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_assigned(immutables).unwrap();
+
+        for id in &ids {
+            assert!(la.contains(id));
+        }
+        assert!(!la.contains(&999));
+    }
+
+    #[test]
+    fn test_extend_assigned_then_extend_stage_register() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let seq = create_sequence(5);
+
+        // Register first 3 as matched prefix
+        let matched = register_blocks_directly(&manager, &seq, 3);
+        let matched_ids: Vec<BlockId> = matched.iter().map(|b| b.block_id()).collect();
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_assigned(matched).unwrap();
+        assert_eq!(la.assigned_count(), 3);
+
+        // Allocate new blocks for the remaining 2 positions
+        let new_blocks = manager.allocate_blocks(2).unwrap();
+        let new_ids: Vec<BlockId> = new_blocks.iter().map(|b| b.block_id()).collect();
+        la.extend_blocks(new_blocks).unwrap();
+
+        // Stage and register the new blocks
+        la.stage(&seq.blocks()[3..5]).unwrap();
+        la.register(&manager);
+
+        assert_eq!(la.assigned_count(), 5);
+        assert_eq!(la.staged_count(), 0);
+        assert_eq!(la.unassigned_count(), 0);
+
+        // First 3 are the matched prefix
+        for (i, expected_id) in matched_ids.iter().enumerate() {
+            let (id, _) = la.get_assigned(i).unwrap();
+            assert_eq!(id, expected_id);
+        }
+        // Last 2 are the newly registered blocks
+        for (i, expected_id) in new_ids.iter().enumerate() {
+            let (id, _) = la.get_assigned(3 + i).unwrap();
+            assert_eq!(id, expected_id);
+        }
+    }
+
+    #[test]
+    fn test_extend_assigned_with_match_blocks() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let seq = create_sequence(3);
+        let hashes = seq.all_sequence_hashes();
+
+        // Populate the manager: allocate → complete → register → drop
+        let mutables = manager.allocate_blocks(3).unwrap();
+        let registered: Vec<ImmutableBlock<TestMeta>> = mutables
+            .into_iter()
+            .zip(seq.blocks().iter())
+            .map(|(m, tb)| manager.register_block(m.complete(tb).unwrap()))
+            .collect();
+        drop(registered);
+
+        // match_blocks retrieves them from the manager's pools
+        let matched = manager.match_blocks(&hashes);
+        assert_eq!(matched.len(), 3);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_assigned(matched).unwrap();
+        assert_eq!(la.assigned_count(), 3);
+
+        // Verify hashes match
+        for (i, expected_hash) in hashes.iter().enumerate() {
+            let (_, imm) = la.get_assigned(i).unwrap();
+            assert_eq!(imm.sequence_hash(), *expected_hash);
+        }
+    }
+
+    #[test]
+    fn test_extend_assigned_empty() {
+        let mut la = LogicalBlockAssignments::<TestMeta>::new();
+        let count = la.extend_assigned(Vec::new()).unwrap();
+        assert_eq!(count, 0);
+        assert!(la.is_empty());
+    }
+
+    // =========================================================================
+    // stage
+    // =========================================================================
+
+    #[test]
+    fn test_stage_basic() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(3).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+        let seq = create_sequence(3);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        let staged = la.stage(seq.blocks()).unwrap();
+        assert_eq!(staged, 3);
+        assert_eq!(la.staged_count(), 3);
+        assert_eq!(la.unassigned_count(), 0);
+
+        // Verify FIFO ordering in staged
+        for (i, expected_id) in ids.iter().enumerate() {
+            let (id, _) = la.get_staged(i).unwrap();
+            assert_eq!(id, expected_id);
+        }
+    }
+
+    #[test]
+    fn test_stage_fifo_drain() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(5).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+        let seq = create_sequence(3);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        // Stage 3 out of 5
+        la.stage(seq.blocks()).unwrap();
+
+        // First 3 should have been drained from unassigned (FIFO)
+        assert_eq!(la.unassigned_count(), 2);
+        let (remaining_0, _) = la.get_unassigned(0).unwrap();
+        let (remaining_1, _) = la.get_unassigned(1).unwrap();
+        assert_eq!(*remaining_0, ids[3]);
+        assert_eq!(*remaining_1, ids[4]);
+    }
+
+    #[test]
+    fn test_stage_partial_fewer_blocks_than_unassigned() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(5).unwrap();
+        let seq = create_sequence(3);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        let staged = la.stage(seq.blocks()).unwrap();
+        assert_eq!(staged, 3);
+        assert_eq!(la.staged_count(), 3);
+        assert_eq!(la.unassigned_count(), 2);
+    }
+
+    #[test]
+    fn test_stage_partial_fewer_unassigned_than_blocks() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(2).unwrap();
+        let seq = create_sequence(5);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        let staged = la.stage(seq.blocks()).unwrap();
+        assert_eq!(staged, 2);
+        assert_eq!(la.staged_count(), 2);
+        assert_eq!(la.unassigned_count(), 0);
+    }
+
+    #[test]
+    fn test_stage_block_size_mismatch_recovery() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(3).unwrap();
+
+        // Sequence with block_size=8 (manager uses block_size=4)
+        let tokens: Vec<u32> = (0..24).collect();
+        let bad_seq = BlockSequence::new(tokens, 8, None);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        let result = la.stage(bad_seq.blocks());
+        assert!(result.is_err());
+
+        match result.unwrap_err() {
+            BlockError::BlockSizeMismatch {
+                expected,
+                actual,
+                block,
+            } => {
+                assert_eq!(expected, 4);
+                assert_eq!(actual, 8);
+                // Block recovered — drop to return to pool
+                drop(block);
+            }
+        }
+
+        // First block was removed from unassigned and returned in error
+        assert_eq!(la.staged_count(), 0);
+        assert_eq!(la.unassigned_count(), 2);
+    }
+
+    #[test]
+    fn test_stage_partial_then_mismatch() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(3).unwrap();
+        let good_seq = create_sequence(2);
+        let bad_tokens: Vec<u32> = (0..8).collect();
+        let bad_seq = BlockSequence::new(bad_tokens, 8, None);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        // Stage first 2 successfully
+        let staged = la.stage(good_seq.blocks()).unwrap();
+        assert_eq!(staged, 2);
+        assert_eq!(la.staged_count(), 2);
+        assert_eq!(la.unassigned_count(), 1);
+
+        // Try to stage 1 more with wrong size → error
+        let result = la.stage(bad_seq.blocks());
+        assert!(result.is_err());
+
+        // The 2 already-staged blocks remain
+        assert_eq!(la.staged_count(), 2);
+        // The failed block was removed from unassigned and returned in error
+        assert_eq!(la.unassigned_count(), 0);
+    }
+
+    // =========================================================================
+    // register
+    // =========================================================================
+
+    #[test]
+    fn test_register_basic() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(3).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+        let seq = create_sequence(3);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+        la.stage(seq.blocks()).unwrap();
+
+        let registered = la.register(&manager);
+        assert_eq!(registered, 3);
+        assert_eq!(la.assigned_count(), 3);
+        assert_eq!(la.staged_count(), 0);
+
+        // Verify ordering preserved
+        for (i, expected_id) in ids.iter().enumerate() {
+            let (id, imm) = la.get_assigned(i).unwrap();
+            assert_eq!(id, expected_id);
+            assert_eq!(imm.block_id(), *expected_id);
+        }
+    }
+
+    #[test]
+    fn test_register_empty_staged() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let mut la = LogicalBlockAssignments::<TestMeta>::new();
+
+        let registered = la.register(&manager);
+        assert_eq!(registered, 0);
+    }
+
+    // =========================================================================
+    // Full pipeline
+    // =========================================================================
+
+    #[test]
+    fn test_full_pipeline_extend_stage_register() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(5).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+        let seq = create_sequence(5);
+        let expected_hashes = seq.all_sequence_hashes();
+
+        let mut la = LogicalBlockAssignments::new();
+
+        // extend
+        la.extend_blocks(blocks).unwrap();
+        assert_eq!(la.unassigned_count(), 5);
+
+        // stage
+        la.stage(seq.blocks()).unwrap();
+        assert_eq!(la.staged_count(), 5);
+        assert_eq!(la.unassigned_count(), 0);
+
+        // register
+        la.register(&manager);
+        assert_eq!(la.assigned_count(), 5);
+        assert_eq!(la.staged_count(), 0);
+
+        // Verify correct blocks in correct order with correct hashes
+        for (i, expected_id) in ids.iter().enumerate() {
+            let (id, immutable) = la.get_assigned(i).unwrap();
+            assert_eq!(id, expected_id);
+            assert_eq!(immutable.block_id(), *expected_id);
+            assert_eq!(immutable.sequence_hash(), expected_hashes[i]);
+        }
+    }
+
+    #[test]
+    fn test_full_pipeline_incremental() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(6).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+        let seq = create_sequence(6);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        // Stage first 3, register them
+        la.stage(&seq.blocks()[..3]).unwrap();
+        la.register(&manager);
+        assert_eq!(la.assigned_count(), 3);
+        assert_eq!(la.unassigned_count(), 3);
+
+        // Stage remaining 3, register them
+        la.stage(&seq.blocks()[..3]).unwrap();
+        la.register(&manager);
+        assert_eq!(la.assigned_count(), 6);
+        assert_eq!(la.unassigned_count(), 0);
+
+        // Verify all 6 in order
+        for (i, expected_id) in ids.iter().enumerate() {
+            let (id, _) = la.get_assigned(i).unwrap();
+            assert_eq!(id, expected_id);
+        }
+    }
+
+    // =========================================================================
+    // contains across all three collections
+    // =========================================================================
+
+    #[test]
+    fn test_contains_across_collections() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(6).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+        let seq = create_sequence(4);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        // Stage first 2, then register them
+        la.stage(&seq.blocks()[..2]).unwrap();
+        la.register(&manager);
+
+        // Stage 2 more
+        la.stage(&seq.blocks()[..2]).unwrap();
+
+        // State: 2 assigned, 2 staged, 2 unassigned
+        assert_eq!(la.assigned_count(), 2);
+        assert_eq!(la.staged_count(), 2);
+        assert_eq!(la.unassigned_count(), 2);
+
+        // All 6 should be contained
+        for id in &ids {
+            assert!(la.contains(id), "block_id {id} should be contained");
+        }
+        assert!(!la.contains(&999));
+    }
+
+    // =========================================================================
+    // clear and drain
+    // =========================================================================
+
+    #[test]
+    fn test_clear() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(3).unwrap();
+        let seq = create_sequence(3);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+        la.stage(seq.blocks()).unwrap();
+        la.register(&manager);
+        assert_eq!(la.assigned_count(), 3);
+
+        la.clear();
+        assert!(la.is_empty());
+    }
+
+    #[test]
+    fn test_take_assigned() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(3).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+        let seq = create_sequence(3);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+        la.stage(seq.blocks()).unwrap();
+        la.register(&manager);
+
+        let drained = la.take_assigned();
+        assert_eq!(drained.len(), 3);
+        assert_eq!(la.assigned_count(), 0);
+
+        for (i, (id, _)) in drained.iter().enumerate() {
+            assert_eq!(*id, ids[i]);
+        }
+    }
+
+    #[test]
+    fn test_take_staged() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(3).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+        let seq = create_sequence(3);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+        la.stage(seq.blocks()).unwrap();
+
+        let drained = la.take_staged();
+        assert_eq!(drained.len(), 3);
+        assert_eq!(la.staged_count(), 0);
+
+        for (i, (id, _)) in drained.iter().enumerate() {
+            assert_eq!(*id, ids[i]);
+        }
+    }
+
+    #[test]
+    fn test_take_unassigned() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(3).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        let drained = la.take_unassigned();
+        assert_eq!(drained.len(), 3);
+        assert_eq!(la.unassigned_count(), 0);
+
+        for (i, (id, _)) in drained.iter().enumerate() {
+            assert_eq!(*id, ids[i]);
+        }
+    }
+
+    // =========================================================================
+    // Negative: extend_assigned duplicate detection
+    // =========================================================================
+
+    #[test]
+    fn test_extend_assigned_duplicate_already_in_assigned() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let seq = create_sequence(3);
+        let immutables = register_blocks_directly(&manager, &seq, 3);
+
+        // Clone before consuming — ImmutableBlock is Clone
+        let clones: Vec<ImmutableBlock<TestMeta>> = immutables.to_vec();
+        let dup_id = clones[0].block_id();
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_assigned(immutables).unwrap();
+
+        // Second call with same block_ids → duplicate error
+        let result = la.extend_assigned(clones);
+        assert!(result.is_err());
+        match result.unwrap_err() {
+            LogicalBlockAssignmentError::DuplicateAssignedBlockId { block_id, blocks } => {
+                assert_eq!(block_id, dup_id);
+                assert_eq!(blocks.len(), 3);
+            }
+            other => panic!("expected DuplicateAssignedBlockId, got: {other:?}"),
+        }
+
+        // Atomic rollback: assigned unchanged
+        assert_eq!(la.assigned_count(), 3);
+    }
+
+    #[test]
+    fn test_extend_assigned_duplicate_within_input_batch() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let seq = create_sequence(1);
+        let immutables = register_blocks_directly(&manager, &seq, 1);
+        let dup = immutables[0].clone();
+        let dup_id = dup.block_id();
+
+        // Two copies of the same block in one batch
+        let batch = vec![immutables.into_iter().next().unwrap(), dup];
+
+        let mut la = LogicalBlockAssignments::new();
+        let result = la.extend_assigned(batch);
+        assert!(result.is_err());
+        match result.unwrap_err() {
+            LogicalBlockAssignmentError::DuplicateAssignedBlockId { block_id, blocks } => {
+                assert_eq!(block_id, dup_id);
+                assert_eq!(blocks.len(), 2);
+            }
+            other => panic!("expected DuplicateAssignedBlockId, got: {other:?}"),
+        }
+
+        // Nothing committed
+        assert!(la.is_empty());
+    }
+
+    #[test]
+    fn test_extend_assigned_disjoint_from_staged() {
+        // Verifies that extend_assigned succeeds when the new block_ids
+        // are disjoint from those already in staged (no collision).
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(3).unwrap();
+        let ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+        let seq = create_sequence(3);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+        la.stage(seq.blocks()).unwrap();
+        // 3 blocks now in staged
+
+        let other_blocks = manager.allocate_blocks(3).unwrap();
+        let immutables: Vec<ImmutableBlock<TestMeta>> = other_blocks
+            .into_iter()
+            .zip(seq.blocks().iter())
+            .map(|(m, tb)| manager.register_block(m.complete(tb).unwrap()))
+            .collect();
+
+        // Different block_ids from the pool — no collision with staged
+        let other_ids: Vec<BlockId> = immutables.iter().map(|b| b.block_id()).collect();
+        assert!(ids.iter().all(|id| !other_ids.contains(id)));
+        la.extend_assigned(immutables).unwrap();
+        assert_eq!(la.assigned_count(), 3);
+        assert_eq!(la.staged_count(), 3);
+    }
+
+    #[test]
+    fn test_extend_assigned_disjoint_from_unassigned() {
+        // Verifies that extend_assigned succeeds when the new block_ids
+        // are disjoint from those already in unassigned (no collision).
+        let manager = create_test_manager::<TestMeta>(10);
+        let seq = create_sequence(3);
+
+        // Put blocks in unassigned
+        let blocks = manager.allocate_blocks(3).unwrap();
+        let unassigned_ids: Vec<BlockId> = blocks.iter().map(|b| b.block_id()).collect();
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        // Register separate blocks to get immutables
+        let other_blocks = manager.allocate_blocks(3).unwrap();
+        let immutables: Vec<ImmutableBlock<TestMeta>> = other_blocks
+            .into_iter()
+            .zip(seq.blocks().iter())
+            .map(|(m, tb)| manager.register_block(m.complete(tb).unwrap()))
+            .collect();
+
+        // Different block_ids from the pool — no collision with unassigned
+        let imm_ids: Vec<BlockId> = immutables.iter().map(|b| b.block_id()).collect();
+        assert!(unassigned_ids.iter().all(|id| !imm_ids.contains(id)));
+        la.extend_assigned(immutables).unwrap();
+        assert_eq!(la.assigned_count(), 3);
+        assert_eq!(la.unassigned_count(), 3);
+    }
+
+    // =========================================================================
+    // Negative: extend_blocks with block_id in staged or assigned
+    // =========================================================================
+
+    #[test]
+    fn test_extend_blocks_id_already_in_assigned() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let seq = create_sequence(3);
+
+        // Register 3 blocks → assigned
+        let blocks = manager.allocate_blocks(3).unwrap();
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+        la.stage(seq.blocks()).unwrap();
+        la.register(&manager);
+        assert_eq!(la.assigned_count(), 3);
+
+        // Those block_ids are now back in the reset pool (RAII from the original
+        // MutableBlocks was consumed by stage/register). Allocating again may
+        // return the same block_ids — but they're still in `assigned`.
+        // With 10 total blocks and 3 in assigned, 7 are free. The 3 assigned
+        // block_ids won't be re-allocated because their RAII guards live in
+        // `la.assigned`. So we can't hit this path with a single manager.
+        //
+        // Verified: extend_blocks can only collide via programming error
+        // (two managers), which is an invalid scenario.
+        assert_eq!(la.assigned_count(), 3);
+    }
+
+    // =========================================================================
+    // Negative: stage edge cases
+    // =========================================================================
+
+    #[test]
+    fn test_stage_empty_unassigned() {
+        let mut la = LogicalBlockAssignments::<TestMeta>::new();
+        let seq = create_sequence(3);
+
+        let staged = la.stage(seq.blocks()).unwrap();
+        assert_eq!(staged, 0);
+        assert_eq!(la.staged_count(), 0);
+    }
+
+    #[test]
+    fn test_stage_empty_sequence_blocks() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(3).unwrap();
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+
+        let staged = la.stage(&[]).unwrap();
+        assert_eq!(staged, 0);
+        assert_eq!(la.staged_count(), 0);
+        assert_eq!(la.unassigned_count(), 3);
+    }
+
+    #[test]
+    fn test_stage_both_empty() {
+        let mut la = LogicalBlockAssignments::<TestMeta>::new();
+
+        let staged = la.stage(&[]).unwrap();
+        assert_eq!(staged, 0);
+        assert!(la.is_empty());
+    }
+
+    // =========================================================================
+    // Negative: positional access out of bounds
+    // =========================================================================
+
+    #[test]
+    fn test_get_assigned_out_of_bounds() {
+        let la = LogicalBlockAssignments::<TestMeta>::new();
+        assert!(la.get_assigned(0).is_none());
+        assert!(la.get_assigned(100).is_none());
+    }
+
+    #[test]
+    fn test_get_staged_out_of_bounds() {
+        let la = LogicalBlockAssignments::<TestMeta>::new();
+        assert!(la.get_staged(0).is_none());
+        assert!(la.get_staged(100).is_none());
+    }
+
+    #[test]
+    fn test_get_unassigned_out_of_bounds() {
+        let la = LogicalBlockAssignments::<TestMeta>::new();
+        assert!(la.get_unassigned(0).is_none());
+        assert!(la.get_unassigned(100).is_none());
+    }
+
+    #[test]
+    fn test_get_out_of_bounds_with_populated_collections() {
+        let manager = create_test_manager::<TestMeta>(10);
+        let blocks = manager.allocate_blocks(2).unwrap();
+        let seq = create_sequence(2);
+
+        let mut la = LogicalBlockAssignments::new();
+        la.extend_blocks(blocks).unwrap();
+        la.stage(seq.blocks()).unwrap();
+        la.register(&manager);
+
+        assert!(la.get_assigned(0).is_some());
+        assert!(la.get_assigned(1).is_some());
+        assert!(la.get_assigned(2).is_none());
+        assert!(la.get_staged(0).is_none());
+        assert!(la.get_unassigned(0).is_none());
+    }
+
+    // =========================================================================
+    // Debug
+    // =========================================================================
+
+    #[test]
+    fn test_debug_impl() {
+        let la = LogicalBlockAssignments::<TestMeta>::new();
+        let debug_str = format!("{la:?}");
+        assert!(debug_str.contains("LogicalBlockAssignments"));
+    }
+}

lib/kvbm-logical/src/sequence/assignments/mod.rs

+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+mod external;
+mod logical;
+
+#[cfg(test)]
+mod tests;
+
+pub use external::{ExternalBlockAssignments, zip_assigned, zip_assigned_pending};
+pub use logical::{LogicalBlockAssignmentError, LogicalBlockAssignments};

lib/kvbm-logical/src/sequence/assignments/tests.rs

+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::sequence::{BlockSequence, BlockSequenceError};
+use crate::{BlockId, KvbmSequenceHashProvider, SequenceHash};
+
+use super::ExternalBlockAssignments;
+
+const TEST_BLOCK_SIZE: u32 = 4;
+
+/// Helper to create a BlockSequence with a given number of complete blocks and optional partial.
+fn create_test_sequence(num_complete_blocks: usize, partial_tokens: usize) -> BlockSequence {
+    let total_tokens = num_complete_blocks * TEST_BLOCK_SIZE as usize + partial_tokens;
+    let tokens: Vec<u32> = (0..total_tokens as u32).collect();
+    BlockSequence::new(tokens, TEST_BLOCK_SIZE, None)
+}
+
+/// Helper to get the expected sequence hashes from a BlockSequence.
+fn get_expected_hashes(seq: &BlockSequence) -> Vec<SequenceHash> {
+    seq.blocks()
+        .iter()
+        .map(|b| b.kvbm_sequence_hash())
+        .collect()
+}
+
+// =========================================================================
+// Test Cases: Aligned sequences (no partial block)
+// =========================================================================
+
+#[test]
+fn test_aligned_0_blocks_0_block_ids() {
+    let seq = create_test_sequence(0, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..0);
+    assert_eq!(assignments.assigned_count(), 0);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_aligned_1_block_0_block_ids() {
+    let seq = create_test_sequence(1, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..0);
+    assert_eq!(assignments.assigned_count(), 0);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_aligned_1_block_1_block_id() {
+    let seq = create_test_sequence(1, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100]).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..1);
+    assert_eq!(assignments.assigned_count(), 1);
+    let (id, hash) = assignments.get_assigned(0).unwrap();
+    assert_eq!(id, 100);
+    assert_eq!(hash, expected_hashes[0]);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_aligned_1_block_2_block_ids() {
+    let seq = create_test_sequence(1, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..1);
+    assert_eq!(assignments.assigned_count(), 1);
+    let (id, hash) = assignments.get_assigned(0).unwrap();
+    assert_eq!(id, 100);
+    assert_eq!(hash, expected_hashes[0]);
+    assert_eq!(assignments.unassigned_count(), 1);
+    assert_eq!(assignments.unassigned_iter().collect::<Vec<_>>(), vec![200]);
+}
+
+#[test]
+fn test_aligned_3_blocks_3_block_ids() {
+    let seq = create_test_sequence(3, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..3);
+    assert_eq!(assignments.assigned_count(), 3);
+    for (i, expected_hash) in expected_hashes.iter().enumerate() {
+        let (id, hash) = assignments.get_assigned(i).unwrap();
+        assert_eq!(id, (i + 1) * 100);
+        assert_eq!(hash, *expected_hash);
+    }
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_aligned_3_blocks_1_block_id() {
+    let seq = create_test_sequence(3, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100]).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..1);
+    assert_eq!(assignments.assigned_count(), 1);
+    let (id, hash) = assignments.get_assigned(0).unwrap();
+    assert_eq!(id, 100);
+    assert_eq!(hash, expected_hashes[0]);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_aligned_3_blocks_5_block_ids() {
+    let seq = create_test_sequence(3, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..3);
+    assert_eq!(assignments.assigned_count(), 3);
+    for (i, expected_hash) in expected_hashes.iter().enumerate() {
+        let (id, hash) = assignments.get_assigned(i).unwrap();
+        assert_eq!(id, (i + 1) * 100);
+        assert_eq!(hash, *expected_hash);
+    }
+    assert_eq!(assignments.unassigned_count(), 2);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![400, 500]
+    );
+}
+
+// =========================================================================
+// Test Cases: Sequences with partial (dangling) block
+// =========================================================================
+
+#[test]
+fn test_partial_0_complete_2_partial_0_block_ids() {
+    let seq = create_test_sequence(0, 2);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..0);
+    assert_eq!(assignments.assigned_count(), 0);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_partial_2_complete_1_partial_2_block_ids() {
+    let seq = create_test_sequence(2, 1);
+    let expected_hashes = get_expected_hashes(&seq);
+    assert_eq!(expected_hashes.len(), 2);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..2);
+    assert_eq!(assignments.assigned_count(), 2);
+    let (id0, hash0) = assignments.get_assigned(0).unwrap();
+    assert_eq!((id0, hash0), (100, expected_hashes[0]));
+    let (id1, hash1) = assignments.get_assigned(1).unwrap();
+    assert_eq!((id1, hash1), (200, expected_hashes[1]));
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_partial_2_complete_3_partial_4_block_ids() {
+    let seq = create_test_sequence(2, 3);
+    let expected_hashes = get_expected_hashes(&seq);
+    assert_eq!(expected_hashes.len(), 2);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..2);
+    assert_eq!(assignments.assigned_count(), 2);
+    assert_eq!(assignments.unassigned_count(), 2);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400]
+    );
+}
+
+#[test]
+fn test_partial_3_complete_2_partial_1_block_id() {
+    let seq = create_test_sequence(3, 2);
+    let expected_hashes = get_expected_hashes(&seq);
+    assert_eq!(expected_hashes.len(), 3);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100]).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..1);
+    assert_eq!(assignments.assigned_count(), 1);
+    let (id, hash) = assignments.get_assigned(0).unwrap();
+    assert_eq!((id, hash), (100, expected_hashes[0]));
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+// =========================================================================
+// Test Cases: Multiple calls (incremental assignment)
+// =========================================================================
+
+#[test]
+fn test_incremental_assignment_aligned() {
+    let seq = create_test_sequence(4, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // First call: assign first 2 blocks
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    let range_1 = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_1, 0..2);
+    assert_eq!(assignments.assigned_count(), 2);
+
+    // Second call: assign next 2 blocks (100, 200 are known prefix → skipped)
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    let range_2 = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_2, 2..4);
+    assert_eq!(assignments.assigned_count(), 4);
+
+    for (i, expected_hash) in expected_hashes.iter().enumerate() {
+        let (id, hash) = assignments.get_assigned(i).unwrap();
+        assert_eq!(id, (i + 1) * 100);
+        assert_eq!(hash, *expected_hash);
+    }
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_incremental_assignment_with_excess() {
+    let seq = create_test_sequence(3, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // First call: assign first 2 blocks
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    let range_1 = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_1, 0..2);
+
+    // Second call: 100, 200 skipped, 300 assigned, 400, 500 unassigned
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+    let range_2 = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_2, 2..3);
+    assert_eq!(assignments.assigned_count(), 3);
+    let (id, hash) = assignments.get_assigned(2).unwrap();
+    assert_eq!((id, hash), (300, expected_hashes[2]));
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![400, 500]
+    );
+}
+
+#[test]
+fn test_incremental_assignment_partial_then_excess() {
+    let seq = create_test_sequence(2, 1);
+    let expected_hashes = get_expected_hashes(&seq);
+    assert_eq!(expected_hashes.len(), 2);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // First call: assign 1 block
+    assignments.extend_block_ids(vec![100]).unwrap();
+    let range_1 = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_1, 0..1);
+
+    // Second call: 100 skipped, 200 assigned, 300, 400 unassigned
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    let range_2 = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_2, 1..2);
+    let (id, hash) = assignments.get_assigned(1).unwrap();
+    assert_eq!((id, hash), (200, expected_hashes[1]));
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400]
+    );
+}
+
+#[test]
+fn test_all_blocks_already_assigned_extra_goes_to_unassigned() {
+    let seq = create_test_sequence(2, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign all blocks
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    let range_1 = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_1, 0..2);
+
+    // All new go to unassigned (100, 200 are skipped)
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    let range_2 = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_2, 2..2);
+    assert_eq!(assignments.assigned_count(), 2);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400]
+    );
+}
+
+// =========================================================================
+// Test Cases: Edge cases
+// =========================================================================
+
+#[test]
+fn test_empty_slot_receives_block_ids() {
+    let seq = create_test_sequence(0, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..0);
+    assert_eq!(assignments.assigned_count(), 0);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![100, 200, 300]
+    );
+}
+
+#[test]
+fn test_only_partial_tokens_receives_block_ids() {
+    let seq = create_test_sequence(0, 3);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..0);
+    assert_eq!(assignments.assigned_count(), 0);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![100, 200]
+    );
+}
+
+#[test]
+fn test_large_sequence_exact_match() {
+    let seq = create_test_sequence(10, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+    let block_ids: Vec<BlockId> = (0..10).map(|i| (i + 1) * 100).collect();
+
+    assignments.extend_block_ids(block_ids).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..10);
+    assert_eq!(assignments.assigned_count(), 10);
+    for (i, expected_hash) in expected_hashes.iter().enumerate() {
+        let (id, hash) = assignments.get_assigned(i).unwrap();
+        assert_eq!(id, (i + 1) * 100);
+        assert_eq!(hash, *expected_hash);
+    }
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_verify_hash_block_id_pairing_order() {
+    let seq = create_test_sequence(5, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments
+        .extend_block_ids(vec![999, 888, 777, 666, 555])
+        .unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..5);
+    assert_eq!(assignments.get_assigned(0).unwrap().0, 999);
+    assert_eq!(assignments.get_assigned(1).unwrap().0, 888);
+    assert_eq!(assignments.get_assigned(2).unwrap().0, 777);
+    assert_eq!(assignments.get_assigned(3).unwrap().0, 666);
+    assert_eq!(assignments.get_assigned(4).unwrap().0, 555);
+
+    for (i, expected_hash) in expected_hashes.iter().enumerate() {
+        assert_eq!(assignments.get_assigned(i).unwrap().1, *expected_hash);
+    }
+}
+
+// =========================================================================
+// Cartesian product test
+// =========================================================================
+
+#[test]
+fn test_cartesian_product_combinations() {
+    let num_blocks_options = [0, 1, 3, 5];
+    let partial_options = [0, 1, 3];
+
+    for &num_blocks in &num_blocks_options {
+        for &partial in &partial_options {
+            let seq = create_test_sequence(num_blocks, partial);
+            let expected_hashes = get_expected_hashes(&seq);
+            let available_blocks = expected_hashes.len();
+
+            // Test with 0 block_ids
+            {
+                let mut assignments = ExternalBlockAssignments::new(0);
+                let range = assignments.assign_pending(seq.blocks()).unwrap();
+                assert_eq!(range, 0..0);
+                assert_eq!(assignments.assigned_count(), 0);
+                assert_eq!(assignments.unassigned_count(), 0);
+            }
+
+            // Test with fewer block_ids than available blocks
+            if available_blocks > 1 {
+                let fewer = available_blocks / 2;
+                let block_ids: Vec<BlockId> = (0..fewer).collect();
+                let mut assignments = ExternalBlockAssignments::new(0);
+                assignments.extend_block_ids(block_ids).unwrap();
+                let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+                assert_eq!(range, 0..fewer);
+                assert_eq!(assignments.assigned_count(), fewer);
+                assert_eq!(assignments.unassigned_count(), 0);
+
+                for (i, expected_hash) in expected_hashes.iter().enumerate().take(fewer) {
+                    let (id, hash) = assignments.get_assigned(i).unwrap();
+                    assert_eq!(id, i);
+                    assert_eq!(hash, *expected_hash);
+                }
+            }
+
+            // Test with exact number of block_ids
+            if available_blocks > 0 {
+                let block_ids: Vec<BlockId> = (0..available_blocks).collect();
+                let mut assignments = ExternalBlockAssignments::new(0);
+                assignments.extend_block_ids(block_ids).unwrap();
+                let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+                assert_eq!(range, 0..available_blocks);
+                assert_eq!(assignments.assigned_count(), available_blocks);
+                assert_eq!(assignments.unassigned_count(), 0);
+            }
+
+            // Test with more block_ids than available blocks
+            {
+                let excess = 3;
+                let total_ids = available_blocks + excess;
+                let block_ids: Vec<BlockId> = (0..total_ids).collect();
+                let mut assignments = ExternalBlockAssignments::new(0);
+                assignments.extend_block_ids(block_ids).unwrap();
+                let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+                assert_eq!(range, 0..available_blocks);
+                assert_eq!(assignments.assigned_count(), available_blocks);
+                assert_eq!(assignments.unassigned_count(), excess);
+
+                let expected_unassigned: Vec<BlockId> = (available_blocks..total_ids).collect();
+                assert_eq!(
+                    assignments.unassigned_iter().collect::<Vec<_>>(),
+                    expected_unassigned
+                );
+            }
+        }
+    }
+}
+
+// =========================================================================
+// Test Cases: Previously unassigned blocks (FIFO behavior)
+// =========================================================================
+
+#[test]
+fn test_unassigned_blocks_applied_before_new_blocks() {
+    let mut seq = create_test_sequence(5, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign 5 blocks + 2 excess
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500, 600, 700])
+        .unwrap();
+    let range_1 = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_1, 0..5);
+    assert_eq!(assignments.assigned_count(), 5);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![600, 700]
+    );
+
+    // Add 2 more complete blocks via token extension
+    for token in 20..28u32 {
+        seq.append_token(token).unwrap();
+    }
+    assert_eq!(seq.blocks().len(), 7);
+
+    // Flush with new block_ids — unassigned (600, 700) should be assigned first
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500, 600, 700, 800, 900])
+        .unwrap();
+    let range_2 = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_2, 5..7);
+    assert_eq!(assignments.assigned_count(), 7);
+
+    // Verify 600, 700 were assigned to positions 5, 6
+    assert_eq!(assignments.get_assigned(5).unwrap().0, 600);
+    assert_eq!(assignments.get_assigned(6).unwrap().0, 700);
+
+    // 800, 900 should be unassigned
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![800, 900]
+    );
+}
+
+#[test]
+fn test_unassigned_blocks_with_new_blocks_all_assigned() {
+    let mut seq = create_test_sequence(3, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign 3 + 1 excess
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.unassigned_iter().collect::<Vec<_>>(), vec![400]);
+
+    // Add 3 more blocks
+    for token in 12..24u32 {
+        seq.append_token(token).unwrap();
+    }
+    assert_eq!(seq.blocks().len(), 6);
+
+    // 400 (unassigned) + 500, 600 (new) should all fit
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500, 600])
+        .unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 3..6);
+    assert_eq!(assignments.assigned_count(), 6);
+    assert_eq!(assignments.get_assigned(3).unwrap().0, 400);
+    assert_eq!(assignments.get_assigned(4).unwrap().0, 500);
+    assert_eq!(assignments.get_assigned(5).unwrap().0, 600);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_unassigned_blocks_no_new_space() {
+    let seq = create_test_sequence(2, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign 2 + 2 excess
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400]
+    );
+
+    // No new space — all new go to unassigned
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500, 600])
+        .unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 2..2);
+    assert_eq!(assignments.assigned_count(), 2);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400, 500, 600]
+    );
+}
+
+#[test]
+fn test_unassigned_blocks_partial_space() {
+    let mut seq = create_test_sequence(3, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign 3 + 2 excess
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![400, 500]
+    );
+
+    // Add 1 more block
+    for token in 12..16u32 {
+        seq.append_token(token).unwrap();
+    }
+    assert_eq!(seq.blocks().len(), 4);
+
+    // Only 1 spot: 400 assigned, 500+new remain unassigned
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500, 600, 700, 800])
+        .unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 3..4);
+    assert_eq!(assignments.get_assigned(3).unwrap().0, 400);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![500, 600, 700, 800]
+    );
+}
+
+#[test]
+fn test_multiple_rounds_of_unassigned_accumulation() {
+    let mut seq = create_test_sequence(2, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Round 1: 2 assigned, 2 unassigned
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400]
+    );
+
+    // Round 2: no space, add 2 more
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500, 600])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400, 500, 600]
+    );
+
+    // Round 3: still no space, add 1 more
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500, 600, 700])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400, 500, 600, 700]
+    );
+
+    // Now add space for 3 more blocks
+    for token in 8..20u32 {
+        seq.append_token(token).unwrap();
+    }
+    assert_eq!(seq.blocks().len(), 5);
+
+    // Flush — all IDs already known, first 3 unassigned (300, 400, 500) get assigned
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 2..5);
+    assert_eq!(assignments.assigned_count(), 5);
+    assert_eq!(assignments.get_assigned(2).unwrap().0, 300);
+    assert_eq!(assignments.get_assigned(3).unwrap().0, 400);
+    assert_eq!(assignments.get_assigned(4).unwrap().0, 500);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![600, 700]
+    );
+}
+
+#[test]
+fn test_unassigned_blocks_ordering_preserved() {
+    let mut seq = create_test_sequence(1, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // 10 assigned, rest unassigned in FIFO order
+    assignments
+        .extend_block_ids(vec![10, 20, 30, 40, 50, 60])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.get_assigned(0).unwrap().0, 10);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![20, 30, 40, 50, 60]
+    );
+
+    // Add 2 more blocks
+    for token in 4..12u32 {
+        seq.append_token(token).unwrap();
+    }
+
+    // Flush — all IDs already known, 20, 30 assigned (FIFO)
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.get_assigned(1).unwrap().0, 20);
+    assert_eq!(assignments.get_assigned(2).unwrap().0, 30);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![40, 50, 60]
+    );
+
+    // Add 1 more block + 1 new ID
+    for token in 12..16u32 {
+        seq.append_token(token).unwrap();
+    }
+    assignments
+        .extend_block_ids(vec![10, 20, 30, 40, 50, 60, 70])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+
+    // 40 should be assigned (FIFO), not 70
+    assert_eq!(assignments.get_assigned(3).unwrap().0, 40);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![50, 60, 70]
+    );
+}
+
+#[test]
+fn test_unassigned_blocks_with_partial_token_block() {
+    let mut seq = create_test_sequence(2, 2);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // 2 assigned, 2 unassigned
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400]
+    );
+
+    // Complete the partial block (need 2 more tokens)
+    seq.append_token(10).unwrap();
+    seq.append_token(11).unwrap();
+    assert_eq!(seq.blocks().len(), 3);
+
+    // 300 should be assigned first (FIFO), 500 is new
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 2..3);
+    assert_eq!(assignments.get_assigned(2).unwrap().0, 300);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![400, 500]
+    );
+}
+
+#[test]
+fn test_unassigned_blocks_exactly_fill_new_space() {
+    let mut seq = create_test_sequence(2, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // 2 assigned, 3 unassigned
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400, 500]
+    );
+
+    // Add exactly 3 more blocks
+    for token in 8..20u32 {
+        seq.append_token(token).unwrap();
+    }
+    assert_eq!(seq.blocks().len(), 5);
+
+    // Flush — all IDs already known, unassigned exactly fill the space
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 2..5);
+    assert_eq!(assignments.assigned_count(), 5);
+    assert_eq!(assignments.get_assigned(2).unwrap().0, 300);
+    assert_eq!(assignments.get_assigned(3).unwrap().0, 400);
+    assert_eq!(assignments.get_assigned(4).unwrap().0, 500);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_empty_unassigned_with_new_blocks() {
+    let mut seq = create_test_sequence(3, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign exactly the right number
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.unassigned_count(), 0);
+
+    // Add more space
+    for token in 12..16u32 {
+        seq.append_token(token).unwrap();
+    }
+
+    // New blocks with no previous unassigned
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 3..4);
+    assert_eq!(assignments.get_assigned(3).unwrap().0, 400);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+// =========================================================================
+// New tests: Prefix validation (OrderingViolation)
+// =========================================================================
+
+#[test]
+fn test_ordering_violation_known_after_new() {
+    let seq = create_test_sequence(3, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign first block
+    assignments.extend_block_ids(vec![100]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+
+    // Now try to extend with [200, 100] — 200 is new, then 100 is known → error
+    let result = assignments.extend_block_ids(vec![200, 100]);
+    assert!(result.is_err());
+    match result.unwrap_err() {
+        BlockSequenceError::OrderingViolation {
+            known_id,
+            new_id,
+            known_index,
+            first_new_index,
+        } => {
+            assert_eq!(known_id, 100);
+            assert_eq!(new_id, 200);
+            assert_eq!(known_index, 1);
+            assert_eq!(first_new_index, 0);
+        }
+        other => panic!("expected OrderingViolation, got: {other:?}"),
+    }
+
+    // Verify atomicity — state unchanged
+    assert_eq!(assignments.assigned_count(), 1);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_ordering_violation_no_state_change() {
+    let seq = create_test_sequence(5, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign first 2
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+
+    // Error case: [300, 100] — 300 new, then 100 known
+    let result = assignments.extend_block_ids(vec![300, 100]);
+    assert!(result.is_err());
+
+    // State should be exactly as before the failed call
+    assert_eq!(assignments.assigned_count(), 2);
+    assert_eq!(assignments.unassigned_count(), 0);
+    assert_eq!(assignments.get_assigned(0).unwrap().0, 100);
+    assert_eq!(assignments.get_assigned(1).unwrap().0, 200);
+}
+
+// =========================================================================
+// New tests: All duplicates (re-extending same IDs is a no-op)
+// =========================================================================
+
+#[test]
+fn test_all_duplicates_noop() {
+    let seq = create_test_sequence(3, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign all 3
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+
+    // Re-extend with same IDs — should be a no-op (all known, skipped)
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+
+    // No new unassigned means assign_pending produces an empty range
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 3..3); // Empty range, no new assignments
+    assert_eq!(assignments.assigned_count(), 3);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_duplicates_with_unassigned() {
+    let seq = create_test_sequence(2, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // 2 assigned + 2 unassigned
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+
+    // Re-extend with same — no-op (300, 400 are known in unassigned)
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 2..2);
+    assert_eq!(assignments.assigned_count(), 2);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![300, 400]
+    );
+}
+
+#[test]
+fn test_shorter_sequence_slice_does_not_shrink_assigned() {
+    let seq = create_test_sequence(5, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign 3 blocks using the full sequence
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.assigned_count(), 3);
+
+    // Re-extend with same IDs but a shorter slice — should be a no-op,
+    // not shrink assigned or move anything to unassigned.
+    let range = assignments.assign_pending(&seq.blocks()[..2]).unwrap();
+    assert_eq!(range, 3..3);
+    assert_eq!(assignments.assigned_count(), 3);
+    assert_eq!(assignments.unassigned_count(), 0);
+
+    // All three original assignments are intact
+    for (i, &expected_hash) in expected_hashes.iter().enumerate().take(3) {
+        let (id, hash) = assignments.get_assigned(i).unwrap();
+        assert_eq!(id, (i + 1) * 100);
+        assert_eq!(hash, expected_hash);
+    }
+}
+
+// =========================================================================
+// New tests: Offset
+// =========================================================================
+
+#[test]
+fn test_offset_assignments() {
+    let seq = create_test_sequence(10, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(5);
+
+    assert_eq!(assignments.offset(), 5);
+    assert_eq!(assignments.next_position(), 5);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    // Should assign to positions 5, 6, 7
+    assert_eq!(range, 0..3);
+    assert_eq!(assignments.assigned_count(), 3);
+    assert_eq!(assignments.next_position(), 8);
+
+    let (id0, hash0) = assignments.get_assigned(0).unwrap();
+    assert_eq!(id0, 100);
+    assert_eq!(hash0, expected_hashes[5]);
+
+    let (id1, hash1) = assignments.get_assigned(1).unwrap();
+    assert_eq!(id1, 200);
+    assert_eq!(hash1, expected_hashes[6]);
+
+    let (id2, hash2) = assignments.get_assigned(2).unwrap();
+    assert_eq!(id2, 300);
+    assert_eq!(hash2, expected_hashes[7]);
+}
+
+#[test]
+fn test_offset_with_excess() {
+    let seq = create_test_sequence(8, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(5);
+
+    // 3 blocks available at offset 5 (positions 5, 6, 7), 5 block_ids → 2 unassigned
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+
+    assert_eq!(range, 0..3);
+    assert_eq!(assignments.assigned_count(), 3);
+    assert_eq!(assignments.get_assigned(0).unwrap().1, expected_hashes[5]);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![400, 500]
+    );
+}
+
+// =========================================================================
+// New tests: Multi-instance (two ExternalBlockAssignments on same blocks)
+// =========================================================================
+
+#[test]
+fn test_multi_instance_different_offsets() {
+    let seq = create_test_sequence(10, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+
+    let mut g1 = ExternalBlockAssignments::new(0);
+    let mut g2 = ExternalBlockAssignments::new(5);
+
+    // G1 assigns positions 0..3
+    g1.extend_block_ids(vec![10, 20, 30]).unwrap();
+    let range_g1 = g1.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_g1, 0..3);
+
+    // G2 assigns positions 5..8
+    g2.extend_block_ids(vec![110, 120, 130]).unwrap();
+    let range_g2 = g2.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range_g2, 0..3);
+
+    // Verify G1 hashes
+    assert_eq!(g1.get_assigned(0).unwrap().1, expected_hashes[0]);
+    assert_eq!(g1.get_assigned(1).unwrap().1, expected_hashes[1]);
+    assert_eq!(g1.get_assigned(2).unwrap().1, expected_hashes[2]);
+
+    // Verify G2 hashes
+    assert_eq!(g2.get_assigned(0).unwrap().1, expected_hashes[5]);
+    assert_eq!(g2.get_assigned(1).unwrap().1, expected_hashes[6]);
+    assert_eq!(g2.get_assigned(2).unwrap().1, expected_hashes[7]);
+
+    // They are independent — no cross-interference
+    assert_eq!(g1.assigned_count(), 3);
+    assert_eq!(g2.assigned_count(), 3);
+}
+
+// =========================================================================
+// New tests: Token extension + flush
+// =========================================================================
+
+#[test]
+fn test_token_extension_then_flush() {
+    let mut seq = create_test_sequence(2, 2);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // 2 blocks available, 3 block_ids → 1 unassigned
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.assigned_count(), 2);
+    assert_eq!(assignments.unassigned_iter().collect::<Vec<_>>(), vec![300]);
+
+    // Complete the partial block
+    seq.append_token(10).unwrap();
+    seq.append_token(11).unwrap();
+    assert_eq!(seq.blocks().len(), 3);
+
+    // Flush — unassigned 300 gets assigned
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 2..3);
+    assert_eq!(assignments.assigned_count(), 3);
+    assert_eq!(assignments.get_assigned(2).unwrap().0, 300);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+#[test]
+fn test_extend_tokens_creates_new_blocks() {
+    let mut seq = create_test_sequence(1, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign 1 block + 2 excess
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.assigned_count(), 1);
+    assert_eq!(
+        assignments.unassigned_iter().collect::<Vec<_>>(),
+        vec![200, 300]
+    );
+
+    // Extend tokens to create 2 more blocks
+    let new_range = seq.extend_tokens((4..12).collect()).unwrap();
+    assert!(new_range.is_some());
+    assert_eq!(seq.blocks().len(), 3);
+
+    // Flush
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 1..3);
+    assert_eq!(assignments.assigned_count(), 3);
+    assert_eq!(assignments.get_assigned(1).unwrap().0, 200);
+    assert_eq!(assignments.get_assigned(2).unwrap().0, 300);
+    assert_eq!(assignments.unassigned_count(), 0);
+}
+
+// =========================================================================
+// New tests: clear()
+// =========================================================================
+
+#[test]
+fn test_clear_preserves_offset() {
+    let seq = create_test_sequence(5, 0);
+    let mut assignments = ExternalBlockAssignments::new(3);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.assigned_count(), 2); // positions 3, 4
+    assert_eq!(assignments.unassigned_count(), 1); // 300
+
+    assignments.clear();
+
+    assert_eq!(assignments.offset(), 3);
+    assert_eq!(assignments.assigned_count(), 0);
+    assert_eq!(assignments.unassigned_count(), 0);
+    assert_eq!(assignments.next_position(), 3);
+}
+
+// =========================================================================
+// New tests: Position mismatch
+// =========================================================================
+
+#[test]
+fn test_assign_with_nonzero_offset() {
+    // Create a sequence with blocks at positions 0, 1, 2.
+    // Use offset=1 so assignments start at sequence_blocks[1].
+    // Verifies that the position validation passes for valid data
+    // and offset-based indexing works correctly.
+    let seq = create_test_sequence(3, 0);
+
+    let mut assignments = ExternalBlockAssignments::new(1);
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    let range = assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 0..2);
+    // Position 1 and 2 should be assigned
+    assert_eq!(
+        assignments.get_assigned(0).unwrap().1,
+        get_expected_hashes(&seq)[1]
+    );
+    assert_eq!(
+        assignments.get_assigned(1).unwrap().1,
+        get_expected_hashes(&seq)[2]
+    );
+}
+
+// =========================================================================
+// New tests: contains()
+// =========================================================================
+
+#[test]
+fn test_contains() {
+    let seq = create_test_sequence(2, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+
+    // 100, 200 are assigned
+    assert!(assignments.contains(&100));
+    assert!(assignments.contains(&200));
+
+    // 300 is unassigned
+    assert!(assignments.contains(&300));
+
+    // 400 is unknown
+    assert!(!assignments.contains(&400));
+}
+
+// =========================================================================
+// New tests: assigned_iter()
+// =========================================================================
+
+#[test]
+fn test_assigned_iter() {
+    let seq = create_test_sequence(3, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+
+    let assigned: Vec<(BlockId, SequenceHash)> = assignments.assigned_iter().collect();
+    assert_eq!(assigned.len(), 3);
+    assert_eq!(assigned[0], (100, expected_hashes[0]));
+    assert_eq!(assigned[1], (200, expected_hashes[1]));
+    assert_eq!(assigned[2], (300, expected_hashes[2]));
+}
+
+// =========================================================================
+// New tests: BlockSequence
+// =========================================================================
+
+#[test]
+fn test_block_sequence_new() {
+    let seq = BlockSequence::new(vec![0, 1, 2, 3, 4, 5, 6, 7], 4, None);
+    assert_eq!(seq.blocks().len(), 2);
+    assert_eq!(seq.block_size(), 4);
+    assert_eq!(seq.total_tokens(), 8);
+}
+
+#[test]
+fn test_block_sequence_all_sequence_hashes() {
+    let seq = BlockSequence::new(vec![0, 1, 2, 3, 4, 5, 6, 7], 4, None);
+    let hashes = seq.all_sequence_hashes();
+    assert_eq!(hashes.len(), 2);
+    // Hashes should match what we'd get from blocks directly
+    let expected: Vec<_> = seq
+        .blocks()
+        .iter()
+        .map(|b| b.kvbm_sequence_hash())
+        .collect();
+    assert_eq!(hashes, expected);
+}
+
+#[test]
+fn test_block_sequence_extend_tokens() {
+    let mut seq = BlockSequence::new(vec![0, 1, 2, 3], 4, None);
+    assert_eq!(seq.blocks().len(), 1);
+
+    let result = seq.extend_tokens(vec![4, 5, 6, 7, 8, 9]).unwrap();
+    assert_eq!(result, Some(1..2)); // One more block completed
+    assert_eq!(seq.blocks().len(), 2);
+    assert_eq!(seq.total_tokens(), 10); // 10 total, 2 partial remaining
+}
+
+#[test]
+fn test_block_sequence_append_token() {
+    let mut seq = BlockSequence::new(vec![0, 1, 2], 4, None);
+    assert_eq!(seq.blocks().len(), 0);
+
+    // Doesn't complete a block
+    let result = seq.append_token(3).unwrap();
+    assert_eq!(result, Some(0)); // Completed first block
+    assert_eq!(seq.blocks().len(), 1);
+}
+
+#[test]
+fn test_block_sequence_with_salt() {
+    let seq1 = BlockSequence::new(vec![0, 1, 2, 3], 4, None);
+    let seq2 = BlockSequence::new(vec![0, 1, 2, 3], 4, Some(42));
+
+    // Different salts should produce different hashes
+    let hashes1 = seq1.all_sequence_hashes();
+    let hashes2 = seq2.all_sequence_hashes();
+    assert_ne!(hashes1, hashes2);
+}
+
+// =========================================================================
+// Positional access methods
+// =========================================================================
+
+#[test]
+fn test_assigned_positions_empty() {
+    let assignments = ExternalBlockAssignments::new(0);
+    assert_eq!(assignments.assigned_positions(), 0..0);
+}
+
+#[test]
+fn test_assigned_positions_with_offset() {
+    let seq = create_test_sequence(10, 0);
+    let mut assignments = ExternalBlockAssignments::new(3);
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.assigned_positions(), 3..5);
+}
+
+#[test]
+fn test_pending_positions_empty() {
+    let assignments = ExternalBlockAssignments::new(0);
+    assert_eq!(assignments.pending_positions(), 0..0);
+}
+
+#[test]
+fn test_pending_positions_with_assigned_and_unassigned() {
+    let seq = create_test_sequence(2, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    // 2 assigned (positions 0,1), 2 pending (positions 2,3)
+    assert_eq!(assignments.pending_positions(), 2..4);
+}
+
+#[test]
+fn test_pending_positions_with_offset() {
+    let seq = create_test_sequence(10, 0);
+    let mut assignments = ExternalBlockAssignments::new(5);
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    // offset=5, 4 block_ids, 5 blocks available at positions 5..10 → 4 assigned
+    // no pending
+    assert_eq!(assignments.assigned_positions(), 5..9);
+    assert_eq!(assignments.pending_positions(), 9..9);
+}
+
+#[test]
+fn test_get_at_position_in_range() {
+    let seq = create_test_sequence(5, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(2);
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    // Assigned at positions 2, 3, 4
+    let (id, hash) = assignments.get_at_position(2).unwrap();
+    assert_eq!(id, 100);
+    assert_eq!(hash, expected_hashes[2]);
+
+    let (id, hash) = assignments.get_at_position(4).unwrap();
+    assert_eq!(id, 300);
+    assert_eq!(hash, expected_hashes[4]);
+}
+
+#[test]
+fn test_get_at_position_out_of_range() {
+    let seq = create_test_sequence(5, 0);
+    let mut assignments = ExternalBlockAssignments::new(2);
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    // Assigned at positions 2, 3
+    assert!(assignments.get_at_position(0).is_none()); // before offset
+    assert!(assignments.get_at_position(1).is_none()); // before offset
+    assert!(assignments.get_at_position(4).is_none()); // past end
+    assert!(assignments.get_at_position(100).is_none()); // way past end
+}
+
+#[test]
+fn test_get_pending_at_position_in_range() {
+    let seq = create_test_sequence(2, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    // 2 assigned, 3 pending at positions 2, 3, 4
+    assert_eq!(assignments.get_pending_at_position(2), Some(300));
+    assert_eq!(assignments.get_pending_at_position(3), Some(400));
+    assert_eq!(assignments.get_pending_at_position(4), Some(500));
+}
+
+#[test]
+fn test_get_pending_at_position_out_of_range() {
+    let seq = create_test_sequence(2, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    // 2 assigned, 2 pending at positions 2, 3
+    assert!(assignments.get_pending_at_position(0).is_none()); // assigned range
+    assert!(assignments.get_pending_at_position(1).is_none()); // assigned range
+    assert!(assignments.get_pending_at_position(4).is_none()); // past end
+    assert!(assignments.get_pending_at_position(100).is_none());
+}
+
+// =========================================================================
+// zip_assigned tests
+// =========================================================================
+
+use super::{zip_assigned, zip_assigned_pending};
+
+#[test]
+fn test_zip_assigned_full_overlap() {
+    let seq = create_test_sequence(5, 0);
+    let mut a = ExternalBlockAssignments::new(0);
+    let mut b = ExternalBlockAssignments::new(0);
+    a.extend_block_ids(vec![10, 20, 30]).unwrap();
+    a.assign_pending(seq.blocks()).unwrap();
+    b.extend_block_ids(vec![110, 120, 130]).unwrap();
+    b.assign_pending(seq.blocks()).unwrap();
+
+    let pairs = zip_assigned(&a, &b);
+    assert_eq!(pairs, vec![(0, 10, 110), (1, 20, 120), (2, 30, 130)]);
+}
+
+#[test]
+fn test_zip_assigned_partial_overlap() {
+    let seq = create_test_sequence(10, 0);
+    let mut a = ExternalBlockAssignments::new(0);
+    let mut b = ExternalBlockAssignments::new(2);
+    a.extend_block_ids(vec![10, 20, 30, 40]).unwrap();
+    a.assign_pending(seq.blocks()).unwrap();
+    b.extend_block_ids(vec![110, 120, 130]).unwrap();
+    b.assign_pending(seq.blocks()).unwrap();
+    // a: positions 0..4, b: positions 2..5 → overlap 2..4
+    let pairs = zip_assigned(&a, &b);
+    assert_eq!(pairs, vec![(2, 30, 110), (3, 40, 120)]);
+}
+
+#[test]
+fn test_zip_assigned_no_overlap() {
+    let seq = create_test_sequence(10, 0);
+    let mut a = ExternalBlockAssignments::new(0);
+    let mut b = ExternalBlockAssignments::new(5);
+    a.extend_block_ids(vec![10, 20]).unwrap();
+    a.assign_pending(seq.blocks()).unwrap();
+    b.extend_block_ids(vec![110, 120]).unwrap();
+    b.assign_pending(seq.blocks()).unwrap();
+    // a: positions 0..2, b: positions 5..7 → no overlap
+    let pairs = zip_assigned(&a, &b);
+    assert!(pairs.is_empty());
+}
+
+#[test]
+fn test_zip_assigned_either_empty() {
+    let seq = create_test_sequence(5, 0);
+    let a = ExternalBlockAssignments::new(0);
+    let mut b = ExternalBlockAssignments::new(0);
+    b.extend_block_ids(vec![110, 120]).unwrap();
+    b.assign_pending(seq.blocks()).unwrap();
+
+    assert!(zip_assigned(&a, &b).is_empty());
+    assert!(zip_assigned(&b, &a).is_empty());
+}
+
+#[test]
+fn test_zip_assigned_both_empty() {
+    let a = ExternalBlockAssignments::new(0);
+    let b = ExternalBlockAssignments::new(0);
+    assert!(zip_assigned(&a, &b).is_empty());
+}
+
+// =========================================================================
+// zip_assigned_pending tests
+// =========================================================================
+
+#[test]
+fn test_zip_assigned_pending_full_overlap() {
+    let seq = create_test_sequence(5, 0);
+    let mut src = ExternalBlockAssignments::new(2);
+    let mut dst = ExternalBlockAssignments::new(0);
+
+    // src: assigned at positions 2, 3, 4
+    src.extend_block_ids(vec![10, 20, 30]).unwrap();
+    src.assign_pending(seq.blocks()).unwrap();
+
+    // dst: only sees first 2 blocks → 2 assigned at 0, 1 + 3 pending at 2, 3, 4
+    dst.extend_block_ids(vec![110, 120, 130, 140, 150]).unwrap();
+    dst.assign_pending(&seq.blocks()[..2]).unwrap();
+    assert_eq!(dst.assigned_count(), 2);
+    assert_eq!(dst.unassigned_count(), 3);
+
+    let pairs = zip_assigned_pending(&src, &dst);
+    assert_eq!(pairs, vec![(2, 10, 130), (3, 20, 140), (4, 30, 150)]);
+}
+
+#[test]
+fn test_zip_assigned_pending_partial_overlap() {
+    let seq = create_test_sequence(10, 0);
+    let mut src = ExternalBlockAssignments::new(2);
+    let mut dst = ExternalBlockAssignments::new(0);
+
+    // src: assigned at 2, 3, 4, 5
+    src.extend_block_ids(vec![10, 20, 30, 40]).unwrap();
+    src.assign_pending(seq.blocks()).unwrap();
+
+    // dst: only sees first 3 blocks → 3 assigned at 0, 1, 2 + 2 pending at 3, 4
+    dst.extend_block_ids(vec![110, 120, 130, 140, 150]).unwrap();
+    dst.assign_pending(&seq.blocks()[..3]).unwrap();
+    assert_eq!(dst.assigned_count(), 3);
+    assert_eq!(dst.unassigned_count(), 2);
+
+    // src assigned: 2..6, dst pending: 3..5 → overlap 3..5
+    let pairs = zip_assigned_pending(&src, &dst);
+    assert_eq!(pairs, vec![(3, 20, 140), (4, 30, 150)]);
+}
+
+#[test]
+fn test_zip_assigned_pending_no_overlap() {
+    let seq = create_test_sequence(10, 0);
+    let mut src = ExternalBlockAssignments::new(0);
+    let mut dst = ExternalBlockAssignments::new(0);
+
+    // src: assigned at 0, 1
+    src.extend_block_ids(vec![10, 20]).unwrap();
+    src.assign_pending(seq.blocks()).unwrap();
+
+    // dst: 5 assigned at 0..5 + 2 pending at 5, 6
+    dst.extend_block_ids(vec![110, 120, 130, 140, 150, 160, 170])
+        .unwrap();
+    dst.assign_pending(seq.blocks()).unwrap();
+
+    // src assigned: 0..2, dst pending: 5..7 → no overlap
+    let pairs = zip_assigned_pending(&src, &dst);
+    assert!(pairs.is_empty());
+}
+
+#[test]
+fn test_zip_assigned_pending_either_empty() {
+    let seq = create_test_sequence(5, 0);
+    let src = ExternalBlockAssignments::new(0); // no assigned
+    let mut dst = ExternalBlockAssignments::new(0);
+    dst.extend_block_ids(vec![110, 120, 130]).unwrap();
+    dst.assign_pending(seq.blocks()).unwrap();
+
+    // src has no assigned blocks
+    assert!(zip_assigned_pending(&src, &dst).is_empty());
+
+    // dst has no pending blocks
+    let mut src2 = ExternalBlockAssignments::new(0);
+    src2.extend_block_ids(vec![10, 20]).unwrap();
+    src2.assign_pending(seq.blocks()).unwrap();
+    let dst2 = ExternalBlockAssignments::new(0);
+    assert!(zip_assigned_pending(&src2, &dst2).is_empty());
+}
+
+// =========================================================================
+// Onboard scenario test (G2 → G1 transfer)
+// =========================================================================
+
+#[test]
+fn test_onboard_scenario_g2_to_g1() {
+    // Sequence: 5 complete blocks (positions 0..5)
+    let seq = create_test_sequence(5, 0);
+
+    // Step 1 — G1 matches positions 0..2
+    let mut g1 = ExternalBlockAssignments::new(0);
+    g1.extend_block_ids(vec![1000, 1001]).unwrap();
+    g1.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(g1.assigned_positions(), 0..2);
+    assert_eq!(g1.assigned_count(), 2);
+
+    // Step 2 — G2 matches positions 2..5
+    let mut g2 = ExternalBlockAssignments::new(2);
+    g2.extend_block_ids(vec![2002, 2003, 2004]).unwrap();
+    g2.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(g2.assigned_positions(), 2..5);
+    assert_eq!(g2.assigned_count(), 3);
+
+    // Step 3 — Queue new block_ids for G1 (don't assign yet — blocks not transferred)
+    g1.extend_block_ids(vec![1000, 1001, 1002, 1003, 1004])
+        .unwrap();
+    assert_eq!(g1.assigned_count(), 2); // still only 2 assigned
+    assert_eq!(g1.unassigned_count(), 3); // 1002, 1003, 1004 pending
+    assert_eq!(g1.pending_positions(), 2..5);
+
+    // Step 4 — Compute onboard pairs (before physical transfer)
+    let pairs = zip_assigned_pending(&g2, &g1);
+    assert_eq!(
+        pairs,
+        vec![(2, 2002, 1002), (3, 2003, 1003), (4, 2004, 1004)]
+    );
+
+    // Step 5 — After physical transfer, assign pending against the full sequence
+    g1.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(g1.assigned_count(), 5);
+    assert_eq!(g1.unassigned_count(), 0);
+    assert_eq!(g1.assigned_positions(), 0..5);
+
+    // Step 6 — Verify assigned overlap matches
+    let assigned_pairs = zip_assigned(&g2, &g1);
+    assert_eq!(
+        assigned_pairs,
+        vec![(2, 2002, 1002), (3, 2003, 1003), (4, 2004, 1004)]
+    );
+}
+
+// =========================================================================
+// New tests: stage_pending + commit_staged (two-step assignment)
+// =========================================================================
+
+#[test]
+fn test_stage_pending_basic() {
+    let seq = create_test_sequence(3, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+
+    // Stage — should move unassigned → staged with hashes
+    let staged_range = assignments.stage_pending(seq.blocks()).unwrap();
+    assert_eq!(staged_range, 0..3);
+    assert_eq!(assignments.staged_count(), 3);
+    assert_eq!(assignments.unassigned_count(), 0);
+    assert_eq!(assignments.assigned_count(), 0);
+
+    // Verify staged contents
+    for (i, expected_hash) in expected_hashes.iter().enumerate() {
+        let (id, hash) = assignments.get_staged(i).unwrap();
+        assert_eq!(id, (i + 1) * 100);
+        assert_eq!(hash, *expected_hash);
+    }
+}
+
+#[test]
+fn test_commit_staged_basic() {
+    let seq = create_test_sequence(3, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+
+    // Stage then commit
+    assignments.stage_pending(seq.blocks()).unwrap();
+    let assigned_range = assignments.commit_staged();
+    assert_eq!(assigned_range, 0..3);
+    assert_eq!(assignments.assigned_count(), 3);
+    assert_eq!(assignments.staged_count(), 0);
+    assert_eq!(assignments.unassigned_count(), 0);
+
+    // Verify assigned contents
+    for (i, expected_hash) in expected_hashes.iter().enumerate() {
+        let (id, hash) = assignments.get_assigned(i).unwrap();
+        assert_eq!(id, (i + 1) * 100);
+        assert_eq!(hash, *expected_hash);
+    }
+}
+
+#[test]
+fn test_stage_pending_partial() {
+    let seq = create_test_sequence(2, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // 5 block_ids but only 2 sequence blocks
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+
+    let staged_range = assignments.stage_pending(seq.blocks()).unwrap();
+    assert_eq!(staged_range, 0..2);
+    assert_eq!(assignments.staged_count(), 2);
+    assert_eq!(assignments.unassigned_count(), 3);
+}
+
+#[test]
+fn test_stage_pending_then_commit_then_stage_more() {
+    let seq = create_test_sequence(5, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+
+    // Stage first 3 (only 3 blocks available at first)
+    let range1 = assignments.stage_pending(&seq.blocks()[..3]).unwrap();
+    assert_eq!(range1, 0..3);
+    assert_eq!(assignments.staged_count(), 3);
+    assert_eq!(assignments.unassigned_count(), 2);
+
+    // Commit them
+    let assigned_range = assignments.commit_staged();
+    assert_eq!(assigned_range, 0..3);
+    assert_eq!(assignments.assigned_count(), 3);
+    assert_eq!(assignments.staged_count(), 0);
+
+    // Stage the remaining 2
+    let range2 = assignments.stage_pending(seq.blocks()).unwrap();
+    assert_eq!(range2, 0..2);
+    assert_eq!(assignments.staged_count(), 2);
+    assert_eq!(assignments.unassigned_count(), 0);
+
+    // Commit them
+    let assigned_range2 = assignments.commit_staged();
+    assert_eq!(assigned_range2, 3..5);
+    assert_eq!(assignments.assigned_count(), 5);
+
+    // Verify all assigned
+    for (i, expected_hash) in expected_hashes.iter().enumerate() {
+        let (id, hash) = assignments.get_assigned(i).unwrap();
+        assert_eq!(id, (i + 1) * 100);
+        assert_eq!(hash, *expected_hash);
+    }
+}
+
+#[test]
+fn test_stage_pending_empty_unassigned() {
+    let seq = create_test_sequence(3, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // No block_ids → nothing to stage
+    let range = assignments.stage_pending(seq.blocks()).unwrap();
+    assert_eq!(range, 0..0);
+    assert_eq!(assignments.staged_count(), 0);
+}
+
+#[test]
+fn test_commit_staged_empty() {
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Nothing staged → empty range
+    let range = assignments.commit_staged();
+    assert_eq!(range, 0..0);
+}
+
+// =========================================================================
+// New tests: staged_positions
+// =========================================================================
+
+#[test]
+fn test_staged_positions_empty() {
+    let assignments = ExternalBlockAssignments::new(0);
+    assert_eq!(assignments.staged_positions(), 0..0);
+}
+
+#[test]
+fn test_staged_positions_after_stage() {
+    let seq = create_test_sequence(5, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Pre-assign first 2
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.assigned_count(), 2);
+
+    // Add 3 more and stage them
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+    assignments.stage_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.staged_count(), 3);
+
+    // Staged positions should be 2..5
+    assert_eq!(assignments.staged_positions(), 2..5);
+    assert_eq!(assignments.assigned_positions(), 0..2);
+}
+
+#[test]
+fn test_staged_positions_with_offset() {
+    let seq = create_test_sequence(10, 0);
+    let mut assignments = ExternalBlockAssignments::new(3);
+
+    assignments.extend_block_ids(vec![100, 200]).unwrap();
+    assignments.assign_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.assigned_count(), 2);
+
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400])
+        .unwrap();
+    assignments.stage_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.staged_count(), 2);
+
+    // offset=3, 2 assigned → staged starts at 5, 2 staged → staged ends at 7
+    assert_eq!(assignments.staged_positions(), 5..7);
+}
+
+// =========================================================================
+// New tests: staged_iter
+// =========================================================================
+
+#[test]
+fn test_staged_iter() {
+    let seq = create_test_sequence(3, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.stage_pending(seq.blocks()).unwrap();
+
+    let staged: Vec<(BlockId, SequenceHash)> = assignments.staged_iter().collect();
+    assert_eq!(staged.len(), 3);
+    assert_eq!(staged[0], (100, expected_hashes[0]));
+    assert_eq!(staged[1], (200, expected_hashes[1]));
+    assert_eq!(staged[2], (300, expected_hashes[2]));
+}
+
+// =========================================================================
+// New tests: take_staged
+// =========================================================================
+
+#[test]
+fn test_take_staged() {
+    let seq = create_test_sequence(3, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+    assignments.stage_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.staged_count(), 3);
+
+    let drained = assignments.take_staged();
+    assert_eq!(drained.len(), 3);
+    assert_eq!(assignments.staged_count(), 0);
+
+    assert_eq!(drained[0], (100, expected_hashes[0]));
+    assert_eq!(drained[1], (200, expected_hashes[1]));
+    assert_eq!(drained[2], (300, expected_hashes[2]));
+}
+
+// =========================================================================
+// New tests: next_position with staged blocks
+// =========================================================================
+
+#[test]
+fn test_next_position_accounts_for_staged() {
+    let seq = create_test_sequence(5, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments.extend_block_ids(vec![100, 200, 300]).unwrap();
+
+    // Before staging
+    assert_eq!(assignments.next_position(), 0);
+
+    // After staging (3 staged, 0 assigned)
+    assignments.stage_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.next_position(), 3);
+    assert_eq!(assignments.assigned_count(), 0);
+    assert_eq!(assignments.staged_count(), 3);
+
+    // After committing (0 staged, 3 assigned)
+    assignments.commit_staged();
+    assert_eq!(assignments.next_position(), 3);
+    assert_eq!(assignments.assigned_count(), 3);
+    assert_eq!(assignments.staged_count(), 0);
+}
+
+#[test]
+fn test_pending_positions_accounts_for_staged() {
+    let seq = create_test_sequence(3, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+
+    // Stage 3, leaving 2 unassigned
+    assignments.stage_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.staged_count(), 3);
+    assert_eq!(assignments.unassigned_count(), 2);
+
+    // Pending positions should be after staged blocks
+    assert_eq!(assignments.pending_positions(), 3..5);
+}
+
+// =========================================================================
+// New tests: extend_assigned (direct insert to assigned)
+// =========================================================================
+
+#[test]
+fn test_extend_assigned_basic() {
+    let seq = create_test_sequence(3, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Build (BlockId, SequenceHash) pairs
+    let items: Vec<(BlockId, SequenceHash)> = vec![
+        (100, expected_hashes[0]),
+        (200, expected_hashes[1]),
+        (300, expected_hashes[2]),
+    ];
+
+    let count = assignments.extend_assigned(items).unwrap();
+    assert_eq!(count, 3);
+    assert_eq!(assignments.assigned_count(), 3);
+    assert_eq!(assignments.staged_count(), 0);
+    assert_eq!(assignments.unassigned_count(), 0);
+
+    for (i, expected_hash) in expected_hashes.iter().enumerate() {
+        let (id, hash) = assignments.get_assigned(i).unwrap();
+        assert_eq!(id, (i + 1) * 100);
+        assert_eq!(hash, *expected_hash);
+    }
+}
+
+#[test]
+fn test_extend_assigned_then_stage_pending() {
+    let seq = create_test_sequence(5, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Direct-assign first 2
+    let items: Vec<(BlockId, SequenceHash)> =
+        vec![(100, expected_hashes[0]), (200, expected_hashes[1])];
+    assignments.extend_assigned(items).unwrap();
+    assert_eq!(assignments.assigned_count(), 2);
+
+    // Enqueue 3 more, stage, commit
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+    assignments.stage_pending(seq.blocks()).unwrap();
+    assignments.commit_staged();
+    assert_eq!(assignments.assigned_count(), 5);
+
+    // Verify all in order
+    for (i, expected_hash) in expected_hashes.iter().enumerate() {
+        let (id, hash) = assignments.get_assigned(i).unwrap();
+        assert_eq!(id, (i + 1) * 100);
+        assert_eq!(hash, *expected_hash);
+    }
+}
+
+#[test]
+fn test_extend_assigned_duplicate_detection() {
+    let seq = create_test_sequence(2, 0);
+    let expected_hashes = get_expected_hashes(&seq);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    // Assign 1 block
+    assignments
+        .extend_assigned(vec![(100, expected_hashes[0])])
+        .unwrap();
+
+    // Try to assign with duplicate block_id → error
+    let result = assignments.extend_assigned(vec![(100, expected_hashes[1])]);
+    assert!(result.is_err());
+    match result.unwrap_err() {
+        BlockSequenceError::DuplicateBlockId { block_id } => {
+            assert_eq!(block_id, 100);
+        }
+        other => panic!("expected DuplicateBlockId, got: {other:?}"),
+    }
+
+    // Original assignment unchanged
+    assert_eq!(assignments.assigned_count(), 1);
+}
+
+#[test]
+fn test_extend_assigned_empty() {
+    let mut assignments = ExternalBlockAssignments::new(0);
+    let count = assignments.extend_assigned(Vec::new()).unwrap();
+    assert_eq!(count, 0);
+    assert_eq!(assignments.assigned_count(), 0);
+}
+
+// =========================================================================
+// New tests: contains checks all three collections
+// =========================================================================
+
+#[test]
+fn test_contains_with_staged() {
+    let seq = create_test_sequence(3, 0);
+    let mut assignments = ExternalBlockAssignments::new(0);
+
+    assignments
+        .extend_block_ids(vec![100, 200, 300, 400, 500])
+        .unwrap();
+
+    // Assign first 2
+    assignments.assign_pending(&seq.blocks()[..2]).unwrap();
+    assert_eq!(assignments.assigned_count(), 2);
+
+    // Stage 1 (block at position 2)
+    assignments.stage_pending(seq.blocks()).unwrap();
+    assert_eq!(assignments.staged_count(), 1);
+    assert_eq!(assignments.unassigned_count(), 2);
+
+    // 100, 200 in assigned
+    assert!(assignments.contains(&100));
+    assert!(assignments.contains(&200));
+    // 300 in staged
+    assert!(assignments.contains(&300));
+    // 400, 500 in unassigned
+    assert!(assignments.contains(&400));
+    assert!(assignments.contains(&500));
+    // 600 nowhere
+    assert!(!assignments.contains(&600));
+}
+
+// =========================================================================
+// New tests: Debug impl
+// =========================================================================
+
+#[test]
+fn test_debug_impl() {
+    let assignments = ExternalBlockAssignments::new(5);
+    let debug_str = format!("{assignments:?}");
+    assert!(debug_str.contains("ExternalBlockAssignments"));
+    assert!(debug_str.contains("offset"));
+}

lib/kvbm-logical/src/sequence/mod.rs

+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! # Sequence Module
+//!
+//! Block assignment tracking for token sequences.
+//!
+//! This module provides two assignment trackers that share a three-phase
+//! lifecycle: **unassigned** (queued), **staged** (paired with token data),
+//! and **assigned** (committed).
+//!
+//! - [`ExternalBlockAssignments`] tracks at the **identity level** — mapping
+//!   `BlockId` to `SequenceHash`.
+//! - [`LogicalBlockAssignments`] tracks at the **guard level** — managing
+//!   RAII block guards through `MutableBlock` to `CompleteBlock` to
+//!   `ImmutableBlock` transitions.
+//!
+//! Both types are backed by the same ordered-collection machinery and expose
+//! similar query, iteration, and mutation APIs.
+//!
+//! ## `ExternalBlockAssignments`
+//!
+//! Identity-level tracking of block IDs paired with sequence hashes.
+//!
+//! ### Basic flow
+//!
+//! Create a sequence from tokens, register block IDs, and assign them
+//! against the completed token blocks:
+//!
+//! ```rust
+//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
+//! use kvbm_logical::{ExternalBlockAssignments, BlockSequence};
+//!
+//! // Create a sequence with 3 complete blocks (4 tokens each).
+//! let tokens: Vec<u32> = (0..12).collect();
+//! let seq = BlockSequence::new(tokens, 4, None);
+//! assert_eq!(seq.blocks().len(), 3);
+//!
+//! // Create assignments starting at offset 0.
+//! let mut assignments = ExternalBlockAssignments::new(0);
+//!
+//! // Register block IDs (e.g., allocated by the scheduler).
+//! assignments.extend_block_ids(vec![10, 20, 30])?;
+//! assert_eq!(assignments.unassigned_count(), 3);
+//!
+//! // Assign: pairs each pending ID with its token block's hash.
+//! let range = assignments.assign_pending(seq.blocks())?;
+//! assert_eq!(range, 0..3);
+//! assert_eq!(assignments.assigned_count(), 3);
+//! assert_eq!(assignments.unassigned_count(), 0);
+//!
+//! // Query by index.
+//! let (id, hash) = assignments.get_assigned(0).unwrap();
+//! assert_eq!(id, 10);
+//! assert_eq!(hash, seq.all_sequence_hashes()[0]);
+//! # Ok(())
+//! # }
+//! ```
+//!
+//! ### Two-step staging
+//!
+//! Use `stage_pending` and `commit_staged` for explicit control over the
+//! staging phase (useful when validation must happen between staging and
+//! committing):
+//!
+//! ```rust
+//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
+//! use kvbm_logical::{ExternalBlockAssignments, BlockSequence};
+//!
+//! let tokens: Vec<u32> = (0..8).collect();
+//! let seq = BlockSequence::new(tokens, 4, None);
+//!
+//! let mut assignments = ExternalBlockAssignments::new(0);
+//! assignments.extend_block_ids(vec![1, 2])?;
+//!
+//! // Stage: pairs IDs with hashes but does not commit.
+//! let staged_range = assignments.stage_pending(seq.blocks())?;
+//! assert_eq!(staged_range, 0..2);
+//! assert_eq!(assignments.staged_count(), 2);
+//! assert_eq!(assignments.assigned_count(), 0);
+//!
+//! // Commit: moves staged into assigned.
+//! let assigned_range = assignments.commit_staged();
+//! assert_eq!(assigned_range, 0..2);
+//! assert_eq!(assignments.assigned_count(), 2);
+//! assert_eq!(assignments.staged_count(), 0);
+//! # Ok(())
+//! # }
+//! ```
+//!
+//!
+//! ## `LogicalBlockAssignments`
+//!
+//! Guard-level tracking through the full block lifecycle. Blocks flow
+//! through `MutableBlock` (unassigned) to `CompleteBlock` (staged) to
+//! `ImmutableBlock` (assigned/registered).
+//!
+//! ### Full pipeline
+//!
+//! Allocate physical blocks from a `BlockManager`, stage them against token
+//! data, register them, and query the result:
+//!
+//! ```rust
+//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
+//! use kvbm_logical::{
+//!     BlockManager, BlockRegistry, BlockSequence,
+//!     LogicalBlockAssignments,
+//!     manager::FrequencyTrackingCapacity,
+//! };
+//!
+//! // Build a manager with 10 blocks of size 4.
+//! let tracker = FrequencyTrackingCapacity::Small.create_tracker();
+//! let registry = BlockRegistry::builder()
+//!     .frequency_tracker(tracker)
+//!     .build();
+//! let manager = BlockManager::<()>::builder()
+//!     .block_count(10)
+//!     .block_size(4)
+//!     .registry(registry)
+//!     .with_lru_backend()
+//!     .build()?;
+//!
+//! // Create a token sequence with 3 complete blocks.
+//! let tokens: Vec<u32> = (0..12).collect();
+//! let seq = BlockSequence::new(tokens, 4, None);
+//!
+//! // Allocate 3 mutable blocks from the manager.
+//! let blocks = manager.allocate_blocks(3).unwrap();
+//! let ids: Vec<usize> = blocks.iter().map(|b| b.block_id()).collect();
+//!
+//! let mut la = LogicalBlockAssignments::new();
+//!
+//! // Extend: adds mutable blocks to the unassigned queue.
+//! la.extend_blocks(blocks)?;
+//! assert_eq!(la.unassigned_count(), 3);
+//!
+//! // Stage: completes each mutable block with its token data.
+//! la.stage(seq.blocks())?;
+//! assert_eq!(la.staged_count(), 3);
+//! assert_eq!(la.unassigned_count(), 0);
+//!
+//! // Register: finalizes staged blocks through the manager.
+//! la.register(&manager);
+//! assert_eq!(la.assigned_count(), 3);
+//! assert_eq!(la.staged_count(), 0);
+//!
+//! // Query assigned blocks.
+//! for i in 0..3 {
+//!     let (id, immutable) = la.get_assigned(i).unwrap();
+//!     assert_eq!(*id, ids[i]);
+//!     assert_eq!(immutable.sequence_hash(), seq.all_sequence_hashes()[i]);
+//! }
+//! # Ok(())
+//! # }
+//! ```
+
+mod assignments;
+mod store;
+
+pub use assignments::{
+    ExternalBlockAssignments, LogicalBlockAssignmentError, LogicalBlockAssignments, zip_assigned,
+    zip_assigned_pending,
+};
+
+use std::ops::Range;
+
+use dynamo_tokens::{Token, TokenBlock, TokenBlockError, TokenBlockSequence, Tokens};
+
+use crate::{BlockId, KvbmSequenceHashProvider, SequenceHash};
+
+/// Errors that can occur in block sequence operations.
+#[derive(Debug, thiserror::Error)]
+pub enum BlockSequenceError {
+    /// A known block_id appeared after an unknown block_id in `extend_block_ids`.
+    #[error(
+        "ordering violation: known block_id {known_id} at index {known_index} \
+         appeared after new block_id {new_id} at index {first_new_index}"
+    )]
+    OrderingViolation {
+        known_id: BlockId,
+        new_id: BlockId,
+        known_index: usize,
+        first_new_index: usize,
+    },
+
+    /// The position embedded in a sequence hash didn't match the expected position.
+    #[error("position mismatch for block_id {block_id}: expected {expected}, actual {actual}")]
+    PositionMismatch {
+        expected: usize,
+        actual: u64,
+        block_id: BlockId,
+    },
+
+    /// A block_id already exists in one of the collections.
+    #[error("duplicate block_id {block_id} already present")]
+    DuplicateBlockId { block_id: BlockId },
+
+    /// Error from underlying token block operations.
+    #[error("token extension error: {0}")]
+    TokenExtension(#[from] TokenBlockError),
+}
+
+/// Owns a `TokenBlockSequence` and provides sequence access and token extension methods.
+///
+/// This is a thin wrapper around `TokenBlockSequence` that provides a convenient API
+/// for the block assignment workflow. It does NOT embed assignments — those are managed
+/// separately by [`ExternalBlockAssignments`].
+#[derive(Debug)]
+pub struct BlockSequence {
+    sequence: TokenBlockSequence,
+}
+
+impl BlockSequence {
+    /// Creates a new `BlockSequence` from tokens, block size, and optional salt hash.
+    pub fn new(tokens: Vec<Token>, block_size: u32, salt_hash: Option<u64>) -> Self {
+        let tokens = Tokens::from(tokens);
+        Self {
+            sequence: TokenBlockSequence::new(tokens, block_size, salt_hash),
+        }
+    }
+
+    /// Returns the completed token blocks.
+    pub fn blocks(&self) -> &[TokenBlock] {
+        self.sequence.blocks()
+    }
+
+    /// Returns the block size.
+    pub fn block_size(&self) -> usize {
+        self.sequence.block_size()
+    }
+
+    /// Returns the total number of tokens (including partial block).
+    pub fn total_tokens(&self) -> usize {
+        self.sequence.total_tokens()
+    }
+
+    /// Returns a reference to the underlying `TokenBlockSequence`.
+    pub fn sequence(&self) -> &TokenBlockSequence {
+        &self.sequence
+    }
+
+    /// Returns all sequence hashes from completed blocks.
+    pub fn all_sequence_hashes(&self) -> Vec<SequenceHash> {
+        self.sequence
+            .blocks()
+            .iter()
+            .map(|b| b.kvbm_sequence_hash())
+            .collect()
+    }
+
+    /// Extends the sequence with tokens, potentially completing blocks.
+    ///
+    /// Returns the range of newly completed block indices, or `None` if no blocks completed.
+    pub fn extend_tokens(
+        &mut self,
+        tokens: Vec<Token>,
+    ) -> Result<Option<Range<usize>>, BlockSequenceError> {
+        let tokens = Tokens::from(tokens);
+        self.sequence.extend(tokens).map_err(Into::into)
+    }
+
+    /// Appends a single token to the sequence.
+    ///
+    /// Returns the index of the completed block if the token completed a block.
+    pub fn append_token(&mut self, token: Token) -> Result<Option<usize>, BlockSequenceError> {
+        self.sequence.append(token).map_err(Into::into)
+    }
+}

lib/kvbm-logical/src/sequence/store.rs

+// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+//! Shared three-collection lifecycle store backing both
+//! [`ExternalBlockAssignments`](super::assignments::ExternalBlockAssignments) and
+//! [`LogicalBlockAssignments`](super::assignments::LogicalBlockAssignments).
+
+use indexmap::IndexMap;
+
+use crate::BlockId;
+
+/// Generic three-phase lifecycle store keyed by [`BlockId`].
+///
+/// Manages three ordered `IndexMap` collections representing the lifecycle
+/// phases: **unassigned** (`U`) → **staged** (`S`) → **assigned** (`A`).
+///
+/// Both [`ExternalBlockAssignments`](super::assignments::ExternalBlockAssignments) and
+/// [`LogicalBlockAssignments`](super::assignments::LogicalBlockAssignments) compose this
+/// type internally and add their own type-specific transition logic on top.
+pub(crate) struct BlockStore<U, S, A> {
+    assigned: IndexMap<BlockId, A>,
+    staged: IndexMap<BlockId, S>,
+    unassigned: IndexMap<BlockId, U>,
+}
+
+impl<U, S, A> BlockStore<U, S, A> {
+    /// Creates an empty store.
+    pub fn new() -> Self {
+        Self {
+            assigned: IndexMap::new(),
+            staged: IndexMap::new(),
+            unassigned: IndexMap::new(),
+        }
+    }
+
+    // -- Counts ---------------------------------------------------------------
+
+    /// Returns the number of assigned entries.
+    pub fn assigned_count(&self) -> usize {
+        self.assigned.len()
+    }
+
+    /// Returns the number of staged entries.
+    pub fn staged_count(&self) -> usize {
+        self.staged.len()
+    }
+
+    /// Returns the number of unassigned entries.
+    pub fn unassigned_count(&self) -> usize {
+        self.unassigned.len()
+    }
+
+    // -- Queries --------------------------------------------------------------
+
+    /// Returns `true` if all three collections are empty.
+    pub fn is_empty(&self) -> bool {
+        self.assigned.is_empty() && self.staged.is_empty() && self.unassigned.is_empty()
+    }
+
+    /// Checks whether a `BlockId` is present in any of the three collections.
+    pub fn contains(&self, block_id: &BlockId) -> bool {
+        self.assigned.contains_key(block_id)
+            || self.staged.contains_key(block_id)
+            || self.unassigned.contains_key(block_id)
+    }
+
+    /// Checks whether a `BlockId` is present in the assigned collection.
+    #[allow(dead_code)]
+    pub fn contains_assigned(&self, block_id: &BlockId) -> bool {
+        self.assigned.contains_key(block_id)
+    }
+
+    /// Checks whether a `BlockId` is present in the staged collection.
+    #[allow(dead_code)]
+    pub fn contains_staged(&self, block_id: &BlockId) -> bool {
+        self.staged.contains_key(block_id)
+    }
+
+    /// Checks whether a `BlockId` is present in the unassigned collection.
+    #[allow(dead_code)]
+    pub fn contains_unassigned(&self, block_id: &BlockId) -> bool {
+        self.unassigned.contains_key(block_id)
+    }
+
+    // -- Index Access ---------------------------------------------------------
+
+    /// Returns the assigned entry at the given index (insertion order).
+    pub fn get_assigned(&self, index: usize) -> Option<(&BlockId, &A)> {
+        self.assigned.get_index(index)
+    }
+
+    /// Returns the staged entry at the given index (staging order).
+    pub fn get_staged(&self, index: usize) -> Option<(&BlockId, &S)> {
+        self.staged.get_index(index)
+    }
+
+    /// Returns the unassigned entry at the given index (FIFO order).
+    pub fn get_unassigned(&self, index: usize) -> Option<(&BlockId, &U)> {
+        self.unassigned.get_index(index)
+    }
+
+    // -- Iteration ------------------------------------------------------------
+
+    /// Iterates over assigned entries in positional order.
+    pub fn assigned_iter(&self) -> impl Iterator<Item = (&BlockId, &A)> {
+        self.assigned.iter()
+    }
+
+    /// Iterates over staged entries in staging order.
+    pub fn staged_iter(&self) -> impl Iterator<Item = (&BlockId, &S)> {
+        self.staged.iter()
+    }
+
+    /// Iterates over unassigned entries in FIFO order.
+    pub fn unassigned_iter(&self) -> impl Iterator<Item = (&BlockId, &U)> {
+        self.unassigned.iter()
+    }
+
+    // -- FIFO Pop -------------------------------------------------------------
+
+    /// Removes and returns the first unassigned entry (FIFO).
+    pub fn shift_unassigned(&mut self) -> Option<(BlockId, U)> {
+        self.unassigned.shift_remove_index(0)
+    }
+
+    /// Removes and returns the first staged entry (FIFO).
+    pub fn shift_staged(&mut self) -> Option<(BlockId, S)> {
+        self.staged.shift_remove_index(0)
+    }
+
+    // -- LIFO Pop -------------------------------------------------------------
+
+    /// Removes and returns the last unassigned entry (LIFO).
+    pub fn pop_unassigned(&mut self) -> Option<(BlockId, U)> {
+        self.unassigned.pop()
+    }
+
+    // -- Insert ---------------------------------------------------------------
+
+    /// Inserts into the assigned collection.
+    pub fn insert_assigned(&mut self, id: BlockId, val: A) {
+        self.assigned.insert(id, val);
+    }
+
+    /// Inserts into the staged collection.
+    pub fn insert_staged(&mut self, id: BlockId, val: S) {
+        self.staged.insert(id, val);
+    }
+
+    /// Inserts into the unassigned collection.
+    pub fn insert_unassigned(&mut self, id: BlockId, val: U) {
+        self.unassigned.insert(id, val);
+    }
+
+    // -- Bulk -----------------------------------------------------------------
+
+    /// Iterates over all block IDs across all three collections in lifecycle
+    /// order: assigned → staged → unassigned.
+    pub fn all_block_ids(&self) -> impl Iterator<Item = &BlockId> {
+        self.assigned
+            .keys()
+            .chain(self.staged.keys())
+            .chain(self.unassigned.keys())
+    }
+
+    /// Validates that none of the given `ids` collide with existing entries
+    /// or with each other.
+    ///
+    /// Returns `Ok(())` if all IDs are unique, or `Err(id)` with the first
+    /// duplicate found.
+    pub fn validate_no_duplicates(
+        &self,
+        ids: impl Iterator<Item = BlockId>,
+        count_hint: usize,
+    ) -> Result<(), BlockId> {
+        let mut seen = indexmap::IndexSet::with_capacity(count_hint);
+        for id in ids {
+            if self.contains(&id) || !seen.insert(id) {
+                return Err(id);
+            }
+        }
+        Ok(())
+    }
+
+    /// Clears all three collections.
+    pub fn clear(&mut self) {
+        self.assigned.clear();
+        self.staged.clear();
+        self.unassigned.clear();
+    }
+
+    /// Takes all assigned entries, returning them as a `Vec`.
+    pub fn take_assigned(&mut self) -> Vec<(BlockId, A)> {
+        std::mem::take(&mut self.assigned).into_iter().collect()
+    }
+
+    /// Takes all staged entries, returning them as a `Vec`.
+    pub fn take_staged(&mut self) -> Vec<(BlockId, S)> {
+        std::mem::take(&mut self.staged).into_iter().collect()
+    }
+
+    /// Takes all unassigned entries, returning them as a `Vec`.
+    pub fn take_unassigned(&mut self) -> Vec<(BlockId, U)> {
+        std::mem::take(&mut self.unassigned).into_iter().collect()
+    }
+}

lib/kvbm-logical/src/testing/config.rs

@@ -66,7 +66,7 @@ pub fn validate_test_block_size(size: usize) -> bool {
 ///
 /// # Example
 /// ```
-/// use dynamo_kvbm_logical::testing::config::generate_test_tokens;
+/// use kvbm_logical::testing::config::generate_test_tokens;
 ///
 /// let tokens = generate_test_tokens(100, 4);
 /// assert_eq!(tokens, vec![100, 101, 102, 103]);
@@ -86,7 +86,7 @@ pub fn generate_test_tokens(base: u32, block_size: usize) -> Vec<u32> {
 ///
 /// # Example
 /// ```
-/// use dynamo_kvbm_logical::testing::config::generate_unique_token_sequences;
+/// use kvbm_logical::testing::config::generate_unique_token_sequences;
 ///
 /// let sequences = generate_unique_token_sequences(2, 4);
 /// assert_eq!(sequences.len(), 2);

lib/kvbm-logical/src/testing/managers.rs

@@ -3,17 +3,9 @@
 
 //! Test manager creation helpers.
 
-use std::marker::PhantomData;
-use std::sync::Arc;
-
-use anyhow::Result;
-
-use crate::SequenceHash;
-use crate::blocks::{BlockMetadata, BlockRegistry};
-use crate::events::EventsManager;
+use crate::blocks::BlockMetadata;
 use crate::manager::{BlockManager, FrequencyTrackingCapacity};
-
-use super::token_blocks;
+use crate::registry::BlockRegistry;
 
 /// Create a basic test manager with LRU backend.
 pub fn create_test_manager<T: BlockMetadata>(block_count: usize) -> BlockManager<T> {
@@ -47,304 +39,3 @@ pub fn create_test_manager_with_block_size<T: BlockMetadata>(
         .build()
         .expect("Should build manager")
 }
-
-/// Builder for creating test BlockRegistry with optional events integration.
-///
-/// # Example
-///
-/// ```ignore
-/// // Simple registry
-/// let registry = TestRegistryBuilder::new().build();
-///
-/// // With events manager
-/// let events_manager = Arc::new(EventsManager::builder().build());
-/// let registry = TestRegistryBuilder::new()
-///     .events_manager(events_manager)
-///     .build();
-///
-/// // With custom frequency tracking
-/// let registry = TestRegistryBuilder::new()
-///     .frequency_tracking(FrequencyTrackingCapacity::Large)
-///     .build();
-/// ```
-#[derive(Default)]
-pub struct TestRegistryBuilder {
-    events_manager: Option<Arc<EventsManager>>,
-    frequency_tracking: FrequencyTrackingCapacity,
-}
-
-impl TestRegistryBuilder {
-    /// Creates a new builder with default settings.
-    pub fn new() -> Self {
-        Self {
-            events_manager: None,
-            frequency_tracking: FrequencyTrackingCapacity::Medium,
-        }
-    }
-
-    /// Sets the events manager for distributed event coordination.
-    pub fn events_manager(mut self, manager: Arc<EventsManager>) -> Self {
-        self.events_manager = Some(manager);
-        self
-    }
-
-    /// Sets the frequency tracking capacity.
-    ///
-    /// Default: Medium
-    pub fn frequency_tracking(mut self, capacity: FrequencyTrackingCapacity) -> Self {
-        self.frequency_tracking = capacity;
-        self
-    }
-
-    /// Builds the BlockRegistry.
-    pub fn build(self) -> BlockRegistry {
-        let mut builder =
-            BlockRegistry::builder().frequency_tracker(self.frequency_tracking.create_tracker());
-
-        if let Some(events_manager) = self.events_manager {
-            builder = builder.event_manager(events_manager);
-        }
-
-        builder.build()
-    }
-}
-
-/// Builder for creating test BlockManagers.
-///
-/// # Example
-///
-/// ```ignore
-/// // Simple manager (creates its own registry)
-/// let manager = TestManagerBuilder::<G1>::new()
-///     .block_count(100)
-///     .block_size(4)
-///     .build();
-///
-/// // With explicit registry (for events integration)
-/// let events_manager = Arc::new(EventsManager::builder().build());
-/// let registry = TestRegistryBuilder::new()
-///     .events_manager(events_manager.clone())
-///     .build();
-/// let manager = TestManagerBuilder::<G1>::new()
-///     .block_count(100)
-///     .block_size(4)
-///     .registry(registry)
-///     .build();
-///
-/// // Convenience: with events manager (creates registry internally)
-/// let manager = TestManagerBuilder::<G1>::new()
-///     .block_count(100)
-///     .block_size(4)
-///     .events_manager(events_manager)
-///     .build();
-/// ```
-pub struct TestManagerBuilder<T: BlockMetadata> {
-    block_count: Option<usize>,
-    block_size: Option<usize>,
-    registry: Option<BlockRegistry>,
-    events_manager: Option<Arc<EventsManager>>,
-    frequency_tracking: FrequencyTrackingCapacity,
-    _phantom: PhantomData<T>,
-}
-
-impl<T: BlockMetadata> Default for TestManagerBuilder<T> {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
-impl<T: BlockMetadata> TestManagerBuilder<T> {
-    /// Creates a new builder with default settings.
-    pub fn new() -> Self {
-        Self {
-            block_count: None,
-            block_size: None,
-            registry: None,
-            events_manager: None,
-            frequency_tracking: FrequencyTrackingCapacity::Medium,
-            _phantom: PhantomData,
-        }
-    }
-
-    /// Sets the number of blocks in the pool.
-    pub fn block_count(mut self, count: usize) -> Self {
-        self.block_count = Some(count);
-        self
-    }
-
-    /// Sets the tokens per block (must be power of 2, 1-1024).
-    pub fn block_size(mut self, size: usize) -> Self {
-        self.block_size = Some(size);
-        self
-    }
-
-    /// Sets the registry to use.
-    ///
-    /// If not set, a registry will be created based on `frequency_tracking`
-    /// and `events_manager` settings.
-    pub fn registry(mut self, registry: BlockRegistry) -> Self {
-        self.registry = Some(registry);
-        self
-    }
-
-    /// Sets the events manager for distributed event coordination.
-    ///
-    /// This is a convenience method that creates a registry with the events manager.
-    /// If you also call `registry()`, this setting is ignored.
-    pub fn events_manager(mut self, manager: Arc<EventsManager>) -> Self {
-        self.events_manager = Some(manager);
-        self
-    }
-
-    /// Sets the frequency tracking capacity for auto-created registry.
-    ///
-    /// Ignored if `registry()` is called.
-    ///
-    /// Default: Medium
-    pub fn frequency_tracking(mut self, capacity: FrequencyTrackingCapacity) -> Self {
-        self.frequency_tracking = capacity;
-        self
-    }
-
-    /// Builds the BlockManager.
-    ///
-    /// # Panics
-    ///
-    /// Panics if `block_count` or `block_size` are not set.
-    pub fn build(self) -> BlockManager<T> {
-        let block_count = self.block_count.expect("block_count is required");
-        let block_size = self.block_size.expect("block_size is required");
-
-        let registry = self.registry.unwrap_or_else(|| {
-            let mut builder =
-                TestRegistryBuilder::new().frequency_tracking(self.frequency_tracking);
-            if let Some(events_manager) = self.events_manager {
-                builder = builder.events_manager(events_manager);
-            }
-            builder.build()
-        });
-
-        BlockManager::<T>::builder()
-            .block_count(block_count)
-            .block_size(block_size)
-            .registry(registry)
-            .with_lru_backend()
-            .build()
-            .expect("Should build test manager")
-    }
-}
-
-/// Populate a BlockManager with token blocks and return their sequence hashes.
-///
-/// This function:
-/// 1. Allocates blocks from the manager
-/// 2. Completes them with provided token blocks
-/// 3. Registers them
-/// 4. Drops the immutable blocks (returns to inactive pool)
-///
-/// # Returns
-/// Vec of sequence hashes for the registered blocks (in order)
-pub fn populate_manager_with_blocks<T: BlockMetadata>(
-    manager: &BlockManager<T>,
-    token_blocks: &[dynamo_tokens::TokenBlock],
-) -> Result<Vec<SequenceHash>> {
-    let blocks = manager
-        .allocate_blocks(token_blocks.len())
-        .ok_or_else(|| anyhow::anyhow!("Failed to allocate {} blocks", token_blocks.len()))?;
-
-    let complete_blocks: Vec<_> = blocks
-        .into_iter()
-        .zip(token_blocks.iter())
-        .map(|(block, token_block)| {
-            block
-                .complete(token_block)
-                .map_err(|e| anyhow::anyhow!("Failed to complete block: {:?}", e))
-        })
-        .collect::<Result<Vec<_>>>()?;
-
-    let seq_hashes: Vec<SequenceHash> = complete_blocks.iter().map(|b| b.sequence_hash()).collect();
-
-    let immutable_blocks = manager.register_blocks(complete_blocks);
-
-    // Drop immutable blocks - they return to inactive pool via RAII
-    drop(immutable_blocks);
-
-    Ok(seq_hashes)
-}
-
-/// Quick setup: create manager and populate with sequential token blocks.
-///
-/// # Arguments
-/// * `block_count` - Number of blocks
-/// * `block_size` - Tokens per block
-/// * `start_token` - Starting token value for sequence
-///
-/// # Returns
-/// (BlockManager, Vec<SequenceHash>)
-pub fn create_and_populate_manager<T: BlockMetadata>(
-    block_count: usize,
-    block_size: usize,
-    start_token: u32,
-    registry: BlockRegistry,
-) -> Result<(BlockManager<T>, Vec<SequenceHash>)> {
-    let manager = TestManagerBuilder::<T>::new()
-        .block_count(block_count)
-        .block_size(block_size)
-        .registry(registry)
-        .build();
-
-    let token_sequence = token_blocks::create_token_sequence(block_count, block_size, start_token);
-    let seq_hashes = populate_manager_with_blocks(&manager, token_sequence.blocks())?;
-
-    Ok((manager, seq_hashes))
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[derive(Clone, Debug)]
-    struct TestMetadata;
-
-    #[test]
-    fn test_create_test_manager() {
-        let manager = TestManagerBuilder::<TestMetadata>::new()
-            .block_count(100)
-            .block_size(16)
-            .build();
-        assert_eq!(manager.total_blocks(), 100);
-        assert_eq!(manager.block_size(), 16);
-        assert_eq!(manager.available_blocks(), 100);
-    }
-
-    #[test]
-    fn test_populate_manager_with_blocks() {
-        let manager = TestManagerBuilder::<TestMetadata>::new()
-            .block_count(50)
-            .block_size(4)
-            .build();
-        let token_seq = token_blocks::create_token_sequence(10, 4, 0);
-
-        let seq_hashes =
-            populate_manager_with_blocks(&manager, token_seq.blocks()).expect("Should populate");
-
-        assert_eq!(seq_hashes.len(), 10);
-        // Blocks should be in inactive pool after population
-        assert_eq!(manager.available_blocks(), 50);
-    }
-
-    #[test]
-    fn test_create_and_populate_manager() {
-        let registry = TestRegistryBuilder::new().build();
-        let (manager, hashes) = create_and_populate_manager::<TestMetadata>(32, 4, 100, registry)
-            .expect("Should create");
-
-        assert_eq!(hashes.len(), 32);
-        assert_eq!(manager.total_blocks(), 32);
-        assert_eq!(manager.available_blocks(), 32);
-
-        // Verify blocks can be matched
-        let matched = manager.match_blocks(&hashes);
-        assert_eq!(matched.len(), 32);
-    }
-}

lib/kvbm-logical/src/testing/mod.rs

@@ -55,15 +55,10 @@ pub const TEST_SALT: u64 = 42;
 // ============================================================================
 
 // pub items — usable by downstream crates via the testing feature
-pub use managers::{
-    TestManagerBuilder, TestRegistryBuilder, create_and_populate_manager, create_test_manager,
-    create_test_manager_with_block_size, populate_manager_with_blocks,
-};
+pub use managers::{create_test_manager, create_test_manager_with_block_size};
 pub use sequences::BlockSequenceBuilder;
 pub use token_blocks::{
-    create_disjoint_sequences, create_iota_token_block, create_sequential_block,
-    create_test_token_block, create_token_block, create_token_sequence, default_request_salt_hash,
-    generate_sequence_hashes, sequential_tokens, tokens_for_id,
+    create_iota_token_block, create_test_token_block, sequential_tokens, tokens_for_id,
 };
 
 // pub(crate) items — internal helpers using internal types

lib/kvbm-logical/src/testing/token_blocks.rs

@@ -3,9 +3,7 @@
 
 //! Token block creation helpers for tests.
 
-use dynamo_tokens::{TokenBlock, TokenBlockSequence, compute_hash_v2};
-
-use crate::{KvbmSequenceHashProvider, SequenceHash};
+use dynamo_tokens::{TokenBlock, TokenBlockSequence};
 
 use super::TEST_SALT;
 
@@ -40,175 +38,3 @@ pub fn sequential_tokens(start: u32, count: usize) -> Vec<u32> {
 pub fn tokens_for_id(id: u64) -> Vec<u32> {
     vec![id as u32, (id + 1) as u32, (id + 2) as u32, (id + 3) as u32]
 }
-
-/// Compute the default salt hash for requests with no salt and no lora.
-///
-/// This matches the hash computed by `Request::new()` when salt=None and lora_name=None.
-pub fn default_request_salt_hash() -> u64 {
-    // Matches Request::new() computation:
-    // SaltPayload { salt: None, lora_name: None } serializes to "{}"
-    compute_hash_v2(b"{}", 0)
-}
-
-/// Create a token block from a slice of tokens.
-///
-/// Uses the default request salt hash to match blocks created by
-/// requests with no salt parameter.
-pub fn create_token_block(tokens: &[u32]) -> TokenBlock {
-    let salt = default_request_salt_hash();
-    let token_sequence = TokenBlockSequence::from_slice(tokens, tokens.len() as u32, Some(salt));
-    if let Some(block) = token_sequence.blocks().first() {
-        block.clone()
-    } else {
-        let mut partial = token_sequence.into_parts().1;
-        partial.commit().expect("Should be able to commit")
-    }
-}
-
-/// Create a token block with sequential tokens starting from `start`.
-///
-/// # Arguments
-/// * `start` - Starting token value
-/// * `count` - Number of tokens to generate
-pub fn create_sequential_block(start: u32, count: usize) -> TokenBlock {
-    let tokens: Vec<u32> = (start..start + count as u32).collect();
-    create_token_block(&tokens)
-}
-
-/// Create a token sequence with multiple blocks.
-///
-/// Uses the default request salt hash to match blocks created by
-/// requests with no salt parameter.
-///
-/// # Arguments
-/// * `num_blocks` - Number of blocks to create
-/// * `block_size` - Tokens per block
-/// * `start_token` - Starting token value
-///
-/// # Returns
-/// A TokenBlockSequence containing the requested blocks.
-pub fn create_token_sequence(
-    num_blocks: usize,
-    block_size: usize,
-    start_token: u32,
-) -> TokenBlockSequence {
-    let salt = default_request_salt_hash();
-    let total_tokens = num_blocks * block_size;
-    let tokens: Vec<u32> = (start_token..start_token + total_tokens as u32).collect();
-    TokenBlockSequence::from_slice(&tokens, block_size as u32, Some(salt))
-}
-
-/// Generate sequence hashes from a token sequence.
-pub fn generate_sequence_hashes(token_sequence: &TokenBlockSequence) -> Vec<SequenceHash> {
-    token_sequence
-        .blocks()
-        .iter()
-        .map(|block| block.kvbm_sequence_hash())
-        .collect()
-}
-
-/// Create multiple disjoint token sequences with gaps between them.
-///
-/// This is useful for testing contiguous subsequence detection, where you need
-/// blocks at non-consecutive positions with gaps between them.
-///
-/// # Arguments
-/// * `segments` - Vec of (num_blocks, start_token) pairs. Each segment creates
-///   consecutive blocks starting at the given token.
-/// * `block_size` - Tokens per block
-///
-/// # Returns
-/// A tuple of (Vec<TokenBlock>, Vec<SequenceHash>) containing all blocks and
-/// their hashes from all segments, sorted by position.
-pub fn create_disjoint_sequences(
-    segments: Vec<(usize, u32)>,
-    block_size: usize,
-) -> (Vec<TokenBlock>, Vec<SequenceHash>) {
-    let mut all_blocks = Vec::new();
-    let mut all_hashes = Vec::new();
-
-    for (num_blocks, start_token) in segments {
-        let token_sequence = create_token_sequence(num_blocks, block_size, start_token);
-        let blocks = token_sequence.blocks().to_vec();
-        let hashes = generate_sequence_hashes(&token_sequence);
-
-        all_blocks.extend(blocks);
-        all_hashes.extend(hashes);
-    }
-
-    // Sort by position to maintain order
-    let mut combined: Vec<_> = all_blocks.into_iter().zip(all_hashes).collect();
-    combined.sort_by_key(|(_, hash)| hash.position());
-
-    combined.into_iter().unzip()
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_create_token_block() {
-        let tokens = vec![1, 2, 3, 4];
-        let block = create_token_block(&tokens);
-        assert_eq!(block.tokens().len(), 4);
-    }
-
-    #[test]
-    fn test_create_sequential_block() {
-        let block = create_sequential_block(100, 4);
-        assert_eq!(block.tokens().len(), 4);
-    }
-
-    #[test]
-    fn test_create_token_sequence() {
-        let sequence = create_token_sequence(10, 4, 0);
-        assert_eq!(sequence.blocks().len(), 10);
-
-        // Verify first block starts at token 0
-        let first_block = &sequence.blocks()[0];
-        assert_eq!(first_block.tokens().len(), 4);
-    }
-
-    #[test]
-    fn test_generate_sequence_hashes() {
-        let sequence = create_token_sequence(5, 4, 100);
-        let hashes = generate_sequence_hashes(&sequence);
-
-        assert_eq!(hashes.len(), 5);
-
-        // Verify hashes are unique
-        let unique_hashes: std::collections::HashSet<_> = hashes.iter().collect();
-        assert_eq!(unique_hashes.len(), 5);
-    }
-
-    #[test]
-    fn test_create_disjoint_sequences() {
-        // Create 3 segments with different token ranges
-        let segments = vec![
-            (2, 0),   // 2 blocks starting at token 0
-            (2, 100), // 2 blocks starting at token 100
-            (3, 200), // 3 blocks starting at token 200
-        ];
-        let block_size = 4;
-
-        let (blocks, hashes) = create_disjoint_sequences(segments, block_size);
-
-        // Should have 7 total blocks
-        assert_eq!(blocks.len(), 7);
-        assert_eq!(hashes.len(), 7);
-
-        // All hashes should be unique (different token content = different hashes)
-        let unique_hashes: std::collections::HashSet<_> = hashes.iter().collect();
-        assert_eq!(unique_hashes.len(), 7);
-
-        // Positions are relative within each segment's TokenBlockSequence
-        assert_eq!(hashes[0].position(), 0);
-        assert_eq!(hashes[1].position(), 0);
-        assert_eq!(hashes[2].position(), 0);
-        assert_eq!(hashes[3].position(), 1);
-        assert_eq!(hashes[4].position(), 1);
-        assert_eq!(hashes[5].position(), 1);
-        assert_eq!(hashes[6].position(), 2);
-    }
-}