jundot--omlx/tests/test_hot_cache.py

# SPDX-License-Identifier: Apache-2.0
"""Tests for the in-memory hot cache tier in PagedSSDCacheManager."""

import threading
import time
from pathlib import Path
from typing import List
from unittest.mock import patch

import pytest

from omlx.cache.paged_ssd_cache import (
    PagedSSDBlockMetadata,
    PagedSSDCacheManager,
    SharedHotCacheBudget,
    _extract_tensor_bytes,
)

try:
    import mlx.core as mx

    HAS_MLX = True
except ImportError:
    HAS_MLX = False


@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
class TestHotCacheDisabled:
    """Verify that hot_cache_max_bytes=0 preserves existing behaviour."""

    @pytest.fixture
    def manager(self, tmp_path):
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=0,
        )
        yield mgr
        mgr.close()

    def test_hot_cache_disabled_by_default(self, manager):
        """hot_cache_max_bytes=0 means hot cache is disabled."""
        assert manager._hot_cache_enabled is False
        assert manager._hot_cache_max_bytes == 0

    def test_save_load_works_without_hot_cache(self, manager):
        """Save/load should work even when hot cache is disabled."""
        block_hash = b"disabled_hot_cache_test"
        cache_data = [
            (mx.zeros((1, 8, 64, 64)), mx.zeros((1, 8, 64, 64))) for _ in range(4)
        ]
        result = manager.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=64,
            model_name="test-model",
            layer_cache_types=["KVCache"] * 4,
        )
        assert result is True
        assert manager.has_block(block_hash)

        loaded = manager.load_block(block_hash)
        assert loaded is not None
        assert len(loaded) == 4

    def test_stats_hot_cache_zero_when_disabled(self, manager):
        """Hot cache stats should be zero when disabled."""
        stats = manager.get_stats()
        assert stats.hot_cache_entries == 0
        assert stats.hot_cache_size_bytes == 0
        assert stats.hot_cache_max_bytes == 0
        assert stats.hot_cache_hits == 0
        assert stats.hot_cache_evictions == 0
        assert stats.hot_cache_promotions == 0


@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
class TestHotCacheEnabled:
    """Test hot cache with in-memory caching active."""

    @pytest.fixture
    def manager(self, tmp_path):
        # 10 MB hot cache — generous for test blocks
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "ssd_cache",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=10 * 1024**2,
        )
        yield mgr
        mgr.close()

    def _make_cache_data(self, num_layers=4, seq_len=32, heads=4, head_dim=32):
        """Create test cache data."""
        return [
            (
                mx.zeros((1, heads, seq_len, head_dim)),
                mx.zeros((1, heads, seq_len, head_dim)),
            )
            for _ in range(num_layers)
        ]

    def _save_block(self, manager, block_hash, num_layers=4, model="test-model"):
        """Save a test block and return True on success."""
        cache_data = self._make_cache_data(num_layers=num_layers)
        return manager.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=64,
            model_name=model,
            layer_cache_types=["KVCache"] * num_layers,
        )

    def test_save_stores_in_hot_cache(self, manager):
        """After save_block(), the entry should be in hot cache."""
        block_hash = b"hot_cache_save_test1"
        self._save_block(manager, block_hash)

        # Verify hot cache has the entry
        entry = manager._hot_cache_get(block_hash)
        assert entry is not None
        assert "tensors_raw" in entry
        assert entry["num_layers"] == 4

    def test_load_from_hot_cache(self, manager):
        """load_block() should return data from hot cache without SSD I/O."""
        block_hash = b"hot_cache_load_test1"
        self._save_block(manager, block_hash)

        loaded = manager.load_block(block_hash)
        assert loaded is not None
        assert len(loaded) == 4

        stats = manager.get_stats()
        assert stats.hot_cache_hits >= 1

    def test_hot_cache_hit_updates_stats(self, manager):
        """Hot cache hit should increment hot_cache_hits counter."""
        block_hash = b"hot_cache_stats_test1"
        self._save_block(manager, block_hash)

        initial_stats = manager.get_stats()
        initial_hits = initial_stats.hot_cache_hits

        manager.load_block(block_hash)
        manager.load_block(block_hash)

        stats = manager.get_stats()
        assert stats.hot_cache_hits >= initial_hits + 2

    def test_hot_cache_size_tracking(self, manager):
        """Hot cache should track total size in bytes."""
        block_hash = b"hot_cache_size_test1"
        self._save_block(manager, block_hash)

        stats = manager.get_stats()
        assert stats.hot_cache_entries == 1
        assert stats.hot_cache_size_bytes > 0
        assert stats.hot_cache_max_bytes == 10 * 1024**2

    def test_delete_block_removes_from_hot_cache(self, manager):
        """delete_block() should remove entry from hot cache."""
        block_hash = b"hot_cache_delete_test"
        self._save_block(manager, block_hash)

        # Verify it's in hot cache
        assert manager._hot_cache_get(block_hash) is not None

        manager.delete_block(block_hash)

        # Verify it's gone from hot cache
        assert manager._hot_cache_get(block_hash) is None

    def test_close_clears_hot_cache(self, tmp_path):
        """close() should clear all hot cache entries."""
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "close_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=10 * 1024**2,
        )
        block_hash = b"hot_cache_close_test1"
        cache_data = self._make_cache_data()
        mgr.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=64,
            model_name="test",
            layer_cache_types=["KVCache"] * 4,
        )
        assert len(mgr._hot_cache) > 0

        mgr.close()

        assert len(mgr._hot_cache) == 0
        assert mgr._hot_cache_total_bytes == 0


@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
class TestHotCacheLRU:
    """Test LRU eviction behaviour of the hot cache."""

    def _make_cache_data(self, num_layers=2, seq_len=16, heads=2, head_dim=16):
        """Create small test cache data."""
        return [
            (
                mx.zeros((1, heads, seq_len, head_dim)),
                mx.zeros((1, heads, seq_len, head_dim)),
            )
            for _ in range(num_layers)
        ]

    def _entry_size(self, num_layers=2, seq_len=16, heads=2, head_dim=16):
        """Estimate the raw byte size of one entry."""
        # Each tensor: 1 * heads * seq_len * head_dim * 4 bytes (float32)
        # 2 tensors (keys + values) per layer, num_layers layers
        return num_layers * 2 * 1 * heads * seq_len * head_dim * 4

    def test_lru_eviction(self, tmp_path):
        """Old entries should be evicted when capacity is exceeded."""
        entry_size = self._entry_size()
        # Allow room for exactly 2 entries
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "lru_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )

        try:
            # Save 3 blocks — the first should be evicted
            for i in range(3):
                block_hash = f"lru_block_{i}".encode()
                cache_data = self._make_cache_data()
                mgr.save_block(
                    block_hash=block_hash,
                    cache_data=cache_data,
                    token_count=16,
                    model_name="test",
                    layer_cache_types=["KVCache"] * 2,
                )

            # Block 0 should have been evicted (LRU)
            assert mgr._hot_cache_get(b"lru_block_0") is None
            # Blocks 1 and 2 should still be in hot cache
            assert mgr._hot_cache_get(b"lru_block_1") is not None
            assert mgr._hot_cache_get(b"lru_block_2") is not None

            stats = mgr.get_stats()
            assert stats.hot_cache_evictions >= 1
        finally:
            mgr.close()

    def test_lru_access_refreshes_order(self, tmp_path):
        """Accessing a block should move it to MRU position."""
        entry_size = self._entry_size()
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "lru_order_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )

        try:
            # Save blocks 0 and 1
            for i in range(2):
                block_hash = f"order_block_{i}".encode()
                cache_data = self._make_cache_data()
                mgr.save_block(
                    block_hash=block_hash,
                    cache_data=cache_data,
                    token_count=16,
                    model_name="test",
                    layer_cache_types=["KVCache"] * 2,
                )

            # Access block 0 to refresh its LRU position
            mgr.load_block(b"order_block_0")

            # Save block 2 — should evict block 1 (LRU), not block 0
            cache_data = self._make_cache_data()
            mgr.save_block(
                block_hash=b"order_block_2",
                cache_data=cache_data,
                token_count=16,
                model_name="test",
                layer_cache_types=["KVCache"] * 2,
            )

            # Block 0 was accessed so should still be present
            assert mgr._hot_cache_get(b"order_block_0") is not None
            # Block 1 was LRU and should be evicted
            assert mgr._hot_cache_get(b"order_block_1") is None
            # Block 2 was just added
            assert mgr._hot_cache_get(b"order_block_2") is not None
        finally:
            mgr.close()


@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
class TestHotCachePromotion:
    """Test promotion from SSD to hot cache on load."""

    def _make_cache_data(self, num_layers=4, seq_len=32, heads=4, head_dim=32):
        return [
            (
                mx.zeros((1, heads, seq_len, head_dim)),
                mx.zeros((1, heads, seq_len, head_dim)),
            )
            for _ in range(num_layers)
        ]

    def test_ssd_load_promotes_to_hot_cache(self, tmp_path):
        """Loading a block from SSD should promote it to hot cache."""
        # Use hot_cache disabled to write directly to SSD first
        mgr_cold = PagedSSDCacheManager(
            cache_dir=tmp_path / "promote_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=0,
        )

        block_hash = b"promote_test_block1"
        cache_data = self._make_cache_data()
        mgr_cold.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=64,
            model_name="test",
            layer_cache_types=["KVCache"] * 4,
        )
        # Wait for background SSD write to complete
        time.sleep(0.5)
        mgr_cold.close()

        # Now open with hot cache enabled — block is on SSD only
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "promote_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=10 * 1024**2,
        )

        try:
            assert mgr._hot_cache_get(block_hash) is None

            # Load from SSD — should promote to hot cache
            loaded = mgr.load_block(block_hash)
            assert loaded is not None
            assert len(loaded) == 4

            # Verify promotion happened
            assert mgr._hot_cache_get(block_hash) is not None
            stats = mgr.get_stats()
            assert stats.hot_cache_promotions >= 1
        finally:
            mgr.close()

    def test_ssd_load_can_skip_hot_cache_promotion(self, tmp_path):
        """SSD loads can reconstruct a request without retaining a hot RAM copy."""
        mgr_cold = PagedSSDCacheManager(
            cache_dir=tmp_path / "skip_promote_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=0,
        )

        block_hash = b"skip_promote_blk_1"
        cache_data = self._make_cache_data()
        mgr_cold.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=64,
            model_name="test",
            layer_cache_types=["KVCache"] * 4,
        )
        mgr_cold.close()

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "skip_promote_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=10 * 1024**2,
        )

        try:
            loaded, metadata = mgr.load_block_with_metadata(
                block_hash,
                promote_to_hot_cache=False,
            )

            assert loaded is not None
            assert metadata is not None
            assert mgr._hot_cache_get(block_hash) is None

            loaded_again, _ = mgr.load_block_with_metadata(block_hash)
            assert loaded_again is not None
            assert mgr._hot_cache_get(block_hash) is not None
        finally:
            mgr.close()

    def test_write_through_save_bypasses_hot_cache_retention(self, tmp_path):
        """Pressure write-through keeps SSD durability without hot-cache RAM."""
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "write_through_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=10 * 1024**2,
        )

        try:
            block_hash = b"write_through_blk1"
            cache_data = self._make_cache_data()
            result = mgr.save_block(
                block_hash=block_hash,
                cache_data=cache_data,
                token_count=64,
                model_name="test",
                layer_cache_types=["KVCache"] * 4,
                hot_cache_write_back=False,
            )

            assert result is True
            assert mgr.has_block(block_hash)
            assert mgr._hot_cache_get(block_hash) is None

            loaded, metadata = mgr.load_block_with_metadata(
                block_hash,
                promote_to_hot_cache=False,
            )
            assert loaded is not None
            assert metadata is not None
            assert mgr._hot_cache_get(block_hash) is None
        finally:
            mgr.close()

    def test_promotion_does_not_happen_when_disabled(self, tmp_path):
        """No promotion when hot cache is disabled."""
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "no_promote_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=0,
        )

        try:
            block_hash = b"no_promote_block1__"
            cache_data = self._make_cache_data()
            mgr.save_block(
                block_hash=block_hash,
                cache_data=cache_data,
                token_count=64,
                model_name="test",
                layer_cache_types=["KVCache"] * 4,
            )

            # Wait for background write
            time.sleep(0.5)

            # Clear the temporary buffer (simulates what happens after write completes)
            mgr._hot_cache_remove(block_hash)

            # Load from SSD
            loaded = mgr.load_block(block_hash)
            assert loaded is not None

            stats = mgr.get_stats()
            assert stats.hot_cache_promotions == 0
        finally:
            mgr.close()


@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
class TestHotCacheCacheTypes:
    """Test hot cache with various cache types (KVCache, CacheList)."""

    @pytest.fixture
    def manager(self, tmp_path):
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "types_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=10 * 1024**2,
        )
        yield mgr
        mgr.close()

    def test_cache_list_blocks(self, manager):
        """Hot cache should handle CacheList blocks correctly."""
        block_hash = b"cache_list_hot_test"

        sub_keys1 = mx.zeros((1, 8, 32, 64))
        sub_values1 = mx.ones((1, 8, 32, 64))
        sub_keys2 = mx.zeros((1, 4, 32, 64))
        sub_values2 = mx.ones((1, 4, 32, 64))

        cache_data = [
            (
                "__cache_list__",
                [(sub_keys1, sub_values1), (sub_keys2, sub_values2)],
            ),
            (mx.zeros((1, 8, 32, 64)), mx.ones((1, 8, 32, 64))),
        ]

        result = manager.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=32,
            model_name="test",
            layer_cache_types=["CacheList", "KVCache"],
        )
        assert result is True

        # Load from hot cache
        loaded = manager.load_block(block_hash)
        assert loaded is not None
        assert len(loaded) == 2
        # First layer is CacheList
        assert isinstance(loaded[0], list)
        assert len(loaded[0]) == 2
        # Second layer is KVCache tuple
        assert isinstance(loaded[1], tuple)


class TestHotCacheConcurrency:
    """Test thread safety of hot cache internal operations."""

    def test_concurrent_put_get(self, tmp_path):
        """Hot cache put/get should be thread-safe under concurrent access."""
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "concurrent_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=50 * 1024**2,
        )

        errors: List[Exception] = []
        num_threads = 8
        ops_per_thread = 20

        def worker(thread_id):
            try:
                for i in range(ops_per_thread):
                    block_hash = f"conc_{thread_id}_{i}____".encode()
                    # Create a fake hot cache entry with raw bytes
                    raw_data = bytes(1024)  # 1KB of zeros
                    entry = {
                        "tensors_raw": {
                            "layer_0_keys": (raw_data, "float32", [1, 2, 16, 8]),
                            "layer_0_values": (raw_data, "float32", [1, 2, 16, 8]),
                        },
                        "file_metadata": {},
                        "num_layers": 1,
                        "layer_cache_types": ["KVCache"],
                        "block_metadata": None,
                    }
                    mgr._hot_cache_put(block_hash, entry)

                # Read back
                for i in range(ops_per_thread):
                    block_hash = f"conc_{thread_id}_{i}____".encode()
                    result = mgr._hot_cache_get(block_hash)
                    # May be None if evicted by another thread, that's OK
            except Exception as e:
                errors.append(e)

        threads = [
            threading.Thread(target=worker, args=(t,)) for t in range(num_threads)
        ]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=30)

        mgr.close()

        assert len(errors) == 0, f"Concurrent errors: {errors}"


class TestHotCacheByteAccounting:
    """Regression tests for raw-byte hot cache size accounting."""

    def _make_raw_entry(
        self,
        tmp_path: Path,
        block_hash: bytes,
        model_name: str = "test-model",
        key_size: int = 128,
        value_size: int = 256,
    ):
        metadata = PagedSSDBlockMetadata(
            block_hash=block_hash,
            file_path=tmp_path / f"{block_hash.hex()}.safetensors",
            file_size=key_size + value_size,
            token_count=16,
            created_at=time.time(),
            last_access=time.time(),
            num_layers=1,
            model_name=model_name,
            layer_cache_types=["KVCache"],
        )
        entry = {
            "tensors_raw": {
                "layer_0_keys": (bytes(key_size), "float32", [1, 1, 1, key_size // 4]),
                "layer_0_values": (
                    bytes(value_size),
                    "float32",
                    [1, 1, 1, value_size // 4],
                ),
            },
            "file_metadata": {},
            "num_layers": 1,
            "layer_cache_types": ["KVCache"],
            "block_metadata": metadata,
        }
        return entry, key_size + value_size

    def test_hot_cache_remove_decrements_raw_entry_bytes(self, tmp_path):
        """Removing a tensors_raw entry should subtract its bytes from the counter."""
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "remove_accounting",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=1024**2,
            hot_cache_only=True,
        )

        try:
            block_hash = b"remove_bytes_test"
            entry, expected_size = self._make_raw_entry(tmp_path, block_hash)

            mgr._hot_cache_put(block_hash, entry)
            assert mgr._hot_cache_total_bytes == expected_size

            mgr._hot_cache_remove(block_hash)

            assert mgr._hot_cache_get(block_hash) is None
            assert mgr._hot_cache_total_bytes == 0
            assert mgr.get_stats().hot_cache_size_bytes == 0
        finally:
            mgr.close()

    def test_get_stats_for_model_reports_raw_hot_cache_bytes(self, tmp_path):
        """Per-model stats should count tensors_raw entries in the hot cache."""
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "model_accounting",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=1024**2,
            hot_cache_only=True,
        )

        try:
            model_a_hash = b"model_a_hot_bytes"
            model_b_hash = b"model_b_hot_bytes"
            model_a_entry, model_a_size = self._make_raw_entry(
                tmp_path, model_a_hash, model_name="model-a"
            )
            model_b_entry, _ = self._make_raw_entry(
                tmp_path, model_b_hash, model_name="model-b"
            )

            mgr._hot_cache_put(model_a_hash, model_a_entry)
            mgr._hot_cache_put(model_b_hash, model_b_entry)

            stats = mgr.get_stats_for_model("model-a")

            assert stats.hot_cache_entries == 1
            assert stats.hot_cache_size_bytes == model_a_size
        finally:
            mgr.close()

    def test_shared_budget_evicts_across_managers_by_global_lru(self, tmp_path):
        """A shared hot cache budget should cap aggregate manager memory."""
        entry, entry_size = self._make_raw_entry(tmp_path, b"budget_size_probe")
        del entry
        budget = SharedHotCacheBudget(entry_size * 2 + 1)
        mgr_a = PagedSSDCacheManager(
            cache_dir=tmp_path / "budget_a",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=budget.max_bytes,
            hot_cache_only=True,
            hot_cache_budget=budget,
        )
        mgr_b = PagedSSDCacheManager(
            cache_dir=tmp_path / "budget_b",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=budget.max_bytes,
            hot_cache_only=True,
            hot_cache_budget=budget,
        )

        try:
            a0, _ = self._make_raw_entry(tmp_path, b"budget_a_0", model_name="model-a")
            b0, _ = self._make_raw_entry(tmp_path, b"budget_b_0", model_name="model-b")
            b1, _ = self._make_raw_entry(tmp_path, b"budget_b_1", model_name="model-b")

            mgr_a._hot_cache_put(b"budget_a_0", a0)
            mgr_b._hot_cache_put(b"budget_b_0", b0)
            assert budget.total_bytes == entry_size * 2

            # Refresh A, so B's first block is the global LRU victim.
            assert mgr_a._hot_cache_get(b"budget_a_0") is not None
            mgr_b._hot_cache_put(b"budget_b_1", b1)

            assert budget.total_bytes == entry_size * 2
            assert mgr_a._hot_cache_get(b"budget_a_0") is not None
            assert mgr_b._hot_cache_get(b"budget_b_0") is None
            assert mgr_b._hot_cache_get(b"budget_b_1") is not None
            assert mgr_a.get_stats().hot_cache_max_bytes == budget.max_bytes
            assert mgr_b.get_stats().hot_cache_max_bytes == budget.max_bytes
        finally:
            mgr_a.close()
            mgr_b.close()

    def test_shared_budget_forgets_owner_on_clear_and_close(self, tmp_path):
        """Clearing or closing a manager must release shared budget bytes."""
        _, entry_size = self._make_raw_entry(tmp_path, b"budget_clear_probe")
        budget = SharedHotCacheBudget(entry_size * 4)
        mgr_a = PagedSSDCacheManager(
            cache_dir=tmp_path / "budget_clear_a",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=budget.max_bytes,
            hot_cache_only=True,
            hot_cache_budget=budget,
        )
        mgr_b = PagedSSDCacheManager(
            cache_dir=tmp_path / "budget_clear_b",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=budget.max_bytes,
            hot_cache_only=True,
            hot_cache_budget=budget,
        )

        try:
            a0, _ = self._make_raw_entry(tmp_path, b"budget_clear_a_0")
            b0, _ = self._make_raw_entry(tmp_path, b"budget_clear_b_0")
            mgr_a._hot_cache_put(b"budget_clear_a_0", a0)
            mgr_b._hot_cache_put(b"budget_clear_b_0", b0)

            assert budget.total_bytes == entry_size * 2
            assert mgr_a.clear_hot_cache() == 1
            assert budget.total_bytes == entry_size

            mgr_b.close()
            assert budget.total_bytes == 0
        finally:
            mgr_a.close()
            mgr_b.close()

    def test_shared_budget_shrink_respects_lru_and_protected_hashes(self, tmp_path):
        """Pressure shrink should skip protected active-request chains."""
        _, entry_size = self._make_raw_entry(tmp_path, b"budget_shrink_probe")
        budget = SharedHotCacheBudget(entry_size * 3)
        mgr_a = PagedSSDCacheManager(
            cache_dir=tmp_path / "budget_shrink_a",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=budget.max_bytes,
            hot_cache_only=True,
            hot_cache_budget=budget,
        )
        mgr_b = PagedSSDCacheManager(
            cache_dir=tmp_path / "budget_shrink_b",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=budget.max_bytes,
            hot_cache_only=True,
            hot_cache_budget=budget,
        )

        try:
            a0, _ = self._make_raw_entry(tmp_path, b"budget_shrink_a0")
            b0, _ = self._make_raw_entry(tmp_path, b"budget_shrink_b0")
            b1, _ = self._make_raw_entry(tmp_path, b"budget_shrink_b1")
            mgr_a._hot_cache_put(b"budget_shrink_a0", a0)
            mgr_b._hot_cache_put(b"budget_shrink_b0", b0)
            mgr_b._hot_cache_put(b"budget_shrink_b1", b1)

            freed = budget.shrink_to(
                entry_size,
                protected_hashes={b"budget_shrink_b0"},
            )

            assert freed == entry_size * 2
            assert budget.total_bytes == entry_size
            assert mgr_a._hot_cache_get(b"budget_shrink_a0") is None
            assert mgr_b._hot_cache_get(b"budget_shrink_b0") is not None
            assert mgr_b._hot_cache_get(b"budget_shrink_b1") is None
        finally:
            mgr_a.close()
            mgr_b.close()

    def test_local_hot_cache_shrink_respects_protected_hashes(self, tmp_path):
        _, entry_size = self._make_raw_entry(tmp_path, b"local_shrink_probe")
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "local_shrink",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=entry_size * 3,
            hot_cache_only=True,
        )

        try:
            h0 = b"local_shrink_0"
            h1 = b"local_shrink_1"
            h2 = b"local_shrink_2"
            e0, _ = self._make_raw_entry(tmp_path, h0)
            e1, _ = self._make_raw_entry(tmp_path, h1)
            e2, _ = self._make_raw_entry(tmp_path, h2)
            mgr._hot_cache_put(h0, e0)
            mgr._hot_cache_put(h1, e1)
            mgr._hot_cache_put(h2, e2)

            freed = mgr.shrink_hot_cache_to(entry_size, protected_hashes={h1})

            assert freed == entry_size * 2
            assert mgr.get_stats().hot_cache_size_bytes == entry_size
            assert mgr._hot_cache_get(h0) is None
            assert mgr._hot_cache_get(h1) is not None
            assert mgr._hot_cache_get(h2) is None
        finally:
            mgr.close()


@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
class TestHotCacheStatsAccuracy:
    """Test that hot cache statistics are accurate."""

    def test_all_stats_counters(self, tmp_path):
        """Verify hot cache stats counters are correctly maintained."""
        entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4  # ~4096 bytes per entry
        # Room for 2 entries
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "stats_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )

        try:

            def save(idx):
                block_hash = f"stats_block_{idx}__".encode()
                cache_data = [
                    (mx.zeros((1, 2, 16, 16)), mx.zeros((1, 2, 16, 16)))
                    for _ in range(2)
                ]
                mgr.save_block(
                    block_hash=block_hash,
                    cache_data=cache_data,
                    token_count=16,
                    model_name="test",
                    layer_cache_types=["KVCache"] * 2,
                )

            # Save 2 blocks: fits in hot cache
            save(0)
            save(1)
            stats = mgr.get_stats()
            assert stats.hot_cache_entries == 2
            assert stats.hot_cache_evictions == 0

            # Save 3rd block: triggers eviction of block 0
            save(2)
            stats = mgr.get_stats()
            assert stats.hot_cache_entries == 2
            assert stats.hot_cache_evictions >= 1

            # Load block 1 (hot cache hit)
            mgr.load_block(b"stats_block_1__")
            stats = mgr.get_stats()
            assert stats.hot_cache_hits >= 1

            # Verify size tracking is positive
            assert stats.hot_cache_size_bytes > 0
            assert stats.hot_cache_max_bytes == max_bytes
        finally:
            mgr.close()


@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
class TestHotCacheWriteBack:
    """Test write-back behavior: no SSD writes until eviction or shutdown."""

    def _make_cache_data(self, num_layers=2, seq_len=16, heads=2, head_dim=16):
        return [
            (
                mx.zeros((1, heads, seq_len, head_dim)),
                mx.zeros((1, heads, seq_len, head_dim)),
            )
            for _ in range(num_layers)
        ]

    def test_save_does_not_write_to_ssd(self, tmp_path):
        """With hot cache enabled, save_block should not create SSD files."""
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "wb_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=10 * 1024**2,
        )

        try:
            block_hash = b"wb_no_ssd_write_t1"
            cache_data = self._make_cache_data()
            mgr.save_block(
                block_hash=block_hash,
                cache_data=cache_data,
                token_count=16,
                model_name="test",
                layer_cache_types=["KVCache"] * 2,
            )

            # Block should be in hot cache
            assert mgr._hot_cache_get(block_hash) is not None

            # No SSD file should exist yet
            time.sleep(0.3)
            ssd_files = list((tmp_path / "wb_test").rglob("*.safetensors"))
            assert len(ssd_files) == 0, f"Unexpected SSD files: {ssd_files}"
        finally:
            mgr.close()

    def test_eviction_writes_to_ssd(self, tmp_path):
        """When hot cache evicts, the evicted block should be written to SSD."""
        entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "wb_evict_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )

        try:
            # Save 2 blocks (fits in hot cache)
            for i in range(2):
                block_hash = f"wb_evict_blk_{i}__".encode()
                cache_data = self._make_cache_data()
                mgr.save_block(
                    block_hash=block_hash,
                    cache_data=cache_data,
                    token_count=16,
                    model_name="test",
                    layer_cache_types=["KVCache"] * 2,
                )

            # No SSD files yet
            time.sleep(0.3)
            ssd_files = list((tmp_path / "wb_evict_test").rglob("*.safetensors"))
            assert len(ssd_files) == 0

            # Save 3rd block → evicts block 0 → should trigger SSD write
            cache_data = self._make_cache_data()
            mgr.save_block(
                block_hash=b"wb_evict_blk_2__",
                cache_data=cache_data,
                token_count=16,
                model_name="test",
                layer_cache_types=["KVCache"] * 2,
            )

            # Wait for background writer to process the evicted block
            time.sleep(0.5)
            ssd_files = list((tmp_path / "wb_evict_test").rglob("*.safetensors"))
            assert len(ssd_files) >= 1, "Evicted block should be written to SSD"
        finally:
            mgr.close()

    def test_close_flushes_hot_cache_to_ssd(self, tmp_path):
        """close() should flush all hot cache entries to SSD."""
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "wb_flush_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=10 * 1024**2,
        )

        block_hashes = []
        for i in range(3):
            block_hash = f"wb_flush_blk_{i}__".encode()
            block_hashes.append(block_hash)
            cache_data = self._make_cache_data()
            mgr.save_block(
                block_hash=block_hash,
                cache_data=cache_data,
                token_count=16,
                model_name="test",
                layer_cache_types=["KVCache"] * 2,
            )

        # No SSD files before close
        time.sleep(0.3)
        ssd_files = list((tmp_path / "wb_flush_test").rglob("*.safetensors"))
        assert len(ssd_files) == 0

        # Close flushes to SSD
        mgr.close()

        ssd_files = list((tmp_path / "wb_flush_test").rglob("*.safetensors"))
        assert (
            len(ssd_files) == 3
        ), f"Expected 3 SSD files after flush, got {len(ssd_files)}"

    def test_close_flushes_all_blocks_with_small_queue(self, tmp_path):
        """close() must flush all hot cache blocks even when more blocks
        exist than the write queue depth.

        Regression test for #1070: put_nowait() in the shutdown flush loop
        lost blocks when the bounded write queue filled up faster than the
        writer thread can drain it.
        """
        queue_depth = 4
        block_count = 12  # 3x the queue depth

        with patch("omlx.cache.paged_ssd_cache._MAX_PENDING_WRITES", queue_depth):
            mgr = PagedSSDCacheManager(
                cache_dir=tmp_path / "wb_queue_full_test",
                max_size_bytes=100 * 1024**2,
                hot_cache_max_bytes=10 * 1024**2,
            )

        for i in range(block_count):
            block_hash = f"wb_qfull_blk_{i:02d}".encode()
            cache_data = self._make_cache_data()
            mgr.save_block(
                block_hash=block_hash,
                cache_data=cache_data,
                token_count=16,
                model_name="test",
                layer_cache_types=["KVCache"] * 2,
            )

        # All blocks in hot cache, no SSD files yet
        time.sleep(0.3)
        ssd_files = list((tmp_path / "wb_queue_full_test").rglob("*.safetensors"))
        assert len(ssd_files) == 0

        mgr.close()

        ssd_files = list((tmp_path / "wb_queue_full_test").rglob("*.safetensors"))
        assert len(ssd_files) == block_count, (
            f"Expected {block_count} SSD files after flush, got {len(ssd_files)}. "
            f"Blocks were likely not flushed during shutdown."
        )


@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
class TestPendingWriteBuffer:
    """Tests for the pending-write buffer that bridges hot cache eviction to SSD write."""

    def _make_cache_data(self, num_layers=2, seq_len=16, heads=2, head_dim=16):
        return [
            (
                mx.zeros((1, heads, seq_len, head_dim)),
                mx.zeros((1, heads, seq_len, head_dim)),
            )
            for _ in range(num_layers)
        ]

    def _save_block(self, manager, block_hash, model="test-model"):
        cache_data = self._make_cache_data()
        return manager.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=16,
            model_name=model,
            layer_cache_types=["KVCache"] * 2,
        )

    def test_evicted_block_readable_from_pending_buffer(self, tmp_path):
        """A block evicted from hot cache is still loadable while SSD write is pending."""
        entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "pending_buf_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )
        try:
            # Save 2 blocks (fills hot cache)
            self._save_block(mgr, b"pending_buf_blk0")
            self._save_block(mgr, b"pending_buf_blk1")

            # Save 3rd block → evicts block 0
            self._save_block(mgr, b"pending_buf_blk2")

            # Block 0 should still be loadable (from pending buffer, not SSD)
            with mgr._hot_cache_lock:
                assert (
                    b"pending_buf_blk0" not in mgr._hot_cache
                ), "Block 0 should have been evicted from hot cache"
            loaded = mgr.load_block(b"pending_buf_blk0")
            assert (
                loaded is not None
            ), "Evicted block should be readable from pending write buffer"
            assert len(loaded) == 2
        finally:
            mgr.close()

    def test_writer_cleanup_empties_buffer(self, tmp_path):
        """After background writer completes, pending buffer entry is removed."""
        entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "writer_cleanup_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )
        try:
            self._save_block(mgr, b"cleanup_test_blk0")
            self._save_block(mgr, b"cleanup_test_blk1")
            # Evict block 0
            self._save_block(mgr, b"cleanup_test_blk2")

            # Block 0 should be in pending buffer
            with mgr._pending_write_hashes_lock:
                assert b"cleanup_test_blk0" in mgr._pending_write_buffers

            # Wait for writer to finish
            time.sleep(1.0)

            # Buffer should be empty after writer cleanup
            with mgr._pending_write_hashes_lock:
                assert b"cleanup_test_blk0" not in mgr._pending_write_buffers
                assert b"cleanup_test_blk0" not in mgr._pending_write_hashes
        finally:
            mgr.close()

    def test_has_block_sees_pending_entries(self, tmp_path):
        """has_block() returns True for blocks in the pending write buffer."""
        entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "has_block_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )
        try:
            self._save_block(mgr, b"has_block_test_b0")
            self._save_block(mgr, b"has_block_test_b1")
            # Evict block 0
            self._save_block(mgr, b"has_block_test_b2")

            with mgr._hot_cache_lock:
                assert b"has_block_test_b0" not in mgr._hot_cache

            assert (
                mgr.has_block(b"has_block_test_b0") is True
            ), "has_block should find blocks in the pending write buffer"
        finally:
            mgr.close()

    def test_delete_block_clears_pending_buffer(self, tmp_path):
        """delete_block() removes the entry from pending write buffer."""
        entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "delete_pending_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )
        try:
            self._save_block(mgr, b"del_pending_blk_0")
            self._save_block(mgr, b"del_pending_blk_1")
            # Evict block 0
            self._save_block(mgr, b"del_pending_blk_2")

            # Block 0 is in pending buffer
            with mgr._pending_write_hashes_lock:
                assert b"del_pending_blk_0" in mgr._pending_write_buffers

            # Delete it
            mgr.delete_block(b"del_pending_blk_0")

            # Should no longer be loadable
            loaded = mgr.load_block(b"del_pending_blk_0")
            assert (
                loaded is None
            ), "Deleted block should not be readable from pending buffer"
        finally:
            mgr.close()

    def test_queue_full_inline_fallback_cleans_pending_buffer(self, tmp_path):
        """Inline fallback removes the block from pending after it reaches SSD."""
        import queue as _queue

        entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "queue_full_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )
        try:
            self._save_block(mgr, b"qf_test_block_00")
            self._save_block(mgr, b"qf_test_block_01")
            self._save_block(mgr, b"qf_test_block_02")
            # This evicts block 1. Force sustained saturation so the bounded
            # wait expires and inline fallback runs.
            with patch.object(mgr._write_queue, "put", side_effect=_queue.Full):
                self._save_block(mgr, b"qf_test_block_03")

            # Block 1 was written inline and no longer needs pending storage.
            with mgr._pending_write_hashes_lock:
                assert (
                    b"qf_test_block_01" not in mgr._pending_write_buffers
                ), "Inline-written block should leave the pending buffer"
                assert (
                    b"qf_test_block_01" not in mgr._pending_write_hashes
                ), "Inline-written block should leave pending hashes"
            stats = mgr.get_stats()
            assert stats.ssd_write_drops == 0
            assert stats.ssd_inline_write_fallbacks == 1
            assert mgr.load_block(b"qf_test_block_01") is not None
        finally:
            mgr.close()

    def test_evicted_block_loadable_with_metadata(self, tmp_path):
        """load_block_with_metadata returns data for pending-buffer blocks."""
        entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "meta_load_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )
        try:
            self._save_block(mgr, b"meta_load_blk_00")
            self._save_block(mgr, b"meta_load_blk_01")
            # Evict block 0
            self._save_block(mgr, b"meta_load_blk_02")

            cache_data, metadata = mgr.load_block_with_metadata(b"meta_load_blk_00")
            assert (
                cache_data is not None
            ), "load_block_with_metadata should serve evicted block from pending buffer"
            assert metadata is not None
            assert metadata["num_layers"] == 2
            assert metadata["token_count"] == 16
            assert metadata["model_name"] == "test-model"
        finally:
            mgr.close()

    def test_end_to_end_lifecycle(self, tmp_path):
        """Full lifecycle: save → evict → pending hit → writer completes → SSD hit."""
        entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4
        max_bytes = entry_size * 2 + 100

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "lifecycle_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )
        try:
            self._save_block(mgr, b"lifecycle_blk_00")
            self._save_block(mgr, b"lifecycle_blk_01")
            # Evict block 0
            self._save_block(mgr, b"lifecycle_blk_02")

            # Phase 1: pending buffer hit (before writer completes)
            loaded = mgr.load_block(b"lifecycle_blk_00")
            assert loaded is not None, "Should load from pending buffer"
            assert mgr._stats["hot_cache_hits"] >= 1

            # Phase 2: wait for writer to complete
            time.sleep(1.0)

            # Buffer should be empty
            with mgr._pending_write_hashes_lock:
                assert b"lifecycle_blk_00" not in mgr._pending_write_buffers

            # Phase 3: SSD hit (block now on disk)
            loaded_ssd = mgr.load_block(b"lifecycle_blk_00")
            assert loaded_ssd is not None, "Should load from SSD after writer completes"
            assert len(loaded_ssd) == 2
        finally:
            mgr.close()


@pytest.mark.skipif(not HAS_MLX, reason="MLX not available")
class TestSSDWriteBackSaturation:
    """Queue saturation falls back to inline writes instead of dropping blocks."""

    # Mirrors TestPendingWriteBuffer._make_cache_data — small dimensions
    # so the per-entry footprint is small and predictable for queue tests.
    def _make_cache_data(self, num_layers=2, seq_len=16, heads=2, head_dim=16):
        return [
            (
                mx.zeros((1, heads, seq_len, head_dim)),
                mx.zeros((1, heads, seq_len, head_dim)),
            )
            for _ in range(num_layers)
        ]

    # Mirrors TestPendingWriteBuffer._save_block — uses 2 layers, token_count=16
    # so the entry_size formula `2 * 2 * 1 * 2 * 16 * 16 * 4` calibrates Test A.
    def _save_block(self, manager, block_hash, model="test-model"):
        cache_data = self._make_cache_data()
        return manager.save_block(
            block_hash=block_hash,
            cache_data=cache_data,
            token_count=16,
            model_name=model,
            layer_cache_types=["KVCache"] * 2,
        )

    def test_paged_ssd_cache_stats_default_and_reset(self):
        """Dataclass: write-back counters default to 0 and reset to 0."""
        from omlx.cache.stats import PagedSSDCacheStats

        # Default is zero.
        stats = PagedSSDCacheStats()
        assert stats.ssd_write_drops == 0
        assert stats.ssd_inline_write_fallbacks == 0
        assert stats.evict_unlink_failures == 0

        # reset() returns it to zero from a non-zero state.
        stats = PagedSSDCacheStats(
            ssd_write_drops=5,
            ssd_inline_write_fallbacks=4,
            evict_unlink_failures=3,
            saves=2,
            loads=3,
        )
        stats.reset()
        assert stats.ssd_write_drops == 0
        assert stats.ssd_inline_write_fallbacks == 0
        assert stats.evict_unlink_failures == 0
        # Verify reset() didn't break the existing fields it already handled.
        assert stats.saves == 0
        assert stats.loads == 0

    def test_ssd_write_back_fields_round_trip_through_get_stats(self, tmp_path):
        """The _stats dict values flow through both stats accessors.

        Force a non-zero value into _stats so a missing pass-through in either
        accessor would leave the dataclass at the default 0 and fail this test.
        The default-0 case is covered by the dataclass-level test in Task 1;
        this test exercises the wiring specifically.
        """
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "wiring_test",
            max_size_bytes=100 * 1024**2,
        )
        try:
            # Directly poke a non-zero value to test the accessor pass-through
            # specifically. The real increment sites have their own tests
            # further down in this class; this test would catch a regression
            # where a future refactor drops the field from get_stats() or
            # get_stats_for_model() but leaves the _stats dict key alone.
            mgr._stats["ssd_write_drops"] = 7
            mgr._stats["ssd_inline_write_fallbacks"] = 5

            # Global stats accessor.
            stats = mgr.get_stats()
            assert stats.ssd_write_drops == 7, (
                "get_stats() must pass _stats['ssd_write_drops'] through "
                "to the dataclass field"
            )
            assert stats.ssd_inline_write_fallbacks == 5, (
                "get_stats() must pass _stats['ssd_inline_write_fallbacks'] "
                "through to the dataclass field"
            )

            # Per-model stats accessor — same wiring, separate code path.
            # Use any model name; the field is global on the manager.
            model_stats = mgr.get_stats_for_model("any-model-name")
            assert model_stats.ssd_write_drops == 7, (
                "get_stats_for_model() must pass _stats['ssd_write_drops'] "
                "through to the dataclass field"
            )
            assert model_stats.ssd_inline_write_fallbacks == 5, (
                "get_stats_for_model() must pass "
                "_stats['ssd_inline_write_fallbacks'] through to the "
                "dataclass field"
            )
        finally:
            mgr.close()

    def test_evict_unlink_failures_round_trips_through_get_stats(self, tmp_path):
        """evict_unlink_failures is visible through both stats accessors."""
        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "unlink_failure_wiring_test",
            max_size_bytes=100 * 1024**2,
        )
        try:
            mgr._stats["evict_unlink_failures"] = 3

            stats = mgr.get_stats()
            assert stats.evict_unlink_failures == 3

            model_stats = mgr.get_stats_for_model("any-model-name")
            assert model_stats.evict_unlink_failures == 3
        finally:
            mgr.close()

    def test_hot_eviction_queue_full_uses_inline_fallback(self, tmp_path):
        """Site 1: hot-cache eviction → put raises queue.Full → inline write.

        Patches the real queue's put to raise queue.Full, guaranteeing the
        fallback path fires on the first eviction without any dependency on the
        writer thread's drain rate. _enqueue_ssd_write uses put(item,
        timeout=...) (not put_nowait) so a transient burst can ride over a
        short writer-backlog window; sustained saturation writes inline.
        """
        import queue as _queue
        from unittest.mock import patch

        # entry_size matches the small-dimension helpers above (num_layers=2,
        # K+V=2 tensors, batch=1, heads=2, seq=16, head_dim=16, float32=4).
        entry_size = 2 * 2 * 1 * 2 * 16 * 16 * 4
        max_bytes = entry_size * 2 + 100  # holds exactly 2 entries

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "drops_hot_eviction_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=max_bytes,
        )
        try:
            with patch.object(mgr._write_queue, "put", side_effect=_queue.Full):
                self._save_block(mgr, b"qf_drop_block_00")
                self._save_block(mgr, b"qf_drop_block_01")
                # save_02 evicts block 00 → _enqueue_ssd_write → put raises
                # queue.Full → inline fallback writes it immediately.
                self._save_block(mgr, b"qf_drop_block_02")

            stats = mgr.get_stats()
            assert stats.ssd_write_drops == 0
            assert stats.ssd_inline_write_fallbacks == 1
            assert stats.saves_persisted == 1
            assert stats.errors == 0

            # Block 00 was the one being enqueued when put raised. It should
            # no longer need the pending buffer because it is already on SSD.
            with mgr._pending_write_hashes_lock:
                assert b"qf_drop_block_00" not in mgr._pending_write_buffers
                assert b"qf_drop_block_00" not in mgr._pending_write_hashes
            metadata = mgr._index.get(b"qf_drop_block_00")
            assert metadata is not None
            assert metadata.file_path.exists()
            assert mgr.load_block(b"qf_drop_block_00") is not None
        finally:
            mgr.close()

    def test_cold_store_sustained_full_uses_inline_fallback(self, tmp_path):
        """Site 2: save_block waits on a full queue before writing inline.

        Hot cache disabled. Even if ``full()`` reports saturation, the cold
        path must still proceed to ``put(timeout=1.0)`` so transient writer
        bursts get a chance to drain. Sustained ``queue.Full`` from the put
        call should write inline and return success.
        """
        import queue as _queue
        from unittest.mock import patch

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "drops_cold_preflight_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=0,  # hot cache disabled → cold-store path
        )
        try:
            calls: list[float | None] = []

            def full_put(*args, timeout=None, **kwargs):
                calls.append(timeout)
                raise _queue.Full

            with (
                patch.object(mgr._write_queue, "full", return_value=True),
                patch.object(mgr._write_queue, "put", side_effect=full_put),
            ):
                cache_data = self._make_cache_data()
                block_hash = b"cold_preflight_drop_00"
                ok = mgr.save_block(
                    block_hash=block_hash,
                    cache_data=cache_data,
                    token_count=16,
                    model_name="test-model",
                    layer_cache_types=["KVCache"] * 2,
                )
                assert ok is True

            assert calls == [1.0]
            stats = mgr.get_stats()
            assert stats.ssd_write_drops == 0
            assert stats.ssd_inline_write_fallbacks == 1
            assert stats.saves == 1
            assert stats.saves_persisted == 1
            assert stats.errors == 0
            assert mgr._index.contains(block_hash)
            with mgr._pending_write_hashes_lock:
                assert block_hash not in mgr._pending_write_hashes
            assert mgr.load_block(block_hash) is not None
        finally:
            mgr.close()

    def test_cold_store_late_queue_full_uses_inline_fallback(self, tmp_path):
        """Site 3: save_block put raises queue.Full after the preflight passes.

        Hot cache disabled. ``put`` is patched to raise queue.Full directly
        (simulating a sustained writer-backlog saturation that materializes
        after the preflight check). The block must still be written inline
        and remain readable.
        """
        import queue as _queue
        from unittest.mock import patch

        mgr = PagedSSDCacheManager(
            cache_dir=tmp_path / "drops_cold_late_exception_test",
            max_size_bytes=100 * 1024**2,
            hot_cache_max_bytes=0,
        )
        try:
            cache_data = self._make_cache_data()
            block_hash = b"cold_late_drop_00"
            with patch.object(mgr._write_queue, "put", side_effect=_queue.Full):
                ok = mgr.save_block(
                    block_hash=block_hash,
                    cache_data=cache_data,
                    token_count=16,
                    model_name="test-model",
                    layer_cache_types=["KVCache"] * 2,
                )
                assert ok is True

            stats = mgr.get_stats()
            assert stats.ssd_write_drops == 0
            assert stats.ssd_inline_write_fallbacks == 1
            assert stats.saves == 1
            assert stats.saves_persisted == 1
            assert mgr._index.contains(block_hash)
            with mgr._pending_write_hashes_lock:
                assert block_hash not in mgr._pending_write_hashes
            assert mgr.load_block(block_hash) is not None
        finally:
            mgr.close()