diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4ff4c43428f..1bccef7dd9e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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() diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 207e966ef28..aebcab75601 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -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: diff --git a/litellm/responses/main.py b/litellm/responses/main.py index b7385c3b3d5..d32d48c51ca 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -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, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 5897b0ac954..145189616af 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -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) diff --git a/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py b/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py new file mode 100644 index 00000000000..8af70d39186 --- /dev/null +++ b/tests/integration/cost_calculation/test_responses_websocket_guardrail_spend.py @@ -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) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index a7b955f5b84..52e95901a16 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -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"] = {"": "alice@example.com"} + return {**data, "input": "email 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": {"": "mallory@example.com"}}, + } + ) + ) + + assert json.loads(guarded.frame)["input"] == "email please" + assert dict(guarded.pii_tokens) == {"": "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 diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 2b78aa98b04..b38808fb57c 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -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 "} + yield { + "type": "response.completed", + "response": { + "id": "resp_1", + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "Sent to "}]} + ], + }, + } + + 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 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 " + reply = json.dumps({"type": "response.output_text.delta", "delta": "Sent to "}) + 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()