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 5d8139f7888..251a1d2834e 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -47,8 +47,8 @@ def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_Manag return ManagedObjectRepository(prisma_client).table -def _is_provider_not_found(error: Exception) -> bool: - return getattr(error, "status_code", None) == 404 +def _is_response_gone_at_provider(error: Exception, provider_response_id: str) -> bool: + return getattr(error, "status_code", None) == 404 and provider_response_id in str(error) class CheckResponsesCost: @@ -196,13 +196,16 @@ class CheckResponsesCost: if model_name: litellm_metadata["model"] = model_name litellm_metadata["model_group"] = model_name # Use same value for model_group + via_router = self._resolve_deployment(responses_id_security) except Exception as e: verbose_proxy_logger.warning( f"Skipping job {unified_object_id} due to error: {e}" ) continue - via_router = self._resolve_deployment(responses_id_security) + provider_response_id = ResponsesAPIRequestUtils.decode_responses_api_response_id( + responses_id_security + ).get("response_id", responses_id_security) try: response = await self._get_response( response_id=responses_id_security, @@ -215,7 +218,7 @@ class CheckResponsesCost: ) except Exception as e: - if via_router and _is_provider_not_found(e): + if via_router and _is_response_gone_at_provider(e, provider_response_id): verbose_proxy_logger.info( f"Response {unified_object_id} no longer available at provider (404), marking stale_expired: {e}" ) diff --git a/tests/unit/proxy/common_utils/test_check_responses_cost.py b/tests/unit/proxy/common_utils/test_check_responses_cost.py index 4e41f9434aa..8ba3f0b000a 100644 --- a/tests/unit/proxy/common_utils/test_check_responses_cost.py +++ b/tests/unit/proxy/common_utils/test_check_responses_cost.py @@ -3,7 +3,9 @@ Unit tests for CheckResponsesCost class """ import asyncio +from collections.abc import Mapping from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest @@ -12,6 +14,19 @@ from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +class _RecordingManagedObjectTable: + def __init__(self, rows: tuple[object, ...]) -> None: + self.rows: Final = rows + self.updates: tuple[tuple[Mapping[str, object], Mapping[str, object]], ...] = () + + async def find_many(self, **query: object) -> tuple[object, ...]: + return self.rows + + async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: + self.updates = (*self.updates, (where, data)) + return len(self.rows) + + class TestCheckResponsesCost: """Test suite for CheckResponsesCost class""" @@ -387,16 +402,12 @@ class TestCheckResponsesCost: mock_llm_router.get_deployment.return_value = {"model_id": "deployment-404"} mock_llm_router.aget_responses = AsyncMock( side_effect=litellm.NotFoundError( - message="Response not found", model="gpt-5", llm_provider="openai" + message="Response with id 'resp_upstream_404' not found.", model="gpt-5", llm_provider="openai" ) ) - 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 - ) + table = _RecordingManagedObjectTable(rows=(mock_job,)) + mock_prisma_client.db.litellm_managedobjecttable = table with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget: await check_responses_cost_instance.check_responses_cost() @@ -404,10 +415,7 @@ class TestCheckResponsesCost: mock_sdk_aget.assert_not_called() mock_llm_router.get_deployment.assert_called_once_with(model_id="deployment-404") assert mock_llm_router.aget_responses.call_args.kwargs["response_id"] == encoded_response_id - update_many = mock_prisma_client.db.litellm_managedobjecttable.update_many - update_many.assert_awaited_once() - assert update_many.call_args.kwargs["where"] == {"id": {"in": ["job-404"]}} - assert update_many.call_args.kwargs["data"] == {"status": "stale_expired"} + assert table.updates == (({"id": {"in": ["job-404"]}}, {"status": "stale_expired"}),) @pytest.mark.asyncio async def test_check_responses_cost_marks_mapped_provider_404_stale_expired( @@ -456,21 +464,14 @@ class TestCheckResponsesCost: mock_llm_router.get_deployment.return_value = {"model_id": "deployment-404-mapped"} mock_llm_router.aget_responses = AsyncMock(side_effect=mapped.value) - 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 - ) + table = _RecordingManagedObjectTable(rows=(mock_job,)) + mock_prisma_client.db.litellm_managedobjecttable = table with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget: await check_responses_cost_instance.check_responses_cost() mock_sdk_aget.assert_not_called() - update_many = mock_prisma_client.db.litellm_managedobjecttable.update_many - update_many.assert_awaited_once() - assert update_many.call_args.kwargs["where"] == {"id": {"in": ["job-404-mapped"]}} - assert update_many.call_args.kwargs["data"] == {"status": "stale_expired"} + assert table.updates == (({"id": {"in": ["job-404-mapped"]}}, {"status": "stale_expired"}),) @pytest.mark.asyncio async def test_check_responses_cost_404_without_router_deployment_keeps_row_for_retry( @@ -494,12 +495,8 @@ class TestCheckResponsesCost: mock_llm_router.get_deployment.return_value = None - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + table = _RecordingManagedObjectTable(rows=(mock_job,)) + mock_prisma_client.db.litellm_managedobjecttable = table with patch( "litellm.aget_responses", @@ -512,8 +509,96 @@ class TestCheckResponsesCost: mock_sdk_aget.assert_awaited_once() mock_llm_router.aget_responses.assert_not_called() - update_many = mock_prisma_client.db.litellm_managedobjecttable.update_many - assert update_many.call_args_list == [] + assert table.updates == () + + @pytest.mark.asyncio + async def test_check_responses_cost_404_not_naming_the_response_keeps_row_for_retry( + self, check_responses_cost_instance, mock_prisma_client, mock_llm_router + ): + """A 404 that does not name the response (a gateway or a reconfigured deployment) is retried, not expired.""" + import litellm + from litellm.responses.utils import ResponsesAPIRequestUtils + + encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="azure", + model_id="deployment-renamed", + response_id="resp_upstream_still_there", + ) + + mock_job = MagicMock() + mock_job.unified_object_id = encoded_response_id + mock_job.created_by = "test-user" + mock_job.id = "job-404-other" + mock_job.file_object = {"model": "azure-gpt-5", "id": encoded_response_id} + + mock_llm_router.get_deployment.return_value = {"model_id": "deployment-renamed"} + mock_llm_router.aget_responses = AsyncMock( + side_effect=litellm.NotFoundError( + message="Resource not found", model="gpt-5", llm_provider="azure" + ) + ) + + table = _RecordingManagedObjectTable(rows=(mock_job,)) + mock_prisma_client.db.litellm_managedobjecttable = table + + with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget: + await check_responses_cost_instance.check_responses_cost() + + mock_sdk_aget.assert_not_called() + assert mock_llm_router.aget_responses.call_args.kwargs["response_id"] == encoded_response_id + assert table.updates == () + + @pytest.mark.asyncio + async def test_check_responses_cost_deployment_lookup_error_skips_only_that_job( + self, check_responses_cost_instance, mock_prisma_client, mock_llm_router + ): + """A deployment lookup that raises skips its own job and the cycle still records the next job.""" + from litellm.responses.utils import ResponsesAPIRequestUtils + + broken_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-broken", response_id="resp_upstream_broken" + ) + healthy_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id( + custom_llm_provider="openai", model_id="deployment-healthy", response_id="resp_upstream_healthy" + ) + + broken_job = MagicMock() + broken_job.unified_object_id = broken_response_id + broken_job.created_by = "test-user" + broken_job.id = "job-broken" + broken_job.file_object = {"model": "gpt-5.5", "id": broken_response_id} + + healthy_job = MagicMock() + healthy_job.unified_object_id = healthy_response_id + healthy_job.created_by = "test-user" + healthy_job.id = "job-healthy" + healthy_job.file_object = {"model": "gpt-5.5", "id": healthy_response_id} + + mock_llm_router.get_deployment.side_effect = [ + Exception("Model invalid format - "), + {"model_id": "deployment-healthy"}, + ] + mock_llm_router.aget_responses = AsyncMock( + return_value=ResponsesAPIResponse( + id=healthy_response_id, + object="response", + status="completed", + created_at=int(datetime.now().timestamp()), + output=[], + usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150), + ) + ) + + table = _RecordingManagedObjectTable(rows=(broken_job, healthy_job)) + mock_prisma_client.db.litellm_managedobjecttable = table + + with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget: + await check_responses_cost_instance.check_responses_cost() + + mock_sdk_aget.assert_not_called() + mock_llm_router.aget_responses.assert_awaited_once() + assert mock_llm_router.aget_responses.call_args.kwargs["response_id"] == healthy_response_id + assert table.updates == (({"id": {"in": ["job-healthy"]}}, {"status": "completed"}),) @pytest.mark.asyncio async def test_check_responses_cost_non_404_error_keeps_row_for_retry( @@ -542,20 +627,15 @@ 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=0 - ) + table = _RecordingManagedObjectTable(rows=(mock_job,)) + mock_prisma_client.db.litellm_managedobjecttable = table with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget: await check_responses_cost_instance.check_responses_cost() mock_sdk_aget.assert_not_called() mock_llm_router.aget_responses.assert_awaited_once() - update_many = mock_prisma_client.db.litellm_managedobjecttable.update_many - assert update_many.call_args_list == [] + assert table.updates == () @pytest.mark.asyncio async def test_check_responses_cost_multiple_jobs(