mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
refactor(responses): give ResponsesIDSecurity a public provider_response_id
Two callers only want the provider's own id behind an advertised one, and both reached past the class to get it out of the decrypt tuple. The poller also typed its rows as the pydantic projection, which has no `id`, while it is handed a Prisma row and reads `job.id` six times. Claude-Session: https://claude.ai/code/session_01RHAjRxNhXTpKHeGMZ1nDKi
This commit is contained in:
parent
2aac6e109d
commit
cacfc47089
4 changed files with 56 additions and 3 deletions
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue