From 49a8ce5e42b20ff8b217198f75aece683da59297 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 9 Sep 2026 17:57:23 -0700 Subject: [PATCH] fix(responses): key a background response's managed row by the provider id model_object_id is documented as "the id returned by the backend API provider", and that is what batches and fine-tuning jobs store there. The background responses create stored the advertised id in it instead, which is encrypted with a fresh nonce on every call, so the row had no stable handle on the generation it describes. The cost poller now reads the provider id straight off the row. Rows written before this still carry the advertised id there, and decrypting is a no-op on an id that is already the provider's, so both shapes resolve through the same call. Claude-Session: https://claude.ai/code/session_01RHAjRxNhXTpKHeGMZ1nDKi --- .../common_utils/check_responses_cost.py | 4 +- .../proxy/response_api_endpoints/endpoints.py | 7 +- .../test_check_responses_cost.py | 120 ++++++++++++++++ .../response_api_endpoints/test_endpoints.py | 134 ++++++++++++++++++ 4 files changed, 262 insertions(+), 3 deletions(-) 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"