diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 30a88c04013..ec00e5cd03a 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -22,7 +22,8 @@ from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN if TYPE_CHECKING: - from litellm.proxy._types import LiteLLM_ManagedObjectTable + from prisma.models import LiteLLM_ManagedObjectTable + from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router @@ -207,7 +208,7 @@ class CheckResponsesCost: model_name = stored_response.get("model", None) # Decrypts rows written before model_object_id held the provider's own id. - responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(job.model_object_id) + responses_id_security = ResponsesIDSecurity().provider_response_id(job.model_object_id) # Prepare metadata with model information for cost tracking litellm_metadata = { diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 7e7f70d6f7e..b499006b64d 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -227,6 +227,11 @@ class ResponsesIDSecurity(CustomLogger): return True return False + def provider_response_id(self, response_id: str) -> str: + """The provider's own id behind an advertised one, returned unchanged when it is not encrypted.""" + original_response_id: Final = self._decrypt_response_id(response_id)[0] + return original_response_id + def _decrypt_response_id(self, response_id: str) -> tuple[str, str | None, str | None]: """ Returns: diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 9f0e4937bba..38399f4bd21 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -67,7 +67,7 @@ async def store_background_response_object( ) return - provider_response_id, _, _ = ResponsesIDSecurity()._decrypt_response_id(response.id) + provider_response_id: Final = ResponsesIDSecurity().provider_response_id(response.id) await managed_files_obj.store_unified_object_id( unified_object_id=response.id, file_object=response, diff --git a/tests/test_litellm/test_responses_id_security.py b/tests/test_litellm/test_responses_id_security.py index a6081670172..44504194686 100644 --- a/tests/test_litellm/test_responses_id_security.py +++ b/tests/test_litellm/test_responses_id_security.py @@ -1094,3 +1094,50 @@ class TestClientSuppliedRetainedIdCannotBypassAuthorization: assert result["response_id"] == "resp_strangerownprovideridcccccccc" assert result["response_id"] != victim_provider_id + + +class TestProviderResponseId: + """The provider's own id behind an advertised one, for callers that only need that.""" + + def test_a_real_encrypted_id_round_trips_to_the_provider_id( + self, responses_id_security, monkeypatch + ): + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-response-ids") + + advertised_id = "resp_" + str( + encrypt_value_helper( + value=SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( + "resp_provider_abc", "user-1", "team-1" + ) + ) + ) + + assert advertised_id != "resp_provider_abc" + assert responses_id_security.provider_response_id(advertised_id) == "resp_provider_abc" + + def test_two_encryptions_of_one_generation_resolve_to_the_same_provider_id( + self, responses_id_security, monkeypatch + ): + """Each advertised id carries a fresh nonce, so only the decrypted id can key a stored row.""" + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-response-ids") + + payload = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( + "resp_provider_abc", "user-1", "team-1" + ) + first = "resp_" + str(encrypt_value_helper(value=payload)) + second = "resp_" + str(encrypt_value_helper(value=payload)) + + assert first != second + assert responses_id_security.provider_response_id(first) == "resp_provider_abc" + assert responses_id_security.provider_response_id(second) == "resp_provider_abc" + + def test_a_raw_provider_id_is_returned_unchanged(self, responses_id_security, monkeypatch): + """Rows written before the provider id was stored hold an encrypted id, so both shapes + have to survive the same call.""" + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-response-ids") + + assert responses_id_security.provider_response_id("resp_provider_abc") == "resp_provider_abc"