From 2509ec892c57f6068c5b5c39fca95152ac13ac9c Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Sat, 15 Aug 2026 17:34:11 -0700 Subject: [PATCH] test(guardrails): cover the remaining native-hook opt-out dispatch sites Adds regression tests for the parallel post-call path, the MCP post-call hook, and the policy engine step, so every read of the opt-out flag fails when removed. --- .../policy_engine/test_pipeline_executor.py | 39 +++++++++++++++++++ .../test_proxy_logging_hook_detection.py | 21 ++++++++++ tests/test_litellm/proxy/test_proxy_utils.py | 30 ++++++++++++++ 3 files changed, 90 insertions(+) diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index adf4e8d47c6..840d93eb12c 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -797,3 +797,42 @@ async def test_step_results_include_duration(): assert result.step_results[0].duration_seconds >= 0 finally: litellm.callbacks = original_callbacks + + +class _PolicyOptOutGuardrail(CustomGuardrail): + """Implements apply_guardrail for the direct endpoint but keeps its native hooks. + + apply_guardrail is defined here rather than inherited because the dispatch check + reads the leaf class __dict__. + """ + + use_native_lifecycle_hooks = True + + def __init__(self): + super().__init__(guardrail_name="policy-opt-out", default_on=True) + self.native_pre_call_ran = False + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.native_pre_call_ran = True + + +@pytest.mark.asyncio +async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch): + guardrail = _PolicyOptOutGuardrail() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + data = {"messages": [{"role": "user", "content": "hi"}]} + outcome, _, _, _ = await PipelineExecutor._run_step( + step=PipelineStep(guardrail="policy-opt-out", on_fail="block", on_pass="allow"), + mode="pre_call", + data=data, + user_api_key_dict=MagicMock(), + call_type="completion", + ) + + assert outcome == "pass" + assert guardrail.native_pre_call_ran is True + assert "guardrail_to_apply" not in data diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index fec524b96f3..133156f9321 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -731,3 +731,24 @@ async def test_post_call_stream_keeps_own_iterator_when_opted_out(monkeypatch): assert own_iterator_ran == ["claude-sonnet-5"] assert delivered == chunks + + +@pytest.mark.asyncio +async def test_parallel_post_call_guardrails_keep_native_hook_when_opted_out(monkeypatch): + """The run_in_parallel post-call path has its own dispatch check, so the opt-out has + to be honored there too.""" + from litellm.types.utils import Choices, Message, ModelResponse + + opted_out = _KeepsNativeHooks(event_hook=GuardrailEventHooks.post_call, default_on=True, run_in_parallel=True) + routed = _AppliesGuardrail(event_hook=GuardrailEventHooks.post_call, default_on=True, run_in_parallel=True) + monkeypatch.setattr(litellm, "callbacks", [opted_out, routed]) + response = ModelResponse(choices=[Choices(message=Message(role="assistant", content="hello"))]) + + await ProxyLogging(user_api_key_cache=DualCache()).post_call_success_hook( + data={"messages": [{"role": "user", "content": "hi"}]}, + response=response, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234"), + ) + + assert opted_out.native_hooks_ran == ["post_call"] + assert routed.native_hooks_ran == [] diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index a4f93e90673..1504c3c3103 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1191,3 +1191,33 @@ async def test_update_data_key_branch_stamps_settings_updated_at(): sent = client.db.litellm_verificationtoken.update.call_args.kwargs["data"] assert sent["models"] == ["gpt-4"] assert before <= sent["settings_updated_at"] <= after + + +@pytest.mark.asyncio +async def test_post_mcp_call_hook_skips_opted_out_guardrail(restore_callbacks): + """A guardrail that keeps its native lifecycle hooks must not have MCP tool results + scanned through the unified path, even though it implements apply_guardrail.""" + from mcp.types import CallToolResult, TextContent + + class _OptedOutMCPGuardrail(_RecordingMCPGuardrail): + # apply_guardrail is redefined rather than inherited because the dispatch check + # reads the leaf class __dict__, so an inherited override would skip for the + # wrong reason and leave the flag untested + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + return await super().apply_guardrail(inputs, request_data, input_type, **kwargs) + + guardrail = _OptedOutMCPGuardrail(event_hook=GuardrailEventHooks.post_mcp_call) + litellm.callbacks = [guardrail] + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False) + + returned = await proxy_logging_obj.post_mcp_call_hook( + response=result, + request_data={"mcp_tool_name": "echo"}, + user_api_key_dict=None, + ) + + assert guardrail.call_count == 0 + assert [item.text for item in returned.content] == ["jane@example.com"]