Repository: BerriAI/litellm
Pull request: Fix VertexAI Credential Caching issue
## Title

Fix VertexAI Credential Caching issue

## Relevant issues

Fixes https://github.com/BerriAI/litellm/issues/7904

## Pre-Submission checklist

**Please complete all items before asking a LiteLLM maintainer to review your PR**

- [x] I have Added testing in the [`tests/litellm/`](https://github.com/BerriAI/litellm/tree/main/tests/litellm) directory, **Adding at least 1 test is a hard requirement** - [see details](https://docs.litellm.ai/docs/extras/contributing_code)
- [ ] I have added a screenshot of my new test passing locally 
- [x] My PR passes all unit tests on (`make test-unit`)[https://docs.litellm.ai/docs/extras/contributing_code]
- [x] My PR's scope is as isolated as possible, it only solves 1 specific problem


## Type

🐛 Bug Fix
🧹 Refactoring
✅ Test

## Changes

- Moves to using a vertex credential map
- Key = combination of credential str + project id 
- Adds more testing to ensure cached credential not sent to wrong project 
- Adds testing to ensure dynamic credentials respected


Problem statement:
[Bug]: Vertex Credentials Are Cached
### What happened?

It seems that when using litellm to make calls to vertex_ai models, the vertex credentials are somehow cached within litellm and the ones passed into `completion()` are not being used on every call

```
>>> litellm.completion(
...     messages=[{'role': 'user', 'content': "Hey, how's it going?"}], 
...     model='vertex_ai/gemini-1.5-flash', 
...     max_tokens=10, 
...     temperature=0.0, 
...     vertex_credentials='sdsdfsdsdfsdfsddsff'
... )

Exception: Unable to load vertex credentials from environment. Got=sdsdfsdsdfsdfsddsff
```

So far so good... but then if I use real working credentials, any subsequent calls with anything passed into `vertex_credentials` still works, when it should not

```
>>> creds = <my actual google json credentials> 
>>> litellm.completion(
...     messages=[{'role': 'user', 'content': "Hey, how's it going?"}], 
...     model='vertex_ai/gemini-1.5-flash', 
...     max_tokens=10, 
...     temperature=0.0, 
...     vertex_credentials=creds
... )

ModelResponse(id=' ................)  # successful response

>>> litellm.completion(
...     messages=[{'role': 'user', 'content': "Hey, how's it going?"}], 
...     model='vertex_ai/gemini-1.5-flash', 
...     max_tokens=10, 
...     temperature=0.0, 
...     vertex_credentials='sdsdfsdsdfsdfsddsff'
... )

ModelResponse(id=' ................)  # successful response when it should not be
```


### Relevant log output

```shell

```

### Are you a ML Ops Team?

Yes

### What LiteLLM version are you on ?

1.47.2

Tried with 1.59.1 as well, same behavior

### Twitter / LinkedIn details

_No response_

Hints from issue discussion:
Thanks for the issue @hakan458 i plan to look into this, this week. 
Hi @krrishdholakia any updates on this? Strange that the key is being cached at all in this case I would expect the `vertex_credentials` argument to take precedence over anything else. If the fix is not obvious, is there perhaps a way I can manually "clear" the credentials that are being cached? 
Seems like in VertexBase , `access_token` is being saved, and once it is, `_ensure_access_token` does nothing with the passed in `vertex_credentials` and just returns `self.access_token` 

Since I want to continue using `litellm.completion` and not use `VertexLLM` directly which has 15 required positional args... I can't think of a way to clear the stored `access_token` from looking at the code. Let me know if you have any ideas. 

This is a blocker for what I am working on at the moment due to this scenario: 

1. We have a web app which multiple users interact with. It uses litellm to make calls to various LLMs
2. User A from organization 1 uses their Vertex credentials successfully
3. User B from organization 2 enters their Vertex credentials, but under the hood litellm is still using creds from User A/Org 1! 
4. OR User B enters garbage credentials like in my bug report, and still under the hood litellm is using creds from User A/Org 1
Hi @hakan458 sorry missed this - hoping to work on this today. 

Appreciate your investigation work on this so far
@krrishdholakia just curious if you've looked into this at all? I might make the change myself. I'm thinking to have _ensure_access_token always fetch the token regardless if its been fetched before. Do you think it should be an optional param? I feel like the expected behavior would be to always use the passed in `vertex_credentials`. Maybe an attempt at optimization was made to not do that every time, but not sure its worth the confusion. 
I made changes and tested locally, I am unable to push a branch to the repo to open a PR, do I need to fork the repo and make changes there? Contribution docs just say to clone and push, but in one step it says to push from fork, unclear 
Forked the repo and we are now using that in our app to workaround this problem - please let me know how to proceed with making a PR in this repo so it can be included in the next release @krrishdholakia 
Ok opened the PR
https://github.com/BerriAI/litellm/pull/8068

Was unable to run tests as described in the PR - relevant packages are not installed via poetry. 
@krrishdholakia still want to merge this PR can you give me some guidance here? We are using forked litellm to have this fix, want to get back to using this repo
Any updates on this issue?
Currently we're unable to use models originating from different Vertex AI projects and credentials due to this.
The first request against LiteLLM determines the credentials to be used and all follow-up requests for models configured with other credentials fail.

Commit message: fix(vertex_llm_base.py): common auth logic across sync + async vertex ai calls

prevents credential caching issue across both flows

Patch to review:
diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py
index 8286cb515f17..df8b1d5ef1b5 100644
--- a/litellm/llms/vertex_ai/vertex_llm_base.py
+++ b/litellm/llms/vertex_ai/vertex_llm_base.py
@@ -6,7 +6,7 @@
 
 import json
 import os
-from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, cast
 
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.asyncify import asyncify
@@ -28,6 +28,10 @@ def __init__(self) -> None:
         self.access_token: Optional[str] = None
         self.refresh_token: Optional[str] = None
         self._credentials: Optional[GoogleCredentialsObject] = None
+        self._credentials_project_mapping: Dict[
+            Tuple[Optional[VERTEX_CREDENTIALS_TYPES], Optional[str]],
+            GoogleCredentialsObject,
+        ] = {}
         self.project_id: Optional[str] = None
         self.async_handler: Optional[AsyncHTTPHandler] = None
 
@@ -128,32 +132,11 @@ def _ensure_access_token(
         """
         if custom_llm_provider == "gemini":
             return "", ""
-        if self.access_token is not None:
-            if project_id is not None:
-                return self.access_token, project_id
-            elif self.project_id is not None:
-                return self.access_token, self.project_id
-
-        if not self._credentials:
-            self._credentials, cred_project_id = self.load_auth(
-                credentials=credentials, project_id=project_id
-            )
-            if not self.project_id:
-                self.project_id = project_id or cred_project_id
         else:
-            if self._credentials.expired or not self._credentials.token:
-                self.refresh_auth(self._credentials)
-
-            if not self.project_id:
-                self.project_id = self._credentials.quota_project_id
-
-        if not self.project_id:
-            raise ValueError("Could not resolve project_id")
-
-        if not self._credentials or not self._credentials.token:
-            raise RuntimeError("Could not resolve API token from the environment")
-
-        return self._credentials.token, project_id or self.project_id
+            return self.get_access_token(
+                credentials=credentials,
+                project_id=project_id,
+            )
 
     def is_using_v1beta1_features(self, optional_params: dict) -> bool:
         """
@@ -259,6 +242,89 @@ def _get_token_and_url(
             url=url,
         )
 
+    def get_access_token(
+        self,
+        credentials: Optional[VERTEX_CREDENTIALS_TYPES],
+        project_id: Optional[str],
+    ) -> Tuple[str, str]:
+        """
+        Get access token and project id
+
+        1. Check if credentials are already in self._credentials_project_mapping
+        2. If not, load credentials and add to self._credentials_project_mapping
+        3. Check if loaded credentials have expired
+        4. If expired, refresh credentials
+        5. Return access token and project id
+        """
+
+        # Convert dict credentials to string for caching
+        cache_credentials = (
+            json.dumps(credentials) if isinstance(credentials, dict) else credentials
+        )
+        credential_cache_key = (cache_credentials, project_id)
+        _credentials: Optional[GoogleCredentialsObject] = None
+
+        verbose_logger.debug(
+            f"Checking cached credentials for project_id: {project_id}"
+        )
+
+        if credential_cache_key in self._credentials_project_mapping:
+            verbose_logger.debug(
+                f"Cached credentials found for project_id: {project_id}."
+            )
+            _credentials = self._credentials_project_mapping[credential_cache_key]
+            verbose_logger.debug("Using cached credentials")
+            credential_project_id = _credentials.quota_project_id
+        else:
+            verbose_logger.debug(
+                f"Credential cache key not found for project_id: {project_id}, loading new credentials"
+            )
+            _credentials, credential_project_id = self.load_auth(
+                credentials=credentials, project_id=project_id
+            )
+
+            if _credentials is None:
+                raise ValueError(
+                    "Could not resolve credentials - either dynamically or from environment, for project_id: {}".format(
+                        project_id
+                    )
+                )
+
+            self._credentials_project_mapping[credential_cache_key] = _credentials
+
+        ## VALIDATE CREDENTIALS
+        verbose_logger.debug(f"Validating credentials for project_id: {project_id}")
+        if (
+            project_id is not None
+            and credential_project_id
+            and credential_project_id != project_id
+        ):
+            raise ValueError(
+                "Could not resolve project_id. Credential project_id: {} does not match requested project_id: {}".format(
+                    _credentials.quota_project_id, project_id
+                )
+            )
+        else:
+            project_id = cast(
+                Optional[str], _credentials.quota_project_id
+            )  # if none set, use credential project_id
+
+        if _credentials.expired:
+            self.refresh_auth(_credentials)
+
+        ## VALIDATION STEP
+        if _credentials.token is None or not isinstance(_credentials.token, str):
+            raise ValueError(
+                "Could not resolve credentials token. Got None or non-string token - {}".format(
+                    _credentials.token
+                )
+            )
+
+        if project_id is None:
+            raise ValueError("Could not resolve project_id")
+
+        return _credentials.token, project_id
+
     async def _ensure_access_token_async(
         self,
         credentials: Optional[VERTEX_CREDENTIALS_TYPES],
@@ -272,38 +338,11 @@ async def _ensure_access_token_async(
         """
         if custom_llm_provider == "gemini":
             return "", ""
-        if self.access_token is not None:
-            if project_id is not None:
-                return self.access_token, project_id
-            elif self.project_id is not None:
-                return self.access_token, self.project_id
-
-        if not self._credentials:
-            try:
-                self._credentials, cred_project_id = await asyncify(self.load_auth)(
-                    credentials=credentials, project_id=project_id
-                )
-            except Exception:
-                verbose_logger.exception(
-                    "Failed to load vertex credentials. Check to see if credentials containing partial/invalid information."
-                )
-                raise
-            if not self.project_id:
-                self.project_id = project_id or cred_project_id
         else:
-            if self._credentials.expired or not self._credentials.token:
-                await asyncify(self.refresh_auth)(self._credentials)
-
-            if not self.project_id:
-                self.project_id = self._credentials.quota_project_id
-
-        if not self.project_id:
-            raise ValueError("Could not resolve project_id")
-
-        if not self._credentials or not self._credentials.token:
-            raise RuntimeError("Could not resolve API token from the environment")
-
-        return self._credentials.token, project_id or self.project_id
+            return await asyncify(self.get_access_token)(
+                credentials=credentials,
+                project_id=project_id,
+            )
 
     def set_headers(
         self, auth_header: Optional[str], extra_headers: Optional[dict]
diff --git a/tests/litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/litellm/llms/vertex_ai/test_vertex_llm_base.py
new file mode 100644
index 000000000000..135dc5b616fd
--- /dev/null
+++ b/tests/litellm/llms/vertex_ai/test_vertex_llm_base.py
@@ -0,0 +1,176 @@
+import os
+import sys
+from unittest.mock import MagicMock, call, patch
+
+import pytest
+
+from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
+
+sys.path.insert(
+    0, os.path.abspath("../../..")
+)  # Adds the parent directory to the system path
+
+import litellm
+from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
+
+
+def run_sync(coro):
+    """Helper to run coroutine synchronously for testing"""
+    import asyncio
+
+    return asyncio.run(coro)
+
+
+class TestVertexBase:
+    @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
+    @pytest.mark.asyncio
+    async def test_credential_project_validation(self, is_async):
+        vertex_base = VertexBase()
+
+        # Mock credentials with project_id "project-1"
+        mock_creds = MagicMock()
+        mock_creds.project_id = "project-1"
+        mock_creds.token = "fake-token-1"
+        mock_creds.expired = False
+        mock_creds.quota_project_id = "project-1"
+
+        # Test case 1: Ensure credentials match project
+        with patch.object(
+            vertex_base, "load_auth", return_value=(mock_creds, "project-1")
+        ):
+            if is_async:
+                token, project = await vertex_base._ensure_access_token_async(
+                    credentials={"type": "service_account", "project_id": "project-1"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            else:
+                token, project = vertex_base._ensure_access_token(
+                    credentials={"type": "service_account", "project_id": "project-1"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            assert project == "project-1"
+            assert token == "fake-token-1"
+
+        # Test case 2: Prevent using credentials from different project
+        with patch.object(
+            vertex_base, "load_auth", return_value=(mock_creds, "project-1")
+        ):
+            with pytest.raises(ValueError, match="Could not resolve project_id"):
+                if is_async:
+                    result = await vertex_base._ensure_access_token_async(
+                        credentials={"type": "service_account"},
+                        project_id="different-project",
+                        custom_llm_provider="vertex_ai",
+                    )
+                else:
+                    result = vertex_base._ensure_access_token(
+                        credentials={"type": "service_account"},
+                        project_id="different-project",
+                        custom_llm_provider="vertex_ai",
+                    )
+                print(f"result: {result}")
+
+    @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
+    @pytest.mark.asyncio
+    async def test_cached_credentials(self, is_async):
+        vertex_base = VertexBase()
+
+        # Initial credentials
+        mock_creds = MagicMock()
+        mock_creds.token = "token-1"
+        mock_creds.expired = False
+        mock_creds.project_id = "project-1"
+        mock_creds.quota_project_id = "project-1"
+
+        # Test initial credential load and caching
+        with patch.object(
+            vertex_base, "load_auth", return_value=(mock_creds, "project-1")
+        ):
+            # First call should load credentials
+            if is_async:
+                token, project = await vertex_base._ensure_access_token_async(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            else:
+                token, project = vertex_base._ensure_access_token(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            assert token == "token-1"
+
+            # Second call should use cached credentials
+            if is_async:
+                token2, project2 = await vertex_base._ensure_access_token_async(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            else:
+                token2, project2 = vertex_base._ensure_access_token(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            assert token2 == "token-1"
+            assert project2 == "project-1"
+
+    @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
+    @pytest.mark.asyncio
+    async def test_credential_refresh(self, is_async):
+        vertex_base = VertexBase()
+
+        # Create expired credentials
+        mock_creds = MagicMock()
+        mock_creds.token = "my-token"
+        mock_creds.expired = True
+        mock_creds.project_id = "project-1"
+        mock_creds.quota_project_id = "project-1"
+
+        with patch.object(
+            vertex_base, "load_auth", return_value=(mock_creds, "project-1")
+        ), patch.object(vertex_base, "refresh_auth") as mock_refresh:
+
+            def mock_refresh_impl(creds):
+                creds.token = "refreshed-token"
+                creds.expired = False
+
+            mock_refresh.side_effect = mock_refresh_impl
+
+            if is_async:
+                token, project = await vertex_base._ensure_access_token_async(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+            else:
+                token, project = vertex_base._ensure_access_token(
+                    credentials={"type": "service_account"},
+                    project_id="project-1",
+                    custom_llm_provider="vertex_ai",
+                )
+
+            assert mock_refresh.called
+            assert token == "refreshed-token"
+            assert not mock_creds.expired
+
+    @pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"])
+    @pytest.mark.asyncio
+    async def test_gemini_credentials(self, is_async):
+        vertex_base = VertexBase()
+
+        # Test that Gemini requests bypass credential checks
+        if is_async:
+            token, project = await vertex_base._ensure_access_token_async(
+                credentials=None, project_id=None, custom_llm_provider="gemini"
+            )
+        else:
+            token, project = vertex_base._ensure_access_token(
+                credentials=None, project_id=None, custom_llm_provider="gemini"
+            )
+        assert token == ""
+        assert project == ""


Code review comments:
diff_hunk: @@ -259,6 +242,89 @@
             url=url,
         )
 
+    def get_access_token(
+        self,
+        credentials: Optional[VERTEX_CREDENTIALS_TYPES],
+        project_id: Optional[str],
+    ) -> Tuple[str, str]:
+        """
+        Get access token and project id
+
+        1. Check if credentials are already in self._credentials_project_mapping
+        2. If not, load credentials and add to self._credentials_project_mapping
+        3. Check if loaded credentials have expired
+        4. If expired, refresh credentials
+        5. Return access token and project id
+        """
+
+        # Convert dict credentials to string for caching
+        cache_credentials = (
+            json.dumps(credentials) if isinstance(credentials, dict) else credentials
+        )
+        credential_cache_key = (cache_credentials, project_id)
+        _credentials: Optional[GoogleCredentialsObject] = None
+
+        verbose_logger.debug(
+            f"Checking cached credentials for project_id: {project_id}"
+        )
+
+        if credential_cache_key in self._credentials_project_mapping:
+            verbose_logger.debug(
+                f"Cached credentials found for project_id: {project_id}."
+            )
+            _credentials = self._credentials_project_mapping[credential_cache_key]
+            verbose_logger.debug("Using cached credentials")
+            credential_project_id = _credentials.quota_project_id
+        else:
+            verbose_logger.debug(
+                f"Credential cache key not found for project_id: {project_id}, loading new credentials"
path: litellm/llms/vertex_ai/vertex_llm_base.py
text: @user1:
## Clear-text logging of sensitive information

This expression logs [sensitive data (secret)](1) as clear text.
This expression logs [sensitive data (secret)](2) as clear text.
This expression logs [sensitive data (secret)](3) as clear text.
This expression logs [sensitive data (secret)](4) as clear text.
This expression logs [sensitive data (secret)](5) as clear text.
This expression logs [sensitive data (secret)](6) as clear text.
This expression logs [sensitive data (secret)](7) as clear text.
This expression logs [sensitive data (secret)](8) as clear text.
This expression logs [sensitive data (secret)](9) as clear text.
This expression logs [sensitive data (secret)](10) as clear text.
This expression logs [sensitive data (secret)](11) as clear text.
This expression logs [sensitive data (secret)](12) as clear text.
This expression logs [sensitive data (secret)](13) as clear text.
This expression logs [sensitive data (secret)](14) as clear text.
This expression logs [sensitive data (secret)](15) as clear text.
This expression logs [sensitive data (secret)](16) as clear text.
This expression logs [sensitive data (secret)](17) as clear text.
This expression logs [sensitive data (secret)](18) as clear text.
This expression logs [sensitive data (secret)](19) as clear text.
This expression logs [sensitive data (secret)](20) as clear text.
This expression logs [sensitive data (secret)](21) as clear text.
This expression logs [sensitive data (secret)](22) as clear text.
This expression logs [sensitive data (secret)](23) as clear text.
This expression logs [sensitive data (secret)](24) as clear text.
This expression logs [sensitive data (secret)](25) as clear text.
This expression logs [sensitive data (secret)](26) as clear text.

[Show more details](https://github.com/BerriAI/litellm/security/code-scanning/2413)