litellm/tests/test_litellm/integrations/test_custom_guardrail.py
Yassin Kortam 3ab500c79e
fix(guardrails): run pre_call hook once for model-level guardrails (#30543)
* fix(guardrails): run pre_call hook once for model-level guardrails

A CustomGuardrail attached to a deployment via litellm_params.guardrails
gets its async_pre_call_hook invoked twice per request: once by the proxy
pre-call loop and again by async_pre_call_deployment_hook after the router
spreads the model-level guardrails into the top-level request kwargs.

Record in request metadata that the proxy pre-call loop already ran a given
guardrail, and have the deployment hook skip it when the marker is present.
Direct-SDK usage never runs the proxy loop, so the deployment hook stays the
sole invocation there and still fires exactly once.

The marker key is stripped from untrusted caller metadata so a request body
cannot suppress a model-only guardrail by pre-seeding it.

* fix(guardrails): mark pre_call dedup on the post-hook request data

Record the exactly-once marker after async_pre_call_hook runs, on the data
object that flows downstream, rather than before it. A guardrail whose hook
returns a brand-new request dict (instead of mutating or spreading the one it
received) would otherwise discard the marker, letting the deployment hook
re-run the guardrail a second time.

(cherry picked from commit 4faeabc254)
2026-06-20 11:49:30 -07:00

1401 lines
54 KiB
Python

from unittest.mock import AsyncMock
import pytest
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
from litellm.types.utils import GuardrailTracingDetail
class TestCustomGuardrailDeploymentHook:
@pytest.mark.asyncio
async def test_async_pre_call_deployment_hook_no_guardrails(self):
"""Test that method returns kwargs unchanged when no guardrails are present"""
custom_guardrail = CustomGuardrail()
# Test with guardrails as None
kwargs = {
"messages": [{"role": "user", "content": "test message"}],
"model": "gpt-3.5-turbo",
"guardrails": None,
}
result = await custom_guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert result == kwargs
# Test with guardrails as non-list
kwargs["guardrails"] = "not_a_list"
result = await custom_guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert result == kwargs
@pytest.mark.asyncio
async def test_async_pre_call_deployment_hook_with_guardrails_and_message_update(
self,
):
"""Test that method processes guardrails and updates messages when result contains messages"""
custom_guardrail = CustomGuardrail()
# Mock the async_pre_call_hook method
mock_result = {"messages": [{"role": "user", "content": "filtered message"}]}
custom_guardrail.async_pre_call_hook = AsyncMock(return_value=mock_result)
original_messages = [{"role": "user", "content": "original message"}]
kwargs = {
"messages": original_messages,
"model": "gpt-3.5-turbo",
"guardrails": ["some_guardrail"],
"user_api_key_user_id": "test_user",
"user_api_key_team_id": "test_team",
"user_api_key_end_user_id": "test_end_user",
"user_api_key_hash": "test_hash",
"user_api_key_request_route": "test_route",
}
result = await custom_guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
# Verify async_pre_call_hook was called with correct parameters
custom_guardrail.async_pre_call_hook.assert_called_once()
call_args = custom_guardrail.async_pre_call_hook.call_args
# Check that UserAPIKeyAuth was created properly
user_api_key_dict = call_args[1]["user_api_key_dict"]
assert isinstance(user_api_key_dict, UserAPIKeyAuth)
assert user_api_key_dict.user_id == "test_user"
assert user_api_key_dict.team_id == "test_team"
assert user_api_key_dict.end_user_id == "test_end_user"
assert user_api_key_dict.api_key == "test_hash"
assert user_api_key_dict.request_route == "test_route"
# Check other parameters
assert call_args[1]["data"] == kwargs
assert call_args[1]["call_type"] == "completion"
# Verify messages were updated in result
assert result["messages"] == mock_result["messages"]
assert result["messages"] != original_messages
@pytest.mark.asyncio
async def test_deployment_hook_skips_when_pre_call_already_ran(self):
"""The deployment hook must not re-run async_pre_call_hook once the proxy
pre-call loop has already run it for this request."""
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="g1", default_on=True)
self.pre_call_count = 0
async def async_pre_call_hook(
self, user_api_key_dict, cache, data, call_type
):
self.pre_call_count += 1
return data
guardrail = CountingGuardrail()
kwargs = {
"messages": [{"role": "user", "content": "hi"}],
"model": "gpt-3.5-turbo",
"guardrails": ["g1"],
"metadata": {},
}
guardrail.mark_pre_call_hook_ran(kwargs)
await guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert guardrail.pre_call_count == 0
@pytest.mark.asyncio
async def test_deployment_hook_runs_when_not_marked(self):
"""Without the proxy marker (direct-SDK usage) the deployment hook is the
only execution path and must still run the guardrail exactly once."""
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="g1", default_on=True)
self.pre_call_count = 0
async def async_pre_call_hook(
self, user_api_key_dict, cache, data, call_type
):
self.pre_call_count += 1
return data
guardrail = CountingGuardrail()
kwargs = {
"messages": [{"role": "user", "content": "hi"}],
"model": "gpt-3.5-turbo",
"guardrails": ["g1"],
"metadata": {},
}
await guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert guardrail.pre_call_count == 1
def test_mark_pre_call_hook_ran_uses_litellm_metadata(self):
"""The marker is recorded in litellm_metadata when that is the metadata
bucket in use, and is then visible to the skip check."""
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
guardrail = CustomGuardrail(guardrail_name="g1")
kwargs = {"litellm_metadata": {}}
guardrail.mark_pre_call_hook_ran(kwargs)
assert kwargs["litellm_metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY]
assert guardrail._pre_call_hook_already_ran(kwargs) is True
@pytest.mark.asyncio
async def test_deployment_hook_ignores_forged_caller_marker(self):
"""A direct-SDK caller controls request metadata but cannot know the
per-process token, so a hand-crafted marker must not suppress a
requested guardrail in async_pre_call_deployment_hook."""
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="g1", default_on=True)
self.pre_call_count = 0
async def async_pre_call_hook(
self, user_api_key_dict, cache, data, call_type
):
self.pre_call_count += 1
return data
guardrail = CountingGuardrail()
kwargs = {
"messages": [{"role": "user", "content": "hi"}],
"model": "gpt-3.5-turbo",
"guardrails": ["g1"],
"metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]},
}
await guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert guardrail.pre_call_count == 1
class TestCustomGuardrailShouldRunGuardrail:
def test_should_run_guardrail_with_litellm_metadata(self):
"""Test that should_run_guardrail works with litellm_metadata pattern"""
from litellm.types.guardrails import GuardrailEventHooks
custom_guardrail = CustomGuardrail(
guardrail_name="test_guardrail",
default_on=False,
event_hook=GuardrailEventHooks.pre_call,
)
# Test with guardrails in litellm_metadata
data = {
"model": "gpt-3.5-turbo",
"litellm_metadata": {"guardrails": ["test_guardrail"]},
}
result = custom_guardrail.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.pre_call
)
assert result is True
def test_should_run_guardrail_with_metadata(self):
"""Test that should_run_guardrail works with metadata pattern"""
from litellm.types.guardrails import GuardrailEventHooks
custom_guardrail = CustomGuardrail(
guardrail_name="test_guardrail",
default_on=False,
event_hook=GuardrailEventHooks.pre_call,
)
# Test with guardrails in metadata
data = {
"model": "gpt-3.5-turbo",
"metadata": {"guardrails": ["test_guardrail"]},
}
result = custom_guardrail.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.pre_call
)
assert result is True
def test_should_run_guardrail_with_root_level_guardrails(self):
"""Test that should_run_guardrail works with root level guardrails"""
from litellm.types.guardrails import GuardrailEventHooks
custom_guardrail = CustomGuardrail(
guardrail_name="test_guardrail",
default_on=False,
event_hook=GuardrailEventHooks.pre_call,
)
# Test with guardrails at root level
data = {"model": "gpt-3.5-turbo", "guardrails": ["test_guardrail"]}
result = custom_guardrail.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.pre_call
)
assert result is True
def test_should_run_guardrail_no_matching_guardrail(self):
"""Test that should_run_guardrail returns False when guardrail name doesn't match"""
from litellm.types.guardrails import GuardrailEventHooks
custom_guardrail = CustomGuardrail(
guardrail_name="test_guardrail",
default_on=False,
event_hook=GuardrailEventHooks.pre_call,
)
# Test with different guardrail name
data = {
"model": "gpt-3.5-turbo",
"litellm_metadata": {"guardrails": ["different_guardrail"]},
}
result = custom_guardrail.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.pre_call
)
assert result is False
def test_should_run_guardrail_with_disable_global_guardrail(self):
"""Test that disable_global_guardrails only works from admin metadata"""
from litellm.types.guardrails import GuardrailEventHooks
custom_guardrail = CustomGuardrail(
guardrail_name="global_guardrail",
default_on=True,
event_hook=GuardrailEventHooks.pre_call,
)
# Test 1: Global guardrail runs by default
data = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
}
result = custom_guardrail.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.pre_call
)
assert result is True, "Global guardrail should run when default_on=True"
# Test 2: User-injected disable at root level is IGNORED
data_with_disable_root = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"disable_global_guardrails": True,
}
result = custom_guardrail.should_run_guardrail(
data=data_with_disable_root, event_type=GuardrailEventHooks.pre_call
)
assert (
result is True
), "User-injected disable_global_guardrails should be ignored"
# Test 3: User-injected disable in metadata is IGNORED
data_with_disable_metadata = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"metadata": {"disable_global_guardrails": True},
}
result = custom_guardrail.should_run_guardrail(
data=data_with_disable_metadata, event_type=GuardrailEventHooks.pre_call
)
assert result is True, "User-injected metadata disable should be ignored"
# Test 4: Admin-configured disable via user_api_key_metadata IS respected
data_with_admin_disable = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
}
result = custom_guardrail.should_run_guardrail(
data=data_with_admin_disable, event_type=GuardrailEventHooks.pre_call
)
assert result is False, "Admin-configured disable should be respected"
# Test 5: Admin config in metadata isn't shadowed by user-supplied litellm_metadata
data_cross_key = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
"litellm_metadata": {"request_tags": ["user-supplied"]},
}
result = custom_guardrail.should_run_guardrail(
data=data_cross_key, event_type=GuardrailEventHooks.pre_call
)
assert (
result is False
), "Admin config in metadata must not be shadowed by user-supplied litellm_metadata"
# Test 6: After the pre-call strip runs, user-injected
# user_api_key_metadata in the non-authoritative metadata key is gone.
# _get_admin_metadata must then surface admin config unchanged.
data_post_strip = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
"litellm_metadata": {}, # post-strip: attacker payload removed
}
result = custom_guardrail.should_run_guardrail(
data=data_post_strip, event_type=GuardrailEventHooks.pre_call
)
assert (
result is False
), "Admin config in metadata must be respected when other metadata key is empty"
def test_should_run_guardrail_with_opted_out_global_guardrails(self):
"""Test that per-guardrail opt-out only works from admin metadata"""
from litellm.types.guardrails import GuardrailEventHooks
custom_guardrail = CustomGuardrail(
guardrail_name="global_guardrail",
default_on=True,
event_hook=GuardrailEventHooks.pre_call,
)
# Test 1: User-injected opt-out at root level is IGNORED
data_root = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"opted_out_global_guardrails": ["global_guardrail"],
}
assert (
custom_guardrail.should_run_guardrail(
data=data_root, event_type=GuardrailEventHooks.pre_call
)
is True
)
# Test 2: User-injected opt-out in metadata is IGNORED
data_metadata = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"metadata": {"opted_out_global_guardrails": ["global_guardrail"]},
}
assert (
custom_guardrail.should_run_guardrail(
data=data_metadata, event_type=GuardrailEventHooks.pre_call
)
is True
)
# Test 4: a different guardrail in the opt-out list → still runs
data_other = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"metadata": {"opted_out_global_guardrails": ["some_other_guardrail"]},
}
assert (
custom_guardrail.should_run_guardrail(
data=data_other, event_type=GuardrailEventHooks.pre_call
)
is True
)
# Test 5: empty opt-out list → still runs
data_empty = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"metadata": {"opted_out_global_guardrails": []},
}
assert (
custom_guardrail.should_run_guardrail(
data=data_empty, event_type=GuardrailEventHooks.pre_call
)
is True
)
# Test 6: malformed value (bool instead of list) → safely ignored, guardrail runs
data_malformed = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"metadata": {"opted_out_global_guardrails": True},
}
assert (
custom_guardrail.should_run_guardrail(
data=data_malformed, event_type=GuardrailEventHooks.pre_call
)
is True
)
def test_should_run_guardrail_opt_out_does_not_affect_non_global(self):
"""Opt-out list only matters for default_on=True guardrails"""
from litellm.types.guardrails import GuardrailEventHooks
non_global = CustomGuardrail(
guardrail_name="opt_in_guardrail",
default_on=False,
event_hook=GuardrailEventHooks.pre_call,
)
# An opt-in guardrail named in opted_out_global_guardrails is still controlled
# by the explicit `guardrails` request list, not by the global opt-out list.
data = {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "test"}],
"metadata": {
"opted_out_global_guardrails": ["opt_in_guardrail"],
"guardrails": ["opt_in_guardrail"],
},
}
assert (
non_global.should_run_guardrail(
data=data, event_type=GuardrailEventHooks.pre_call
)
is True
)
class TestApplyGuardrailCheck:
def test_apply_guardrail_check_only_on_direct_implementation(self):
"""
Test that "apply_guardrail" in type(callback).__dict__ only returns True
when the object's own class implements the method, not when it's inherited
from a parent class.
This is critical for properly routing guardrail handling to the unified
guardrail handler vs the guardrail's own implementation.
"""
# Parent class with apply_guardrail (CustomGuardrail already has it)
class ParentGuardrail(CustomGuardrail):
"""Parent that inherits apply_guardrail from CustomGuardrail"""
pass
# Child class that only inherits apply_guardrail (doesn't override)
class ChildGuardrailWithoutOverride(ParentGuardrail):
"""Child that only inherits apply_guardrail"""
pass
# Child class that overrides apply_guardrail
class ChildGuardrailWithOverride(ParentGuardrail):
"""Child that overrides apply_guardrail"""
async def apply_guardrail(self, text, language=None, entities=None):
return f"modified: {text}"
# Instantiate the classes
parent_instance = ParentGuardrail()
child_without_override = ChildGuardrailWithoutOverride()
child_with_override = ChildGuardrailWithOverride()
# Test: CustomGuardrail itself has apply_guardrail in its __dict__
assert (
"apply_guardrail" in type(CustomGuardrail()).__dict__
), "CustomGuardrail should have apply_guardrail in its own __dict__"
# Test: ParentGuardrail inherits but doesn't override, so it should NOT be in __dict__
assert (
"apply_guardrail" not in type(parent_instance).__dict__
), "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)"
# Test: ChildGuardrailWithoutOverride only inherits, should NOT be in __dict__
assert (
"apply_guardrail" not in type(child_without_override).__dict__
), "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)"
# Test: ChildGuardrailWithOverride overrides the method, SHOULD be in __dict__
assert (
"apply_guardrail" in type(child_with_override).__dict__
), "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)"
# Verify that all instances still have the method via inheritance (hasattr)
assert hasattr(
parent_instance, "apply_guardrail"
), "All instances should have apply_guardrail via inheritance"
assert hasattr(
child_without_override, "apply_guardrail"
), "All instances should have apply_guardrail via inheritance"
assert hasattr(
child_with_override, "apply_guardrail"
), "All instances should have apply_guardrail via inheritance"
class TestGuardrailLoggingAggregation:
def _make_guardrail(self):
from litellm.types.guardrails import GuardrailEventHooks
return CustomGuardrail(
guardrail_name="test_guardrail",
event_hook=GuardrailEventHooks.pre_call,
)
def _invoke_add_log(self, request_data: dict) -> None:
guardrail = self._make_guardrail()
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"result": "ok"},
request_data=request_data,
guardrail_status="success",
start_time=1.0,
end_time=2.0,
duration=1.0,
masked_entity_count={"EMAIL": 1},
guardrail_provider="presidio",
)
def test_appends_to_existing_metadata_list(self):
request_data = {
"metadata": {
"standard_logging_guardrail_information": [
{"guardrail_name": "existing_guardrail"}
]
}
}
self._invoke_add_log(request_data)
info = request_data["metadata"]["standard_logging_guardrail_information"]
assert isinstance(info, list)
assert len(info) == 2
assert info[0]["guardrail_name"] == "existing_guardrail"
assert info[1]["guardrail_name"] == "test_guardrail"
def test_converts_existing_metadata_dict_to_list(self):
request_data = {
"metadata": {
"standard_logging_guardrail_information": {"guardrail_name": "legacy"}
}
}
self._invoke_add_log(request_data)
info = request_data["metadata"]["standard_logging_guardrail_information"]
assert isinstance(info, list)
assert len(info) == 2
assert info[0]["guardrail_name"] == "legacy"
assert info[1]["guardrail_name"] == "test_guardrail"
def test_appends_to_litellm_metadata(self):
request_data = {
"litellm_metadata": {
"standard_logging_guardrail_information": [
{"guardrail_name": "litellm_existing"}
]
}
}
self._invoke_add_log(request_data)
info = request_data["litellm_metadata"][
"standard_logging_guardrail_information"
]
assert isinstance(info, list)
assert len(info) == 2
assert info[1]["guardrail_name"] == "test_guardrail"
class TestGuardrailOtelSpanEmission:
"""Recording a guardrail emits its otel span inline, so every guardrail
execution produces a span — including the pass-through allow path that never
reaches a post-call hook."""
def _make_guardrail(self):
from litellm.types.guardrails import GuardrailEventHooks
return CustomGuardrail(
guardrail_name="emit_guard",
event_hook=GuardrailEventHooks.pre_call,
)
def _record(self, guardrail, request_data):
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"result": "ok"},
request_data=request_data,
guardrail_status="success",
start_time=1.0,
end_time=2.0,
duration=1.0,
)
def test_emits_span_for_recorded_entry(self, monkeypatch):
captured = []
monkeypatch.setattr(
"litellm.integrations.otel.logger.emit_guardrail_span",
captured.append,
)
request_data = {"metadata": {}}
self._record(self._make_guardrail(), request_data)
assert len(captured) == 1
emitted = captured[0]
recorded = request_data["metadata"]["standard_logging_guardrail_information"][
-1
]
assert emitted is recorded
assert emitted["guardrail_name"] == "emit_guard"
assert emitted["start_time"] == 1.0
assert emitted["end_time"] == 2.0
def test_span_emission_failure_does_not_break_recording(self, monkeypatch):
def _boom(_entry):
raise RuntimeError("otel exporter down")
monkeypatch.setattr(
"litellm.integrations.otel.logger.emit_guardrail_span", _boom
)
request_data = {"metadata": {}}
self._record(self._make_guardrail(), request_data)
info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(info) == 1
assert info[0]["guardrail_name"] == "emit_guard"
class TestGuardrailSensitiveFieldStripping:
"""Tests that secret_fields is stripped from guardrail responses before logging.
Matches the pattern used by Langfuse and Arize integrations which also
pop("secret_fields") to prevent raw Authorization headers from being persisted.
"""
def _make_guardrail(self):
from litellm.types.guardrails import GuardrailEventHooks
return CustomGuardrail(
guardrail_name="test_guardrail",
event_hook=GuardrailEventHooks.pre_call,
)
def test_secret_fields_stripped_from_guardrail_response(self):
"""Ensure secret_fields (containing raw Authorization headers) is not persisted."""
guardrail = self._make_guardrail()
request_data = {"metadata": {}}
guardrail_response_with_secrets = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "hello"}],
"secret_fields": {
"raw_headers": {
"authorization": "Bearer sk-live-secret-key-12345",
"content-type": "application/json",
}
},
"proxy_server_request": {"url": "http://localhost:4000/chat/completions"},
}
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=guardrail_response_with_secrets,
request_data=request_data,
guardrail_status="success",
duration=1.0,
)
info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(info) == 1
logged_response = info[0]["guardrail_response"]
# secret_fields must be stripped
assert "secret_fields" not in logged_response
# Other fields should be preserved
assert "model" in logged_response
assert "messages" in logged_response
assert "proxy_server_request" in logged_response
def test_string_guardrail_response_not_affected(self):
"""String responses (e.g. 'allow', 'deny') should pass through unchanged."""
guardrail = self._make_guardrail()
request_data = {"metadata": {}}
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response="allow",
request_data=request_data,
guardrail_status="success",
duration=0.5,
)
info = request_data["metadata"]["standard_logging_guardrail_information"]
assert info[0]["guardrail_response"] == "allow"
def test_no_authorization_header_in_logged_response(self):
"""Verify no plaintext Authorization header ends up in the logged guardrail response."""
import json
guardrail = self._make_guardrail()
request_data = {"metadata": {}}
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={
"model": "gpt-4",
"secret_fields": {
"raw_headers": {
"authorization": "Bearer sk-live-SHOULD-NOT-APPEAR",
}
},
},
request_data=request_data,
guardrail_status="success",
duration=1.0,
)
logged_response = request_data["metadata"][
"standard_logging_guardrail_information"
][0]["guardrail_response"]
assert "secret_fields" not in logged_response
assert "sk-live-SHOULD-NOT-APPEAR" not in json.dumps(logged_response)
def test_secret_fields_stripped_from_list_dict_response(self):
"""Ensure secret_fields is stripped from List[dict] guardrail responses too."""
guardrail = self._make_guardrail()
request_data = {"metadata": {}}
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=[
{
"result": "ok",
"secret_fields": {
"raw_headers": {"authorization": "Bearer sk-secret"}
},
},
{"result": "also_ok"},
],
request_data=request_data,
guardrail_status="success",
duration=1.0,
)
import json
serialized = json.dumps(request_data)
assert "secret_fields" not in serialized
assert "sk-secret" not in serialized
class TestCustomGuardrailPassthroughSupport:
"""Tests for passthrough endpoint guardrail support - Issue fixes."""
@pytest.mark.asyncio
async def test_async_post_call_success_deployment_hook_with_httpx_response(self):
"""
Test that async_post_call_success_deployment_hook handles raw httpx.Response objects
from passthrough endpoints without crashing with TypeError.
This tests Fix #3: TypeError: TypedDict does not support instance and class checks
"""
import httpx
custom_guardrail = CustomGuardrail()
# Mock the async_post_call_success_hook to return None (guardrail didn't modify response)
custom_guardrail.async_post_call_success_hook = AsyncMock(return_value=None)
# Create a mock httpx.Response object (typical passthrough response)
mock_response = AsyncMock(spec=httpx.Response)
mock_response.status_code = 200
mock_response.text = "Mock response"
request_data = {
"guardrails": ["test_guardrail"],
"user_api_key_user_id": "test_user",
"user_api_key_team_id": "test_team",
"user_api_key_end_user_id": "test_end_user",
"user_api_key_hash": "test_hash",
"user_api_key_request_route": "passthrough_route",
}
# This should not raise TypeError: TypedDict does not support instance and class checks
result = await custom_guardrail.async_post_call_success_deployment_hook(
request_data=request_data,
response=mock_response,
call_type=CallTypes.allm_passthrough_route,
)
# When result is None, should return the original response
assert result == mock_response
@pytest.mark.asyncio
async def test_async_post_call_success_deployment_hook_with_none_call_type(self):
"""
Test that async_post_call_success_deployment_hook handles None call_type gracefully.
This ensures that even if call_type is None (before fix #1), the guardrail doesn't crash.
"""
custom_guardrail = CustomGuardrail()
# Mock the async_post_call_success_hook to return None
custom_guardrail.async_post_call_success_hook = AsyncMock(return_value=None)
mock_response = AsyncMock()
request_data = {
"guardrails": ["test_guardrail"],
"user_api_key_user_id": "test_user",
}
# Call with None call_type - should not crash
result = await custom_guardrail.async_post_call_success_deployment_hook(
request_data=request_data,
response=mock_response,
call_type=None,
)
# Should return the original response when result is None
assert result == mock_response
def test_is_valid_response_type_with_none(self):
"""
Test _is_valid_response_type helper method correctly identifies None as invalid.
This is part of Fix #3: Safely handling TypedDict types that don't support isinstance checks.
"""
custom_guardrail = CustomGuardrail()
# None should be invalid
assert custom_guardrail._is_valid_response_type(None) is False
def test_is_valid_response_type_with_typeddict_error(self):
"""
Test _is_valid_response_type gracefully handles TypeError from TypedDict.
This tests Fix #3: When isinstance() is called with TypedDict types, it raises TypeError.
The method should catch this and allow the response through.
"""
from litellm.types.utils import ModelResponse
custom_guardrail = CustomGuardrail()
# Create a valid LiteLLM response object
response = ModelResponse(
id="test-id",
choices=[],
created=0,
model="test-model",
object="chat.completion",
)
# This should return True (it's a valid response type or TypeError is caught)
result = custom_guardrail._is_valid_response_type(response)
assert result is True
class TestEventTypeLogging:
"""Tests for event_type logging in guardrail information."""
@pytest.mark.asyncio
async def test_log_guardrail_information_infers_event_type_from_async_pre_call_hook(
self,
):
"""
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.pre_call
from async_pre_call_hook function name.
"""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="test_event_type_guardrail",
event_hook=[
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
],
)
@log_guardrail_information
async def async_pre_call_hook(self, data: dict, **kwargs):
return {"result": "pre_call_executed"}
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.async_pre_call_hook(data=request_data)
# Check that the guardrail_mode was set to pre_call (not the full list)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
@pytest.mark.asyncio
async def test_log_guardrail_information_infers_event_type_from_async_post_call_success_hook(
self,
):
"""
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call
from async_post_call_success_hook function name.
"""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="test_event_type_guardrail",
event_hook=[
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
],
)
@log_guardrail_information
async def async_post_call_success_hook(self, data: dict, **kwargs):
return {"result": "post_call_executed"}
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.async_post_call_success_hook(data=request_data)
# Check that the guardrail_mode was set to post_call (not the full list)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
@pytest.mark.asyncio
async def test_log_guardrail_information_infers_event_type_from_async_moderation_hook(
self,
):
"""
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.during_call
from async_moderation_hook function name.
"""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="test_event_type_guardrail",
event_hook=[
GuardrailEventHooks.during_call,
GuardrailEventHooks.post_call,
],
)
@log_guardrail_information
async def async_moderation_hook(self, data: dict, **kwargs):
return {"result": "moderation_executed"}
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.async_moderation_hook(data=request_data)
# Check that the guardrail_mode was set to during_call (not the full list)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.during_call
@pytest.mark.asyncio
async def test_log_guardrail_information_infers_event_type_from_async_post_call_streaming_hook(
self,
):
"""
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call
from async_post_call_streaming_hook function name.
"""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="test_event_type_guardrail",
event_hook=[
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
],
)
@log_guardrail_information
async def async_post_call_streaming_hook(self, data: dict, **kwargs):
return {"result": "streaming_executed"}
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.async_post_call_streaming_hook(data=request_data)
# Check that the guardrail_mode was set to post_call (not the full list)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
@pytest.mark.asyncio
async def test_log_guardrail_information_returns_none_for_unknown_function_name(
self,
):
"""
Test that log_guardrail_information decorator returns None for event_type
when function name doesn't match known patterns, and falls back to self.event_hook.
"""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="test_event_type_guardrail",
event_hook=GuardrailEventHooks.pre_call,
)
@log_guardrail_information
async def some_other_hook(self, data: dict, **kwargs):
return {"result": "other_hook_executed"}
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.some_other_hook(data=request_data)
# Check that the guardrail_mode falls back to self.event_hook
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
def test_add_standard_logging_uses_event_type_over_event_hook(self):
"""
Test that add_standard_logging_guardrail_information_to_request_data
prioritizes event_type parameter over self.event_hook.
"""
from litellm.types.guardrails import GuardrailEventHooks
guardrail = CustomGuardrail(
guardrail_name="test_guardrail",
event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
)
request_data = {"metadata": {}}
# Call with explicit event_type
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"result": "ok"},
request_data=request_data,
guardrail_status="success",
event_type=GuardrailEventHooks.post_call,
)
# Should use the provided event_type (post_call), not the full event_hook list
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
@pytest.mark.asyncio
async def test_log_guardrail_information_skips_auto_record_if_function_already_recorded(
self,
):
"""When a wrapped guardrail function records its own entry directly
(e.g. block_code_execution.apply_guardrail records a rich
``[detections...]`` payload), the decorator must NOT also append its
own ``"allow"``/raw-response entry — otherwise every backend
(OTEL spans, Datadog, Langfuse, spend logs) double-records one
logical guardrail invocation."""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="block-code",
event_hook=GuardrailEventHooks.pre_call,
)
@log_guardrail_information
async def apply_guardrail(self, inputs, request_data, **kwargs):
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=[{"action_taken": "block"}],
request_data=request_data,
guardrail_status="success",
event_type=GuardrailEventHooks.pre_call,
)
return inputs
guardrail = TestGuardrail()
request_data = {"metadata": {}}
await guardrail.apply_guardrail(
inputs={"texts": ["x"]}, request_data=request_data
)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1, (
f"Decorator must not double-record when the wrapped function "
f"already appended its own entry; got {len(logged_info)} entries"
)
assert logged_info[0]["guardrail_response"] == [{"action_taken": "block"}]
@pytest.mark.asyncio
async def test_log_guardrail_information_skips_auto_record_on_exception_if_function_already_recorded(
self,
):
"""Same as above on the failure path: if the wrapped function
appended an entry in its ``finally`` block before re-raising, the
decorator must just re-raise without auto-recording on top."""
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class TestGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="block-code",
event_hook=GuardrailEventHooks.pre_call,
)
@log_guardrail_information
async def apply_guardrail(self, inputs, request_data, **kwargs):
try:
raise ValueError("blocked")
finally:
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=[{"action_taken": "block"}],
request_data=request_data,
guardrail_status="guardrail_intervened",
event_type=GuardrailEventHooks.pre_call,
)
guardrail = TestGuardrail()
request_data = {"metadata": {}}
with pytest.raises(ValueError, match="blocked"):
await guardrail.apply_guardrail(
inputs={"texts": ["x"]}, request_data=request_data
)
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_status"] == "guardrail_intervened"
def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none(
self,
):
"""
Test that add_standard_logging_guardrail_information_to_request_data
falls back to self.event_hook when event_type is None.
"""
from litellm.types.guardrails import GuardrailEventHooks
guardrail = CustomGuardrail(
guardrail_name="test_guardrail",
event_hook=GuardrailEventHooks.pre_call,
)
request_data = {"metadata": {}}
# Call with event_type=None
guardrail.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"result": "ok"},
request_data=request_data,
guardrail_status="success",
event_type=None,
)
# Should fall back to self.event_hook
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(logged_info) == 1
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
class TestTracingFieldsPopulation:
"""Verify add_standard_logging_guardrail_information_to_request_data passes tracing_detail fields."""
def test_new_fields_set_on_slg(self):
cg = CustomGuardrail(guardrail_name="test-rail")
request_data = {"metadata": {}}
cg.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"result": "ok"},
request_data=request_data,
guardrail_status="success",
tracing_detail=GuardrailTracingDetail(
guardrail_id="rail-123",
policy_template="EU AI Act Article 5",
detection_method="regex",
confidence_score=0.95,
match_details=[{"type": "pattern", "action_taken": "BLOCK"}],
patterns_checked=12,
alert_recipients=["admin@example.com"],
),
)
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(slg_list) == 1
slg = slg_list[0]
assert slg["guardrail_id"] == "rail-123"
assert slg["policy_template"] == "EU AI Act Article 5"
assert slg["detection_method"] == "regex"
assert slg["confidence_score"] == 0.95
assert slg["patterns_checked"] == 12
assert slg["alert_recipients"] == ["admin@example.com"]
assert len(slg["match_details"]) == 1
def test_new_fields_default_to_absent(self):
"""When tracing_detail is not passed, new fields are absent from the SLG dict."""
cg = CustomGuardrail(guardrail_name="test-rail")
request_data = {"metadata": {}}
cg.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response="ok",
request_data=request_data,
guardrail_status="success",
)
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
assert slg.get("guardrail_id") is None
assert slg.get("policy_template") is None
assert slg.get("confidence_score") is None
def test_multiple_guardrails_with_different_policies(self):
"""One request, multiple guardrails each with own policy_template."""
cg1 = CustomGuardrail(guardrail_name="rail-1")
cg2 = CustomGuardrail(guardrail_name="rail-2")
request_data = {"metadata": {}}
cg1.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response="ok",
request_data=request_data,
guardrail_status="success",
tracing_detail=GuardrailTracingDetail(policy_template="GDPR"),
)
cg2.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response="blocked",
request_data=request_data,
guardrail_status="guardrail_intervened",
tracing_detail=GuardrailTracingDetail(
policy_template="EU AI Act Article 5"
),
)
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(slg_list) == 2
assert slg_list[0]["policy_template"] == "GDPR"
assert slg_list[1]["policy_template"] == "EU AI Act Article 5"
def test_classification_field_passed_through(self):
"""Classification dict for LLM-judge guardrails is passed through."""
cg = CustomGuardrail(guardrail_name="judge-rail")
request_data = {"metadata": {}}
classification = {
"flagged": True,
"category": "workplace_emotion_recognition",
"article_reference": "Article 5(1)(f)",
"confidence": 0.94,
"reason": "Request asks to analyze employee sentiment",
}
cg.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response="blocked",
request_data=request_data,
guardrail_status="guardrail_intervened",
tracing_detail=GuardrailTracingDetail(
classification=classification,
detection_method="llm-judge",
confidence_score=0.94,
),
)
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
assert slg["classification"] == classification
assert slg["detection_method"] == "llm-judge"
assert slg["confidence_score"] == 0.94
class TestCustomGuardrailSpendLogMatchRedaction:
"""Guardrail JSON persisted via standard_logging must not contain raw match spans."""
def test_add_standard_logging_redacts_nested_match(self):
cg = CustomGuardrail(guardrail_name="test-rail")
raw = {
"assessments": [
{
"sensitiveInformationPolicy": {
"piiEntities": [
{"type": "NAME", "match": "GG", "action": "BLOCKED"}
]
}
}
]
}
request_data: dict = {"metadata": {}}
cg.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=raw,
request_data=request_data,
guardrail_status="guardrail_intervened",
)
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
assert (
slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][
"piiEntities"
][0]["match"]
== "[REDACTED]"
)
assert (
raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
"match"
]
== "GG"
)
def test_add_standard_logging_redacts_regex_field(self):
cg = CustomGuardrail(guardrail_name="test-rail")
raw = {"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]}
request_data: dict = {"metadata": {}}
cg.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=raw,
request_data=request_data,
guardrail_status="success",
)
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
assert slg["guardrail_response"]["filters"][0]["regex"] == "[REDACTED]"
assert raw["filters"][0]["regex"] == r"\d{3}-\d{2}-\d{4}"
class TestGuardrailInterventionClassification:
"""A routing decision is a deliberate guardrail intervention, not a failure."""
def test_sensitive_data_route_exception_is_intervention(self):
from litellm.exceptions import SensitiveDataRouteException
exc = SensitiveDataRouteException(
route_to_model="on-prem-model",
session_id="sess-1",
guardrail_name="pii-rail",
)
assert CustomGuardrail._is_guardrail_intervention(exc) is True
@pytest.mark.asyncio
async def test_routing_logged_as_intervened_not_failed(self):
from litellm.exceptions import SensitiveDataRouteException
from litellm.integrations.custom_guardrail import log_guardrail_information
from litellm.types.guardrails import GuardrailEventHooks
class RoutingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="pii-rail",
event_hook=GuardrailEventHooks.pre_call,
)
@log_guardrail_information
async def async_pre_call_hook(self, data, **kwargs):
raise SensitiveDataRouteException(
route_to_model="on-prem-model",
session_id="sess-1",
guardrail_name=self.guardrail_name,
)
guardrail = RoutingGuardrail()
request_data: dict = {"metadata": {}}
with pytest.raises(SensitiveDataRouteException):
await guardrail.async_pre_call_hook(data=request_data)
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
assert slg["guardrail_status"] == "guardrail_intervened"