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 c756832a7b7..47a6f38e8cb 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -238,8 +238,8 @@ class CheckResponsesCost: stored_response = job.file_object model_name = stored_response.get("model", None) - # Decrypt the response ID - responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(unified_object_id) + # Decrypts rows written before model_object_id held the provider's own id. + responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(job.model_object_id) # Prepare metadata with model information for cost tracking litellm_metadata = { diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index f320e2c65fb..95f3615026b 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -379,6 +379,10 @@ async def responses_api( if managed_files_obj and llm_router: try: + from litellm.proxy.hooks.responses_id_security import ( + ResponsesIDSecurity, + ) + # Get the actual deployment model_id from hidden params hidden_params: Final = getattr(response, "_hidden_params", {}) or {} model_id: Final = hidden_params.get("model_id", None) @@ -389,12 +393,13 @@ async def responses_api( response.id, ) raise Exception("No model_id found in response hidden params") + provider_response_id, _, _ = ResponsesIDSecurity()._decrypt_response_id(response.id) # Store in managed objects table await managed_files_obj.store_unified_object_id( unified_object_id=response.id, file_object=response, litellm_parent_otel_span=None, - model_object_id=response.id, + model_object_id=provider_response_id, file_purpose="response", user_api_key_dict=user_api_key_dict, persist_attribution=True, diff --git a/tests/proxy_unit_tests/test_check_responses_cost.py b/tests/proxy_unit_tests/test_check_responses_cost.py index bbf2b88e5b7..351dc6a4cac 100644 --- a/tests/proxy_unit_tests/test_check_responses_cost.py +++ b/tests/proxy_unit_tests/test_check_responses_cost.py @@ -142,6 +142,7 @@ class TestCheckResponsesCost: # Mock job with response ID mock_job = MagicMock() mock_job.unified_object_id = "resp_test_123" + mock_job.model_object_id = "resp_test_123" mock_job.created_by = "test-user" mock_job.id = "job-123" mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_123"} @@ -185,6 +186,7 @@ class TestCheckResponsesCost: # Mock job mock_job = MagicMock() mock_job.unified_object_id = "resp_test_456" + mock_job.model_object_id = "resp_test_456" mock_job.created_by = "test-user" mock_job.id = "job-456" mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_456"} @@ -224,6 +226,7 @@ class TestCheckResponsesCost: # Mock job mock_job = MagicMock() mock_job.unified_object_id = "resp_test_789" + mock_job.model_object_id = "resp_test_789" mock_job.created_by = "test-user" mock_job.id = "job-789" mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_789"} @@ -263,6 +266,7 @@ class TestCheckResponsesCost: # Mock job mock_job = MagicMock() mock_job.unified_object_id = "resp_test_in_progress" + mock_job.model_object_id = "resp_test_in_progress" mock_job.created_by = "test-user" mock_job.id = "job-in-progress" mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_in_progress"} @@ -304,6 +308,7 @@ class TestCheckResponsesCost: # Mock job mock_job = MagicMock() mock_job.unified_object_id = "resp_test_queued" + mock_job.model_object_id = "resp_test_queued" mock_job.created_by = "test-user" mock_job.id = "job-queued" mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_queued"} @@ -345,6 +350,7 @@ class TestCheckResponsesCost: # Mock job mock_job = MagicMock() mock_job.unified_object_id = "resp_test_error" + mock_job.model_object_id = "resp_test_error" mock_job.created_by = "test-user" mock_job.id = "job-error" mock_job.file_object = {"model": "gpt-4o", "id": "resp_test_error"} @@ -379,18 +385,21 @@ class TestCheckResponsesCost: # Mock multiple jobs mock_job1 = MagicMock() mock_job1.unified_object_id = "resp_test_1" + mock_job1.model_object_id = "resp_test_1" mock_job1.created_by = "user1" mock_job1.id = "job-1" mock_job1.file_object = {"model": "gpt-4o", "id": "resp_test_1"} mock_job2 = MagicMock() mock_job2.unified_object_id = "resp_test_2" + mock_job2.model_object_id = "resp_test_2" mock_job2.created_by = "user2" mock_job2.id = "job-2" mock_job2.file_object = {"model": "gpt-4o", "id": "resp_test_2"} mock_job3 = MagicMock() mock_job3.unified_object_id = "resp_test_3" + mock_job3.model_object_id = "resp_test_3" mock_job3.created_by = "user3" mock_job3.id = "job-3" mock_job3.file_object = {"model": "gpt-4o", "id": "resp_test_3"} @@ -470,6 +479,7 @@ class TestCheckResponsesCost: mock_job = MagicMock() mock_job.unified_object_id = encoded_response_id + mock_job.model_object_id = encoded_response_id mock_job.created_by = "test-user" mock_job.id = "job-router" mock_job.file_object = {"model": "azure-gpt-5", "id": encoded_response_id} @@ -541,6 +551,7 @@ class TestCheckResponsesCost: mock_job = MagicMock() mock_job.unified_object_id = encrypted_response_id + mock_job.model_object_id = encrypted_response_id mock_job.created_by = "test-user" mock_job.id = "job-encrypted" mock_job.file_object = {"model": "gpt-5", "id": encrypted_response_id} @@ -586,6 +597,7 @@ class TestCheckResponsesCost: """Ids that carry no deployment info can't be routed, so fall back to the SDK.""" mock_job = MagicMock() mock_job.unified_object_id = "resp_plain_upstream_id" + mock_job.model_object_id = "resp_plain_upstream_id" mock_job.created_by = "test-user" mock_job.id = "job-plain" mock_job.file_object = {"model": "gpt-5", "id": "resp_plain_upstream_id"} @@ -635,6 +647,7 @@ class TestCheckResponsesCost: mock_job = MagicMock() mock_job.unified_object_id = encoded_response_id + mock_job.model_object_id = encoded_response_id mock_job.created_by = "test-user" mock_job.id = "job-missing-deployment" mock_job.file_object = {"model": "gpt-5", "id": encoded_response_id} @@ -677,6 +690,7 @@ class TestCheckResponsesCost: """'incomplete' is terminal in the Responses API, so the row must not stay queued.""" mock_job = MagicMock() mock_job.unified_object_id = "resp_test_incomplete" + mock_job.model_object_id = "resp_test_incomplete" mock_job.created_by = "test-user" mock_job.id = "job-incomplete" mock_job.file_object = {"model": "gpt-5", "id": "resp_test_incomplete"} @@ -710,6 +724,7 @@ class TestCheckResponsesCost: """When file_object has no 'model' key, model_name is None and metadata skips model fields.""" mock_job = MagicMock() mock_job.unified_object_id = "resp_test_no_model" + mock_job.model_object_id = "resp_test_no_model" mock_job.created_by = "test-user" mock_job.team_id = None mock_job.api_key = None @@ -752,6 +767,7 @@ class TestCheckResponsesCost: mock_job = MagicMock() mock_job.unified_object_id = "resp_test_billed" + mock_job.model_object_id = "resp_test_billed" mock_job.created_by = "test-user" mock_job.team_id = "team-billed" mock_job.api_key = "sk-billed" @@ -788,6 +804,7 @@ class TestCheckResponsesCost: spend log, so losing the claim has to skip the read entirely or the job is billed twice.""" mock_job = MagicMock() mock_job.unified_object_id = "resp_test_claimed_elsewhere" + mock_job.model_object_id = "resp_test_claimed_elsewhere" mock_job.created_by = "test-user" mock_job.id = "job-claimed-elsewhere" mock_job.file_object = {"model": "gpt-5", "id": "resp_test_claimed_elsewhere"} @@ -828,6 +845,7 @@ class TestCheckResponsesCost: mock_job = MagicMock() mock_job.unified_object_id = "resp_test_abandoned" + mock_job.model_object_id = "resp_test_abandoned" mock_job.created_by = "test-user" mock_job.id = "job-abandoned" mock_job.file_object = {"model": "gpt-5", "id": "resp_test_abandoned"} @@ -869,6 +887,7 @@ class TestCheckResponsesCost: afterwards is what stops a second pod reading and billing the same row again.""" mock_job = MagicMock() mock_job.unified_object_id = "resp_test_ordering" + mock_job.model_object_id = "resp_test_ordering" mock_job.created_by = "test-user" mock_job.id = "job-ordering" mock_job.file_object = {"model": "gpt-5", "id": "resp_test_ordering"} @@ -919,6 +938,7 @@ class TestCheckResponsesCost: back to batch_processed=False; holding the claim retires it before it is ever billed.""" mock_job = MagicMock() mock_job.unified_object_id = "resp_test_still_running" + mock_job.model_object_id = "resp_test_still_running" mock_job.created_by = "test-user" mock_job.id = "job-still-running" mock_job.file_object = {"model": "gpt-5", "id": "resp_test_still_running"} @@ -959,6 +979,7 @@ class TestCheckResponsesCost: retired unbilled and no later poll cycle ever retries it.""" mock_job = MagicMock() mock_job.unified_object_id = "resp_test_read_error" + mock_job.model_object_id = "resp_test_read_error" mock_job.created_by = "test-user" mock_job.id = "job-read-error" mock_job.file_object = {"model": "gpt-5", "id": "resp_test_read_error"} @@ -993,12 +1014,14 @@ class TestCheckResponsesCost: has to be read and billed in the same cycle.""" mock_job1 = MagicMock() mock_job1.unified_object_id = "resp_test_first" + mock_job1.model_object_id = "resp_test_first" mock_job1.created_by = "user1" mock_job1.id = "job-first" mock_job1.file_object = {"model": "gpt-5", "id": "resp_test_first"} mock_job2 = MagicMock() mock_job2.unified_object_id = "resp_test_second" + mock_job2.model_object_id = "resp_test_second" mock_job2.created_by = "user2" mock_job2.id = "job-second" mock_job2.file_object = {"model": "gpt-5", "id": "resp_test_second"} @@ -1078,6 +1101,7 @@ class TestCheckResponsesCost: (which is what bills it) and the row is still marked completed.""" mock_job = MagicMock() mock_job.unified_object_id = "resp_test_old_schema" + mock_job.model_object_id = "resp_test_old_schema" mock_job.created_by = "test-user" mock_job.id = "job-old-schema" mock_job.file_object = {"model": "gpt-5", "id": "resp_test_old_schema"} @@ -1122,12 +1146,14 @@ class TestCheckResponsesCost: are already read and billed, so losing their write loses their usage for good.""" mock_job1 = MagicMock() mock_job1.unified_object_id = "resp_test_persist_fails" + mock_job1.model_object_id = "resp_test_persist_fails" mock_job1.created_by = "user1" mock_job1.id = "job-persist-fails" mock_job1.file_object = {"model": "gpt-5", "id": "resp_test_persist_fails"} mock_job2 = MagicMock() mock_job2.unified_object_id = "resp_test_persist_works" + mock_job2.model_object_id = "resp_test_persist_works" mock_job2.created_by = "user2" mock_job2.id = "job-persist-works" mock_job2.file_object = {"model": "gpt-5", "id": "resp_test_persist_works"} @@ -1157,3 +1183,97 @@ class TestCheckResponsesCost: "job-persist-fails", "job-persist-works", ] + + @pytest.mark.asyncio + async def test_poller_fetches_the_provider_id_from_model_object_id( + self, check_responses_cost_instance, mock_prisma_client, mock_llm_router, monkeypatch + ): + """The row's provider id drives the fetch, not the nonce-encrypted advertised id. + + A background create advertises a freshly encrypted id per call, so unified_object_id + is not a stable handle on the generation. + """ + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.types.utils import SpecialEnums + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-response-ids") + + provider_response_id = "resp_provider_stable_1" + stale_advertised_id = "resp_" + str( + encrypt_value_helper( + value=SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( + "resp_some_other_encoding", "test-user", "test-team" + ) + ) + ) + + mock_job = MagicMock() + mock_job.unified_object_id = stale_advertised_id + mock_job.model_object_id = provider_response_id + mock_job.created_by = "test-user" + mock_job.id = "job-provider-id" + mock_job.file_object = {"model": "gpt-5", "id": stale_advertised_id} + + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + + mock_response = ResponsesAPIResponse( + id=provider_response_id, + object="response", + status="completed", + created_at=int(datetime.now().timestamp()), + output=[], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=5, total_tokens=15), + ) + + with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: + mock_aget.return_value = mock_response + await check_responses_cost_instance.check_responses_cost() + + assert mock_aget.call_args[1]["response_id"] == provider_response_id + assert _completed_job_ids(mock_prisma_client) == ["job-provider-id"] + + @pytest.mark.asyncio + async def test_poller_still_reads_rows_written_before_the_provider_id_was_stored( + self, check_responses_cost_instance, mock_prisma_client, mock_llm_router, monkeypatch + ): + """Rows created earlier carry the encrypted advertised id in both columns.""" + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.types.utils import SpecialEnums + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-response-ids") + + provider_response_id = "resp_legacy_upstream_9" + legacy_id = "resp_" + str( + encrypt_value_helper( + value=SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( + provider_response_id, "test-user", "test-team" + ) + ) + ) + + mock_job = MagicMock() + mock_job.unified_object_id = legacy_id + mock_job.model_object_id = legacy_id + mock_job.created_by = "test-user" + mock_job.id = "job-legacy" + mock_job.file_object = {"model": "gpt-5", "id": legacy_id} + + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + + mock_response = ResponsesAPIResponse( + id=provider_response_id, + object="response", + status="completed", + created_at=int(datetime.now().timestamp()), + output=[], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=5, total_tokens=15), + ) + + with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: + mock_aget.return_value = mock_response + await check_responses_cost_instance.check_responses_cost() + + assert mock_aget.call_args[1]["response_id"] == provider_response_id + assert _completed_job_ids(mock_prisma_client) == ["job-legacy"] diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index d7010de6405..f120f88bba6 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -1959,3 +1959,137 @@ class TestResponsesInputTokens: assert response.status_code == 429, response.text assert response.json()["error"]["message"] == "rate limited" + + +class TestBackgroundResponseManagedObjectId: + """The managed row for a background response must be keyed by the provider's own id. + + The advertised ``response.id`` is encrypted with a fresh nonce per call, so storing it + in ``model_object_id`` leaves the row with no stable lookup key and every later read + of the same generation looks like a new object. + """ + + @staticmethod + def _encrypted_id(provider_response_id: str) -> str: + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.types.utils import SpecialEnums + + managed_id = SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( + provider_response_id, "u-1", "t-1" + ) + return f"resp_{encrypt_value_helper(value=managed_id)}" + + async def _store_call_for(self, provider_response_id: str) -> dict: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import responses_api + from litellm.types.llms.openai import ResponsesAPIResponse + + advertised_id = self._encrypted_id(provider_response_id) + assert advertised_id != self._encrypted_id(provider_response_id), ( + "advertised ids must be nonce-encrypted, otherwise this regression cannot occur" + ) + + response = ResponsesAPIResponse( + id=advertised_id, + created_at=0, + model="gpt-4o", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + status="queued", + ) + response._hidden_params = {"model_id": "deployment-1"} + + managed_files_obj = MagicMock() + managed_files_obj.store_unified_object_id = AsyncMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.get_proxy_hook = MagicMock(return_value=managed_files_obj) + + with patch( + "litellm.proxy.proxy_server._read_request_body", + AsyncMock(return_value={"model": "gpt-4o", "input": "hi", "background": True}), + ), patch("litellm.proxy.proxy_server.polling_via_cache_enabled", False), patch( + "litellm.proxy.proxy_server.llm_router", MagicMock() + ), patch( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ), patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.base_process_llm_request", + AsyncMock(return_value=response), + ): + await responses_api( + request=MagicMock(), + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", user_id="u-1", team_id="t-1"), + ) + + managed_files_obj.store_unified_object_id.assert_awaited_once() + return managed_files_obj.store_unified_object_id.await_args.kwargs + + @pytest.mark.asyncio + async def test_model_object_id_is_the_provider_response_id(self, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-regression-salt") + provider_response_id = "resp_provider68abc123" + + kwargs = await self._store_call_for(provider_response_id) + + assert kwargs["model_object_id"] == provider_response_id + assert kwargs["unified_object_id"] != provider_response_id + assert kwargs["unified_object_id"] == kwargs["file_object"].id + + @pytest.mark.asyncio + async def test_two_background_creates_are_distinguishable_by_provider_id(self, monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-regression-salt") + + first = await self._store_call_for("resp_providerAAA") + second = await self._store_call_for("resp_providerBBB") + + assert first["model_object_id"] == "resp_providerAAA" + assert second["model_object_id"] == "resp_providerBBB" + + @pytest.mark.asyncio + async def test_unencrypted_advertised_id_is_stored_as_is(self, monkeypatch): + """With response-id security disabled the advertised id is already the provider's.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import responses_api + from litellm.types.llms.openai import ResponsesAPIResponse + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-regression-salt") + response = ResponsesAPIResponse( + id="resp_rawprovider999", + created_at=0, + model="gpt-4o", + object="response", + output=[], + parallel_tool_calls=False, + tool_choice="auto", + tools=[], + status="queued", + ) + response._hidden_params = {"model_id": "deployment-1"} + + managed_files_obj = MagicMock() + managed_files_obj.store_unified_object_id = AsyncMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.get_proxy_hook = MagicMock(return_value=managed_files_obj) + + with patch( + "litellm.proxy.proxy_server._read_request_body", + AsyncMock(return_value={"model": "gpt-4o", "input": "hi", "background": True}), + ), patch("litellm.proxy.proxy_server.polling_via_cache_enabled", False), patch( + "litellm.proxy.proxy_server.llm_router", MagicMock() + ), patch( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ), patch( + "litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing.base_process_llm_request", + AsyncMock(return_value=response), + ): + await responses_api( + request=MagicMock(), + fastapi_response=MagicMock(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", user_id="u-1", team_id="t-1"), + ) + + kwargs = managed_files_obj.store_unified_object_id.await_args.kwargs + assert kwargs["model_object_id"] == "resp_rawprovider999"