diff --git a/litellm/integrations/custom_batch_logger.py b/litellm/integrations/custom_batch_logger.py index 0549c3aeaaf..abaacffc7f7 100644 --- a/litellm/integrations/custom_batch_logger.py +++ b/litellm/integrations/custom_batch_logger.py @@ -53,7 +53,20 @@ class CustomBatchLogger(CustomLogger): verbose_logger.debug( "CustomLogger: Flushing batch of %s events", len(self.log_queue) ) - await self.async_send_batch() + try: + await self.async_send_batch() + except Exception: + # If the underlying batch send raised, do NOT drop the + # in-flight events. They will be retried on the next flush. + # Most existing async_send_batch implementations swallow + # their own errors, so this only affects loggers that opt + # in to surfacing failures (e.g. Rubrik). + verbose_logger.exception( + "CustomLogger: async_send_batch raised; preserving " + "%s events in queue for retry", + log_queue_length, + ) + return if self.preserve_events_added_during_flush: del self.log_queue[:log_queue_length] else: diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index ff2d811e74a..607653d91f2 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -367,6 +367,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Batch logging --------------------------------------------------------- async def _log_batch_to_rubrik(self, data): + # NOTE: this method intentionally re-raises on failure so the parent + # CustomBatchLogger.flush_queue keeps the unsent events in the queue + # for the next flush attempt instead of silently dropping them. try: response = await self.async_httpx_client.post( url=self.logging_endpoint, @@ -378,8 +381,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): verbose_logger.exception( f"Rubrik HTTP Error: {e.response.status_code} - {e.response.text}" ) + raise except Exception: verbose_logger.exception("Rubrik Layer Error") + raise async def async_send_batch(self): """Handles sending batches of responses to Rubrik.""" diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py index 414ae05d721..13b41a7f660 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/tests/test_litellm/integrations/test_rubrik.py @@ -269,7 +269,11 @@ class TestBatchLogging: assert handler.log_queue == [{"msg": "c"}] - async def test_log_batch_error_does_not_crash(self, handler): + async def test_log_batch_error_does_not_crash_and_preserves_events(self, handler): + """A failed batch send must not crash the caller AND must preserve the + original events in the queue so they can be retried on the next flush. + Previously the events were silently dropped on HTTP 5xx / network errors. + """ handler.log_queue = [{"msg": "a"}] mock_response = Mock() mock_response.status_code = 500 @@ -282,7 +286,39 @@ class TestBatchLogging: handler.async_httpx_client = AsyncMock() handler.async_httpx_client.post = AsyncMock(return_value=mock_response) await handler.flush_queue() - assert len(handler.log_queue) == 0 + assert handler.log_queue == [{"msg": "a"}] + + async def test_log_batch_network_error_preserves_events(self, handler): + """Network/timeout errors must also preserve the in-flight events.""" + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = AsyncMock( + side_effect=httpx.TimeoutException("timeout") + ) + await handler.flush_queue() + assert handler.log_queue == [{"msg": "a"}, {"msg": "b"}] + + async def test_log_batch_failure_preserves_events_added_during_send(self, handler): + """Failure must preserve both the snapshot AND events appended mid-flush.""" + handler.log_queue = [{"msg": "a"}, {"msg": "b"}] + + async def mock_post(*_args, **_kwargs): + handler.log_queue.append({"msg": "c"}) + mock_response = Mock() + mock_response.status_code = 500 + mock_response.text = "boom" + mock_response.raise_for_status = Mock( + side_effect=httpx.HTTPStatusError( + "err", request=Mock(), response=mock_response + ) + ) + return mock_response + + handler.async_httpx_client = AsyncMock() + handler.async_httpx_client.post = mock_post + + await handler.flush_queue() + assert handler.log_queue == [{"msg": "a"}, {"msg": "b"}, {"msg": "c"}] async def test_system_prompt_prepended_to_messages(self, handler): kwargs = {