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.
This commit is contained in:
Yucheng Zhu 2026-08-15 17:34:11 -07:00
parent 44ef388694
commit 2509ec892c
3 changed files with 90 additions and 0 deletions

View file

@ -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

View file

@ -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 == []

View file

@ -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"]