scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(responses): run the guardrails a Responses WebSocket frame selects

Browse source
This commit is contained in:
tavren27 2026-10-27 23:55:31 -07:00
parent c7c895db45
commit e8b8bf73c4
3 changed files with 196 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm/proxy/litellm_pre_call_utils.py
View file

@ -4220,7 +4220,7 @@ def _policy_match_tags(request_data: Mapping[str, object], metadata_variable_nam
async def add_guardrails_from_policy_engine(
data: dict,
data: dict[str, object], # mutable-ok: pops the request's policies and records the resolved guardrails in place
metadata_variable_name: str,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
@ -4263,10 +4263,11 @@ async def add_guardrails_from_policy_engine(
all_tags: Final = _policy_match_tags(data, metadata_variable_name)
_team_alias: Final = user_api_key_dict.team_alias
_key_alias: Final = user_api_key_dict.key_alias
_model: Final = data.get("model")
context: Final = PolicyMatchContext(
team_alias=_team_alias if isinstance(_team_alias, str) else None,
key_alias=_key_alias if isinstance(_key_alias, str) else None,
model=data.get("model"),
model=_model if isinstance(_model, str) else None,
tags=all_tags,
)

89
litellm/proxy/response_api_endpoints/endpoints.py
View file

@ -1,5 +1,6 @@
import asyncio
import contextlib
import copy
import json
import uuid
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence
@ -75,6 +76,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry
from litellm.proxy.litellm_pre_call_utils import (
_CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list
_UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list
add_guardrails_from_policy_engine,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
request_tags_from_request_data,
@ -2017,8 +2019,12 @@ _PII_TOKENS: Final = TypeAdapter(dict[str, str])
_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
_GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools")
_CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata")
# the guardrail list is merged from key, team and the admission frame; a later frame cannot shrink it
_CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies")
_PROXY_ONLY_FRAME_FIELDS: Final = (*_CLIENT_GUARDRAIL_SELECTORS, "include_guardrail_response")
_NESTED_BODY_ONLY_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS)
_GUARDRAIL_NAMES: Final = TypeAdapter(list[str])
_OBJECTS: Final = TypeAdapter(list[object])
_TURN_COPIED_CONTAINERS: Final = (list, dict)
class _PreCallGuardrailWalk(Protocol):
@ -2057,21 +2063,21 @@ def _responses_ws_frame_guard(
async def guard(frame: str) -> GuardedResponsesWsFrame:
frame_obj: Final = _WS_FRAME.validate_json(frame)
nested: Final = frame_obj.get("response")
if isinstance(nested, dict) and any(key in frame_obj for key in _GUARDED_FRAME_FIELDS):
if isinstance(nested, dict) and any(key in frame_obj for key in _NESTED_BODY_ONLY_FIELDS):
raise ValueError(
"response.create carries input, instructions or tools both at the top level and under response"
"response.create carries input, instructions, tools or a guardrail field both at the top level "
"and under response"
)
body: Final = _WS_FRAME.validate_python(nested) if isinstance(nested, dict) else frame_obj
frame_data: Final = {
**{
key: _without_untrusted_metadata(key, value)
for key, value in body.items()
if key != "type" and key not in _CLIENT_GUARDRAIL_SELECTORS
if key != "type" and key not in _PROXY_ONLY_FRAME_FIELDS
},
"model": model,
"litellm_call_id": str(uuid.uuid4()),
# shallow: the bucket carries proxy-owned objects (the OTEL parent span) that cannot be deep-copied
metadata_key: {**bucket},
metadata_key: await _metadata_with_frame_guardrails(bucket, body, model, metadata_key, user_api_key_dict),
}
try:
guarded_data: Final = await proxy_logging_obj.pre_call_hook(
@ -2089,7 +2095,11 @@ def _responses_ws_frame_guard(
raise rerouted
# a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them
guarded_fields: Final = {key: guarded_data[key] for key in _GUARDED_FRAME_FIELDS if key in guarded_data}
kept: Final = {key: value for key, value in body.items() if key not in _GUARDED_FRAME_FIELDS}
kept: Final = {
key: value
for key, value in body.items()
if key not in _GUARDED_FRAME_FIELDS and key not in _PROXY_ONLY_FRAME_FIELDS
}
rewritten: Final = (
{**frame_obj, "response": {**kept, **guarded_fields}}
if isinstance(nested, dict)
@ -2106,6 +2116,71 @@ def _responses_ws_frame_guard(
return guard
async def _metadata_with_frame_guardrails(
bucket: Mapping[str, object],
body: Mapping[str, object],
model: object,
metadata_key: str,
user_api_key_dict: UserAPIKeyAuth,
) -> Mapping[str, object]:
"""A frame selects guardrails the way an HTTP ``/v1/responses`` body does, adding to the key, team and policy
guardrails the socket was admitted with, so a frame can never shrink that list"""
metadata: Final = (
await _metadata_with_frame_policies(bucket, body["policies"], model, metadata_key, user_api_key_dict)
if "policies" in body
else bucket
)
existing_guardrails: Final = metadata.get("guardrails")
existing_config: Final = metadata.get("guardrail_config")
return {
**metadata,
**(
{
"guardrails": [
*(_OBJECTS.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else ()),
*_GUARDRAIL_NAMES.validate_python(body["guardrails"]),
]
}
if "guardrails" in body
else {}
),
**(
{
"guardrail_config": {
**(_WS_FRAME.validate_python(existing_config) if isinstance(existing_config, dict) else {}),
**_WS_FRAME.validate_python(body["guardrail_config"]),
}
}
if "guardrail_config" in body
else {}
),
}
async def _metadata_with_frame_policies(
bucket: Mapping[str, object],
policies: object,
model: object,
metadata_key: str,
user_api_key_dict: UserAPIKeyAuth,
) -> Mapping[str, object]:
selection: Final = {
"model": model,
"policies": policies,
metadata_key: {key: _turn_copy(value) for key, value in bucket.items()},
}
await add_guardrails_from_policy_engine(
data=selection, metadata_variable_name=metadata_key, user_api_key_dict=user_api_key_dict
)
return _WS_FRAME.validate_python(selection[metadata_key])
def _turn_copy(value: object) -> object:
"""Shallow, since the bucket carries proxy-owned objects (the OTEL parent span) that cannot be copied, and the
policy engine writes its lists in place"""
return copy.copy(value) if type(value) in _TURN_COPIED_CONTAINERS else value
def _without_untrusted_metadata(key: str, value: object) -> object:
"""Like the HTTP path, drop the proxy-owned ``user_api_key_*`` and guardrail control keys a client may forge."""
if key not in _CLIENT_METADATA_KEYS or not isinstance(value, dict):

113
tests/unit/proxy/response_api_endpoints/test_endpoints.py
View file

@ -2226,6 +2226,112 @@ class TestResponsesWSFirstFrameModelAuth:
assert "guardrails" not in (scanned.get("metadata") or {})
assert scanned["litellm_metadata"]["guardrails"] == ["key-required"]
@pytest.mark.asyncio
@pytest.mark.parametrize("nested", [False, True])
async def test_frame_guard_runs_the_guardrails_a_frame_selects_on_top_of_the_admitted_ones(self, nested: bool):
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
scans: list[dict] = []
async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
scans.append(data)
return data
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
admitted = {
"model": "gpt-4o-mini",
"litellm_metadata": {
"guardrails": ["key-required"],
"guardrail_config": {"key-required": {"threshold": 1}},
},
}
guard = _responses_ws_frame_guard(admitted, MagicMock(), proxy_logging_obj, MagicMock())
selection = {
"guardrails": ["frame-pick"],
"guardrail_config": {"frame-pick": {"threshold": 2}},
"include_guardrail_response": True,
}
body = {"input": "hi", "store": False, **selection}
frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body}
selected = await guard(json.dumps(frame))
plain = await guard(json.dumps({"type": "response.create", "input": "again"}))
assert scans[0]["litellm_metadata"]["guardrails"] == ["key-required", "frame-pick"]
assert scans[0]["litellm_metadata"]["guardrail_config"] == {
"key-required": {"threshold": 1},
"frame-pick": {"threshold": 2},
}
assert not set(selection) & set(scans[0])
unselected = {"input": "hi", "store": False}
assert json.loads(selected.frame) == (
{"type": "response.create", "response": unselected} if nested else {"type": "response.create", **unselected}
)
assert scans[1]["litellm_metadata"]["guardrails"] == ["key-required"]
assert scans[1]["litellm_metadata"]["guardrail_config"] == {"key-required": {"threshold": 1}}
assert json.loads(plain.frame) == {"type": "response.create", "input": "again"}
assert admitted["litellm_metadata"] == {
"guardrails": ["key-required"],
"guardrail_config": {"key-required": {"threshold": 1}},
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"selection", [{"guardrails": "frame-pick"}, {"guardrails": [{"name": "frame-pick"}]}, {"guardrail_config": []}]
)
async def test_frame_guard_refuses_a_frame_whose_guardrail_selection_is_malformed(self, selection):
from pydantic import ValidationError
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"])
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}},
MagicMock(),
proxy_logging_obj,
MagicMock(),
)
with pytest.raises(ValidationError):
await guard(json.dumps({"type": "response.create", "input": "hi", **selection}))
proxy_logging_obj.pre_call_hook.assert_not_awaited()
@pytest.mark.asyncio
async def test_frame_guard_runs_the_guardrails_of_the_policies_a_frame_selects(self):
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"])
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}},
UserAPIKeyAuth(api_key="hashed"),
proxy_logging_obj,
MagicMock(),
)
policy_registry = get_policy_registry()
policy_registry._policies = {}
policy_registry._policies_by_id = {}
policy_registry._initialized = False
try:
policy_registry.add_policy("frame-policy", Policy(guardrails=PolicyGuardrails(add=["policy-pick"])))
guarded = await guard(json.dumps({"type": "response.create", "input": "hi", "policies": ["frame-policy"]}))
finally:
policy_registry._policies = {}
policy_registry._policies_by_id = {}
policy_registry._initialized = False
scanned = proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]
assert scanned["litellm_metadata"]["guardrails"] == ["key-required", "policy-pick"]
assert "policies" not in scanned
assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"}
@pytest.mark.asyncio
async def test_frame_guard_copies_an_admitted_bucket_that_holds_the_otel_parent_span(self):
import threading
@ -2367,7 +2473,10 @@ class TestResponsesWSFirstFrameModelAuth:
assert dict(guarded.pii_tokens) == {"<EMAIL_ADDRESS_1>": "alice@example.com"}
@pytest.mark.asyncio
async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self):
@pytest.mark.parametrize(
"top_level", [{"input": "forbidden"}, {"guardrails": ["frame-pick"]}, {"include_guardrail_response": True}]
)
async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self, top_level):
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
proxy_logging_obj = MagicMock()
@ -2377,7 +2486,7 @@ class TestResponsesWSFirstFrameModelAuth:
)
with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"):
await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}}))
await guard(json.dumps({"type": "response.create", **top_level, "response": {}}))
proxy_logging_obj.pre_call_hook.assert_not_awaited()