diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 21d12c8f720..c4c15c40d1e 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -5,10 +5,11 @@ This hook uses the DBSpendUpdateWriter to batch-write response IDs to the databa instead of writing immediately on each request. """ -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Callable, Mapping from typing import TYPE_CHECKING, Any, Final, cast from fastapi import HTTPException +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_logger import CustomLogger @@ -32,6 +33,44 @@ _RESPONSES_API_PROVIDER_PREFIX: Final = "/openai" _RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) +_RESPONSE_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _response_payload(response_obj: object) -> Mapping[str, object] | None: + try: + return _RESPONSE_PAYLOAD_ADAPTER.validate_python(response_obj) + except ValidationError: + return None + + +def _rewrite_advertised_id( + event: BaseLiteLLMOpenAIResponseObject, + rewrite: Callable[[str], str], +) -> BaseLiteLLMOpenAIResponseObject: + event_id: Final = getattr(event, "id", None) + if isinstance(event_id, str) and event_id.startswith("resp_"): + setattr(event, "id", rewrite(event_id)) + return event + + nested: Final = getattr(event, "response", None) + if isinstance(nested, ResponsesAPIResponse): + setattr(nested, "id", rewrite(nested.id)) + setattr(event, "response", nested) + return event + + payload: Final = _response_payload(nested) + if payload is None: + return event + + payload_id: Final = payload.get("id") + if not isinstance(payload_id, str): + return event + + rewritten: Final = {**payload, "id": rewrite(payload_id)} # mutable-ok: pydantic cannot serialize a frozen map + setattr(event, "response", rewritten) + return event + + def _is_responses_api_create_route(request_route: str | None) -> bool: if request_route is None: return False @@ -196,10 +235,6 @@ class ResponsesIDSecurity(CustomLogger): user_api_key_dict: "UserAPIKeyAuth", request_cache: dict[str, str] | None = None, ) -> BaseLiteLLMOpenAIResponseObject: - # encrypt the response id using the symmetric key - # encrypt the response id, and encode the user id and response id in base64 - - # Check if signing key is available signing_key: Final = self._get_signing_key() if signing_key is None: verbose_proxy_logger.debug( @@ -210,43 +245,22 @@ class ResponsesIDSecurity(CustomLogger): ) return response - response_id: Final = getattr(response, "id", None) - response_obj: Final = getattr(response, "response", None) + def encrypt(original_id: str) -> str: + cached: Final = request_cache.get(original_id) if request_cache is not None else None + if cached is not None: + return cached - if response_id and isinstance(response_id, str) and response_id.startswith("resp_"): - # Check request-scoped cache first (for streaming consistency) - if request_cache is not None and response_id in request_cache: - setattr(response, "id", request_cache[response_id]) - else: - encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( - response_id, - user_api_key_dict.user_id or "", - user_api_key_dict.team_id or "", - ) + managed_id: Final = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( + original_id, + user_api_key_dict.user_id or "", + user_api_key_dict.team_id or "", + ) + encrypted_id: Final = f"resp_{encrypt_value_helper(value=managed_id)}" + if request_cache is not None: + request_cache[original_id] = encrypted_id + return encrypted_id - encoded_user_id_and_response_id = encrypt_value_helper(value=encrypted_response_id) - encrypted_id = f"resp_{encoded_user_id_and_response_id}" - if request_cache is not None: - request_cache[response_id] = encrypted_id - setattr(response, "id", encrypted_id) - - elif response_obj and isinstance(response_obj, ResponsesAPIResponse): - # Check request-scoped cache first (for streaming consistency) - if request_cache is not None and response_obj.id in request_cache: - setattr(response_obj, "id", request_cache[response_obj.id]) - else: - encrypted_response_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( - response_obj.id, - user_api_key_dict.user_id or "", - user_api_key_dict.team_id or "", - ) - encoded_user_id_and_response_id = encrypt_value_helper(value=encrypted_response_id) - encrypted_id = f"resp_{encoded_user_id_and_response_id}" - if request_cache is not None: - request_cache[response_obj.id] = encrypted_id - setattr(response_obj, "id", encrypted_id) - setattr(response, "response", response_obj) - return response + return _rewrite_advertised_id(response, encrypt) async def async_post_call_success_hook( self, diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 12c3cf46373..24c0ff6b181 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -42,7 +42,7 @@ "limit": 52 }, "B010": { - "limit": 190 + "limit": 187 }, "B018": { "limit": 2 diff --git a/tests/test_litellm/test_responses_id_security.py b/tests/test_litellm/test_responses_id_security.py index 763ee4dac00..d35b9563888 100644 --- a/tests/test_litellm/test_responses_id_security.py +++ b/tests/test_litellm/test_responses_id_security.py @@ -14,7 +14,9 @@ from litellm.proxy.hooks.responses_id_security import ( _is_responses_api_create_route, ) from litellm.types.llms.openai import ( + GenericEvent, ResponseCompletedEvent, + ResponseCreatedEvent, ResponsesAPIResponse, ResponsesAPIStreamEvents, ) @@ -691,6 +693,126 @@ class TestAsyncPostCallStreamingIteratorHook: assert not responses_id_security._is_encrypted_response_id(streamed_id) +class TestStreamedGenericEventIdEncryption: + """A background stream carries event types with no typed model, which arrive as + GenericEvent holding a plain dict. Those used to skip encryption while their typed + siblings were encrypted, so one stream advertised two ids and the unencrypted one + skipped the ownership check. Asserts the property rather than one event type: every + id a client can see is the same encrypted id, and the raw one appears in no frame.""" + + RAW_ID = "resp_rawprovider123" + + @staticmethod + async def _agen(chunks): + for chunk in chunks: + yield chunk + + @classmethod + def _typed_event(cls, event_type): + return { + ResponsesAPIStreamEvents.RESPONSE_CREATED: ResponseCreatedEvent, + ResponsesAPIStreamEvents.RESPONSE_COMPLETED: ResponseCompletedEvent, + }[event_type]( + type=event_type, + response=ResponsesAPIResponse( + id=cls.RAW_ID, + created_at=0, + model="gpt-5.1", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + ), + ) + + @classmethod + def _background_stream(cls): + return [ + cls._typed_event(ResponsesAPIStreamEvents.RESPONSE_CREATED), + GenericEvent( + type="response.queued", + response={"id": cls.RAW_ID, "status": "queued"}, + ), + GenericEvent(type="keepalive"), + GenericEvent( + type="response.some_event_openai_adds_later", + response={"id": cls.RAW_ID, "status": "in_progress"}, + ), + cls._typed_event(ResponsesAPIStreamEvents.RESPONSE_COMPLETED), + ] + + @staticmethod + def _advertised_ids(events): + nested = (getattr(event, "response", None) for event in events) + return [ + payload["id"] if isinstance(payload, dict) else payload.id + for payload in nested + if payload is not None + ] + [ + event.id for event in events if isinstance(getattr(event, "id", None), str) + ] + + async def _drain(self, responses_id_security, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-abcdefghij") + + mock_auth = MagicMock() + mock_auth.user_id = "user-a" + mock_auth.team_id = "team-a" + mock_auth.request_route = "/v1/responses" + + return [ + out + async for out in responses_id_security.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_auth, + response=self._agen(self._background_stream()), + request_data={}, + ) + ] + + @pytest.mark.asyncio + async def test_every_event_advertises_the_same_encrypted_id( + self, responses_id_security, monkeypatch + ): + events = await self._drain(responses_id_security, monkeypatch) + advertised = self._advertised_ids(events) + + assert len(advertised) == 4 + assert len(set(advertised)) == 1 + + streamed_id = advertised[0] + assert streamed_id != self.RAW_ID + assert responses_id_security._is_encrypted_response_id(streamed_id) + assert responses_id_security._decrypt_response_id(streamed_id) == ( + self.RAW_ID, + "user-a", + "team-a", + ) + + @pytest.mark.asyncio + async def test_raw_provider_id_never_reaches_the_client( + self, responses_id_security, monkeypatch + ): + events = await self._drain(responses_id_security, monkeypatch) + + assert [self.RAW_ID in event.model_dump_json() for event in events] == [ + False + ] * len(events) + + @pytest.mark.asyncio + async def test_sibling_fields_survive_the_rewrite( + self, responses_id_security, monkeypatch + ): + _, queued, keepalive, later, _ = await self._drain( + responses_id_security, monkeypatch + ) + + assert queued.response["status"] == "queued" + assert later.response["status"] == "in_progress" + assert keepalive.type == "keepalive" + assert getattr(keepalive, "response", None) is None + + class TestAsyncPostCallSuccessHook: """Test async_post_call_success_hook function""" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 80b4531439d..2f85128b4b6 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16476 + "limit": 16470 }, "LIT011": { - "limit": 5518 + "limit": 5516 }, "LIT012": { "limit": 4489