From a737f8625d0d2acdffca3efb93a042a23f718a09 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 03:28:52 -0700 Subject: [PATCH] fix(guardrails): remove the module-global translation mapping that leaked between tests 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. --- .../unified_guardrail/unified_guardrail.py | 54 +++++--------- .../test_bedrock_guardrails.py | 26 +++---- .../test_unified_guardrail.py | 70 +++++++++---------- .../test_passthrough_post_call_guardrails.py | 4 +- .../proxy/test_blocked_response_usage.py | 6 +- 5 files changed, 64 insertions(+), 96 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index c6b8df1b493..9029e926b35 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 953e3de1519..479d1f2d4b2 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -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) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index a28a2a71613..5846d655069 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -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): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index 9d1975513a1..c7696079adc 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -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, diff --git a/tests/test_litellm/proxy/test_blocked_response_usage.py b/tests/test_litellm/proxy/test_blocked_response_usage.py index 37aea8fe3aa..4f20f35e94b 100644 --- a/tests/test_litellm/proxy/test_blocked_response_usage.py +++ b/tests/test_litellm/proxy/test_blocked_response_usage.py @@ -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, },