litellm/tests/unit/router_utils/test_cooldown_cache.py

"""
Unit tests for CooldownCache exception masking functionality
"""

import time
from unittest.mock import MagicMock

import pytest

# Add the parent directory to the system path
from litellm._internal_context import current_service_target
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue


class TestCooldownCacheExceptionMasking:
    """Test suite for CooldownCache exception masking functionality"""

    @pytest.fixture
    def cooldown_cache(self):
        """Create a CooldownCache instance for testing"""
        mock_dual_cache = MagicMock(spec=DualCache)
        return CooldownCache(cache=mock_dual_cache, default_cooldown_time=60.0)

    def test_exception_masker_initialization(self, cooldown_cache):
        """Test that the exception masker is properly initialized"""
        assert isinstance(cooldown_cache.exception_masker, SensitiveDataMasker)
        assert cooldown_cache.exception_masker.visible_prefix == 50
        assert cooldown_cache.exception_masker.visible_suffix == 0
        assert cooldown_cache.exception_masker.mask_char == "*"

    def test_short_exception_string_not_masked(self, cooldown_cache):
        """Test that short exception strings are not masked"""
        short_exception = "Short error"
        model_id = "test-model"
        exception_status = 500
        cooldown_time = 30.0

        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
            model_id=model_id,
            original_exception=Exception(short_exception),
            exception_status=exception_status,
            cooldown_time=cooldown_time,
        )

        # Short exception should not be masked
        assert cooldown_data["exception_received"] == short_exception
        assert cooldown_key == f"deployment:{model_id}:cooldown"

    def test_long_exception_string_masked(self, cooldown_cache):
        """Test that long exception strings are properly masked"""
        # Create a long exception string that simulates prompt leakage
        long_exception = (
            "litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occurred - "
            "No deployments available for selected model, Try again in 5 seconds. "
            "Passed model=anthropic_claude_sonnet_4_v1_0. pre-call-checks=False, "
            "cooldown_list=[('deepseek_r1-eastus', {'exception_received': "
            "'litellm.RateLimitError: RateLimitError: Azure_aiException - "
            '{"error":{"code":"Invalid input","status":422,"message":"invalid input error",'
            '"details":[{"type":"model_attributes_type","loc":["body"],'
            '"msg":"Tell me a story about a dragon and a princess in a magical kingdom '
            "where the dragon is actually protecting the princess from an evil wizard "
            'who wants to steal her magical powers and use them to conquer the world"}]}'
        )

        model_id = "test-model"
        exception_status = 429
        cooldown_time = 60.0

        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
            model_id=model_id,
            original_exception=Exception(long_exception),
            exception_status=exception_status,
            cooldown_time=cooldown_time,
        )

        masked_exception = cooldown_data["exception_received"]

        # Should start with first 50 characters
        assert masked_exception.startswith(long_exception[:50])

        # Should contain masking characters
        assert "*" in masked_exception

        # Should be same length (prefix + asterisks)
        assert len(masked_exception) == len(long_exception)

        # Should not contain the sensitive prompt content
        assert "Tell me a story about a dragon" not in masked_exception
        assert "magical kingdom" not in masked_exception

        # Should preserve the error type information at the beginning (first 50 chars)
        assert masked_exception.startswith("litellm.proxy.proxy_server._handle_llm_api_excepti")

    def test_exception_with_api_keys_masked(self, cooldown_cache):
        """Test that API keys in exceptions are properly masked"""
        exception_with_key = (
            "Authentication failed with api_key=sk-9876567890abcdefghijklmnopqrstuvwxyz "
            "and token=bearer_token_123456789 for model gpt-4"
        )

        model_id = "test-model"
        exception_status = 401
        cooldown_time = 30.0

        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
            model_id=model_id,
            original_exception=Exception(exception_with_key),
            exception_status=exception_status,
            cooldown_time=cooldown_time,
        )

        masked_exception = cooldown_data["exception_received"]

        # Should mask the sensitive content while preserving structure
        assert masked_exception.startswith("Authentication failed with api_key=sk-98765678")
        assert "*" in masked_exception
        assert len(masked_exception) == len(exception_with_key)

    def test_cooldown_data_structure(self, cooldown_cache):
        """Test that the cooldown data structure is correctly formed"""
        exception_msg = "Test exception for structure validation"
        model_id = "test-model"
        exception_status = 500
        cooldown_time = 45.0

        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
            model_id=model_id,
            original_exception=Exception(exception_msg),
            exception_status=exception_status,
            cooldown_time=cooldown_time,
        )

        # Verify cooldown data structure
        assert isinstance(cooldown_data, dict)
        assert "exception_received" in cooldown_data
        assert "status_code" in cooldown_data
        assert "timestamp" in cooldown_data
        assert "cooldown_time" in cooldown_data

        # Verify data types
        assert isinstance(cooldown_data["exception_received"], str)
        assert isinstance(cooldown_data["status_code"], str)
        assert isinstance(cooldown_data["timestamp"], float)
        assert isinstance(cooldown_data["cooldown_time"], float)

        # Verify values
        assert cooldown_data["status_code"] == str(exception_status)
        assert cooldown_data["cooldown_time"] == cooldown_time
        assert cooldown_data["exception_received"] == exception_msg

    def test_exception_object_conversion(self, cooldown_cache):
        """Test that different exception types are properly converted to strings"""
        # Test with different exception types
        exceptions = [
            ValueError("Invalid value provided"),
            KeyError("Missing required key"),
            RuntimeError("Runtime error occurred"),
            Exception("Generic exception"),
        ]

        for exc in exceptions:
            model_id = f"test-model-{exc.__class__.__name__}"

            cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
                model_id=model_id,
                original_exception=exc,
                exception_status=500,
                cooldown_time=30.0,
            )

            # Should successfully convert exception to string
            assert isinstance(cooldown_data["exception_received"], str)
            assert str(exc) == cooldown_data["exception_received"]  # Short exceptions not masked

    def test_masking_preserves_error_debugging_info(self, cooldown_cache):
        """Test that masking preserves essential debugging information"""
        debugging_exception = (
            "RateLimitError: Rate limit exceeded for model gpt-4. "
            "Current usage: 1000 tokens/minute. Limit: 500 tokens/minute. "
            "Request details: model=gpt-4, user_id=user123, "
            "prompt='Write a comprehensive analysis of the economic implications "
            "of artificial intelligence adoption in the healthcare sector, including "
            "potential cost savings, job displacement, and regulatory challenges'"
        )

        model_id = "gpt-4-deployment"
        exception_status = 429
        cooldown_time = 120.0

        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
            model_id=model_id,
            original_exception=Exception(debugging_exception),
            exception_status=exception_status,
            cooldown_time=cooldown_time,
        )

        masked_exception = cooldown_data["exception_received"]

        # Should preserve error type and initial debugging info (first 50 chars)
        assert masked_exception.startswith("RateLimitError: Rate limit exceeded for model gpt-")

        # Should mask the prompt content
        assert "Write a comprehensive analysis" not in masked_exception
        assert "healthcare sector" not in masked_exception

        # Should contain masking indicator
        assert "*" in masked_exception

    def test_error_handling_in_common_add_cooldown_logic(self, cooldown_cache):
        """Test error handling in the _common_add_cooldown_logic method"""
        # This test ensures that edge cases are properly handled
        model_id = "test-model"

        # Test with None exception (edge case) - should be handled gracefully
        cooldown_key, cooldown_data = cooldown_cache._common_add_cooldown_logic(
            model_id=model_id,
            original_exception=None,
            exception_status=500,
            cooldown_time=30.0,
        )

        # Should handle None by converting to string
        assert cooldown_data["exception_received"] == "None"
        assert cooldown_key == f"deployment:{model_id}:cooldown"

    def test_custom_masker_settings(self):
        """Test that custom masker settings work correctly"""
        mock_dual_cache = MagicMock(spec=DualCache)

        # Create cooldown cache and verify default settings
        cache = CooldownCache(cache=mock_dual_cache, default_cooldown_time=60.0)

        # Test that we can access and verify the masker configuration
        assert cache.exception_masker.visible_prefix == 50
        assert cache.exception_masker.visible_suffix == 0
        assert cache.exception_masker.mask_char == "*"

        # Test masking behavior with these settings
        long_string = "A" * 100  # 100 character string
        masked = cache.exception_masker.mask_value(long_string)

        # Should show first 50 characters, then all asterisks
        expected = "A" * 50 + "*" * 50
        assert masked == expected


class TestCooldownCacheTTLCorrection:
    def _make_cooldown_cache(self) -> CooldownCache:
        in_memory = InMemoryCache()
        dual_cache = DualCache(in_memory_cache=in_memory)
        return CooldownCache(cache=dual_cache, default_cooldown_time=60.0)

    def test_expired_entry_evicted_and_not_returned(self):
        """
        An entry with timestamp+cooldown_time in the past must be evicted from
        in-memory cache and excluded from the active cooldown list.
        """
        cc = self._make_cooldown_cache()
        model_id = "expired-deployment"
        key = CooldownCache.get_cooldown_cache_key(model_id)

        expired_value: CooldownCacheValue = {
            "exception_received": "Rate limit",
            "status_code": "429",
            "timestamp": time.time() - 120.0,
            "cooldown_time": 60.0,
        }
        cc.in_memory_cache.set_cache(key, expired_value, ttl=600)

        active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)

        assert active == [], "Expired cooldown entry must not appear in active cooldowns"
        assert cc.in_memory_cache.get_cache(key) is None, "Expired entry must be evicted from in-memory cache"

    def test_active_entry_is_returned(self):
        """
        An entry whose cooldown window has not elapsed must appear in the active list.
        """
        cc = self._make_cooldown_cache()
        model_id = "active-deployment"
        key = CooldownCache.get_cooldown_cache_key(model_id)

        active_value: CooldownCacheValue = {
            "exception_received": "Rate limit",
            "status_code": "429",
            "timestamp": time.time(),
            "cooldown_time": 60.0,
        }
        cc.in_memory_cache.set_cache(key, active_value, ttl=60)

        active = cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)

        assert len(active) == 1
        assert active[0][0] == model_id

    def test_ttl_corrected_when_in_memory_expiry_far_exceeds_remaining(self):
        """
        When DualCache backfills from Redis using the default 600s TTL, the in-memory
        TTL must be corrected to min(remaining, 60) seconds.
        """
        cc = self._make_cooldown_cache()
        model_id = "backfilled-deployment"
        key = CooldownCache.get_cooldown_cache_key(model_id)

        remaining = 30.0
        value: CooldownCacheValue = {
            "exception_received": "Rate limit",
            "status_code": "429",
            "timestamp": time.time() - (60.0 - remaining),
            "cooldown_time": 60.0,
        }
        cc.in_memory_cache.set_cache(key, value, ttl=600)

        before_expiry = cc.in_memory_cache.ttl_dict.get(key)
        assert before_expiry is not None

        cc.get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)

        after_expiry = cc.in_memory_cache.ttl_dict.get(key)
        assert after_expiry is not None
        corrected_remaining = after_expiry - time.time()
        assert corrected_remaining <= 60.0, "Corrected TTL must not exceed 60s"
        assert corrected_remaining > 0, "Corrected TTL must be positive (cooldown still active)"

    @pytest.mark.asyncio
    async def test_async_expired_entry_evicted(self):
        """
        Async path must also evict expired entries.
        """
        cc = self._make_cooldown_cache()
        model_id = "async-expired"
        key = CooldownCache.get_cooldown_cache_key(model_id)

        expired_value: CooldownCacheValue = {
            "exception_received": "Rate limit",
            "status_code": "429",
            "timestamp": time.time() - 120.0,
            "cooldown_time": 60.0,
        }
        cc.in_memory_cache.set_cache(key, expired_value, ttl=600)

        active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)

        assert active == [], "Expired entry must not appear in async active cooldowns"
        assert cc.in_memory_cache.get_cache(key) is None

    @pytest.mark.asyncio
    async def test_async_active_entry_is_returned(self):
        """
        Async counterpart of test_active_entry_is_returned: an entry whose cooldown
        window has not elapsed must appear in the async active list too.
        """
        cc = self._make_cooldown_cache()
        model_id = "async-active-deployment"
        key = CooldownCache.get_cooldown_cache_key(model_id)

        active_value: CooldownCacheValue = {
            "exception_received": "Rate limit",
            "status_code": "429",
            "timestamp": time.time(),
            "cooldown_time": 60.0,
        }
        cc.in_memory_cache.set_cache(key, active_value, ttl=60)

        active = await cc.async_get_active_cooldowns(model_ids=[model_id], parent_otel_span=None)

        assert len(active) == 1
        assert active[0][0] == model_id


class TestCorrectedActiveCooldown:
    def _make_cooldown_cache(self) -> CooldownCache:
        in_memory = InMemoryCache()
        dual_cache = DualCache(in_memory_cache=in_memory)
        return CooldownCache(cache=dual_cache, default_cooldown_time=60.0)

    def _entry(self, timestamp: float, cooldown_time: float) -> CooldownCacheValue:
        return CooldownCacheValue(
            exception_received="Rate limit",
            status_code="429",
            timestamp=timestamp,
            cooldown_time=cooldown_time,
        )

    def test_expired_entry_returns_none_and_evicts(self):
        cc = self._make_cooldown_cache()
        key = "deployment:expired-dep:cooldown"
        entry = self._entry(timestamp=time.time() - 120.0, cooldown_time=60.0)
        cc.in_memory_cache.set_cache(key, dict(entry), ttl=600)

        result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())

        assert result is None
        assert cc.in_memory_cache.get_cache(key) is None

    def test_active_entry_within_window_returns_value(self):
        cc = self._make_cooldown_cache()
        key = "deployment:active-dep:cooldown"
        entry = self._entry(timestamp=time.time(), cooldown_time=60.0)
        cc.in_memory_cache.set_cache(key, dict(entry), ttl=60)

        result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())

        assert result is not None
        assert result["status_code"] == "429"

    def test_inflated_ttl_is_corrected(self):
        cc = self._make_cooldown_cache()
        key = "deployment:backfilled-dep:cooldown"
        remaining = 30.0
        entry = self._entry(timestamp=time.time() - (60.0 - remaining), cooldown_time=60.0)
        cc.in_memory_cache.set_cache(key, dict(entry), ttl=600)

        result = cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())

        assert result is not None
        corrected_expiry = cc.in_memory_cache.ttl_dict.get(key)
        assert corrected_expiry is not None
        assert corrected_expiry - time.time() <= 60.0

    def test_normal_ttl_not_modified(self):
        cc = self._make_cooldown_cache()
        key = "deployment:normal-dep:cooldown"
        entry = self._entry(timestamp=time.time(), cooldown_time=60.0)
        cc.in_memory_cache.set_cache(key, dict(entry), ttl=60)
        original_expiry = cc.in_memory_cache.ttl_dict.get(key)

        cc._corrected_active_cooldown(key, dict(entry), current_time=time.time())

        after_expiry = cc.in_memory_cache.ttl_dict.get(key)
        assert after_expiry == original_expiry


class SharedRedisDouble:
    """
    In-process stand-in for RedisCache, shared by several DualCache instances so that
    tests can model two proxy replicas talking to one Redis.
    """

    def __init__(self) -> None:
        self.store: dict = {}  # mutable-ok: stands in for Redis' own mutable keyspace

    def set_cache(self, key, value, **kwargs):
        self.store[key] = value

    async def async_set_cache(self, key, value, **kwargs):
        self.store[key] = value

    def batch_get_cache(self, key_list, parent_otel_span=None, **kwargs):
        return {key: self.store.get(key) for key in key_list}

    async def async_batch_get_cache(self, key_list, parent_otel_span=None, **kwargs):
        return {key: self.store.get(key) for key in key_list}


class TestCooldownPropagationBetweenReplicas:
    """
    A cooldown written by one replica has to reach its siblings quickly. The router's own
    DualCache re-reads a key that is missing from memory only every 10s, so cooldown reads
    get their own cache with a much shorter Redis read interval.
    """

    def _make_replica(self, redis: SharedRedisDouble, read_interval: float | None = None) -> CooldownCache:
        router_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis)
        if read_interval is None:
            return CooldownCache(cache=router_cache, default_cooldown_time=60.0)
        return CooldownCache(
            cache=router_cache,
            default_cooldown_time=60.0,
            redis_read_interval_seconds=read_interval,
        )

    @pytest.mark.asyncio
    async def test_sibling_replica_sees_cooldown_within_configured_read_interval(self):
        redis = SharedRedisDouble()
        replica_a = self._make_replica(redis, read_interval=0.25)
        replica_b = self._make_replica(redis, read_interval=0.25)
        model_id = "shared-deployment"

        assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == []

        replica_a.add_deployment_to_cooldown(
            model_id=model_id,
            original_exception=Exception("Internal server error"),
            exception_status=500,
            cooldown_time=60.0,
        )

        time.sleep(0.3)

        active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None)
        assert [model_id] == [entry[0] for entry in active], (
            "sibling replica must pick up a cooldown written by another replica within the read interval"
        )

    @pytest.mark.asyncio
    async def test_sibling_replica_sees_cooldown_within_default_read_interval(self):
        redis = SharedRedisDouble()
        replica_a = self._make_replica(redis)
        replica_b = self._make_replica(redis)
        model_id = "default-interval-deployment"

        assert await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None) == []

        replica_a.add_deployment_to_cooldown(
            model_id=model_id,
            original_exception=Exception("Internal server error"),
            exception_status=500,
            cooldown_time=60.0,
        )

        time.sleep(1.2)

        active = await replica_b.async_get_active_cooldowns([model_id], parent_otel_span=None)
        assert [model_id] == [entry[0] for entry in active], (
            "the shipped default read interval must let a sibling replica see a cooldown about a second later"
        )

    def test_sync_read_path_sees_sibling_cooldown_within_read_interval(self):
        redis = SharedRedisDouble()
        replica_a = self._make_replica(redis, read_interval=0.25)
        replica_b = self._make_replica(redis, read_interval=0.25)
        model_id = "sync-shared-deployment"

        assert replica_b.get_active_cooldowns([model_id], parent_otel_span=None) == []

        replica_a.add_deployment_to_cooldown(
            model_id=model_id,
            original_exception=Exception("Internal server error"),
            exception_status=500,
            cooldown_time=60.0,
        )

        time.sleep(0.3)

        active = replica_b.get_active_cooldowns([model_id], parent_otel_span=None)
        assert [model_id] == [entry[0] for entry in active]

    @pytest.mark.asyncio
    async def test_redis_attached_after_construction_is_still_used(self):
        redis = SharedRedisDouble()
        router_cache = DualCache(in_memory_cache=InMemoryCache())
        writer = CooldownCache(cache=router_cache, default_cooldown_time=60.0, redis_read_interval_seconds=0.25)
        router_cache.attach_redis_cache(redis)
        reader = self._make_replica(redis, read_interval=0.25)
        model_id = "late-redis-deployment"

        writer.add_deployment_to_cooldown(
            model_id=model_id,
            original_exception=Exception("Internal server error"),
            exception_status=500,
            cooldown_time=60.0,
        )

        active = await reader.async_get_active_cooldowns([model_id], parent_otel_span=None)
        assert [model_id] == [entry[0] for entry in active], (
            "a router that wires Redis after building its cooldown cache must still publish cooldowns to it"
        )


class TestCooldownSurvivesUnrelatedCacheTraffic:
    @pytest.mark.asyncio
    async def test_unrelated_router_cache_writes_do_not_evict_active_cooldown(self):
        router_cache = DualCache(in_memory_cache=InMemoryCache())
        cc = CooldownCache(cache=router_cache, default_cooldown_time=60.0)
        model_id = "busy-router-deployment"

        cc.add_deployment_to_cooldown(
            model_id=model_id,
            original_exception=Exception("Internal server error"),
            exception_status=500,
            cooldown_time=30.0,
        )

        for i in range(400):
            router_cache.set_cache(key=f"unrelated-router-key-{i}", value={"n": i})

        active = await cc.async_get_active_cooldowns([model_id], parent_otel_span=None)
        assert [model_id] == [entry[0] for entry in active], (
            "unrelated router cache traffic must not evict a cooldown that is still running"
        )


class TestCooldownStoreCallsDeclareTheirKeyFamily:
    """Every cooldown store call runs inside ``service_target("router_cooldowns")`` so the
    Redis service spans read ``redis.set router_cooldowns`` / ``redis.mget router_cooldowns``,
    the sync paths included (the async MGET already did)."""

    def _cooldown_cache_with_recording_store(self, seen: list[tuple[str, str | None]]) -> CooldownCache:
        cc = CooldownCache(cache=DualCache(in_memory_cache=InMemoryCache()), default_cooldown_time=60.0)
        store = MagicMock()

        def _set_cache(**_kwargs):
            seen.append(("set", current_service_target()))

        def _batch_get_cache(**_kwargs):
            seen.append(("mget", current_service_target()))
            return []

        store.set_cache.side_effect = _set_cache
        store.batch_get_cache.side_effect = _batch_get_cache
        cc._cooldown_store = store
        return cc

    def test_sync_cooldown_write_runs_under_router_cooldowns(self):
        seen: list[tuple[str, str | None]] = []
        cc = self._cooldown_cache_with_recording_store(seen)

        cc.add_deployment_to_cooldown(
            model_id="dep-1",
            original_exception=Exception("Internal server error"),
            exception_status=500,
            cooldown_time=30.0,
        )

        assert seen == [("set", "router_cooldowns")]
        assert current_service_target() is None

    def test_sync_cooldown_reads_run_under_router_cooldowns(self):
        seen: list[tuple[str, str | None]] = []
        cc = self._cooldown_cache_with_recording_store(seen)

        assert cc.get_active_cooldowns(["dep-1"], parent_otel_span=None) == []
        assert cc.get_min_cooldown(["dep-1"], parent_otel_span=None) == 60.0

        assert seen == [("mget", "router_cooldowns"), ("mget", "router_cooldowns")]