scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

fix(cooldown_cache.py): mask error string to avoid leaking sensitive prompt data

Browse source 
Fixes https://github.com/BerriAI/litellm/issues/13329
This commit is contained in:
Krrish Dholakia 2025-08-11 18:37:57 -07:00
parent dab3acdd26
commit bc9d0484e4
4 changed files with 283 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/litellm_core_utils/sensitive_data_masker.py
View file

@ -33,7 +33,12 @@ class SensitiveDataMasker:
value_str = str(value)
masked_length = len(value_str) - (self.visible_prefix + self.visible_suffix)
return f"{value_str[:self.visible_prefix]}{self.mask_char * masked_length}{value_str[-self.visible_suffix:]}"
# Handle the case where visible_suffix is 0 to avoid showing the entire string
if self.visible_suffix == 0:
return f"{value_str[:self.visible_prefix]}{self.mask_char * masked_length}"
else:
return f"{value_str[:self.visible_prefix]}{self.mask_char * masked_length}{value_str[-self.visible_suffix:]}"
def is_sensitive_key(self, key: str) -> bool:
key_lower = str(key).lower()

11
litellm/router_utils/cooldown_cache.py
View file

@ -8,6 +8,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypedDict, Union
from litellm import verbose_logger
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -29,6 +30,12 @@ class CooldownCache:
self.cache = cache
self.default_cooldown_time = default_cooldown_time
self.in_memory_cache = InMemoryCache()
# Initialize the masker with custom settings for exception strings
self.exception_masker = SensitiveDataMasker(
visible_prefix=50, # Show first 50 characters
visible_suffix=0, # Show last 0 characters
mask_char="*", # Use * for masking
)
def _common_add_cooldown_logic(
self, model_id: str, original_exception, exception_status, cooldown_time: float
@ -39,7 +46,9 @@ class CooldownCache:
# Store the cooldown information for the deployment separately
cooldown_data = CooldownCacheValue(
exception_received=str(original_exception),
exception_received=self.exception_masker._mask_value(
str(original_exception)
),
status_code=str(exception_status),
timestamp=current_time,
cooldown_time=cooldown_time,

20
litellm/router_utils/cooldown_handlers.py
View file

@ -118,16 +118,16 @@ def _should_run_cooldown_logic(
"Should Not Run Cooldown Logic: deployment id is none or model group can't be found."
)
return False
#########################################################
# If time_to_cooldown is 0 or 0.0000000, don't run cooldown logic
#########################################################
if time_to_cooldown is not None and math.isclose(
a=time_to_cooldown,
b=0.0,
abs_tol=1e-9
a=time_to_cooldown, b=0.0, abs_tol=1e-9
):
verbose_router_logger.debug("Should Not Run Cooldown Logic: time_to_cooldown is effectively 0")
verbose_router_logger.debug(
"Should Not Run Cooldown Logic: time_to_cooldown is effectively 0"
)
return False
if litellm_router_instance.disable_cooldowns:
@ -275,8 +275,8 @@ def _set_cooldown_deployments(
if (
_should_run_cooldown_logic(
litellm_router_instance=litellm_router_instance,
deployment=deployment,
exception_status=exception_status,
deployment=deployment,
exception_status=exception_status,
original_exception=original_exception,
time_to_cooldown=time_to_cooldown,
)
@ -290,9 +290,9 @@ def _set_cooldown_deployments(
verbose_router_logger.debug(f"Attempting to add {deployment} to cooldown list")
if _should_cooldown_deployment(
litellm_router_instance=litellm_router_instance,
deployment=deployment,
exception_status=exception_status,
litellm_router_instance=litellm_router_instance,
deployment=deployment,
exception_status=exception_status,
original_exception=original_exception,
):
litellm_router_instance.cooldown_cache.add_deployment_to_cooldown(

257
tests/test_litellm/router_utils/test_cooldown_cache.py Normal file
View file

@ -0,0 +1,257 @@
"""
Unit tests for CooldownCache exception masking functionality
"""
import os
import sys
from unittest.mock import MagicMock
import pytest
# Add the parent directory to the system path
sys.path.insert(0, os.path.abspath("../../.."))
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-1234567890abcdefghijklmnopqrstuvwxyz "
"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-12345678"
)
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