Fix Rubrik batch drain and SSE recovery mutation

litellm/completion_extras/litellm_responses_transformation/transformation.py

@@ -670,9 +670,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
         except (TypeError, ValueError):
             output_index = len(recovered_text_only_items)
 
-        item = recovered_output_items.get(
+        if output_index in recovered_output_items:
-            output_index
+            return
-        ) or recovered_text_only_items.get(output_index)
+
+        item = recovered_text_only_items.get(output_index)
         if item is None:
             item = {
                 "type": "message",

litellm/integrations/rubrik.py

@@ -384,23 +384,22 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
         if not self.log_queue:
             return
 
-        await self._log_batch_to_rubrik(
+        log_queue_snapshot = list(self.log_queue)
-            data=list(self.log_queue),
+        verbose_logger.debug(
+            "Rubrik: Flushing batch of %s events", len(log_queue_snapshot)
         )
+        await self._log_batch_to_rubrik(
+            data=log_queue_snapshot,
+        )
+        del self.log_queue[: len(log_queue_snapshot)]
+        self.last_flush_time = time.time()
 
     async def flush_queue(self):
         if self.flush_lock is None:
             return
 
         async with self.flush_lock:
-            if self.log_queue:
+            await self.async_send_batch()
-                log_queue_snapshot = list(self.log_queue)
-                verbose_logger.debug(
-                    "Rubrik: Flushing batch of %s events", len(log_queue_snapshot)
-                )
-                await self._log_batch_to_rubrik(data=log_queue_snapshot)
-                del self.log_queue[: len(log_queue_snapshot)]
-                self.last_flush_time = time.time()
 
     # -- Tool blocking service -------------------------------------------------
 

tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py

@@ -705,6 +705,44 @@ def test_transform_response_recovers_output_item_done_from_whitespace_padded_raw
     assert result.choices[0].message.content == "Recovered from padded output item"
 
 
+def test_transform_response_preserves_output_item_when_text_done_arrives_later():
+    from litellm.completion_extras.litellm_responses_transformation.transformation import (
+        LiteLLMResponsesTransformationHandler,
+    )
+
+    handler = LiteLLMResponsesTransformationHandler()
+
+    raw_sse = "\n".join(
+        [
+            'data: {"type":"response.output_item.done","output_index":0,"item":{"type":"message","id":"msg_from_item","role":"assistant","status":"completed","content":[{"type":"output_text","text":"Complete output item text","annotations":[]}]}}',
+            'data: {"type":"response.output_text.done","output_index":0,"content_index":0,"item_id":"msg_from_stream","text":"Late text event"}',
+            'data: {"type":"response.completed","response":{"id":"resp_from_stream","object":"response","created_at":1760144904,"status":"completed","model":"gpt-5.4","output":[]}}',
+            "data: [DONE]",
+            "",
+        ]
+    )
+
+    raw_response = _make_empty_responses_api_response()
+    model_response = _make_empty_model_response()
+    logging_obj = Mock()
+    logging_obj.model_call_details = {"original_response": raw_sse}
+
+    result = handler.transform_response(
+        model="gpt-5.4",
+        raw_response=raw_response,
+        model_response=model_response,
+        logging_obj=logging_obj,
+        request_data={"model": "gpt-5.4"},
+        messages=[{"role": "user", "content": "Reply with exactly: ok"}],
+        optional_params={},
+        litellm_params={},
+        encoding=Mock(),
+    )
+
+    assert len(result.choices) == 1
+    assert result.choices[0].message.content == "Complete output item text"
+
+
 def test_transform_response_prefers_completed_output_from_raw_sse():
     from litellm.completion_extras.litellm_responses_transformation.transformation import (
         LiteLLMResponsesTransformationHandler,

tests/test_litellm/integrations/test_rubrik.py

@@ -236,6 +236,22 @@ class TestBatchLogging:
 
         assert handler.log_queue == [{"msg": "c"}]
 
+    async def test_async_send_batch_drains_sent_events(self, handler):
+        handler.log_queue = [{"msg": "a"}, {"msg": "b"}]
+
+        async def mock_post(*_args, **_kwargs):
+            handler.log_queue.append({"msg": "c"})
+            mock_response = Mock()
+            mock_response.raise_for_status = Mock()
+            return mock_response
+
+        handler.async_httpx_client = AsyncMock()
+        handler.async_httpx_client.post = mock_post
+
+        await handler.async_send_batch()
+
+        assert handler.log_queue == [{"msg": "c"}]
+
     async def test_log_batch_error_does_not_crash(self, handler):
         handler.log_queue = [{"msg": "a"}]
         mock_response = Mock()