Merge branch 'litellm_yj_may1' into codex/integration-host-credential-guard

.github/workflows/check-lazy-openapi-snapshot.yml

@@ -1,75 +0,0 @@
-name: Check Lazy OpenAPI Snapshot
-
-on:
-  pull_request:
-    branches:
-      - main
-      - litellm_internal_staging
-      - "litellm_**"
-
-permissions:
-  contents: read
-  checks: write
-
-concurrency:
-  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
-  cancel-in-progress: true
-
-jobs:
-  verify:
-    runs-on: ubuntu-latest
-    timeout-minutes: 10
-    steps:
-      - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
-        with:
-          persist-credentials: false
-
-      - name: Set up Python
-        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
-        with:
-          python-version: "3.12"
-
-      - name: Set up uv
-        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
-        with:
-          version: "0.10.9"
-
-      - name: Cache uv dependencies
-        uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
-        with:
-          path: |
-            ~/.cache/uv
-            .venv
-          key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
-          restore-keys: |
-            ${{ runner.os }}-uv-
-
-      - name: Install dependencies
-        run: uv sync --frozen --all-groups --all-extras
-
-      - name: Regenerate snapshot to /tmp
-        id: regen
-        run: |
-          cp litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.committed.json
-          uv run --no-sync python -m litellm.proxy._lazy_openapi_snapshot
-          mv litellm/proxy/_lazy_openapi_snapshot.json /tmp/snapshot.fresh.json
-          mv /tmp/snapshot.committed.json litellm/proxy/_lazy_openapi_snapshot.json
-
-      - name: Compare
-        id: diff
-        continue-on-error: true
-        run: |
-          diff -q /tmp/snapshot.fresh.json litellm/proxy/_lazy_openapi_snapshot.json
-
-      - name: Mark neutral if drift
-        if: steps.diff.outcome == 'failure'
-        uses: LouisBrunner/checks-action@6b626ffbad7cc56fd58627f774b9067e6118af23 # v2.0.0
-        with:
-          token: ${{ secrets.GITHUB_TOKEN }}
-          name: lazy-openapi-snapshot
-          conclusion: neutral
-          output: |
-            {
-              "title": "Lazy openapi snapshot is stale",
-              "summary": "Run `python -m litellm.proxy._lazy_openapi_snapshot` and commit the regenerated `litellm/proxy/_lazy_openapi_snapshot.json`. Not blocking — the snapshot will regenerate at release if not committed."
-            }

.gitignore

@@ -90,7 +90,6 @@ test.py
 litellm_config.yaml
 !.github/observatory/litellm_config.yaml
 .cursor
-.vscode/launch.json
 litellm/proxy/to_delete_loadtest_work/*
 update_model_cost_map.py
 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
@@ -100,4 +99,5 @@ STABILIZATION_TODO.md
 **/test-results
 **/playwright-report
 **/*.storageState.json
-**/coverage
+**/coverage
+test-config

README.md

@@ -68,7 +68,7 @@ Managing LLM calls across providers gets complicated fast — different SDKs, au
     <td><img height="60" alt="Stripe" src="https://github.com/user-attachments/assets/f7296d4f-9fbd-460d-9d05-e4df31697c4b" /></td>
     <td><img height="60" alt="image" src="https://github.com/user-attachments/assets/436fca71-988b-40bb-b5fe-8450c80fdbd0" /></td>
     <td><img height="60" alt="Google ADK" src="https://github.com/user-attachments/assets/caf270a2-5aee-45c4-8222-41a2070c4f19" /></td>
-    <td><img height="60" alt="Greptile" src="https://github.com/user-attachments/assets/0be4bd8a-7cfa-48d3-9090-f415fe948280" /></td>
+    <td><img height="60" alt="Greptile" src="https://github.com/user-attachments/assets/3db0ae72-0843-4005-a56d-bba1dde2193d" /></td>
     <td><img height="60" alt="OpenHands" src="https://github.com/user-attachments/assets/a6150c4c-149e-4cae-888b-8b92be6e003f" /></td>
     <td><h2>Netflix</h2></td>
     <td><img height="60" alt="OpenAI Agents SDK" src="https://github.com/user-attachments/assets/c02f7be0-8c2e-4d27-aea7-7c024bfaebc0" /></td>

enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py

@@ -857,10 +857,16 @@ async def project_info(
                 where={"team_id": project.team_id}
             )
             if team:
-                is_team_member = (
-                    user_api_key_dict.user_id in team.admins
-                    or user_api_key_dict.user_id in team.members
-                )
+                caller_user_id = user_api_key_dict.user_id
+                for m in team.members_with_roles or []:
+                    m_user_id = (
+                        m.get("user_id")
+                        if isinstance(m, dict)
+                        else getattr(m, "user_id", None)
+                    )
+                    if m_user_id == caller_user_id:
+                        is_team_member = True
+                        break
 
         if not (is_admin or is_team_member):
             raise HTTPException(
@@ -911,20 +917,20 @@ async def list_projects(
                 include={"litellm_budget_table": True, "object_permission": True}
             )
         else:
-            # Get projects for teams the user belongs to
-            user_teams = await prisma_client.db.litellm_teamtable.find_many(
-                where={
-                    "OR": [
-                        {"members": {"has": user_api_key_dict.user_id}},
-                        {"admins": {"has": user_api_key_dict.user_id}},
-                    ]
-                }
+            # Look up the user's team memberships via the reverse-index on
+            # LiteLLM_UserTable.teams (maintained by team_member_add alongside
+            # members_with_roles). This avoids a full scan of all team rows.
+            user_record = await prisma_client.db.litellm_usertable.find_unique(
+                where={"user_id": user_api_key_dict.user_id},
+            )
+            user_team_ids = (
+                user_record.teams
+                if user_record is not None and user_record.teams
+                else []
             )
 
-            team_ids = [team.team_id for team in user_teams]
-
             projects = await prisma_client.db.litellm_projecttable.find_many(
-                where={"team_id": {"in": team_ids}},
+                where={"team_id": {"in": user_team_ids}},
                 include={"litellm_budget_table": True, "object_permission": True},
             )
 

litellm-proxy-extras/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litellm-proxy-extras"
-version = "0.4.69"
+version = "0.4.70"
 description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,7 +26,7 @@ required-version = ">=0.10.9"
 module-root = ""
 
 [tool.commitizen]
-version = "0.4.69"
+version = "0.4.70"
 version_files = [
     "pyproject.toml:^version",
     "../pyproject.toml:litellm-proxy-extras==",

litellm/caching/caching.py

@@ -432,9 +432,10 @@ class Cache:
             str: The final hashed cache key with the redis namespace.
         """
         dynamic_cache_control: DynamicCacheControl = kwargs.get("cache", {})
+        metadata = kwargs.get("metadata") or {}
         namespace = (
             dynamic_cache_control.get("namespace")
-            or kwargs.get("metadata", {}).get("redis_namespace")
+            or metadata.get("redis_namespace")
             or self.namespace
         )
         if namespace:

litellm/caching/caching_handler.py

@@ -87,6 +87,18 @@ class CachingHandlerResponse(BaseModel):
 in_memory_cache_obj = InMemoryCache()
 
 
+def _should_defer_streaming_cache_hit_callbacks(*, kwargs: Dict[str, Any]) -> bool:
+    """
+    When stream=True, do not run success callbacks at cache-hit time.
+
+    Cached chat/text completion replay uses CustomStreamWrapper; cached Responses
+    replay uses CachedResponsesAPIStreamingIterator. Both invoke logging success
+    handlers when the stream finishes; firing them here too would double-count
+    spend and callback records.
+    """
+    return kwargs.get("stream", False) is True
+
+
 class LLMCachingHandler:
     def __init__(
         self,
@@ -99,6 +111,7 @@ class LLMCachingHandler:
         self.async_streaming_chunks: List[ModelResponse] = []
         self.sync_streaming_chunks: List[ModelResponse] = []
         self.request_kwargs = request_kwargs
+        self.preset_cache_key: Optional[str] = None
         self.original_function = original_function
         self.start_time = start_time
         if litellm.cache is not None and isinstance(litellm.cache.cache, RedisCache):
@@ -206,7 +219,7 @@ class LLMCachingHandler:
                         custom_llm_provider=kwargs.get("custom_llm_provider", None),
                         args=args,
                     )
-                    if kwargs.get("stream", False) is False:
+                    if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
                         # LOG SUCCESS
                         self._async_log_cache_hit_on_callbacks(
                             logging_obj=logging_obj,
@@ -215,11 +228,12 @@ class LLMCachingHandler:
                             end_time=end_time,
                             cache_hit=cache_hit,
                         )
-                    cache_key = litellm.cache.get_cache_key(**kwargs)
-                    if (
-                        isinstance(cached_result, BaseModel)
-                        or isinstance(cached_result, CustomStreamWrapper)
-                    ) and hasattr(cached_result, "_hidden_params"):
+                    cache_key = (
+                        self.preset_cache_key
+                        or self.request_kwargs.get("cache_key")
+                        or litellm.cache.get_cache_key(**self.request_kwargs)
+                    )
+                    if hasattr(cached_result, "_hidden_params"):
                         cached_result._hidden_params["cache_key"] = cache_key  # type: ignore
                     return CachingHandlerResponse(cached_result=cached_result)
                 elif (
@@ -265,8 +279,6 @@ class LLMCachingHandler:
         kwargs: Dict[str, Any],
         args: Optional[Tuple[Any, ...]] = None,
     ) -> CachingHandlerResponse:
-        from litellm.utils import CustomStreamWrapper
-
         cached_result: Optional[Any] = None
 
         # Check if caching should be performed BEFORE doing expensive kwargs copy
@@ -282,6 +294,11 @@ class LLMCachingHandler:
                     args,
                 )
             )
+            if new_kwargs.get("metadata") is None:
+                new_kwargs.pop("metadata", None)
+            if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
+                new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
+            self.request_kwargs = new_kwargs
             print_verbose("Checking Sync Cache")
             cached_result = litellm.cache.get_cache(**new_kwargs)
             if cached_result is not None:
@@ -322,17 +339,19 @@ class LLMCachingHandler:
                         is_async=False,
                     )
 
-                    logging_obj.handle_sync_success_callbacks_for_async_calls(
-                        result=cached_result,
-                        start_time=start_time,
-                        end_time=end_time,
-                        cache_hit=cache_hit,
+                    if not _should_defer_streaming_cache_hit_callbacks(kwargs=kwargs):
+                        logging_obj.handle_sync_success_callbacks_for_async_calls(
+                            result=cached_result,
+                            start_time=start_time,
+                            end_time=end_time,
+                            cache_hit=cache_hit,
+                        )
+                    cache_key = (
+                        self.preset_cache_key
+                        or self.request_kwargs.get("cache_key")
+                        or litellm.cache.get_cache_key(**self.request_kwargs)
                     )
-                    cache_key = litellm.cache.get_cache_key(**kwargs)
-                    if (
-                        isinstance(cached_result, BaseModel)
-                        or isinstance(cached_result, CustomStreamWrapper)
-                    ) and hasattr(cached_result, "_hidden_params"):
+                    if hasattr(cached_result, "_hidden_params"):
                         cached_result._hidden_params["cache_key"] = cache_key  # type: ignore
                     return CachingHandlerResponse(cached_result=cached_result)
         return CachingHandlerResponse(cached_result=cached_result)
@@ -686,6 +705,11 @@ class LLMCachingHandler:
                 args,
             )
         )
+        if new_kwargs.get("metadata") is None:
+            new_kwargs.pop("metadata", None)
+        if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs:
+            new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs)
+        self.request_kwargs = new_kwargs
         cached_result: Optional[Any] = None
         if call_type == CallTypes.aembedding.value:
             if isinstance(new_kwargs["input"], str):
@@ -710,14 +734,26 @@ class LLMCachingHandler:
                 if all(result is None for result in cached_result):
                     cached_result = None
         else:
+            request_kwargs = new_kwargs.copy()
+            request_cache_key = request_kwargs.pop("cache_key", None)
             if litellm.cache._supports_async() is True:
                 ## check if dual cache is supported ##
+                self.preset_cache_key = (
+                    request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
+                )
                 cached_result = await litellm.cache.async_get_cache(
-                    dynamic_cache_object=self.dual_cache, **new_kwargs
+                    dynamic_cache_object=self.dual_cache,
+                    cache_key=self.preset_cache_key,
+                    **request_kwargs,
                 )
             else:  # fallback for caches that don't support async
+                self.preset_cache_key = (
+                    request_cache_key or litellm.cache.get_cache_key(**request_kwargs)
+                )
                 cached_result = litellm.cache.get_cache(
-                    dynamic_cache_object=self.dual_cache, **new_kwargs
+                    dynamic_cache_object=self.dual_cache,
+                    cache_key=self.preset_cache_key,
+                    **request_kwargs,
                 )
         return cached_result
 
@@ -825,8 +861,27 @@ class LLMCachingHandler:
         elif (call_type == "aresponses" or call_type == "responses") and isinstance(
             cached_result, dict
         ):
-            # Convert cached dict back to ResponsesAPIResponse object
-            cached_result = ResponsesAPIResponse(**cached_result)
+            from litellm.responses.streaming_iterator import (
+                CachedResponsesAPIStreamingIterator,
+            )
+
+            response_obj = ResponsesAPIResponse(**cached_result)
+            if (
+                hasattr(response_obj, "_hidden_params")
+                and response_obj._hidden_params is not None
+                and isinstance(response_obj._hidden_params, dict)
+            ):
+                response_obj._hidden_params["cache_hit"] = True
+
+            if kwargs.get("stream", False) is True:
+                cached_result = CachedResponsesAPIStreamingIterator(
+                    response=response_obj,
+                    logging_obj=logging_obj,
+                    request_data=kwargs,
+                    call_type=call_type,
+                )
+            else:
+                cached_result = response_obj
 
         if (
             hasattr(cached_result, "_hidden_params")

litellm/caching/dual_cache.py

@@ -92,6 +92,25 @@ class DualCache(BaseCache):
         if default_redis_ttl is not None:
             self.default_redis_ttl = default_redis_ttl
 
+    def attach_redis_cache(
+        self,
+        redis_cache: Optional[RedisCache] = None,
+        *,
+        default_redis_ttl: Optional[float] = None,
+    ) -> None:
+        """
+        Attach a Redis backend if this DualCache does not already have one.
+
+        No-op when ``redis_cache`` is None or when Redis was already set (constructor
+        or a prior attach). Use this for lazy wiring after a shared Redis client exists.
+        Does not backfill in-memory-only keys to Redis.
+        """
+        if redis_cache is None or self.redis_cache is not None:
+            return
+        self.redis_cache = redis_cache
+        if default_redis_ttl is not None:
+            self.default_redis_ttl = default_redis_ttl
+
     def set_cache(self, key, value, local_only: bool = False, **kwargs):
         # Update both Redis and in-memory cache
         try:
@@ -392,6 +411,7 @@ class DualCache(BaseCache):
         value: float,
         parent_otel_span: Optional[Span] = None,
         local_only: bool = False,
+        refresh_ttl: bool = False,
         **kwargs,
     ) -> Optional[float]:
         """
@@ -399,6 +419,9 @@ class DualCache(BaseCache):
 
         Value - float - the value you want to increment by
 
+        Refresh_ttl - bool - if True, resets the Redis TTL on every write.
+        Default False preserves window-style semantics.
+
         Returns - the incremented value, or None if no cache backend is
         available (in_memory_cache is None and Redis failed/is absent).
         """
@@ -415,6 +438,7 @@ class DualCache(BaseCache):
                     value,
                     parent_otel_span=parent_otel_span,
                     ttl=kwargs.get("ttl", None),
+                    refresh_ttl=refresh_ttl,
                 )
 
             return result

litellm/caching/redis_cache.py

@@ -551,6 +551,13 @@ class RedisCache(BaseCache):
     async def async_set_cache(self, key, value, **kwargs):
         from redis.asyncio import Redis
 
+        if key is None:
+            verbose_logger.debug(
+                "LiteLLM Redis Caching: async set() skipped — key is None, value=%r",
+                value,
+            )
+            return None
+
         start_time = time.time()
         try:
             _redis_client: Redis = self.init_async_client()  # type: ignore
@@ -569,8 +576,9 @@ class RedisCache(BaseCache):
                 )
             )
             verbose_logger.error(
-                "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, Writing value=%s",
+                "LiteLLM Redis Caching: async set() - Got exception from REDIS %s, key=%r, value=%r",
                 str(e),
+                key,
                 value,
             )
             raise e
@@ -824,6 +832,7 @@ class RedisCache(BaseCache):
         value: float,
         ttl: Optional[int] = None,
         parent_otel_span: Optional[Span] = None,
+        refresh_ttl: bool = False,
     ) -> float:
         from redis.asyncio import Redis
 
@@ -834,11 +843,12 @@ class RedisCache(BaseCache):
         try:
             result = await _redis_client.incrbyfloat(name=key, amount=value)
             if _used_ttl is not None:
-                # check if key already has ttl, if not -> set ttl
-                current_ttl = await _redis_client.ttl(key)
-                if current_ttl == -1:
-                    # Key has no expiration
+                if refresh_ttl:
                     await _redis_client.expire(key, _used_ttl)
+                else:
+                    current_ttl = await _redis_client.ttl(key)
+                    if current_ttl == -1:
+                        await _redis_client.expire(key, _used_ttl)
 
             ## LOGGING ##
             end_time = time.time()

litellm/constants.py

@@ -1425,6 +1425,7 @@ LITELLM_PROXY_ADMIN_NAME = "default_user_id"
 LITELLM_CLI_SOURCE_IDENTIFIER = "litellm-cli"
 LITELLM_CLI_SESSION_TOKEN_PREFIX = "litellm-session-token"
 CLI_SSO_SESSION_CACHE_KEY_PREFIX = "cli_sso_session"
+CLI_SSO_SESSION_TTL_SECONDS = 600
 CLI_JWT_TOKEN_NAME = "cli-jwt-token"
 # Support both CLI_JWT_EXPIRATION_HOURS and LITELLM_CLI_JWT_EXPIRATION_HOURS for backwards compatibility
 CLI_JWT_EXPIRATION_HOURS = int(

litellm/integrations/arize/arize_phoenix_client.py

@@ -2,11 +2,23 @@
 Arize Phoenix API client for fetching prompt versions from Arize Phoenix.
 """
 
+import urllib.parse
 from typing import Any, Dict, Optional
 
 from litellm.llms.custom_httpx.http_handler import HTTPHandler
 
 
+def _sanitize_id(identifier: str) -> str:
+    """Reject path traversal characters and URL-encode the identifier."""
+    if any(c in identifier for c in ("/", "\\", "#", "?")):
+        raise ValueError(
+            f"Invalid identifier {identifier!r}: contains disallowed characters"
+        )
+    if ".." in identifier:
+        raise ValueError(f"Invalid identifier {identifier!r}: path traversal detected")
+    return urllib.parse.quote(identifier, safe="")
+
+
 class ArizePhoenixClient:
     """
     Client for interacting with Arize Phoenix API to fetch prompt versions.
@@ -53,7 +65,8 @@ class ArizePhoenixClient:
         Returns:
             Dictionary containing prompt version data, or None if not found
         """
-        url = f"{self.api_base}/v1/prompt_versions/{prompt_version_id}"
+        safe_id = _sanitize_id(prompt_version_id)
+        url = f"{self.api_base}/v1/prompt_versions/{safe_id}"
 
         try:
             # Use the underlying httpx client directly to avoid query param extraction

litellm/integrations/bitbucket/bitbucket_client.py

@@ -3,11 +3,27 @@ BitBucket API client for fetching .prompt files from BitBucket repositories.
 """
 
 import base64
+import urllib.parse
 from typing import Any, Dict, List, Optional
 
 from litellm.llms.custom_httpx.http_handler import HTTPHandler
 
 
+def _sanitize_file_path(file_path: str) -> str:
+    """Reject path traversal and URL-encode each path segment."""
+    if "#" in file_path or "?" in file_path:
+        raise ValueError(
+            f"Invalid file path {file_path!r}: contains URL special characters"
+        )
+    parts = file_path.split("/")
+    for part in parts:
+        if part == "..":
+            raise ValueError(
+                f"Invalid file path {file_path!r}: path traversal detected"
+            )
+    return "/".join(urllib.parse.quote(part, safe="") for part in parts)
+
+
 class BitBucketClient:
     """
     Client for interacting with BitBucket API to fetch .prompt files.
@@ -72,7 +88,8 @@ class BitBucketClient:
         Returns:
             File content as string, or None if file not found
         """
-        url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{file_path}"
+        safe_path = _sanitize_file_path(file_path)
+        url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}"
 
         try:
             response = self.http_handler.get(url, headers=self.headers)
@@ -119,7 +136,8 @@ class BitBucketClient:
         Returns:
             List of file paths
         """
-        url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{directory_path}"
+        safe_dir = _sanitize_file_path(directory_path) if directory_path else ""
+        url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_dir}"
 
         try:
             response = self.http_handler.get(url, headers=self.headers)
@@ -211,7 +229,8 @@ class BitBucketClient:
         Returns:
             Dictionary containing file metadata, or None if file not found
         """
-        url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{file_path}"
+        safe_path = _sanitize_file_path(file_path)
+        url = f"{self.base_url}/repositories/{self.workspace}/{self.repository}/src/{self.branch}/{safe_path}"
 
         try:
             # Use GET with Range header to get just the headers (HEAD equivalent)

litellm/integrations/prometheus.py

@@ -265,6 +265,7 @@ class PrometheusLogger(CustomLogger):
             ########################################
             # LiteLLM Virtual API KEY metrics
             ########################################
+
             # Remaining MODEL RPM limit for API Key
             self.litellm_remaining_api_key_requests_for_model = self._gauge_factory(
                 "litellm_remaining_api_key_requests_for_model",

litellm/litellm_core_utils/cli_token_utils.py

@@ -31,15 +31,23 @@ def load_cli_token() -> Optional[dict]:
         return None
 
 
-def get_litellm_gateway_api_key() -> Optional[str]:
+def get_litellm_gateway_api_key(
+    expected_base_url: Optional[str] = None,
+) -> Optional[str]:
     """
     Get the stored CLI API key for use with LiteLLM SDK.
 
     This function reads the token file created by `litellm-proxy login`
     and returns the API key for use in Python scripts.
 
+    Args:
+        expected_base_url: When provided, the key is only returned if it was
+            originally issued for this URL. Pass the target server URL to
+            prevent credential leakage when the client is pointed at a
+            different (possibly malicious) server.
+
     Returns:
-        str: The API key if found, None otherwise
+        str: The API key if found (and origin matches), None otherwise
 
     Example:
         >>> import litellm
@@ -53,6 +61,10 @@ def get_litellm_gateway_api_key() -> Optional[str]:
         >>>     )
     """
     token_data = load_cli_token()
-    if token_data and "key" in token_data:
-        return token_data["key"]
-    return None
+    if not token_data or "key" not in token_data:
+        return None
+    if expected_base_url is not None:
+        stored_url = token_data.get("base_url")
+        if stored_url != expected_base_url.rstrip("/"):
+            return None
+    return token_data["key"]

litellm/litellm_core_utils/litellm_logging.py

@@ -4725,7 +4725,7 @@ class StandardLoggingPayloadSetup:
         ):
             for key, value in litellm_params["metadata"].items():
                 # Skip non-serializable objects like UserAPIKeyAuth
-                if key == "user_api_key_auth":
+                if key in {"user_api_key_auth", "user_api_key_budget_reservation"}:
                     continue
                 merged_metadata[key] = value
 

litellm/litellm_core_utils/prompt_templates/factory.py

@@ -4582,6 +4582,11 @@ class BedrockConverseMessagesProcessor:
                                     message=cast(ChatCompletionFileObject, element)
                                 )
                                 _parts.append(_part)
+                            elif element["type"] == "document":
+                                _part = BedrockConverseMessagesProcessor._process_document_message(
+                                    element
+                                )
+                                _parts.append(_part)
                             _cache_point_block = (
                                 litellm.AmazonConverseConfig()._get_cache_point_block(
                                     message_block=cast(
@@ -4864,6 +4869,44 @@ class BedrockConverseMessagesProcessor:
             image_url=cast(str, file_id or file_data), format=format
         )
 
+    @staticmethod
+    def _process_document_message(element: dict) -> BedrockContentBlock:
+        """Convert a document content block to a Bedrock DocumentBlock.
+
+        Handles the Anthropic-style document format:
+        {"type": "document", "source": {"type": "base64", "media_type": "application/pdf", "data": "..."}}
+        """
+        source = element["source"]
+        source_type = source.get("type")
+        if source_type != "base64":
+            raise ValueError(
+                f"Bedrock Converse only supports base64-encoded document sources, got '{source_type}'. "
+                "Please convert the document to base64 before sending to Bedrock."
+            )
+        media_type: str = source["media_type"]
+        data: str = source["data"]
+        doc_format = BedrockImageProcessor._validate_format(
+            mime_type=media_type, image_format=media_type.split("/")[1]
+        )
+
+        # Deterministic name using the same hashing pattern as _create_bedrock_block
+        HASH_SAMPLE_BYTES = 64 * 1024
+        normalized = "".join(data.split()).encode("utf-8")
+        sample = normalized[:HASH_SAMPLE_BYTES]
+        hasher = hashlib.sha256()
+        hasher.update(sample)
+        hasher.update(str(len(normalized)).encode("utf-8"))
+        content_hash = hasher.hexdigest()[:16]
+        document_name = f"Document_{content_hash}_{doc_format}"
+
+        return BedrockContentBlock(
+            document=BedrockDocumentBlock(
+                source=BedrockSourceBlock(bytes=data),
+                format=doc_format,
+                name=document_name,
+            )
+        )
+
     @staticmethod
     def add_thinking_blocks_to_assistant_content(
         thinking_blocks: List[BedrockContentBlock],
@@ -4961,6 +5004,11 @@ def _bedrock_converse_messages_pt(  # noqa: PLR0915
                                 )
                             )
                             _parts.append(_part)
+                        elif element["type"] == "document":
+                            _part = BedrockConverseMessagesProcessor._process_document_message(
+                                element
+                            )
+                            _parts.append(_part)
                         _cache_point_block = (
                             litellm.AmazonConverseConfig()._get_cache_point_block(
                                 message_block=cast(

litellm/litellm_core_utils/streaming_handler.py

@@ -2244,7 +2244,7 @@ class CustomStreamWrapper:
                 asyncio.create_task(
                     self.logging_obj.async_failure_handler(e, traceback_exception)
                 )
-            raise e
+            self._handle_stream_fallback_error(e)
         except Exception as e:
             traceback_exception = traceback.format_exc()
             if self.logging_obj is not None:

litellm/litellm_core_utils/url_utils.py

@@ -199,6 +199,47 @@ def validate_url(url: str) -> Tuple[str, str]:
     return rewritten, host_header
 
 
+def assert_same_origin(candidate_url: str, expected_url: str) -> None:
+    """Verify ``candidate_url`` shares scheme, host, and port with ``expected_url``.
+
+    Use when an upstream API returns a URL meant for follow-up requests
+    (e.g. an async-job polling URL that will be hit with the operator's
+    API key in the headers). The upstream is trusted because the operator
+    configured ``api_base``, but the URL it hands back must actually point
+    back at the same origin or we'd be blindly forwarding credentials
+    wherever the upstream told us to.
+
+    Hostnames are compared case-insensitively. Default ports are made
+    explicit (HTTP→80, HTTPS→443) so ``https://api.example.com:443/...``
+    and ``https://api.example.com/...`` are treated as the same origin.
+
+    Error messages identify *which* component mismatched but never echo
+    the operator's ``expected`` host or the candidate's hostname back to
+    the caller — in the SSRF threat model the caller is the attacker,
+    and reflecting host info would be a secondary leak of operator
+    infrastructure details.
+    """
+    candidate = urlparse(candidate_url)
+    expected = urlparse(expected_url)
+
+    if candidate.scheme not in _ALLOWED_SCHEMES:
+        raise SSRFError("URL scheme is not allowed")
+
+    if candidate.scheme != expected.scheme:
+        raise SSRFError("Origin mismatch on scheme")
+
+    candidate_host = _normalize_host(candidate.hostname or "")
+    expected_host = _normalize_host(expected.hostname or "")
+    if not candidate_host or candidate_host != expected_host:
+        raise SSRFError("Origin mismatch on host")
+
+    default_port = 443 if candidate.scheme == "https" else 80
+    candidate_port = candidate.port if candidate.port is not None else default_port
+    expected_port = expected.port if expected.port is not None else default_port
+    if candidate_port != expected_port:
+        raise SSRFError("Origin mismatch on port")
+
+
 _MAX_REDIRECTS = 10
 
 

litellm/llms/anthropic/chat/transformation.py

@@ -1553,25 +1553,43 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
             )
         data["output_config"] = output_config
 
-    def _transform_response_for_json_mode(
+    def _resolve_json_mode_non_streaming(
         self,
         json_mode: Optional[bool],
         tool_calls: List[ChatCompletionToolCallChunk],
-    ) -> Optional[LitellmMessage]:
-        _message: Optional[LitellmMessage] = None
-        if json_mode is True and len(tool_calls) == 1:
-            # check if tool name is the default tool name
-            json_mode_content_str: Optional[str] = None
-            if (
-                "name" in tool_calls[0]["function"]
-                and tool_calls[0]["function"]["name"] == RESPONSE_FORMAT_TOOL_NAME
-            ):
-                json_mode_content_str = tool_calls[0]["function"].get("arguments")
-            if json_mode_content_str is not None:
-                _message = AnthropicConfig._convert_tool_response_to_message(
-                    tool_calls=tool_calls,
-                )
-        return _message
+    ) -> Tuple[
+        Optional[LitellmMessage],
+        List[ChatCompletionToolCallChunk],
+        Optional[str],
+    ]:
+        """Strip internal response_format tool calls; merge payload into content when mixed with user tools."""
+        if json_mode is not True or not tool_calls:
+            return None, tool_calls, None
+
+        json_indices = [
+            i
+            for i, t in enumerate(tool_calls)
+            if t.get("function", {}).get("name") == RESPONSE_FORMAT_TOOL_NAME
+        ]
+        if not json_indices:
+            return None, tool_calls, None
+
+        if len(json_indices) == len(tool_calls):
+            json_tool = tool_calls[json_indices[0]]
+            if json_tool.get("function", {}).get("arguments") is None:
+                return None, tool_calls, None
+            _message = AnthropicConfig._convert_tool_response_to_message(
+                tool_calls=[json_tool]
+            )
+            return _message, [], None
+
+        first_json = tool_calls[json_indices[0]]
+        json_msg = AnthropicConfig._convert_tool_response_to_message([first_json])
+        extra_content: Optional[str] = (
+            json_msg.content if json_msg is not None else None
+        )
+        filtered_tools = [t for i, t in enumerate(tool_calls) if i not in json_indices]
+        return None, filtered_tools, extra_content
 
     def extract_response_content(self, completion_response: dict) -> Tuple[
         str,
@@ -1931,19 +1949,27 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
             tool_calls,
         )
 
+        json_mode_message, tool_calls_for_message, json_extra_content = (
+            self._resolve_json_mode_non_streaming(
+                json_mode=json_mode,
+                tool_calls=tool_calls,
+            )
+        )
+        merged_text = text_content or ""
+        if json_extra_content:
+            merged_text = (
+                merged_text + json_extra_content if merged_text else json_extra_content
+            )
+
         _message = litellm.Message(
-            tool_calls=tool_calls,
-            content=text_content or None,
+            tool_calls=tool_calls_for_message,
+            content=merged_text or None,
             provider_specific_fields=provider_specific_fields,
             thinking_blocks=thinking_blocks,
             reasoning_content=reasoning_content,
         )
         _message.provider_specific_fields = provider_specific_fields
 
-        json_mode_message = self._transform_response_for_json_mode(
-            json_mode=json_mode,
-            tool_calls=tool_calls,
-        )
         if json_mode_message is not None:
             completion_response["stop_reason"] = "stop"
             _message = json_mode_message

litellm/llms/azure/azure.py

@@ -16,6 +16,7 @@ import litellm
 from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.logging_utils import track_llm_api_timing
+from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     HTTPHandler,
@@ -898,6 +899,17 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
                 operation_location_url = response.headers["operation-location"]
             else:
                 raise AzureOpenAIError(status_code=500, message=response.text)
+            # Reject polling URLs that don't share an origin with ``api_base``.
+            # Without this an upstream-controlled or attacker-controlled
+            # value would receive the operator's Azure API key in the
+            # request headers below. VERIA-51.
+            try:
+                assert_same_origin(operation_location_url, api_base)
+            except SSRFError as ssrf_err:
+                raise AzureOpenAIError(
+                    status_code=502,
+                    message=f"Rejected polling URL: {ssrf_err}",
+                )
             response = await async_handler.get(
                 url=operation_location_url,
                 headers=headers,
@@ -908,8 +920,13 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
             timeout_secs: int = AZURE_OPERATION_POLLING_TIMEOUT
             start_time = time.time()
             if "status" not in response.json():
-                raise Exception(
-                    "Expected 'status' in response. Got={}".format(response.json())
+                # Don't reflect the raw response body — when the polling
+                # URL points at an internal JSON API (cloud metadata
+                # service etc.) reflecting it here turns Blind SSRF into
+                # Full-Read SSRF. VERIA-51.
+                raise AzureOpenAIError(
+                    status_code=502,
+                    message="Polling response missing 'status' field",
                 )
             while response.json()["status"] not in ["succeeded", "failed"]:
                 if time.time() - start_time > timeout_secs:
@@ -1009,6 +1026,13 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
                 operation_location_url = response.headers["operation-location"]
             else:
                 raise AzureOpenAIError(status_code=500, message=response.text)
+            try:
+                assert_same_origin(operation_location_url, api_base)
+            except SSRFError as ssrf_err:
+                raise AzureOpenAIError(
+                    status_code=502,
+                    message=f"Rejected polling URL: {ssrf_err}",
+                )
             response = sync_handler.get(
                 url=operation_location_url,
                 headers=headers,
@@ -1019,8 +1043,9 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
             timeout_secs: int = AZURE_OPERATION_POLLING_TIMEOUT
             start_time = time.time()
             if "status" not in response.json():
-                raise Exception(
-                    "Expected 'status' in response. Got={}".format(response.json())
+                raise AzureOpenAIError(
+                    status_code=502,
+                    message="Polling response missing 'status' field",
                 )
             while response.json()["status"] not in ["succeeded", "failed"]:
                 if time.time() - start_time > timeout_secs:

litellm/llms/azure_ai/ocr/document_intelligence/transformation.py

@@ -17,6 +17,7 @@ from urllib.parse import quote
 import httpx
 
 from litellm._logging import verbose_logger
+from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin
 from litellm.constants import (
     AZURE_DOCUMENT_INTELLIGENCE_API_VERSION,
     AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI,
@@ -599,6 +600,16 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
                         "Azure Document Intelligence returned 202 but no Operation-Location header found"
                     )
 
+                # Reject cross-origin polling URLs — the auth headers
+                # below would otherwise leak to whatever URL the upstream
+                # (or an attacker-controlled upstream) returns. VERIA-51.
+                try:
+                    assert_same_origin(operation_url, str(raw_response.request.url))
+                except SSRFError as ssrf_err:
+                    raise ValueError(
+                        f"Azure Document Intelligence: rejected polling URL ({ssrf_err})"
+                    )
+
                 # Get headers for polling (need auth)
                 poll_headers = {
                     "Ocp-Apim-Subscription-Key": raw_response.request.headers.get(
@@ -711,6 +722,14 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig):
                         "Azure Document Intelligence returned 202 but no Operation-Location header found"
                     )
 
+                # Reject cross-origin polling URLs (see sync path). VERIA-51.
+                try:
+                    assert_same_origin(operation_url, str(raw_response.request.url))
+                except SSRFError as ssrf_err:
+                    raise ValueError(
+                        f"Azure Document Intelligence: rejected polling URL ({ssrf_err})"
+                    )
+
                 # Get headers for polling (need auth)
                 poll_headers = {
                     "Ocp-Apim-Subscription-Key": raw_response.request.headers.get(

litellm/llms/base_llm/rerank/transformation.py

@@ -33,6 +33,7 @@ class BaseRerankConfig(ABC):
         model: str,
         optional_rerank_params: Dict,
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         return {}
 

litellm/llms/black_forest_labs/image_edit/handler.py

@@ -15,6 +15,7 @@ import httpx
 import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     HTTPHandler,
@@ -331,6 +332,17 @@ class BlackForestLabsImageEdit:
                 message="No polling_url in BFL response",
             )
 
+        # Reject cross-origin polling URLs — the ``x-key`` auth header
+        # would otherwise leak to whatever URL the upstream returns.
+        # VERIA-51.
+        try:
+            assert_same_origin(polling_url, str(initial_response.request.url))
+        except SSRFError as ssrf_err:
+            raise BlackForestLabsError(
+                status_code=502,
+                message=f"Rejected polling URL: {ssrf_err}",
+            )
+
         # Get just the auth header for polling
         polling_headers = {"x-key": headers.get("x-key", "")}
 
@@ -416,6 +428,17 @@ class BlackForestLabsImageEdit:
                 message="No polling_url in BFL response",
             )
 
+        # Reject cross-origin polling URLs — the ``x-key`` auth header
+        # would otherwise leak to whatever URL the upstream returns.
+        # VERIA-51.
+        try:
+            assert_same_origin(polling_url, str(initial_response.request.url))
+        except SSRFError as ssrf_err:
+            raise BlackForestLabsError(
+                status_code=502,
+                message=f"Rejected polling URL: {ssrf_err}",
+            )
+
         # Get just the auth header for polling
         polling_headers = {"x-key": headers.get("x-key", "")}
 

litellm/llms/black_forest_labs/image_generation/handler.py

@@ -15,6 +15,7 @@ import httpx
 import litellm
 from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin
 from litellm.llms.custom_httpx.http_handler import (
     AsyncHTTPHandler,
     HTTPHandler,
@@ -317,6 +318,17 @@ class BlackForestLabsImageGeneration:
                 message="No polling_url in BFL response",
             )
 
+        # Reject cross-origin polling URLs — the ``x-key`` auth header
+        # would otherwise leak to whatever URL the upstream returns.
+        # VERIA-51.
+        try:
+            assert_same_origin(polling_url, str(initial_response.request.url))
+        except SSRFError as ssrf_err:
+            raise BlackForestLabsError(
+                status_code=502,
+                message=f"Rejected polling URL: {ssrf_err}",
+            )
+
         # Get just the auth header for polling
         polling_headers = {"x-key": headers.get("x-key", "")}
 
@@ -402,6 +414,17 @@ class BlackForestLabsImageGeneration:
                 message="No polling_url in BFL response",
             )
 
+        # Reject cross-origin polling URLs — the ``x-key`` auth header
+        # would otherwise leak to whatever URL the upstream returns.
+        # VERIA-51.
+        try:
+            assert_same_origin(polling_url, str(initial_response.request.url))
+        except SSRFError as ssrf_err:
+            raise BlackForestLabsError(
+                status_code=502,
+                message=f"Rejected polling URL: {ssrf_err}",
+            )
+
         # Get just the auth header for polling
         polling_headers = {"x-key": headers.get("x-key", "")}
 

litellm/llms/cohere/rerank/transformation.py

@@ -111,6 +111,7 @@ class CohereRerankConfig(BaseRerankConfig):
         model: str,
         optional_rerank_params: Dict,
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         if "query" not in optional_rerank_params:
             raise ValueError("query is required for Cohere rerank")

litellm/llms/cohere/rerank_v2/transformation.py

@@ -71,6 +71,7 @@ class CohereRerankV2Config(CohereRerankConfig):
         model: str,
         optional_rerank_params: Dict,
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         if "query" not in optional_rerank_params:
             raise ValueError("query is required for Cohere rerank")

litellm/llms/custom_httpx/llm_http_handler.py

@@ -1007,6 +1007,7 @@ class BaseLLMHTTPHandler:
         api_key: Optional[str] = None,
         api_base: Optional[str] = None,
         client: Optional[Union[HTTPHandler, AsyncHTTPHandler]] = None,
+        litellm_params: Optional[Dict[str, Any]] = None,
     ) -> RerankResponse:
         # get config from model, custom llm provider
         headers = provider_config.validate_environment(
@@ -1026,6 +1027,7 @@ class BaseLLMHTTPHandler:
             model=model,
             optional_rerank_params=optional_rerank_params,
             headers=headers,
+            litellm_params=litellm_params,
         )
 
         ## LOGGING
@@ -2535,10 +2537,16 @@ class BaseLLMHTTPHandler:
             },
         )
 
+        delete_kwargs: Dict[str, Any] = {
+            "url": url,
+            "headers": headers,
+            "timeout": timeout,
+        }
+        if data:
+            delete_kwargs["json"] = data
+
         try:
-            response = await async_httpx_client.delete(
-                url=url, headers=headers, json=data, timeout=timeout
-            )
+            response = await async_httpx_client.delete(**delete_kwargs)
 
         except Exception as e:
             raise self._handle_error(
@@ -2619,10 +2627,16 @@ class BaseLLMHTTPHandler:
             },
         )
 
+        delete_kwargs: Dict[str, Any] = {
+            "url": url,
+            "headers": headers,
+            "timeout": timeout,
+        }
+        if data:
+            delete_kwargs["json"] = data
+
         try:
-            response = sync_httpx_client.delete(
-                url=url, headers=headers, json=data, timeout=timeout
-            )
+            response = sync_httpx_client.delete(**delete_kwargs)
 
         except Exception as e:
             raise self._handle_error(

litellm/llms/deepinfra/rerank/transformation.py

@@ -132,6 +132,7 @@ class DeepinfraRerankConfig(BaseRerankConfig):
         model: str,
         optional_rerank_params: Dict,
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         # Convert OptionalRerankParams to dict as expected by parent class
         if optional_rerank_params is None:

litellm/llms/fireworks_ai/rerank/transformation.py

@@ -127,6 +127,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig):
         model: str,
         optional_rerank_params: Dict,
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         """
         Transform request to Fireworks AI rerank format

litellm/llms/hosted_vllm/rerank/transformation.py

@@ -121,6 +121,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
         model: str,
         optional_rerank_params: Dict,
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         if "query" not in optional_rerank_params:
             raise ValueError("query is required for Hosted VLLM rerank")

litellm/llms/huggingface/rerank/transformation.py

@@ -146,6 +146,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig):
         model: str,
         optional_rerank_params: Union[OptionalRerankParams, dict],
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         if "query" not in optional_rerank_params:
             raise ValueError("query is required for HuggingFace rerank")

litellm/llms/jina_ai/rerank/transformation.py

@@ -74,7 +74,11 @@ class JinaAIRerankConfig(BaseRerankConfig):
         return cleaned_base
 
     def transform_rerank_request(
-        self, model: str, optional_rerank_params: Dict, headers: Dict
+        self,
+        model: str,
+        optional_rerank_params: Dict,
+        headers: Dict,
+        litellm_params: Optional[dict] = None,
     ) -> Dict:
         return {"model": model, **optional_rerank_params}
 

litellm/llms/nvidia_nim/rerank/ranking_transformation.py

@@ -66,6 +66,7 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
         model: str,
         optional_rerank_params: Dict,
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         """
         Transform request, using clean model name without 'ranking/' prefix.
@@ -75,4 +76,5 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig):
             model=clean_model,
             optional_rerank_params=optional_rerank_params,
             headers=headers,
+            litellm_params=litellm_params,
         )

litellm/llms/nvidia_nim/rerank/transformation.py

@@ -177,6 +177,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig):
         model: str,
         optional_rerank_params: Dict,
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         """
         Transform request to Nvidia NIM format.

litellm/llms/vertex_ai/common_utils.py

@@ -27,6 +27,53 @@ class VertexAIError(BaseLLMException):
         super().__init__(message=message, status_code=status_code, headers=headers)
 
 
+def vertex_request_labels_from_litellm_params(
+    litellm_params: Optional[dict],
+) -> Optional[Dict[str, str]]:
+    """
+    Build Vertex/GCP billing labels from LiteLLM user metadata on ``litellm_params``:
+    ``metadata`` (``completion(..., metadata=...)``) or ``litellm_metadata``,
+    using ``requester_metadata`` string key-value pairs (same convention as Gemini).
+    ``metadata`` is tried first when both are present.
+    """
+    if not litellm_params:
+        return None
+    for key in ("metadata", "litellm_metadata"):
+        if key not in litellm_params:
+            continue
+        metadata = litellm_params[key]
+        if metadata is None or not isinstance(metadata, dict):
+            continue
+        if "requester_metadata" not in metadata:
+            continue
+        rm = metadata["requester_metadata"]
+        if not isinstance(rm, dict):
+            continue
+        labels = {k: v for k, v in rm.items() if isinstance(v, str)}
+        if labels:
+            return labels
+    return None
+
+
+def pop_vertex_request_labels(
+    optional_params: Optional[dict],
+    litellm_params: Optional[dict],
+) -> Optional[Dict[str, str]]:
+    """
+    Resolve labels from optional ``labels`` (Gemini-style) and/or
+    ``litellm_params["metadata"]`` / ``litellm_params["litellm_metadata"]``
+    (``requester_metadata``). Pops ``labels`` from optional_params when present.
+    """
+    labels: Optional[Dict[str, str]] = None
+    if optional_params is not None and "labels" in optional_params:
+        raw = optional_params.pop("labels")
+        if isinstance(raw, dict):
+            labels = {k: v for k, v in raw.items() if isinstance(v, str)}
+    if not labels:
+        labels = vertex_request_labels_from_litellm_params(litellm_params)
+    return labels if labels else None
+
+
 class VertexAIModelRoute(str, Enum):
     """Enum for Vertex AI model routing"""
 
@@ -50,7 +97,7 @@ def get_vertex_ai_model_route(
     Determine which handler to use for a Vertex AI model based on the model name.
 
     Args:
-        model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "openai/gpt-oss-120b")
+        model: The model name (e.g., "llama3-405b", "gemini-pro", "gemma/gemma-3-12b-it", "xai/grok-4.1-fast-non-reasoning")
         litellm_params: Optional litellm parameters dict that may contain base_model for routing
 
     Returns:
@@ -66,7 +113,7 @@ def get_vertex_ai_model_route(
         >>> get_vertex_ai_model_route("gemma/gemma-3-12b-it")
         VertexAIModelRoute.GEMMA
 
-        >>> get_vertex_ai_model_route("openai/gpt-oss-120b")
+        >>> get_vertex_ai_model_route("xai/grok-4.1-fast-non-reasoning")
         VertexAIModelRoute.MODEL_GARDEN
 
         >>> get_vertex_ai_model_route("1234567890", {"api_base": "http://10.96.32.8"})
@@ -102,8 +149,11 @@ def get_vertex_ai_model_route(
     if "gemma/" in model:
         return VertexAIModelRoute.GEMMA
 
-    # Check for model garden openai models
-    if "openai" in model:
+    # Check for model garden OpenAI-compatible publisher models.
+    # Examples:
+    # - openai/gpt-oss-120b-maas
+    # - xai/grok-4.1-fast-non-reasoning
+    if "openai" in model or model.startswith("xai/"):
         return VertexAIModelRoute.MODEL_GARDEN
 
     # Check for gemini models
@@ -209,8 +259,8 @@ def get_vertex_base_model_name(model: str) -> str:
         >>> get_vertex_base_model_name("gemma/gemma-3-12b-it")
         "gemma-3-12b-it"
 
-        >>> get_vertex_base_model_name("openai/gpt-oss-120b")
-        "gpt-oss-120b"
+        >>> get_vertex_base_model_name("xai/grok-4.1-fast-non-reasoning")
+        "grok-4.1-fast-non-reasoning"
 
         >>> get_vertex_base_model_name("1234567890")
         "1234567890"

litellm/llms/vertex_ai/gemini/transformation.py

@@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
     response_schema_prompt,
 )
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
+from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
 from litellm.types.files import (
     get_file_mime_type_for_file_type,
     get_file_type_from_extension,
@@ -714,16 +715,8 @@ def _transform_request_body(  # noqa: PLR0915
         optional_params.pop("output_config", None)
         config_fields = GenerationConfig.__annotations__.keys()
 
-        # If the LiteLLM client sends Gemini-supported parameter "labels", add it
-        # as "labels" field to the request sent to the Gemini backend.
-        labels: Optional[dict[str, str]] = optional_params.pop("labels", None)
-        # If the LiteLLM client sends OpenAI-supported parameter "metadata", add it
-        # as "labels" field to the request sent to the Gemini backend.
-        if labels is None and "metadata" in litellm_params:
-            metadata = litellm_params["metadata"]
-            if metadata is not None and "requester_metadata" in metadata:
-                rm = metadata["requester_metadata"]
-                labels = {k: v for k, v in rm.items() if isinstance(v, str)}
+        # labels: optional explicit param and/or metadata.requester_metadata (OpenAI metadata)
+        labels = pop_vertex_request_labels(optional_params, litellm_params)
 
         filtered_params = {
             k: v

litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py

@@ -3,7 +3,7 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint
 """
 
 import json
-from typing import Any, Dict, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 import httpx
 
@@ -13,8 +13,8 @@ from litellm.llms.custom_httpx.http_handler import (
     HTTPHandler,
     get_async_httpx_client,
 )
-from litellm.types.llms.openai import EmbeddingInput
 from litellm.types.llms.vertex_ai import (
+    GeminiEmbeddingInput,
     VertexAIBatchEmbeddingsRequestBody,
     VertexAIBatchEmbeddingsResponseObject,
 )
@@ -23,7 +23,6 @@ from litellm.types.utils import EmbeddingResponse
 from ..gemini.vertex_and_google_ai_studio_gemini import VertexLLM
 from .batch_embed_content_transformation import (
     _is_file_reference,
-    _is_multimodal_input,
     process_embed_content_response,
     process_response,
     transform_openai_input_gemini_content,
@@ -32,9 +31,24 @@ from .batch_embed_content_transformation import (
 
 
 class GoogleBatchEmbeddings(VertexLLM):
+    @staticmethod
+    def _flatten_and_detect_file_refs(
+        input: GeminiEmbeddingInput,
+    ) -> Tuple[List[str], bool]:
+        """Flatten nested input lists and detect file references."""
+        input_list = [input] if isinstance(input, str) else input
+        flat_elements = [
+            e
+            for item in input_list
+            for e in (item if isinstance(item, list) else [item])
+            if isinstance(e, str)
+        ]
+        has_file_refs = any(_is_file_reference(e) for e in flat_elements)
+        return flat_elements, has_file_refs
+
     def _resolve_file_references(
         self,
-        input: EmbeddingInput,
+        input: GeminiEmbeddingInput,
         api_key: str,
         sync_handler: HTTPHandler,
     ) -> Dict[str, Dict[str, str]]:
@@ -42,7 +56,7 @@ class GoogleBatchEmbeddings(VertexLLM):
         Resolve Gemini file references (files/...) to get mime_type and uri.
 
         Args:
-            input: EmbeddingInput that may contain file references
+            input: GeminiEmbeddingInput that may contain file references
             api_key: Gemini API key
             sync_handler: HTTP client
 
@@ -73,7 +87,7 @@ class GoogleBatchEmbeddings(VertexLLM):
 
     async def _async_resolve_file_references(
         self,
-        input: EmbeddingInput,
+        input: GeminiEmbeddingInput,
         api_key: str,
         async_handler: AsyncHTTPHandler,
     ) -> Dict[str, Dict[str, str]]:
@@ -81,7 +95,7 @@ class GoogleBatchEmbeddings(VertexLLM):
         Async version of _resolve_file_references.
 
         Args:
-            input: EmbeddingInput that may contain file references
+            input: GeminiEmbeddingInput that may contain file references
             api_key: Gemini API key
             async_handler: Async HTTP client
 
@@ -110,10 +124,10 @@ class GoogleBatchEmbeddings(VertexLLM):
 
         return resolved_files
 
-    def batch_embeddings(
+    def batch_embeddings(  # noqa: PLR0915
         self,
         model: str,
-        input: EmbeddingInput,
+        input: GeminiEmbeddingInput,
         print_verbose,
         model_response: EmbeddingResponse,
         custom_llm_provider: Literal["gemini", "vertex_ai"],
@@ -151,8 +165,7 @@ class GoogleBatchEmbeddings(VertexLLM):
 
         optional_params = optional_params or {}
 
-        is_multimodal = _is_multimodal_input(input)
-        use_embed_content = is_multimodal or (custom_llm_provider == "vertex_ai")
+        use_embed_content = custom_llm_provider == "vertex_ai"
         mode: Literal["embedding", "batch_embedding"]
         if use_embed_content:
             mode = "embedding"
@@ -215,8 +228,22 @@ class GoogleBatchEmbeddings(VertexLLM):
                 resolved_files=resolved_files,
             )
         else:
+            flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input)
+            if has_file_refs and not api_key:
+                raise ValueError(
+                    "An API key is required to resolve Gemini file references (files/...). "
+                    "Pass api_key= or set GEMINI_API_KEY."
+                )
+            resolved_files = {}
+            if api_key and has_file_refs:
+                resolved_files = self._resolve_file_references(
+                    input=flat_elements, api_key=api_key, sync_handler=sync_handler
+                )
             request_data = transform_openai_input_gemini_content(
-                input=input, model=model, optional_params=optional_params
+                input=input,
+                model=model,
+                optional_params=optional_params,
+                resolved_files=resolved_files,
             )
 
         ## LOGGING
@@ -264,7 +291,7 @@ class GoogleBatchEmbeddings(VertexLLM):
         url: str,
         data: Optional[Union[VertexAIBatchEmbeddingsRequestBody, dict]],
         model_response: EmbeddingResponse,
-        input: EmbeddingInput,
+        input: GeminiEmbeddingInput,
         timeout: Optional[Union[float, httpx.Timeout]],
         headers={},
         client: Optional[AsyncHTTPHandler] = None,
@@ -303,8 +330,22 @@ class GoogleBatchEmbeddings(VertexLLM):
                 resolved_files=resolved_files,
             )
         else:
+            flat_elements, has_file_refs = self._flatten_and_detect_file_refs(input)
+            if has_file_refs and not api_key:
+                raise ValueError(
+                    "An API key is required to resolve Gemini file references (files/...). "
+                    "Pass api_key= or set GEMINI_API_KEY."
+                )
+            resolved_files = {}
+            if api_key and has_file_refs:
+                resolved_files = await self._async_resolve_file_references(
+                    input=flat_elements, api_key=api_key, async_handler=async_handler
+                )
             data = transform_openai_input_gemini_content(
-                input=input, model=model, optional_params=optional_params or {}
+                input=input,
+                model=model,
+                optional_params=optional_params or {},
+                resolved_files=resolved_files,
             )
 
         ## LOGGING

litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py

@@ -6,12 +6,12 @@ Why separate file? Make it easy to see how transformation works
 
 from typing import Dict, List, Optional, Tuple
 
-from litellm.types.llms.openai import EmbeddingInput
 from litellm.types.llms.vertex_ai import (
     BlobType,
     ContentType,
     EmbedContentRequest,
     FileDataType,
+    GeminiEmbeddingInput,
     PartType,
     VertexAIBatchEmbeddingsRequestBody,
     VertexAIBatchEmbeddingsResponseObject,
@@ -114,33 +114,77 @@ def _parse_data_url(data_url: str) -> Tuple[str, str]:
     return media_type, base64_data
 
 
-def _is_multimodal_input(input: EmbeddingInput) -> bool:
+def _is_multimodal_input(input: GeminiEmbeddingInput) -> bool:
     """
-    Check if the input contains multimodal data (data URIs, file references, or GCS URLs).
+    Check if the input contains multimodal data (data URIs, file references,
+    GCS URLs, or nested lists for combined embeddings).
 
     Args:
-        input: EmbeddingInput (str or List[str])
+        input: GeminiEmbeddingInput — str, List[str], or List[List[str]] for combined embeddings
 
     Returns:
-        bool: True if any element is a data URI, file reference, or GCS URL
+        bool: True if any element is multimodal or a nested list
     """
     if isinstance(input, str):
-        input_list = [input]
-    else:
-        input_list = input
+        return _is_multimodal_element(input)
 
-    for element in input_list:
-        if isinstance(element, str):
-            if element.startswith("data:") and ";base64," in element:
-                return True
-            if _is_file_reference(element):
-                return True
-            if _is_gcs_url(element):
+    for element in input:
+        if isinstance(element, list):
+            if any(
+                _is_multimodal_element(sub) for sub in element if isinstance(sub, str)
+            ):
                 return True
+        elif isinstance(element, str) and _is_multimodal_element(element):
+            return True
 
     return False
 
 
+def _is_multimodal_element(element: str) -> bool:
+    """Check if a single string element is multimodal."""
+    if element.startswith("data:") and ";base64," in element:
+        return True
+    if _is_file_reference(element):
+        return True
+    if _is_gcs_url(element):
+        return True
+    return False
+
+
+def _build_part_for_input(
+    element: str,
+    resolved_files: Optional[Dict[str, Dict[str, str]]] = None,
+) -> PartType:
+    """
+    Build a single PartType for an input element, handling text, data URIs,
+    file references, and GCS URLs.
+    """
+    resolved_files = resolved_files or {}
+
+    if element.startswith("data:") and ";base64," in element:
+        mime_type, base64_data = _parse_data_url(element)
+        blob: BlobType = {"mime_type": mime_type, "data": base64_data}
+        return PartType(inline_data=blob)
+    elif _is_gcs_url(element):
+        mime_type = _infer_mime_type_from_gcs_url(element)
+        file_data: FileDataType = {
+            "mime_type": mime_type,
+            "file_uri": element,
+        }
+        return PartType(file_data=file_data)
+    elif _is_file_reference(element):
+        if element not in resolved_files:
+            raise ValueError(f"File reference {element} not resolved")
+        file_info = resolved_files[element]
+        file_data_ref: FileDataType = {
+            "mime_type": file_info["mime_type"],
+            "file_uri": file_info["uri"],
+        }
+        return PartType(file_data=file_data_ref)
+    else:
+        return PartType(text=element)
+
+
 _SUPPORTED_EMBED_PARAMS = {"outputDimensionality", "taskType", "title"}
 
 
@@ -155,37 +199,60 @@ def _filter_embed_params(optional_params: dict) -> dict:
 
 
 def transform_openai_input_gemini_content(
-    input: EmbeddingInput, model: str, optional_params: dict
+    input: GeminiEmbeddingInput,
+    model: str,
+    optional_params: dict,
+    resolved_files: Optional[Dict[str, Dict[str, str]]] = None,
 ) -> VertexAIBatchEmbeddingsRequestBody:
     """
-    The content to embed. Only the parts.text fields will be counted.
+    Transform OpenAI embedding input to Gemini batchEmbedContents format.
+
+    Each input element becomes a separate EmbedContentRequest, supporting
+    text, data URIs, file references, and GCS URLs.
+
+    If an element is a list (nested input), all sub-elements are combined
+    into a single content with multiple parts, producing one combined
+    embedding for the group.
+
+    Examples:
+        input=["text", "image"]         → 2 separate embeddings
+        input=[["text", "image"]]       → 1 combined embedding
+        input=[["text", "image"], "x"]  → 2 embeddings (1 combined + 1 separate)
     """
     gemini_model_name = "models/{}".format(model)
 
     gemini_params = _filter_embed_params(optional_params)
 
+    input_list = [input] if isinstance(input, str) else input
     requests: List[EmbedContentRequest] = []
-    if isinstance(input, str):
+
+    for element in input_list:
+        if isinstance(element, list):
+            if not element:
+                raise ValueError("Nested input list must not be empty")
+            for sub in element:
+                if not isinstance(sub, str):
+                    raise ValueError(
+                        f"Elements inside a nested input list must be strings, got {type(sub)}"
+                    )
+            parts = [
+                _build_part_for_input(sub, resolved_files=resolved_files)
+                for sub in element
+            ]
+        else:
+            parts = [_build_part_for_input(element, resolved_files=resolved_files)]
         request = EmbedContentRequest(
             model=gemini_model_name,
-            content=ContentType(parts=[PartType(text=input)]),
+            content=ContentType(parts=parts),
             **gemini_params,
         )
         requests.append(request)
-    else:
-        for i in input:
-            request = EmbedContentRequest(
-                model=gemini_model_name,
-                content=ContentType(parts=[PartType(text=i)]),
-                **gemini_params,
-            )
-            requests.append(request)
 
     return VertexAIBatchEmbeddingsRequestBody(requests=requests)
 
 
 def transform_openai_input_gemini_embed_content(
-    input: EmbeddingInput,
+    input: GeminiEmbeddingInput,
     model: str,
     optional_params: dict,
     resolved_files: Optional[Dict[str, Dict[str, str]]] = None,
@@ -194,7 +261,7 @@ def transform_openai_input_gemini_embed_content(
     Transform OpenAI embedding input to Gemini embedContent format (multimodal).
 
     Args:
-        input: EmbeddingInput (str or List[str]) with text, data URIs, or file references
+        input: GeminiEmbeddingInput with text, data URIs, or file references
         model: Model name
         optional_params: Additional parameters (taskType, outputDimensionality, etc.)
         resolved_files: Dict mapping file names (files/abc) to {mime_type, uri}
@@ -210,31 +277,14 @@ def transform_openai_input_gemini_embed_content(
     parts: List[PartType] = []
 
     for element in input_list:
+        if isinstance(element, list):
+            raise ValueError(
+                "Nested (combined) embeddings are not supported on the embedContent path. "
+                "Use the batchEmbedContents path or pass a flat list instead."
+            )
         if not isinstance(element, str):
             raise ValueError(f"Unsupported input type: {type(element)}")
-
-        if element.startswith("data:") and ";base64," in element:
-            mime_type, base64_data = _parse_data_url(element)
-            blob: BlobType = {"mime_type": mime_type, "data": base64_data}
-            parts.append(PartType(inline_data=blob))
-        elif _is_gcs_url(element):
-            mime_type = _infer_mime_type_from_gcs_url(element)
-            file_data: FileDataType = {
-                "mime_type": mime_type,
-                "file_uri": element,
-            }
-            parts.append(PartType(file_data=file_data))
-        elif _is_file_reference(element):
-            if element not in resolved_files:
-                raise ValueError(f"File reference {element} not resolved")
-            file_info = resolved_files[element]
-            file_data_ref: FileDataType = {
-                "mime_type": file_info["mime_type"],
-                "file_uri": file_info["uri"],
-            }
-            parts.append(PartType(file_data=file_data_ref))
-        else:
-            parts.append(PartType(text=element))
+        parts.append(_build_part_for_input(element, resolved_files=resolved_files))
 
     request_body: dict = {
         "content": ContentType(parts=parts),
@@ -245,7 +295,7 @@ def transform_openai_input_gemini_embed_content(
 
 
 def process_embed_content_response(
-    input: EmbeddingInput,
+    input: GeminiEmbeddingInput,
     model_response: EmbeddingResponse,
     model: str,
     response_json: dict,
@@ -291,7 +341,7 @@ def process_embed_content_response(
 
 
 def process_response(
-    input: EmbeddingInput,
+    input: GeminiEmbeddingInput,
     model_response: EmbeddingResponse,
     model: str,
     _predictions: VertexAIBatchEmbeddingsResponseObject,
@@ -308,8 +358,29 @@ def process_response(
     model_response.data = openai_embeddings
     model_response.model = model
 
-    input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
-    prompt_tokens = token_counter(model=model, text=input_text)
+    has_nested = isinstance(input, list) and any(isinstance(e, list) for e in input)
+    if _is_multimodal_input(input) or has_nested:
+        input_list = input if isinstance(input, list) else [input]
+        text_elements: List[str] = []
+        for e in input_list:
+            if isinstance(e, list):
+                text_elements.extend(
+                    sub
+                    for sub in e
+                    if isinstance(sub, str) and not _is_multimodal_element(sub)
+                )
+            elif isinstance(e, str) and not _is_multimodal_element(e):
+                text_elements.append(e)
+        if text_elements:
+            input_text = get_formatted_prompt(
+                data={"input": text_elements}, call_type="embedding"
+            )
+            prompt_tokens = token_counter(model=model, text=input_text)
+        else:
+            prompt_tokens = 0
+    else:
+        input_text = get_formatted_prompt(data={"input": input}, call_type="embedding")
+        prompt_tokens = token_counter(model=model, text=input_text)
     model_response.usage = Usage(
         prompt_tokens=prompt_tokens, total_tokens=prompt_tokens
     )

litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py

@@ -7,7 +7,10 @@ import litellm
 from litellm.llms.base_llm.image_generation.transformation import (
     BaseImageGenerationConfig,
 )
-from litellm.llms.vertex_ai.common_utils import get_vertex_base_url
+from litellm.llms.vertex_ai.common_utils import (
+    get_vertex_base_url,
+    pop_vertex_request_labels,
+)
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.llms.openai import (
@@ -203,13 +206,16 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
             "sampleCount": 1,
         }
 
-        # Merge with optional params
+        labels = pop_vertex_request_labels(optional_params, litellm_params)
+        # Merge with optional params (after popping labels so they are not sent as Imagen parameters)
         parameters = {**default_params, **optional_params}
 
-        request_body = {
+        request_body: dict = {
             "instances": [{"prompt": prompt}],
             "parameters": parameters,
         }
+        if labels:
+            request_body["labels"] = labels
 
         return request_body
 

litellm/llms/vertex_ai/rerank/transformation.py

@@ -11,12 +11,15 @@ import httpx
 import litellm
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
+from litellm.llms.vertex_ai.common_utils import (
+    vertex_request_labels_from_litellm_params,
+)
 from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
 from litellm.secret_managers.main import get_secret_str
 from litellm.types.rerank import (
+    RerankBilledUnits,
     RerankResponse,
     RerankResponseMeta,
-    RerankBilledUnits,
     RerankResponseResult,
 )
 
@@ -109,6 +112,7 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
         model: str,
         optional_rerank_params: Dict,
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         """
         Transform the request from Cohere format to Vertex AI Discovery Engine format
@@ -145,6 +149,10 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase):
         # When return_documents is False, we want to ignore record details (return only IDs)
         request_data["ignoreRecordDetailsInResponse"] = not return_documents
 
+        user_labels = vertex_request_labels_from_litellm_params(litellm_params)
+        if user_labels:
+            request_data["userLabels"] = user_labels
+
         return request_data
 
     def transform_rerank_response(

litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py

@@ -1,4 +1,4 @@
-from typing import Literal, Optional, Union
+from typing import Dict, Literal, Optional, Union
 
 import httpx
 
@@ -44,6 +44,7 @@ class VertexEmbedding(VertexBase):
         vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES] = None,
         gemini_api_key: Optional[str] = None,
         extra_headers: Optional[dict] = None,
+        litellm_params: Optional[Dict] = None,
     ) -> EmbeddingResponse:
         if aembedding is True:
             return self.async_embedding(  # type: ignore
@@ -61,6 +62,7 @@ class VertexEmbedding(VertexBase):
                 vertex_credentials=vertex_credentials,
                 gemini_api_key=gemini_api_key,
                 extra_headers=extra_headers,
+                litellm_params=litellm_params,
             )
 
         should_use_v1beta1_features = self.is_using_v1beta1_features(
@@ -92,7 +94,10 @@ class VertexEmbedding(VertexBase):
         headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers)
         vertex_request: VertexEmbeddingRequest = (
             litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request(
-                input=input, optional_params=optional_params, model=model
+                input=input,
+                optional_params=optional_params,
+                model=model,
+                litellm_params=litellm_params,
             )
         )
 
@@ -156,6 +161,7 @@ class VertexEmbedding(VertexBase):
         gemini_api_key: Optional[str] = None,
         extra_headers: Optional[dict] = None,
         encoding=None,
+        litellm_params: Optional[Dict] = None,
     ) -> EmbeddingResponse:
         """
         Async embedding implementation
@@ -188,7 +194,10 @@ class VertexEmbedding(VertexBase):
         headers = self.set_headers(auth_header=auth_header, extra_headers=extra_headers)
         vertex_request: VertexEmbeddingRequest = (
             litellm.vertexAITextEmbeddingConfig.transform_openai_request_to_vertex_embedding_request(
-                input=input, optional_params=optional_params, model=model
+                input=input,
+                optional_params=optional_params,
+                model=model,
+                litellm_params=litellm_params,
             )
         )
 

litellm/llms/vertex_ai/vertex_embeddings/transformation.py

@@ -3,6 +3,7 @@ from typing import List, Literal, Optional, Union
 
 from pydantic import BaseModel
 
+from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
 from litellm.types.utils import EmbeddingResponse, Usage
 
 from .types import *
@@ -100,7 +101,11 @@ class VertexAITextEmbeddingConfig(BaseModel):
         return optional_params
 
     def transform_openai_request_to_vertex_embedding_request(
-        self, input: Union[list, str], optional_params: dict, model: str
+        self,
+        input: Union[list, str],
+        optional_params: dict,
+        model: str,
+        litellm_params: Optional[dict] = None,
     ) -> VertexEmbeddingRequest:
         """
         Transforms an openai request to a vertex embedding request.
@@ -108,16 +113,26 @@ class VertexAITextEmbeddingConfig(BaseModel):
         # Import here to avoid circular import issues with litellm.__init__
         from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig
 
+        labels = pop_vertex_request_labels(optional_params, litellm_params)
+
         if model.isdigit():
-            return self._transform_openai_request_to_fine_tuned_embedding_request(
-                input, optional_params, model
+            vertex_request = (
+                self._transform_openai_request_to_fine_tuned_embedding_request(
+                    input, optional_params, model
+                )
             )
+            if labels:
+                vertex_request["labels"] = labels
+            return vertex_request
         if VertexBGEConfig.is_bge_model(model):
-            return VertexBGEConfig.transform_request(
+            vertex_request = VertexBGEConfig.transform_request(
                 input=input, optional_params=optional_params, model=model
             )
+            if labels:
+                vertex_request["labels"] = labels
+            return vertex_request
 
-        vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest()
+        vertex_request = VertexEmbeddingRequest()
         vertex_text_embedding_input_list: List[TextEmbeddingInput] = []
         task_type: Optional[TaskType] = optional_params.get("task_type")
         title = optional_params.get("title")
@@ -133,6 +148,8 @@ class VertexAITextEmbeddingConfig(BaseModel):
 
         vertex_request["instances"] = vertex_text_embedding_input_list
         vertex_request["parameters"] = EmbeddingParameters(**optional_params)
+        if labels:
+            vertex_request["labels"] = labels
 
         return vertex_request
 

litellm/llms/vertex_ai/vertex_embeddings/types.py

@@ -3,7 +3,7 @@ Types for Vertex Embeddings Requests
 """
 
 from enum import Enum
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from typing_extensions import TypedDict
 
@@ -56,6 +56,7 @@ class VertexEmbeddingRequest(TypedDict, total=False):
         List[TextEmbeddingFineTunedInput],
     ]
     parameters: Optional[Union[EmbeddingParameters, TextEmbeddingFineTunedParameters]]
+    labels: Optional[Dict[str, str]]
 
 
 # Example usage:

litellm/llms/vertex_ai/vertex_model_garden/main.py

@@ -27,6 +27,17 @@ from ..common_utils import VertexAIError, get_vertex_base_model_name
 from ..vertex_llm_base import VertexBase
 
 
+def _vertex_model_garden_model_id_in_json_body(model: str) -> bool:
+    """
+    Vertex catalog / publisher models are addressed as publisher/model (e.g.
+    xai/grok-4.1-fast-reasoning) on the shared OpenAPI URL, with the id in the JSON body.
+
+    Deployed Model Garden endpoints are typically a single segment (often numeric)
+    and use .../endpoints/{ENDPOINT_ID}/chat/completions with an empty model field.
+    """
+    return "/" in model
+
+
 def create_vertex_url(
     vertex_location: str,
     vertex_project: str,
@@ -34,8 +45,13 @@ def create_vertex_url(
     model: str,
     api_base: Optional[str] = None,
 ) -> str:
-    """Return the base url for the vertex garden models"""
+    """Return the api base for vertex model garden (without /chat/completions)."""
     base_url = get_vertex_base_url(vertex_location)
+    if _vertex_model_garden_model_id_in_json_body(model):
+        return (
+            f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}"
+            "/endpoints/openapi"
+        )
     return f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/endpoints/{model}"
 
 
@@ -129,7 +145,10 @@ class VertexAIModelGardenModels(VertexBase):
                 vertex_location=vertex_location or "us-central1",
                 vertex_api_version="v1beta1",
             )
-            model = ""
+            # Publisher/catalog models: model id must be sent in the JSON body (OpenAPI route).
+            # Single-segment endpoint ids: model is encoded in the URL path; body model stays empty.
+            if not _vertex_model_garden_model_id_in_json_body(model):
+                model = ""
             return openai_like_chat_completions.completion(
                 model=model,
                 messages=messages,

litellm/llms/voyage/rerank/transformation.py

@@ -67,7 +67,11 @@ class VoyageRerankConfig(BaseRerankConfig):
         return api_base
 
     def transform_rerank_request(
-        self, model: str, optional_rerank_params: Dict, headers: Dict
+        self,
+        model: str,
+        optional_rerank_params: Dict,
+        headers: Dict,
+        litellm_params: Optional[dict] = None,
     ) -> Dict:
         return {"model": model, **optional_rerank_params}
 

litellm/llms/watsonx/rerank/transformation.py

@@ -143,6 +143,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig):
         model: str,
         optional_rerank_params: Dict,
         headers: dict,
+        litellm_params: Optional[dict] = None,
     ) -> dict:
         """
         Transform request to IBM watsonx.ai rerank format

litellm/llms/xai/chat/transformation.py

@@ -43,6 +43,7 @@ class XAIChatConfig(OpenAIGPTConfig):
             "logprobs",
             "max_tokens",
             "n",
+            "parallel_tool_calls",
             "presence_penalty",
             "response_format",
             "seed",

litellm/main.py

@@ -5311,6 +5311,7 @@ def embedding(  # noqa: PLR0915
                     api_key=api_key,
                     api_base=api_base,
                     client=client,
+                    litellm_params=litellm_params_dict,
                 )
         elif custom_llm_provider == "oobabooga":
             response = oobabooga.embedding(

litellm/model_prices_and_context_window_backup.json

@@ -33337,6 +33337,72 @@
         "source": "https://console.cloud.google.com/vertex-ai/publishers/openai/model-garden/gpt-oss-120b-maas",
         "supports_reasoning": true
     },
+    "vertex_ai/xai/grok-4.1-fast-non-reasoning": {
+        "cache_read_input_token_cost": 5e-08,
+        "input_cost_per_token": 2e-07,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 5e-07,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "vertex_ai/xai/grok-4.1-fast-reasoning": {
+        "cache_read_input_token_cost": 5e-08,
+        "input_cost_per_token": 2e-07,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 5e-07,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "vertex_ai/xai/grok-4.20-non-reasoning": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
+    "vertex_ai/xai/grok-4.20-reasoning": {
+        "cache_read_input_token_cost": 2e-07,
+        "input_cost_per_token": 2e-06,
+        "litellm_provider": "vertex_ai",
+        "max_input_tokens": 2000000,
+        "max_output_tokens": 2000000,
+        "max_tokens": 2000000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-06,
+        "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_web_search": true
+    },
     "vertex_ai/qwen/qwen3-235b-a22b-instruct-2507-maas": {
         "input_cost_per_token": 2.5e-07,
         "litellm_provider": "vertex_ai-qwen_models",

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -169,6 +169,37 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]:
 class MCPServerManager:
     _STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$")
 
+    @staticmethod
+    def _resolve_oauth2_flow(
+        *,
+        auth_type: Optional[MCPAuthType],
+        oauth2_flow: Optional[str],
+        token_url: Optional[str],
+        authorization_url: Optional[str],
+        client_id: Optional[str],
+        client_secret: Optional[str],
+    ) -> Optional[Literal["client_credentials", "authorization_code"]]:
+        """Infer oauth2_flow for legacy records that omit the field.
+
+        DB rows created before oauth2_flow support may have OAuth2 client
+        credentials + token_url but a null oauth2_flow. Treat these as M2M,
+        unless authorization_url is present (interactive OAuth).
+        """
+        if oauth2_flow in ("client_credentials", "authorization_code"):
+            return cast(
+                Literal["client_credentials", "authorization_code"], oauth2_flow
+            )
+        if oauth2_flow:
+            # Ignore unknown/untyped values and continue legacy inference.
+            return None
+        if auth_type != MCPAuth.oauth2:
+            return None
+        if authorization_url:
+            return None
+        if token_url and client_id and client_secret:
+            return "client_credentials"
+        return None
+
     def __init__(self):
         self.registry: Dict[str, MCPServer] = {}
         self.config_mcp_servers: Dict[str, MCPServer] = {}
@@ -342,7 +373,14 @@ class MCPServerManager:
                 # oauth specific fields
                 client_id=server_config.get("client_id", None),
                 client_secret=server_config.get("client_secret", None),
-                oauth2_flow=server_config.get("oauth2_flow", None),
+                oauth2_flow=self._resolve_oauth2_flow(
+                    auth_type=auth_type,
+                    oauth2_flow=server_config.get("oauth2_flow", None),
+                    token_url=resolved_token_url,
+                    authorization_url=resolved_authorization_url,
+                    client_id=server_config.get("client_id", None),
+                    client_secret=server_config.get("client_secret", None),
+                ),
                 scopes=resolved_scopes,
                 authorization_url=resolved_authorization_url,
                 token_url=resolved_token_url,
@@ -679,7 +717,17 @@ class MCPServerManager:
             client_id=client_id_value or getattr(mcp_server, "client_id", None),
             client_secret=client_secret_value
             or getattr(mcp_server, "client_secret", None),
-            oauth2_flow=getattr(mcp_server, "oauth2_flow", None),
+            oauth2_flow=self._resolve_oauth2_flow(
+                auth_type=auth_type,
+                oauth2_flow=getattr(mcp_server, "oauth2_flow", None),
+                token_url=mcp_server.token_url
+                or getattr(mcp_oauth_metadata, "token_url", None),
+                authorization_url=mcp_server.authorization_url
+                or getattr(mcp_oauth_metadata, "authorization_url", None),
+                client_id=client_id_value or getattr(mcp_server, "client_id", None),
+                client_secret=client_secret_value
+                or getattr(mcp_server, "client_secret", None),
+            ),
             scopes=resolved_scopes,
             authorization_url=mcp_server.authorization_url
             or getattr(mcp_oauth_metadata, "authorization_url", None),
@@ -2426,7 +2474,7 @@ class MCPServerManager:
             )
         )
 
-    async def _call_regular_mcp_tool(
+    async def _call_regular_mcp_tool(  # noqa: PLR0915
         self,
         mcp_server: MCPServer,
         original_tool_name: str,
@@ -2489,7 +2537,11 @@ class MCPServerManager:
         # oauth2 headers
         extra_headers: Optional[Dict[str, str]] = None
         if mcp_server.auth_type == MCPAuth.oauth2:
-            extra_headers = oauth2_headers
+            if mcp_server.has_client_credentials:
+                # For M2M OAuth servers, Authorization must come from token fetch.
+                extra_headers = None
+            else:
+                extra_headers = oauth2_headers
 
         if mcp_server.extra_headers and raw_headers:
             if extra_headers is None:
@@ -2501,6 +2553,11 @@ class MCPServerManager:
             for header in mcp_server.extra_headers:
                 if not isinstance(header, str):
                     continue
+                if (
+                    mcp_server.has_client_credentials
+                    and header.lower() == "authorization"
+                ):
+                    continue
                 header_value = normalized_raw_headers.get(header.lower())
                 if header_value is None:
                     continue
@@ -2536,6 +2593,10 @@ class MCPServerManager:
                     )
             extra_headers.update(hook_extra_headers)
 
+        # Reset to None if no headers were actually added
+        if extra_headers is not None and len(extra_headers) == 0:
+            extra_headers = None
+
         stdio_env = self._build_stdio_env(mcp_server, raw_headers)
 
         client = await self._create_mcp_client(

litellm/proxy/_experimental/mcp_server/server.py

@@ -153,6 +153,7 @@ if MCP_AVAILABLE:
         MCPAuthenticatedUser,
     )
     from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
+        MCPServerManager,
         global_mcp_server_manager,
     )
     from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
@@ -900,6 +901,20 @@ if MCP_AVAILABLE:
                 allowed_mcp_server_id
             )
             if mcp_server is not None:
+                # Apply oauth2_flow resolution for legacy DB rows where it may be NULL
+                resolved_flow = MCPServerManager._resolve_oauth2_flow(
+                    auth_type=mcp_server.auth_type,
+                    oauth2_flow=mcp_server.oauth2_flow,
+                    token_url=mcp_server.token_url,
+                    authorization_url=mcp_server.authorization_url,
+                    client_id=mcp_server.client_id,
+                    client_secret=mcp_server.client_secret,
+                )
+                if resolved_flow and resolved_flow != mcp_server.oauth2_flow:
+                    # Create a new instance with the resolved flow for this request
+                    mcp_server = mcp_server.model_copy(
+                        update={"oauth2_flow": resolved_flow}
+                    )
                 allowed_mcp_servers.append(mcp_server)
 
         if mcp_servers is not None:
@@ -1100,8 +1115,13 @@ if MCP_AVAILABLE:
 
         extra_headers: Optional[Dict[str, str]] = None
         if server.auth_type == MCPAuth.oauth2:
-            # Copy to avoid mutating the original dict (important for parallel fetching)
-            extra_headers = oauth2_headers.copy() if oauth2_headers else None
+            # For OAuth2 M2M servers, upstream Authorization must come from
+            # client_credentials token fetch, never from caller headers.
+            if server.has_client_credentials:
+                extra_headers = None
+            else:
+                # Copy to avoid mutating the original dict (important for parallel fetching)
+                extra_headers = oauth2_headers.copy() if oauth2_headers else None
 
         if server.extra_headers and raw_headers:
             if extra_headers is None:
@@ -1114,11 +1134,17 @@ if MCP_AVAILABLE:
             for header in server.extra_headers:
                 if not isinstance(header, str):
                     continue
+                if server.has_client_credentials and header.lower() == "authorization":
+                    continue
                 header_value = normalized_raw_headers.get(header.lower())
                 if header_value is None:
                     continue
                 extra_headers[header] = header_value
 
+        # Reset to None if no headers were actually added
+        if extra_headers is not None and len(extra_headers) == 0:
+            extra_headers = None
+
         if server_auth_header is None:
             server_auth_header = mcp_auth_header
 
@@ -1377,11 +1403,19 @@ if MCP_AVAILABLE:
                     spend_meta["per_server_tool_counts"] = per_server_tool_counts
 
                 end_time = datetime.now()
-                await litellm_logging_obj.async_success_handler(
-                    result=all_tools,
-                    start_time=list_tools_start_time,
-                    end_time=end_time,
-                )
+                try:
+                    await litellm_logging_obj.async_success_handler(
+                        result=all_tools,
+                        start_time=list_tools_start_time,
+                        end_time=end_time,
+                    )
+                except Exception as log_exc:
+                    # list_tools responses must not be dropped due to non-blocking
+                    # observability/serialization failures.
+                    verbose_logger.warning(
+                        "MCP list_tools success logging failed (continuing): %s",
+                        log_exc,
+                    )
 
             verbose_logger.info(
                 f"Successfully fetched {len(all_tools)} tools total from all MCP servers"

litellm/proxy/_lazy_openapi_snapshot.json

@@ -14008,7 +14008,7 @@
       "/mcp-rest/test/connection": {
         "post": {
           "description": "Test if we can connect to the provided MCP server before adding it",
-          "operationId": "test_connection_mcp_rest_test_connection_post",
+          "operationId": "test_connection_mcp_rest_test_connection_post_2",
           "requestBody": {
             "content": {
               "application/json": {
@@ -14053,7 +14053,7 @@
       "/mcp-rest/test/tools/list": {
         "post": {
           "description": "Preview tools available from MCP server before adding it",
-          "operationId": "test_tools_list_mcp_rest_test_tools_list_post",
+          "operationId": "test_tools_list_mcp_rest_test_tools_list_post_2",
           "requestBody": {
             "content": {
               "application/json": {
@@ -14098,7 +14098,7 @@
       "/mcp-rest/tools/call": {
         "post": {
           "description": "REST API to call a specific MCP tool with the provided arguments",
-          "operationId": "call_tool_rest_api_mcp_rest_tools_call_post",
+          "operationId": "call_tool_rest_api_mcp_rest_tools_call_post_2",
           "responses": {
             "200": {
               "content": {
@@ -14123,7 +14123,7 @@
       "/mcp-rest/tools/list": {
         "get": {
           "description": "List all available tools with information about the server they belong to.\n\nExample response:\n{\n    \"tools\": [\n        {\n            \"name\": \"create_zap\",\n            \"description\": \"Create a new zap\",\n            \"inputSchema\": \"tool_input_schema\",\n            \"mcp_info\": {\n                \"server_name\": \"zapier\",\n                \"logo_url\": \"https://www.zapier.com/logo.png\",\n            }\n        }\n    ],\n    \"error\": null,\n    \"message\": \"Successfully retrieved tools\"\n}",
-          "operationId": "list_tool_rest_api_mcp_rest_tools_list_get",
+          "operationId": "list_tool_rest_api_mcp_rest_tools_list_get_2",
           "parameters": [
             {
               "description": "The server id to list tools for",
@@ -21896,7 +21896,7 @@
       "/policies/usage/overview": {
         "get": {
           "description": "Return policy performance overview for the dashboard.",
-          "operationId": "policies_usage_overview_policies_usage_overview_get",
+          "operationId": "policies_usage_overview_policies_usage_overview_get_2",
           "parameters": [
             {
               "description": "YYYY-MM-DD",
@@ -22521,7 +22521,7 @@
       "/policies/attachments/estimate-impact": {
         "post": {
           "description": "Estimate how many keys and teams would be affected by a policy attachment.\n\nUse this before creating an attachment to preview the blast radius.\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/policies/attachments/estimate-impact\" \\\n    -H \"Authorization: Bearer <your_api_key>\" \\\n    -H \"Content-Type: application/json\" \\\n    -d '{\n        \"policy_name\": \"hipaa-compliance\",\n        \"tags\": [\"healthcare\", \"health-*\"]\n    }'\n```",
-          "operationId": "estimate_attachment_impact_policies_attachments_estimate_impact_post",
+          "operationId": "estimate_attachment_impact_policies_attachments_estimate_impact_post_2",
           "requestBody": {
             "content": {
               "application/json": {
@@ -22568,7 +22568,7 @@
       "/policies/resolve": {
         "post": {
           "description": "Resolve which policies and guardrails apply for a given context.\n\nUse this endpoint to debug \"what guardrails would apply to a request\nwith this team/key/model/tags combination?\"\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/policies/resolve\" \\\n    -H \"Authorization: Bearer <your_api_key>\" \\\n    -H \"Content-Type: application/json\" \\\n    -d '{\n        \"tags\": [\"healthcare\"],\n        \"model\": \"gpt-4\"\n    }'\n```",
-          "operationId": "resolve_policies_for_context_policies_resolve_post",
+          "operationId": "resolve_policies_for_context_policies_resolve_post_2",
           "parameters": [
             {
               "description": "Force a DB sync before resolving. Default uses in-memory cache.",
@@ -28329,7 +28329,7 @@
       "/v1/vector_stores": {
         "get": {
           "description": "List vector stores.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/list",
-          "operationId": "vector_store_list_v1_vector_stores_get",
+          "operationId": "vector_store_list_v1_vector_stores_get_2",
           "parameters": [
             {
               "in": "query",
@@ -28430,7 +28430,7 @@
         },
         "post": {
           "description": "Create a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/create\n\nSupports target_model_names parameter for creating vector stores across multiple models:\n```json\n{\n    \"name\": \"my-vector-store\",\n    \"target_model_names\": \"gpt-4,gemini-2.0\"\n}\n```",
-          "operationId": "vector_store_create_v1_vector_stores_post",
+          "operationId": "vector_store_create_v1_vector_stores_post_2",
           "responses": {
             "200": {
               "content": {
@@ -28455,7 +28455,7 @@
       "/v1/vector_stores/{vector_store_id}": {
         "delete": {
           "description": "Delete a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/delete",
-          "operationId": "vector_store_delete_v1_vector_stores__vector_store_id__delete",
+          "operationId": "vector_store_delete_v1_vector_stores__vector_store_id__delete_2",
           "parameters": [
             {
               "in": "path",
@@ -28499,7 +28499,7 @@
         },
         "get": {
           "description": "Retrieve a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/retrieve",
-          "operationId": "vector_store_retrieve_v1_vector_stores__vector_store_id__get",
+          "operationId": "vector_store_retrieve_v1_vector_stores__vector_store_id__get_2",
           "parameters": [
             {
               "in": "path",
@@ -28543,7 +28543,7 @@
         },
         "post": {
           "description": "Update a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/modify",
-          "operationId": "vector_store_update_v1_vector_stores__vector_store_id__post",
+          "operationId": "vector_store_update_v1_vector_stores__vector_store_id__post_2",
           "parameters": [
             {
               "in": "path",
@@ -28588,7 +28588,7 @@
       },
       "/v1/vector_stores/{vector_store_id}/files": {
         "get": {
-          "operationId": "vector_store_file_list_v1_vector_stores__vector_store_id__files_get",
+          "operationId": "vector_store_file_list_v1_vector_stores__vector_store_id__files_get_2",
           "parameters": [
             {
               "in": "path",
@@ -28631,7 +28631,7 @@
           ]
         },
         "post": {
-          "operationId": "vector_store_file_create_v1_vector_stores__vector_store_id__files_post",
+          "operationId": "vector_store_file_create_v1_vector_stores__vector_store_id__files_post_2",
           "parameters": [
             {
               "in": "path",
@@ -28676,7 +28676,7 @@
       },
       "/v1/vector_stores/{vector_store_id}/files/{file_id}": {
         "delete": {
-          "operationId": "vector_store_file_delete_v1_vector_stores__vector_store_id__files__file_id__delete",
+          "operationId": "vector_store_file_delete_v1_vector_stores__vector_store_id__files__file_id__delete_2",
           "parameters": [
             {
               "in": "path",
@@ -28728,7 +28728,7 @@
           ]
         },
         "get": {
-          "operationId": "vector_store_file_retrieve_v1_vector_stores__vector_store_id__files__file_id__get",
+          "operationId": "vector_store_file_retrieve_v1_vector_stores__vector_store_id__files__file_id__get_2",
           "parameters": [
             {
               "in": "path",
@@ -28780,7 +28780,7 @@
           ]
         },
         "post": {
-          "operationId": "vector_store_file_update_v1_vector_stores__vector_store_id__files__file_id__post",
+          "operationId": "vector_store_file_update_v1_vector_stores__vector_store_id__files__file_id__post_2",
           "parameters": [
             {
               "in": "path",
@@ -28834,7 +28834,7 @@
       },
       "/v1/vector_stores/{vector_store_id}/files/{file_id}/content": {
         "get": {
-          "operationId": "vector_store_file_content_v1_vector_stores__vector_store_id__files__file_id__content_get",
+          "operationId": "vector_store_file_content_v1_vector_stores__vector_store_id__files__file_id__content_get_2",
           "parameters": [
             {
               "in": "path",
@@ -28889,7 +28889,7 @@
       "/v1/vector_stores/{vector_store_id}/search": {
         "post": {
           "description": "Search a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/search",
-          "operationId": "vector_store_search_v1_vector_stores__vector_store_id__search_post",
+          "operationId": "vector_store_search_v1_vector_stores__vector_store_id__search_post_2",
           "parameters": [
             {
               "in": "path",
@@ -28935,7 +28935,7 @@
       "/vector_stores": {
         "get": {
           "description": "List vector stores.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/list",
-          "operationId": "vector_store_list_vector_stores_get",
+          "operationId": "vector_store_list_vector_stores_get_2",
           "parameters": [
             {
               "in": "query",
@@ -29036,7 +29036,7 @@
         },
         "post": {
           "description": "Create a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/create\n\nSupports target_model_names parameter for creating vector stores across multiple models:\n```json\n{\n    \"name\": \"my-vector-store\",\n    \"target_model_names\": \"gpt-4,gemini-2.0\"\n}\n```",
-          "operationId": "vector_store_create_vector_stores_post",
+          "operationId": "vector_store_create_vector_stores_post_2",
           "responses": {
             "200": {
               "content": {
@@ -29061,7 +29061,7 @@
       "/vector_stores/{vector_store_id}": {
         "delete": {
           "description": "Delete a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/delete",
-          "operationId": "vector_store_delete_vector_stores__vector_store_id__delete",
+          "operationId": "vector_store_delete_vector_stores__vector_store_id__delete_2",
           "parameters": [
             {
               "in": "path",
@@ -29105,7 +29105,7 @@
         },
         "get": {
           "description": "Retrieve a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/retrieve",
-          "operationId": "vector_store_retrieve_vector_stores__vector_store_id__get",
+          "operationId": "vector_store_retrieve_vector_stores__vector_store_id__get_2",
           "parameters": [
             {
               "in": "path",
@@ -29149,7 +29149,7 @@
         },
         "post": {
           "description": "Update a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/modify",
-          "operationId": "vector_store_update_vector_stores__vector_store_id__post",
+          "operationId": "vector_store_update_vector_stores__vector_store_id__post_2",
           "parameters": [
             {
               "in": "path",
@@ -29194,7 +29194,7 @@
       },
       "/vector_stores/{vector_store_id}/files": {
         "get": {
-          "operationId": "vector_store_file_list_vector_stores__vector_store_id__files_get",
+          "operationId": "vector_store_file_list_vector_stores__vector_store_id__files_get_2",
           "parameters": [
             {
               "in": "path",
@@ -29237,7 +29237,7 @@
           ]
         },
         "post": {
-          "operationId": "vector_store_file_create_vector_stores__vector_store_id__files_post",
+          "operationId": "vector_store_file_create_vector_stores__vector_store_id__files_post_2",
           "parameters": [
             {
               "in": "path",
@@ -29282,7 +29282,7 @@
       },
       "/vector_stores/{vector_store_id}/files/{file_id}": {
         "delete": {
-          "operationId": "vector_store_file_delete_vector_stores__vector_store_id__files__file_id__delete",
+          "operationId": "vector_store_file_delete_vector_stores__vector_store_id__files__file_id__delete_2",
           "parameters": [
             {
               "in": "path",
@@ -29334,7 +29334,7 @@
           ]
         },
         "get": {
-          "operationId": "vector_store_file_retrieve_vector_stores__vector_store_id__files__file_id__get",
+          "operationId": "vector_store_file_retrieve_vector_stores__vector_store_id__files__file_id__get_2",
           "parameters": [
             {
               "in": "path",
@@ -29386,7 +29386,7 @@
           ]
         },
         "post": {
-          "operationId": "vector_store_file_update_vector_stores__vector_store_id__files__file_id__post",
+          "operationId": "vector_store_file_update_vector_stores__vector_store_id__files__file_id__post_2",
           "parameters": [
             {
               "in": "path",
@@ -29440,7 +29440,7 @@
       },
       "/vector_stores/{vector_store_id}/files/{file_id}/content": {
         "get": {
-          "operationId": "vector_store_file_content_vector_stores__vector_store_id__files__file_id__content_get",
+          "operationId": "vector_store_file_content_vector_stores__vector_store_id__files__file_id__content_get_2",
           "parameters": [
             {
               "in": "path",
@@ -29495,7 +29495,7 @@
       "/vector_stores/{vector_store_id}/search": {
         "post": {
           "description": "Search a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/search",
-          "operationId": "vector_store_search_vector_stores__vector_store_id__search_post",
+          "operationId": "vector_store_search_vector_stores__vector_store_id__search_post_2",
           "parameters": [
             {
               "in": "path",

litellm/proxy/_lazy_openapi_snapshot.py

@@ -11,9 +11,32 @@ import json
 import re
 import sys
 from pathlib import Path
-from typing import Dict, Iterable, Optional
+from typing import Dict, Optional, Set
 
 SNAPSHOT_FILE = Path(__file__).parent / "_lazy_openapi_snapshot.json"
+HTTP_METHOD_SUFFIXES = {
+    "delete",
+    "get",
+    "head",
+    "options",
+    "patch",
+    "post",
+    "put",
+    "trace",
+}
+
+
+def _stabilize_multi_method_route_ids(routes) -> None:
+    """FastAPI derives route IDs from a set of methods; make snapshots stable."""
+
+    for route in routes:
+        methods = sorted(getattr(route, "methods", None) or [])
+        if len(methods) <= 1 or not getattr(route, "path_format", None):
+            continue
+
+        operation_id = f"{route.name}{route.path_format}"
+        operation_id = re.sub(r"\W", "_", operation_id)
+        route.unique_id = f"{operation_id}_{methods[0].lower()}"
 
 
 def load_snapshot() -> Optional[Dict[str, Dict]]:
@@ -26,22 +49,37 @@ def load_snapshot() -> Optional[Dict[str, Dict]]:
         return None
 
 
-def _stable_generate_unique_id(route) -> str:
-    operation_id = f"{route.name}{route.path_format}"
-    operation_id = re.sub(r"\W", "_", operation_id)
-    methods = sorted(route.methods or [])
-    if not methods:
-        return operation_id
-    return f"{operation_id}_{methods[0].lower()}"
+def _normalize_operation_ids(paths: Dict[str, Dict]) -> None:
+    """Make FastAPI-generated operation IDs stable for multi-method routes.
 
+    FastAPI derives the default operation ID suffix from the first item in the
+    route's methods set. For routes registered with several HTTP methods, that
+    set iteration order can vary between processes, which makes the snapshot
+    drift even when no routes changed.
+    """
+    for path_ops in paths.values():
+        if not isinstance(path_ops, dict):
+            continue
 
-def _set_stable_operation_ids(routes: Iterable) -> None:
-    for route in routes:
-        if getattr(route, "operation_id", None) is not None:
+        methods = {method for method in path_ops if method in HTTP_METHODS}
+        if not methods:
             continue
-        if getattr(route, "methods", None) is None:
-            continue
-        route.operation_id = _stable_generate_unique_id(route)
+
+        for method, operation in path_ops.items():
+            if method not in HTTP_METHODS or not isinstance(operation, dict):
+                continue
+
+            operation_id = operation.get("operationId")
+            if not isinstance(operation_id, str):
+                continue
+
+            for suffix in methods:
+                suffix_token = f"_{suffix}"
+                if operation_id.endswith(suffix_token):
+                    operation["operationId"] = (
+                        operation_id[: -len(suffix_token)] + f"_{method}"
+                    )
+                    break
 
 
 def generate_snapshot() -> Dict[str, Dict]:
@@ -50,7 +88,7 @@ def generate_snapshot() -> Dict[str, Dict]:
     from fastapi.openapi.utils import get_openapi
 
     from litellm.proxy._lazy_features import LAZY_FEATURES
-    from litellm.proxy.proxy_server import app
+    from litellm.proxy.proxy_server import app, ensure_unique_openapi_operation_ids
 
     for feat in LAZY_FEATURES:
         if feat.module_path in sys.modules:
@@ -62,6 +100,7 @@ def generate_snapshot() -> Dict[str, Dict]:
             sys.stderr.write(f"warning: skip {feat.name}: {exc}\n")
 
     fragments: Dict[str, Dict] = {}
+    used_operation_ids: Set[str] = set()
     for feat in LAZY_FEATURES:
         feat_routes = [
             r
@@ -70,15 +109,26 @@ def generate_snapshot() -> Dict[str, Dict]:
         ]
         if not feat_routes:
             continue
-        _set_stable_operation_ids(feat_routes)
+        _stabilize_multi_method_route_ids(feat_routes)
         full = get_openapi(title=app.title, version=app.version, routes=feat_routes)
+        paths = full.get("paths", {})
+        _normalize_operation_ids(paths)
         # Group all of a feature's routes under one tag.
         for path_ops in full.get("paths", {}).values():
-            for op in path_ops.values():
+            for method, op in path_ops.items():
                 if isinstance(op, dict):
+                    operation_id = op.get("operationId")
+                    if isinstance(operation_id, str):
+                        for suffix in HTTP_METHOD_SUFFIXES:
+                            if operation_id.endswith(f"_{suffix}"):
+                                op["operationId"] = (
+                                    operation_id[: -len(suffix)] + method
+                                )
+                                break
                     op["tags"] = [feat.name]
+        full = ensure_unique_openapi_operation_ids(full, used_operation_ids)
         fragments[feat.name] = {
-            "paths": full.get("paths", {}),
+            "paths": paths,
             "components": {"schemas": full.get("components", {}).get("schemas", {})},
         }
     return fragments

litellm/proxy/_types.py

@@ -668,6 +668,8 @@ class LiteLLMRoutes(enum.Enum):
             "/models/{model_id}",
             "/guardrails/list",
             "/v2/guardrails/list",
+            "/project/list",
+            "/project/info",
         ]
         + spend_tracking_routes
         + key_management_routes
@@ -692,6 +694,9 @@ class LiteLLMRoutes(enum.Enum):
         "/model/{model_id}/update",
         "/prompt/list",
         "/prompt/info",
+        # Project read routes - endpoint scopes results to caller's teams (non-admin)
+        "/project/list",
+        "/project/info",
         # Invitation routes - org/team admins checked in endpoint via _user_has_admin_privileges
         "/invitation/new",
         "/invitation/delete",
@@ -2161,8 +2166,8 @@ class PassThroughGenericEndpoint(LiteLLMPydanticObjectBase):
         description="The USD cost per request to the target endpoint. This is used to calculate the cost of the request to the target endpoint.",
     )
     auth: bool = Field(
-        default=False,
-        description="Whether authentication is required for the pass-through endpoint. If True, requests to the endpoint will require a valid LiteLLM API key.",
+        default=True,
+        description="Whether authentication is required for the pass-through endpoint. Defaults to True so a pass-through silently created without an explicit value still requires a valid LiteLLM API key — set to False only if the endpoint is meant to be a public forwarder (e.g. an unauthenticated webhook target).",
     )
     guardrails: Optional[PassThroughGuardrailsConfig] = Field(
         default=None,
@@ -2574,6 +2579,7 @@ class UserAPIKeyAuth(
     user_spend: Optional[float] = None
     user_max_budget: Optional[float] = None
     request_route: Optional[str] = None
+    budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True)
     user: Optional[Any] = None  # Expanded user object when expand=user is used
     created_by_user: Optional[Any] = (
         None  # Expanded created_by user when expand=user is used

litellm/proxy/auth/auth_checks.py

@@ -12,14 +12,13 @@ Run checks for:
 import asyncio
 import re
 import time
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast
 
 from fastapi import HTTPException, Request, status
 from pydantic import BaseModel
 
 import litellm
 from litellm._logging import verbose_proxy_logger
-from litellm.caching.caching import DualCache
 from litellm.caching.dual_cache import LimitedSizeOrderedDict
 from litellm.constants import (
     CLI_JWT_EXPIRATION_HOURS,
@@ -61,11 +60,17 @@ from litellm.proxy._types import (
     UserAPIKeyAuth,
 )
 from litellm.proxy.auth.route_checks import RouteChecks
+from litellm.proxy.common_utils.http_parsing_utils import (
+    _safe_get_request_headers,
+    _safe_get_request_query_params,
+)
 from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
 from litellm.proxy.guardrails.tool_name_extraction import (
     TOOL_CAPABLE_CALL_TYPES,
     extract_request_tool_names,
 )
+from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
 from litellm.proxy.route_llm_request import route_request
 from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
 from litellm.router import Router
@@ -485,7 +490,10 @@ async def common_checks(  # noqa: PLR0915
     from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
 
     _model: Optional[Union[str, List[str]]] = get_model_from_request(
-        request_body, route
+        request_data=request_body,
+        route=route,
+        request_headers=_safe_get_request_headers(request=request),
+        request_query_params=_safe_get_request_query_params(request=request),
     )
 
     # 1. If team is blocked
@@ -655,13 +663,7 @@ async def common_checks(  # noqa: PLR0915
             end_user_object is not None
             and end_user_object.litellm_budget_table is not None
         ):
-            end_user_budget = end_user_object.litellm_budget_table.max_budget
-            if end_user_budget is not None and end_user_object.spend > end_user_budget:
-                raise litellm.BudgetExceededError(
-                    current_cost=end_user_object.spend,
-                    max_budget=end_user_budget,
-                    message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}",
-                )
+            await _check_end_user_budget(end_user_obj=end_user_object, route=route)
 
     _enforce_user_param_check(general_settings, request, request_body, route)
     _reject_clientside_metadata_tags_check(general_settings, request_body, route)
@@ -852,7 +854,7 @@ def get_actual_routes(allowed_routes: list) -> list:
 
 async def get_default_end_user_budget(
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span] = None,
 ) -> Optional[LiteLLM_BudgetTable]:
     """
@@ -875,9 +877,12 @@ async def get_default_end_user_budget(
     cache_key = f"default_end_user_budget:{litellm.max_end_user_budget_id}"
 
     # Check cache first
-    cached_budget = await user_api_key_cache.async_get_cache(key=cache_key)
+    cached_budget = await user_api_key_cache.async_get_cache(
+        key=cache_key,
+        model_type=LiteLLM_BudgetTable,
+    )
     if cached_budget is not None:
-        return LiteLLM_BudgetTable(**cached_budget)
+        return cached_budget
 
     # Fetch from database
     try:
@@ -891,14 +896,16 @@ async def get_default_end_user_budget(
             )
             return None
 
+        _budget_obj = LiteLLM_BudgetTable(**budget_record.dict())
         # Cache the budget for 60 seconds
         await user_api_key_cache.async_set_cache(
             key=cache_key,
-            value=budget_record.dict(),
+            value=_budget_obj,
+            model_type=LiteLLM_BudgetTable,
             ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
         )
 
-        return LiteLLM_BudgetTable(**budget_record.dict())
+        return _budget_obj
 
     except Exception as e:
         verbose_proxy_logger.error(f"Error fetching default end user budget: {str(e)}")
@@ -909,7 +916,7 @@ async def get_default_end_user_budget(
 async def get_team_member_default_budget(
     budget_id: str,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
 ) -> Optional[LiteLLM_BudgetTable]:
     """
     Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
@@ -966,7 +973,7 @@ async def get_team_member_default_budget(
 async def _apply_default_budget_to_end_user(
     end_user_obj: LiteLLM_EndUserTable,
     prisma_client: PrismaClient,
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span] = None,
 ) -> LiteLLM_EndUserTable:
     """
@@ -1006,7 +1013,7 @@ async def _apply_default_budget_to_end_user(
     return end_user_obj
 
 
-def _check_end_user_budget(
+async def _check_end_user_budget(
     end_user_obj: LiteLLM_EndUserTable,
     route: str,
 ) -> None:
@@ -1027,11 +1034,20 @@ def _check_end_user_budget(
         return
 
     end_user_budget = end_user_obj.litellm_budget_table.max_budget
-    if end_user_budget is not None and end_user_obj.spend > end_user_budget:
+    if end_user_budget is None:
+        return
+
+    from litellm.proxy.proxy_server import get_current_spend
+
+    end_user_spend = await get_current_spend(
+        counter_key=f"spend:end_user:{end_user_obj.user_id}",
+        fallback_spend=end_user_obj.spend or 0.0,
+    )
+    if end_user_spend > end_user_budget:
         raise litellm.BudgetExceededError(
-            current_cost=end_user_obj.spend,
+            current_cost=end_user_spend,
             max_budget=end_user_budget,
-            message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_obj.spend}, Budget={end_user_budget}",
+            message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}",
         )
 
 
@@ -1039,7 +1055,7 @@ def _check_end_user_budget(
 async def get_end_user_object(
     end_user_id: Optional[str],
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     route: str,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
@@ -1070,10 +1086,12 @@ async def get_end_user_object(
     _key = "end_user_id:{}".format(end_user_id)
 
     # Check cache first
-    cached_user_obj = await user_api_key_cache.async_get_cache(key=_key)
+    cached_user_obj = await user_api_key_cache.async_get_cache(
+        key=_key,
+        model_type=LiteLLM_EndUserTable,
+    )
     if cached_user_obj is not None:
-        return_obj = LiteLLM_EndUserTable(**cached_user_obj)
-
+        return_obj = cached_user_obj
         # Apply default budget if needed
         return_obj = await _apply_default_budget_to_end_user(
             end_user_obj=return_obj,
@@ -1083,7 +1101,7 @@ async def get_end_user_object(
         )
 
         # Check budget limits
-        _check_end_user_budget(end_user_obj=return_obj, route=route)
+        await _check_end_user_budget(end_user_obj=return_obj, route=route)
 
         return return_obj
 
@@ -1108,13 +1126,15 @@ async def get_end_user_object(
             parent_otel_span=parent_otel_span,
         )
 
-        # Save to cache (always store as dict for consistency)
+        # Save to cache
         await user_api_key_cache.async_set_cache(
-            key="end_user_id:{}".format(end_user_id), value=_response.dict()
+            key="end_user_id:{}".format(end_user_id),
+            value=_response,
+            model_type=LiteLLM_EndUserTable,
         )
 
         # Check budget limits
-        _check_end_user_budget(end_user_obj=_response, route=route)
+        await _check_end_user_budget(end_user_obj=_response, route=route)
 
         return _response
 
@@ -1128,7 +1148,7 @@ async def get_end_user_object(
 async def get_tag_objects_batch(
     tag_names: List[str],
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> Dict[str, LiteLLM_TagTable]:
@@ -1161,12 +1181,12 @@ async def get_tag_objects_batch(
     # Try to get all tags from cache first
     for tag_name in tag_names:
         cache_key = f"tag:{tag_name}"
-        cached_tag = await user_api_key_cache.async_get_cache(key=cache_key)
+        cached_tag = await user_api_key_cache.async_get_cache(
+            key=cache_key,
+            model_type=LiteLLM_TagTable,
+        )
         if cached_tag is not None:
-            if isinstance(cached_tag, dict):
-                tag_objects[tag_name] = LiteLLM_TagTable(**cached_tag)
-            else:
-                tag_objects[tag_name] = cached_tag
+            tag_objects[tag_name] = cached_tag
         else:
             uncached_tags.append(tag_name)
 
@@ -1182,11 +1202,13 @@ async def get_tag_objects_batch(
             for db_tag in db_tags:
                 tag_name = db_tag.tag_name
                 cache_key = f"tag:{tag_name}"
-                # Cache with default TTL (same as end_user objects)
+                _tag_obj = LiteLLM_TagTable(**db_tag.dict())
                 await user_api_key_cache.async_set_cache(
-                    key=cache_key, value=db_tag.dict()
+                    key=cache_key,
+                    value=_tag_obj,
+                    model_type=LiteLLM_TagTable,
                 )
-                tag_objects[tag_name] = LiteLLM_TagTable(**db_tag.dict())
+                tag_objects[tag_name] = _tag_obj
         except Exception as e:
             verbose_proxy_logger.debug(f"Error batch fetching tags from database: {e}")
 
@@ -1197,7 +1219,7 @@ async def get_tag_objects_batch(
 async def get_tag_object(
     tag_name: Optional[str],
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> Optional[LiteLLM_TagTable]:
@@ -1236,7 +1258,7 @@ async def get_team_membership(
     user_id: str,
     team_id: str,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> Optional["LiteLLM_TeamMembership"]:
@@ -1256,9 +1278,12 @@ async def get_team_membership(
     _key = "team_membership:{}:{}".format(user_id, team_id)
 
     # check if in cache
-    cached_membership_obj = await user_api_key_cache.async_get_cache(key=_key)
+    cached_membership_obj = await user_api_key_cache.async_get_cache(
+        key=_key,
+        model_type=LiteLLM_TeamMembership,
+    )
     if cached_membership_obj is not None:
-        return LiteLLM_TeamMembership(**cached_membership_obj)
+        return cached_membership_obj
 
     # else, check db
     try:
@@ -1270,10 +1295,12 @@ async def get_team_membership(
         if response is None:
             return None
 
-        # save the team membership object to cache (store as dict)
-        await user_api_key_cache.async_set_cache(key=_key, value=response.dict())
-
         _response = LiteLLM_TeamMembership(**response.dict())
+        await user_api_key_cache.async_set_cache(
+            key=_key,
+            value=_response,
+            model_type=LiteLLM_TeamMembership,
+        )
 
         return _response
     except Exception:
@@ -1441,7 +1468,7 @@ async def _get_fuzzy_user_object(
 async def get_user_object(
     user_id: Optional[str],
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     user_id_upsert: bool,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
@@ -1460,12 +1487,12 @@ async def get_user_object(
 
     # check if in cache
     if not check_db_only:
-        cached_user_obj = await user_api_key_cache.async_get_cache(key=user_id)
+        cached_user_obj = await user_api_key_cache.async_get_cache(
+            key=user_id,
+            model_type=LiteLLM_UserTable,
+        )
         if cached_user_obj is not None:
-            if isinstance(cached_user_obj, dict):
-                return LiteLLM_UserTable(**cached_user_obj)
-            elif isinstance(cached_user_obj, LiteLLM_UserTable):
-                return cached_user_obj
+            return cached_user_obj
     # else, check db
     if prisma_client is None:
         raise Exception("No db connected")
@@ -1527,7 +1554,8 @@ async def get_user_object(
         # save the user object to cache
         await user_api_key_cache.async_set_cache(
             key=user_id,
-            value=response_dict,
+            value=_response,
+            model_type=LiteLLM_UserTable,
             ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
         )
 
@@ -1548,13 +1576,21 @@ async def get_user_object(
 
 async def _cache_management_object(
     key: str,
-    value: BaseModel,
-    user_api_key_cache: DualCache,
+    value: Union[BaseModel, Dict[str, Any]],
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: Optional[ProxyLogging],
+    *,
+    model_type: Type[BaseModel],
 ):
+    """
+    Persist management objects via ``UserApiKeyCache`` (in-memory + optional Redis).
+
+    ``UserApiKeyCache`` serializes with ``model_type`` so Redis and in-memory stay aligned.
+    """
     await user_api_key_cache.async_set_cache(
         key=key,
         value=value,
+        model_type=model_type,
         ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
     )
 
@@ -1562,7 +1598,7 @@ async def _cache_management_object(
 async def _cache_team_object(
     team_id: str,
     team_table: LiteLLM_TeamTableCachedObj,
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: Optional[ProxyLogging],
 ):
     key = "team_id:{}".format(team_id)
@@ -1575,13 +1611,14 @@ async def _cache_team_object(
         value=team_table,
         user_api_key_cache=user_api_key_cache,
         proxy_logging_obj=proxy_logging_obj,
+        model_type=LiteLLM_TeamTableCachedObj,
     )
 
 
 async def _cache_key_object(
     hashed_token: str,
     user_api_key_obj: UserAPIKeyAuth,
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: Optional[ProxyLogging],
 ):
     key = hashed_token
@@ -1589,17 +1626,21 @@ async def _cache_key_object(
     ## CACHE REFRESH TIME
     user_api_key_obj.last_refreshed_at = time.time()
 
+    cached_key_obj = _copy_user_api_key_auth_for_cache(
+        user_api_key_obj=user_api_key_obj
+    )
     await _cache_management_object(
         key=key,
-        value=user_api_key_obj,
+        value=cached_key_obj,
         user_api_key_cache=user_api_key_cache,
         proxy_logging_obj=proxy_logging_obj,
+        model_type=UserAPIKeyAuth,
     )
 
 
 async def _delete_cache_key_object(
     hashed_token: str,
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: Optional[ProxyLogging],
 ):
     key = hashed_token
@@ -1647,7 +1688,7 @@ async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient):
 async def _get_team_object_from_user_api_key_cache(
     team_id: str,
     prisma_client: PrismaClient,
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     last_db_access_time: LimitedSizeOrderedDict,
     db_cache_expiry: int,
     proxy_logging_obj: Optional[ProxyLogging],
@@ -1708,38 +1749,38 @@ async def _get_team_object_from_user_api_key_cache(
 async def _get_team_object_from_cache(
     key: str,
     proxy_logging_obj: Optional[ProxyLogging],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span],
 ) -> Optional[LiteLLM_TeamTableCachedObj]:
-    cached_team_obj: Optional[LiteLLM_TeamTableCachedObj] = None
-
-    ## CHECK REDIS CACHE ##
+    ## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ##
     if (
         proxy_logging_obj is not None
         and proxy_logging_obj.internal_usage_cache.dual_cache
     ):
-        cached_team_obj = (
+        cached_raw = (
             await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache(
                 key=key, parent_otel_span=parent_otel_span
             )
         )
+        if cached_raw is not None:
+            from_internal = CacheCodec.deserialize(
+                cached_raw, LiteLLM_TeamTableCachedObj
+            )
+            if from_internal is not None:
+                return from_internal
 
-    if cached_team_obj is None:
-        cached_team_obj = await user_api_key_cache.async_get_cache(key=key)
-
-    if cached_team_obj is not None:
-        if isinstance(cached_team_obj, dict):
-            return LiteLLM_TeamTableCachedObj(**cached_team_obj)
-        elif isinstance(cached_team_obj, LiteLLM_TeamTableCachedObj):
-            return cached_team_obj
-
-    return None
+    decoded = await user_api_key_cache.async_get_cache(
+        key=key,
+        parent_otel_span=parent_otel_span,
+        model_type=LiteLLM_TeamTableCachedObj,
+    )
+    return decoded
 
 
 async def get_team_object(
     team_id: str,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
     check_cache_only: Optional[bool] = None,
@@ -1805,20 +1846,21 @@ async def get_team_object(
 async def _cache_access_object(
     access_group_id: str,
     access_group_table: LiteLLM_AccessGroupTable,
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ):
     key = "access_group_id:{}".format(access_group_id)
     await user_api_key_cache.async_set_cache(
         key=key,
         value=access_group_table,
+        model_type=LiteLLM_AccessGroupTable,
         ttl=DEFAULT_ACCESS_GROUP_CACHE_TTL,
     )
 
 
 async def _delete_cache_access_object(
     access_group_id: str,
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ):
     key = "access_group_id:{}".format(access_group_id)
@@ -1836,7 +1878,7 @@ async def _delete_cache_access_object(
 async def get_access_object(
     access_group_id: str,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> LiteLLM_AccessGroupTable:
     """
@@ -1858,13 +1900,12 @@ async def get_access_object(
 
     key = "access_group_id:{}".format(access_group_id)
 
-    # Always check cache first
-    cached_access_obj = await user_api_key_cache.async_get_cache(key=key)
+    cached_access_obj = await user_api_key_cache.async_get_cache(
+        key=key,
+        model_type=LiteLLM_AccessGroupTable,
+    )
     if cached_access_obj is not None:
-        if isinstance(cached_access_obj, dict):
-            return LiteLLM_AccessGroupTable(**cached_access_obj)
-        elif isinstance(cached_access_obj, LiteLLM_AccessGroupTable):
-            return cached_access_obj
+        return cached_access_obj
 
     # Not in cache - fetch from DB
     try:
@@ -1910,7 +1951,7 @@ async def get_access_object(
 async def get_team_object_by_alias(
     team_alias: str,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional["Span"] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> LiteLLM_TeamTableCachedObj:
@@ -1992,6 +2033,7 @@ async def get_team_object_by_alias(
         await user_api_key_cache.async_set_cache(
             key=cache_key,
             value=team_obj,
+            model_type=LiteLLM_TeamTableCachedObj,
             ttl=DEFAULT_IN_MEMORY_TTL,
         )
         # Also cache by team_id for consistency
@@ -1999,6 +2041,7 @@ async def get_team_object_by_alias(
         await user_api_key_cache.async_set_cache(
             key=team_id_cache_key,
             value=team_obj,
+            model_type=LiteLLM_TeamTableCachedObj,
             ttl=DEFAULT_IN_MEMORY_TTL,
         )
 
@@ -2020,7 +2063,7 @@ async def get_team_object_by_alias(
 async def get_org_object_by_alias(
     org_alias: str,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional["Span"] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> Optional[LiteLLM_OrganizationTable]:
@@ -2047,12 +2090,12 @@ async def get_org_object_by_alias(
 
     # Check cache first (keyed by alias)
     cache_key = "org_alias:{}".format(org_alias)
-    cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key)
+    cached_org_obj = await user_api_key_cache.async_get_cache(
+        key=cache_key,
+        model_type=LiteLLM_OrganizationTable,
+    )
     if cached_org_obj is not None:
-        if isinstance(cached_org_obj, dict):
-            return LiteLLM_OrganizationTable(**cached_org_obj)
-        elif isinstance(cached_org_obj, LiteLLM_OrganizationTable):
-            return cached_org_obj
+        return cached_org_obj
 
     # Query database by organization_alias
     try:
@@ -2082,13 +2125,15 @@ async def get_org_object_by_alias(
         # Cache the result
         await user_api_key_cache.async_set_cache(
             key=cache_key,
-            value=org_obj.model_dump(),
+            value=org_obj,
+            model_type=LiteLLM_OrganizationTable,
             ttl=DEFAULT_IN_MEMORY_TTL,
         )
         # Also cache by org_id for consistency
         await user_api_key_cache.async_set_cache(
             key="org_id:{}".format(org_obj.organization_id),
-            value=org_obj.model_dump(),
+            value=org_obj,
+            model_type=LiteLLM_OrganizationTable,
             ttl=DEFAULT_IN_MEMORY_TTL,
         )
 
@@ -2291,7 +2336,7 @@ async def get_jwt_key_mapping_object(
 async def get_key_object(
     hashed_token: str,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
     check_cache_only: Optional[bool] = None,
@@ -2309,15 +2354,14 @@ async def get_key_object(
     # check if in cache
     key = hashed_token
 
-    cached_key_obj: Optional[UserAPIKeyAuth] = await user_api_key_cache.async_get_cache(
-        key=key
+    # Same flow as before: use cache only when we have a hit we can turn into UserAPIKeyAuth
+    # (dict from Redis / model_dump, or UserAPIKeyAuth from in-memory). Otherwise fall through to DB.
+    user_api_key_auth = await user_api_key_cache.async_get_cache(
+        key=key,
+        model_type=UserAPIKeyAuth,
     )
-
-    if cached_key_obj is not None:
-        if isinstance(cached_key_obj, dict):
-            return UserAPIKeyAuth(**cached_key_obj)
-        elif isinstance(cached_key_obj, UserAPIKeyAuth):
-            return cached_key_obj
+    if user_api_key_auth is not None:
+        return _copy_user_api_key_auth_for_cache(user_api_key_obj=user_api_key_auth)
 
     if check_cache_only:
         raise Exception(
@@ -2370,11 +2414,21 @@ async def get_key_object(
     return _response
 
 
+def _copy_user_api_key_auth_for_cache(
+    user_api_key_obj: UserAPIKeyAuth,
+) -> UserAPIKeyAuth:
+    copied_key_obj = user_api_key_obj.model_copy()
+    copied_key_obj.budget_reservation = None
+    copied_key_obj.parent_otel_span = None
+    copied_key_obj.request_route = None
+    return copied_key_obj
+
+
 @log_db_metrics
 async def get_object_permission(
     object_permission_id: str,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> Optional[LiteLLM_ObjectPermissionTable]:
@@ -2390,12 +2444,12 @@ async def get_object_permission(
 
     # check if in cache
     key = "object_permission_id:{}".format(object_permission_id)
-    cached_obj_permission = await user_api_key_cache.async_get_cache(key=key)
-    if cached_obj_permission is not None:
-        if isinstance(cached_obj_permission, dict):
-            return LiteLLM_ObjectPermissionTable(**cached_obj_permission)
-        elif isinstance(cached_obj_permission, LiteLLM_ObjectPermissionTable):
-            return cached_obj_permission
+    deserialized_perm = await user_api_key_cache.async_get_cache(
+        key=key,
+        model_type=LiteLLM_ObjectPermissionTable,
+    )
+    if deserialized_perm is not None:
+        return deserialized_perm
 
     # else, check db
     try:
@@ -2406,14 +2460,15 @@ async def get_object_permission(
         if response is None:
             return None
 
-        # save the object permission to cache
+        _perm_obj = LiteLLM_ObjectPermissionTable(**response.dict())
         await user_api_key_cache.async_set_cache(
             key=key,
-            value=response.model_dump(),
+            value=_perm_obj,
+            model_type=LiteLLM_ObjectPermissionTable,
             ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
         )
 
-        return LiteLLM_ObjectPermissionTable(**response.dict())
+        return _perm_obj
     except Exception:
         return None
 
@@ -2422,7 +2477,7 @@ async def get_object_permission(
 async def get_managed_vector_store_rows_by_uuids(
     uuids: List[str],
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> List[LiteLLM_ManagedVectorStoresTable]:
@@ -2442,14 +2497,12 @@ async def get_managed_vector_store_rows_by_uuids(
 
     for uuid in uuids:
         key = "managed_vector_store_id:{}".format(uuid)
-        cached = await user_api_key_cache.async_get_cache(key=key)
-        if cached is not None:
-            if isinstance(cached, dict):
-                result.append(LiteLLM_ManagedVectorStoresTable(**cached))
-            elif isinstance(cached, LiteLLM_ManagedVectorStoresTable):
-                result.append(cached)
-            else:
-                cache_misses.append(uuid)
+        deserialized_vs = await user_api_key_cache.async_get_cache(
+            key=key,
+            model_type=LiteLLM_ManagedVectorStoresTable,
+        )
+        if deserialized_vs is not None:
+            result.append(deserialized_vs)
         else:
             cache_misses.append(uuid)
 
@@ -2475,7 +2528,8 @@ async def get_managed_vector_store_rows_by_uuids(
         key = "managed_vector_store_id:{}".format(cached_obj.vector_store_id)
         await user_api_key_cache.async_set_cache(
             key=key,
-            value=row_dict,
+            value=cached_obj,
+            model_type=LiteLLM_ManagedVectorStoresTable,
             ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
         )
         result.append(cached_obj)
@@ -2487,7 +2541,7 @@ async def get_managed_vector_store_rows_by_uuids(
 async def get_org_object(
     org_id: str,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
     include_budget_table: bool = False,
@@ -2518,12 +2572,12 @@ async def get_org_object(
         cache_key = "org_id:{}:with_budget".format(org_id)
 
     # check if in cache
-    cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key)
-    if cached_org_obj is not None:
-        if isinstance(cached_org_obj, dict):
-            return LiteLLM_OrganizationTable(**cached_org_obj)
-        elif isinstance(cached_org_obj, LiteLLM_OrganizationTable):
-            return cached_org_obj
+    deserialized_org = await user_api_key_cache.async_get_cache(
+        key=cache_key,
+        model_type=LiteLLM_OrganizationTable,
+    )
+    if deserialized_org is not None:
+        return deserialized_org
     # else, check db
     try:
         query_kwargs: Dict[str, Any] = {"where": {"organization_id": org_id}}
@@ -2537,16 +2591,16 @@ async def get_org_object(
         if response is None:
             raise Exception
 
+        _org_obj = LiteLLM_OrganizationTable(**response.model_dump())
         # Cache the result
         await user_api_key_cache.async_set_cache(
             key=cache_key,
-            value=(
-                response.model_dump() if hasattr(response, "model_dump") else response
-            ),
+            value=_org_obj,
+            model_type=LiteLLM_OrganizationTable,
             ttl=DEFAULT_IN_MEMORY_TTL,
         )
 
-        return response
+        return _org_obj
     except Exception:
         raise Exception(
             f"Organization doesn't exist in db. Organization={org_id}. Create organization via `/organization/new` call."
@@ -2559,7 +2613,7 @@ async def _get_resources_from_access_groups(
         "access_model_names", "access_mcp_server_ids", "access_agent_ids"
     ],
     prisma_client: Optional[PrismaClient] = None,
-    user_api_key_cache: Optional[DualCache] = None,
+    user_api_key_cache: Optional[UserApiKeyCache] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> List[str]:
     """
@@ -2617,7 +2671,7 @@ async def _get_resources_from_access_groups(
 async def _get_models_from_access_groups(
     access_group_ids: List[str],
     prisma_client: Optional[PrismaClient] = None,
-    user_api_key_cache: Optional[DualCache] = None,
+    user_api_key_cache: Optional[UserApiKeyCache] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> List[str]:
     """
@@ -2636,7 +2690,7 @@ async def _get_models_from_access_groups(
 async def _get_mcp_server_ids_from_access_groups(
     access_group_ids: List[str],
     prisma_client: Optional[PrismaClient] = None,
-    user_api_key_cache: Optional[DualCache] = None,
+    user_api_key_cache: Optional[UserApiKeyCache] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> List[str]:
     """
@@ -2655,7 +2709,7 @@ async def _get_mcp_server_ids_from_access_groups(
 async def _get_agent_ids_from_access_groups(
     access_group_ids: List[str],
     prisma_client: Optional[PrismaClient] = None,
-    user_api_key_cache: Optional[DualCache] = None,
+    user_api_key_cache: Optional[UserApiKeyCache] = None,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> List[str]:
     """
@@ -3379,7 +3433,7 @@ async def _check_team_member_budget(
     user_object: Optional[LiteLLM_UserTable],
     valid_token: Optional[UserAPIKeyAuth],
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
 ):
     """Check if team member is over their max budget within the team."""
@@ -3447,7 +3501,7 @@ async def _check_team_member_model_access(
     valid_token: UserAPIKeyAuth,
     llm_router: Optional[Router],
     prisma_client: Optional["PrismaClient"],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
 ) -> None:
     """
@@ -3754,7 +3808,7 @@ async def _project_soft_budget_check(
 async def get_project_object(
     project_id: str,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: Optional[ProxyLogging] = None,
 ) -> Optional[LiteLLM_ProjectTableCachedObj]:
     """
@@ -3769,12 +3823,12 @@ async def get_project_object(
 
     # Check cache first
     cache_key = "project_id:{}".format(project_id)
-    cached_obj = await user_api_key_cache.async_get_cache(key=cache_key)
-    if cached_obj is not None:
-        if isinstance(cached_obj, dict):
-            return LiteLLM_ProjectTableCachedObj(**cached_obj)
-        elif isinstance(cached_obj, LiteLLM_ProjectTableCachedObj):
-            return cached_obj
+    deserialized_project = await user_api_key_cache.async_get_cache(
+        key=cache_key,
+        model_type=LiteLLM_ProjectTableCachedObj,
+    )
+    if deserialized_project is not None:
+        return deserialized_project
 
     # Fetch from DB
     project_row = await prisma_client.db.litellm_projecttable.find_unique(
@@ -3793,6 +3847,7 @@ async def get_project_object(
         value=project_obj,
         user_api_key_cache=user_api_key_cache,
         proxy_logging_obj=proxy_logging_obj,
+        model_type=LiteLLM_ProjectTableCachedObj,
     )
 
     return project_obj
@@ -3802,7 +3857,7 @@ async def _organization_max_budget_check(
     valid_token: Optional[UserAPIKeyAuth],
     team_object: Optional[LiteLLM_TeamTable],
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
 ):
     """
@@ -3896,7 +3951,7 @@ async def _organization_max_budget_check(
 async def _tag_max_budget_check(
     request_body: dict,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     proxy_logging_obj: ProxyLogging,
     valid_token: Optional[UserAPIKeyAuth],
 ):
@@ -3935,13 +3990,19 @@ async def _tag_max_budget_check(
         if (
             tag_object.litellm_budget_table is not None
             and tag_object.litellm_budget_table.max_budget is not None
-            and tag_object.spend is not None
-            and tag_object.spend > tag_object.litellm_budget_table.max_budget
         ):
+            from litellm.proxy.proxy_server import get_current_spend
+
+            tag_spend = await get_current_spend(
+                counter_key=f"spend:tag:{tag_name}",
+                fallback_spend=tag_object.spend or 0.0,
+            )
+            if tag_spend <= tag_object.litellm_budget_table.max_budget:
+                continue
             raise litellm.BudgetExceededError(
-                current_cost=tag_object.spend,
+                current_cost=tag_spend,
                 max_budget=tag_object.litellm_budget_table.max_budget,
-                message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_object.spend}, Max budget: {tag_object.litellm_budget_table.max_budget}",
+                message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}",
             )
 
 

litellm/proxy/auth/auth_utils.py

@@ -2,7 +2,7 @@ import os
 import re
 import sys
 from functools import lru_cache
-from typing import Any, List, Optional, Tuple
+from typing import Any, Dict, List, Mapping, Optional, Tuple, Union
 
 from fastapi import HTTPException, Request, status
 
@@ -167,6 +167,81 @@ def _allow_model_level_clientside_configurable_parameters(
     )
 
 
+# Config dicts whose entries are spread as ``**dict`` into outbound LLM
+# API calls. ``litellm_embedding_config`` is consumed by the Milvus
+# vector store transformer; future nested-config keys with the same
+# threat shape should be added here.
+_NESTED_CONFIG_KEYS: Tuple[str, ...] = ("litellm_embedding_config",)
+
+# Banned root-level params. Same list applies to every entry in
+# ``_NESTED_CONFIG_KEYS`` because those dicts get spread as ``**kwargs``
+# into the same outbound calls.
+_BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = (
+    "api_base",
+    "base_url",
+    "user_config",
+    "aws_sts_endpoint",
+    "aws_web_identity_token",
+    "aws_role_name",
+    "vertex_credentials",
+    # Endpoint-targeting fields that retarget the outbound request or
+    # an observability callback. An attacker-controlled value either
+    # exfiltrates the request payload (incl. messages + admin-set
+    # tokens) to the attacker's host, or coerces the proxy into
+    # authenticating against the attacker's host with admin secrets.
+    "aws_bedrock_runtime_endpoint",
+    "langsmith_base_url",
+    "langfuse_host",
+    "posthog_host",
+    "braintrust_host",
+    "slack_webhook_url",
+    # Provider-specific endpoint overrides that flow into the outbound
+    # request via ``optional_params``. Same threat as ``api_base``:
+    # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker
+    # S3; ``sagemaker_base_url`` redirects all SageMaker traffic;
+    # ``deployment_url`` redirects SAP deployments.
+    "s3_endpoint_url",
+    "sagemaker_base_url",
+    "deployment_url",
+)
+
+
+def _check_banned_params(
+    body: dict,
+    general_settings: dict,
+    llm_router: Optional[Router],
+    model: str,
+) -> None:
+    """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in.
+
+    Shared between the root-level check and the nested-config check so a
+    new banned param only needs to be added in one place.
+    """
+    for param in _BANNED_REQUEST_BODY_PARAMS:
+        if param not in body:
+            continue
+        if general_settings.get("allow_client_side_credentials") is True:
+            return
+        if (
+            _allow_model_level_clientside_configurable_parameters(
+                model=model,
+                param=param,
+                request_body_value=body[param],
+                llm_router=llm_router,
+            )
+            is True
+        ):
+            return
+        raise ValueError(
+            f"Rejected Request: {param} is not allowed in request body. "
+            "Clientside passthrough requires explicit admin opt-in via "
+            "either `general_settings.allow_client_side_credentials = true` "
+            "(proxy-wide) or `configurable_clientside_auth_params` on the "
+            "deployment in your proxy config.yaml. "
+            "Relevant Issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997",
+        )
+
+
 def is_request_body_safe(
     request_body: dict, general_settings: dict, llm_router: Optional[Router], model: str
 ) -> bool:
@@ -175,72 +250,31 @@ def is_request_body_safe(
 
     A malicious user can set the api_base to their own domain and invoke POST /chat/completions to intercept and steal the OpenAI API key.
     Relevant issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997
+
+    The blocklist is enforced unconditionally. Legitimate clientside
+    credential / endpoint passthrough goes through one of the two
+    explicit admin opt-ins (``general_settings.allow_client_side_credentials``
+    proxy-wide or ``configurable_clientside_auth_params`` per deployment).
+    Historically there was a third, *implicit*, *caller-controlled* path:
+    ``check_complete_credentials`` returned True when the caller supplied
+    any non-empty ``api_key``, which made the entire blocklist a no-op.
+    That bypass turned every missing entry on the blocklist into an
+    exploitable SSRF / credential-exfil hole — see GHSA-jh89-88fc-qrfp,
+    GHSA-3frq-6r6h-7j64, and the chain of veria-admin findings (Dv_m860l,
+    b_yRJeQ5, stN90yjP, LBlyOAc8, U2TD78kg). Removed: the blocklist now
+    has a single, predictable failure mode for missing entries (a 400),
+    not a credential leak.
+
+    Iterative single-level descent into ``_NESTED_CONFIG_KEYS`` (rather
+    than recursion) covers nested-config attacks like Milvus's
+    ``litellm_embedding_config.api_base`` (VERIA-6) without exposing a
+    recursion-depth DoS surface.
     """
-    banned_params = [
-        "api_base",
-        "base_url",
-        "user_config",
-        "aws_sts_endpoint",
-        "aws_web_identity_token",
-        "aws_role_name",
-        "vertex_credentials",
-        # Endpoint-targeting fields that retarget the outbound request or
-        # an observability callback. An attacker-controlled value either
-        # exfiltrates the request payload (incl. messages + admin-set
-        # tokens) to the attacker's host, or coerces the proxy into
-        # authenticating against the attacker's host with admin secrets.
-        "aws_bedrock_runtime_endpoint",
-        "langsmith_base_url",
-        "langfuse_host",
-        "posthog_host",
-        "braintrust_host",
-        "slack_webhook_url",
-        # Provider-specific endpoint overrides that flow into the outbound
-        # request via ``optional_params``. Same threat as ``api_base``:
-        # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker
-        # S3; ``sagemaker_base_url`` redirects all SageMaker traffic;
-        # ``deployment_url`` redirects SAP deployments.
-        "s3_endpoint_url",
-        "sagemaker_base_url",
-        "deployment_url",
-    ]
-
-    # The blocklist is enforced unconditionally. Legitimate clientside
-    # credential / endpoint passthrough goes through one of the two
-    # explicit admin opt-ins (``general_settings.allow_client_side_credentials``
-    # proxy-wide or ``configurable_clientside_auth_params`` per deployment).
-    # Historically there was a third, *implicit*, *caller-controlled* path:
-    # ``check_complete_credentials`` returned True when the caller supplied
-    # any non-empty ``api_key``, which made the entire blocklist a no-op.
-    # That bypass turned every missing entry on the blocklist into an
-    # exploitable SSRF / credential-exfil hole — see GHSA-jh89-88fc-qrfp,
-    # GHSA-3frq-6r6h-7j64, and the chain of veria-admin findings (Dv_m860l,
-    # b_yRJeQ5, stN90yjP, LBlyOAc8, U2TD78kg). Removed: the blocklist now
-    # has a single, predictable failure mode for missing entries (a 400),
-    # not a credential leak.
-    for param in banned_params:
-        if param in request_body:
-            if general_settings.get("allow_client_side_credentials") is True:
-                return True
-            elif (
-                _allow_model_level_clientside_configurable_parameters(
-                    model=model,
-                    param=param,
-                    request_body_value=request_body[param],
-                    llm_router=llm_router,
-                )
-                is True
-            ):
-                return True
-            raise ValueError(
-                f"Rejected Request: {param} is not allowed in request body. "
-                "Clientside passthrough requires explicit admin opt-in via "
-                "either `general_settings.allow_client_side_credentials = true` "
-                "(proxy-wide) or `configurable_clientside_auth_params` on the "
-                "deployment in your proxy config.yaml. "
-                "Relevant Issue: https://huntr.com/bounties/4001e1a2-7b7a-4776-a3ae-e6692ec3d997",
-            )
-
+    _check_banned_params(request_body, general_settings, llm_router, model)
+    for nested_key in _NESTED_CONFIG_KEYS:
+        nested = request_body.get(nested_key)
+        if isinstance(nested, dict):
+            _check_banned_params(nested, general_settings, llm_router, model)
     return True
 
 
@@ -942,20 +976,257 @@ def get_end_user_id_from_request_body(
     return None
 
 
-def get_model_from_request(
-    request_data: dict, route: str
-) -> Optional[Union[str, List[str]]]:
-    # First try to get model from request_data
-    model = request_data.get("model") or request_data.get("target_model_names")
+MODEL_ROUTING_HEADER_NAME = "x-litellm-model"
+_MODEL_ROUTING_ROUTE_MARKERS = (
+    "/files",
+    "/batches",
+    "/vector_stores",
+    "/skills",
+    "/evals",
+    "/fine_tuning",
+    "/videos",
+)
+_MODEL_ROUTING_HEADER_OR_QUERY_ROUTE_MARKERS = (
+    "/files",
+    "/batches",
+    "/skills",
+    "/evals",
+)
+_MODEL_ROUTING_QUERY_TARGET_MODEL_ROUTE_MARKERS = (
+    "/files",
+    "/batches",
+    "/fine_tuning",
+)
+_MODEL_ROUTING_BODY_TARGET_MODEL_ROUTE_MARKERS = (
+    "/files",
+    "/batches",
+    "/vector_stores",
+)
+_MODEL_ROUTING_COMPLETION_MODEL_ROUTE_MARKERS = ("/evals",)
+_MODEL_ROUTING_ID_FIELDS = (
+    "file_id",
+    "input_file_id",
+    "output_file_id",
+    "error_file_id",
+    "batch_id",
+    "fine_tuning_job_id",
+    "training_file",
+    "validation_file",
+    "vector_store_id",
+    "video_id",
+    "character_id",
+)
 
-    if model is not None:
-        model_names = model.split(",")
-        if len(model_names) == 1:
-            model = model_names[0].strip()
+
+def _append_model_candidates(candidates: List[str], value: Any) -> None:
+    if value is None:
+        return
+
+    values = value if isinstance(value, (list, tuple, set)) else [value]
+    for item in values:
+        if item is None:
+            continue
+        if isinstance(item, str):
+            model_names = [model.strip() for model in item.split(",")]
         else:
-            model = [m.strip() for m in model_names]
+            model_names = [str(item).strip()]
+        candidates.extend(model for model in model_names if model)
 
-    # If model not in request_data, try to extract from route
+
+def _dedupe_model_candidates(candidates: List[str]) -> List[str]:
+    deduped: List[str] = []
+    for model in candidates:
+        if model not in deduped:
+            deduped.append(model)
+    return deduped
+
+
+def _get_case_insensitive_mapping_value(
+    mapping: Optional[Mapping[str, Any]], key: str
+) -> Any:
+    if not mapping:
+        return None
+    if key in mapping:
+        return mapping[key]
+    key_lower = key.lower()
+    for mapping_key, value in mapping.items():
+        if str(mapping_key).lower() == key_lower:
+            return value
+    return None
+
+
+def _route_matches_any_marker(route: str, markers: Tuple[str, ...]) -> bool:
+    normalized_route = route.lower()
+    return any(marker in normalized_route for marker in markers)
+
+
+def _route_uses_model_routing_sources(route: str) -> bool:
+    return _route_matches_any_marker(route=route, markers=_MODEL_ROUTING_ROUTE_MARKERS)
+
+
+def _extract_models_from_managed_resource_id(
+    resource_id: Any, resource_id_field: Optional[str] = None
+) -> List[str]:
+    if not isinstance(resource_id, str) or not resource_id:
+        return []
+
+    candidates: List[str] = []
+
+    try:
+        from litellm.proxy.openai_files_endpoints.common_utils import (
+            _is_base64_encoded_unified_file_id,
+            decode_model_from_file_id,
+            get_model_id_from_unified_batch_id,
+            get_models_from_unified_file_id,
+        )
+
+        _append_model_candidates(
+            candidates=candidates, value=decode_model_from_file_id(resource_id)
+        )
+        unified_file_id = _is_base64_encoded_unified_file_id(resource_id)
+        if unified_file_id:
+            _append_model_candidates(
+                candidates=candidates,
+                value=get_models_from_unified_file_id(unified_file_id),
+            )
+            _append_model_candidates(
+                candidates=candidates,
+                value=get_model_id_from_unified_batch_id(unified_file_id),
+            )
+    except Exception as e:
+        verbose_proxy_logger.debug(
+            "Unable to extract model from managed file/batch ID: %s", str(e)
+        )
+
+    try:
+        from litellm.llms.base_llm.managed_resources.utils import parse_unified_id
+
+        parsed_id = parse_unified_id(resource_id)
+        if parsed_id:
+            _append_model_candidates(
+                candidates=candidates, value=parsed_id.get("model_id")
+            )
+            _append_model_candidates(
+                candidates=candidates, value=parsed_id.get("target_model_names")
+            )
+    except Exception as e:
+        verbose_proxy_logger.debug(
+            "Unable to extract model from unified managed resource ID: %s", str(e)
+        )
+
+    if resource_id_field in ("video_id", "character_id"):
+        try:
+            from litellm.types.videos.utils import (
+                decode_character_id_with_provider,
+                decode_video_id_with_provider,
+            )
+
+            if resource_id_field == "video_id":
+                _append_model_candidates(
+                    candidates=candidates,
+                    value=decode_video_id_with_provider(resource_id).get("model_id"),
+                )
+            else:
+                _append_model_candidates(
+                    candidates=candidates,
+                    value=decode_character_id_with_provider(resource_id).get(
+                        "model_id"
+                    ),
+                )
+        except Exception as e:
+            verbose_proxy_logger.debug(
+                "Unable to extract model from managed video/character ID: %s", str(e)
+            )
+
+    return _dedupe_model_candidates(candidates)
+
+
+def _extract_model_candidates_from_request(
+    request_data: dict,
+    route: str,
+    request_headers: Optional[Mapping[str, Any]] = None,
+    request_query_params: Optional[Mapping[str, Any]] = None,
+) -> List[str]:
+    candidates: List[str] = []
+    uses_model_routing_sources = _route_uses_model_routing_sources(route=route)
+    uses_header_or_query_model_sources = _route_matches_any_marker(
+        route=route, markers=_MODEL_ROUTING_HEADER_OR_QUERY_ROUTE_MARKERS
+    )
+    uses_query_target_model_sources = _route_matches_any_marker(
+        route=route, markers=_MODEL_ROUTING_QUERY_TARGET_MODEL_ROUTE_MARKERS
+    )
+    uses_body_target_model_sources = _route_matches_any_marker(
+        route=route, markers=_MODEL_ROUTING_BODY_TARGET_MODEL_ROUTE_MARKERS
+    )
+    uses_completion_model_sources = _route_matches_any_marker(
+        route=route, markers=_MODEL_ROUTING_COMPLETION_MODEL_ROUTE_MARKERS
+    )
+
+    body_model = request_data.get("model")
+    _append_model_candidates(candidates, body_model)
+    if uses_body_target_model_sources or not body_model:
+        _append_model_candidates(candidates, request_data.get("target_model_names"))
+    if uses_completion_model_sources and isinstance(
+        request_data.get("completion"), dict
+    ):
+        _append_model_candidates(candidates, request_data["completion"].get("model"))
+
+    if uses_model_routing_sources:
+        if uses_header_or_query_model_sources:
+            _append_model_candidates(
+                candidates,
+                _get_case_insensitive_mapping_value(request_query_params, "model"),
+            )
+            _append_model_candidates(
+                candidates,
+                _get_case_insensitive_mapping_value(
+                    request_headers, MODEL_ROUTING_HEADER_NAME
+                ),
+            )
+        if uses_query_target_model_sources:
+            _append_model_candidates(
+                candidates,
+                _get_case_insensitive_mapping_value(
+                    request_query_params, "target_model_names"
+                ),
+            )
+
+        for field in _MODEL_ROUTING_ID_FIELDS:
+            _append_model_candidates(
+                candidates,
+                _extract_models_from_managed_resource_id(
+                    request_data.get(field), resource_id_field=field
+                ),
+            )
+
+    return _dedupe_model_candidates(candidates)
+
+
+def _format_model_candidates(
+    candidates: List[str],
+) -> Optional[Union[str, List[str]]]:
+    if not candidates:
+        return None
+    if len(candidates) == 1:
+        return candidates[0]
+    return candidates
+
+
+def get_model_from_request(
+    request_data: dict,
+    route: str,
+    request_headers: Optional[Mapping[str, Any]] = None,
+    request_query_params: Optional[Mapping[str, Any]] = None,
+) -> Optional[Union[str, List[str]]]:
+    candidates = _extract_model_candidates_from_request(
+        request_data=request_data,
+        route=route,
+        request_headers=request_headers,
+        request_query_params=request_query_params,
+    )
+    model = _format_model_candidates(candidates)
+
+    # If no explicit model was found, try to extract from route
     if model is None:
         # Parse model from route that follows the pattern /openai/deployments/{model}/*
         match = re.match(r"/openai/deployments/([^/]+)", route)

litellm/proxy/auth/handle_jwt.py

@@ -6,6 +6,8 @@ Currently only supports admin.
 JWT token must have 'litellm_proxy_admin' in scope.
 """
 
+from __future__ import annotations
+
 import fnmatch
 import hashlib
 import os
@@ -20,7 +22,6 @@ import jwt
 from jwt.api_jwk import PyJWK
 
 from litellm._logging import verbose_proxy_logger
-from litellm.caching.caching import DualCache
 from litellm.constants import DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL
 from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
 from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
@@ -46,6 +47,7 @@ from litellm.proxy._types import (
 )
 from litellm.proxy.auth.auth_checks import can_team_access_model
 from litellm.proxy.auth.route_checks import RouteChecks
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
 from litellm.proxy.utils import PrismaClient, ProxyLogging
 
 from .auth_checks import (
@@ -73,7 +75,7 @@ class JWTHandler:
     """
 
     prisma_client: Optional[PrismaClient]
-    user_api_key_cache: DualCache
+    user_api_key_cache: UserApiKeyCache
     # Supported algos: https://pyjwt.readthedocs.io/en/stable/algorithms.html
     # "Warning: Make sure not to mix symmetric and asymmetric algorithms that interpret
     #   the key in different ways (e.g. HS* and RS*)."
@@ -99,7 +101,7 @@ class JWTHandler:
     def update_environment(
         self,
         prisma_client: Optional[PrismaClient],
-        user_api_key_cache: DualCache,
+        user_api_key_cache: UserApiKeyCache,
         litellm_jwtauth: LiteLLM_JWTAuth,
         leeway: int = 0,
     ) -> None:
@@ -952,7 +954,7 @@ class JWTAuthManager:
         jwt_handler: JWTHandler,
         jwt_valid_token: dict,
         prisma_client: Optional[PrismaClient],
-        user_api_key_cache: DualCache,
+        user_api_key_cache: UserApiKeyCache,
         parent_otel_span: Optional[Span],
         proxy_logging_obj: ProxyLogging,
     ) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]:
@@ -1045,7 +1047,7 @@ class JWTAuthManager:
         route: str,
         jwt_handler: JWTHandler,
         prisma_client: Optional[PrismaClient],
-        user_api_key_cache: DualCache,
+        user_api_key_cache: UserApiKeyCache,
         parent_otel_span: Optional[Span],
         proxy_logging_obj: ProxyLogging,
     ) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]:
@@ -1133,7 +1135,7 @@ class JWTAuthManager:
         valid_user_email: Optional[bool],
         jwt_handler: JWTHandler,
         prisma_client: Optional[PrismaClient],
-        user_api_key_cache: DualCache,
+        user_api_key_cache: UserApiKeyCache,
         parent_otel_span: Optional[Span],
         proxy_logging_obj: ProxyLogging,
         route: str,
@@ -1349,7 +1351,7 @@ class JWTAuthManager:
         jwt_valid_token: dict,
         user_object: Optional[LiteLLM_UserTable],
         prisma_client: Optional[PrismaClient],
-        user_api_key_cache: Optional[DualCache] = None,
+        user_api_key_cache: Optional[UserApiKeyCache] = None,
     ) -> None:
         """
         Sync user role and team memberships with JWT claims
@@ -1377,7 +1379,8 @@ class JWTAuthManager:
             if user_api_key_cache is not None:
                 await user_api_key_cache.async_set_cache(
                     key=user_object.user_id,
-                    value=user_object.model_dump(),
+                    value=user_object,
+                    model_type=LiteLLM_UserTable,
                     ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
                 )
 
@@ -1400,7 +1403,8 @@ class JWTAuthManager:
             if user_api_key_cache is not None:
                 await user_api_key_cache.async_set_cache(
                     key=user_object.user_id,
-                    value=user_object.model_dump(),
+                    value=user_object,
+                    model_type=LiteLLM_UserTable,
                     ttl=DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
                 )
         return None
@@ -1412,7 +1416,7 @@ class JWTAuthManager:
         request_headers: Optional[dict],
         jwt_handler: JWTHandler,
         prisma_client: Optional[PrismaClient],
-        user_api_key_cache: DualCache,
+        user_api_key_cache: UserApiKeyCache,
         parent_otel_span: Optional[Span],
         proxy_logging_obj: ProxyLogging,
     ) -> None:
@@ -1456,7 +1460,7 @@ class JWTAuthManager:
         user_object: Optional[LiteLLM_UserTable],
         user_id: Optional[str],
         prisma_client: Optional[PrismaClient],
-        user_api_key_cache: DualCache,
+        user_api_key_cache: UserApiKeyCache,
         parent_otel_span: Optional[Span],
         proxy_logging_obj: ProxyLogging,
         team_id_upsert: Optional[bool],
@@ -1514,7 +1518,7 @@ class JWTAuthManager:
         general_settings: dict,
         route: str,
         prisma_client: Optional[PrismaClient],
-        user_api_key_cache: DualCache,
+        user_api_key_cache: UserApiKeyCache,
         parent_otel_span: Optional[Span],
         proxy_logging_obj: ProxyLogging,
         request_headers: Optional[dict] = None,

litellm/proxy/auth/user_api_key_auth.py

@@ -11,7 +11,7 @@ import asyncio
 import re
 import secrets
 from datetime import datetime, timezone
-from typing import Any, List, Optional, Tuple, cast
+from typing import Any, List, Optional, Tuple, Union, cast
 
 import fastapi
 from fastapi import HTTPException, Request, WebSocket, status
@@ -20,7 +20,6 @@ from fastapi.security.api_key import APIKeyHeader
 import litellm
 from litellm._logging import verbose_logger, verbose_proxy_logger
 from litellm._service_logger import ServiceLogging
-from litellm.caching import DualCache
 from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
 from litellm.litellm_core_utils.dd_tracing import tracer
 from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
@@ -60,9 +59,11 @@ from litellm.proxy.auth.oauth2_check import Oauth2Handler
 from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request
 from litellm.proxy.auth.route_checks import RouteChecks
 from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
 from litellm.proxy.common_utils.http_parsing_utils import (
     _read_request_body,
     _safe_get_request_headers,
+    _safe_get_request_query_params,
     populate_request_with_path_params,
 )
 from litellm.proxy.common_utils.realtime_utils import _realtime_request_body
@@ -118,6 +119,29 @@ azure_apim_header = APIKeyHeader(
 )
 
 
+def _get_model_from_request_context(
+    request_data: dict,
+    route: str,
+    request: Optional[Request],
+) -> Optional[Union[str, List[str]]]:
+    return get_model_from_request(
+        request_data=request_data,
+        route=route,
+        request_headers=_safe_get_request_headers(request=request),
+        request_query_params=_safe_get_request_query_params(request=request),
+    )
+
+
+def _get_model_names_for_budget_checks(
+    model: Optional[Union[str, List[str]]],
+) -> List[str]:
+    if model is None:
+        return []
+    if isinstance(model, str):
+        return [model]
+    return model
+
+
 def _get_bearer_token_or_received_api_key(api_key: str) -> str:
     if api_key.startswith("Bearer "):  # ensure Bearer token passed in
         api_key = api_key.replace("Bearer ", "")  # extract the token
@@ -329,7 +353,7 @@ _global_spend_coordinator = EventDrivenCacheCoordinator(log_prefix="[GLOBAL SPEN
 
 async def _fetch_global_spend_with_event_coordination(
     cache_key: str,
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     prisma_client: PrismaClient,
 ) -> Optional[float]:
     """
@@ -345,14 +369,14 @@ async def _fetch_global_spend_with_event_coordination(
 
     return await _global_spend_coordinator.get_or_load(
         cache_key=cache_key,
-        cache=user_api_key_cache,
+        cache=user_api_key_cache,  # pyright: ignore[reportArgumentType]
         load_fn=_load_global_spend,
     )
 
 
 async def get_global_proxy_spend(
     litellm_proxy_admin_name: str,
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     prisma_client: Optional[PrismaClient],
     token: str,
     proxy_logging_obj: ProxyLogging,
@@ -473,7 +497,12 @@ async def check_api_key_for_custom_headers_or_pass_through_endpoints(
         for endpoint in pass_through_endpoints:
             if isinstance(endpoint, dict) and endpoint.get("path", "") == route:
                 ## IF AUTH DISABLED
-                if endpoint.get("auth") is not True:
+                # Default to True: a config dict with no ``auth`` key
+                # otherwise produced an unauthenticated forwarder. The
+                # Pydantic ``PassThroughGenericEndpoint.auth`` default
+                # is also True, but raw config dicts skip that path —
+                # so this runtime check has to default to True too.
+                if endpoint.get("auth", True) is not True:
                     return UserAPIKeyAuth()
                 ## IF AUTH ENABLED
                 ### IF CUSTOM PARSER REQUIRED
@@ -505,7 +534,7 @@ async def _resolve_jwt_to_virtual_key(
     jwt_claims: dict,
     jwt_handler: JWTHandler,
     prisma_client: Optional[PrismaClient],
-    user_api_key_cache: DualCache,
+    user_api_key_cache: UserApiKeyCache,
     parent_otel_span: Optional[Span],
     proxy_logging_obj: ProxyLogging,
 ) -> Optional[UserAPIKeyAuth]:
@@ -879,7 +908,11 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                     )
 
                     # Check if model has zero cost - if so, skip all budget checks
-                    model = get_model_from_request(request_data, route)
+                    model = _get_model_from_request_context(
+                        request_data=request_data,
+                        route=route,
+                        request=request,
+                    )
                     skip_budget_checks = False
                     if model is not None and llm_router is not None:
                         from litellm.proxy.auth.auth_checks import _is_model_cost_zero
@@ -1107,9 +1140,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
             is_master_key_valid = False
 
         ## VALIDATE MASTER KEY ##
-        try:
-            assert isinstance(master_key, str)
-        except Exception:
+        if not isinstance(master_key, str):
             raise HTTPException(
                 status_code=500,
                 detail={
@@ -1179,11 +1210,15 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                     if len(api_key) > 8
                     else "****"
                 )
-                assert api_key.startswith(
-                    "sk-"
-                ), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
-                    _masked_key
-                )  # prevent token hashes from being used
+                if not api_key.startswith("sk-"):
+                    raise HTTPException(
+                        status_code=status.HTTP_401_UNAUTHORIZED,
+                        detail=(
+                            "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format(
+                                _masked_key
+                            )
+                        ),
+                    )  # prevent token hashes from being used
             else:
                 verbose_logger.warning(
                     "litellm.proxy.proxy_server.user_api_key_auth(): Warning - Key is not a string. Got type={}".format(
@@ -1247,6 +1282,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                 valid_token=valid_token,
                 request_data=request_data,
                 route=route,
+                request=request,
                 llm_model_list=llm_model_list,
                 llm_router=llm_router,
             )
@@ -1272,7 +1308,11 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                     user_obj = None
 
             # Check 2a. Check if model has zero cost - if so, skip all budget checks
-            model = get_model_from_request(request_data, route)
+            model = _get_model_from_request_context(
+                request_data=request_data,
+                route=route,
+                request=request,
+            )
             skip_budget_checks = False
             if model is not None and llm_router is not None:
                 from litellm.proxy.auth.auth_checks import _is_model_cost_zero
@@ -1291,7 +1331,8 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                     _cache_key = f"{valid_token.team_id}_{valid_token.user_id}"
 
                     team_member_info = await user_api_key_cache.async_get_cache(
-                        key=_cache_key
+                        key=_cache_key,
+                        model_type=LiteLLM_TeamMembership,
                     )
                     if team_member_info is None:
                         # read from DB
@@ -1299,18 +1340,23 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                         _team_id = valid_token.team_id
 
                         if _user_id is not None and _team_id is not None:
-                            team_member_info = await prisma_client.db.litellm_teammembership.find_first(
+                            _db_member = await prisma_client.db.litellm_teammembership.find_first(
                                 where={
                                     "user_id": _user_id,
                                     "team_id": _team_id,
                                 },  # type: ignore
                                 include={"litellm_budget_table": True},
                             )
-                            await user_api_key_cache.async_set_cache(
-                                key=_cache_key,
-                                value=team_member_info,
-                                ttl=5,
-                            )
+                            if _db_member is not None:
+                                team_member_info = LiteLLM_TeamMembership(
+                                    **_db_member.dict()
+                                )
+                                await user_api_key_cache.async_set_cache(
+                                    key=_cache_key,
+                                    value=team_member_info,
+                                    model_type=LiteLLM_TeamMembership,
+                                    ttl=5,
+                                )
 
                     if (
                         team_member_info is not None
@@ -1390,21 +1436,29 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
 
                     # Check 5. Token Model Spend is under Model budget
                     max_budget_per_model = valid_token.model_max_budget
-                    current_model = request_data.get("model", None)
+                    current_model = _get_model_from_request_context(
+                        request_data=request_data,
+                        route=route,
+                        request=request,
+                    )
+                    current_models = _get_model_names_for_budget_checks(
+                        model=current_model
+                    )
 
                     if (
                         max_budget_per_model is not None
                         and isinstance(max_budget_per_model, dict)
                         and len(max_budget_per_model) > 0
                         and prisma_client is not None
-                        and current_model is not None
+                        and current_models
                         and valid_token.token is not None
                     ):
                         ## GET THE SPEND FOR THIS MODEL
-                        await model_max_budget_limiter.is_key_within_model_budget(
-                            user_api_key_dict=valid_token,
-                            model=current_model,
-                        )
+                        for model_name in current_models:
+                            await model_max_budget_limiter.is_key_within_model_budget(
+                                user_api_key_dict=valid_token,
+                                model=model_name,
+                            )
 
                     # Check 5b. End-user model max budget
                     end_user_mmb = valid_token.end_user_model_max_budget
@@ -1412,14 +1466,15 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                         end_user_mmb is not None
                         and isinstance(end_user_mmb, dict)
                         and len(end_user_mmb) > 0
-                        and current_model is not None
+                        and current_models
                         and valid_token.end_user_id is not None
                     ):
-                        await model_max_budget_limiter.is_end_user_within_model_budget(
-                            end_user_id=valid_token.end_user_id,
-                            end_user_model_max_budget=end_user_mmb,
-                            model=current_model,
-                        )
+                        for model_name in current_models:
+                            await model_max_budget_limiter.is_end_user_within_model_budget(
+                                end_user_id=valid_token.end_user_id,
+                                end_user_model_max_budget=end_user_mmb,
+                                model=model_name,
+                            )
 
             # Check 6: Additional Common Checks across jwt + key auth
             if valid_token.team_id is not None:
@@ -1457,9 +1512,13 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
             else:
                 valid_token.team_object_permission = None
 
-            await user_api_key_cache.async_set_cache(
-                key=valid_token.team_id, value=_team_obj
-            )  # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py
+            # Only cache when the key is a real team_id (non-team keys must not use key=None).
+            if valid_token.team_id is not None and _team_obj is not None:
+                await user_api_key_cache.async_set_cache(
+                    key=valid_token.team_id,
+                    value=_team_obj,
+                    model_type=LiteLLM_TeamTableCachedObj,
+                )  # save team table in cache - used for tpm/rpm limiting - tpm_rpm_limiter.py
 
             # Fetch project object if key belongs to a project
             _project_obj = None
@@ -1845,10 +1904,12 @@ async def _run_centralized_common_checks(
         user_api_key_auth_obj.project_metadata = project_object.metadata
         user_api_key_auth_obj.project_alias = project_object.project_alias
 
-    skip_budget_checks = False
-    model = get_model_from_request(request_data, route)
-    if model is not None and llm_router is not None:
-        skip_budget_checks = _is_model_cost_zero(model=model, llm_router=llm_router)
+    skip_budget_checks = _should_skip_budget_checks(
+        request_data=request_data,
+        route=route,
+        request=request,
+        llm_router=llm_router,
+    )
 
     _ = await common_checks(
         request=request,
@@ -1866,6 +1927,21 @@ async def _run_centralized_common_checks(
         project_object=project_object,
     )
 
+    await _reserve_budget_after_common_checks(
+        user_api_key_auth_obj=user_api_key_auth_obj,
+        request_data=request_data,
+        route=route,
+        llm_router=llm_router,
+        team_object=team_object,
+        user_object=user_object,
+        end_user_id=end_user_id,
+        end_user_object=end_user_object,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+        skip_budget_checks=skip_budget_checks,
+    )
+
 
 async def _noop_none() -> None:
     """Sentinel coroutine for asyncio.gather when a fetch is unnecessary
@@ -1873,6 +1949,59 @@ async def _noop_none() -> None:
     return None
 
 
+async def _reserve_budget_after_common_checks(
+    user_api_key_auth_obj: UserAPIKeyAuth,
+    request_data: dict,
+    route: str,
+    llm_router: Optional[Any],
+    team_object: Optional[LiteLLM_TeamTableCachedObj],
+    user_object: Optional[LiteLLM_UserTable],
+    prisma_client: Optional[PrismaClient],
+    user_api_key_cache: UserApiKeyCache,
+    proxy_logging_obj: ProxyLogging,
+    skip_budget_checks: bool,
+    end_user_id: Optional[str] = None,
+    end_user_object: Optional[LiteLLM_EndUserTable] = None,
+) -> None:
+    user_api_key_auth_obj.budget_reservation = None
+    if skip_budget_checks:
+        return
+
+    from litellm.proxy.spend_tracking.budget_reservation import (
+        reserve_budget_for_request,
+    )
+
+    user_api_key_auth_obj.budget_reservation = await reserve_budget_for_request(
+        request_body=request_data,
+        route=route,
+        llm_router=llm_router,
+        valid_token=user_api_key_auth_obj,
+        team_object=team_object,
+        user_object=user_object,
+        prisma_client=prisma_client,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+        end_user_id=end_user_id,
+        end_user_object=end_user_object,
+    )
+
+
+def _should_skip_budget_checks(
+    request_data: dict,
+    route: str,
+    request: Optional[Request],
+    llm_router: Optional[Any],
+) -> bool:
+    model = _get_model_from_request_context(
+        request_data=request_data,
+        route=route,
+        request=request,
+    )
+    if model is not None and llm_router is not None:
+        return _is_model_cost_zero(model=model, llm_router=llm_router)
+    return False
+
+
 @tracer.wrap()
 async def user_api_key_auth(
     request: Request,
@@ -1910,6 +2039,7 @@ async def user_api_key_auth(
         request_data=request_data,
         custom_litellm_key_header=custom_litellm_key_header,
     )
+    user_api_key_auth_obj.budget_reservation = None
 
     ## ENSURE DISABLE ROUTE WORKS ACROSS ALL USER AUTH FLOWS ##
     RouteChecks.should_call_route(route=route, valid_token=user_api_key_auth_obj)
@@ -2117,6 +2247,7 @@ async def _enforce_key_and_fallback_model_access(
     valid_token: UserAPIKeyAuth,
     request_data: dict,
     route: str,
+    request: Optional[Request],
     llm_model_list: Optional[list],
     llm_router: Optional[Any],
 ) -> None:
@@ -2135,7 +2266,11 @@ async def _enforce_key_and_fallback_model_access(
     ):
         pass
     else:
-        model = get_model_from_request(request_data, route)
+        model = _get_model_from_request_context(
+            request_data=request_data,
+            route=route,
+            request=request,
+        )
         fallback_models = cast(
             Optional[List[ALL_FALLBACK_MODEL_VALUES]],
             request_data.get("fallbacks", None),
@@ -2222,11 +2357,17 @@ async def _run_post_custom_auth_checks(
             valid_token=valid_token,
             request_data=request_data,
             route=route,
+            request=request,
             llm_model_list=llm_model_list,
             llm_router=llm_router,
         )
 
-    current_model = request_data.get("model", None)
+    current_model = _get_model_from_request_context(
+        request_data=request_data,
+        route=route,
+        request=request,
+    )
+    current_models = _get_model_names_for_budget_checks(model=current_model)
 
     # 3. Check key-level model_max_budget
     max_budget_per_model = valid_token.model_max_budget
@@ -2234,13 +2375,14 @@ async def _run_post_custom_auth_checks(
         max_budget_per_model is not None
         and isinstance(max_budget_per_model, dict)
         and len(max_budget_per_model) > 0
-        and current_model is not None
+        and current_models
         and valid_token.token is not None
     ):
-        await model_max_budget_limiter.is_key_within_model_budget(
-            user_api_key_dict=valid_token,
-            model=current_model,
-        )
+        for model_name in current_models:
+            await model_max_budget_limiter.is_key_within_model_budget(
+                user_api_key_dict=valid_token,
+                model=model_name,
+            )
 
     # 4. Check end-user model_max_budget
     end_user_mmb = valid_token.end_user_model_max_budget
@@ -2248,14 +2390,15 @@ async def _run_post_custom_auth_checks(
         end_user_mmb is not None
         and isinstance(end_user_mmb, dict)
         and len(end_user_mmb) > 0
-        and current_model is not None
+        and current_models
         and valid_token.end_user_id is not None
     ):
-        await model_max_budget_limiter.is_end_user_within_model_budget(
-            end_user_id=valid_token.end_user_id,
-            end_user_model_max_budget=end_user_mmb,
-            model=current_model,
-        )
+        for model_name in current_models:
+            await model_max_budget_limiter.is_end_user_within_model_budget(
+                end_user_id=valid_token.end_user_id,
+                end_user_model_max_budget=end_user_mmb,
+                model=model_name,
+            )
 
     # team / user / end_user / project context objects are fetched by
     # the centralized common_checks gate in user_api_key_auth after

litellm/proxy/client/README.md

@@ -313,23 +313,24 @@ sequenceDiagram
     participant Proxy as LiteLLM Proxy
     participant SSO as SSO Provider
     
-    CLI->>CLI: Generate key ID (sk-uuid)
-    CLI->>Browser: Open /sso/key/generate?source=litellm-cli&key=sk-uuid
+    CLI->>Proxy: POST /sso/cli/start
+    Proxy->>CLI: Return login_id, poll_secret, user_code
+    CLI->>Browser: Open /sso/key/generate?source=litellm-cli&key=login_id
     
-    Browser->>Proxy: GET /sso/key/generate?source=litellm-cli&key=sk-uuid
-    Proxy->>Proxy: Set cli_state = litellm-session-token:sk-uuid
-    Proxy->>SSO: Redirect with state=litellm-session-token:sk-uuid
+    Browser->>Proxy: GET /sso/key/generate?source=litellm-cli&key=login_id
+    Proxy->>Proxy: Set cli_state = litellm-session-token:login_id
+    Proxy->>SSO: Redirect with state=litellm-session-token:login_id
     
     SSO->>Browser: Show login page
     Browser->>SSO: User authenticates
-    SSO->>Proxy: Redirect to /sso/callback?state=litellm-session-token:sk-uuid
+    SSO->>Proxy: Redirect to /sso/callback?state=litellm-session-token:login_id
     
     Proxy->>Proxy: Check if state starts with "litellm-session-token:"
-    Proxy->>Proxy: Generate API key with ID=sk-uuid
-    Proxy->>Browser: Show success page
+    Proxy->>Browser: Prompt for user_code
+    Browser->>Proxy: POST /sso/cli/complete/login_id
     
-    CLI->>Proxy: Poll /sso/cli/poll/sk-uuid
-    Proxy->>CLI: Return {"status": "ready", "key": "sk-uuid"}
+    CLI->>Proxy: Poll /sso/cli/poll/login_id with poll_secret header
+    Proxy->>CLI: Return {"status": "ready", "key": "jwt"}
     CLI->>CLI: Save key to ~/.litellm/token.json
 ```
 
@@ -343,13 +344,13 @@ The CLI provides three authentication commands:
 
 ### Authentication Flow Steps
 
-1. **Generate Session ID**: CLI generates a unique key ID (`sk-{uuid}`)
-2. **Open Browser**: CLI opens browser to `/sso/key/generate` with CLI source and key parameters
-3. **SSO Redirect**: Proxy sets the formatted state (`litellm-session-token:sk-uuid`) as OAuth state parameter and redirects to SSO provider
+1. **Start Session**: CLI creates a short-lived login session with `/sso/cli/start`
+2. **Open Browser**: CLI opens browser to `/sso/key/generate` with CLI source and login ID parameters
+3. **SSO Redirect**: Proxy sets the formatted state (`litellm-session-token:{login_id}`) as OAuth state parameter and redirects to SSO provider
 4. **User Authentication**: User completes SSO authentication in browser
 5. **Callback Processing**: SSO provider redirects back to proxy with state parameter
-6. **Key Generation**: Proxy detects CLI login (state starts with "litellm-session-token:") and generates API key with pre-specified ID
-7. **Polling**: CLI polls `/sso/cli/poll/{key_id}` endpoint until key is ready
+6. **User Code Verification**: Browser confirms the verification code shown in the CLI
+7. **Polling**: CLI polls `/sso/cli/poll/{login_id}` with the polling secret header until the JWT is ready
 8. **Token Storage**: CLI saves the authentication token to `~/.litellm/token.json`
 
 ### Benefits of This Approach
@@ -357,7 +358,7 @@ The CLI provides three authentication commands:
 - **No Local Server**: No need to run a local callback server
 - **Standard OAuth**: Uses OAuth 2.0 state parameter correctly
 - **Remote Compatible**: Works with remote proxy servers
-- **Secure**: Uses UUID session identifiers
+- **Secure**: Keeps the polling secret out of the browser handoff
 - **Simple Setup**: No additional OAuth redirect URL configuration needed
 
 ### Token Storage

litellm/proxy/client/cli/commands/auth.py

@@ -5,6 +5,7 @@ import time
 import webbrowser
 from pathlib import Path
 from typing import Any, Dict, List, Optional
+from urllib.parse import urlencode
 
 import click
 import requests
@@ -52,12 +53,16 @@ def clear_token() -> None:
         os.remove(token_file)
 
 
-def get_stored_api_key() -> Optional[str]:
-    """Get the stored API key from token file"""
-    # Use the SDK-level utility
+def get_stored_api_key(expected_base_url: Optional[str] = None) -> Optional[str]:
+    """Get the stored API key from token file.
+
+    If expected_base_url is provided, the key is only returned when it was
+    originally issued for that URL. This prevents credential leakage when the
+    CLI is pointed at a different (possibly malicious) server.
+    """
     from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key
 
-    return get_litellm_gateway_api_key()
+    return get_litellm_gateway_api_key(expected_base_url=expected_base_url)
 
 
 # Team selection utilities
@@ -241,7 +246,7 @@ def prompt_team_selection(teams: List[Dict[str, Any]]) -> Optional[Dict[str, Any
 
 
 def prompt_team_selection_fallback(
-    teams: List[Dict[str, Any]]
+    teams: List[Dict[str, Any]],
 ) -> Optional[Dict[str, Any]]:
     """Fallback team selection for non-interactive environments"""
     if not teams:
@@ -279,6 +284,7 @@ def prompt_team_selection_fallback(
 def _poll_for_ready_data(
     url: str,
     *,
+    headers: Optional[Dict[str, str]] = None,
     total_timeout: int = 300,
     poll_interval: int = 2,
     request_timeout: int = 10,
@@ -291,7 +297,10 @@ def _poll_for_ready_data(
 ) -> Optional[Dict[str, Any]]:
     for attempt in range(total_timeout // poll_interval):
         try:
-            response = requests.get(url, timeout=request_timeout)
+            request_kwargs: Dict[str, Any] = {"timeout": request_timeout}
+            if headers is not None:
+                request_kwargs["headers"] = headers
+            response = requests.get(url, **request_kwargs)
             if response.status_code == 200:
                 data = response.json()
                 status = data.get("status")
@@ -346,7 +355,23 @@ def _normalize_teams(teams, team_details):
     return []
 
 
-def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
+def _start_cli_sso_flow(base_url: str) -> Dict[str, Any]:
+    response = requests.post(f"{base_url}/sso/cli/start", timeout=10)
+    response.raise_for_status()
+    data = response.json()
+    required_fields = ("login_id", "poll_secret", "user_code")
+    if not all(isinstance(data.get(field), str) for field in required_fields):
+        raise ValueError("Invalid CLI SSO start response")
+    return data
+
+
+def _get_cli_sso_poll_headers(poll_secret: str) -> Dict[str, str]:
+    return {"x-litellm-cli-poll-secret": poll_secret}
+
+
+def _poll_for_authentication(
+    base_url: str, key_id: str, poll_secret: str
+) -> Optional[dict]:
     """
     Poll the server for authentication completion and handle team selection.
 
@@ -356,6 +381,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
     poll_url = f"{base_url}/sso/cli/poll/{key_id}"
     data = _poll_for_ready_data(
         poll_url,
+        headers=_get_cli_sso_poll_headers(poll_secret),
         pending_message="Still waiting for authentication...",
     )
     if not data:
@@ -373,6 +399,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
         jwt_with_team = _handle_team_selection_during_polling(
             base_url=base_url,
             key_id=key_id,
+            poll_secret=poll_secret,
             teams=normalized_teams,
         )
 
@@ -410,7 +437,7 @@ def _poll_for_authentication(base_url: str, key_id: str) -> Optional[dict]:
 
 
 def _handle_team_selection_during_polling(
-    base_url: str, key_id: str, teams: List[Dict[str, Any]]
+    base_url: str, key_id: str, poll_secret: str, teams: List[Dict[str, Any]]
 ) -> Optional[str]:
     """
     Handle team selection and re-poll with selected team_id.
@@ -441,6 +468,7 @@ def _handle_team_selection_during_polling(
     poll_url = f"{base_url}/sso/cli/poll/{key_id}?team_id={team_id}"
     data = _poll_for_ready_data(
         poll_url,
+        headers=_get_cli_sso_poll_headers(poll_secret),
         pending_message="Still waiting for team authentication...",
         other_status_message="Waiting for team authentication to complete...",
         http_error_log_every=10,
@@ -514,29 +542,24 @@ def _render_and_prompt_for_team_selection(teams: List[Dict[str, Any]]) -> Option
 @click.pass_context
 def login(ctx: click.Context):
     """Login to LiteLLM proxy using SSO authentication"""
-    from litellm._uuid import uuid
     from litellm.constants import LITELLM_CLI_SOURCE_IDENTIFIER
     from litellm.proxy.client.cli.interface import show_commands
 
     base_url = ctx.obj["base_url"]
 
-    # Check if we have an existing key to regenerate
-    existing_key = get_stored_api_key()
-
-    # Generate unique key ID for this login session
-    key_id = f"sk-{str(uuid.uuid4())}"
-
     try:
-        # Construct SSO login URL with CLI source and pre-generated key
-        sso_url = f"{base_url}/sso/key/generate?source={LITELLM_CLI_SOURCE_IDENTIFIER}&key={key_id}"
+        cli_sso_flow = _start_cli_sso_flow(base_url=base_url)
+        key_id = cli_sso_flow["login_id"]
+        poll_secret = cli_sso_flow["poll_secret"]
+        user_code = cli_sso_flow["user_code"]
 
-        # If we have an existing key, include it as a parameter to the login endpoint
-        # The server will encode it in the OAuth state parameter for the SSO flow
-        if existing_key:
-            sso_url += f"&existing_key={existing_key}"
+        sso_url = f"{base_url}/sso/key/generate?" + urlencode(
+            {"source": LITELLM_CLI_SOURCE_IDENTIFIER, "key": key_id}
+        )
 
         click.echo(f"Opening browser to: {sso_url}")
         click.echo("Please complete the SSO authentication in your browser...")
+        click.echo(f"Verification code: {user_code}")
         click.echo(f"Session ID: {key_id}")
 
         # Open browser
@@ -545,15 +568,19 @@ def login(ctx: click.Context):
         # Poll for authentication completion
         click.echo("Waiting for authentication...")
 
-        auth_result = _poll_for_authentication(base_url=base_url, key_id=key_id)
+        auth_result = _poll_for_authentication(
+            base_url=base_url, key_id=key_id, poll_secret=poll_secret
+        )
 
         if auth_result:
             api_key = auth_result["api_key"]
             user_id = auth_result["user_id"]
 
-            # Save token data (simplified for CLI - we just need the key)
+            # Save token data. base_url is stored so we can verify origin
+            # before reusing the key on a subsequent CLI invocation.
             save_token(
                 {
+                    "base_url": base_url.rstrip("/"),
                     "key": api_key,
                     "user_id": user_id or "cli-user",
                     "user_email": "unknown",

litellm/proxy/client/cli/main.py

@@ -74,9 +74,10 @@ def cli(ctx: click.Context, base_url: str, api_key: Optional[str]) -> None:
     """LiteLLM Proxy CLI - Manage your LiteLLM proxy server"""
     ctx.ensure_object(dict)
 
-    # If no API key provided via flag or environment variable, try to load from saved token
+    # If no API key provided via flag or environment variable, try to load from saved token.
+    # Pass base_url so we only use the stored key when it was issued for this server.
     if api_key is None:
-        api_key = get_stored_api_key()
+        api_key = get_stored_api_key(expected_base_url=base_url)
 
     ctx.obj["base_url"] = base_url
     ctx.obj["api_key"] = api_key

litellm/proxy/client/client.py

@@ -28,12 +28,17 @@ class Client:
             api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token.
             timeout: Request timeout in seconds (default: 30)
         """
-        self._base_url = base_url.rstrip("/")  # Remove trailing slash if present
-        self._api_key = get_litellm_gateway_api_key() or api_key
+        self._base_url = base_url.rstrip("/")
+        # Only use the stored CLI key when it was issued for this server.
+        self._api_key = api_key or get_litellm_gateway_api_key(
+            expected_base_url=self._base_url
+        )
 
         # Initialize resource clients
 
-        self.http = HTTPClient(base_url=base_url, api_key=api_key, timeout=timeout)
+        self.http = HTTPClient(
+            base_url=base_url, api_key=self._api_key, timeout=timeout
+        )
         self.models = ModelsManagementClient(
             base_url=self._base_url, api_key=self._api_key
         )

litellm/proxy/common_request_processing.py

@@ -97,6 +97,55 @@ def _serialize_http_exception_detail(
     return str(detail), None
 
 
+def _collect_response_file_search_vector_store_ids(data: Dict[str, Any]) -> set[str]:
+    vector_store_ids: set[str] = set()
+    tools = data.get("tools")
+    if not isinstance(tools, list):
+        return vector_store_ids
+
+    for tool in tools:
+        if not isinstance(tool, dict) or tool.get("type") != "file_search":
+            continue
+        ids = tool.get("vector_store_ids") or []
+        if not isinstance(ids, list):
+            raise HTTPException(
+                status_code=400,
+                detail={
+                    "error": "file_search.vector_store_ids must be a list of strings"
+                },
+            )
+        for vector_store_id in ids:
+            if not isinstance(vector_store_id, str) or not vector_store_id:
+                raise HTTPException(
+                    status_code=400,
+                    detail={
+                        "error": "file_search.vector_store_ids must be a list of strings"
+                    },
+                )
+            vector_store_ids.add(vector_store_id)
+
+    return vector_store_ids
+
+
+async def _authorize_response_file_search_vector_stores(
+    data: Dict[str, Any],
+    user_api_key_dict: UserAPIKeyAuth,
+) -> None:
+    vector_store_ids = _collect_response_file_search_vector_store_ids(data)
+    if not vector_store_ids:
+        return
+
+    from litellm.proxy.vector_store_endpoints.utils import (
+        assert_user_can_access_vector_store_id,
+    )
+
+    for vector_store_id in sorted(vector_store_ids):
+        await assert_user_can_access_vector_store_id(
+            vector_store_id=vector_store_id,
+            user_api_key_dict=user_api_key_dict,
+        )
+
+
 async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional[int]:
     """Parses an event line and returns an error code if present, else None."""
     event_line = (
@@ -744,6 +793,11 @@ class ProxyBaseLLMRequestProcessing:
             "aingest",
             "aretrieve_container",
             "adelete_container",
+            "aupload_container_file",
+            "alist_container_files",
+            "aretrieve_container_file",
+            "adelete_container_file",
+            "aretrieve_container_file_content",
             "acreate_skill",
             "alist_skills",
             "aget_skill",
@@ -786,6 +840,11 @@ class ProxyBaseLLMRequestProcessing:
             version=version,
             proxy_config=proxy_config,
         )
+        if route_type in {"aresponses", "_aresponses_websocket"}:
+            await _authorize_response_file_search_vector_stores(
+                data=self.data,
+                user_api_key_dict=user_api_key_dict,
+            )
 
         # Calculate request queue time after add_litellm_data_to_request
         # which sets arrival_time in proxy_server_request
@@ -1001,6 +1060,11 @@ class ProxyBaseLLMRequestProcessing:
             "aingest",
             "aretrieve_container",
             "adelete_container",
+            "aupload_container_file",
+            "alist_container_files",
+            "aretrieve_container_file",
+            "adelete_container_file",
+            "aretrieve_container_file_content",
             "acreate_skill",
             "alist_skills",
             "aget_skill",

litellm/proxy/common_utils/cache_coordinator.py

@@ -20,11 +20,27 @@ T = TypeVar("T")
 
 
 class AsyncCacheProtocol(Protocol):
-    """Protocol for cache backends used by EventDrivenCacheCoordinator."""
+    """Protocol for cache backends used by EventDrivenCacheCoordinator.
 
-    async def async_get_cache(self, key: str, **kwargs: Any) -> Any: ...
+    Matches ``DualCache`` / ``UserApiKeyCache`` call shapes (explicit optional params
+    before ``**kwargs``), not only ``(key, **kwargs)``, so overloads validate.
+    """
 
-    async def async_set_cache(self, key: str, value: Any, **kwargs: Any) -> Any: ...
+    async def async_get_cache(
+        self,
+        key: str,
+        parent_otel_span: Any = None,
+        local_only: bool = False,
+        **kwargs: Any,
+    ) -> Any: ...
+
+    async def async_set_cache(
+        self,
+        key: str,
+        value: Any,
+        local_only: bool = False,
+        **kwargs: Any,
+    ) -> Any: ...
 
 
 class EventDrivenCacheCoordinator:
@@ -36,6 +52,9 @@ class EventDrivenCacheCoordinator:
     - Other requests: wait for the signal, then read from cache.
 
     Create one instance per resource (e.g. one for global spend, one for feature flags).
+
+    Args:
+        log_prefix: Prefix for debug log messages.
     """
 
     def __init__(self, log_prefix: str = "[CACHE]"):

litellm/proxy/common_utils/cache_pydantic_utils.py

@@ -0,0 +1,93 @@
+"""
+DualCache presents a single API for reads and writes, but the two backends behave
+differently: the in-memory layer can store arbitrary Python objects (including live
+``BaseModel`` instances), while Redis persists strings and therefore needs JSON-safe
+payloads (``json.dumps`` on the Redis side).
+
+Call sites therefore see cache ``value`` / ``cached`` as effectively ``Any``: the same
+key may deserialize to a model on one process (memory hit) or to a ``dict`` after a
+Redis round-trip. ``CacheCodec`` centralizes encode/decode at that boundary:
+``CacheCodec.serialize`` before ``set``, ``CacheCodec.deserialize`` after ``get``
+when you need a typed ``BaseModel``.
+
+``dataclasses`` are not supported: only ``dict`` and Pydantic ``BaseModel`` inputs
+are encoded; pass a Pydantic model or convert with e.g. ``dataclasses.asdict`` first.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Type, TypeVar
+
+from pydantic import BaseModel, ValidationError
+
+from litellm._logging import verbose_proxy_logger
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class CacheCodec:
+    """
+    Encode/decode Pydantic models for DualCache (memory vs Redis safe payloads).
+
+    Dataclasses are not supported yet (only ``dict`` and ``BaseModel``).
+
+    Use ``serialize`` with ``model_type`` when writing so the same schema is used
+    as on read (``deserialize``). Pass ``model_type`` whenever you know it
+    (validates ``dict`` payloads and normalizes ``BaseModel`` instances).
+    """
+
+    @staticmethod
+    def serialize(value: Any, model_type: Optional[Type[T]] = None) -> Any:
+        """
+        Encode a value for DualCache / Redis (``json.dumps``-safe).
+
+        If ``model_type`` is set, the payload is validated with that model, then
+        ``model_dump(mode="json", exclude_none=True)`` — symmetric with ``deserialize``.
+
+        If the value is already an instance of ``model_type`` (or a subclass),
+        ``model_validate`` is skipped to avoid an unnecessary Pydantic copy — the
+        value is dumped directly.
+
+        If ``model_type`` is omitted, any ``BaseModel`` is dumped as above; other
+        values (e.g. plain ``dict``) are returned unchanged.
+        """
+        if model_type is not None:
+            if isinstance(value, model_type):
+                # Already the right type: dump directly, skip re-validation copy.
+                return value.model_dump(mode="json", exclude_none=True)
+            if isinstance(value, (dict, BaseModel)):
+                return model_type.model_validate(value).model_dump(
+                    mode="json", exclude_none=True
+                )
+            return value
+        if isinstance(value, BaseModel):
+            return value.model_dump(mode="json", exclude_none=True)
+        return value
+
+    @staticmethod
+    def deserialize(cached: Any, model_type: Type[T]) -> Optional[T]:
+        """
+        Decode a cache entry to ``model_type``.
+
+        - ``None`` → ``None``
+        - Already an instance of ``model_type`` (including subclasses) → returned as-is
+        - ``dict`` → ``model_type.model_validate(...)``; on ``ValidationError``,
+          logs a warning and returns ``None`` (treat as cache miss; avoids serving
+          malformed or schema-drifted entries)
+        - Any other type → ``None`` (caller should treat as cache miss or log)
+        """
+        if cached is None:
+            return None
+        if isinstance(cached, model_type):
+            return cached
+        if isinstance(cached, dict):
+            try:
+                return model_type.model_validate(cached)
+            except ValidationError as e:
+                verbose_proxy_logger.warning(
+                    "CacheCodec.deserialize: validation failed for %s (%s)",
+                    model_type.__name__,
+                    e,
+                )
+                return None
+        return None

litellm/proxy/common_utils/expired_ui_session_key_cleanup_manager.py

@@ -8,7 +8,7 @@ from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional
 
 from litellm._logging import verbose_proxy_logger
-from litellm.caching import DualCache
+from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
 from litellm.constants import (
     EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME,
     LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_BATCH_SIZE,
@@ -31,7 +31,7 @@ class ExpiredUISessionKeyCleanupManager:
     def __init__(
         self,
         prisma_client: PrismaClient,
-        user_api_key_cache: DualCache,
+        user_api_key_cache: UserApiKeyCache,
         pod_lock_manager=None,
     ):
         self.prisma_client = prisma_client