From 5ef4feba948333bc4d1a8bcec99518253db525ad Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 9 Nov 2026 17:04:59 -0800 Subject: [PATCH] fix(proxy): check s3 media on every Responses WebSocket frame on the LTS line The managed Responses WebSocket handler calls litellm.aresponses once per response.create frame, and the HTTP s3 media check only saw the first frame as a raw JSON string. A non-admin key could send an s3 object from an unlisted bucket, or one under LiteLLM-managed storage, to Bedrock over WS /v1/responses The WS route now hands the handler a frame check bound to the connection's key, general_settings, managed files hook and router. The handler runs it on the params it is about to forward for every frame, the first one included, and a refused frame gets an invalid_request_error frame with the 403 message while the socket keeps serving later frames --- litellm/llms/custom_httpx/llm_http_handler.py | 3 + .../proxy/response_api_endpoints/endpoints.py | 32 +++++++- litellm/responses/main.py | 4 +- litellm/responses/streaming_iterator.py | 13 +++- .../response_api_endpoints/test_endpoints.py | 75 +++++++++++++++++++ 5 files changed, 124 insertions(+), 3 deletions(-) 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):