fix: address three bug detection findings

- vertex_llm_base: guard background refresh's cache write with an
  identity check so a stale write cannot overwrite a credentials
  reference replaced by a concurrent reauthentication path.
- router: make shared backend mode preservation directional - only
  preserve when an existing 'responses' mode would be downgraded to
  'chat', or when the deployment mode is None (which would otherwise
  clear the existing mode). Legitimate upgrades now apply.
- rubrik: remove unused preserve_events_added_during_flush attribute;
  RubrikLogger overrides flush_queue, so the base-class flag never
  applied. Drop the test that exercised the parent path on a Rubrik
  instance since it does not reflect real flush behavior.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

litellm/integrations/rubrik.py

@@ -39,8 +39,6 @@ _WEBHOOK_PATH_LOGGING_BATCH = "/v1/litellm/batch"
 
 
 class RubrikLogger(CustomGuardrail, CustomBatchLogger):
-    preserve_events_added_during_flush = True
-
     def __init__(
         self,
         api_key: str | None = None,

litellm/llms/vertex_ai/vertex_llm_base.py

@@ -507,10 +507,18 @@ class VertexBase:
         try:
             verbose_logger.debug("Background proactive credential refresh")
             await asyncify(self.refresh_auth)(credentials)
-            self._credentials_project_mapping[credential_cache_key] = (
-                credentials,
-                credential_project_id,
-            )
+            # Only update the cache if it still points at the credentials
+            # object we just refreshed. The per-key async lock is not held
+            # here, so a concurrent INVALID path may have already replaced
+            # this entry (e.g. via _handle_reauthentication_async, which
+            # creates a fresh credentials object). In that case our write
+            # would clobber the newer entry with a stale reference.
+            cached_creds, _ = self._unpack_cached_credentials(credential_cache_key)
+            if cached_creds is credentials:
+                self._credentials_project_mapping[credential_cache_key] = (
+                    credentials,
+                    credential_project_id,
+                )
         except Exception:
             verbose_logger.debug(
                 "Background credential refresh failed, will retry on next request",

litellm/router.py

@@ -7323,15 +7323,22 @@ class Router:
             _existing_shared_mode = (
                 cast(Optional[dict], litellm.model_cost.get(_model_name, {})) or {}
             ).get("mode")
-            if (
-                _existing_shared_mode is not None
-                and _shared_model_info.get("mode") != _existing_shared_mode
-            ):
-                # Keep the built-in bridge mode stable for shared backend keys.
-                # 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.
-                _deployment_mode = _shared_model_info.get("mode")
+            _deployment_mode = _shared_model_info.get("mode")
+            # Keep the built-in bridge mode stable for shared backend keys.
+            # 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.
+            # Only preserve in that specific direction so legitimate upgrades
+            # (e.g. chat -> responses) and unrelated mode changes still apply,
+            # and so a missing deployment mode does not silently clear the
+            # existing shared backend mode.
+            _is_responses_to_chat_downgrade = (
+                _existing_shared_mode == "responses" and _deployment_mode == "chat"
+            )
+            _would_clear_existing_mode = (
+                _existing_shared_mode is not None and _deployment_mode is None
+            )
+            if _is_responses_to_chat_downgrade or _would_clear_existing_mode:
                 if _deployment_mode is not None:
                     verbose_router_logger.warning(
                         "Router: preserving existing mode=%s for shared backend "

tests/test_litellm/integrations/test_rubrik.py

@@ -12,7 +12,6 @@ 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
@@ -282,22 +281,6 @@ class TestBatchLogging:
 
         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_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.