From 9ee908089b5edcfa7654ddacd4bef68073583610 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 20 May 2026 22:04:16 +0000 Subject: [PATCH] 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 --- litellm/integrations/rubrik.py | 2 -- litellm/llms/vertex_ai/vertex_llm_base.py | 16 +++++++++--- litellm/router.py | 25 ++++++++++++------- .../test_litellm/integrations/test_rubrik.py | 17 ------------- 4 files changed, 28 insertions(+), 32 deletions(-) diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 8477523a2a7..eb68fa2db7d 100644 --- a/litellm/integrations/rubrik.py +++ b/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, diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 5a3695a3536..efa27c5d477 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/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", diff --git a/litellm/router.py b/litellm/router.py index ba2abe83dc8..265e51cca67 100644 --- a/litellm/router.py +++ b/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 " diff --git a/tests/test_litellm/integrations/test_rubrik.py b/tests/test_litellm/integrations/test_rubrik.py index b681a8ed1ff..596f939d0c2 100644 --- a/tests/test_litellm/integrations/test_rubrik.py +++ b/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.