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
|
assert result.step_results[0].duration_seconds >= 0
|
||||||
finally:
|
finally:
|
||||||
litellm.callbacks = original_callbacks
|
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 own_iterator_ran == ["claude-sonnet-5"]
|
||||||
assert delivered == chunks
|
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"]
|
sent = client.db.litellm_verificationtoken.update.call_args.kwargs["data"]
|
||||||
assert sent["models"] == ["gpt-4"]
|
assert sent["models"] == ["gpt-4"]
|
||||||
assert before <= sent["settings_updated_at"] <= after
|
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