diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index b57ff91f650..6f8ef830ee6 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -129,7 +129,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports _check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_and_merge_model_level_guardrails, ) -from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect +from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task from litellm.router import Router from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.router_utils.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX @@ -1991,9 +1991,15 @@ class _ClientStreamFailureReport: self.logging_obj.on_detached_stream_failure = self.report if self.held is None: return + await asyncio.shield( + spawn_logging_task( + asyncio.get_running_loop(), self._report_held(self.held), task_name="stream failure report" + ) + ) + + async def _report_held(self, failure: Exception) -> None: try: - with anyio.CancelScope(shield=True): - await self.report(self.held) + await self.report(failure) except Exception as hook_failure: # noqa: BLE001 # a failing failure hook must not skip the stream cleanup verbose_proxy_logger.warning( "Stream failure hook raised: %s(%s)", type(hook_failure).__name__, sanitize_for_log(hook_failure) @@ -4549,16 +4555,18 @@ class ProxyBaseLLMRequestProcessing: seal: Final = "" if seal_open_frame is None else seal_open_frame(recent_tail) yield seal + error_frame if seal else error_frame finally: - await release_stream_failure() - await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( - request=request, - request_data=request_data, - response=response, - stream_completed=stream_completed, - client_disconnected=client_disconnected, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, - ) + try: + await release_stream_failure() + finally: + await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup( + request=request, + request_data=request_data, + response=response, + stream_completed=stream_completed, + client_disconnected=client_disconnected, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) @staticmethod def async_sse_data_generator( diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 8ec2d2563d8..22a97a07d3b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -41,7 +41,7 @@ from litellm.exceptions import MidStreamFallbackError, RateLimitError from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base -from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation +from litellm.litellm_core_utils.internal_call_metadata import requests_background, without_budget_reservation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers @@ -974,7 +974,7 @@ class BaseResponsesAPIStreamingIterator: generated_text: Final = self._generated_content + self._generated_tool_arguments if ( not generated_text - or self.request_data.get("background") is True + or requests_background(self.request_data) or self.logging_obj.model_call_details.get("combined_usage_object") is not None ): return @@ -990,7 +990,12 @@ class BaseResponsesAPIStreamingIterator: usage: Final = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the success path above uses estimate ) - self.logging_obj.record_assembled_response_for_failure(ModelResponse(model=self.model, usage=usage)) + served_tier: Final[object] = getattr(self._created_response, "service_tier", None) + self.logging_obj.record_assembled_response_for_failure( + ModelResponse( + model=self.model, usage=usage, service_tier=served_tier if isinstance(served_tier, str) else None + ) + ) def _handle_failure(self, exception: Exception): """ @@ -1068,9 +1073,8 @@ class BaseResponsesAPIStreamingIterator: Its response is the response.created payload (it carries the served service_tier) plus the text and tool arguments relayed so far, with counted usage since the provider never sent any. None when nothing was relayed yet, - the stream already reached its terminal event and billed itself, its failure - log already billed the relayed output, or it runs in the background, where - retrieving it bills the whole run. + the stream already reached its terminal event and billed itself, or its failure + log already billed the relayed output. """ created: Final = self._created_response generated: Final = self._generated_content + self._generated_tool_arguments @@ -1096,9 +1100,7 @@ class BaseResponsesAPIStreamingIterator: return ResponseIncompleteEvent(type=ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, response=partial) def _billed_without_disconnect(self) -> bool: - return ( - self.completed_response is not None or self._failure_handled or self.request_data.get("background") is True - ) + return self.completed_response is not None or self._failure_handled def _note_yielded_event(self, event: ResponsesAPIStreamingResponse) -> None: self._yielded_first_chunk = True diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 4bca44dbe3c..ebef1fbbef7 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -10676,12 +10676,19 @@ class TestResolveLitellmCallId: class _FailureHookRecorder: """Stands in for ProxyLogging.post_call_failure_hook, recording what the detached-failure closure hands it.""" - def __init__(self, raises: Optional[Exception] = None): + def __init__(self, raises: Optional[Exception] = None, gate: Optional[asyncio.Event] = None): self.calls = [] + self.entered = asyncio.Event() + self.reported = asyncio.Event() self._raises = raises + self._gate = gate async def post_call_failure_hook(self, **kwargs): + self.entered.set() + if self._gate is not None: + await self._gate.wait() self.calls.append(kwargs) + self.reported.set() if self._raises is not None: raise self._raises @@ -10782,7 +10789,7 @@ class TestDetachedStreamFailureHook: assert [call["original_exception"] for call in recorder.calls] == [failure] @staticmethod - def _armed_stream(recorder: _FailureHookRecorder, iterator_hook_for): + def _armed_stream(recorder: _FailureHookRecorder, iterator_hook_for, response: Optional[MagicMock] = None): from litellm.proxy._types import UserAPIKeyAuth as RealUserAPIKeyAuth logging_obj = LiteLLMLoggingObj( @@ -10806,7 +10813,7 @@ class TestDetachedStreamFailureHook: proxy_logging_obj=proxy_logging_obj, ) stream = ProxyBaseLLMRequestProcessing.async_streaming_data_generator( - response=MagicMock(), + response=MagicMock() if response is None else response, user_api_key_dict=user_api_key_dict, request_data=processing_obj.data, proxy_logging_obj=proxy_logging_obj, @@ -10875,6 +10882,30 @@ class TestDetachedStreamFailureHook: assert frames == ["data: chunk\n\n"] * 3 assert [call["original_exception"] for call in recorder.calls] == [upstream_failure] + @pytest.mark.asyncio + async def test_failure_report_and_stream_cleanup_survive_the_client_task_being_cancelled_mid_report(self): + gate = asyncio.Event() + recorder = _FailureHookRecorder(gate=gate) + upstream_failure = RuntimeError("upstream Responses stream failed") + response = MagicMock() + response.aclose = AsyncMock() + + _, stream = self._armed_stream(recorder, self._hook_reporting(upstream_failure), response=response) + + async def client_reads_the_whole_stream(): + return [frame async for frame in stream] + + client = asyncio.create_task(client_reads_the_whole_stream()) + await asyncio.wait_for(recorder.entered.wait(), timeout=5) + client.cancel() + with pytest.raises(asyncio.CancelledError): + await client + + response.aclose.assert_awaited_once() + gate.set() + await asyncio.wait_for(recorder.reported.wait(), timeout=5) + assert [call["original_exception"] for call in recorder.calls] == [upstream_failure] + @pytest.mark.asyncio async def test_failure_reported_by_the_stream_reaches_the_failure_hook_once_when_the_client_leaves_early(self): recorder = _FailureHookRecorder() diff --git a/tests/unit/responses/test_responses_streaming_iterator.py b/tests/unit/responses/test_responses_streaming_iterator.py index 9f64540d6ba..c572b4151b2 100644 --- a/tests/unit/responses/test_responses_streaming_iterator.py +++ b/tests/unit/responses/test_responses_streaming_iterator.py @@ -270,7 +270,7 @@ async def _failure_payload_after_mid_stream_timeout( monkeypatch: pytest.MonkeyPatch, request_input: str | list[dict[str, object]], upstream: _UpstreamThatTimesOutAfterThreeDeltas | None = None, - background: bool = False, + background_fields: Mapping[str, object] | None = None, ) -> tuple[FailureRecorder, StandardLoggingPayload]: upstream_response: Final = upstream or _UpstreamThatTimesOutAfterThreeDeltas() recorder = FailureRecorder() @@ -290,7 +290,7 @@ async def _failure_payload_after_mid_stream_timeout( responses_api_provider_config=OpenAIResponsesAPIConfig(), logging_obj=logging_obj, custom_llm_provider="openai", - request_data={"model": "gpt-5.4-nano", "input": request_input, "stream": True, "background": background}, + request_data={"model": "gpt-5.4-nano", "input": request_input, "stream": True, **(background_fields or {})}, ) received = _ChunkRecorder() @@ -336,11 +336,18 @@ async def test_mid_stream_read_timeout_bills_the_tool_call_arguments_already_str assert payload["response_cost"] > 0 +@pytest.mark.parametrize( + "background_fields", + [{"background": True}, {"extra_body": {"background": True}}], + ids=["top_level", "extra_body"], +) @pytest.mark.asyncio async def test_mid_stream_failure_of_a_background_response_bills_nothing_since_retrieving_it_bills_the_whole_run( - monkeypatch: pytest.MonkeyPatch, + monkeypatch: pytest.MonkeyPatch, background_fields: Mapping[str, object] ) -> None: - _, payload = await _failure_payload_after_mid_stream_timeout(monkeypatch, _SEA_PROMPT, background=True) + _, payload = await _failure_payload_after_mid_stream_timeout( + monkeypatch, _SEA_PROMPT, background_fields=background_fields + ) assert payload["status"] == "failure" assert payload["completion_tokens"] == 0 @@ -462,6 +469,55 @@ class _UpstreamThatTimesOutAfterRelayingText: raise httpx.ReadTimeout("Timeout on reading data from socket") +class _UpstreamServedAtTierThatTimesOut: + headers: dict[str, str] = {} + + def __init__(self, served_tier: str) -> None: + self.served_tier = served_tier + + async def aiter_bytes(self) -> AsyncIterator[bytes]: + created: Final = { + **_CREATED_WITH_PRIORITY_TIER, + "response": {**_CREATED_RESPONSE_WITH_PRIORITY_TIER, "service_tier": self.served_tier}, + } + for event in (created, _text_delta("Once upon a time", 1)): + yield f"data: {json.dumps(event)}\n\n".encode() + raise httpx.ReadTimeout("Timeout on reading data from socket") + + +async def _mid_stream_failure_cost( + monkeypatch: pytest.MonkeyPatch, requested_tier: str | None, served_tier: str +) -> float: + recorder: Final = FailureRecorder() + monkeypatch.setattr(litellm, "failure_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_failure_callback", [recorder]) + logging_obj: Final = _make_logging_obj() + logging_obj.update_environment_variables( + model="gpt-5.4-nano", + optional_params={"stream": True} if requested_tier is None else {"stream": True, "service_tier": requested_tier}, + litellm_params={"custom_llm_provider": "openai", "aresponses": True}, + ) + iterator: Final = _live_iterator(logging_obj, upstream=_UpstreamServedAtTierThatTimesOut(served_tier)) + + with pytest.raises(litellm.Timeout): + await _ChunkRecorder().drain(iterator) + await iterator._await_pending_logging() + + (payload,) = recorder.failure_payloads + return payload["response_cost"] + + +@pytest.mark.asyncio +async def test_mid_stream_failure_bills_the_relayed_text_at_the_tier_the_provider_served( + monkeypatch: pytest.MonkeyPatch, +) -> None: + flex_requested_default_served: Final = await _mid_stream_failure_cost(monkeypatch, "flex", "default") + nothing_requested_default_served: Final = await _mid_stream_failure_cost(monkeypatch, None, "default") + + assert nothing_requested_default_served > 0 + assert flex_requested_default_served == nothing_requested_default_served + + def test_partial_response_on_disconnect_bills_the_relayed_text_at_the_served_service_tier(): iterator: Final = _live_iterator(_make_logging_obj()) for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time ", 1), _text_delta("a fox slept.", 2)): @@ -516,9 +572,12 @@ async def test_partial_response_on_disconnect_is_none_once_the_failure_log_bille await iterator._await_pending_logging() -def test_partial_response_on_disconnect_is_none_for_a_background_response(): +def test_partial_response_on_disconnect_reports_the_relayed_text_of_a_background_response(): iterator: Final = _live_iterator(_make_logging_obj(), background=True) for event in (_CREATED_WITH_PRIORITY_TIER, _text_delta("Once upon a time", 1)): iterator._process_chunk(json.dumps(event)) - assert iterator.partial_response_on_disconnect() is None + incomplete_event: Final = iterator.partial_response_on_disconnect() + + assert incomplete_event is not None + assert incomplete_event.response.output[0].content[0].text == "Once upon a time"