diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index c042264c034..646a52384d9 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -119,6 +119,7 @@ from litellm.responses.streaming_iterator import ( ProjectQuotaCallback, ResponsesAPIStreamingIterator, ResponsesWebSocketStreaming, + ResponsesWsFrameCheck, SyncResponsesAPIStreamingIterator, ) from litellm.types.containers.main import ( @@ -6851,6 +6852,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str | None = None, first_message: str | None = None, request_defaults: ResponsesWebSocketRequestDefaults | None = None, + frame_check: ResponsesWsFrameCheck | None = None, **kwargs: object, ) -> Exception | None: """ @@ -6883,6 +6885,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, first_message=first_message, quota_callbacks=_ws_quota_callbacks, + frame_check=frame_check, **kwargs, ) await handler.run() diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 63062569af4..9b6838dad58 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -51,6 +51,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, cost_headers, create_response, + refuse_s3_media_the_key_may_not_read, resolve_counted_tool_fields, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id, with_call_id @@ -71,6 +72,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing +from litellm.responses.streaming_iterator import ResponsesWsFrameCheck from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -1904,6 +1906,25 @@ async def _enforce_responses_ws_first_frame_model_auth( ) +def _s3_media_frame_check( + admitted_model: object, + user_api_key_dict: UserAPIKeyAuth, + general_settings: Mapping[str, object], + managed_files: object, + llm_router: Router | None, +) -> ResponsesWsFrameCheck: + async def refusal(params: Mapping[str, object], /) -> str | None: + try: + await refuse_s3_media_the_key_may_not_read( + {**params, "model": admitted_model}, user_api_key_dict, general_settings, managed_files, llm_router + ) + except HTTPException as refused: + return str(refused.detail) + return None + + return refusal + + async def _responses_websocket_session( websocket: WebSocket, model: str | None, @@ -2000,7 +2021,16 @@ 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_check=_s3_media_frame_check( + data.get("model"), + user_api_key_dict, + general_settings, + proxy_logging_obj.get_proxy_hook("managed_files"), + llm_router, + ), + **_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 a858b91c4c8..0d04558322b 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -77,7 +77,7 @@ if TYPE_CHECKING: else: MCPTool = Any -from .streaming_iterator import BaseResponsesAPIStreamingIterator +from .streaming_iterator import BaseResponsesAPIStreamingIterator, ResponsesWsFrameCheck __all__ = ( "acancel_responses", @@ -2488,6 +2488,7 @@ async def _aresponses_websocket( api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, + responses_ws_frame_check: ResponsesWsFrameCheck | None = None, **kwargs, ) -> Exception | None: """ @@ -2580,5 +2581,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_check=responses_ws_frame_check, **remaining_kwargs, ) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index cb745384403..61abeb1d74b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -66,6 +66,10 @@ if TYPE_CHECKING: from litellm.types.router import LiteLLM_Params +class ResponsesWsFrameCheck(Protocol): + async def __call__(self, params: Mapping[str, object], /) -> str | None: ... + + class ProjectQuotaCallback(Protocol): async def enforce_project_io_token_quota_for_frame( self, @@ -2723,6 +2727,7 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, + frame_check: ResponsesWsFrameCheck | None = None, **kwargs: object, ) -> None: self.websocket = websocket @@ -2741,6 +2746,7 @@ class ManagedResponsesWebSocketHandler: self._connection_provider = self._resolve_provider(model) or custom_llm_provider self.first_message = first_message self.quota_callbacks: tuple[ProjectQuotaCallback, ...] = tuple(quota_callbacks) if quota_callbacks else () + self.frame_check = frame_check # Carry through safe pass-through kwargs (e.g. extra_headers) self.extra_kwargs: dict[str, object] = {k: v for k, v in kwargs.items() if k not in _MANAGED_WS_SKIP_KWARGS} # In-memory session history: response_id → full accumulated message list. @@ -3144,6 +3150,12 @@ class ManagedResponsesWebSocketHandler: verbose_logger.debug("ManagedResponsesWS: error sending warmup ack: %s", exc) return + call_kwargs: Final = self._build_base_call_kwargs(msg_obj) + refusal: Final = None if self.frame_check is None else await self.frame_check(call_kwargs) + if refusal is not None: + await self._send_error(refusal, "invalid_request_error") + return + try: await _enforce_frame_project_quota( self.quota_callbacks, self.user_api_key_dict, self.model_group or self.model, raw_message @@ -3152,7 +3164,6 @@ class ManagedResponsesWebSocketHandler: await self._send_error(str(e), error_type="rate_limit_exceeded") return - call_kwargs: Final = self._build_base_call_kwargs(msg_obj) call_kwargs["stream"] = True # A frame that repeats the connection's public alias (model_group) must diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index bde570a8f61..077d7e2aa3c 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1612,6 +1612,81 @@ class TestManagedResponsesWSFirstMessage: assert processed == [subsequent] +class TestResponsesWSS3MediaFrameCheck: + @pytest.mark.asyncio + async def test_managed_socket_refuses_a_non_admin_s3_image_on_every_frame_and_forwards_a_proxy_admin_frame( + self, monkeypatch: pytest.MonkeyPatch + ): + from litellm import Router + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import _responses_websocket_session + + image_uri: Final = "s3://media-bucket/dir/cat.png" + frame: Final = json.dumps( + { + "type": "response.create", + "input": [{"role": "user", "content": [{"type": "input_image", "image_url": image_uri}]}], + } + ) + forwarded_event: Final = {"type": "response.created", "response": {"id": "resp_admin"}} + provider_inputs: Final[list[object]] = [] + + async def fake_aresponses(**kwargs: object): + provider_inputs.append(kwargs["input"]) + + async def events(): + yield forwarded_event + + return events() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + test_router: Final = Router( + model_list=[{"model_name": "nova", "litellm_params": {"model": "bedrock/us.amazon.nova-2-lite-v1:0"}}] + ) + + async def route_to_test_router(*, data: dict[str, Any], **_kwargs: object): + return test_router._aresponses_websocket(**data) + + async def frames_sent_back(user_role: LitellmUserRoles) -> list[object]: + ws = MagicMock() + ws.scope = {"headers": []} + ws.url = "ws://testserver/v1/responses" + ws.receive_text = AsyncMock(side_effect=[frame, frame, Exception("disconnect")]) + ws.send_text = AsyncMock() + ws.close = AsyncMock() + with ( + patch( # test-quality-ok: the pre-call processor needs a live proxy and is not what this test covers + "litellm.proxy.response_api_endpoints.endpoints.ProxyBaseLLMRequestProcessing", + side_effect=lambda data: MagicMock( + common_processing_pre_call_logic=AsyncMock(return_value=(data, MagicMock())) + ), + ), + patch( # test-quality-ok: routes to a router holding one Bedrock deployment instead of the proxy's + "litellm.proxy.route_llm_request.route_request", + side_effect=route_to_test_router, + ), + ): + await _responses_websocket_session( + websocket=ws, + model="nova", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="user-a", user_role=user_role), + ) + return [json.loads(call.args[0]) for call in ws.send_text.await_args_list] + + refused: Final = await frames_sent_back(LitellmUserRoles.INTERNAL_USER) + + assert provider_inputs == [] + assert [event["type"] for event in refused] == ["error", "error"] + assert all("allowed_client_side_s3_buckets" in event["error"]["message"] for event in refused) + assert all(image_uri in event["error"]["message"] for event in refused) + + admitted: Final = await frames_sent_back(LitellmUserRoles.PROXY_ADMIN) + + assert admitted == [forwarded_event, forwarded_event] + assert len(provider_inputs) == 2 + assert all(image_uri in json.dumps(provider_input) for provider_input in provider_inputs) + + class TestResponsesWSStreamingFirstMessage: @pytest.mark.asyncio async def test_client_to_backend_replays_first_message(self):