mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
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:
parent
44ef388694
commit
2509ec892c
3 changed files with 90 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 == []
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue