diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 7a8820a8785..2a939decb15 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1816,6 +1816,54 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, ) + def _invalidate_vertex_credentials_on_anthropic_http_error( + self, + e: Exception, + anthropic_messages_provider_config: BaseAnthropicMessagesConfig, + litellm_params: GenericLiteLLMParams, + custom_llm_provider: str, + ) -> None: + """ + If Vertex AI returns 401 on Anthropic /messages, clear cached OAuth tokens + on the provider config so a router retry (or the next request) refreshes them. + """ + if custom_llm_provider != "vertex_ai": + return + status: Optional[int] = None + if isinstance(e, httpx.HTTPStatusError) and e.response is not None: + status = e.response.status_code + else: + sc = getattr(e, "status_code", None) + if sc is not None: + try: + status = int(sc) + except (TypeError, ValueError): + status = None + if status != 401: + return + + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + if not isinstance(anthropic_messages_provider_config, VertexBase): + return + + lp: Dict[str, Any] + if isinstance(litellm_params, dict): + lp = dict(litellm_params) + else: + model_dump = getattr(litellm_params, "model_dump", None) + if callable(model_dump): + lp = model_dump(exclude_none=False) + else: + lp = dict(litellm_params) # type: ignore[arg-type] + + creds = VertexBase.safe_get_vertex_ai_credentials(lp) + project = VertexBase.safe_get_vertex_ai_project(lp) + anthropic_messages_provider_config.invalidate_credentials( + credentials=creds, + project_id=project, + ) + async def async_anthropic_messages_handler( self, model: str, @@ -1965,6 +2013,12 @@ class BaseLLMHTTPHandler: ) response.raise_for_status() except Exception as e: + self._invalidate_vertex_credentials_on_anthropic_http_error( + e=e, + anthropic_messages_provider_config=anthropic_messages_provider_config, + litellm_params=litellm_params, + custom_llm_provider=custom_llm_provider, + ) raise self._handle_error( e=e, provider_config=anthropic_messages_provider_config ) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 68d8f0d046d..10cdd4741ce 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -557,6 +557,40 @@ class VertexBase: # Re-raise the original error for better context raise error + def invalidate_credentials( + self, + credentials: Optional[VERTEX_CREDENTIALS_TYPES], + project_id: Optional[str], + ) -> None: + """ + Drop cached OAuth credentials for this (credentials, project_id) pair. + + Used when the Vertex / Google API returns 401 so the next get_access_token() + reloads credentials instead of reusing an access token Google rejects while + google-auth may still consider non-expired (see #23512). + """ + cache_credentials = ( + json.dumps(credentials) if isinstance(credentials, dict) else credentials + ) + key = (cache_credentials, project_id) + if key in self._credentials_project_mapping: + verbose_logger.debug( + "Invalidating cached Vertex credentials for project_id=%s (API 401)", + project_id, + ) + del self._credentials_project_mapping[key] + # get_access_token may register (creds, None) then duplicate under resolved + # project_id; clear the None entry when invalidating a resolved project. + if project_id is not None: + none_key = (cache_credentials, None) + if none_key in self._credentials_project_mapping: + verbose_logger.debug( + "Invalidating cached Vertex credentials for project_id=None " + "(paired with resolved project_id=%s)", + project_id, + ) + del self._credentials_project_mapping[none_key] + def get_access_token( self, credentials: Optional[VERTEX_CREDENTIALS_TYPES], diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index a3512bc6e7b..a9d1faeeb04 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,7 +1,8 @@ import os import sys -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch +import httpx import pytest sys.path.insert( @@ -288,3 +289,104 @@ async def test_async_anthropic_messages_handler_header_priority(): assert captured_headers["X-Forwarded-Only"] == "keep" assert captured_headers["X-Extra-Only"] == "also-keep" assert captured_headers["X-Provider-Only"] == "keep-this-too" + + +def _http_401_error() -> httpx.HTTPStatusError: + req = httpx.Request("POST", "https://aiplatform.googleapis.com/v1/fake") + resp = httpx.Response(401, request=req, json={"error": {"code": 401}}) + return httpx.HTTPStatusError("unauthorized", request=req, response=resp) + + +def test_invalidate_vertex_on_anthropic_http_error_skips_non_vertex_provider(): + """401 must not call Vertex cache invalidation when custom_llm_provider is not vertex_ai.""" + handler = BaseLLMHTTPHandler() + mock_config = MagicMock() + handler._invalidate_vertex_credentials_on_anthropic_http_error( + e=_http_401_error(), + anthropic_messages_provider_config=mock_config, + litellm_params=GenericLiteLLMParams(), + custom_llm_provider="anthropic", + ) + mock_config.invalidate_credentials.assert_not_called() + + +def test_invalidate_vertex_on_anthropic_http_error_skips_non_vertex_base_config(): + """Non-VertexBase provider configs are ignored even for vertex_ai (defensive).""" + handler = BaseLLMHTTPHandler() + mock_config = Mock(spec=["invalidate_credentials"]) + handler._invalidate_vertex_credentials_on_anthropic_http_error( + e=_http_401_error(), + anthropic_messages_provider_config=mock_config, + litellm_params=GenericLiteLLMParams( + vertex_project="proj-x", + vertex_credentials={"type": "service_account"}, + ), + custom_llm_provider="vertex_ai", + ) + mock_config.invalidate_credentials.assert_not_called() + + +def test_invalidate_vertex_on_anthropic_http_error_skips_non_401(): + handler = BaseLLMHTTPHandler() + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( + VertexAIPartnerModelsAnthropicMessagesConfig, + ) + + config = VertexAIPartnerModelsAnthropicMessagesConfig() + with patch.object(config, "invalidate_credentials") as mock_inv: + req = httpx.Request("POST", "https://example.com") + resp = httpx.Response(429, request=req) + err = httpx.HTTPStatusError("rate limit", request=req, response=resp) + handler._invalidate_vertex_credentials_on_anthropic_http_error( + e=err, + anthropic_messages_provider_config=config, + litellm_params=GenericLiteLLMParams(vertex_project="p"), + custom_llm_provider="vertex_ai", + ) + mock_inv.assert_not_called() + + +def test_invalidate_vertex_on_anthropic_http_error_calls_invalidate_on_401(): + """Vertex Anthropic + httpx 401 → invalidate_credentials with safe_get params.""" + handler = BaseLLMHTTPHandler() + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( + VertexAIPartnerModelsAnthropicMessagesConfig, + ) + + creds = {"type": "service_account", "project_id": "unit-test-proj"} + config = VertexAIPartnerModelsAnthropicMessagesConfig() + with patch.object(config, "invalidate_credentials") as mock_inv: + handler._invalidate_vertex_credentials_on_anthropic_http_error( + e=_http_401_error(), + anthropic_messages_provider_config=config, + litellm_params=GenericLiteLLMParams( + vertex_project="unit-test-proj", + vertex_location="global", + vertex_credentials=creds, + ), + custom_llm_provider="vertex_ai", + ) + mock_inv.assert_called_once_with( + credentials=creds, + project_id="unit-test-proj", + ) + + +def test_invalidate_vertex_on_anthropic_http_error_401_via_status_code_attr(): + """Some transports attach status_code without being httpx.HTTPStatusError.""" + handler = BaseLLMHTTPHandler() + from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( + VertexAIPartnerModelsAnthropicMessagesConfig, + ) + + config = VertexAIPartnerModelsAnthropicMessagesConfig() + err = Exception("wrapped") + setattr(err, "status_code", 401) + with patch.object(config, "invalidate_credentials") as mock_inv: + handler._invalidate_vertex_credentials_on_anthropic_http_error( + e=err, + anthropic_messages_provider_config=config, + litellm_params=GenericLiteLLMParams(vertex_project="p2"), + custom_llm_provider="vertex_ai", + ) + mock_inv.assert_called_once_with(credentials=None, project_id="p2") diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 78caf4b9778..e5e21e0093e 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -120,6 +120,69 @@ class TestVertexBase: assert token2 == "token-1" assert project2 == "project-1" + def test_invalidate_credentials_clears_project_and_none_keys(self): + """ + After get_access_token(project_id=None), both (creds, None) and (creds, resolved) + may exist; invalidate_credentials(project_id=resolved) must remove both (#23512). + """ + vertex_base = VertexBase() + credentials = {"type": "service_account", "project_id": "resolved-proj"} + cache_creds = json.dumps(credentials) + mock_creds = MagicMock() + key_none = (cache_creds, None) + key_resolved = (cache_creds, "resolved-proj") + vertex_base._credentials_project_mapping[key_none] = ( + mock_creds, + "resolved-proj", + ) + vertex_base._credentials_project_mapping[key_resolved] = ( + mock_creds, + "resolved-proj", + ) + + vertex_base.invalidate_credentials( + credentials=credentials, + project_id="resolved-proj", + ) + + assert key_none not in vertex_base._credentials_project_mapping + assert key_resolved not in vertex_base._credentials_project_mapping + + def test_invalidate_credentials_removes_single_cache_entry(self): + """When only (creds, project_id) exists, invalidate removes that entry.""" + vertex_base = VertexBase() + credentials = {"type": "service_account", "project_id": "p-only"} + cache_creds = json.dumps(credentials) + mock_creds = MagicMock() + key = (cache_creds, "p-only") + vertex_base._credentials_project_mapping[key] = (mock_creds, "p-only") + + vertex_base.invalidate_credentials(credentials=credentials, project_id="p-only") + + assert key not in vertex_base._credentials_project_mapping + + def test_invalidate_credentials_with_none_project_id(self): + """invalidate_credentials(..., project_id=None) only removes the (creds, None) key.""" + vertex_base = VertexBase() + credentials = {"type": "service_account"} + cache_creds = json.dumps(credentials) + mock_creds = MagicMock() + key_none = (cache_creds, None) + key_other = (cache_creds, "other-project") + vertex_base._credentials_project_mapping[key_none] = (mock_creds, "x") + vertex_base._credentials_project_mapping[key_other] = (mock_creds, "x") + + vertex_base.invalidate_credentials(credentials=credentials, project_id=None) + + assert key_none not in vertex_base._credentials_project_mapping + assert key_other in vertex_base._credentials_project_mapping + + def test_invalidate_credentials_noop_when_key_missing(self): + vertex_base = VertexBase() + credentials = {"type": "service_account", "project_id": "missing"} + vertex_base.invalidate_credentials(credentials=credentials, project_id="missing") + assert vertex_base._credentials_project_mapping == {} + @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"]) @pytest.mark.asyncio async def test_credential_refresh(self, is_async):