fix(rubrik): preserve events on batch send failure

Previously, _log_batch_to_rubrik swallowed all HTTP errors and exceptions,
and the parent flush_queue unconditionally drained the queue afterwards.
On Rubrik 5xx responses, network errors, or timeouts the in-flight events
were silently dropped without ever being delivered.

- Re-raise from _log_batch_to_rubrik so failures surface to the caller.
- In CustomBatchLogger.flush_queue, catch exceptions from async_send_batch
  and leave the queue intact for retry on the next flush. Existing loggers
  that override flush_queue (e.g. Datadog) or that swallow their own errors
  inside async_send_batch (e.g. Langsmith, GCS, Argilla) are unaffected.
- Tests now assert events are preserved on HTTP errors, network errors,
  and that mid-flush appended events are also preserved on failure.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>

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:

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."""

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 = {