scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-18 04:44:37 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

fix(responses): run pre_call guardrails on every Responses WebSocket turn and bill each turn its own guardrail cost (#47013)

Browse source 
* fix(responses): run pre_call guardrails on every Responses WebSocket turn and bill each turn its own guardrail cost

* test(responses): keep the websocket unit test file on main's formatting

* fix(responses): type the websocket frame guard for the lint gates

* fix(responses): pass the websocket frame guard as a named parameter and mark the guardrail dict edits

* fix(responses): bill a guardrail-rejected Responses WebSocket frame as a failed request with its guardrail cost

* fix(responses): type the recorded guardrail records as Mapping for the lint gate

* fix(responses): bill a rejected websocket frame through the spend logger only so the connection keeps its rate-limit slot

* fix(responses): guard nested response.create frames, strip forged user_api_key_* frame metadata, copy guarded tools back and give each rejected frame its own spend row id

* fix(responses): mark the spend logger's bare-dict failure hook for the type gate

* fix(responses): reject response.create frames that carry input at both levels and keep billing of a rejected frame non-blocking

* fix(responses): name the non-blocking billing catch for the strict lint gate

* fix(responses): keep the admitted guardrail list on every websocket frame, a frame cannot replace it with its own guardrails list

* test(responses): wait for the key spend total before asserting it in the websocket guardrail cost test

* fix(responses): copy the admitted metadata shallowly per websocket frame, the OTEL parent span cannot be deep-copied

* fix(responses): reject a WebSocket turn a pre_call guardrail rerouted to another model

* fix(responses): bill a rerouted WS turn and carry Presidio pii_tokens from the frame guard to the socket unmasker

* fix(responses): default pii_tokens through default_factory, Python 3.13 rejects a mappingproxy dataclass default

* fix(responses): run the per-frame guardrail walk on managed Responses WebSockets too

* fix(responses): unmask managed Responses WebSocket replies with the frame guard's Presidio pii_tokens

* fix(responses): strip client guardrail cost records from WS frames so a blocked frame bills the real scan

* fix(responses): forward a WS frame without the fields a guardrail deleted

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-10-17 11:00:51 -07:00 • committed by GitHub
parent dd8b2853a0
commit 571ef3e2ae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 1225 additions and 133 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

64
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -216,6 +216,7 @@ if TYPE_CHECKING:
FakeAnthropicMessagesStreamIterator,
)
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
from litellm.responses.streaming_iterator import ResponsesWsFrameGuard
from litellm.types.google_genai.main import GenerateContentResponse
from litellm.types.llms.openai_evals import (
CancelEvalResponse,
@ -6867,6 +6868,7 @@ class BaseLLMHTTPHandler:
custom_llm_provider: str | None = None,
first_message: str | None = None,
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
frame_guard: "ResponsesWsFrameGuard | None" = None,
**kwargs: object,
) -> Exception | None:
"""
@ -6882,6 +6884,35 @@ class BaseLLMHTTPHandler:
"""
_ws_quota_callbacks: Final = _collect_ws_project_quota_callbacks()
_ws_guardrail_callbacks: list = []
_ws_output_guardrail_callbacks: list = []
try:
import litellm as _litellm
# Use duck-typing so any guardrail that exposes the PII
# masking interface works, not just _OPTIONAL_PresidioPIIMasking.
# This avoids a layering violation (SDK importing from proxy).
_ws_guardrail_callbacks = [
cb
for cb in _litellm.callbacks
if callable(getattr(cb, "check_pii", None))
and callable(getattr(cb, "get_presidio_settings_from_request_data", None))
and callable(getattr(cb, "_unmask_pii_text", None))
and getattr(cb, "output_parse_pii", False)
]
_ws_output_guardrail_callbacks = [
cb
for cb in _litellm.callbacks
if callable(getattr(cb, "check_pii", None))
and callable(getattr(cb, "get_presidio_settings_from_request_data", None))
and getattr(cb, "apply_to_output", False)
]
except Exception as _guardrail_exc:
verbose_logger.warning(
"Responses WebSocket: failed to collect guardrail callbacks — PII masking will be skipped. Error: %s",
_guardrail_exc,
)
if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket():
from litellm.responses.streaming_iterator import (
ManagedResponsesWebSocketHandler,
@ -6899,6 +6930,8 @@ class BaseLLMHTTPHandler:
custom_llm_provider=custom_llm_provider,
first_message=first_message,
quota_callbacks=_ws_quota_callbacks,
guardrail_callbacks=_ws_guardrail_callbacks, # pyright: ignore[reportUnknownArgumentType] # duck-typed callbacks from litellm.callbacks
frame_guard=frame_guard,
**kwargs,
)
await handler.run()
@ -6980,36 +7013,6 @@ class BaseLLMHTTPHandler:
if litellm_metadata:
_request_data["litellm_metadata"] = litellm_metadata
_ws_guardrail_callbacks: list = []
_ws_output_guardrail_callbacks: list = []
try:
import litellm as _litellm
# Use duck-typing so any guardrail that exposes the PII
# masking interface works, not just _OPTIONAL_PresidioPIIMasking.
# This avoids a layering violation (SDK importing from proxy).
_ws_guardrail_callbacks = [
cb
for cb in _litellm.callbacks
if callable(getattr(cb, "check_pii", None))
and callable(getattr(cb, "get_presidio_settings_from_request_data", None))
and callable(getattr(cb, "_unmask_pii_text", None))
and getattr(cb, "output_parse_pii", False)
]
_ws_output_guardrail_callbacks = [
cb
for cb in _litellm.callbacks
if callable(getattr(cb, "check_pii", None))
and callable(getattr(cb, "get_presidio_settings_from_request_data", None))
and getattr(cb, "apply_to_output", False)
]
except Exception as _guardrail_exc:
verbose_logger.warning(
"Responses WebSocket: failed to collect guardrail "
"callbacks — PII masking will be skipped. Error: %s",
_guardrail_exc,
)
streaming: Final = ResponsesWebSocketStreaming(
websocket=websocket,
backend_ws=cast(ClientConnection, backend_ws),
@ -7023,6 +7026,7 @@ class BaseLLMHTTPHandler:
authorized_model=model,
custom_llm_provider=custom_llm_provider,
request_defaults=request_defaults,
frame_guard=frame_guard,
)
return await streaming.bidirectional_forward()

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

@ -1,6 +1,7 @@
import asyncio
import contextlib
import json
import uuid
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from enum import Enum
@ -28,10 +29,11 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError
from starlette.websockets import WebSocket, WebSocketDisconnect
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES
from litellm.integrations.custom_guardrail import ModifyResponseException
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
from litellm.litellm_core_utils.tool_fields import first_validation_error
@ -58,11 +60,17 @@ from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # leg
read_request_body,
safe_set_request_parsed_body,
)
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decrypted_response_id
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
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import (
request_tags_from_request_data,
)
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard
from litellm.router import Router
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.llms.openai import (
@ -1810,6 +1818,143 @@ async def _enforce_responses_ws_first_frame_model_auth(
)
_WS_FRAME: Final = TypeAdapter(dict[str, object])
_GUARDRAIL_INFORMATION: Final = TypeAdapter(list[dict[str, object]])
_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")
class _PreCallGuardrailWalk(Protocol):
async def pre_call_hook(
self,
*,
user_api_key_dict: UserAPIKeyAuth,
data: dict[str, object], # mutable-ok: ProxyLogging.pre_call_hook hands the guardrails a dict they edit
call_type: Literal["aresponses"],
guardrails_only: bool,
) -> dict[str, object]: ... # mutable-ok: ProxyLogging.pre_call_hook returns the guardrails' own dict
def _responses_ws_frame_guard(
admitted: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: _PreCallGuardrailWalk
) -> ResponsesWsFrameGuard:
"""Socket admission skips guardrails (the first frame's ``input`` is not part of the admitted request), so every
``response.create`` frame gets the pre_call guardrail walk the HTTP ``/v1/responses`` path runs on its body."""
metadata_key: Final = get_metadata_variable_name_from_kwargs(admitted)
admitted_bucket: Final = admitted.get(metadata_key)
bucket: Final = {
key: value
for key, value in (
_WS_FRAME.validate_python(admitted_bucket) if isinstance(admitted_bucket, dict) else {}
).items()
if key != _GUARDRAIL_INFORMATION_KEY
}
model: Final = admitted.get("model")
# a blocked frame is billed as its own failed request; only the first turn settles the admission reservation
rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None})
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):
raise ValueError(
"response.create carries input, instructions or tools 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
},
"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},
}
try:
guarded_data: Final = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_dict, data=frame_data, call_type="aresponses", guardrails_only=True
)
except Exception as e:
await _bill_rejected_frame(rejected_frame_auth, e, frame_data)
raise
if guarded_data.get("model") != model:
rerouted: Final = ValueError(
f"a guardrail rerouted this turn to {guarded_data.get('model')!r}; an open Responses WebSocket"
f" is bound to {model!r}, send this turn over HTTP /v1/responses"
)
await _bill_rejected_frame(rejected_frame_auth, rerouted, frame_data)
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}
rewritten: Final = (
{**frame_obj, "response": {**kept, **guarded_fields}}
if isinstance(nested, dict)
else {**kept, **guarded_fields}
)
return GuardedResponsesWsFrame(
frame=json.dumps(rewritten),
guardrail_information=_recorded_guardrail_information(guarded_data.get(metadata_key)),
pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")),
)
return guard
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):
return value
return {
name: item
for name, item in _WS_FRAME.validate_python(value).items()
if not name.startswith("user_api_key_")
and name not in _UNTRUSTED_METADATA_CONTROL_FIELDS
and name not in _CLIENT_PRICING_METADATA_FIELDS
and name not in ("headers", "pii_tokens", *_CLIENT_GUARDRAIL_SELECTORS)
}
async def _bill_rejected_frame(
user_api_key_dict: UserAPIKeyAuth,
exception: Exception,
frame_data: dict[str, object], # mutable-ok: the spend logger stamps the failure row on the request dict
) -> None:
"""Only the spend logger sees a rejected frame; the other failure hooks (the rate limiter's slot release,
alerting) belong to the live connection, which stays open. Billing never replaces the guardrail's own error."""
for callback in litellm.callbacks:
if isinstance(callback, ProxyDBLogger):
try:
await callback.async_post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # bare dict
request_data=frame_data, original_exception=exception, user_api_key_dict=user_api_key_dict
)
except Exception: # noqa: BLE001 # billing never replaces the guardrail's own error
verbose_proxy_logger.exception("[Non-Blocking] Responses WebSocket: billing a rejected frame failed")
def _recorded_pii_tokens(metadata: object) -> Mapping[str, str]:
"""Presidio stores the tokens it masked a frame with in ``metadata.pii_tokens``; the socket unmasks with them."""
recorded: Final = _WS_FRAME.validate_python(metadata).get("pii_tokens") if isinstance(metadata, dict) else None
return MappingProxyType(_PII_TOKENS.validate_python(recorded)) if isinstance(recorded, dict) else _NO_PII_TOKENS
_NO_PII_TOKENS: Final[Mapping[str, str]] = MappingProxyType({})
def _recorded_guardrail_information(bucket: object) -> tuple[Mapping[str, object], ...]:
recorded: Final = (
_WS_FRAME.validate_python(bucket).get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None
)
return tuple(_GUARDRAIL_INFORMATION.validate_python(recorded)) if recorded else ()
async def _responses_websocket_session(
websocket: WebSocket,
model: str | None,
@ -1885,6 +2030,7 @@ async def _responses_websocket_session(
user_api_base=user_api_base,
model=resolved_model,
route_type="_aresponses_websocket",
skip_guardrails=True,
)
except Exception as e:
verbose_proxy_logger.exception("Responses WebSocket pre-call error")
@ -1906,7 +2052,10 @@ async def _responses_websocket_session(
return
routed_data: Final = dict(
data, user_api_key_dict=user_api_key_dict, **_routing_hints_from_first_ws_frame(first_message)
data,
user_api_key_dict=user_api_key_dict,
responses_ws_frame_guard=_responses_ws_frame_guard(data, user_api_key_dict, proxy_logging_obj),
**_routing_hints_from_first_ws_frame(first_message),
)
# Phase 2: route to upstream provider
try:

3
litellm/responses/main.py
View file

@ -62,6 +62,7 @@ from litellm.types.llms.openai import (
# Handle ResponseText import with fallback
if TYPE_CHECKING:
from litellm.responses.streaming_iterator import ResponsesWsFrameGuard
from litellm.types.llms.openai import ResponseText
else:
ResponseText = str # Fallback for ResponseText import
@ -2598,6 +2599,7 @@ async def _aresponses_websocket(
api_base: str | None = None,
api_key: str | None = None,
timeout: float | None = None,
responses_ws_frame_guard: "ResponsesWsFrameGuard | None" = None,
**kwargs,
) -> Exception | None:
"""
@ -2690,5 +2692,6 @@ async def _aresponses_websocket(
litellm_metadata=_build_litellm_metadata_for_ws(kwargs),
custom_llm_provider=_custom_llm_provider,
request_defaults=_build_responses_websocket_request_defaults(deployment_kwargs),
frame_guard=responses_ws_frame_guard,
**remaining_kwargs,
)

280
litellm/responses/streaming_iterator.py
View file

@ -7,12 +7,23 @@ import time
import traceback
import uuid
from collections.abc import AsyncIterable, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence
from dataclasses import dataclass
from dataclasses import dataclass, field
from datetime import datetime
from functools import lru_cache
from itertools import chain
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Protocol, cast, overload, runtime_checkable
from typing import (
TYPE_CHECKING,
Any,
Final,
Literal,
NoReturn,
Protocol,
TypeAlias,
cast,
overload,
runtime_checkable,
)
import httpx
from openai._streaming import SSEDecoder
@ -1932,6 +1943,73 @@ RESPONSES_WS_MASKABLE_TEXT_BLOCK_TYPES: Final = frozenset({"input_text", "output
_RESPONSES_WS_FAILURE_EVENT_TYPES: Final = frozenset({"error", "response.failed"})
_PER_TURN_METADATA_KEYS: Final = frozenset({"metadata", "litellm_metadata"})
_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
_PII_DELTA_EVENT_TYPES: Final = frozenset(
{
"response.output_text.delta",
"response.reasoning_summary_text.delta",
"response.refusal.delta",
"response.function_call_arguments.delta",
}
)
def _unmask_responses_event(response_str: str, pii_tokens: Mapping[str, str], unmask_pii_text: _UnmasksPiiText) -> str:
"""Put the original PII back into one client-bound event: ``response.completed`` output text, or a delta."""
try:
evt_obj: Final = _load_json_object(response_str)
except (json.JSONDecodeError, TypeError):
return response_str
event_type: Final = evt_obj.get("type")
if event_type == "response.completed":
modified = False
response_obj: Final = evt_obj.get("response")
if not _is_json_object(response_obj):
return response_str
output_items: Final = response_obj.get("output")
for output_item in output_items if _is_json_array(output_items) else []:
if not _is_json_object(output_item):
continue
content = output_item.get("content")
if not _is_json_array(content):
continue
for content_block in content:
if not _is_json_object(content_block):
continue
text = content_block.get("text")
if isinstance(text, str):
unmasked = unmask_pii_text(text, pii_tokens)
if unmasked != text:
content_block["text"] = unmasked
modified = True
return json.dumps(evt_obj) if modified else response_str
if event_type in _PII_DELTA_EVENT_TYPES:
delta: Final = evt_obj.get("delta")
if isinstance(delta, str):
unmasked = unmask_pii_text(delta, pii_tokens)
if unmasked != delta:
evt_obj["delta"] = unmasked
return json.dumps(evt_obj)
return response_str
@dataclass(frozen=True, slots=True)
class GuardedResponsesWsFrame:
"""A ``response.create`` frame after the proxy's pre_call guardrails ran on it: the frame to forward (guardrails
may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, and the
Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply."""
frame: str
guardrail_information: tuple[Mapping[str, object], ...] = ()
pii_tokens: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({}))
ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]]
_RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES
@ -1964,6 +2042,30 @@ def _user_texts_in_input_item(item: object) -> tuple[str, ...]:
)
def _set_turn_guardrail_information(
turn_logging_obj: LiteLLMLoggingObj, guardrail_information: tuple[Mapping[str, object], ...]
) -> None:
"""A turn bills the guardrails that ran on its own ``response.create`` frame, never the records the
connection-level logging object carries from an earlier frame."""
litellm_params: Final = _json_object_or_none(
turn_logging_obj.litellm_params # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped
)
if litellm_params is None:
return
for key in _PER_TURN_METADATA_KEYS:
_replace_guardrail_information(litellm_params.get(key), guardrail_information)
def _replace_guardrail_information(bucket: object, guardrail_information: tuple[Mapping[str, object], ...]) -> None:
if not _is_json_object(bucket):
return
bucket.pop(_GUARDRAIL_INFORMATION_KEY, None)
if guardrail_information:
bucket[_GUARDRAIL_INFORMATION_KEY] = [ # rebind-ok: the turn copy's own metadata bucket
dict(entry) for entry in guardrail_information
]
def _logged_user_messages(client_frame: object) -> tuple[Mapping[str, object], ...]:
frame: Final = _json_object_or_none(client_frame)
if frame is None or frame.get("type") != "response.create":
@ -2092,11 +2194,13 @@ class ResponsesWebSocketStreaming:
authorized_model: str | None = None,
custom_llm_provider: str | None = None,
request_defaults: ResponsesWebSocketRequestDefaults | None = None,
frame_guard: ResponsesWsFrameGuard | None = None,
):
self.websocket = websocket
self.backend_ws = backend_ws
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.frame_guard: ResponsesWsFrameGuard | None = frame_guard
self.request_data: dict[str, object] = request_data or {}
litellm_metadata: Final = self.request_data.get("litellm_metadata")
self.litellm_metadata: dict[str, object] = litellm_metadata if _is_json_object(litellm_metadata) else {}
@ -2112,12 +2216,15 @@ class ResponsesWebSocketStreaming:
self.authorized_model: str | None = authorized_model
self.request_defaults: ResponsesWebSocketRequestDefaults | None = request_defaults
def _start_turn(self, client_frame: str | None) -> None:
def _start_turn(
self, client_frame: str | None, guardrail_information: tuple[Mapping[str, object], ...] = ()
) -> None:
start_time: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time
turn_logging_obj: Final = self.logging_obj.copy_for_new_call(
litellm_call_id=str(uuid.uuid4()), start_time=start_time
)
self._detach_budget_reservation()
_set_turn_guardrail_information(turn_logging_obj, guardrail_information)
turn_logging_obj.model_call_details["messages"] = list(_logged_user_messages(client_frame))
turn_logging_obj.pre_call(input=client_frame, api_key="")
self._open_turns = (*self._open_turns, _ResponsesWsTurn(logging_obj=turn_logging_obj))
@ -2136,10 +2243,46 @@ class ResponsesWebSocketStreaming:
for key, value in litellm_params.items()
}
def _record_client_frame(self, client_frame: str) -> None:
def _record_client_frame(
self, client_frame: str, guardrail_information: tuple[Mapping[str, object], ...] = ()
) -> None:
frame: Final = _json_object_or_none(client_frame)
if frame is not None and frame.get("type") == "response.create":
self._start_turn(client_frame)
self._start_turn(client_frame, guardrail_information)
async def _guard_frame(self, message: str) -> GuardedResponsesWsFrame | None:
"""Run the proxy's pre_call guardrails on one ``response.create`` frame.
Returns None when a guardrail rejects the frame: the client gets an ``error`` event and the
frame is dropped while the connection stays open for the next turn.
"""
frame: Final = _json_object_or_none(message)
if self.frame_guard is None or frame is None or frame.get("type") != "response.create":
return GuardedResponsesWsFrame(frame=message)
try:
return await self.frame_guard(message)
except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event
verbose_logger.info("Responses WebSocket: pre_call guardrail rejected a response.create frame: %s", e)
await self._send_error_event("invalid_request_error", str(e))
return None
async def _forward_guarded_frame(self, message: str) -> None:
guarded: Final = await self._guard_frame(message)
if guarded is None:
return
if guarded.pii_tokens:
self._remember_pii_tokens(guarded.pii_tokens)
masked: Final = await self._mask_response_create(guarded.frame)
self._record_client_frame(masked, guarded.guardrail_information)
await self.backend_ws.send(masked)
def _remember_pii_tokens(self, pii_tokens: Mapping[str, str]) -> None:
metadata: Final = self.request_data.get("metadata")
known: Final = metadata.get("pii_tokens") if _is_json_object(metadata) else None
self.request_data["metadata"] = {
**(metadata if _is_json_object(metadata) else {}),
"pii_tokens": {**(known if _is_str_mapping(known) else {}), **pii_tokens},
}
def _record_backend_event(self, event: str) -> None:
event_obj: Final = _json_object_or_none(event)
@ -2383,14 +2526,7 @@ class ResponsesWebSocketStreaming:
return json.dumps(msg_obj) if modified else message
# Delta event types whose ``delta`` field may contain PII tokens.
_DELTA_EVENT_TYPES = frozenset(
{
"response.output_text.delta",
"response.reasoning_summary_text.delta",
"response.refusal.delta",
"response.function_call_arguments.delta",
}
)
_DELTA_EVENT_TYPES = _PII_DELTA_EVENT_TYPES
# Terminal events that carry the full output text or tool-call arguments
# already delivered by ``response.completed``. Suppressed when output masking
@ -2430,47 +2566,9 @@ class ResponsesWebSocketStreaming:
if not pii_tokens:
return response_str
try:
evt_obj: Final = _load_json_object(response_str)
except (json.JSONDecodeError, TypeError):
return response_str
cb: Final = self.guardrail_callbacks[0]
unmask_pii_text: Final[_UnmasksPiiText] = getattr(cb, _UNMASK_PII_TEXT_ATTR)
event_type: Final = evt_obj.get("type")
if event_type == "response.completed":
modified = False
response_obj: Final = evt_obj.get("response")
if not _is_json_object(response_obj):
return response_str
output_items: Final = response_obj.get("output")
for output_item in output_items if _is_json_array(output_items) else []:
if not _is_json_object(output_item):
continue
content = output_item.get("content")
if not _is_json_array(content):
continue
for content_block in content:
if not _is_json_object(content_block):
continue
text = content_block.get("text")
if isinstance(text, str):
unmasked = unmask_pii_text(text, pii_tokens)
if unmasked != text:
content_block["text"] = unmasked
modified = True
return json.dumps(evt_obj) if modified else response_str
if event_type in self._DELTA_EVENT_TYPES:
delta: Final = evt_obj.get("delta")
if isinstance(delta, str):
unmasked = unmask_pii_text(delta, pii_tokens)
if unmasked != delta:
evt_obj["delta"] = unmasked
return json.dumps(evt_obj)
return response_str
return _unmask_responses_event(response_str, pii_tokens, unmask_pii_text)
async def _mask_response_completed(self, response_str: str) -> str:
"""
@ -2564,38 +2662,29 @@ class ResponsesWebSocketStreaming:
self.quota_callbacks, self.user_api_key_dict, self.authorized_model, message
)
except RateLimitError as e:
try:
await self.websocket.send_text(
json.dumps(
{
"type": "error",
"error": {
"type": "rate_limit_exceeded",
"message": str(e),
},
}
)
)
except Exception: # noqa: BLE001, S110 # client may already be gone
pass
await self._send_error_event("rate_limit_exceeded", str(e))
return False
return True
async def _send_error_event(self, error_type: str, message: str) -> None:
try:
await self.websocket.send_text(
json.dumps({"type": "error", "error": {"type": error_type, "message": message}})
)
except Exception: # noqa: BLE001, S110 # client may already be gone
pass
async def client_to_backend(self) -> None:
"""Forward response.create events from client to backend."""
try:
if self.first_message is not None and await self._enforce_or_reject_frame(self.first_message):
masked_first: Final = await self._mask_response_create(self.first_message)
self._record_client_frame(masked_first)
await self.backend_ws.send(masked_first)
await self._forward_guarded_frame(self.first_message)
while True:
message = await self.websocket.receive_text()
if not await self._enforce_or_reject_frame(message):
continue
masked = await self._mask_response_create(message)
self._record_client_frame(masked)
await self.backend_ws.send(masked)
await self._forward_guarded_frame(message)
except Exception as e:
verbose_logger.debug("Responses WS client_to_backend ended: %s", e)
@ -2672,10 +2761,15 @@ class ManagedResponsesWebSocketHandler:
custom_llm_provider: str | None = None,
first_message: str | None = None,
quota_callbacks: Sequence[ProjectQuotaCallback] | None = None,
frame_guard: ResponsesWsFrameGuard | None = None,
guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None,
**kwargs: object,
) -> None:
self.websocket = websocket
self.model = model
self.frame_guard = frame_guard
self.guardrail_callbacks: tuple[PresidioGuardrailCallback, ...] = tuple(guardrail_callbacks or ())
self._pii_tokens: Mapping[str, str] = {}
self.logging_obj = logging_obj
self.user_api_key_dict = user_api_key_dict
self.litellm_metadata: Mapping[str, object] = litellm_metadata or {}
@ -2983,6 +3077,27 @@ class ManagedResponsesWebSocketHandler:
if self.litellm_metadata:
call_kwargs["litellm_metadata"] = dict(self.litellm_metadata)
async def _guard_frame(self, raw_message: str) -> GuardedResponsesWsFrame | None:
"""Same per-frame pre_call guardrail walk as the native socket: a rejected frame errors and is not sent."""
if self.frame_guard is None:
return GuardedResponsesWsFrame(frame=raw_message)
try:
return await self.frame_guard(raw_message)
except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event
verbose_logger.info("ManagedResponsesWS: pre_call guardrail rejected a response.create frame: %s", e)
await self._send_error(str(e), error_type="invalid_request_error")
return None
@staticmethod
def _metadata_with_turn_guardrails(
existing: object, guardrail_information: tuple[Mapping[str, object], ...]
) -> Mapping[str, object]:
"""Each turn is its own ``litellm.aresponses`` call, so its guardrail records ride on that call's metadata."""
return {
**(existing if _is_json_object(existing) else {}),
_GUARDRAIL_INFORMATION_KEY: [dict(record) for record in guardrail_information],
}
@staticmethod
def _update_proxy_request(call_kwargs: dict[str, Any], model: str) -> None:
"""Update proxy_server_request body so spend logs record the full request."""
@ -3026,12 +3141,19 @@ class ManagedResponsesWebSocketHandler:
except Exception:
pass
try:
await self.websocket.send_text(serialized)
await self.websocket.send_text(self._unmask_for_client(serialized))
except Exception as send_exc:
verbose_logger.debug("ManagedResponsesWS: error sending chunk to client: %s", send_exc)
return terminal_event # Client disconnected
return terminal_event
def _unmask_for_client(self, serialized: str) -> str:
"""The provider and the turn history keep Presidio's tokens, the client gets the original text back."""
if not self._pii_tokens or not self.guardrail_callbacks:
return serialized
unmask_pii_text: Final[_UnmasksPiiText] = getattr(self.guardrail_callbacks[0], _UNMASK_PII_TEXT_ATTR)
return _unmask_responses_event(serialized, self._pii_tokens, unmask_pii_text)
def _save_turn_history(
self,
terminal_event: _MutableJsonObject | None,
@ -3101,7 +3223,13 @@ class ManagedResponsesWebSocketHandler:
await self._send_error(str(e), error_type="rate_limit_exceeded")
return
call_kwargs: Final = self._build_base_call_kwargs(msg_obj)
guarded: Final = await self._guard_frame(raw_message)
if guarded is None:
return
if guarded.pii_tokens:
self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens}
call_kwargs: Final = self._build_base_call_kwargs(_load_json_object(guarded.frame))
call_kwargs["stream"] = True
# A frame that repeats the connection's public alias (model_group) must
@ -3121,6 +3249,10 @@ class ManagedResponsesWebSocketHandler:
self._apply_history(call_kwargs, previous_response_id, current_messages, prior_history)
self._inject_credentials(call_kwargs, model=model)
if guarded.guardrail_information:
call_kwargs["litellm_metadata"] = self._metadata_with_turn_guardrails(
call_kwargs.get("litellm_metadata"), guarded.guardrail_information
)
self._update_proxy_request(call_kwargs, requested_model or self.model_group or model)
call_kwargs.update(self.extra_kwargs)

168
tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py Normal file
View file

@ -0,0 +1,168 @@
"""A native Responses WebSocket runs the proxy's pre_call guardrails on every ``response.create`` frame, bills each
turn's spend row with the guardrail cost of that turn alone, and bills a blocked frame's guardrail scan as a failed
turn (LIT-9744)"""
import asyncio
import hashlib
import json
import uuid
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import pytest
import websockets
import yaml
from integration._support.client import Gateway, eventually, gateway_from_environment
from integration._support.database import read_rows
from integration._support.process import graceful_stop_seconds, owned_proxy
from integration._support.responses_ws_peer import PROVIDER_MODEL, USAGE, create, responses_peer, turn
from integration._support.tls import write_self_signed_cert
from pydantic import JsonValue
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
_INPUT_COST_PER_TOKEN: Final = 0.002
_OUTPUT_COST_PER_TOKEN: Final = 0.004
_COST_PER_WORD: Final = 0.1
_MARKER: Final = " [guarded]"
_BLOCKED_WORD: Final = "forbidden"
_GUARDRAIL_MODULE: Final = f"""
from litellm.integrations.custom_guardrail import CustomGuardrail
class TurnGuardrail(CustomGuardrail):
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
blocks = [block for item in data.get("input") or [] for block in item.get("content") or []]
words = sum(len(block["text"].split()) for block in blocks)
blocked = any({_BLOCKED_WORD!r} in block["text"] for block in blocks)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={{"words": words}},
request_data=data,
guardrail_status="guardrail_intervened" if blocked else "success",
tracing_detail={{"guardrail_cost": round({_COST_PER_WORD} * words, 10)}},
)
if blocked:
raise ValueError("turn guardrail blocked the request")
for block in blocks:
block["text"] = block["text"] + {_MARKER!r}
return data
"""
@dataclass(frozen=True, slots=True)
class _Rig:
gateway: Gateway
cert: tuple[Path, Path]
def _config(directory: Path) -> Path:
(directory / "turn_guardrail.py").write_text(_GUARDRAIL_MODULE)
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["guardrails"] = [
{
"guardrail_name": "turn-guardrail",
"litellm_params": {"guardrail": "turn_guardrail.TurnGuardrail", "mode": "pre_call", "default_on": True},
}
]
path: Final = directory / "responses-ws-guardrail.yaml"
path.write_text(yaml.safe_dump(configuration))
return path
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
directory: Final = tmp_path_factory.mktemp("responses-ws-guardrail")
cert: Final = write_self_signed_cert(directory)
with (
gateway_from_environment() as shared,
owned_proxy(shared, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_config(directory)) as owned,
):
yield _Rig(owned, cert)
def _nested(frame: str) -> str:
"""The ``{"type": "response.create", "response": {...}}`` wire shape the WS layer accepts next to the flat one."""
flat: Final = json.loads(frame)
return json.dumps({"type": flat["type"], "response": {key: value for key, value in flat.items() if key != "type"}})
async def _session(
proxy_url: str, key: str, model: str, frames: tuple[str, ...]
) -> tuple[tuple[dict[str, JsonValue], ...], ...]:
proxy: Final = proxy_url.rstrip("/").replace("http://", "ws://")
async with websockets.connect(
f"{proxy}/v1/responses?model={model}",
additional_headers={"Authorization": f"Bearer {key}"},
open_timeout=10,
) as connection:
return tuple([await turn(connection, frame) for frame in frames])
def _spend_rows(key: str) -> list[dict[str, JsonValue]]:
return read_rows(
'SELECT spend, status FROM "LiteLLM_SpendLogs" WHERE api_key = %s ORDER BY "startTime"',
(hashlib.sha256(key.encode()).hexdigest(),),
)
def _key_spend(key: str) -> float:
(row,) = read_rows(
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (hashlib.sha256(key.encode()).hexdigest(),)
)
return float(row["spend"])
def _llm_cost() -> float:
return USAGE["input_tokens"] * _INPUT_COST_PER_TOKEN + USAGE["output_tokens"] * _OUTPUT_COST_PER_TOKEN
def test_every_turn_runs_the_pre_call_guardrail_and_bills_only_its_own_guardrail_cost(rig: _Rig) -> None:
with responses_peer(rig.cert) as peer, rig.gateway.scenario() as scenario:
model: Final = scenario.model(
model=f"openai/{PROVIDER_MODEL}",
api_base=peer.url,
input_cost_per_token=_INPUT_COST_PER_TOKEN,
output_cost_per_token=_OUTPUT_COST_PER_TOKEN,
)
key: Final = scenario.key(models=[model])
tag: Final = uuid.uuid4().hex
texts: Final = (
f"one {tag}",
f"one two {tag}",
f"one two three {tag}",
f"{_BLOCKED_WORD} {tag}",
f"one {tag}",
f"{_BLOCKED_WORD} again {tag}",
)
frames: Final = tuple(create(model, text) for text in texts[:-1]) + (_nested(create(model, texts[-1])),)
turns: Final = asyncio.run(_session(str(rig.gateway.client.base_url), key, model, frames))
assert [events[-1]["type"] for events in turns] == [
"response.completed",
"response.completed",
"response.completed",
"error",
"response.completed",
"error",
], [[event["type"] for event in events] for events in turns]
assert "blocked" in str(turns[3][-1]["error"]), turns[3]
forwarded: Final = tuple(peer.frames.get_nowait() for _ in range(peer.frames.qsize()))
assert [frame["input"][0]["content"][0]["text"] for frame in forwarded] == [
f"{text}{_MARKER}" for text in texts if _BLOCKED_WORD not in text
], forwarded
rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 6, seconds=70)
assert [row["status"] for row in rows] == ["success", "success", "success", "failure", "success", "failure"], (
rows
)
expected: Final = [
(0.0 if _BLOCKED_WORD in text else _llm_cost()) + _COST_PER_WORD * len(text.split()) for text in texts
]
assert [float(row["spend"]) for row in rows] == pytest.approx(expected), rows
total: Final = sum(expected)
assert eventually(
lambda: _key_spend(key), lambda spend: abs(spend - total) < 1e-6, seconds=70
) == pytest.approx(total)

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

@ -480,7 +480,12 @@ def _advertised_cost_split(headers: Mapping[str, str]) -> list[tuple[float, floa
@pytest.mark.asyncio
@pytest.mark.parametrize(
"background, stored, status, expected_billed_events",
[(True, True, "completed", 1), (True, True, "in_progress", 0), (True, False, "completed", 0), (False, True, "completed", 0)],
[
(True, True, "completed", 1),
(True, True, "in_progress", 0),
(True, False, "completed", 0),
(False, True, "completed", 0),
],
)
async def test_reading_a_finished_background_response_bills_it(
monkeypatch: pytest.MonkeyPatch, background: bool, stored: bool, status: str, expected_billed_events: int
@ -531,8 +536,12 @@ async def test_reading_a_finished_background_response_bills_it(
prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.prisma_client", prisma_client
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.llm_router", MagicMock()
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")),
):
await _bill_finished_background_response(
@ -662,9 +671,15 @@ async def _read_twice_on_two_pods(
router.aget_responses = AsyncMock(side_effect=AssertionError("billing must not read the provider again"))
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.prisma_client", prisma_client
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.llm_router", router
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.aget_responses", side_effect=AssertionError("billing must not read the provider again")),
):
await asyncio.gather(
@ -702,7 +717,9 @@ async def test_a_finished_response_the_create_could_not_store_is_billed_once_und
monkeypatch,
table,
response_id,
UserAPIKeyAuth(api_key="sk-reader", user_id="bob", team_id="team-beta", org_id="org-gamma", end_user_id="customer-42"),
UserAPIKeyAuth(
api_key="sk-reader", user_id="bob", team_id="team-beta", org_id="org-gamma", end_user_id="customer-42"
),
)
(billed,) = payloads
@ -746,7 +763,9 @@ async def test_a_finished_read_never_bills_a_response_already_billed_or_billed_b
)
table: Final = _UniqueManagedObjectTable(stored)
payloads: Final = await _read_twice_on_two_pods(monkeypatch, table, response_id, UserAPIKeyAuth(api_key="sk-reader"))
payloads: Final = await _read_twice_on_two_pods(
monkeypatch, table, response_id, UserAPIKeyAuth(api_key="sk-reader")
)
assert payloads == ()
assert table.rows == stored
@ -878,9 +897,15 @@ async def test_a_finished_read_or_cancel_bills_usage_only_the_configured_deploym
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.prisma_client", stored.prisma_client
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.llm_router", router
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.aget_responses", side_effect=AssertionError("the configured deployment answers billing reads")),
):
await getattr(endpoints, endpoint_name)(
@ -941,9 +966,15 @@ async def test_deleting_a_finished_unbilled_background_response_bills_it_before_
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.prisma_client", stored.prisma_client
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.llm_router", router
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging
), # test-quality-ok: the endpoint reads the proxy singletons at call time
):
result: Final = await delete_response(
response_id=stored.response_id,
@ -989,7 +1020,9 @@ async def test_a_delete_the_proxy_refuses_never_fetches_or_bills_the_response(ca
caller: Final = UserAPIKeyAuth(api_key="sk-caller", user_role=caller_role)
async def pipeline(**_: object) -> DeleteResponseResult:
await security.async_pre_call_hook(caller, MagicMock(), {"response_id": stored.response_id}, "adelete_responses")
await security.async_pre_call_hook(
caller, MagicMock(), {"response_id": stored.response_id}, "adelete_responses"
)
return DeleteResponseResult(id=stored.response_id, object="response", deleted=True)
processor: Final = MagicMock()
@ -1002,9 +1035,15 @@ async def test_a_delete_the_proxy_refuses_never_fetches_or_bills_the_response(ca
"litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing",
return_value=processor,
),
patch("litellm.proxy.proxy_server.prisma_client", stored.prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.prisma_client", stored.prisma_client
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.llm_router", router
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.proxy_logging_obj", stored.proxy_logging
), # test-quality-ok: the endpoint reads the proxy singletons at call time
):
with contextlib.suppress(HTTPException):
await delete_response(
@ -1058,9 +1097,15 @@ def _capture_spend_events(monkeypatch: pytest.MonkeyPatch) -> _SpendEvents:
@contextlib.contextmanager
def _proxy_over(prisma_client: MagicMock, proxy_logging: MagicMock, router: MagicMock):
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.prisma_client", prisma_client
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.llm_router", router
), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging
), # test-quality-ok: the endpoint reads the proxy singletons at call time
):
yield
@ -1211,7 +1256,9 @@ async def test_a_streamed_background_create_is_billed_wherever_it_sets_backgroun
"litellm.proxy.response_polling.polling_handler.should_use_polling_for_request",
return_value=False,
),
patch("litellm.proxy.proxy_server.select_data_generator", provider_events), # test-quality-ok: the endpoint reads the proxy singletons at call time
patch(
"litellm.proxy.proxy_server.select_data_generator", provider_events
), # test-quality-ok: the endpoint reads the proxy singletons at call time
):
await responses_api(
request=_request_for("POST", "/v1/responses", body),
@ -1329,7 +1376,12 @@ async def test_background_create_the_provider_confirms_is_stored_under_its_provi
}
request = Request(
{"type": "http", "method": "POST", "path": "/v1/responses", "headers": [(b"content-type", b"application/json")]},
{
"type": "http",
"method": "POST",
"path": "/v1/responses",
"headers": [(b"content-type", b"application/json")],
},
receive,
)
@ -1918,9 +1970,327 @@ class TestResponsesWSFirstFrameModelAuth:
assert routed["input"] == replayed_input
assert routed["previous_response_id"] == "resp_prev"
assert processor.common_processing_pre_call_logic.await_args.kwargs["model"] == "gpt-4o-mini"
assert processor.common_processing_pre_call_logic.await_args.kwargs["skip_guardrails"] is True
assert callable(routed["responses_ws_frame_guard"])
assert mock_route_request.await_args.kwargs["route_type"] == "_aresponses_websocket"
ws.close.assert_not_awaited()
@pytest.mark.asyncio
async def test_frame_guard_runs_the_guardrail_walk_on_the_frame_body_and_returns_its_records(self):
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
assert call_type == "aresponses" and guardrails_only is True
assert "type" not in data and data["model"] == "gpt-4o-mini"
assert data["litellm_metadata"]["user_api_key_hash"] == "hash"
assert "standard_logging_guardrail_information" not in data["litellm_metadata"]
data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_name": "turn"}]
return {**data, "input": data["input"] + " [guarded]"}
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
admitted = {
"model": "gpt-4o-mini",
"litellm_metadata": {
"user_api_key_hash": "hash",
"standard_logging_guardrail_information": [{"guardrail_name": "admission"}],
},
}
guard = _responses_ws_frame_guard(admitted, MagicMock(), proxy_logging_obj)
first = await guard(json.dumps({"type": "response.create", "model": "alias", "input": "hi"}))
second = await guard(json.dumps({"type": "response.create", "input": "again"}))
assert json.loads(first.frame) == {"type": "response.create", "model": "alias", "input": "hi [guarded]"}
assert first.guardrail_information == ({"guardrail_name": "turn"},)
assert json.loads(second.frame)["input"] == "again [guarded]"
assert admitted["litellm_metadata"]["standard_logging_guardrail_information"] == [
{"guardrail_name": "admission"}
]
@pytest.mark.asyncio
async def test_frame_guard_bills_a_rejected_frame_as_a_failed_request(self):
import litellm
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}]
raise ValueError("blocked")
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
proxy_logging_obj.post_call_failure_hook = AsyncMock()
spend_logger = MagicMock(spec=ProxyDBLogger)
spend_logger.async_post_call_failure_hook = AsyncMock()
rate_limiter = MagicMock()
rate_limiter.async_post_call_failure_hook = AsyncMock()
user_api_key_dict = MagicMock()
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {}}, user_api_key_dict, proxy_logging_obj
)
with (
patch.object(litellm, "callbacks", [rate_limiter, spend_logger]),
pytest.raises(ValueError, match="blocked"),
):
await guard(json.dumps({"type": "response.create", "input": "hi"}))
proxy_logging_obj.post_call_failure_hook.assert_not_awaited()
rate_limiter.async_post_call_failure_hook.assert_not_awaited()
spend_logger.async_post_call_failure_hook.assert_awaited_once()
failure = spend_logger.async_post_call_failure_hook.await_args.kwargs
user_api_key_dict.model_copy.assert_called_once_with(update={"budget_reservation": None})
assert failure["user_api_key_dict"] is user_api_key_dict.model_copy.return_value
assert str(failure["original_exception"]) == "blocked"
assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [
{"guardrail_cost": 0.5}
]
assert failure["request_data"]["input"] == "hi"
assert len(failure["request_data"]["litellm_call_id"]) == 36
@pytest.mark.asyncio
async def test_frame_guard_guards_the_nested_response_body_and_copies_back_tools(self):
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
return {**data, "input": data["input"] + " [guarded]", "tools": data["tools"][:1]}
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj
)
tools = [{"type": "function", "name": "a"}, {"type": "function", "name": "b"}]
guarded = await guard(json.dumps({"type": "response.create", "response": {"input": "hi", "tools": tools}}))
assert proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["input"] == "hi"
assert json.loads(guarded.frame) == {
"type": "response.create",
"response": {"input": "hi [guarded]", "tools": tools[:1]},
}
@pytest.mark.asyncio
async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self):
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": {}}, MagicMock(), proxy_logging_obj
)
metadata = {
"user_api_key_metadata": {"disable_global_guardrails": True},
"disable_global_guardrails": True,
"headers": {"x-litellm-tags": "forged"},
"tag": "kept",
}
await guard(json.dumps({"type": "response.create", "input": "hi", "metadata": metadata}))
assert proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["metadata"] == {"tag": "kept"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"frame",
[
{"type": "response.create", "input": "forbidden", "guardrails": []},
{"type": "response.create", "input": "forbidden", "metadata": {"guardrails": []}},
{"type": "response.create", "response": {"input": "forbidden", "guardrails": [], "policies": []}},
],
)
async def test_frame_guard_keeps_the_admitted_guardrail_list_when_a_frame_tries_to_replace_it(self, frame):
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,
)
await guard(json.dumps(frame))
scanned = proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]
assert "guardrails" not in scanned and "policies" not in scanned
assert "guardrails" not in (scanned.get("metadata") or {})
assert scanned["litellm_metadata"]["guardrails"] == ["key-required"]
@pytest.mark.asyncio
async def test_frame_guard_copies_an_admitted_bucket_that_holds_the_otel_parent_span(self):
import threading
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"])
span = threading.RLock()
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {"litellm_parent_otel_span": span}},
MagicMock(),
proxy_logging_obj,
)
guarded = await guard(json.dumps({"type": "response.create", "input": "hi"}))
assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"}
assert (
proxy_logging_obj.pre_call_hook.await_args.kwargs["data"]["litellm_metadata"]["litellm_parent_otel_span"]
is span
)
@pytest.mark.asyncio
async def test_frame_guard_rejects_a_turn_a_guardrail_rerouted_to_another_model(self):
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}]
return {**data, "model": "onprem-llama"}
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
spend_logger = MagicMock(spec=ProxyDBLogger)
spend_logger.async_post_call_failure_hook = AsyncMock()
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj
)
with (
patch.object(litellm, "callbacks", [spend_logger]),
pytest.raises(ValueError, match="rerouted this turn to 'onprem-llama'"),
):
await guard(json.dumps({"type": "response.create", "input": "my SSN is 123-45-6789"}))
failure = spend_logger.async_post_call_failure_hook.await_args.kwargs
assert "rerouted this turn to 'onprem-llama'" in str(failure["original_exception"])
assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [
{"guardrail_cost": 0.5}
]
@pytest.mark.asyncio
@pytest.mark.parametrize("nested", [False, True])
async def test_frame_guard_drops_a_client_s_guardrail_cost_records_so_a_blocked_frame_bills_the_real_scan(
self, nested: bool
):
import litellm
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}]
raise ValueError("blocked")
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
spend_logger = MagicMock(spec=ProxyDBLogger)
spend_logger.async_post_call_failure_hook = AsyncMock()
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj
)
body = {
"input": "forbidden",
"metadata": {"trace": "t1", "standard_logging_guardrail_information": [], "model_info": {"x": 1}},
}
frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body}
with patch.object(litellm, "callbacks", [spend_logger]), pytest.raises(ValueError, match="blocked"):
await guard(json.dumps(frame))
billed = spend_logger.async_post_call_failure_hook.await_args.kwargs["request_data"]
assert billed["metadata"] == {"trace": "t1"}
assert billed["litellm_metadata"]["standard_logging_guardrail_information"] == [{"guardrail_cost": 0.5}]
@pytest.mark.asyncio
@pytest.mark.parametrize("nested", [False, True])
async def test_frame_guard_forwards_a_frame_without_the_instructions_a_guardrail_deleted(self, nested: bool):
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
return {key: value for key, value in data.items() if key != "instructions"}
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj
)
body = {"instructions": "SECRET system prompt", "input": "hi", "store": False}
frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body}
with patch.object(litellm, "callbacks", []):
guarded = await guard(json.dumps(frame))
forwarded = json.loads(guarded.frame)
assert forwarded["type"] == "response.create"
sent_body = forwarded["response"] if nested else {k: v for k, v in forwarded.items() if k != "type"}
assert sent_body == {"input": "hi", "store": False}
@pytest.mark.asyncio
async def test_frame_guard_hands_back_the_pii_tokens_presidio_masked_with_and_drops_forged_ones(self):
from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard
async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only):
assert data["metadata"] == {"trace": "t1"}, "a client cannot seed pii_tokens"
data["metadata"]["pii_tokens"] = {"<EMAIL_ADDRESS_1>": "alice@example.com"}
return {**data, "input": "email <EMAIL_ADDRESS_1> please"}
proxy_logging_obj = MagicMock()
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook)
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj
)
with patch.object(litellm, "callbacks", []):
guarded = await guard(
json.dumps(
{
"type": "response.create",
"input": "email alice@example.com please",
"metadata": {"trace": "t1", "pii_tokens": {"<EMAIL_ADDRESS_1>": "mallory@example.com"}},
}
)
)
assert json.loads(guarded.frame)["input"] == "email <EMAIL_ADDRESS_1> please"
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):
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": {}}, MagicMock(), proxy_logging_obj
)
with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"):
await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}}))
proxy_logging_obj.pre_call_hook.assert_not_awaited()
@pytest.mark.asyncio
async def test_frame_guard_keeps_the_guardrail_error_when_billing_the_rejected_frame_fails(self):
import litellm
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
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=ValueError("blocked"))
spend_logger = MagicMock(spec=ProxyDBLogger)
spend_logger.async_post_call_failure_hook = AsyncMock(side_effect=RuntimeError("db down"))
guard = _responses_ws_frame_guard(
{"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj
)
with patch.object(litellm, "callbacks", [spend_logger]), pytest.raises(ValueError, match="blocked"):
await guard(json.dumps({"type": "response.create", "input": "hi"}))
spend_logger.async_post_call_failure_hook.assert_awaited_once()
@pytest.mark.asyncio
@pytest.mark.parametrize("provider_rejected", [True, False])
async def test_endpoint_books_a_provider_rejected_connection_as_a_failed_request(self, provider_rejected: bool):
@ -3614,16 +3984,18 @@ class TestResponsesInputTokens:
model_list=[
{
"model_name": "gpt-5.2",
"litellm_params": {"model": "openai/gpt-5.2", "api_key": "sk-test", "additional_drop_params": ["tools"]},
"litellm_params": {
"model": "openai/gpt-5.2",
"api_key": "sk-test",
"additional_drop_params": ["tools"],
},
},
{"model_name": "backup", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-test"}},
],
fallbacks=[{"gpt-5.2": ["backup"]}],
)
response, counter = self._post_input_tokens(
{"model": "gpt-5.2", "input": "hi", "tools": 5}, llm_router=router
)
response, counter = self._post_input_tokens({"model": "gpt-5.2", "input": "hi", "tools": 5}, llm_router=router)
assert response.status_code == 200, response.text
assert counter.await_args.kwargs["request"].tools is None

264
tests/unit/responses/test_responses_websocket_all_providers.py
View file

@ -285,6 +285,151 @@ class TestManagedWebSocketHandlerIntegration:
assert captured["model"] == "bedrock_mantle/openai.gpt-5.5"
@pytest.mark.asyncio
async def test_managed_socket_runs_the_frame_guard_on_every_turn(self, monkeypatch):
import json
from unittest.mock import AsyncMock, MagicMock
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ManagedResponsesWebSocketHandler
async def guard(frame: str) -> GuardedResponsesWsFrame:
body = json.loads(frame)
if "forbidden" in body["input"]:
raise ValueError("blocked by guardrail")
return GuardedResponsesWsFrame(
frame=json.dumps({**body, "input": body["input"].upper()}),
guardrail_information=({"guardrail_name": "g", "guardrail_cost": 0.5},),
)
async def empty_stream():
return
yield {}
aresponses: Final = AsyncMock(side_effect=lambda **kwargs: empty_stream())
monkeypatch.setattr(litellm, "aresponses", aresponses)
websocket: Final = MagicMock()
websocket.send_text = AsyncMock()
handler: Final = ManagedResponsesWebSocketHandler(
websocket=websocket,
model="anthropic/claude-sonnet-4-5",
logging_obj=Logging(
model="anthropic/claude-sonnet-4-5",
messages=[],
stream=True,
call_type="aresponses",
start_time=0,
litellm_call_id="test-id",
function_id="test-func",
),
litellm_metadata={"user_api_key_hash": "k1"},
frame_guard=guard,
)
await handler._process_response_create(json.dumps({"type": "response.create", "input": "say forbidden"}))
await handler._process_response_create(json.dumps({"type": "response.create", "input": "hello"}))
aresponses.assert_awaited_once()
error_event: Final = json.loads(websocket.send_text.await_args_list[0].args[0])
assert error_event["error"] == {"type": "invalid_request_error", "message": "blocked by guardrail"}
sent: Final = aresponses.await_args.kwargs
assert sent["input"] == "HELLO"
assert sent["litellm_metadata"]["user_api_key_hash"] == "k1"
assert sent["litellm_metadata"]["standard_logging_guardrail_information"] == [
{"guardrail_name": "g", "guardrail_cost": 0.5}
]
@pytest.mark.asyncio
async def test_managed_socket_unmasks_the_reply_with_the_frame_guard_s_pii_tokens(self, monkeypatch):
import json
from unittest.mock import AsyncMock, MagicMock
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ManagedResponsesWebSocketHandler
presidio = _FakeWSGuardrail()
async def guard(frame: str) -> GuardedResponsesWsFrame:
body = json.loads(frame)
masked = await presidio.check_pii(body["input"], True, None, walk := {})
return GuardedResponsesWsFrame(
frame=json.dumps({**body, "input": masked}), pii_tokens=walk["metadata"]["pii_tokens"]
)
provider_saw: dict = {}
async def fake_aresponses(**kwargs):
provider_saw.update(kwargs)
async def stream():
yield {"type": "response.output_text.delta", "delta": "Sent to <EMAIL_ADDRESS_1>"}
yield {
"type": "response.completed",
"response": {
"id": "resp_1",
"output": [
{"type": "message", "content": [{"type": "output_text", "text": "Sent to <EMAIL_ADDRESS_1>"}]}
],
},
}
return stream()
monkeypatch.setattr(litellm, "aresponses", fake_aresponses)
websocket: Final = MagicMock()
websocket.send_text = AsyncMock()
handler: Final = ManagedResponsesWebSocketHandler(
websocket=websocket,
model="anthropic/claude-sonnet-4-5",
logging_obj=Logging(
model="anthropic/claude-sonnet-4-5",
messages=[],
stream=True,
call_type="aresponses",
start_time=0,
litellm_call_id="test-id",
function_id="test-func",
),
frame_guard=guard,
guardrail_callbacks=[presidio],
)
await handler._process_response_create(
json.dumps({"type": "response.create", "input": "email alice@example.com please"})
)
assert provider_saw["input"] == "email <EMAIL_ADDRESS_1> please"
sent: Final = [json.loads(call.args[0]) for call in websocket.send_text.await_args_list]
assert sent[0]["delta"] == "Sent to alice@example.com"
assert sent[1]["response"]["output"][0]["content"][0]["text"] == "Sent to alice@example.com"
@pytest.mark.asyncio
async def test_http_handler_hands_the_frame_guard_to_a_managed_socket(self):
import json
from unittest.mock import AsyncMock, MagicMock
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
guard: Final = AsyncMock(side_effect=ValueError("blocked by guardrail"))
websocket: Final = MagicMock()
websocket.send_text = AsyncMock()
websocket.receive_text = AsyncMock(side_effect=ConnectionError("client left"))
await BaseLLMHTTPHandler().async_responses_websocket(
model="anthropic/claude-sonnet-4-5",
websocket=websocket,
logging_obj=MagicMock(),
responses_api_provider_config=None,
first_message=json.dumps({"type": "response.create", "input": "say forbidden"}),
frame_guard=guard,
)
guard.assert_awaited_once()
error_event: Final = json.loads(websocket.send_text.await_args.args[0])
assert error_event["error"]["message"] == "blocked by guardrail"
@pytest.mark.parametrize("terminal_type", ["response.completed", "response.incomplete"])
@pytest.mark.asyncio
async def test_truncated_turn_is_kept_as_history_for_the_next_turn(self, monkeypatch, terminal_type: str):
@ -1319,6 +1464,97 @@ class TestWebSocketProjectQuotaEnforcement:
mock_websocket.send_text.assert_called_once()
assert "rate_limit_exceeded" in mock_websocket.send_text.call_args[0][0]
@pytest.mark.asyncio
async def test_native_handler_runs_the_frame_guard_on_every_turn_and_bills_each_turn_its_own_guardrail_records(
self,
):
from unittest.mock import AsyncMock, MagicMock
from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWebSocketStreaming
def connection_logging():
logging_obj = MagicMock()
logging_obj.litellm_params = {
"metadata": {
"standard_logging_guardrail_information": [{"guardrail_name": "admission", "guardrail_cost": 9.0}]
}
}
logging_obj.model_call_details = {}
logging_obj.copy_for_new_call.side_effect = lambda **_: connection_logging()
return logging_obj
async def guard(frame: str) -> GuardedResponsesWsFrame:
body = json.loads(frame)
return GuardedResponsesWsFrame(
frame=json.dumps({**body, "input": body["input"] + " [guarded]"}),
guardrail_information=({"guardrail_name": "turn", "guardrail_cost": 0.1 * len(body["input"])},),
)
frames = [json.dumps({"type": "response.create", "input": text}) for text in ("hi", "hello")]
mock_backend_ws = MagicMock()
mock_backend_ws.send = AsyncMock()
mock_websocket = MagicMock()
mock_websocket.send_text = AsyncMock()
mock_websocket.receive_text = AsyncMock(side_effect=[frames[1], RuntimeError("client closed")])
handler = ResponsesWebSocketStreaming(
websocket=mock_websocket,
backend_ws=mock_backend_ws,
logging_obj=connection_logging(),
first_message=frames[0],
frame_guard=guard,
)
await handler.client_to_backend()
sent = [json.loads(call.args[0]) for call in mock_backend_ws.send.call_args_list]
assert [frame["input"] for frame in sent] == ["hi [guarded]", "hello [guarded]"]
assert [
turn.logging_obj.litellm_params["metadata"]["standard_logging_guardrail_information"]
for turn in handler._open_turns
] == [
[{"guardrail_name": "turn", "guardrail_cost": pytest.approx(0.2)}],
[{"guardrail_name": "turn", "guardrail_cost": pytest.approx(0.5)}],
]
@pytest.mark.asyncio
async def test_native_handler_drops_a_frame_the_guard_rejects_and_keeps_the_socket_open(self):
from unittest.mock import AsyncMock, MagicMock
from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWebSocketStreaming
async def guard(frame: str) -> GuardedResponsesWsFrame:
if "forbidden" in frame:
raise ValueError("turn guardrail blocked the request")
return GuardedResponsesWsFrame(frame=frame)
blocked = json.dumps({"type": "response.create", "input": "forbidden"})
allowed = json.dumps({"type": "response.create", "input": "fine"})
mock_backend_ws = MagicMock()
mock_backend_ws.send = AsyncMock()
mock_websocket = MagicMock()
mock_websocket.send_text = AsyncMock()
mock_websocket.receive_text = AsyncMock(side_effect=[allowed, RuntimeError("client closed")])
logging_obj = MagicMock()
logging_obj.copy_for_new_call.return_value = logging_obj
handler = ResponsesWebSocketStreaming(
websocket=mock_websocket,
backend_ws=mock_backend_ws,
logging_obj=logging_obj,
first_message=blocked,
frame_guard=guard,
)
await handler.client_to_backend()
error_event = json.loads(mock_websocket.send_text.call_args[0][0])
assert error_event["type"] == "error"
assert error_event["error"] == {
"type": "invalid_request_error",
"message": "turn guardrail blocked the request",
}
assert [json.loads(call.args[0])["input"] for call in mock_backend_ws.send.call_args_list] == ["fine"]
assert len(handler._open_turns) == 1
@pytest.mark.asyncio
async def test_native_handler_forwards_frame_allowed_by_quota_callback(self):
from unittest.mock import AsyncMock, MagicMock
@ -1905,6 +2141,34 @@ def _make_streaming(**kwargs):
class TestNativeWebSocketGuardrailMasking:
"""Exercises the input/output PII masking hooks on ResponsesWebSocketStreaming."""
@pytest.mark.asyncio
async def test_a_reply_to_a_frame_the_guard_masked_is_unmasked_with_the_guard_s_tokens(self):
from unittest.mock import AsyncMock, MagicMock
from litellm.responses.streaming_iterator import GuardedResponsesWsFrame
guardrail = _FakeWSGuardrail()
async def guard(frame: str) -> GuardedResponsesWsFrame:
masked = await guardrail.check_pii(json.loads(frame)["input"], True, None, walk := {})
return GuardedResponsesWsFrame(
frame=json.dumps({"type": "response.create", "input": masked}),
pii_tokens=walk["metadata"]["pii_tokens"],
)
backend_ws = MagicMock()
backend_ws.send = AsyncMock()
handler = _make_streaming(
backend_ws=backend_ws, request_data={}, guardrail_callbacks=[guardrail], frame_guard=guard
)
handler.logging_obj.copy_for_new_call.return_value = handler.logging_obj
await handler._forward_guarded_frame(json.dumps({"type": "response.create", "input": "email alice@example.com"}))
assert json.loads(backend_ws.send.await_args.args[0])["input"] == "email <EMAIL_ADDRESS_1>"
reply = json.dumps({"type": "response.output_text.delta", "delta": "Sent to <EMAIL_ADDRESS_1>"})
assert json.loads(handler._unmask_response_event(reply))["delta"] == "Sent to alice@example.com"
@pytest.mark.asyncio
async def test_mask_response_create_flat_string_input(self):
guardrail = _FakeWSGuardrail()