From f3635c86bfbc8baa58e322c90eea068a89933c43 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Tue, 5 May 2026 23:52:34 +0000
Subject: [PATCH] fix(rubrik): override flush_queue so a single snapshot drives
 send and drain

Previously RubrikLogger relied on CustomBatchLogger.flush_queue, which
captured len(self.log_queue) separately from the snapshot taken inside
async_send_batch. Although both happen without an intervening await today
(so they agree in practice), they are semantically disconnected: a future
refactor that adds an await between the two captures, or that changes the
async_send_batch contract, could cause the parent to delete a different
number of items than were actually sent and trigger duplicate deliveries
to Rubrik.

Override flush_queue on RubrikLogger so a single snapshot drives both the
HTTP POST and the queue truncation. async_send_batch is preserved for
direct callers/tests but no longer participates in the canonical flush
path. Existing tests (including the one that explicitly invokes the base
CustomBatchLogger.flush_queue path) still pass.

Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
---
 litellm/integrations/rubrik.py | 35 +++++++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py
index 607653d91f2..7ffd89b4802 100644
--- a/litellm/integrations/rubrik.py
+++ b/litellm/integrations/rubrik.py
@@ -387,7 +387,13 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
             raise
 
     async def async_send_batch(self):
-        """Handles sending batches of responses to Rubrik."""
+        """Handles sending batches of responses to Rubrik.
+
+        Note: the canonical flush path is :meth:`flush_queue`, which takes a
+        single snapshot used for both sending and queue draining. This method
+        is kept for direct callers / tests; it intentionally does NOT remove
+        events from the queue.
+        """
         if not self.log_queue:
             return
 
@@ -399,6 +405,33 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
             data=log_queue_snapshot,
         )
 
+    async def flush_queue(self):
+        """Snapshot, send, and drain in one consistent step.
+
+        Overrides the base implementation so the same snapshot drives both
+        the HTTP send and the queue truncation. This avoids the subtle
+        coupling where the base class captures `len(self.log_queue)`
+        separately from the snapshot taken inside `async_send_batch`,
+        which could otherwise drift in a future refactor and cause
+        duplicate deliveries to Rubrik.
+        """
+        if self.flush_lock is None:
+            return
+
+        async with self.flush_lock:
+            if not self.log_queue:
+                return
+            snapshot = list(self.log_queue)
+            verbose_logger.debug("Rubrik: Flushing batch of %s events", len(snapshot))
+            try:
+                await self._log_batch_to_rubrik(data=snapshot)
+            except Exception:
+                # Already logged with traceback inside _log_batch_to_rubrik.
+                # Preserve the in-flight events for retry on the next flush.
+                return
+            del self.log_queue[: len(snapshot)]
+            self.last_flush_time = time.time()
+
     # -- Tool blocking service -------------------------------------------------
 
     async def _post_to_tool_blocking_service(