diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 95f3615026b..9f0e4937bba 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -39,10 +39,47 @@ from litellm.types.responses.main import DeleteResponseResult from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: + from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles + from litellm.router import Router router: Final = APIRouter() + +async def store_background_response_object( + response: ResponsesAPIResponse, + managed_files_obj: "_PROXY_LiteLLMManagedFiles", + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Record a queued background response so the cost poller can find and bill it. + + ``model_object_id`` carries the provider's own id because the advertised ``response.id`` + is re-encrypted with a fresh nonce on every call, leaving the row no stable handle on + the generation it describes. + """ + from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity + + hidden_params: Final = getattr(response, "_hidden_params", {}) or {} + if not hidden_params.get("model_id"): + verbose_proxy_logger.warning( + "No model_id found in response hidden params for response %s, skipping managed object storage", + response.id, + ) + return + + provider_response_id, _, _ = ResponsesIDSecurity()._decrypt_response_id(response.id) + await managed_files_obj.store_unified_object_id( + unified_object_id=response.id, + file_object=response, + litellm_parent_otel_span=None, + model_object_id=provider_response_id, + file_purpose="response", + user_api_key_dict=user_api_key_dict, + persist_attribution=True, + ) + verbose_proxy_logger.info("Stored background response %s in managed objects table", response.id) + + _user_api_key_auth_dep: Final = Depends(user_api_key_auth) _RESPONSES_TAGS: Final[list[str | Enum]] = ["responses"] # mutable-ok: fastapi's route signature requires list tags @@ -366,54 +403,29 @@ async def responses_api( ) # Store in managed objects table if background mode is enabled - if data.get("background") and isinstance(response, ResponsesAPIResponse): - if response.status in ["queued", "in_progress"]: - from litellm_enterprise.proxy.hooks.managed_files import ( - _PROXY_LiteLLMManagedFiles, - ) + if ( + data.get("background") + and isinstance(response, ResponsesAPIResponse) + and response.status in ("queued", "in_progress") + ): + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) - managed_files_obj: Final = cast( - _PROXY_LiteLLMManagedFiles | None, - proxy_logging_obj.get_proxy_hook("managed_files"), - ) + managed_files_obj: Final = cast( + _PROXY_LiteLLMManagedFiles | None, + proxy_logging_obj.get_proxy_hook("managed_files"), + ) - 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) - - if not model_id: - verbose_proxy_logger.warning( - "No model_id found in response hidden params for response %s, skipping managed object storage", - 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=provider_response_id, - file_purpose="response", - user_api_key_dict=user_api_key_dict, - persist_attribution=True, - ) - - verbose_proxy_logger.info( - "Stored background response %s in managed objects table with unified_id=%s", - response.id, - response.id, - ) - except Exception as e: - verbose_proxy_logger.error( - "Failed to store background response in managed objects table: %s", e - ) + if managed_files_obj and llm_router: + try: + await store_background_response_object( + response=response, + managed_files_obj=managed_files_obj, + user_api_key_dict=user_api_key_dict, + ) + except Exception as e: + verbose_proxy_logger.error("Failed to store background response in managed objects table: %s", e) return response except ModifyResponseException as e: diff --git a/tests/proxy_unit_tests/test_check_responses_cost.py b/tests/proxy_unit_tests/test_check_responses_cost.py index 351dc6a4cac..4768cdc576d 100644 --- a/tests/proxy_unit_tests/test_check_responses_cost.py +++ b/tests/proxy_unit_tests/test_check_responses_cost.py @@ -1191,18 +1191,23 @@ class TestCheckResponsesCost: """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. + is no handle on the generation. """ from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper + from litellm.responses.utils import ResponsesAPIRequestUtils 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" + provider_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", + model_id="deployment-xyz", + response_id="resp_upstream_stable", + ) 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" + "resp_a_previous_encoding", "test-user", "test-team" ) ) ) @@ -1216,21 +1221,20 @@ class TestCheckResponsesCost: 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), + mock_llm_router.aget_responses = AsyncMock( + return_value=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() + await check_responses_cost_instance.check_responses_cost() - assert mock_aget.call_args[1]["response_id"] == provider_response_id + assert mock_llm_router.aget_responses.call_args[1]["response_id"] == provider_response_id assert _completed_job_ids(mock_prisma_client) == ["job-provider-id"] @pytest.mark.asyncio @@ -1239,11 +1243,16 @@ class TestCheckResponsesCost: ): """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.responses.utils import ResponsesAPIRequestUtils 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" + provider_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", + model_id="deployment-xyz", + response_id="resp_legacy_upstream", + ) legacy_id = "resp_" + str( encrypt_value_helper( value=SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format( @@ -1261,19 +1270,18 @@ class TestCheckResponsesCost: 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), + mock_llm_router.aget_responses = AsyncMock( + return_value=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() + await check_responses_cost_instance.check_responses_cost() - assert mock_aget.call_args[1]["response_id"] == provider_response_id + assert mock_llm_router.aget_responses.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 f120f88bba6..f528dc6ba58 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -1962,11 +1962,11 @@ class TestResponsesInputTokens: class TestBackgroundResponseManagedObjectId: - """The managed row for a background response must be keyed by the provider's own id. + """The managed row for a background response is 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. + in ``model_object_id`` leaves the row with no stable handle on the generation and every + later read of the same generation looks like a new object. """ @staticmethod @@ -1979,16 +1979,10 @@ class TestBackgroundResponseManagedObjectId: ) 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 + @staticmethod + def _queued_response(advertised_id: str, model_id: str | None = "deployment-1"): 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, @@ -2000,50 +1994,60 @@ class TestBackgroundResponseManagedObjectId: tools=[], status="queued", ) - response._hidden_params = {"model_id": "deployment-1"} + response._hidden_params = {"model_id": model_id} if model_id else {} + return response + + async def _stored_kwargs(self, advertised_id: str, model_id: str | None = "deployment-1"): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.response_api_endpoints.endpoints import ( + store_background_response_object, + ) 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 + await store_background_response_object( + response=self._queued_response(advertised_id, model_id), + managed_files_obj=managed_files_obj, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", user_id="u-1", team_id="t-1"), + ) + return managed_files_obj.store_unified_object_id @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" + 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" + ) - kwargs = await self._store_call_for(provider_response_id) + store = await self._stored_kwargs(advertised_id) + store.assert_awaited_once() + kwargs = store.await_args.kwargs 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 + assert kwargs["unified_object_id"] == advertised_id + assert kwargs["file_object"].id == advertised_id @pytest.mark.asyncio - async def test_two_background_creates_are_distinguishable_by_provider_id(self, monkeypatch): + async def test_two_creates_of_one_generation_share_a_provider_id(self, monkeypatch): + """Re-encrypting the same generation must not look like a second object.""" + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-regression-salt") + provider_response_id = "resp_provider_same_gen" + + first = (await self._stored_kwargs(self._encrypted_id(provider_response_id))).await_args.kwargs + second = (await self._stored_kwargs(self._encrypted_id(provider_response_id))).await_args.kwargs + + assert first["unified_object_id"] != second["unified_object_id"] + assert first["model_object_id"] == second["model_object_id"] == provider_response_id + + @pytest.mark.asyncio + async def test_distinct_generations_keep_distinct_provider_ids(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") + first = (await self._stored_kwargs(self._encrypted_id("resp_providerAAA"))).await_args.kwargs + second = (await self._stored_kwargs(self._encrypted_id("resp_providerBBB"))).await_args.kwargs assert first["model_object_id"] == "resp_providerAAA" assert second["model_object_id"] == "resp_providerBBB" @@ -2051,45 +2055,17 @@ class TestBackgroundResponseManagedObjectId: @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) + store = await self._stored_kwargs("resp_rawprovider999") - 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"), - ) + assert store.await_args.kwargs["model_object_id"] == "resp_rawprovider999" - kwargs = managed_files_obj.store_unified_object_id.await_args.kwargs - assert kwargs["model_object_id"] == "resp_rawprovider999" + @pytest.mark.asyncio + async def test_response_without_a_deployment_is_not_stored(self, monkeypatch): + """No model_id means the poller could never route the read, so no row is written.""" + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-regression-salt") + + store = await self._stored_kwargs(self._encrypted_id("resp_no_deployment"), model_id=None) + + store.assert_not_awaited()