scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix(policy_engine): keep post_call pipeline guardrail logging and reject background bypass

Browse source 
Post_call pipelines run step hooks against a copied request dict, so guardrail
writes into the metadata bucket (applied_guardrails for the response header,
standard_logging_guardrail_information for spend logs) were dropped when the
guardrail was the first writer. Merge those writes back onto the request on the
post_call allow path, keeping the request payload and the executor's per-step
guardrails activation flag out of it.

Background /v1/responses requests dodge the streaming 400: pre_call sees stream
unset, then the polling task forces stream=true with pre-call logic skipped and
the streaming branch returns before post_call_success_hook, silently bypassing
post_call pipelines. Reject background=true at pre_call the same way as
stream=true.

Also pin the run_in_parallel pipeline-managed exclusion in both hook loops with
regression tests.
This commit is contained in:
mateo-berri 2026-08-29 00:14:23 -07:00
parent 55569729b0
commit 996019cd23
2 changed files with 187 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
litellm/proxy/utils.py
View file

@ -457,8 +457,37 @@ def _pipeline_managed_guardrail_names(
)
def _merge_pipeline_metadata_bucket(data: dict, bucket_key: str, modified_bucket_value: object) -> None:
if not isinstance(modified_bucket_value, dict):
return
modified_bucket: Final = cast("dict[str, object]", modified_bucket_value) # cast-ok: metadata buckets are str-keyed
surviving_writes: Final = {key: value for key, value in modified_bucket.items() if key != "guardrails"}
existing_bucket: Final = data.get(bucket_key)
if isinstance(existing_bucket, dict):
cast("dict[str, object]", existing_bucket).update(surviving_writes) # cast-ok: metadata buckets are str-keyed
else:
data[bucket_key] = surviving_writes
def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, object]) -> None:
"""
Copy metadata-bucket writes from a pipeline's working copy back onto the request.
Post_call pipelines run step hooks against a copied request dict so the payload
already sent upstream stays untouched, but hooks record proxy-internal logging
state in the metadata buckets (``applied_guardrails`` for response headers,
``standard_logging_guardrail_information`` for spend logs), and those writes
must reach the request dict the proxy keeps reading after the pipeline returns.
The ``guardrails`` key is the executor's per-step activation flag for
``should_run_guardrail``, not a hook write, so it stays in the working copy.
"""
for bucket_key in ("metadata", "litellm_metadata"):
_merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key))
def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None:
if data.get("stream") is not True:
if data.get("stream") is not True and data.get("background") is not True:
return
post_call_policies: Final = tuple(
policy_name for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call"
@ -470,9 +499,10 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None
detail={
"error": {
"message": (
"Policies with post_call guardrail pipelines cannot govern streaming responses yet: "
f"{', '.join(post_call_policies)}. Retry with stream=false, or move these policies' output "
"guardrails from pipeline steps to guardrails.add, which scans streamed output."
"Policies with post_call guardrail pipelines cannot govern streaming or background "
f"responses yet: {', '.join(post_call_policies)}. Retry with stream=false and "
"background=false, or move these policies' output guardrails from pipeline steps to "
"guardrails.add, which scans streamed output."
),
"type": "guardrail_pipeline_error",
"policies": list(post_call_policies),
@ -1667,11 +1697,16 @@ class ProxyLogging:
Returns data dict if allowed, raises on block/modify_response.
``original_response`` is set on the post_call path, where the request
payload (already sent upstream) must stay untouched; a replacement
response carried in ``modified_data`` is adopted by the caller.
response carried in ``modified_data`` is adopted by the caller, and
metadata-bucket writes (applied guardrails, guardrail logging info)
are merged back so headers and spend logs still see them.
"""
if result.terminal_action == "allow":
if result.modified_data is not None and original_response is None:
data.update(result.modified_data)
if result.modified_data is not None:
if original_response is None:
data.update(result.modified_data)
else:
_merge_pipeline_metadata_writes(data, result.modified_data)
return data
if result.terminal_action == "block":

148
tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -23,6 +23,7 @@ from litellm.integrations.custom_guardrail import (
ModifyResponseException,
)
from litellm.integrations.prometheus import PrometheusLogger
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.policy_engine.pipeline_types import (
@ -1139,18 +1140,132 @@ def test_handle_pipeline_result_modify_response_carries_original_response():
assert info.value.original_response is response
def test_handle_pipeline_result_allow_discards_modifications_on_post_call():
def test_handle_pipeline_result_allow_on_post_call_keeps_metadata_writes_only():
data = {"a": 1, "metadata": {"guardrails": ["other"]}}
result = MagicMock()
result.terminal_action = "allow"
result.modified_data = {"metadata": {"guardrails": ["gr-post"]}, "response": object()}
result.modified_data = {
"a": 2,
"metadata": {"guardrails": ["other"], "applied_guardrails": ["gr-post"]},
"response": object(),
}
out = ProxyLogging._handle_pipeline_result(
result=result, data=data, policy_name="p", original_response=litellm.ModelResponse()
)
assert out is data
assert data == {"a": 1, "metadata": {"guardrails": ["other"]}}
assert data["a"] == 1
assert "response" not in data
assert data["metadata"] == {"guardrails": ["other"], "applied_guardrails": ["gr-post"]}
@pytest.mark.asyncio
async def test_post_call_pipeline_guardrail_metadata_writes_reach_request_data(
proxy_logging, make_user_api_key_auth, monkeypatch
):
class HeaderWritingGuardrail(CustomGuardrail):
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name="gr-post")
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"verdict": "pass"},
request_data=data,
guardrail_status="success",
)
return None
monkeypatch.setattr(
litellm,
"callbacks",
[HeaderWritingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)],
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data()
await proxy_logging.post_call_success_hook(
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
)
assert data["metadata"]["applied_guardrails"] == ["gr-post"]
slg_entries = data["metadata"]["standard_logging_guardrail_information"]
assert len(slg_entries) == 1
assert slg_entries[0]["guardrail_name"] == "gr-post"
@pytest.mark.asyncio
async def test_post_call_pipeline_managed_parallel_guardrail_runs_exactly_once(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Dict[str, Any] = {"count": 0}
class CountingGuardrail(CustomGuardrail):
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
seen["count"] += 1
return None
monkeypatch.setattr(
litellm,
"callbacks",
[
CountingGuardrail(
guardrail_name="gr-post",
event_hook=GuardrailEventHooks.post_call,
default_on=True,
run_in_parallel=True,
)
],
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data()
await proxy_logging.post_call_success_hook(
data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth()
)
assert seen["count"] == 1
@pytest.mark.asyncio
async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Dict[str, Any] = {"count": 0}
class CountingGuardrail(CustomGuardrail):
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
seen["count"] += 1
return data
pre_call_pipeline = GuardrailPipeline(
mode="pre_call",
steps=[PipelineStep(guardrail="gr-pre", on_pass="allow", on_fail="block")],
)
monkeypatch.setattr(
litellm,
"callbacks",
[
CountingGuardrail(
guardrail_name="gr-pre",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
run_in_parallel=True,
)
],
)
data = {
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"metadata": {
"_guardrail_pipelines": [("request-governance", pre_call_pipeline)],
"_pipeline_managed_guardrails": {"gr-pre"},
},
}
await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion"
)
assert seen["count"] == 1
@pytest.mark.asyncio
@ -1173,6 +1288,26 @@ async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline(
assert "stream=false" in info.value.detail["error"]["message"]
@pytest.mark.asyncio
async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline(
proxy_logging, make_user_api_key_auth, monkeypatch
):
monkeypatch.setattr(litellm, "callbacks", [])
data = _post_call_pipeline_data(background=True)
with pytest.raises(HTTPException) as info:
await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(),
data=data,
call_type="aresponses",
guardrails_only=True,
)
assert info.value.status_code == 400
assert info.value.detail["error"]["policies"] == ["response-governance"]
assert "background=false" in info.value.detail["error"]["message"]
def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call():
post_call = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g", on_fail="block")])
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")])
@ -1183,6 +1318,12 @@ def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_c
)
is None
)
assert (
_raise_for_streaming_post_call_pipelines(
{"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}
)
is None
)
assert _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}) is None
assert (
_raise_for_streaming_post_call_pipelines(
@ -1191,3 +1332,4 @@ def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_c
is None
)
assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None
assert _raise_for_streaming_post_call_pipelines({"background": True}) is None