diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7b3189b861c..1313b7f60b4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -308,27 +308,28 @@ def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]: ) -def _collect_ws_response_id_authorizer() -> ResponseIdAuthorizer | None: - """Duck-type discover the proxy hook that owns Responses id authorization, so +def _collect_ws_response_id_authorizers() -> tuple[ResponseIdAuthorizer, ...]: + """Duck-type discover the proxy hooks that own Responses id authorization, so the WebSocket surface refuses a ``previous_response_id`` the connection's key does not own through the very same step the HTTP routes run. Uses duck-typing on ``litellm.callbacks`` (rather than importing the proxy hook directly) to avoid a layering violation (SDK importing from the proxy - layer). Without the proxy there is no key to authorize and no hook to find. + layer). Every match is returned rather than the first, because config-loaded + callbacks are registered ahead of the proxy hooks: taking the first match + would let an unrelated callback answering to the same method name silently + stand in for the hook that owns this check. Without the proxy there is no key + to authorize and no hook to find. """ import litellm as _litellm callbacks: Final = cast( # cast-ok: callback registry is inspected before protocol use Sequence[object], _litellm.callbacks ) - return next( - ( - cast(ResponseIdAuthorizer, callback) # cast-ok: required callback method is callable - for callback in callbacks - if callable(getattr(callback, "response_id_ownership_refusal", None)) - ), - None, + return tuple( + cast(ResponseIdAuthorizer, callback) # cast-ok: required callback method is callable + for callback in callbacks + if callable(getattr(callback, "response_id_ownership_refusal", None)) ) @@ -6454,7 +6455,7 @@ class BaseLLMHTTPHandler: - Forwards events over the websocket connection """ _ws_quota_callbacks: Final = _collect_ws_project_quota_callbacks() - _ws_response_id_authorizer: Final = _collect_ws_response_id_authorizer() + _ws_response_id_authorizers: Final = _collect_ws_response_id_authorizers() if responses_api_provider_config is None or not responses_api_provider_config.supports_native_websocket(): from litellm.responses.streaming_iterator import ( @@ -6473,7 +6474,7 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, first_message=first_message, quota_callbacks=_ws_quota_callbacks, - response_id_authorizer=_ws_response_id_authorizer, + response_id_authorizers=_ws_response_id_authorizers, **kwargs, ) await handler.run() @@ -6596,7 +6597,7 @@ class BaseLLMHTTPHandler: output_guardrail_callbacks=_ws_output_guardrail_callbacks, quota_callbacks=_ws_quota_callbacks, authorized_model=model, - response_id_authorizer=_ws_response_id_authorizer, + response_id_authorizers=_ws_response_id_authorizers, ) await streaming.bidirectional_forward() diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 8a11639bcdd..7e12302e7ff 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1542,38 +1542,60 @@ async def _enforce_frame_project_quota( ) -def _frame_previous_response_id(raw_message: str) -> str | None: - """Read ``previous_response_id`` off a ``response.create`` frame, handling both - wire shapes: +def _frame_previous_response_ids(raw_message: str) -> tuple[str, ...]: + """Read every ``previous_response_id`` a ``response.create`` frame carries, across + both wire shapes: flat: {"type": "response.create", "previous_response_id": "..."} nested: {"type": "response.create", "response": {"previous_response_id": "..."}} + + A frame may carry both placements at once, and the native relay forwards the frame + to the provider close to verbatim, so every placement is reported rather than only + the one this proxy would itself read. ``_enforce_authorized_model`` defends the + ``model`` field across both placements for the same reason. """ try: msg_obj: Final = _load_json_value(raw_message) except (json.JSONDecodeError, TypeError): - return None + return () if not _is_json_object(msg_obj) or msg_obj.get("type") != "response.create": - return None + return () nested: Final = msg_obj.get("response") - params: Final[Mapping[str, object]] = nested if _is_json_object(nested) and nested else msg_obj - return _optional_str(params.get("previous_response_id")) + nested_params: Final[Mapping[str, object]] = nested if _is_json_object(nested) else {} + return tuple( + dict.fromkeys( + candidate + for candidate in ( + _optional_str(msg_obj.get("previous_response_id")), + _optional_str(nested_params.get("previous_response_id")), + ) + if candidate is not None + ) + ) -def _previous_response_id_refusal( - authorizer: ResponseIdAuthorizer | None, +def _previous_response_ids_refusal( + authorizers: Sequence[ResponseIdAuthorizer], user_api_key_dict: UserAPIKeyAuth | None, - previous_response_id: str, + previous_response_ids: Sequence[str], ) -> str | None: """Run the shared Responses id authorization step, the one the HTTP routes run, - against an id this connection was not itself handed. + against ids this connection was not itself handed. - Returns the refusal message when the connection's key may not address the id, - and None when it may. Without an authorizer or an authenticated key there is - no proxy in front of this socket and nothing to authorize against. + Returns the first refusal message when the connection's key may not address one + of the ids, and None when it may address all of them. Every discovered authorizer + is consulted, so an unrelated callback answering to the same method name can only + add refusals, never shadow the proxy hook that owns this check. Without an + authorizer or an authenticated key there is no proxy in front of this socket and + nothing to authorize against. """ - if authorizer is None or user_api_key_dict is None: + if not authorizers or user_api_key_dict is None: return None - return authorizer.response_id_ownership_refusal(previous_response_id, user_api_key_dict) + refusals: Final = ( + authorizer.response_id_ownership_refusal(previous_response_id, user_api_key_dict) + for previous_response_id in previous_response_ids + for authorizer in authorizers + ) + return next((refusal for refusal in refusals if refusal is not None), None) RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ @@ -1612,7 +1634,7 @@ class ResponsesWebSocketStreaming: output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, authorized_model: str | None = None, - response_id_authorizer: ResponseIdAuthorizer | None = None, + response_id_authorizers: Sequence[ResponseIdAuthorizer] | None = None, ): self.websocket = websocket self.backend_ws = backend_ws @@ -1628,7 +1650,9 @@ class ResponsesWebSocketStreaming: # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model - self.response_id_authorizer: ResponseIdAuthorizer | None = response_id_authorizer + self.response_id_authorizers: tuple[ResponseIdAuthorizer, ...] = ( + tuple(response_id_authorizers) if response_id_authorizers else () + ) # Response ids this connection handed the client. A connection carries # exactly one authenticated key, so continuing one of these needs no # further authorization; any other id does. @@ -2073,11 +2097,17 @@ class ResponsesWebSocketStreaming: def _unauthorized_previous_response_id(self, message: str) -> str | None: """Return why this connection's key may not continue the frame's - ``previous_response_id``, or None when it may.""" - previous_response_id: Final = _frame_previous_response_id(message) - if previous_response_id is None or previous_response_id in self._issued_response_ids: - return None - return _previous_response_id_refusal(self.response_id_authorizer, self.user_api_key_dict, previous_response_id) + ``previous_response_id``, or None when it may. + + The frame reaches the provider close to verbatim, so every placement the + frame carries is checked, not only the one this proxy would read itself. + """ + unowned_ids: Final = tuple( + previous_response_id + for previous_response_id in _frame_previous_response_ids(message) + if previous_response_id not in self._issued_response_ids + ) + return _previous_response_ids_refusal(self.response_id_authorizers, self.user_api_key_dict, unowned_ids) async def _enforce_or_reject_frame(self, message: str) -> bool: """Run the per-frame ownership and project quota checks. @@ -2193,7 +2223,7 @@ class ManagedResponsesWebSocketHandler: custom_llm_provider: str | None = None, first_message: str | None = None, quota_callbacks: Sequence[ProjectQuotaCallback] | None = None, - response_id_authorizer: ResponseIdAuthorizer | None = None, + response_id_authorizers: Sequence[ResponseIdAuthorizer] | None = None, **kwargs: object, ) -> None: self.websocket = websocket @@ -2212,7 +2242,9 @@ 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.response_id_authorizer: ResponseIdAuthorizer | None = response_id_authorizer + self.response_id_authorizers: tuple[ResponseIdAuthorizer, ...] = ( + tuple(response_id_authorizers) if response_id_authorizers else () + ) # 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. @@ -2465,8 +2497,8 @@ class ManagedResponsesWebSocketHandler: # No history on this connection, so the id was issued elsewhere: it goes # through the same authorization step the HTTP routes run before the # session is reconstructed from anywhere else. - refusal: Final = _previous_response_id_refusal( - self.response_id_authorizer, self.user_api_key_dict, previous_response_id + refusal: Final = _previous_response_ids_refusal( + self.response_id_authorizers, self.user_api_key_dict, (previous_response_id,) ) if refusal is not None: return refusal diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 3a3f0e6cad0..146573e6a23 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -22,7 +22,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import ( BaseLLMHTTPHandler, _collect_ws_project_quota_callbacks, - _collect_ws_response_id_authorizer, + _collect_ws_response_id_authorizers, _google_genai_streaming_hidden_params, _has_pre_call_deployment_hook, _rust_responses_websocket_enabled, @@ -2784,10 +2784,18 @@ def test_websocket_surface_finds_the_proxy_hook_that_authorizes_response_ids(mon plain, decoy, hook = _PlainLogger(), _NotCallableAttribute(), ResponsesIDSecurity() monkeypatch.setattr(litellm, "callbacks", [plain, decoy]) - assert _collect_ws_response_id_authorizer() is None + assert _collect_ws_response_id_authorizers() == () monkeypatch.setattr(litellm, "callbacks", [plain, decoy, hook]) - assert _collect_ws_response_id_authorizer() is hook + assert _collect_ws_response_id_authorizers() == (hook,) + + class _Impostor: + def response_id_ownership_refusal(self, response_id, user_api_key_dict): + return None + + impostor = _Impostor() + monkeypatch.setattr(litellm, "callbacks", [impostor, hook]) + assert _collect_ws_response_id_authorizers() == (impostor, hook) @pytest.mark.asyncio diff --git a/tests/test_litellm/responses/test_responses_websocket_all_providers.py b/tests/test_litellm/responses/test_responses_websocket_all_providers.py index b2782409002..d4ce9a89369 100644 --- a/tests/test_litellm/responses/test_responses_websocket_all_providers.py +++ b/tests/test_litellm/responses/test_responses_websocket_all_providers.py @@ -2655,6 +2655,29 @@ class _CountingAuthorizer: return self.inner.response_id_ownership_refusal(response_id, user_api_key_dict) +def _native_streaming(monkeypatch, **kwargs): + from unittest.mock import AsyncMock, MagicMock + + from litellm.responses.streaming_iterator import ResponsesWebSocketStreaming + + monkeypatch.setenv("LITELLM_SALT_KEY", _WS_UNIT_TEST_SALT_KEY) + mock_backend_ws = MagicMock() + mock_backend_ws.send = AsyncMock() + mock_websocket = MagicMock() + mock_websocket.send_text = AsyncMock() + return ResponsesWebSocketStreaming( + websocket=mock_websocket, + backend_ws=mock_backend_ws, + logging_obj=MagicMock(), + authorized_model="gpt-4o", + **kwargs, + ) + + +def _sent_error_type(mock_websocket): + return json.loads(mock_websocket.send_text.call_args[0][0])["error"]["type"] + + def _managed_handler(monkeypatch, **kwargs): from unittest.mock import AsyncMock, MagicMock @@ -2712,7 +2735,7 @@ class TestWebSocketPreviousResponseIdOwnership: handler, mock_websocket = _managed_handler( monkeypatch, user_api_key_dict=_ws_auth(), - response_id_authorizer=_ws_id_authorizer(), + response_id_authorizers=[_ws_id_authorizer()], ) await handler._process_response_create( @@ -2736,7 +2759,7 @@ class TestWebSocketPreviousResponseIdOwnership: handler, _ = _managed_handler( monkeypatch, user_api_key_dict=_ws_auth(), - response_id_authorizer=authorizer, + response_id_authorizers=[authorizer], ) handler._store_history(_WS_FABRICATED_PROVIDER_ID, [{"role": "user", "content": "turn one"}]) @@ -2759,7 +2782,7 @@ class TestWebSocketPreviousResponseIdOwnership: handler, _ = _managed_handler( monkeypatch, user_api_key_dict=_ws_auth(), - response_id_authorizer=_ws_id_authorizer({"allow_unmanaged_response_ids": True}), + response_id_authorizers=[_ws_id_authorizer({"allow_unmanaged_response_ids": True})], ) await handler._process_response_create( @@ -2810,7 +2833,7 @@ class TestWebSocketPreviousResponseIdOwnership: logging_obj=MagicMock(), user_api_key_dict=_ws_auth(), authorized_model="gpt-4o", - response_id_authorizer=_ws_id_authorizer(), + response_id_authorizers=[_ws_id_authorizer()], ) allowed = await handler._enforce_or_reject_frame( @@ -2844,7 +2867,7 @@ class TestWebSocketPreviousResponseIdOwnership: logging_obj=MagicMock(), user_api_key_dict=_ws_auth(), authorized_model="gpt-4o", - response_id_authorizer=authorizer, + response_id_authorizers=[authorizer], ) handler._record_issued_response_id( json.dumps( @@ -2881,7 +2904,7 @@ class TestWebSocketPreviousResponseIdOwnership: logging_obj=MagicMock(), user_api_key_dict=_ws_auth(), authorized_model="gpt-4o", - response_id_authorizer=_ws_id_authorizer({"allow_unmanaged_response_ids": True}), + response_id_authorizers=[_ws_id_authorizer({"allow_unmanaged_response_ids": True})], ) allowed = await handler._enforce_or_reject_frame( @@ -2901,27 +2924,127 @@ class TestWebSocketPreviousResponseIdOwnership: [ ( {"type": "response.create", "previous_response_id": _WS_FABRICATED_STRANGER_ID}, - _WS_FABRICATED_STRANGER_ID, + (_WS_FABRICATED_STRANGER_ID,), ), ( { "type": "response.create", "response": {"previous_response_id": _WS_FABRICATED_STRANGER_ID}, }, - _WS_FABRICATED_STRANGER_ID, + (_WS_FABRICATED_STRANGER_ID,), ), - ({"type": "response.create", "input": "hi"}, None), - ({"type": "response.cancel", "previous_response_id": _WS_FABRICATED_STRANGER_ID}, None), - ({"type": "response.create", "previous_response_id": {"forged": "object"}}, None), + ( + { + "type": "response.create", + "previous_response_id": _WS_FABRICATED_STRANGER_ID, + "response": {"input": "continue please"}, + }, + (_WS_FABRICATED_STRANGER_ID,), + ), + ( + { + "type": "response.create", + "previous_response_id": _WS_FABRICATED_STRANGER_ID, + "response": {"previous_response_id": _WS_FABRICATED_PROVIDER_ID}, + }, + (_WS_FABRICATED_STRANGER_ID, _WS_FABRICATED_PROVIDER_ID), + ), + ( + { + "type": "response.create", + "previous_response_id": _WS_FABRICATED_STRANGER_ID, + "response": {"previous_response_id": _WS_FABRICATED_STRANGER_ID}, + }, + (_WS_FABRICATED_STRANGER_ID,), + ), + ({"type": "response.create", "input": "hi"}, ()), + ({"type": "response.cancel", "previous_response_id": _WS_FABRICATED_STRANGER_ID}, ()), + ({"type": "response.create", "previous_response_id": {"forged": "object"}}, ()), ], ) - def test_previous_response_id_is_read_from_both_wire_shapes(self, frame, expected): - from litellm.responses.streaming_iterator import _frame_previous_response_id + def test_previous_response_ids_are_read_from_every_wire_placement(self, frame, expected): + from litellm.responses.streaming_iterator import _frame_previous_response_ids - assert _frame_previous_response_id(json.dumps(frame)) == expected + assert _frame_previous_response_ids(json.dumps(frame)) == expected def test_malformed_frames_read_as_carrying_no_previous_response_id(self): - from litellm.responses.streaming_iterator import _frame_previous_response_id + from litellm.responses.streaming_iterator import _frame_previous_response_ids - assert _frame_previous_response_id("not json at all") is None - assert _frame_previous_response_id(json.dumps(["not", "an", "object"])) is None + assert _frame_previous_response_ids("not json at all") == () + assert _frame_previous_response_ids(json.dumps(["not", "an", "object"])) == () + + @pytest.mark.asyncio + async def test_native_relay_refuses_a_stranger_id_hidden_beside_a_nested_response(self, monkeypatch): + """A frame carrying a top-level ``previous_response_id`` next to a non-empty + ``response`` object reaches the provider socket verbatim, so the id is refused + even though this proxy would read the nested placement for its own call.""" + streaming = _native_streaming( + monkeypatch, + user_api_key_dict=_ws_auth(), + response_id_authorizers=[_ws_id_authorizer()], + ) + + allowed = await streaming._enforce_or_reject_frame( + json.dumps( + { + "type": "response.create", + "previous_response_id": _WS_FABRICATED_STRANGER_ID, + "response": {"input": "continue please"}, + } + ) + ) + + assert allowed is False + streaming.backend_ws.send.assert_not_called() + assert _sent_error_type(streaming.websocket) == "permission_denied" + + @pytest.mark.asyncio + async def test_managed_handler_never_forwards_a_stranger_id_beside_a_nested_response(self, monkeypatch): + """The managed handler rebuilds the upstream call from the nested params, so a + top-level ``previous_response_id`` sibling is dropped rather than forwarded.""" + captured = _capture_aresponses(monkeypatch) + handler, _ = _managed_handler( + monkeypatch, + user_api_key_dict=_ws_auth(), + response_id_authorizers=[_ws_id_authorizer()], + ) + + await handler._process_response_create( + json.dumps( + { + "type": "response.create", + "previous_response_id": _WS_FABRICATED_STRANGER_ID, + "response": {"input": "continue please"}, + } + ) + ) + + assert captured.get("called") is True + assert "previous_response_id" not in captured + + @pytest.mark.asyncio + async def test_a_callback_answering_to_the_same_name_cannot_shadow_the_proxy_hook(self, monkeypatch): + """Config-loaded callbacks are registered ahead of the proxy hooks, so the + surface consults every discovered authorizer rather than only the first.""" + + class _PermissiveImpostor: + def response_id_ownership_refusal(self, response_id, user_api_key_dict): + return None + + streaming = _native_streaming( + monkeypatch, + user_api_key_dict=_ws_auth(), + response_id_authorizers=[_PermissiveImpostor(), _ws_id_authorizer()], + ) + + allowed = await streaming._enforce_or_reject_frame( + json.dumps( + { + "type": "response.create", + "previous_response_id": _WS_FABRICATED_STRANGER_ID, + } + ) + ) + + assert allowed is False + streaming.backend_ws.send.assert_not_called()