scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

fix(guardrails): remove the module-global translation mapping that leaked between tests

Browse source 
The unified guardrail cached the endpoint translation mappings in its own
module global on top of the loader's cache in litellm/llms. Tests wrote to
that second copy directly, so a teardown that restored a stale snapshot left
a test double installed for every later test on the same xdist worker, and
proxy-endpoints went red on whichever guardrail streaming test happened to
land after it.

Read through load_guardrail_translation_mappings() at each call site and give
the tests one seam to patch, so pytest owns every restore.
This commit is contained in:
mateo-berri 2026-09-03 03:28:52 -07:00
parent 658f50663d
commit a737f8625d
5 changed files with 64 additions and 96 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -121,9 +121,6 @@ def _a2a_jsonrpc_error_chunk(exc: HTTPException, request_id: str | None) -> Mapp
}
endpoint_guardrail_translation_mappings = None
def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> None:
"""Populate data['litellm_metadata'] from user_api_key_dict if absent."""
if "litellm_metadata" not in data:
@ -164,7 +161,6 @@ class UnifiedLLMGuardrails(CustomLogger):
Use this if you want to MODIFY the input
"""
global endpoint_guardrail_translation_mappings
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
@ -186,18 +182,15 @@ class UnifiedLLMGuardrails(CustomLogger):
)
return data
if endpoint_guardrail_translation_mappings is None:
endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
mappings: Final = load_guardrail_translation_mappings()
try:
if CallTypes(call_type) not in endpoint_guardrail_translation_mappings:
if CallTypes(call_type) not in mappings:
return data
except ValueError:
return data # handle unmapped call types
endpoint_translation: Final = _as_endpoint_translation(
endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
)
endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
_ensure_litellm_metadata(data, user_api_key_dict)
@ -222,8 +215,6 @@ class UnifiedLLMGuardrails(CustomLogger):
This can NOT modify the input, only used to reject or accept a call before going to LLM API
"""
global endpoint_guardrail_translation_mappings
verbose_proxy_logger.debug("Running UnifiedLLMGuardrails moderation hook")
guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None)
@ -241,14 +232,11 @@ class UnifiedLLMGuardrails(CustomLogger):
)
return data
if endpoint_guardrail_translation_mappings is None:
endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings:
mappings: Final = load_guardrail_translation_mappings()
if call_type is not None and CallTypes(call_type) not in mappings:
return data
endpoint_translation: Final = _as_endpoint_translation(
endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
)
endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
_ensure_litellm_metadata(data, user_api_key_dict)
@ -271,7 +259,6 @@ class UnifiedLLMGuardrails(CustomLogger):
Uses Enkrypt AI guardrails to check the response for policy violations, PII, and injection attacks
"""
global endpoint_guardrail_translation_mappings
# Local import avoids a module-level cyclic import with
# litellm.integrations.custom_guardrail.
from litellm.integrations.custom_guardrail import ModifyResponseException
@ -319,10 +306,9 @@ class UnifiedLLMGuardrails(CustomLogger):
)
return response
if endpoint_guardrail_translation_mappings is None:
endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
mappings: Final = load_guardrail_translation_mappings()
if CallTypes(call_type) not in endpoint_guardrail_translation_mappings:
if CallTypes(call_type) not in mappings:
verbose_proxy_logger.warning(
"Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; "
"skipping post-call scanning.",
@ -332,9 +318,7 @@ class UnifiedLLMGuardrails(CustomLogger):
)
return response
endpoint_translation: Final = _as_endpoint_translation(
endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
)
endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
try:
response = await endpoint_translation.process_output_response(
@ -906,8 +890,6 @@ class UnifiedLLMGuardrails(CustomLogger):
sampling_rate=1 means every chunk, sampling_rate=5 means every 5th chunk, etc.
"""
global endpoint_guardrail_translation_mappings
# Local import avoids a module-level cyclic import with
# litellm.integrations.custom_guardrail.
from litellm.integrations.custom_guardrail import ModifyResponseException
@ -978,9 +960,7 @@ class UnifiedLLMGuardrails(CustomLogger):
yield item
return
# Initialize translation mappings if needed
if endpoint_guardrail_translation_mappings is None:
endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
mappings: Final = load_guardrail_translation_mappings()
# Streaming text transformation (incremental_diff) diverges enough from the
# block_only path that it runs as its own iterator. It requires a route we
@ -989,7 +969,7 @@ class UnifiedLLMGuardrails(CustomLogger):
if streaming_transform_mode == "incremental_diff":
transform_call_type: Final = self._resolve_transform_call_type(
user_api_key_dict=user_api_key_dict,
mappings=endpoint_guardrail_translation_mappings,
mappings=mappings,
)
if transform_call_type is not None:
async for transformed_item in self._run_incremental_transform_stream(
@ -1000,7 +980,7 @@ class UnifiedLLMGuardrails(CustomLogger):
call_type=transform_call_type,
sampling_rate=sampling_rate,
end_of_stream_only=end_of_stream_only,
mappings=endpoint_guardrail_translation_mappings,
mappings=mappings,
):
yield transformed_item
return
@ -1037,7 +1017,7 @@ class UnifiedLLMGuardrails(CustomLogger):
call_type = _infer_call_type(call_type=None, completion_response=item)
# If call type not supported, just pass through all chunks
if call_type is None or CallTypes(call_type) not in endpoint_guardrail_translation_mappings:
if call_type is None or CallTypes(call_type) not in mappings:
yield item
async for remaining_item in response:
yield remaining_item
@ -1049,7 +1029,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# moderation runs below.
if end_of_stream_only:
if not buffer_until_moderated:
endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
endpoint_translation = mappings[CallTypes(call_type)]()
stream_has_ended = hasattr(
endpoint_translation, "_check_streaming_has_ended"
) and endpoint_translation._check_streaming_has_ended(responses_so_far)
@ -1063,7 +1043,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# Process chunk based on sampling rate
if chunk_counter % sampling_rate == 0:
endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
endpoint_translation = mappings[CallTypes(call_type)]()
scan_key = endpoint_translation.get_streaming_scan_key(responses_so_far)
if _is_redundant_scan(scan_key, last_scan_key):
verbose_proxy_logger.debug(
@ -1143,14 +1123,14 @@ class UnifiedLLMGuardrails(CustomLogger):
yield item
# Stream has ended - do final processing with all collected chunks
if call_type is not None and CallTypes(call_type) in endpoint_guardrail_translation_mappings:
if call_type is not None and CallTypes(call_type) in mappings:
verbose_proxy_logger.debug(
"Processing final streaming response with all %s chunks for guardrail %s",
len(responses_so_far),
guardrail_to_apply.guardrail_name,
)
endpoint_translation = endpoint_guardrail_translation_mappings[CallTypes(call_type)]()
endpoint_translation = mappings[CallTypes(call_type)]()
# When buffering, snapshot the original chunks before moderation.
# A shallow copy suffices: end-of-stream

26
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

@ -5527,10 +5527,6 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca
error frame instead. The finish chunk is withheld while the end-of-stream
scan runs, so on a block it is dropped rather than relayed before the
frame."""
from litellm.llms import load_guardrail_translation_mappings
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import (
unified_guardrail as unified_module,
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -5569,20 +5565,16 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca
yield _chunk("the forbidden ")
yield _chunk("topic answer", finish_reason="stop")
unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
try:
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
mock_api.side_effect = guardrail._get_http_exception_for_blocked_guardrail(blocked_response)
with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api:
mock_api.side_effect = guardrail._get_http_exception_for_blocked_guardrail(blocked_response)
out = []
async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions"),
response=_mock_stream(),
request_data={"guardrail_to_apply": guardrail, "model": "gpt-4"},
):
out.append(item)
finally:
unified_module.endpoint_guardrail_translation_mappings = None
out = []
async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions"),
response=_mock_stream(),
request_data={"guardrail_to_apply": guardrail, "model": "gpt-4"},
):
out.append(item)
assert len(out) == 2
assert isinstance(out[0], ModelResponseStream)

70
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -75,19 +75,29 @@ class _NoopTranslation(BaseTranslation):
return response
def _patch_translation_mappings(monkeypatch, mappings):
"""Point the unified guardrail at ``mappings`` for one test, restored by pytest.
Every override goes through this one seam: competing writers to the same state
are what leaked a stale handler map into unrelated test files (LIT-6834).
"""
monkeypatch.setattr(unified_module, "load_guardrail_translation_mappings", lambda: mappings)
@pytest.fixture(autouse=True)
def _inject_mcp_handler_mapping():
def _inject_mcp_handler_mapping(monkeypatch):
"""Inject MCP handler mapping so the unified guardrail can run inside tests."""
unified_module.endpoint_guardrail_translation_mappings = {
CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler,
CallTypes.anthropic_messages: _NoopTranslation,
CallTypes.ocr: OCRHandler,
CallTypes.aocr: OCRHandler,
CallTypes.responses: OpenAIResponsesHandler,
CallTypes.aresponses: OpenAIResponsesHandler,
}
yield
unified_module.endpoint_guardrail_translation_mappings = None
_patch_translation_mappings(
monkeypatch,
{
CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler,
CallTypes.anthropic_messages: _NoopTranslation,
CallTypes.ocr: OCRHandler,
CallTypes.aocr: OCRHandler,
CallTypes.responses: OpenAIResponsesHandler,
CallTypes.aresponses: OpenAIResponsesHandler,
},
)
class TestUnifiedLLMGuardrails:
@ -396,7 +406,7 @@ class TestUnifiedLLMGuardrails:
class TestAsyncPostCallStreamingIteratorHook:
@pytest.mark.asyncio
async def test_streaming_content_not_lost_on_sampled_chunks(self):
async def test_streaming_content_not_lost_on_sampled_chunks(self, monkeypatch):
"""
Verify that every chunk's content is preserved in the output stream.
@ -442,10 +452,7 @@ class TestUnifiedLLMGuardrails:
return responses_so_far
# Override the mapping to use our content-clearing translation
unified_module.endpoint_guardrail_translation_mappings = {
CallTypes.acompletion: _ContentClearingTranslation,
}
_patch_translation_mappings(monkeypatch, {CallTypes.acompletion: _ContentClearingTranslation})
handler = UnifiedLLMGuardrails()
guardrail = RecordingGuardrail()
@ -885,12 +892,8 @@ class TestStreamingTransform:
completions streaming surface."""
@pytest.fixture(autouse=True)
def _use_openai_handler_mapping(self):
unified_module.endpoint_guardrail_translation_mappings = {
CallTypes.acompletion: OpenAIChatCompletionsHandler,
}
yield
unified_module.endpoint_guardrail_translation_mappings = None
def _use_openai_handler_mapping(self, monkeypatch):
_patch_translation_mappings(monkeypatch, {CallTypes.acompletion: OpenAIChatCompletionsHandler})
@pytest.mark.asyncio
async def test_block_only_drops_text_rewrites(self):
@ -1719,6 +1722,10 @@ class TestAppliedGuardrailsReflectsExecution:
decision and marks itself only when it actually ran (LIT-4650). Ordinary
guardrails are still auto-marked by the hook after dispatch."""
@pytest.fixture(autouse=True)
def _use_texts_only_mapping(self, monkeypatch):
_patch_translation_mappings(monkeypatch, {CallTypes.pass_through: _TextsOnlyTranslation})
@staticmethod
def _data(guardrail):
return {
@ -1728,7 +1735,6 @@ class TestAppliedGuardrailsReflectsExecution:
}
async def _run(self, guardrail):
unified_module.endpoint_guardrail_translation_mappings = {CallTypes.pass_through: _TextsOnlyTranslation}
data = self._data(guardrail)
await UnifiedLLMGuardrails().async_pre_call_hook(
user_api_key_dict=None,
@ -1830,10 +1836,8 @@ class TestStreamingHttpErrorFrames:
silently truncates the SSE stream (PR #38722 defect 1)."""
@pytest.fixture(autouse=True)
def _use_real_mappings(self):
unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
yield
unified_module.endpoint_guardrail_translation_mappings = None
def _use_real_mappings(self, monkeypatch):
_patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings())
@pytest.mark.asyncio
async def test_chat_eos_block_emits_data_error_frame(self):
@ -1938,10 +1942,8 @@ class TestStreamingGuardrailInformationBucket:
guardrail_information write was diverted and /spend/logs showed null."""
@pytest.fixture(autouse=True)
def _use_real_mappings(self):
unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings()
yield
unified_module.endpoint_guardrail_translation_mappings = None
def _use_real_mappings(self, monkeypatch):
_patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings())
@pytest.mark.asyncio
async def test_chat_eos_scan_writes_guardrail_information_to_metadata(self):
@ -2038,11 +2040,7 @@ class TestStreamingScanDedup:
@pytest.fixture(autouse=True)
def _use_real_mappings(self, monkeypatch):
monkeypatch.setattr(
unified_module,
"endpoint_guardrail_translation_mappings",
load_guardrail_translation_mappings(),
)
_patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings())
@pytest.mark.asyncio
async def test_chat_terminal_chunk_on_sampled_index_is_scanned_once(self):

4
tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py
View file

@ -292,8 +292,8 @@ class TestUnifiedGuardrailCallTypeResolution:
with patch.object(
unified_guardrail_module,
"endpoint_guardrail_translation_mappings",
{CallTypes.pass_through: mock_handler_class},
"load_guardrail_translation_mappings",
lambda: {CallTypes.pass_through: mock_handler_class},
):
result = await unified.async_post_call_success_hook(
data=data,

6
tests/test_litellm/proxy/test_blocked_response_usage.py
View file

@ -68,12 +68,10 @@ async def test_success_hook_attaches_original_response_on_block():
user_api_key_dict = UserAPIKeyAuth(api_key="test", request_route="/chat/completions")
data = {"guardrail_to_apply": guardrail, "model": "gpt-4o"}
# Inject our translation for the inferred call type (the module global is
# cached across tests, so patch it directly rather than the loader).
with patch.object(
ug,
"endpoint_guardrail_translation_mappings",
{
"load_guardrail_translation_mappings",
lambda: {
CallTypes.acompletion: lambda: translation,
CallTypes.completion: lambda: translation,
},