Fix router mode preservation and Rubrik batch flushing

litellm/integrations/custom_batch_logger.py

@@ -14,6 +14,8 @@ from litellm.integrations.custom_logger import CustomLogger
 
 
 class CustomBatchLogger(CustomLogger):
+    preserve_events_added_during_flush = False
+
     def __init__(
         self,
         flush_lock: Optional[asyncio.Lock] = None,
@@ -47,11 +49,15 @@ class CustomBatchLogger(CustomLogger):
 
         async with self.flush_lock:
             if self.log_queue:
+                log_queue_length = len(self.log_queue)
                 verbose_logger.debug(
                     "CustomLogger: Flushing batch of %s events", len(self.log_queue)
                 )
                 await self.async_send_batch()
-                self.log_queue.clear()
+                if self.preserve_events_added_during_flush:
+                    del self.log_queue[:log_queue_length]
+                else:
+                    self.log_queue.clear()
                 self.last_flush_time = time.time()
 
     async def async_send_batch(self, *args, **kwargs):

litellm/integrations/rubrik.py

@@ -48,6 +48,8 @@ class BlockedToolsResult:
 
 
 class RubrikLogger(CustomGuardrail, CustomBatchLogger):
+    preserve_events_added_during_flush = True
+
     def __init__(
         self,
         api_key: str | None = None,
@@ -391,15 +393,6 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
         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:
-            await self.async_send_batch()
 
     # -- Tool blocking service -------------------------------------------------
 

litellm/router.py

@@ -6809,7 +6809,7 @@ class Router:
                 # Multiple aliases can point at the same provider/model backend,
                 # but their deployment-level overrides should not downgrade the
                 # backend from responses -> chat via last-write-wins registration.
-                _shared_model_info.pop("mode", None)
+                _shared_model_info["mode"] = _existing_shared_mode
             litellm.register_model(
                 model_cost={
                     _model_name: _shared_model_info,

tests/test_litellm/integrations/test_rubrik.py

@@ -12,6 +12,7 @@ from unittest.mock import AsyncMock, Mock, patch
 import httpx
 import pytest
 
+from litellm.integrations.custom_batch_logger import CustomBatchLogger
 from litellm.integrations.custom_guardrail import ModifyResponseException
 from litellm.integrations.rubrik import RubrikLogger
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
@@ -236,7 +237,7 @@ class TestBatchLogging:
 
         assert handler.log_queue == [{"msg": "c"}]
 
-    async def test_async_send_batch_drains_sent_events(self, handler):
+    async def test_async_send_batch_does_not_drain_events(self, handler):
         handler.log_queue = [{"msg": "a"}, {"msg": "b"}]
 
         async def mock_post(*_args, **_kwargs):
@@ -250,6 +251,22 @@ class TestBatchLogging:
 
         await handler.async_send_batch()
 
+        assert handler.log_queue == [{"msg": "a"}, {"msg": "b"}, {"msg": "c"}]
+
+    async def test_parent_flush_queue_preserves_events_added_during_send(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 CustomBatchLogger.flush_queue(handler)
+
         assert handler.log_queue == [{"msg": "c"}]
 
     async def test_log_batch_error_does_not_crash(self, handler):

tests/test_litellm/test_router_model_cost_isolation.py

@@ -335,6 +335,7 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override():
             )
 
         assert litellm.model_cost[backend_model]["mode"] == "responses"
+        assert "mode" in litellm.model_cost[backend_model]
 
         bridge_model_info, bridge_model = responses_api_bridge_check(
             model="gpt-5.4",