diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3b3c6a6ce29..c9e70b7db73 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1612,6 +1612,35 @@ class Logging(LiteLLMLoggingBaseClass): **kwargs, ) + async def dispatch_failure_handlers( + self, + exception: Exception, + traceback_exception: str, + prefer_async_handlers: bool = False, + ) -> None: + """Route failure logging to async and/or sync handlers for this request. + + Mirrors ``dispatch_success_handlers``: the sync ``failure_handler`` never runs + concurrently with ``async_failure_handler`` on the shared logging object, so the + two paths cannot mutate it at the same time. ``prefer_async_handlers`` only + bypasses the sync-SDK-only shortcut (e.g. ``async for`` on a stream from + ``completion()``); legacy string callbacks still run via + ``executor.submit(failure_handler)`` when configured. + """ + litellm_params = self.model_call_details.get("litellm_params", {}) or {} + sync_sdk = self._is_sync_litellm_request(litellm_params) + passthrough = self.call_type == CallTypes.pass_through.value + if sync_sdk and not prefer_async_handlers and not passthrough: + self.failure_handler(exception, traceback_exception) + return + + await self.async_failure_handler(exception, traceback_exception) + + if not self._should_run_sync_failure_callbacks_for_async_calls(): + return + + executor.submit(self.failure_handler, exception, traceback_exception) + def should_run_logging( self, event_type: Literal["async_success", "sync_success", "async_failure", "sync_failure"], @@ -3076,6 +3105,24 @@ class Logging(LiteLLMLoggingBaseClass): _filtered_success_callbacks = self._remove_internal_litellm_callbacks(_filtered_success_callbacks) return len(_filtered_success_callbacks) > 0 + def _should_run_sync_failure_callbacks_for_async_calls(self) -> bool: + """ + Returns: + - bool: True if sync failure callbacks should be run for async calls. eg. `langfuse`, `s3` + + Mirrors ``_should_run_sync_callbacks_for_async_calls`` but reads the failure + callback lists. Gating the legacy sync ``failure_handler`` on the success lists + would drop sync failure callbacks for any caller that configures only failure + callbacks, so streaming errors would be logged nowhere. + """ + _combined_sync_callbacks = self.get_combined_callback_list( + dynamic_success_callbacks=self.dynamic_failure_callbacks, + global_callbacks=litellm.failure_callback, + ) + _filtered_failure_callbacks = self._remove_internal_custom_logger_callbacks(_combined_sync_callbacks) + _filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks) + return len(_filtered_failure_callbacks) > 0 + def get_combined_callback_list(self, dynamic_success_callbacks: Optional[List], global_callbacks: List) -> List: if dynamic_success_callbacks is None: return list(global_callbacks) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f518cbaadea..60dbf7c644a 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2008,12 +2008,9 @@ class CustomStreamWrapper: if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING - threading.Thread( - target=self.logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() # log response - # Handle any exceptions that might occur during streaming - asyncio.create_task(self.logging_obj.async_failure_handler(e, traceback_exception)) + asyncio.create_task( + self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) + ) self._handle_stream_fallback_error(e) except (httpx.ReadError, httpx.RemoteProtocolError) as e: if self.received_finish_reason is None: @@ -2122,13 +2119,8 @@ class CustomStreamWrapper: if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING - threading.Thread( - target=self.logging_obj.failure_handler, - args=(e, traceback_exception), - ).start() # log response - # Handle any exceptions that might occur during streaming asyncio.create_task( - self.logging_obj.async_failure_handler(e, traceback_exception) # type: ignore + self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) ) self._handle_stream_fallback_error(e) diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 5bffda126fe..edc257f4c3f 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -1111,6 +1111,171 @@ async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through litellm._async_success_callback = original_async_callbacks +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_prefer_async_does_not_submit_sync_handler( + logging_obj, +): + """prefer_async_handlers must await async_failure_handler and never submit the sync failure_handler. + + Submitting the sync ``failure_handler`` while awaiting ``async_failure_handler`` + lets both mutate the shared logging_obj at once, which is the concurrent-mutation + crash this dispatch guard exists to prevent. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object( + logging_obj, "async_failure_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch.object( + logging_obj, + "_should_run_sync_failure_callbacks_for_async_calls", + return_value=False, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + + mock_async.assert_awaited_once_with(exception, traceback_exception) + mock_sync.assert_not_called() + mock_submit.assert_not_called() + + +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_async_completes_before_sync_submit( + logging_obj, +): + """The async failure handler must fully finish before the legacy sync handler is scheduled. + + Ordering proves there is no window where both handlers touch the shared + logging_obj concurrently: the sync submit only happens after the await returns. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + events: list[str] = [] + + async def _async_failure(exc, tb, **kwargs): + events.append("async_start") + await asyncio.sleep(0) + events.append("async_end") + + def _submit(*args, **kwargs): + events.append("sync_submit") + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object(logging_obj, "async_failure_handler", side_effect=_async_failure), + patch.object(logging_obj, "failure_handler", new_callable=MagicMock), + patch.object( + logging_obj, + "_should_run_sync_failure_callbacks_for_async_calls", + return_value=True, + ), + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit", + side_effect=_submit, + ), + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + + assert events == ["async_start", "async_end", "sync_submit"] + + +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_submits_sync_handler_for_failure_only_callbacks( + logging_obj, +): + """A sync failure callback must still run when only failure callbacks are configured. + + The legacy thread-based path always submitted the sync failure_handler, so gating it on + the success callback list would silently drop failure logging for any deployment that + registers only failure callbacks and no success callbacks. This drives the real predicate + (unmocked), so gating the sync failure handler on the success list fails this test. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + + def _sync_failure_callback(*args, **kwargs): + return None + + logging_obj.model_call_details["litellm_params"] = {} + logging_obj.dynamic_success_callbacks = None + logging_obj.dynamic_failure_callbacks = None + + with ( + patch.object(litellm, "success_callback", []), + patch.object(litellm, "failure_callback", [_sync_failure_callback]), + patch.object(logging_obj, "async_failure_handler", new_callable=AsyncMock), + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=True, + ) + + mock_submit.assert_called_once_with(mock_sync, exception, traceback_exception) + + +@pytest.mark.asyncio +async def test_dispatch_failure_handlers_sync_sdk_shortcut_runs_sync_handler_inline( + logging_obj, +): + """A sync-SDK request (prefer_async_handlers=False) runs failure_handler inline. + + ``async for`` over a stream from ``completion()`` passes prefer_async_handlers=True; a + plain sync request leaves it False, so the legacy sync handler runs directly and the + async handler is never awaited, matching dispatch_success_handlers. + """ + exception = ValueError("boom") + traceback_exception = "traceback" + + logging_obj.model_call_details["litellm_params"] = {} + + with ( + patch.object( + logging_obj, "async_failure_handler", new_callable=AsyncMock + ) as mock_async, + patch.object( + logging_obj, "failure_handler", new_callable=MagicMock + ) as mock_sync, + patch( + "litellm.litellm_core_utils.litellm_logging.executor.submit" + ) as mock_submit, + ): + await logging_obj.dispatch_failure_handlers( + exception, + traceback_exception, + prefer_async_handlers=False, + ) + + mock_sync.assert_called_once_with(exception, traceback_exception) + mock_async.assert_not_awaited() + mock_submit.assert_not_called() + + def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj): """Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False.""" import datetime diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 95e8e6561f1..3723976191b 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2087,6 +2087,9 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped(): async def async_failure_handler(self, *args, **kwargs): return None + async def dispatch_failure_handlers(self, *args, **kwargs): + return None + async def failing_make_call(client=None, **kwargs): raise VertexAIError(status_code=400, message="bad input", headers={})