Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_fix_stale_member_search_results

# Conflicts:
#	ui/litellm-dashboard/src/components/organisms/create_key_button.tsx

basedpyright-code-budget.json

@@ -1,12 +1,12 @@
 {
   "reportAny": {
-    "limit": 22343
+    "limit": 19955
   },
   "reportArgumentType": {
-    "limit": 2578
+    "limit": 2566
   },
   "reportAssignmentType": {
-    "limit": 323
+    "limit": 320
   },
   "reportAttributeAccessIssue": {
     "limit": 488
@@ -24,7 +24,7 @@
     "limit": 19
   },
   "reportExplicitAny": {
-    "limit": 6991
+    "limit": 6049
   },
   "reportFunctionMemberAccess": {
     "limit": 7
@@ -54,10 +54,10 @@
     "limit": 0
   },
   "reportMissingParameterType": {
-    "limit": 5681
+    "limit": 5663
   },
   "reportMissingTypeArgument": {
-    "limit": 15605
+    "limit": 15557
   },
   "reportMissingTypeStubs": {
     "limit": 40
@@ -99,19 +99,19 @@
     "limit": 0
   },
   "reportUnknownArgumentType": {
-    "limit": 44709
+    "limit": 44655
   },
   "reportUnknownLambdaType": {
-    "limit": 112
+    "limit": 109
   },
   "reportUnknownMemberType": {
-    "limit": 39154
+    "limit": 39043
   },
   "reportUnknownParameterType": {
-    "limit": 19944
+    "limit": 19887
   },
   "reportUnknownVariableType": {
-    "limit": 30772
+    "limit": 30574
   },
   "reportUnnecessaryCast": {
     "limit": 117
@@ -123,7 +123,7 @@
     "limit": 5
   },
   "reportUnnecessaryIsInstance": {
-    "limit": 851
+    "limit": 836
   },
   "reportUntypedBaseClass": {
     "limit": 0

enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py

@@ -10,7 +10,8 @@ All /vector_store management endpoints
 
 import copy
 import json
-from typing import List, Optional
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Final, List, Optional, Protocol
 
 from fastapi import APIRouter, Depends, HTTPException
 
@@ -32,9 +33,35 @@ from litellm.types.vector_stores import (
 )
 from litellm.vector_stores.vector_store_registry import VectorStoreRegistry
 
+if TYPE_CHECKING:
+    from litellm.proxy.utils import PrismaClient
+
 router = APIRouter()
 
 
+class ManagedVectorStoreRow(Protocol):
+    """A ``litellm_managedvectorstorestable`` row as returned by Prisma."""
+
+    def model_dump(self) -> LiteLLM_ManagedVectorStore: ...
+
+
+class ManagedVectorStoreTable(Protocol):
+    """The Prisma actions namespace for ``litellm_managedvectorstorestable``."""
+
+    async def find_unique(self, where: Mapping[str, str | None]) -> ManagedVectorStoreRow | None: ...
+
+    async def create(self, data: Mapping[str, object]) -> ManagedVectorStoreRow: ...
+
+    async def delete(self, where: Mapping[str, str | None]) -> ManagedVectorStoreRow | None: ...
+
+    async def update(self, where: Mapping[str, str | None], data: Mapping[str, object]) -> ManagedVectorStoreRow: ...
+
+
+def managed_vector_store_table(prisma_client: "PrismaClient") -> ManagedVectorStoreTable:
+    """The Prisma table actions for managed vector stores, behind a typed surface."""
+    return prisma_client.db.litellm_managedvectorstorestable
+
+
 ########################################################
 # Management Endpoints
 ########################################################
@@ -66,7 +93,7 @@ async def new_vector_store(
     try:
         # Check if vector store already exists
         existing_vector_store = (
-            await prisma_client.db.litellm_managedvectorstorestable.find_unique(
+            await managed_vector_store_table(prisma_client).find_unique(
                 where={"vector_store_id": vector_store.get("vector_store_id")}
             )
         )
@@ -92,7 +119,7 @@ async def new_vector_store(
             del vector_store["litellm_params"]
 
         _new_vector_store = (
-            await prisma_client.db.litellm_managedvectorstorestable.create(
+            await managed_vector_store_table(prisma_client).create(
                 data={
                     **vector_store,
                     "litellm_params": litellm_params_json,
@@ -213,7 +240,7 @@ async def delete_vector_store(
     try:
         # Check if vector store exists
         existing_vector_store = (
-            await prisma_client.db.litellm_managedvectorstorestable.find_unique(
+            await managed_vector_store_table(prisma_client).find_unique(
                 where={"vector_store_id": data.vector_store_id}
             )
         )
@@ -224,7 +251,7 @@ async def delete_vector_store(
             )
 
         # Delete vector store
-        await prisma_client.db.litellm_managedvectorstorestable.delete(
+        await managed_vector_store_table(prisma_client).delete(
             where={"vector_store_id": data.vector_store_id}
         )
 
@@ -288,7 +315,7 @@ async def get_vector_store_info(
                 return {"vector_store": vector_store_pydantic_obj}
 
         vector_store = (
-            await prisma_client.db.litellm_managedvectorstorestable.find_unique(
+            await managed_vector_store_table(prisma_client).find_unique(
                 where={"vector_store_id": data.vector_store_id}
             )
         )
@@ -298,7 +325,7 @@ async def get_vector_store_info(
                 detail=f"Vector store with ID {data.vector_store_id} not found",
             )
 
-        vector_store_dict = vector_store.model_dump()  # type: ignore[attr-defined]
+        vector_store_dict = vector_store.model_dump()
         return {"vector_store": vector_store_dict}
     except Exception as e:
         verbose_proxy_logger.exception(f"Error getting vector store info: {str(e)}")
@@ -322,13 +349,13 @@ async def update_vector_store(
 
     try:
         update_data = data.model_dump(exclude_unset=True)
-        vector_store_id = update_data.pop("vector_store_id")
+        vector_store_id: Final[str] = update_data.pop("vector_store_id")
         if update_data.get("vector_store_metadata") is not None:
             update_data["vector_store_metadata"] = safe_dumps(
                 update_data["vector_store_metadata"]
             )
 
-        updated = await prisma_client.db.litellm_managedvectorstorestable.update(
+        updated = await managed_vector_store_table(prisma_client).update(
             where={"vector_store_id": vector_store_id},
             data=update_data,
         )

litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py

@@ -7,9 +7,10 @@ import hashlib
 import json
 import time
 from collections.abc import AsyncIterator
-from typing import Any, Final, NamedTuple, cast
+from typing import Any, Final, NamedTuple, Protocol
 
 import httpx
+from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 from litellm._logging import verbose_logger
 from litellm.a2a_protocol.providers.watsonx_orchestrate.transformation import (
@@ -38,11 +39,59 @@ class WXORequestParams(NamedTuple):
     thread_id: str | None
 
 
+class WXOLitellmParams(TypedDict, total=False):
+    """litellm_params keys read when routing an A2A request to watsonx Orchestrate."""
+
+    cp4d_host: ReadOnly[str]
+    instance_id: ReadOnly[str]
+    wxo_agent_id: ReadOnly[str]
+    api_key: ReadOnly[str]
+    username: ReadOnly[str | None]
+    auth_mode: ReadOnly[str]
+    thread_id: ReadOnly[str | None]
+
+
+class _IBMCloudTokenBody(TypedDict):
+    """Fields read from the IBM Cloud IAM token response."""
+
+    access_token: ReadOnly[str]
+    expires_in: ReadOnly[NotRequired[int]]
+
+
+class _CP4DTokenBody(TypedDict):
+    """Fields read from the CP4D authorize response."""
+
+    token: ReadOnly[str]
+    expiration: ReadOnly[NotRequired[float]]
+
+
+class _WXORun(TypedDict, total=False):
+    """Fields the handler reads from a WXO run object or run event."""
+
+    status: ReadOnly[str]
+    run_id: ReadOnly[str]
+    id: ReadOnly[str]
+
+
+class _SSELineSource(Protocol):
+    def aiter_lines(self) -> AsyncIterator[str]: ...
+
+
+class _WXOView(TypedDict, total=False):
+    """Typed reads of otherwise untyped watsonx Orchestrate and httpx values."""
+
+    ibm_cloud_token: ReadOnly[_IBMCloudTokenBody]
+    cp4d_token: ReadOnly[_CP4DTokenBody]
+    run: ReadOnly[_WXORun]
+    content_type: ReadOnly[str]
+    sse_source: ReadOnly[_SSELineSource]
+
+
 class WatsonxOrchestrateHandler:
     @staticmethod
     def _http_client(timeout: float = 90.0) -> AsyncHTTPHandler:
         return get_async_httpx_client(
-            llm_provider=cast(Any, httpxSpecialProvider.A2AProvider),
+            llm_provider=httpxSpecialProvider.A2AProvider,
             params={"timeout": timeout},
         )
 
@@ -57,7 +106,7 @@ class WatsonxOrchestrateHandler:
         return hashlib.sha256(material.encode()).hexdigest()
 
     @staticmethod
-    def _cp4d_token_ttl_seconds(expiration: Any, now_wall: float | None = None) -> int:
+    def _cp4d_token_ttl_seconds(expiration: float, now_wall: float | None = None) -> int:
         # CP4D returns expiration as absolute Unix epoch seconds, not a duration.
         expires_at: Final = int(expiration)
         wall: Final = now_wall if now_wall is not None else time.time()
@@ -90,9 +139,9 @@ class WatsonxOrchestrateHandler:
                 headers={"Content-Type": "application/x-www-form-urlencoded"},
             )
             response.raise_for_status()
-            payload = response.json()
-            token = str(payload["access_token"])
-            ttl_s = int(payload.get("expires_in", 3600))
+            iam_payload: Final[_WXOView] = {"ibm_cloud_token": response.json()}
+            token = str(iam_payload["ibm_cloud_token"]["access_token"])
+            ttl_s = int(iam_payload["ibm_cloud_token"].get("expires_in", 3600))
         else:
             if not username:
                 raise ValueError("'username' is required in litellm_params when auth_mode='cp4d'")
@@ -103,9 +152,9 @@ class WatsonxOrchestrateHandler:
                 headers={"Content-Type": "application/json"},
             )
             response.raise_for_status()
-            payload = response.json()
-            token = str(payload["token"])
-            expiration: Final = payload.get("expiration")
+            cp4d_payload: Final[_WXOView] = {"cp4d_token": response.json()}
+            token = str(cp4d_payload["cp4d_token"]["token"])
+            expiration: Final = cp4d_payload["cp4d_token"].get("expiration")
             if expiration is None:
                 ttl_s = 3600
             else:
@@ -118,6 +167,16 @@ class WatsonxOrchestrateHandler:
                 del _token_cache[stale_key]
         return token
 
+    @staticmethod
+    def _run_body(response: httpx.Response) -> _WXORun:
+        view: Final[_WXOView] = {"run": response.json()}
+        return view["run"]
+
+    @staticmethod
+    def _decode_run_event(payload: str | bytes) -> _WXORun:
+        view: Final[_WXOView] = {"run": json.loads(payload)}
+        return view["run"]
+
     @staticmethod
     async def _poll_run(
         base_url: str,
@@ -126,14 +185,14 @@ class WatsonxOrchestrateHandler:
         client: AsyncHTTPHandler,
         max_attempts: int = _MAX_POLL_ATTEMPTS,
         interval_s: float = _POLL_INTERVAL_S,
-    ) -> dict[str, Any]:
+    ) -> _WXORun:
         url: Final = f"{base_url}/v1/orchestrate/runs/{run_id}"
 
         for attempt in range(max_attempts):
             await asyncio.sleep(interval_s)
             response = await client.get(url, headers=auth_headers)
             response.raise_for_status()
-            result: dict[str, Any] = response.json()
+            result = WatsonxOrchestrateHandler._run_body(response)
             status = result.get("status", "")
             verbose_logger.debug("WXO: Poll %s/%s run='%s' status='%s'", attempt + 1, max_attempts, run_id, status)
             if status in WatsonxOrchestrateTransformation.TERMINAL_STATES:
@@ -145,11 +204,11 @@ class WatsonxOrchestrateHandler:
 
     @staticmethod
     async def _get_successful_run_data(
-        run_data: dict[str, Any],
+        run_data: _WXORun,
         base_url: str,
         auth_headers: dict[str, str],
         client: AsyncHTTPHandler,
-    ) -> dict[str, Any]:
+    ) -> _WXORun:
         status = run_data.get("status", "")
         if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES:
             run_id: Final = run_data.get("run_id") or run_data.get("id") or ""
@@ -170,15 +229,16 @@ class WatsonxOrchestrateHandler:
 
     @staticmethod
     async def _accumulate_wxo_sse_text(response: Any) -> str:
+        source: Final[_WXOView] = {"sse_source": response}
         accumulated_text = ""
-        async for line in response.aiter_lines():
+        async for line in source["sse_source"].aiter_lines():
             if not line.startswith("data:"):
                 continue
             data_str = line[5:].strip()
             if not data_str or data_str == "[DONE]":
                 continue
             try:
-                event = json.loads(data_str)
+                event = WatsonxOrchestrateHandler._decode_run_event(data_str)
             except json.JSONDecodeError:
                 continue
             chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result(event)
@@ -187,7 +247,7 @@ class WatsonxOrchestrateHandler:
         return accumulated_text
 
     @staticmethod
-    def _extract_litellm_params(litellm_params: dict[str, Any]) -> WXORequestParams:
+    def _extract_litellm_params(litellm_params: WXOLitellmParams) -> WXORequestParams:
         cp4d_host: Final = litellm_params.get("cp4d_host") or ""
         instance_id: Final = litellm_params.get("instance_id") or ""
         wxo_agent_id: Final = litellm_params.get("wxo_agent_id") or ""
@@ -215,9 +275,9 @@ class WatsonxOrchestrateHandler:
     @staticmethod
     async def handle_non_streaming(
         request_id: str,
-        params: dict[str, Any],
-        litellm_params: dict[str, Any],
-    ) -> dict[str, Any]:
+        params: dict[str, object],
+        litellm_params: WXOLitellmParams,
+    ) -> dict[str, object]:
         wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
 
         client: Final = WatsonxOrchestrateHandler._http_client(timeout=90.0)
@@ -246,7 +306,8 @@ class WatsonxOrchestrateHandler:
             headers=auth_headers,
         )
         run_response.raise_for_status()
-        run_data: dict[str, Any] = run_response.json()
+        started: Final[_WXOView] = {"run": run_response.json()}
+        run_data: _WXORun = started["run"]
 
         run_data = await WatsonxOrchestrateHandler._get_successful_run_data(
             run_data=run_data,
@@ -261,11 +322,11 @@ class WatsonxOrchestrateHandler:
     @staticmethod
     async def handle_streaming(
         request_id: str,
-        params: dict[str, Any],
-        litellm_params: dict[str, Any],
+        params: dict[str, object],
+        litellm_params: WXOLitellmParams,
         chunk_size: int = 50,
         delay_ms: int = 10,
-    ) -> AsyncIterator[dict[str, Any]]:
+    ) -> AsyncIterator[dict[str, object]]:
         wxo: Final = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params)
 
         client: Final = WatsonxOrchestrateHandler._http_client(timeout=120.0)
@@ -316,10 +377,11 @@ class WatsonxOrchestrateHandler:
                 yield chunk
             return
 
-        content_type: Final = response.headers.get("content-type", "").lower()
+        header_view: Final[_WXOView] = {"content_type": response.headers.get("content-type", "")}
+        content_type: Final = header_view["content_type"].lower()
         if "text/event-stream" not in content_type:
             response_body: Final = await response.aread()
-            result = json.loads(response_body)
+            result = WatsonxOrchestrateHandler._decode_run_event(response_body)
             result = await WatsonxOrchestrateHandler._get_successful_run_data(
                 run_data=result,
                 base_url=base_url,

litellm/caching/caching_handler.py

@@ -18,7 +18,7 @@ import asyncio
 import datetime
 import inspect
 import time
-from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator
+from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Mapping
 from typing import TYPE_CHECKING, Any, Final, Optional, TypeVar
 
 from pydantic import BaseModel
@@ -106,7 +106,7 @@ def _is_chat_completion_cached_dict(cached_result: dict) -> bool:
     return "choices" in cached_result
 
 
-def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, Any]) -> bool:
+def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, object]) -> bool:
     """
     When stream=True, do not run success callbacks at cache-hit time.
 
@@ -119,11 +119,21 @@ def _should_defer_streaming_cache_hit_callbacks(*, kwargs: dict[str, Any]) -> bo
     return kwargs.get("stream", False) is True
 
 
+def _prompt_tokens_details_as_mapping(details: "PromptTokensDetailsWrapper") -> Mapping[str, object]:
+    """Dump prompt token details to an opaque field mapping, tolerating non-pydantic stand-ins."""
+    return details.model_dump(exclude_none=True) if hasattr(details, "model_dump") else {}
+
+
+def _request_cache_key(request_kwargs: Mapping[str, Any]) -> str | None:
+    """Read the caller-supplied ``cache_key`` off the request kwargs."""
+    return request_kwargs.get("cache_key", None)
+
+
 class LLMCachingHandler:
     def __init__(
         self,
         original_function: Callable,
-        request_kwargs: dict[str, Any],
+        request_kwargs: dict[str, object],
         start_time: datetime.datetime,
     ):
         from litellm.caching import DualCache, RedisCache
@@ -150,7 +160,7 @@ class LLMCachingHandler:
         start_time: datetime.datetime,
         call_type: str,
         kwargs: dict[str, Any],
-        args: tuple[Any, ...] | None = None,
+        args: tuple[object, ...] | None = None,
     ) -> CachingHandlerResponse | None:
         """
         Internal method to get from the cache.
@@ -289,7 +299,7 @@ class LLMCachingHandler:
         start_time: datetime.datetime,
         call_type: str,
         kwargs: dict[str, Any],
-        args: tuple[Any, ...] | None = None,
+        args: tuple[object, ...] | None = None,
     ) -> CachingHandlerResponse:
         cached_result: Any | None = None
 
@@ -366,7 +376,7 @@ class LLMCachingHandler:
                     return CachingHandlerResponse(cached_result=cached_result)
         return CachingHandlerResponse(cached_result=cached_result)
 
-    def handle_kwargs_input_list_or_str(self, kwargs: dict[str, Any]) -> list[str]:
+    def handle_kwargs_input_list_or_str(self, kwargs: dict[str, object]) -> list[str]:
         """
         Handles the input of kwargs['input'] being a list or a string
         """
@@ -548,8 +558,8 @@ class LLMCachingHandler:
         if details2 is None:
             return details1
 
-        dict1: Final = details1.model_dump(exclude_none=True) if hasattr(details1, "model_dump") else {}
-        dict2: Final = details2.model_dump(exclude_none=True) if hasattr(details2, "model_dump") else {}
+        dict1: Final = _prompt_tokens_details_as_mapping(details1)
+        dict2: Final = _prompt_tokens_details_as_mapping(details2)
 
         merged: Final[dict] = {}
         for key in set(dict1.keys()) | set(dict2.keys()):
@@ -671,7 +681,9 @@ class LLMCachingHandler:
             cache_hit=cache_hit,
         )
 
-    async def _retrieve_from_cache(self, call_type: str, kwargs: dict[str, Any], args: tuple[Any, ...]) -> Any | None:
+    async def _retrieve_from_cache(
+        self, call_type: str, kwargs: dict[str, object], args: tuple[object, ...]
+    ) -> Any | None:
         """
         Internal method to
         - get cache key
@@ -727,7 +739,8 @@ class LLMCachingHandler:
                     cached_result = None
         else:
             request_kwargs: Final = new_kwargs.copy()
-            request_cache_key: Final = request_kwargs.pop("cache_key", None)
+            request_cache_key: Final = _request_cache_key(request_kwargs)
+            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)
@@ -749,10 +762,10 @@ class LLMCachingHandler:
         self,
         cached_result: Any,
         call_type: str,
-        kwargs: dict[str, Any],
+        kwargs: dict[str, object],
         logging_obj: LiteLLMLoggingObj,
         model: str,
-        args: tuple[Any, ...],
+        args: tuple[object, ...],
         custom_llm_provider: str | None = None,
     ) -> (
         ModelResponse
@@ -948,7 +961,7 @@ class LLMCachingHandler:
         result: Any,
         original_function: Callable,
         kwargs: dict[str, Any],
-        args: tuple[Any, ...] | None = None,
+        args: tuple[object, ...] | None = None,
     ):
         """
         Internal method to check the type of the result & cache used and adds the result to the cache accordingly
@@ -1013,8 +1026,8 @@ class LLMCachingHandler:
     def sync_set_cache(
         self,
         result: Any,
-        kwargs: dict[str, Any],
-        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, object],
+        args: tuple[object, ...] | None = None,
     ):
         """
         Sync internal method to add the result to the cache
@@ -1204,8 +1217,8 @@ class LLMCachingHandler:
 
 def convert_args_to_kwargs(
     original_function: Callable,
-    args: tuple[Any, ...] | None = None,
-) -> dict[str, Any]:
+    args: tuple[object, ...] | None = None,
+) -> dict[str, object]:
     # Get the signature of the original function
     signature: Final = inspect.signature(original_function)
 

litellm/cost_calculator.py

@@ -102,6 +102,7 @@ from litellm.types.utils import (
     LlmProviders,
     LlmProvidersSet,
     ModelInfo,
+    PromptTokensDetailsWrapper,
     ServiceTier,
     StandardBuiltInToolsParams,
     TranscriptionUsageDurationObject,
@@ -286,7 +287,7 @@ def _transcription_usage_has_token_details(
 
     prompt_tokens_val: Final = getattr(usage_block, "prompt_tokens", 0) or 0
     completion_tokens_val: Final = getattr(usage_block, "completion_tokens", 0) or 0
-    prompt_details: Final = getattr(usage_block, "prompt_tokens_details", None)
+    prompt_details: Final[PromptTokensDetailsWrapper | None] = getattr(usage_block, "prompt_tokens_details", None)
 
     if prompt_details is not None:
         audio_token_count: Final = getattr(prompt_details, "audio_tokens", 0) or 0
@@ -375,7 +376,7 @@ def cost_per_token(
     _is_anthropic_style = False
 
     if usage_object is not None:
-        _pt_details: Final = getattr(usage_object, "prompt_tokens_details", None)
+        _pt_details: Final[PromptTokensDetailsWrapper | None] = getattr(usage_object, "prompt_tokens_details", None)
         if _pt_details is not None:
             _cache_read_tokens = float(getattr(_pt_details, "cached_tokens", 0) or 0)
             # OpenAI-compatible providers report cache-write tokens under
@@ -385,8 +386,8 @@ def cost_per_token(
                 getattr(_pt_details, "cache_write_tokens", 0) or getattr(_pt_details, "cache_creation_tokens", 0) or 0
             )
 
-        _anthropic_read: Final = getattr(usage_object, "cache_read_input_tokens", None)
-        _anthropic_create: Final = getattr(usage_object, "cache_creation_input_tokens", None)
+        _anthropic_read: Final[int | None] = getattr(usage_object, "cache_read_input_tokens", None)
+        _anthropic_create: Final[int | None] = getattr(usage_object, "cache_creation_input_tokens", None)
         if _anthropic_read is not None or _anthropic_create is not None:
             _is_anthropic_style = True
             if _anthropic_read is not None:
@@ -703,7 +704,7 @@ def get_replicate_completion_pricing(completion_response: dict, total_time=0.0):
     return a100_80gb_price_per_second_public * total_time / 1000
 
 
-def has_hidden_params(obj: Any) -> bool:
+def has_hidden_params(obj: object) -> bool:
     return hasattr(obj, "_hidden_params")
 
 
@@ -728,7 +729,7 @@ def _get_provider_for_cost_calc(
 
 def _select_model_name_for_cost_calc(
     model: str | None,
-    completion_response: Any | None,
+    completion_response: object | None,
     base_model: str | None = None,
     custom_pricing: bool | None = None,
     custom_llm_provider: str | None = None,
@@ -804,7 +805,7 @@ def _model_contains_known_llm_provider(model: str) -> bool:
     return _provider_prefix in LlmProvidersSet
 
 
-def _get_response_model(completion_response: Any) -> str | None:
+def _get_response_model(completion_response: object) -> str | None:
     """
     Extract the model name from a completion response object.
 
@@ -866,8 +867,18 @@ def _normalize_service_tier(service_tier: object) -> str | None:
     return service_tier
 
 
+def _extract_service_tier(source: object) -> str | None:
+    """Read a raw ``service_tier`` off a response body or usage object, dict or pydantic model alike."""
+    if isinstance(source, BaseModel):
+        return getattr(source, "service_tier", None)
+    elif isinstance(source, dict):
+        return source.get("service_tier")
+
+    return None
+
+
 def _get_usage_object(
-    completion_response: Any,
+    completion_response: object,
 ) -> Usage | None:
     usage_obj: Final = cast(
         Usage | ResponseAPIUsage | dict | BaseModel,
@@ -1110,7 +1121,7 @@ def _store_cost_breakdown_in_logging_obj(
 
 
 def completion_cost(
-    completion_response=None,
+    completion_response: object | None = None,
     model: str | None = None,
     prompt="",
     messages: list = [],
@@ -1197,19 +1208,13 @@ def completion_cost(
 
         # Extract service_tier from completion_response if not provided
         if service_tier is None and completion_response is not None:
-            if isinstance(completion_response, BaseModel):
-                service_tier = getattr(completion_response, "service_tier", None)
-            elif isinstance(completion_response, dict):
-                service_tier = completion_response.get("service_tier")
+            service_tier = _extract_service_tier(completion_response)
 
         service_tier = _normalize_service_tier(service_tier)
 
         # Extract service_tier from usage object if not provided
         if service_tier is None and cost_per_token_usage_object is not None:
-            if isinstance(cost_per_token_usage_object, BaseModel):
-                service_tier = getattr(cost_per_token_usage_object, "service_tier", None)
-            elif isinstance(cost_per_token_usage_object, dict):
-                service_tier = cost_per_token_usage_object.get("service_tier")
+            service_tier = _extract_service_tier(cost_per_token_usage_object)
 
         service_tier = _normalize_service_tier(service_tier)
 
@@ -1412,7 +1417,7 @@ def completion_cost(
                     if completion_response is not None and isinstance(completion_response, RerankResponse):
                         meta_obj = completion_response.meta
                         if meta_obj is not None:
-                            billed_units = meta_obj.get("billed_units", {}) or {}
+                            billed_units: RerankBilledUnits = meta_obj.get("billed_units") or {}
                         else:
                             billed_units = {}
 
@@ -1801,7 +1806,7 @@ def response_cost_calculator(
 def ocr_cost(
     model: str,
     custom_llm_provider: str | None,
-    response: Any | None = None,
+    response: object | None = None,
 ) -> tuple[float, float]:
     """
     Args:

litellm/google_genai/adapters/transformation.py

@@ -1,5 +1,5 @@
 import json
-from collections.abc import AsyncIterator, Iterator
+from collections.abc import AsyncIterator, Iterator, Sequence
 from typing import Any, Final, TypedDict, cast
 
 from typing_extensions import ReadOnly
@@ -27,6 +27,7 @@ from litellm.types.utils import (
     ModelResponse,
     ModelResponseStream,
     StreamingChoices,
+    Usage,
 )
 
 
@@ -43,6 +44,29 @@ class _GenAIPart(TypedDict, total=False):
     functionCall: ReadOnly[dict[str, object]]
 
 
+class _GenAIFunctionDeclaration(TypedDict, total=False):
+    name: ReadOnly[str]
+    description: ReadOnly[str]
+    parametersJsonSchema: ReadOnly[dict[str, object]]
+
+
+class _GenAITool(TypedDict, total=False):
+    functionDeclarations: ReadOnly[list[_GenAIFunctionDeclaration]]
+
+
+class _GenAIFunctionCallingConfig(TypedDict, total=False):
+    mode: ReadOnly[str]
+
+
+class _GenAIToolConfig(TypedDict, total=False):
+    functionCallingConfig: ReadOnly[_GenAIFunctionCallingConfig]
+
+
+def _decode_tool_call_arguments(raw_arguments: str) -> object:
+    """Decode a tool call's JSON-encoded arguments into the value Google GenAI expects."""
+    return json.loads(raw_arguments)
+
+
 class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
     """
     Wrapper for streaming Google GenAI generate_content responses.
@@ -51,7 +75,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
 
     sent_first_chunk: bool = False
     # State tracking for accumulating partial tool calls
-    accumulated_tool_calls: dict[str, dict[str, str]]
+    accumulated_tool_calls: dict[int, dict[str, str]]
 
     def __init__(self, completion_stream: object):
         self.sent_first_chunk = False
@@ -108,7 +132,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
                         try:
                             # For tool calls with no arguments, accumulated_args will be "", which is not valid JSON.
                             # We default to an empty JSON object in this case.
-                            parsed_args = json.loads(tool_call_data["arguments"] or "{}")
+                            parsed_args = _decode_tool_call_arguments(tool_call_data["arguments"] or "{}")
                             function_call_part: _GenAIPart = {
                                 "functionCall": {
                                     "name": tool_call_data["name"] or "undefined_tool_name",
@@ -319,7 +343,7 @@ class GoogleGenAIAdapter:
 
     def _transform_google_genai_tools_to_openai(
         self,
-        tools: list[dict[str, Any]],
+        tools: Sequence[_GenAITool],
     ) -> list[ChatCompletionToolParam]:
         """Transform Google GenAI tools to OpenAI tools format"""
         openai_tools: Final[list[dict[str, object]]] = []
@@ -346,7 +370,7 @@ class GoogleGenAIAdapter:
 
     def _transform_google_genai_tool_config_to_openai(
         self,
-        tool_config: dict[str, Any],
+        tool_config: _GenAIToolConfig,
     ) -> ChatCompletionToolChoiceValues | None:
         """Transform Google GenAI tool_config to OpenAI tool_choice"""
         function_calling_config: Final = tool_config.get("functionCallingConfig", {})
@@ -563,7 +587,7 @@ class GoogleGenAIAdapter:
                 parts = self._transform_openai_delta_to_google_genai_parts_with_accumulation(choice.delta, wrapper)
             else:
                 parts = []
-            finish_reason = getattr(choice, "finish_reason", None)
+            finish_reason: str | None = getattr(choice, "finish_reason", None)
         else:
             # Fallback for generic choice objects
             message_content: Final = getattr(choice, "delta", {}).get("content", "")
@@ -625,7 +649,11 @@ class GoogleGenAIAdapter:
             for tool_call in message.tool_calls:
                 if hasattr(tool_call, "function") and tool_call.function:
                     try:
-                        args = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {}
+                        args = (
+                            _decode_tool_call_arguments(tool_call.function.arguments)
+                            if tool_call.function.arguments
+                            else {}
+                        )
                     except json.JSONDecodeError:
                         args = {}
 
@@ -661,7 +689,7 @@ class GoogleGenAIAdapter:
                 continue
 
             # 3. Use `index` as the primary key for accumulation
-            tool_call_index = getattr(tool_call, "index", None)
+            tool_call_index: int | None = getattr(tool_call, "index", None)
             if tool_call_index is None:
                 continue  # Index is essential for tracking streaming tool calls
 
@@ -695,7 +723,7 @@ class GoogleGenAIAdapter:
             # 5. Attempt to parse arguments even if name hasn't arrived.
             try:
                 # Attempt to parse the accumulated arguments string
-                parsed_args = json.loads(accumulated_args)
+                parsed_args = _decode_tool_call_arguments(accumulated_args)
 
                 # If parsing succeeds, but we don't have a name yet, wait.
                 # The part will be created by a later chunk that brings the name.
@@ -729,7 +757,7 @@ class GoogleGenAIAdapter:
 
         return mapping.get(finish_reason, "STOP")
 
-    def _map_usage(self, usage: Any) -> dict[str, int]:
+    def _map_usage(self, usage: Usage | None) -> dict[str, int]:
         """Map OpenAI usage to Google GenAI usage format"""
         return {
             "promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0,

litellm/integrations/galileo.py

@@ -60,13 +60,13 @@ class LLMResponse(BaseModel):
         default=None,
         description="Total cost of the LLM call in USD as computed by LiteLLM.",
     )
-    output_logprobs: dict[str, Any] | None = Field(
+    output_logprobs: dict[str, object] | None = Field(
         default=None,
         description="Optional. When available, logprobs are used to compute Uncertainty.",
     )
     created_at: str = Field(..., description='timestamp constructed in "%Y-%m-%dT%H:%M:%S" format')
     tags: list[str] | None = None
-    user_metadata: dict[str, Any] | None = None
+    user_metadata: dict[str, object] | None = None
 
 
 class GalileoObserve(CustomLogger):
@@ -238,13 +238,13 @@ class GalileoObserve(CustomLogger):
         return created_at
 
     @staticmethod
-    def _token_metrics_from_record(record: Mapping[str, Any]) -> dict[str, Any]:
+    def _token_metrics_from_record(record: Mapping[str, Any]) -> dict[str, object]:
         num_input_tokens: Final = int(record.get("num_input_tokens") or 0)
         num_output_tokens: Final = int(record.get("num_output_tokens") or 0)
         num_total_tokens = int(record.get("num_total_tokens") or 0)
         if num_total_tokens == 0 and (num_input_tokens or num_output_tokens):
             num_total_tokens = num_input_tokens + num_output_tokens
-        metrics: Final[dict[str, Any]] = {
+        metrics: Final[dict[str, object]] = {
             "num_input_tokens": num_input_tokens,
             "num_output_tokens": num_output_tokens,
             "num_total_tokens": num_total_tokens,
@@ -260,10 +260,10 @@ class GalileoObserve(CustomLogger):
         *,
         trace_id: str,
         span_id: str,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         created_at: Final = GalileoObserve._normalize_created_at(record.get("created_at", ""))
 
-        span: Final[dict[str, Any]] = {
+        span: Final[dict[str, object]] = {
             "type": "llm",
             "id": span_id,
             "trace_id": trace_id,
@@ -287,7 +287,7 @@ class GalileoObserve(CustomLogger):
         return span
 
     @staticmethod
-    def _record_to_v2_trace(record: Mapping[str, Any]) -> dict[str, Any]:
+    def _record_to_v2_trace(record: Mapping[str, Any]) -> dict[str, object]:
         trace_id: Final = str(uuid.uuid4())
         span_id: Final = str(uuid.uuid4())
         created_at: Final = GalileoObserve._normalize_created_at(record.get("created_at", ""))
@@ -307,8 +307,8 @@ class GalileoObserve(CustomLogger):
             "spans": [GalileoObserve._record_to_v2_span(record, trace_id=trace_id, span_id=span_id)],
         }
 
-    def _build_traces_payload(self, records: Sequence[Mapping[str, Any]]) -> dict[str, Any]:
-        payload: Final[dict[str, Any]] = {
+    def _build_traces_payload(self, records: Sequence[Mapping[str, object]]) -> dict[str, object]:
+        payload: Final[dict[str, object]] = {
             "traces": [self._record_to_v2_trace(record) for record in records],
             "logging_method": "api_direct",
             "reliable": False,
@@ -318,7 +318,7 @@ class GalileoObserve(CustomLogger):
             payload["log_stream_id"] = self.log_stream_id
         return payload
 
-    def _get_ingest_request(self) -> tuple[str, dict[str, Any]] | None:
+    def _get_ingest_request(self) -> tuple[str, dict[str, object]] | None:
         if not self.base_url or not self.project_id:
             return None
 
@@ -427,9 +427,9 @@ class GalileoObserve(CustomLogger):
             pass
 
     @staticmethod
-    def _build_prompt(kwargs: Mapping[str, Any]) -> dict[str, Any]:
+    def _build_prompt(kwargs: Mapping[str, Any]) -> dict[str, object]:
         optional_params: Final[Mapping[str, object]] = kwargs.get("optional_params", {}) or {}
-        prompt: Final[dict[str, Any]] = {"messages": kwargs.get("messages")}
+        prompt: Final[dict[str, object]] = {"messages": kwargs.get("messages")}
         if optional_params.get("functions") is not None:
             prompt["functions"] = optional_params["functions"]
         if optional_params.get("tools") is not None:
@@ -451,7 +451,7 @@ class GalileoObserve(CustomLogger):
         return json.dumps(value, default=_json_default)
 
     @staticmethod
-    def _prompt_to_input_text(prompt: Mapping[str, Any]) -> str:
+    def _prompt_to_input_text(prompt: Mapping[str, object]) -> str:
         messages: Final[object] = prompt.get("messages")
         if messages is not None:
             text: Final = GalileoObserve._input_text_from_messages(messages)
@@ -464,7 +464,7 @@ class GalileoObserve(CustomLogger):
         if response_obj.choices and len(response_obj.choices) > 0:
             message: Final = response_obj["choices"][0]["message"]
             if hasattr(message, "json"):
-                message_json: Final = message.json()
+                message_json: Final[object] = message.json()
                 if isinstance(message_json, str):
                     return json.loads(message_json)
                 return message_json
@@ -488,7 +488,7 @@ class GalileoObserve(CustomLogger):
         return None
 
     @staticmethod
-    def _langfuse_style_rerank_prompt(kwargs: Mapping[str, object]) -> dict[str, Any]:
+    def _langfuse_style_rerank_prompt(kwargs: Mapping[str, object]) -> dict[str, object]:
         """Match Langfuse rerank input: prompt = {"messages": kwargs.get("messages")}."""
         return {"messages": kwargs.get("messages")}
 

litellm/integrations/langfuse/langfuse.py

@@ -5,7 +5,7 @@ import traceback
 from collections.abc import Callable, Iterable, Mapping
 from datetime import datetime
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast
 
 from packaging.version import Version
 
@@ -49,10 +49,21 @@ else:
 
 
 _DENIED_STEERING_KEYS: Final = frozenset({"headers", "endpoint", "caching_groups", "previous_models"})
-_NO_METADATA: Final[Mapping[str, Any]] = MappingProxyType({})
+_NO_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
 _REDACTED_PROXY_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "cookie", "referer"})
 
 
+def _object_mapping(value: object) -> Mapping[str, object] | None:
+    """Return ``value`` as an opaque mapping when it is a dict."""
+    return value if isinstance(value, dict) else None
+
+
+class _UsageObject(Protocol):
+    """Token-count surface the Langfuse logger reads off a response usage payload."""
+
+    def get(self, key: Literal["cache_creation_input_tokens", "cache_read_input_tokens"], /) -> int | None: ...
+
+
 def _extract_cache_read_input_tokens(usage_obj) -> int:
     """
     Extract cache_read_input_tokens from usage object.
@@ -82,6 +93,11 @@ def _extract_cache_read_input_tokens(usage_obj) -> int:
     return cache_read_input_tokens
 
 
+def _logging_id(start_time: datetime | None, response_obj: object) -> str | None:
+    """Typed view of the timestamped response id Langfuse uses as the generation id."""
+    return litellm.utils.get_logging_id(start_time, response_obj)
+
+
 def _as_steering_flag(value: object) -> bool:
     """A string ``str_to_bool`` does not recognise falls back to its truthiness."""
     if isinstance(value, str):
@@ -222,7 +238,7 @@ class LangFuseLogger:
         return langfuse_client
 
     @staticmethod
-    def add_metadata_from_header(litellm_params: dict, metadata: dict) -> dict:
+    def add_metadata_from_header(litellm_params: dict, metadata: dict) -> dict[str, object]:
         """
         Adds metadata from proxy request headers to Langfuse logging if keys start with "langfuse_"
         and overwrites litellm_params.metadata if already included.
@@ -494,7 +510,7 @@ class LangFuseLogger:
     def _log_langfuse_v2(
         self,
         user_id: str | None,
-        metadata: dict,
+        metadata: dict[str, object],
         litellm_params: dict,
         output: str | dict | list | None,
         start_time: datetime | None,
@@ -519,7 +535,7 @@ class LangFuseLogger:
                 else []
             )
 
-            allowlisted_metadata: Final[StandardLoggingMetadata | dict[str, Any]] = (
+            allowlisted_metadata: Final[StandardLoggingMetadata | Mapping[str, object]] = (
                 standard_logging_object["metadata"] if standard_logging_object is not None else _NO_METADATA
             )
             end_user_id: Final = allowlisted_metadata.get("user_api_key_end_user_id", None)
@@ -531,11 +547,12 @@ class LangFuseLogger:
             # Clean Metadata before logging - never log raw metadata
             # the raw metadata can contain circular references which leads to infinite recursion
             # we clean out all extra litellm metadata params before logging
-            clean_metadata: dict[str, Any] = {}
+            clean_metadata: dict[str, object] = {}
             if prompt_management_metadata is not None:
                 clean_metadata["prompt_management_metadata"] = prompt_management_metadata
-            if isinstance(metadata, dict):
-                for key, value in metadata.items():
+            metadata_entries: Final = _object_mapping(metadata)
+            if metadata_entries is not None:
+                for key, value in metadata_entries.items():
                     # generate langfuse tags - Default Tags sent to Langfuse from LiteLLM Proxy
                     if (
                         litellm.langfuse_default_tags is not None
@@ -705,8 +722,8 @@ class LangFuseLogger:
             usage_details = None
             if response_obj is not None:
                 if hasattr(response_obj, "id") and response_obj.get("id", None) is not None:
-                    generation_id = litellm.utils.get_logging_id(start_time, response_obj)
-                _usage_obj: Final = getattr(response_obj, "usage", None)
+                    generation_id = _logging_id(start_time, response_obj)
+                _usage_obj: Final[_UsageObject | None] = getattr(response_obj, "usage", None)
 
                 if _usage_obj:
                     # Safely get usage values, defaulting None to 0 for Langfuse compatibility.
@@ -811,7 +828,7 @@ class LangFuseLogger:
     @staticmethod
     def _get_chat_content_for_langfuse(
         response_obj: ModelResponse,
-    ):
+    ) -> str | None:
         """
         Get the chat content for Langfuse logging
         """
@@ -1078,7 +1095,7 @@ def log_provider_specific_information_as_span(
         None
     """
 
-    _hidden_params: Final = clean_metadata.get("hidden_params", None)
+    _hidden_params: Final[Mapping[str, object] | None] = clean_metadata.get("hidden_params", None)
     if _hidden_params is None:
         return
 

litellm/integrations/otel/logger.py

@@ -62,8 +62,13 @@ from litellm.integrations.otel.plumbing.providers import (
 from litellm.integrations.otel.plumbing.routing import TenantTracerCache
 
 if TYPE_CHECKING:
+    from opentelemetry.metrics import MeterProvider
+
+    from litellm.caching.dual_cache import DualCache
     from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.types.services import ServiceLoggerPayload
     from litellm.types.utils import (
+        CallTypesLiteral,
         StandardLoggingGuardrailInformation,
         StandardLoggingPayload,
     )
@@ -140,7 +145,7 @@ class OpenTelemetryV2(CustomLogger):
         callback_name: str | None = None,
         tracer_provider: TracerProvider | None = None,
         logger_provider: LoggerProvider | None = None,
-        meter_provider: Any | None = None,
+        meter_provider: "MeterProvider | None" = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
@@ -162,7 +167,7 @@ class OpenTelemetryV2(CustomLogger):
         self._open_llm_calls: OrderedDict[str, _LLMCallSpan] = OrderedDict()
         self._init_otel_logger_on_litellm_proxy()
 
-    def _init_metrics(self, meter_provider: Any | None) -> "GenAIMetricRecorder | None":
+    def _init_metrics(self, meter_provider: "MeterProvider | None") -> "GenAIMetricRecorder | None":
         """Create the six GenAI histograms when metrics are enabled, else ``None``.
 
         ``meter_provider`` is an explicit override (tests inject one); otherwise the
@@ -340,7 +345,7 @@ class OpenTelemetryV2(CustomLogger):
 
     def _emit_mcp_tool_call(
         self,
-        kwargs: Mapping[str, Any],
+        kwargs: Mapping[str, object],
         start_time: datetime | float | None,
         end_time: datetime | float | None,
     ) -> bool:
@@ -417,7 +422,7 @@ class OpenTelemetryV2(CustomLogger):
 
     def _close_llm_call(
         self,
-        kwargs: Mapping[str, Any],
+        kwargs: Mapping[str, object],
         start_time: datetime | float | None,
         end_time: datetime | float | None,
     ) -> Span | None:
@@ -474,7 +479,7 @@ class OpenTelemetryV2(CustomLogger):
 
     async def async_service_success_hook(
         self,
-        payload: Any,
+        payload: "ServiceLoggerPayload",
         parent_otel_span: Span | None = None,
         start_time: datetime | float | None = None,
         end_time: datetime | float | None = None,
@@ -491,7 +496,7 @@ class OpenTelemetryV2(CustomLogger):
 
     async def async_service_failure_hook(
         self,
-        payload: Any,
+        payload: "ServiceLoggerPayload",
         error: str | None = "",
         parent_otel_span: Span | None = None,
         start_time: datetime | float | None = None,
@@ -509,7 +514,7 @@ class OpenTelemetryV2(CustomLogger):
 
     def _emit_service(
         self,
-        payload: Any,
+        payload: "ServiceLoggerPayload",
         *,
         parent_otel_span: Span | None,
         start_time: datetime | float | None,
@@ -559,7 +564,7 @@ class OpenTelemetryV2(CustomLogger):
     #  / errors are the FastAPI instrumentor's job, so we don't touch it here.
     # ====================================================================== #
 
-    def seed_request_identity(self, user_api_key_dict: Any, model: Any = None) -> None:
+    def seed_request_identity(self, user_api_key_dict: object, model: str | None = None) -> None:
         """Attach request-identity Baggage to the current context + server span.
 
         Seeding identity into Baggage makes **every** span emitted afterwards for
@@ -615,10 +620,10 @@ class OpenTelemetryV2(CustomLogger):
 
     async def async_pre_call_hook(
         self,
-        user_api_key_dict: Any,
-        cache: Any,
+        user_api_key_dict: "UserAPIKeyAuth",
+        cache: "DualCache",
         data: dict,
-        call_type: Any,
+        call_type: "CallTypesLiteral",
     ) -> dict:
         self.seed_request_identity(
             user_api_key_dict,
@@ -790,7 +795,7 @@ def emit_guardrail_span(entry: "StandardLoggingGuardrailInformation") -> None:
         pass
 
 
-def seed_request_identity(user_api_key_dict: Any, model: Any = None) -> None:
+def seed_request_identity(user_api_key_dict: object, model: str | None = None) -> None:
     logger: Final = _registered_v2_logger()
     if logger is not None:
         logger.seed_request_identity(user_api_key_dict, model=model)

litellm/integrations/prometheus.py

@@ -9,7 +9,9 @@ import os
 import sys
 from collections.abc import Awaitable, Callable, Mapping, Sequence
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Final, Literal, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast
+
+from pydantic import BaseModel
 
 import litellm
 from litellm._logging import print_verbose, verbose_logger
@@ -38,6 +40,7 @@ from litellm.proxy._types import (
     LiteLLM_UserTable,
     UserAPIKeyAuth,
 )
+from litellm.repositories.base_repository import BaseRepository
 from litellm.repositories.organization_repository import OrganizationRepository
 from litellm.repositories.team_repository import TeamRepository
 from litellm.repositories.user_repository import UserRepository
@@ -58,6 +61,9 @@ if TYPE_CHECKING:
 else:
     AsyncIOScheduler = Any
 
+_BudgetRowT: Final = TypeVar("_BudgetRowT")
+_TableRowT: Final = TypeVar("_TableRowT", bound=BaseModel)
+
 _DEFAULT_BUDGET_METRICS_PER_REQUEST_TIMEOUT: Final = 5.0
 
 _NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset(
@@ -73,6 +79,36 @@ _NON_ENUM_METRIC_LABELS: Final[frozenset[str]] = frozenset(
 )
 
 
+class _PaginatedPrismaTable(Protocol[_TableRowT]):
+    """The slice of a prisma table action surface used for budget-metric pagination."""
+
+    async def find_many(
+        self,
+        *,
+        skip: int,
+        take: int,
+        order: Mapping[str, str],
+        include: Mapping[str, bool] | None = None,
+    ) -> list[_TableRowT]: ...
+
+    async def count(self) -> int: ...
+
+
+def _paginated_table(repository: BaseRepository[_TableRowT]) -> _PaginatedPrismaTable[_TableRowT]:
+    """View a repository's prisma table through the pagination surface budget metrics need."""
+    return repository.table
+
+
+class _OrgBudgetRow(Protocol):
+    """The budget columns joined onto an organization row."""
+
+    @property
+    def max_budget(self) -> float | None: ...
+
+    @property
+    def budget_reset_at(self) -> datetime | None: ...
+
+
 class _ExcludedLabelMetric:
     """Proxies a prometheus metric whose declared ``labelnames`` had globally
     excluded labels removed, dropping those labels from every ``labels(...)``
@@ -1531,7 +1567,7 @@ class PrometheusLogger(CustomLogger):
 
         cache_creation_detail_tokens: Final = PrometheusLogger._resolve_cache_write_tokens(prompt_details)
 
-        detail_metrics: Final[list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]]] = [
+        detail_metrics: Final[list[tuple[Any, DEFINED_PROMETHEUS_METRICS, object]]] = [
             (
                 self.litellm_input_cached_tokens_metric,
                 "litellm_input_cached_tokens_metric",
@@ -1584,7 +1620,7 @@ class PrometheusLogger(CustomLogger):
         if not isinstance(usage_object, dict):
             return
 
-        media_metrics: Final[list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]]] = [
+        media_metrics: Final[list[tuple[Any, DEFINED_PROMETHEUS_METRICS, object]]] = [
             (
                 self.litellm_video_duration_seconds_metric,
                 "litellm_video_duration_seconds_metric",
@@ -1606,7 +1642,7 @@ class PrometheusLogger(CustomLogger):
 
     def _inc_sparse_usage_counters(
         self,
-        counters_with_values: list[tuple[Any, DEFINED_PROMETHEUS_METRICS, Any]],
+        counters_with_values: Sequence[tuple[Any, DEFINED_PROMETHEUS_METRICS, object]],
         enum_values: UserAPIKeyLabelValues,
         label_context: PrometheusLabelFactoryContext | None = None,
     ) -> None:
@@ -2133,7 +2169,7 @@ class PrometheusLogger(CustomLogger):
     def _extract_status_code(
         self,
         kwargs: dict | None = None,
-        enum_values: Any | None = None,
+        enum_values: UserAPIKeyLabelValues | None = None,
         exception: Exception | None = None,
     ) -> int | None:
         """
@@ -2151,7 +2187,7 @@ class PrometheusLogger(CustomLogger):
         Returns:
             Status code as integer if found, None otherwise
         """
-        status_code = None
+        status_code: int | None = None
 
         # Try from enum_values first (most common in our callbacks)
         if enum_values and hasattr(enum_values, "status_code") and enum_values.status_code:
@@ -2225,8 +2261,8 @@ class PrometheusLogger(CustomLogger):
     def _should_skip_metrics_for_invalid_key(
         self,
         kwargs: dict | None = None,
-        user_api_key_dict: Any | None = None,
-        enum_values: Any | None = None,
+        user_api_key_dict: UserAPIKeyAuth | None = None,
+        enum_values: UserAPIKeyLabelValues | None = None,
         standard_logging_payload: dict | StandardLoggingPayload | None = None,
         exception: Exception | None = None,
     ) -> bool:
@@ -2391,7 +2427,7 @@ class PrometheusLogger(CustomLogger):
         for all successful requests (both streaming and non-streaming).
         """
 
-    def _safe_get(self, obj: Any, key: str, default: Any = None) -> Any:
+    def _safe_get(self, obj: Any, key: str, default: object = None) -> Any:
         """Get value from dict or Pydantic model."""
         if obj is None:
             return default
@@ -3273,8 +3309,8 @@ class PrometheusLogger(CustomLogger):
 
     async def _initialize_budget_metrics(
         self,
-        data_fetch_function: Callable[..., Awaitable[tuple[list[Any], int | None]]],
-        set_metrics_function: Callable[[list[Any]], Awaitable[None]],
+        data_fetch_function: Callable[..., Awaitable[tuple[list[_BudgetRowT], int | None]]],
+        set_metrics_function: Callable[[list[_BudgetRowT]], Awaitable[None]],
         data_type: Literal["teams", "keys", "users", "orgs"],
     ):
         """
@@ -3393,12 +3429,12 @@ class PrometheusLogger(CustomLogger):
 
         async def fetch_users(page_size: int, page: int) -> tuple[list[LiteLLM_UserTable], int | None]:
             skip: Final = (page - 1) * page_size
-            users: Final = await UserRepository(prisma_client).table.find_many(
+            users: Final = await _paginated_table(UserRepository(prisma_client)).find_many(
                 skip=skip,
                 take=page_size,
                 order={"created_at": "desc"},
             )
-            total_count: Final = await UserRepository(prisma_client).table.count()
+            total_count: Final = await _paginated_table(UserRepository(prisma_client)).count()
             return users, total_count
 
         await self._initialize_budget_metrics(
@@ -3419,13 +3455,13 @@ class PrometheusLogger(CustomLogger):
 
         async def fetch_orgs(page_size: int, page: int) -> tuple[list, int | None]:
             skip: Final = (page - 1) * page_size
-            orgs: Final = await OrganizationRepository(prisma_client).table.find_many(
+            orgs: Final = await _paginated_table(OrganizationRepository(prisma_client)).find_many(
                 skip=skip,
                 take=page_size,
                 order={"created_at": "desc"},
                 include={"litellm_budget_table": True},
             )
-            total_count: Final = await OrganizationRepository(prisma_client).table.count()
+            total_count: Final = await _paginated_table(OrganizationRepository(prisma_client)).count()
             return orgs, total_count
 
         await self._initialize_budget_metrics(
@@ -3488,7 +3524,7 @@ class PrometheusLogger(CustomLogger):
 
         try:
             # Get total user count
-            total_users: Final = await UserRepository(prisma_client).table.count()
+            total_users: Final = await _paginated_table(UserRepository(prisma_client)).count()
             self.litellm_total_users_metric.set(total_users)
             verbose_logger.debug("Prometheus: set litellm_total_users to %s", total_users)
 
@@ -3497,13 +3533,13 @@ class PrometheusLogger(CustomLogger):
             verbose_logger.debug("Prometheus: set litellm_active_users to %s", billable_users)
 
             # Get total team count
-            total_teams: Final = await TeamRepository(prisma_client).table.count()
+            total_teams: Final = await _paginated_table(TeamRepository(prisma_client)).count()
             self.litellm_teams_count_metric.set(total_teams)
             verbose_logger.debug("Prometheus: set litellm_teams_count to %s", total_teams)
         except Exception as e:
             verbose_logger.exception("Error initializing user/team count metrics: %s", e)
 
-    async def _set_key_list_budget_metrics(self, keys: list[str | UserAPIKeyAuth]):
+    async def _set_key_list_budget_metrics(self, keys: list[str | UserAPIKeyAuth | LiteLLM_DeletedVerificationToken]):
         """Helper function to set budget metrics for a list of keys"""
         for key in keys:
             if isinstance(key, UserAPIKeyAuth):
@@ -3522,7 +3558,7 @@ class PrometheusLogger(CustomLogger):
     async def _set_org_list_budget_metrics(self, orgs: list):
         """Helper function to set budget metrics for a list of orgs"""
         for org in orgs:
-            budget_table = getattr(org, "litellm_budget_table", None)
+            budget_table: _OrgBudgetRow | None = getattr(org, "litellm_budget_table", None)
             self._set_org_budget_metrics(
                 org_id=org.organization_id or "",
                 org_alias=org.organization_alias or "",
@@ -4051,6 +4087,11 @@ class PrometheusLogger(CustomLogger):
         verbose_proxy_logger.debug("Starting Prometheus Metrics on /metrics (no authentication)")
 
 
+def _label_source(enum_values: UserAPIKeyLabelValues) -> Mapping[str, object]:
+    """Flatten the label values into the opaque name/value mapping the label filters read."""
+    return enum_values.model_dump()
+
+
 def _prometheus_labels_from_context(
     supported_enum_labels: list[str],
     ctx: PrometheusLabelFactoryContext,
@@ -4098,7 +4139,7 @@ def prometheus_label_factory(
         return _prometheus_labels_from_context(supported_enum_labels, label_context)
 
     # Extract dictionary from Pydantic object
-    enum_dict: Final = enum_values.model_dump()
+    enum_dict: Final = _label_source(enum_values)
 
     # Filter supported labels and sanitize values to prevent breaking
     # the Prometheus text format (e.g. U+2028 Line Separator in label values)
@@ -4154,7 +4195,7 @@ def get_custom_labels_from_metadata(metadata: dict) -> dict[str, str]:
 
         keys_parts = key.split(".")
         # Traverse through the dictionary using the parts
-        value: Any = metadata
+        value: object = metadata
         for part in keys_parts:
             if isinstance(value, dict):
                 value = value.get(part, None)  # Get the value, return None if not found
@@ -4171,7 +4212,7 @@ def get_custom_labels_from_metadata(metadata: dict) -> dict[str, str]:
 
 def _get_combined_custom_metadata_from_standard_logging_payload(
     standard_logging_payload: dict | None,
-) -> dict[str, Any]:
+) -> dict[str, object]:
     """
     Combine the metadata sources that can supply custom Prometheus labels.
 

litellm/integrations/websearch_interception/handler.py

@@ -12,6 +12,8 @@ import uuid
 from collections.abc import AsyncIterator, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
 
+from typing_extensions import ReadOnly
+
 import litellm
 from litellm._logging import verbose_logger
 from litellm.anthropic_interface import messages as anthropic_messages
@@ -90,6 +92,20 @@ class _SearchToolConfig(TypedDict, total=False):
     litellm_params: Mapping[str, object] | None
 
 
+class _DeploymentKwargsView(TypedDict):
+    """Typed reads of the untyped request kwargs seen by the deployment hook."""
+
+    custom_llm_provider: ReadOnly[str]
+    litellm_params: ReadOnly[Mapping[str, object]]
+    model: ReadOnly[str]
+
+
+class _UserAuthView(TypedDict):
+    """Typed read of the optional team attached to the caller's auth object."""
+
+    team_id: ReadOnly[str | None]
+
+
 class WebSearchInterceptionLogger(CustomLogger):
     """
     CustomLogger that intercepts WebSearch tool calls for models that don't
@@ -265,7 +281,9 @@ class WebSearchInterceptionLogger(CustomLogger):
         )
         return response
 
-    async def async_pre_call_deployment_hook(self, kwargs: dict[str, Any], call_type: CallTypes | None) -> dict | None:
+    async def async_pre_call_deployment_hook(
+        self, kwargs: dict[str, Any], call_type: CallTypes | None
+    ) -> dict[str, object] | None:
         """
         Pre-call hook to convert native Anthropic web_search tools to regular tools.
 
@@ -275,12 +293,17 @@ class WebSearchInterceptionLogger(CustomLogger):
         """
         # Check if this is for an enabled provider
         # Try top-level kwargs first, then nested litellm_params, then derive from model name
-        custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get(
+        kwargs_view: Final[_DeploymentKwargsView] = {
+            "custom_llm_provider": kwargs.get("custom_llm_provider", ""),
+            "litellm_params": kwargs.get("litellm_params", {}),
+            "model": kwargs.get("model", ""),
+        }
+        custom_llm_provider = kwargs_view["custom_llm_provider"] or kwargs_view["litellm_params"].get(
             "custom_llm_provider", ""
         )
         if not custom_llm_provider:
             try:
-                _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs.get("model", ""))
+                _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs_view["model"])
             except Exception:
                 custom_llm_provider = ""
         if custom_llm_provider not in self.enabled_providers:
@@ -1422,7 +1445,8 @@ class WebSearchInterceptionLogger(CustomLogger):
             valid_token=user_api_key_auth,
         )
 
-        team_id: Final = getattr(user_api_key_auth, "team_id", None)
+        auth_view: Final[_UserAuthView] = {"team_id": getattr(user_api_key_auth, "team_id", None)}
+        team_id: Final = auth_view["team_id"]
         if team_id:
             from litellm.proxy.proxy_server import (
                 prisma_client,

litellm/litellm_core_utils/realtime_streaming.py

@@ -1,7 +1,9 @@
 import asyncio
 import json
 from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Final, Protocol, cast
+from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
+
+from typing_extensions import ReadOnly
 
 import litellm
 from litellm._logging import verbose_logger
@@ -32,13 +34,52 @@ class _ClientWebSocketExceptions(Protocol):
     ConnectionClosed: type[Exception]
 
 
-class _ClientWebSocket(Protocol):
+class _ASGIScope(TypedDict, total=False):
+    """The part of an ASGI connection scope this module reads."""
+
+    headers: ReadOnly[Sequence[tuple[bytes | str, bytes | str]]]
+
+
+class _ClientEventItem(TypedDict, total=False):
+    """The ``item`` payload of a client ``conversation.item.create`` frame."""
+
+    type: ReadOnly[str]
+    role: ReadOnly[str]
+    output: ReadOnly[object]
+    content: ReadOnly[Sequence[object]]
+
+
+class _ClientEventFrame(TypedDict, total=False):
+    """The fields the proxy reads from a client realtime frame."""
+
+    type: ReadOnly[str]
+    item: ReadOnly[_ClientEventItem]
+    session: ReadOnly[Mapping[str, object]]
+
+
+class _ResponseDoneBody(TypedDict, total=False):
+    """The ``response`` body of a ``response.done`` event, as read for spend logging."""
+
+    output: ReadOnly[Sequence[Mapping[str, object]]]
+
+
+class _ScopedWebSocket(Protocol):
+    @property
+    def scope(self) -> _ASGIScope: ...
+
+
+class _ClientWebSocket(_ScopedWebSocket, Protocol):
     exceptions: _ClientWebSocketExceptions
 
     async def send_text(self, data: str) -> None: ...
     async def receive_text(self) -> str: ...
 
 
+def _decode_json_object(payload: str) -> Mapping[str, object]:
+    """Decode a realtime frame into its top-level field mapping."""
+    return json.loads(payload)
+
+
 class RealtimeEventNormalizer(Protocol):
     def should_drop(self, event: object) -> bool: ...
     def normalize(self, event: dict) -> dict: ...
@@ -294,7 +335,7 @@ class RealTimeStreaming:
         try:
             if event_obj.get("type") != "response.done":
                 return
-            response: Final = cast(dict[str, Any], event_obj.get("response", {}))
+            response: Final = cast(_ResponseDoneBody, event_obj.get("response", {}))
             item: Mapping[str, object]
             for item in response.get("output", []):
                 if item.get("type") == "function_call":
@@ -353,7 +394,7 @@ class RealTimeStreaming:
             sent = False
             for msg in transformed:
                 try:
-                    msg_obj = json.loads(msg)
+                    msg_obj = _decode_json_object(msg)
                 except (json.JSONDecodeError, TypeError):
                     msg_obj = None
                 if isinstance(msg_obj, dict) and self.provider_config.is_setup_message(msg_obj):
@@ -399,7 +440,7 @@ class RealTimeStreaming:
             return message
 
         try:
-            message_obj: Final[Mapping[str, object]] = json.loads(message)
+            message_obj: Final = _decode_json_object(message)
         except (json.JSONDecodeError, TypeError):
             return message
 
@@ -468,7 +509,7 @@ class RealTimeStreaming:
 
         for message in messages:
             try:
-                msg_type = json.loads(message).get("type")
+                msg_type = _decode_json_object(message).get("type")
             except (json.JSONDecodeError, TypeError):
                 collapsed.extend(pending_appends)
                 pending_appends = []
@@ -502,14 +543,14 @@ class RealTimeStreaming:
         if self._backend_setup_complete and not self._flushing_pending_messages_until_setup:
             return False
         try:
-            msg_obj: Final[Mapping[str, object]] = json.loads(message)
+            msg_obj: Final = _decode_json_object(message)
         except (json.JSONDecodeError, TypeError):
             return False
         return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES
 
     def _buffer_pending_message_until_setup(self, message: str) -> None:
         try:
-            msg_type = json.loads(message).get("type")
+            msg_type = _decode_json_object(message).get("type")
         except (json.JSONDecodeError, TypeError):
             msg_type = None
 
@@ -602,7 +643,7 @@ class RealTimeStreaming:
         ``return_new_content_delta_events`` modality lookup, ...).
         """
         try:
-            message_obj: Final = json.loads(transformed_message)
+            message_obj: Final = _decode_json_object(transformed_message)
             if "setup" in message_obj:
                 self.session_configuration_request = transformed_message
         except (json.JSONDecodeError, TypeError):
@@ -930,7 +971,7 @@ class RealTimeStreaming:
     def _parse_backend_event(raw_response: str) -> dict[str, object] | None:
         """Parse a backend frame once. Returns None for non-JSON or non-object frames."""
         try:
-            event: Final = json.loads(raw_response)
+            event: Final = _decode_json_object(raw_response)
         except (json.JSONDecodeError, TypeError):
             return None
         return event if isinstance(event, dict) else None
@@ -1030,14 +1071,14 @@ class RealTimeStreaming:
             await self.log_messages()
 
     @staticmethod
-    def _detect_beta_header(websocket: Any) -> bool:
+    def _detect_beta_header(websocket: _ScopedWebSocket) -> bool:
         """Return True if the client sent 'OpenAI-Beta: realtime=v1'.
 
         Checks the raw ASGI scope headers so it works for both FastAPI WebSocket
         objects and any test doubles that expose a .scope dict.
         """
         try:
-            headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = websocket.scope.get("headers", [])
+            headers: Final = websocket.scope.get("headers", [])
             for name, value in headers:
                 if isinstance(name, bytes):
                     name = name.decode("latin-1")
@@ -1183,6 +1224,7 @@ class RealTimeStreaming:
         return item
 
     async def client_ack_messages(self):
+        client_event: _ClientEventFrame
         try:
             while True:
                 message = await self.websocket.receive_text()
@@ -1194,11 +1236,12 @@ class RealTimeStreaming:
                     from litellm.types.guardrails import GuardrailEventHooks
 
                     msg_obj = json.loads(message)
-                    msg_type = msg_obj.get("type")
+                    client_event = msg_obj
+                    msg_type = client_event.get("type")
 
                     if msg_type == "conversation.item.create":
                         # Check user text messages for prompt injection
-                        item = msg_obj.get("item", {})
+                        item = client_event.get("item", {})
                         # Check function_call_output first so a client cannot
                         # bypass the tool-result guardrail by also setting
                         # role="user" on a function_call_output item.
@@ -1297,7 +1340,7 @@ class RealTimeStreaming:
                         and not self._guardrail_turn_detection_update_sent
                         and self._has_audio_transcription_guardrails()
                     ):
-                        session: object = msg_obj.setdefault("session", {})
+                        session: Mapping[str, object] | None = msg_obj.setdefault("session", {})
                         if isinstance(session, dict):
                             existing_td = session.get("turn_detection")
                             if not isinstance(existing_td, dict):
@@ -1324,7 +1367,7 @@ class RealTimeStreaming:
                         and not guardrail_turn_detection_injected
                         and self._has_audio_transcription_guardrails()
                     ):
-                        session = msg_obj.get("session")
+                        session = client_event.get("session")
                         if isinstance(session, dict):
                             td_overridden = False
                             flat_td = session.get("turn_detection")
@@ -1367,14 +1410,14 @@ class RealTimeStreaming:
                     # the upstream is in GA mode. Beta upstreams expect the flat
                     # session shape unchanged.
                     if msg_type == "session.update" and not self._backend_uses_beta_protocol:
-                        session = msg_obj.get("session", {})
+                        session = client_event.get("session", {})
                         if isinstance(session, dict):
                             session = self._remap_beta_session_to_ga(session)
                             msg_obj["session"] = session
                             message = json.dumps(msg_obj)
 
                     if msg_type == "session.update" and self._event_normalizer:
-                        session = msg_obj.get("session")
+                        session = client_event.get("session")
                         if isinstance(session, dict):
                             msg_obj["session"] = self._event_normalizer.patch_outgoing_session(session)
                             message = json.dumps(msg_obj)

litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py

@@ -4,7 +4,7 @@ Handler for the Anthropic v1/messages -> OpenAI Responses API path.
 Used when the target model is an OpenAI or Azure model.
 """
 
-from collections.abc import AsyncIterator, Coroutine
+from collections.abc import AsyncIterator, Coroutine, Mapping
 from typing import Any, Final
 
 import litellm
@@ -25,6 +25,11 @@ from .transformation import LiteLLMAnthropicToResponsesAPIAdapter
 _ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter()
 
 
+def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, object]:
+    """The litellm-specific kwargs forwarded verbatim onto the Responses API request."""
+    return extra_kwargs or {}
+
+
 def _build_responses_kwargs(
     *,
     max_tokens: int,
@@ -100,7 +105,7 @@ def _build_responses_kwargs(
 
     # Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.)
     excluded: Final = {"anthropic_messages"}
-    for key, value in (extra_kwargs or {}).items():
+    for key, value in _forwarded_kwargs(extra_kwargs).items():
         if key == "litellm_logging_obj" and value is not None:
             from litellm.litellm_core_utils.litellm_logging import (
                 Logging as LiteLLMLoggingObject,

litellm/llms/bedrock/passthrough/guardrail_translation/handler.py

@@ -1,4 +1,7 @@
-from typing import TYPE_CHECKING, Any, Final, Optional
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Final, Optional, Protocol, TypeAlias
+
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
 from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
@@ -40,7 +43,7 @@ def _generic_passthrough_handler() -> BaseTranslation:
 _StringHolder = tuple[Any, str | int]
 
 
-def _collect_strings(node: Any, holders: list[_StringHolder]) -> None:
+def _collect_strings(node: object, holders: list[_StringHolder]) -> None:
     """
     Record a (container, key) holder for every non-empty string value nested
     under an arbitrary JSON node, so prompt content a caller hides in fields
@@ -48,7 +51,7 @@ def _collect_strings(node: Any, holders: list[_StringHolder]) -> None:
     and can be written back in place. Iterative to avoid unbounded recursion
     on deeply nested payloads.
     """
-    stack: Final[list[Any]] = [node]
+    stack: Final[list[object]] = [node]
     while stack:
         current = stack.pop()
         if isinstance(current, dict):
@@ -129,7 +132,7 @@ def _extract_converse_texts(
 
 
 def _extract_converse_output_texts(
-    content_blocks: list[Any],
+    content_blocks: Sequence[object],
 ) -> tuple[list[str], list[_StringHolder]]:
     """
     Collect user-visible text from Bedrock Converse output content blocks.
@@ -178,10 +181,34 @@ def _write_back_texts(
         container[key] = guardrailed_texts[idx]
 
 
-_DeltaHolder = tuple[Any, Any, str | int]
+_GroupKey: TypeAlias = str | tuple[str, int]
 
 
-def _collect_stream_delta_text_holders(delta: Any) -> list[_DeltaHolder]:
+class _TextContainer(Protocol):
+    """JSON object whose ``key`` entry holds a guardrailable text string."""
+
+    def __getitem__(self, key: str, /) -> str: ...
+
+    def __setitem__(self, key: str, value: str, /) -> None: ...
+
+
+_DeltaHolder = tuple[_GroupKey, _TextContainer, str]
+
+
+class _StreamFrame(TypedDict):
+    """One raw event-stream frame plus the guardrailable texts it carries."""
+
+    raw: ReadOnly[bytes]
+    texts: ReadOnly[Sequence[tuple[_GroupKey, str]]]
+
+
+def _unpack_uint32(buffer: bytes) -> int:
+    import struct
+
+    return struct.unpack("!I", buffer)[0]
+
+
+def _collect_stream_delta_text_holders(delta: object) -> list[_DeltaHolder]:
     """
     Collect the user-visible text strings a Bedrock Converse ``contentBlockDelta``
     can carry, matching the coverage of the non-streaming output handler.
@@ -238,11 +265,11 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
 
         from botocore.eventstream import EventStreamBuffer
 
-        frames: Final[list[dict]] = []
+        frames: Final[list[_StreamFrame]] = []
         offset = 0
 
         while offset + 16 <= len(body_bytes):
-            total_length = struct.unpack("!I", body_bytes[offset : offset + 4])[0]
+            total_length = _unpack_uint32(body_bytes[offset : offset + 4])
             if total_length < 16 or offset + total_length > len(body_bytes):
                 break
             frame_raw = body_bytes[offset : offset + total_length]
@@ -263,10 +290,10 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
                 frames.append({"raw": frame_raw, "texts": []})
                 continue
 
-            texts: list[tuple[Any, str]] = []
+            texts: list[tuple[_GroupKey, str]] = []
             if event_type == "contentBlockDelta":
                 try:
-                    payload_dict = _json.loads(payload_bytes)
+                    payload_dict: dict[str, object] = _json.loads(payload_bytes)
                     texts = [
                         (group_key, container[key])
                         for group_key, container, key in _collect_stream_delta_text_holders(payload_dict.get("delta"))
@@ -282,9 +309,9 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
 
         trailing_bytes: Final = body_bytes[offset:]
 
-        group_order: Final[list[Any]] = []
-        group_members: Final[dict[Any, list[tuple[int, int]]]] = {}
-        group_texts: Final[dict[Any, list[str]]] = {}
+        group_order: Final[list[_GroupKey]] = []
+        group_members: Final[dict[_GroupKey, list[tuple[int, int]]]] = {}
+        group_texts: Final[dict[_GroupKey, list[str]]] = {}
         for frame_idx, frame in enumerate(frames):
             for local_idx, (group_key, text) in enumerate(frame["texts"]):
                 if group_key not in group_members:
@@ -351,8 +378,8 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
                 continue
 
             frame_raw = frame["raw"]
-            orig_total = struct.unpack("!I", frame_raw[0:4])[0]
-            orig_hdrs_len = struct.unpack("!I", frame_raw[4:8])[0]
+            orig_total = _unpack_uint32(frame_raw[0:4])
+            orig_hdrs_len = _unpack_uint32(frame_raw[4:8])
             headers_bytes = frame_raw[12 : 12 + orig_hdrs_len]
 
             try:
@@ -386,7 +413,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
         data: dict,
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
-    ) -> Any:
+    ) -> Mapping[str, object]:
         endpoint: Final = data.get("endpoint", "")
         body: Final = data.get("data")
 
@@ -428,12 +455,12 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
 
     async def process_output_response(
         self,
-        response: Any,
+        response: object,
         guardrail_to_apply: "CustomGuardrail",
         litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
-        user_api_key_dict: Any | None = None,
+        user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
         request_data: dict | None = None,
-    ) -> Any:
+    ) -> object:
         endpoint: Final = (request_data or {}).get("endpoint", "")
         if endpoint and not _is_converse_endpoint(endpoint):
             return await _generic_passthrough_handler().process_output_response(

litellm/llms/gemini/vector_stores/transformation.py

@@ -5,9 +5,11 @@ Implements the transformation between LiteLLM's unified vector store API
 and Google Gemini's File Search API.
 """
 
+from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final
 
 import httpx
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
 from litellm.llms.gemini.common_utils import (
@@ -35,6 +37,61 @@ else:
     LiteLLMLoggingObj = Any
 
 
+class GeminiRetrievedContext(TypedDict, total=False):
+    """Passage Gemini retrieved from a File Search store."""
+
+    text: ReadOnly[str]
+    uri: ReadOnly[str]
+    title: ReadOnly[str]
+
+
+class GeminiGroundingChunk(TypedDict, total=False):
+    """One source Gemini grounded its answer on."""
+
+    retrievedContext: ReadOnly[GeminiRetrievedContext]
+
+
+class GeminiGroundingSegment(TypedDict, total=False):
+    """Span of the generated answer a grounding support refers to."""
+
+    text: ReadOnly[str]
+
+
+class GeminiGroundingSupport(TypedDict, total=False):
+    """Citation linking an answer span to the grounding chunks that back it."""
+
+    segment: ReadOnly[GeminiGroundingSegment]
+    groundingChunkIndices: ReadOnly[Sequence[int]]
+    confidenceScores: ReadOnly[Sequence[float]]
+
+
+class GeminiFileSearchGroundingMetadata(TypedDict, total=False):
+    """Grounding metadata Gemini returns for a File Search candidate."""
+
+    groundingChunks: ReadOnly[Sequence[GeminiGroundingChunk]]
+    groundingSupports: ReadOnly[Sequence[GeminiGroundingSupport]]
+
+
+class GeminiFileSearchCandidate(TypedDict, total=False):
+    """One candidate of a Gemini File Search ``generateContent`` response."""
+
+    groundingMetadata: ReadOnly[GeminiFileSearchGroundingMetadata]
+
+
+class GeminiFileSearchResponse(TypedDict, total=False):
+    """Body of a ``generateContent`` call made with the File Search tool."""
+
+    candidates: ReadOnly[Sequence[GeminiFileSearchCandidate]]
+
+
+class GeminiFileSearchStore(TypedDict, total=False):
+    """Body of a Gemini ``fileSearchStores`` create response."""
+
+    name: ReadOnly[str]
+    displayName: ReadOnly[str]
+    createTime: ReadOnly[str]
+
+
 class GeminiVectorStoreConfig(BaseVectorStoreConfig):
     """
     Vector store configuration for Google Gemini File Search.
@@ -110,7 +167,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
         api_base: str,
         litellm_logging_obj: LiteLLMLoggingObj,
         litellm_params: dict,
-        extra_body: dict[str, Any] | None = None,
+        extra_body: Mapping[str, object] | None = None,
     ) -> tuple[str, dict]:
         """
         Transform search request to Gemini's generateContent format.
@@ -133,7 +190,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
         url: Final = f"{api_base}/models/{model}:generateContent"
 
         # Build file_search tool configuration (using snake_case as per Gemini docs)
-        file_search_config: Final[dict[str, Any]] = {"file_search_store_names": [vector_store_id]}
+        file_search_config: Final[dict[str, object]] = {"file_search_store_names": [vector_store_id]}
 
         # Add metadata filter if provided
         metadata_filter: Final = vector_store_search_optional_params.get("filters")
@@ -178,7 +235,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
         Extracts grounding metadata and citations from the response.
         """
         try:
-            response_data: Final = response.json()
+            response_data: Final[GeminiFileSearchResponse] = response.json()
             results: Final[list[VectorStoreSearchResult]] = []
 
             # Extract candidates and grounding metadata
@@ -246,7 +303,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
                             )
                         )
 
-            query: Final = litellm_logging_obj.model_call_details.get("query", "")
+            query: Final[str] = litellm_logging_obj.model_call_details.get("query", "")
 
             return VectorStoreSearchResponse(
                 object="vector_store.search_results.page",
@@ -273,7 +330,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
 
         # API key is passed via x-goog-api-key header (set in validate_environment)
 
-        request_body: Final[dict[str, Any]] = {}
+        request_body: Final[dict[str, object]] = {}
 
         # Add display name if provided
         name: Final = vector_store_create_optional_params.get("name")
@@ -287,7 +344,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
         Transform Gemini's fileSearchStore response to standard format.
         """
         try:
-            response_data: Final = response.json()
+            response_data: Final[GeminiFileSearchStore] = response.json()
 
             # Extract store name (format: fileSearchStores/xxxxxxx)
             store_name: Final = response_data.get("name", "")

litellm/llms/nvidia_riva/audio_transcription/handler.py

@@ -26,7 +26,9 @@ without the optional STT extras installed.
 
 import asyncio
 import inspect
-from typing import TYPE_CHECKING, Any, Final
+from collections.abc import Callable, Iterable
+from types import ModuleType
+from typing import TYPE_CHECKING, Any, Final, Protocol
 
 from litellm.litellm_core_utils.audio_utils.utils import (
     get_audio_file_name,
@@ -62,6 +64,45 @@ _DEFAULT_CHUNK_BYTES: Final = _DEFAULT_CHUNK_SAMPLES * 2  # int16 = 2 bytes/samp
 _RIVA_INSTALL_HINT = "NVIDIA Riva client is not installed. Install with `pip install 'litellm[stt-nvidia-riva]'`."
 
 
+class _RivaAuth(Protocol):
+    """Opaque ``riva.client.Auth`` handle."""
+
+
+class _AsrService(Protocol):
+    @property
+    def streaming_response_generator(self) -> Callable[..., Iterable[object]]: ...
+
+
+class _EndpointingConfig(Protocol):
+    """Opaque ``EndpointingConfig`` protobuf message."""
+
+
+class _EndpointingConfigField(Protocol):
+    CopyFrom: Callable[[_EndpointingConfig], None]
+
+
+class _RecognitionConfig(Protocol):
+    @property
+    def endpointing_config(self) -> _EndpointingConfigField: ...
+
+
+class _StreamingRecognitionConfig(Protocol):
+    """Opaque ``StreamingRecognitionConfig`` protobuf message."""
+
+
+class _AudioEncoding(Protocol):
+    @property
+    def LINEAR_PCM(self) -> object: ...
+
+
+def _auth_factory(riva_module: ModuleType) -> Callable[..., _RivaAuth]:
+    return riva_module.Auth
+
+
+def _audio_encoding(riva_asr_module: ModuleType) -> _AudioEncoding:
+    return riva_asr_module.AudioEncoding
+
+
 class NvidiaRivaAudioTranscription:
     """Sync + async entry point for Riva ASR."""
 
@@ -206,7 +247,9 @@ class NvidiaRivaAudioTranscription:
             riva_asr_module=riva_asr_module,
             recognition_config_dict=recognition_config_dict,
         )
-        streaming_config = riva_asr_module.StreamingRecognitionConfig(config=recognition_config, interim_results=False)
+        streaming_config: Final[_StreamingRecognitionConfig] = riva_asr_module.StreamingRecognitionConfig(
+            config=recognition_config, interim_results=False
+        )
 
         logging_obj.pre_call(
             input=None,
@@ -223,9 +266,9 @@ class NvidiaRivaAudioTranscription:
         )
 
         try:
-            asr_service: Final = riva_module.ASRService(auth_obj)
+            asr_service: Final[_AsrService] = riva_module.ASRService(auth_obj)
             audio_chunks: Final = self._iter_audio_chunks(resampled.pcm_bytes)
-            stream_kwargs: Final[dict[str, Any]] = {
+            stream_kwargs: Final[dict[str, object]] = {
                 "audio_chunks": audio_chunks,
                 "streaming_config": streaming_config,
             }
@@ -274,11 +317,11 @@ class NvidiaRivaAudioTranscription:
 
     def _construct_auth(
         self,
-        riva_module: Any,
+        riva_module: ModuleType,
         api_base: str,
         api_key: str | None,
         optional_params: dict,
-    ) -> Any:
+    ) -> _RivaAuth:
         """
         Build a ``riva.client.Auth`` object.
 
@@ -300,20 +343,22 @@ class NvidiaRivaAudioTranscription:
             metadata.append(("authorization", f"Bearer {api_key}"))
 
         try:
-            return riva_module.Auth(uri=api_base, use_ssl=use_ssl, metadata_args=metadata)
+            return _auth_factory(riva_module)(uri=api_base, use_ssl=use_ssl, metadata_args=metadata)
         except TypeError:
             # Older riva-client signatures used positional-only args.
-            return riva_module.Auth(None, use_ssl, api_base, metadata)
+            return _auth_factory(riva_module)(None, use_ssl, api_base, metadata)
 
-    def _build_recognition_config_proto(self, riva_asr_module: Any, recognition_config_dict: dict[str, Any]):
+    def _build_recognition_config_proto(
+        self, riva_asr_module: ModuleType, recognition_config_dict: dict[str, Any]
+    ) -> _RecognitionConfig:
         encoding_name: Final = (recognition_config_dict.get("encoding") or "LINEAR_PCM").upper()
-        encoding_enum: Final = getattr(
-            riva_asr_module.AudioEncoding,
+        encoding_enum: Final[object] = getattr(
+            _audio_encoding(riva_asr_module),
             encoding_name,
-            riva_asr_module.AudioEncoding.LINEAR_PCM,
+            _audio_encoding(riva_asr_module).LINEAR_PCM,
         )
 
-        config: Final = riva_asr_module.RecognitionConfig(
+        config: Final[_RecognitionConfig] = riva_asr_module.RecognitionConfig(
             encoding=encoding_enum,
             sample_rate_hertz=int(recognition_config_dict["sample_rate_hertz"]),
             language_code=recognition_config_dict["language_code"],
@@ -329,7 +374,7 @@ class NvidiaRivaAudioTranscription:
         endpointing: Final = recognition_config_dict.get("endpointing_config")
         if isinstance(endpointing, dict) and endpointing:
             try:
-                ep: Final = riva_asr_module.EndpointingConfig(**endpointing)
+                ep: Final[_EndpointingConfig] = riva_asr_module.EndpointingConfig(**endpointing)
                 config.endpointing_config.CopyFrom(ep)
             except Exception:
                 # If the user supplied an unknown EndpointingConfig field
@@ -340,7 +385,7 @@ class NvidiaRivaAudioTranscription:
         return config
 
     @staticmethod
-    def _supports_timeout_kwarg(callable_obj: Any) -> bool:
+    def _supports_timeout_kwarg(callable_obj: Callable[..., object]) -> bool:
         try:
             sig: Final = inspect.signature(callable_obj)
         except (TypeError, ValueError):
@@ -359,14 +404,14 @@ class NvidiaRivaAudioTranscription:
             yield chunk
 
     @staticmethod
-    def _collect_final_results(stream) -> list[dict[str, Any]]:
+    def _collect_final_results(stream) -> list[dict[str, object]]:
         """
         Walk the gRPC stream, ignore empty / non-final chunks, and return a
         list of normalized final-result dicts. Matching the user's note: the
         ``id`` blocks with no ``results`` are streaming heartbeats and must
         be skipped.
         """
-        final_results: Final[list[dict[str, Any]]] = []
+        final_results: Final[list[dict[str, object]]] = []
         for response in stream:
             results = getattr(response, "results", None) or []
             for result in results:
@@ -391,7 +436,7 @@ class NvidiaRivaAudioTranscription:
         return final_results
 
 
-def _import_riva():
+def _import_riva() -> tuple[ModuleType, ModuleType]:
     """
     Lazy import of ``riva.client`` and ``riva.client.proto.riva_asr_pb2``.
 

litellm/llms/oci/chat/cohere.py

@@ -84,9 +84,9 @@ def adapt_messages_to_cohere_standard(
             tool_calls_raw: Any = msg.get("tool_calls") or []
             for tc in tool_calls_raw:
                 tc_id = tc.get("id", "")
-                raw_args: Any = tc.get("function", {}).get("arguments", "{}")
+                raw_args = tc.get("function", {}).get("arguments", "{}")
                 try:
-                    params: dict[str, Any] = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
+                    params: dict[str, object] = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
                 except json.JSONDecodeError:
                     params = {}
                 tool_call_lookup[tc_id] = CohereToolCall(
@@ -111,10 +111,10 @@ def adapt_messages_to_cohere_standard(
         if role == "assistant" and msg.get("tool_calls"):
             tool_calls = []
             for tc in msg["tool_calls"]:  # pyright: ignore[reportOptionalIterable]  # truthiness check above rules out None
-                raw_arguments: Any = tc.get("function", {}).get("arguments", {})
+                raw_arguments = tc.get("function", {}).get("arguments", {})
                 if isinstance(raw_arguments, str):
                     try:
-                        arguments: dict[str, Any] = json.loads(raw_arguments)
+                        arguments: dict[str, object] = json.loads(raw_arguments)
                     except json.JSONDecodeError:
                         arguments = {}
                 else:
@@ -211,7 +211,7 @@ def handle_cohere_response(
     response_text: Final = cohere_response.chatResponse.text
     finish_reason: Final = _normalize_oci_finish_reason(cohere_response.chatResponse.finishReason)
 
-    tool_calls: list[dict[str, Any]] | None = None
+    tool_calls: list[dict[str, object]] | None = None
     if cohere_response.chatResponse.toolCalls:
         tool_calls = [
             {
@@ -232,7 +232,7 @@ def handle_cohere_response(
     # ``"tool_calls" in message`` (rather than truthiness) incorrectly conclude
     # that tool calls were attempted. Matches the generic handler's behaviour,
     # which only sets ``message.tool_calls`` when tool calls are present.
-    message: Final[dict[str, Any]] = {"role": "assistant", "content": content}
+    message: Final[dict[str, object]] = {"role": "assistant", "content": content}
     if tool_calls is not None:
         message["tool_calls"] = tool_calls
 
@@ -317,7 +317,7 @@ def handle_cohere_stream_chunk(
     # passing them through is the only chance to surface them.
     cohere_tool_calls = None if (is_terminal_consolidation and prior_tool_calls_emitted) else typed_chunk.toolCalls
 
-    tool_calls: list[dict[str, Any]] | None = None
+    tool_calls: list[dict[str, object]] | None = None
     if cohere_tool_calls:
         tool_calls = [
             {

litellm/llms/openai/responses/guardrail_translation/handler.py

@@ -28,10 +28,13 @@ Output: response.output is List[GenericResponseOutputItem] where each has:
     - text: str
 """
 
+from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Final, Union, cast
 
 from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
+from openai.types.responses.tool_param import FunctionToolParam
 from pydantic import BaseModel
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
 from litellm.completion_extras.litellm_responses_transformation.transformation import (
@@ -45,6 +48,7 @@ from litellm.types.llms.openai import (
     AllMessageValues,
     ChatCompletionToolCallChunk,
     ChatCompletionToolParam,
+    OpenAIMcpServerTool,
     ResponsesAPIStreamEvents,
 )
 from litellm.types.responses.main import (
@@ -56,10 +60,26 @@ from litellm.types.utils import GenericGuardrailAPIInputs
 
 if TYPE_CHECKING:
     from litellm.integrations.custom_guardrail import CustomGuardrail
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+    from litellm.proxy._types import UserAPIKeyAuth
     from litellm.types.llms.openai import ResponseInputParam
     from litellm.types.utils import ResponsesAPIResponse
 
 
+class ResponseOutputEnvelope(TypedDict, total=False):
+    """Dict form of a Responses API response, as far as guardrail write-back reads it."""
+
+    output: ReadOnly[Sequence[object]]
+    model: ReadOnly[str | None]
+
+
+class ResponsesStreamChunk(TypedDict, total=False):
+    """Responses API streaming event, as far as the accumulated-stream helpers read it."""
+
+    type: ReadOnly[str]
+    text: ReadOnly[str]
+
+
 class OpenAIResponsesHandler(BaseTranslation):
     """
     Handler for processing OpenAI Responses API with guardrails.
@@ -91,8 +111,8 @@ class OpenAIResponsesHandler(BaseTranslation):
         self,
         data: dict,
         guardrail_to_apply: "CustomGuardrail",
-        litellm_logging_obj: Any | None = None,
-    ) -> Any:
+        litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
+    ) -> dict[str, object]:
         """
         Process input by applying guardrails to text content.
 
@@ -108,7 +128,7 @@ class OpenAIResponsesHandler(BaseTranslation):
         # Handle simple string input
         if isinstance(input_data, str):
             inputs = GenericGuardrailAPIInputs(texts=[input_data])
-            original_tools: list[dict[str, Any]] = []
+            original_tools: list[dict[str, object]] = []
 
             # Extract and transform tools if present
             if "tools" in data and data["tools"]:
@@ -142,7 +162,7 @@ class OpenAIResponsesHandler(BaseTranslation):
         texts_to_check: Final[list[str]] = []
         images_to_check: Final[list[str]] = []
         task_mappings: Final[list[tuple[int, int | None]]] = []
-        original_tools_list: Final[list[dict[str, Any]]] = list(data.get("tools") or [])
+        original_tools_list: Final[list[dict[str, object]]] = list(data.get("tools") or [])
 
         # Step 1: Extract all text content, images, and tools
         for msg_idx, message in enumerate(input_data):
@@ -211,7 +231,7 @@ class OpenAIResponsesHandler(BaseTranslation):
 
     def _extract_and_transform_tools(
         self,
-        tools: list[dict[str, Any]],
+        tools: list[FunctionToolParam | OpenAIMcpServerTool],
         tools_to_check: list[ChatCompletionToolParam],
     ) -> None:
         """
@@ -228,7 +248,7 @@ class OpenAIResponsesHandler(BaseTranslation):
             ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools(tools)
             tools_to_check.extend(cast(list[ChatCompletionToolParam], transformed_tools))
 
-    def _remap_tools_to_responses_api_format(self, guardrailed_tools: list[Any]) -> list[dict[str, Any]]:
+    def _remap_tools_to_responses_api_format(self, guardrailed_tools: list[Any]) -> list[dict[str, object]]:
         """
         Remap guardrail-returned tools (Chat Completion format) back to
         Responses API request tool format.
@@ -239,9 +259,9 @@ class OpenAIResponsesHandler(BaseTranslation):
 
     def _merge_tools_after_guardrail(
         self,
-        original_tools: list[dict[str, Any]],
-        remapped: list[dict[str, Any]],
-    ) -> list[dict[str, Any]]:
+        original_tools: list[dict[str, object]],
+        remapped: list[dict[str, object]],
+    ) -> list[dict[str, object]]:
         """
         Merge remapped guardrailed tools with original tools that were not sent
         to the guardrail (e.g. web_search, web_search_preview), preserving order.
@@ -250,7 +270,7 @@ class OpenAIResponsesHandler(BaseTranslation):
         """
         if not original_tools:
             return remapped
-        result: Final[list[dict[str, Any]]] = []
+        result: Final[list[dict[str, object]]] = []
         j = 0
         for tool in original_tools:
             if isinstance(tool, dict) and tool.get("type") in (
@@ -269,8 +289,8 @@ class OpenAIResponsesHandler(BaseTranslation):
     def _apply_guardrailed_tools_to_data(
         self,
         data: dict,
-        original_tools: list[dict[str, Any]],
-        guardrailed_tools: list[Any] | None,
+        original_tools: list[dict[str, object]],
+        guardrailed_tools: list[ChatCompletionToolParam] | None,
     ) -> None:
         """Remap guardrailed tools to Responses API format and merge with original, then set data['tools']."""
         if guardrailed_tools is not None:
@@ -279,7 +299,7 @@ class OpenAIResponsesHandler(BaseTranslation):
 
     def _extract_input_text_and_images(
         self,
-        message: Any,  # Can be Dict[str, Any] or ResponseInputParam
+        message: Any,
         msg_idx: int,
         texts_to_check: list[str],
         images_to_check: list[str],
@@ -348,12 +368,12 @@ class OpenAIResponsesHandler(BaseTranslation):
 
     async def process_output_response(
         self,
-        response: "ResponsesAPIResponse",
+        response: Union["ResponsesAPIResponse", ResponseOutputEnvelope],
         guardrail_to_apply: "CustomGuardrail",
-        litellm_logging_obj: Any | None = None,
-        user_api_key_dict: Any | None = None,
+        litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
+        user_api_key_dict: "UserAPIKeyAuth | None" = None,
         request_data: dict | None = None,
-    ) -> Any:
+    ) -> Union["ResponsesAPIResponse", ResponseOutputEnvelope]:
         """
         Process output response by applying guardrails to text content and tool calls.
 
@@ -381,6 +401,7 @@ class OpenAIResponsesHandler(BaseTranslation):
         # Track (output_item_index, content_index) for each text
 
         # Handle both dict and Pydantic object responses
+        response_output: Sequence[object]
         if isinstance(response, dict):
             response_output = response.get("output", [])
         elif hasattr(response, "output"):
@@ -426,7 +447,7 @@ class OpenAIResponsesHandler(BaseTranslation):
             if tool_calls_to_check:
                 inputs["tool_calls"] = tool_calls_to_check
             # Include model information from the response if available
-            response_model = None
+            response_model: str | None = None
             if isinstance(response, dict):
                 response_model = response.get("model")
             elif hasattr(response, "model"):
@@ -458,8 +479,8 @@ class OpenAIResponsesHandler(BaseTranslation):
         self,
         responses_so_far: list[Any],
         guardrail_to_apply: "CustomGuardrail",
-        litellm_logging_obj: Any | None = None,
-        user_api_key_dict: Any | None = None,
+        litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
+        user_api_key_dict: "UserAPIKeyAuth | None" = None,
         request_data: dict | None = None,
     ) -> list[Any]:
         """
@@ -488,10 +509,10 @@ class OpenAIResponsesHandler(BaseTranslation):
         # final chunk; iterate output items, apply guardrail, write back.     #
         # ------------------------------------------------------------------ #
         if final_chunk.get("type") == "response.completed":
-            response_obj: Final = final_chunk.get("response") or {}
+            response_obj: Final[ResponseOutputEnvelope] = final_chunk.get("response") or {}
             if not hasattr(response_obj, "get"):
                 return responses_so_far
-            outputs: Final[list[Any]] = response_obj.get("output") or []
+            outputs: Final[Sequence[object]] = response_obj.get("output") or []
 
             texts_to_check: Final[list[str]] = []
             tool_calls_to_check: Final[list[ChatCompletionToolCallChunk]] = []
@@ -586,7 +607,7 @@ class OpenAIResponsesHandler(BaseTranslation):
             )
         return responses_so_far
 
-    def _check_streaming_has_ended(self, responses_so_far: list[Any]) -> bool:
+    def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool:
         """
         Check if the streaming has ended.
         """
@@ -599,7 +620,7 @@ class OpenAIResponsesHandler(BaseTranslation):
         }
         return responses_so_far[-1].get("type") in terminal_types
 
-    def get_streaming_string_so_far(self, responses_so_far: list[Any]) -> str:
+    def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str:
         """
         Get the string so far from the responses so far.
         """
@@ -641,7 +662,7 @@ class OpenAIResponsesHandler(BaseTranslation):
 
     def _extract_output_text_and_images(
         self,
-        output_item: Any,
+        output_item: object,
         output_idx: int,
         texts_to_check: list[str],
         images_to_check: list[str],
@@ -724,7 +745,7 @@ class OpenAIResponsesHandler(BaseTranslation):
 
     async def _apply_guardrail_responses_to_output(
         self,
-        response: Union["ResponsesAPIResponse", dict[Any, Any]],
+        response: Union["ResponsesAPIResponse", ResponseOutputEnvelope],
         responses: list[str],
         task_mappings: list[tuple[int, int]],
     ) -> None:

litellm/llms/snowflake/chat/transformation.py

@@ -9,9 +9,11 @@ Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api
 """
 
 import json
-from typing import TYPE_CHECKING, Any, Final
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict
 
 import httpx
+from typing_extensions import ReadOnly
 
 from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
 from litellm.types.utils import (
@@ -44,6 +46,47 @@ _CLAUDE_MODEL_PREFIXES: Final = (
 )
 
 
+class _AnthropicContentBlock(TypedDict, total=False):
+    type: ReadOnly[str]
+    text: ReadOnly[str]
+    id: ReadOnly[str]
+    name: ReadOnly[str]
+    input: ReadOnly[Mapping[str, object]]
+
+
+class _AnthropicUsageBlock(TypedDict, total=False):
+    input_tokens: ReadOnly[int]
+    output_tokens: ReadOnly[int]
+
+
+class _AnthropicMessagesResponse(TypedDict, total=False):
+    id: ReadOnly[str]
+    model: ReadOnly[str]
+    stop_reason: ReadOnly[str]
+    content: ReadOnly[Sequence[_AnthropicContentBlock]]
+    usage: ReadOnly[_AnthropicUsageBlock]
+
+
+class _ChatCompletionsResponse(Protocol):
+    """Response view that decodes the Cortex chat-completions body as a field mapping."""
+
+    def json(self) -> Mapping[str, object]: ...
+
+
+class _MessagesResponse(Protocol):
+    """Response view that decodes the Cortex messages body in Anthropic shape."""
+
+    def json(self) -> _AnthropicMessagesResponse: ...
+
+
+def _decoded_chat_completions(response: _ChatCompletionsResponse) -> Mapping[str, object]:
+    return response.json()
+
+
+def _decoded_messages(response: _MessagesResponse) -> _AnthropicMessagesResponse:
+    return response.json()
+
+
 def _is_claude_model(model: str) -> bool:
     """Return True if model name (after stripping snowflake/ prefix) is a Claude model."""
     name: Final = model.lower().removeprefix("snowflake/")
@@ -129,7 +172,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
         for tool in tools:
             if tool.get("type") == "function" and "function" in tool:
                 func = tool["function"]
-                anthropic_tool: dict[str, Any] = {
+                anthropic_tool: dict[str, object] = {
                     "name": func.get("name", ""),
                 }
                 if "description" in func:
@@ -173,7 +216,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
             elif role == "assistant":
                 tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None)
                 if tool_calls:
-                    content_blocks: list[dict[str, Any]] = []
+                    content_blocks: list[dict[str, object]] = []
                     if content:
                         content_blocks.append({"type": "text", "text": content})
                     for tc in tool_calls:
@@ -310,7 +353,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
 
         model_name: Final = model.removeprefix("snowflake/")
 
-        body: Final[dict[str, Any]] = {
+        body: Final[dict[str, object]] = {
             "model": model_name,
             "messages": conversation,
             "stream": stream,
@@ -336,7 +379,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
         messages: list[AllMessageValues],
         optional_params: dict,
         litellm_params: dict,
-        encoding: Any,
+        encoding: object,
         api_key: str | None = None,
         json_mode: bool | None = None,
     ) -> ModelResponse:
@@ -356,7 +399,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
         messages: list[AllMessageValues],
     ) -> ModelResponse:
         """Parse standard OpenAI chat completions response."""
-        response_json: Final = raw_response.json()
+        response_json: Final = _decoded_chat_completions(raw_response)
 
         logging_obj.post_call(
             input=messages,
@@ -383,7 +426,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
         messages: list[AllMessageValues],
     ) -> ModelResponse:
         """Parse Anthropic Messages response into OpenAI format."""
-        response_json: Final = raw_response.json()
+        response_json: Final = _decoded_messages(raw_response)
 
         logging_obj.post_call(
             input=messages,
@@ -447,10 +490,10 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
 
     def get_model_response_iterator(
         self,
-        streaming_response: Any,
+        streaming_response: object,
         sync_stream: bool,
         json_mode: bool | None = False,
-    ) -> Any:
+    ) -> "SnowflakeStreamingHandler":
         return SnowflakeStreamingHandler(
             streaming_response=streaming_response,
             sync_stream=sync_stream,
@@ -468,7 +511,7 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator):
 
     def __init__(
         self,
-        streaming_response: Any,
+        streaming_response: object,
         sync_stream: bool,
         json_mode: bool | None = False,
     ):

litellm/llms/soniox/audio_transcription/handler.py

@@ -18,10 +18,11 @@ handler (analogous to the OpenAI / Azure transcription handlers).
 import asyncio
 import math
 import time
-from collections.abc import Coroutine
+from collections.abc import Coroutine, Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final
 
 import httpx
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm.litellm_core_utils.audio_utils.utils import (
     get_audio_file_name,
@@ -57,6 +58,49 @@ else:
     LiteLLMLoggingObj = Any
 
 
+class _TranscriptionMeta(TypedDict, total=False):
+    """Fields the handler reads from a Soniox transcription object."""
+
+    status: ReadOnly[str]
+    error_message: ReadOnly[str]
+    error_type: ReadOnly[str]
+    audio_duration_ms: ReadOnly[float]
+
+
+class _IdentifiedResource(TypedDict):
+    """Soniox create/upload response, carrying the new resource id."""
+
+    id: ReadOnly[str]
+
+
+class _SonioxErrorBody(TypedDict, total=False):
+    """Fields the handler reads from a Soniox error response body."""
+
+    error_message: ReadOnly[object]
+    error: ReadOnly[object]
+
+
+class _SonioxJsonView(TypedDict, total=False):
+    """Typed reads of decoded Soniox JSON response bodies."""
+
+    resource: ReadOnly[_IdentifiedResource]
+    transcription: ReadOnly[_TranscriptionMeta]
+    transcript: ReadOnly[Mapping[str, object]]
+    error: ReadOnly[_SonioxErrorBody]
+
+
+class _HandlerOptions(TypedDict):
+    """Handler-only options pulled out of ``optional_params``."""
+
+    poll_interval: ReadOnly[float]
+    max_attempts: ReadOnly[int]
+    cleanup: ReadOnly[Sequence[str]]
+    filename_override: ReadOnly[str | None]
+    audio_url: ReadOnly[str | None]
+    file_id: ReadOnly[str | None]
+    response_format: ReadOnly[str | None]
+
+
 class SonioxAudioTranscriptionHandler:
     """Orchestrates the Soniox async transcription flow."""
 
@@ -78,9 +122,9 @@ class SonioxAudioTranscriptionHandler:
         api_base: str | None,
         client: HTTPHandler | AsyncHTTPHandler | None = None,
         atranscription: bool = False,
-        headers: dict[str, Any] | None = None,
+        headers: dict[str, str] | None = None,
         provider_config: SonioxAudioTranscriptionConfig | None = None,
-    ) -> TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse]:
+    ) -> TranscriptionResponse | Coroutine[object, object, TranscriptionResponse]:
         """Sync/async dispatch for Soniox transcription requests.
 
         Note: ``max_retries`` is accepted for signature compatibility with
@@ -134,12 +178,12 @@ class SonioxAudioTranscriptionHandler:
         api_key: str | None,
         api_base: str | None,
         provider_config: SonioxAudioTranscriptionConfig,
-        headers: dict[str, Any],
+        headers: dict[str, str],
     ) -> tuple[
         dict[str, str],  # auth headers
         str,  # api_base (no trailing slash)
-        dict[str, Any],  # body for POST /v1/transcriptions (without file_id/audio_url)
-        dict[str, Any],  # handler-only options (poll interval, cleanup, ...)
+        dict[str, object],  # body for POST /v1/transcriptions (without file_id/audio_url)
+        _HandlerOptions,  # handler-only options (poll interval, cleanup, ...)
     ]:
         # Validate env -> auth headers.
         auth_headers: Final = provider_config.validate_environment(
@@ -184,32 +228,31 @@ class SonioxAudioTranscriptionHandler:
         clamped_poll_interval: Final = max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL))
         clamped_max_attempts: Final = max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS))
 
-        handler_opts: Final[dict[str, Any]] = {
+        # response_format is handled by LiteLLM post-processing, not Soniox.
+        handler_opts: Final[_HandlerOptions] = {
             "poll_interval": clamped_poll_interval,
             "max_attempts": clamped_max_attempts,
             "cleanup": cleanup,
             "filename_override": filename_override,
             "audio_url": params.pop("audio_url", None),
             "file_id": params.pop("file_id", None),
+            "response_format": params.pop("response_format", None),
         }
 
         # Soniox does not accept `language` directly; map_openai_params should
         # already have translated it, but drop any leftover to be safe.
         params.pop("language", None)
 
-        # response_format is handled by LiteLLM post-processing, not Soniox.
-        handler_opts["response_format"] = params.pop("response_format", None)
-
         return auth_headers, base_url, params, handler_opts
 
     def _build_create_body(
         self,
         model: str,
-        optional_params: dict,
-        handler_opts: dict[str, Any],
+        optional_params: Mapping[str, object],
+        handler_opts: _HandlerOptions,
         file_id: str | None,
-    ) -> dict[str, Any]:
-        body: Final[dict[str, Any]] = {"model": model}
+    ) -> dict[str, object]:
+        body: Final[dict[str, object]] = {"model": model}
         # Soniox-native passthrough fields
         for key, value in optional_params.items():
             if value is None:
@@ -224,7 +267,7 @@ class SonioxAudioTranscriptionHandler:
         return body
 
     @staticmethod
-    def _redact_body_for_logging(body: dict[str, Any]) -> dict[str, Any]:
+    def _redact_body_for_logging(body: dict[str, object]) -> dict[str, object]:
         """Return a shallow copy of ``body`` with secret fields redacted.
 
         Soniox's create-transcription body can include
@@ -248,7 +291,7 @@ class SonioxAudioTranscriptionHandler:
         logging_obj: LiteLLMLoggingObj,
         api_key: str | None,
         api_base: str,
-        body: dict[str, Any],
+        body: dict[str, object],
     ) -> None:
         try:
             logging_obj.pre_call(
@@ -270,8 +313,8 @@ class SonioxAudioTranscriptionHandler:
         logging_obj: LiteLLMLoggingObj,
         audio_file: FileTypes | None,
         api_key: str | None,
-        body: dict[str, Any],
-        original_response: Any,
+        body: dict[str, object],
+        original_response: Mapping[str, object],
     ) -> None:
         try:
             logging_obj.post_call(
@@ -285,6 +328,11 @@ class SonioxAudioTranscriptionHandler:
             # observability integration must never break a real Soniox call.
             pass
 
+    @staticmethod
+    def _transcription_meta(response: httpx.Response) -> _TranscriptionMeta:
+        polled: Final[_SonioxJsonView] = {"transcription": response.json()}
+        return polled["transcription"]
+
     @staticmethod
     def _raise_for_response(
         response: httpx.Response,
@@ -293,8 +341,8 @@ class SonioxAudioTranscriptionHandler:
     ) -> None:
         if response.status_code >= 400:
             try:
-                payload: Final = response.json()
-                message = payload.get("error_message") or payload.get("error") or response.text
+                payload: Final[_SonioxJsonView] = {"error": response.json()}
+                message = payload["error"].get("error_message") or payload["error"].get("error") or response.text
             except Exception:
                 message = response.text
             raise provider_config.get_error_class(
@@ -319,7 +367,7 @@ class SonioxAudioTranscriptionHandler:
         api_key: str | None,
         api_base: str | None,
         client: HTTPHandler | None,
-        headers: dict[str, Any],
+        headers: dict[str, str],
         provider_config: SonioxAudioTranscriptionConfig,
     ) -> TranscriptionResponse:
         auth_headers, base_url, opt_params, handler_opts = self._prepare(
@@ -378,7 +426,8 @@ class SonioxAudioTranscriptionHandler:
                 timeout=timeout,
             )
             self._raise_for_response(create_resp, provider_config, "create transcription")
-            transcription_id = create_resp.json()["id"]
+            created: Final[_SonioxJsonView] = {"resource": create_resp.json()}
+            transcription_id = created["resource"]["id"]
 
             transcription_meta: Final = self._sync_poll_until_completed(
                 http_client=http_client,
@@ -397,9 +446,9 @@ class SonioxAudioTranscriptionHandler:
                 timeout=timeout,
             )
             self._raise_for_response(transcript_resp, provider_config, "fetch transcript")
-            transcript: Final = transcript_resp.json()
+            fetched: Final[_SonioxJsonView] = {"transcript": transcript_resp.json()}
 
-            payload: Final = {"transcription": transcription_meta, "transcript": transcript}
+            payload: Final = {"transcription": transcription_meta, "transcript": fetched["transcript"]}
             response: Final = provider_config._build_response_from_payload(
                 payload,
                 model_response=model_response,
@@ -454,7 +503,8 @@ class SonioxAudioTranscriptionHandler:
             timeout=timeout,
         )
         self._raise_for_response(resp, provider_config, "upload file")
-        return resp.json()["id"]
+        uploaded: Final[_SonioxJsonView] = {"resource": resp.json()}
+        return uploaded["resource"]["id"]
 
     def _sync_poll_until_completed(
         self,
@@ -466,7 +516,7 @@ class SonioxAudioTranscriptionHandler:
         max_attempts: int,
         timeout: float,
         provider_config: SonioxAudioTranscriptionConfig,
-    ) -> dict[str, Any]:
+    ) -> _TranscriptionMeta:
         for _ in range(max_attempts):
             resp = http_client.get(
                 url=f"{base_url}/v1/transcriptions/{transcription_id}",
@@ -474,7 +524,7 @@ class SonioxAudioTranscriptionHandler:
                 timeout=timeout,
             )
             self._raise_for_response(resp, provider_config, "poll transcription")
-            data = resp.json()
+            data = self._transcription_meta(resp)
             status = data.get("status")
             if status == "completed":
                 return data
@@ -502,7 +552,7 @@ class SonioxAudioTranscriptionHandler:
         http_client: HTTPHandler,
         base_url: str,
         auth_headers: dict[str, str],
-        cleanup: list[str],
+        cleanup: Sequence[str],
         file_id_to_cleanup: str | None,
         transcription_id: str | None,
         timeout: float,
@@ -548,7 +598,7 @@ class SonioxAudioTranscriptionHandler:
         api_key: str | None,
         api_base: str | None,
         client: AsyncHTTPHandler | None,
-        headers: dict[str, Any],
+        headers: dict[str, str],
         provider_config: SonioxAudioTranscriptionConfig,
     ) -> TranscriptionResponse:
         import litellm
@@ -610,7 +660,8 @@ class SonioxAudioTranscriptionHandler:
                 timeout=timeout,
             )
             self._raise_for_response(create_resp, provider_config, "create transcription")
-            transcription_id = create_resp.json()["id"]
+            created: Final[_SonioxJsonView] = {"resource": create_resp.json()}
+            transcription_id = created["resource"]["id"]
 
             transcription_meta: Final = await self._async_poll_until_completed(
                 http_client=http_client,
@@ -629,9 +680,9 @@ class SonioxAudioTranscriptionHandler:
                 timeout=timeout,
             )
             self._raise_for_response(transcript_resp, provider_config, "fetch transcript")
-            transcript: Final = transcript_resp.json()
+            fetched: Final[_SonioxJsonView] = {"transcript": transcript_resp.json()}
 
-            payload: Final = {"transcription": transcription_meta, "transcript": transcript}
+            payload: Final = {"transcription": transcription_meta, "transcript": fetched["transcript"]}
             response: Final = provider_config._build_response_from_payload(
                 payload,
                 model_response=model_response,
@@ -685,7 +736,8 @@ class SonioxAudioTranscriptionHandler:
             timeout=timeout,
         )
         self._raise_for_response(resp, provider_config, "upload file")
-        return resp.json()["id"]
+        uploaded: Final[_SonioxJsonView] = {"resource": resp.json()}
+        return uploaded["resource"]["id"]
 
     async def _async_poll_until_completed(
         self,
@@ -697,7 +749,7 @@ class SonioxAudioTranscriptionHandler:
         max_attempts: int,
         timeout: float,
         provider_config: SonioxAudioTranscriptionConfig,
-    ) -> dict[str, Any]:
+    ) -> _TranscriptionMeta:
         for _ in range(max_attempts):
             resp = await http_client.get(
                 url=f"{base_url}/v1/transcriptions/{transcription_id}",
@@ -705,7 +757,7 @@ class SonioxAudioTranscriptionHandler:
                 timeout=timeout,
             )
             self._raise_for_response(resp, provider_config, "poll transcription")
-            data = resp.json()
+            data = self._transcription_meta(resp)
             status = data.get("status")
             if status == "completed":
                 return data
@@ -733,7 +785,7 @@ class SonioxAudioTranscriptionHandler:
         http_client: AsyncHTTPHandler,
         base_url: str,
         auth_headers: dict[str, str],
-        cleanup: list[str],
+        cleanup: Sequence[str],
         file_id_to_cleanup: str | None,
         transcription_id: str | None,
         timeout: float,

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -123,6 +123,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
     iter_known_server_prefixes,
     iter_known_tool_name_spellings,
     logging_safe_mcp_headers,
+    lookup_mcp_server_auth_in_headers,
     match_known_server_prefix,
     match_known_tool_name,
     merge_mcp_headers,
@@ -873,6 +874,53 @@ def _openapi_forwarded_extra_headers(
     return forwarded or None
 
 
+def _resolve_openapi_tool_auth(
+    mcp_server: MCPServer,
+    mcp_auth_header: str | None,
+    mcp_server_auth_headers: Mapping[str, str | dict[str, str]] | None,  # mutable-ok: sink shape
+    raw_headers: dict[str, str] | None,  # mutable-ok: sink takes a concrete dict
+    user_api_key_auth: UserAPIKeyAuth | None,
+) -> tuple[str | None, dict[str, str] | None, str | dict[str, str] | None]:  # mutable-ok: sink shapes
+    """The caller's upstream credential for one ``spec_path`` server, for both OpenAPI dispatch arms.
+
+    A per-server ``x-mcp-{alias}-authorization`` wins over the deprecated global / BYOK
+    ``mcp_auth_header``, the same precedence ``_call_regular_mcp_tool`` applies, so the OpenAPI and
+    managed paths cannot disagree about which credential is authoritative. The two kinds are not
+    interchangeable: a per-server value is already a complete header value and is forwarded verbatim,
+    while a BYOK credential is a raw secret that takes the server's auth-type prefix. Formatting the
+    former would ship ``Bearer Bearer <token>``.
+
+    Returns the ``Authorization`` value to inject, the extra headers to forward, and the credential to
+    hand ``resolve_openapi_upstream_auth``, whose passthrough arm reads it via
+    ``_passthrough_token_from_mcp_auth_header``. The per-server Authorization travels only in the
+    credential, never also in the forwarded headers, because the resolver pops Authorization out of
+    those and would otherwise have two sources to reconcile.
+    """
+    forwarded: Final = _openapi_forwarded_extra_headers(mcp_server, raw_headers, user_api_key_auth)
+    per_server: Final = (
+        lookup_mcp_server_auth_in_headers(
+            mcp_server_auth_headers,
+            alias=mcp_server.alias,
+            server_name=mcp_server.server_name,
+        )
+        if mcp_server_auth_headers
+        else None
+    )
+
+    if isinstance(per_server, dict):
+        authorization: Final = next((v for k, v in per_server.items() if k.lower() == "authorization"), None)
+        merged: Final = merge_mcp_headers(extra_headers=forwarded, static_headers=_without_authorization(per_server))
+        if authorization is None:
+            byok: Final = _format_byok_openapi_auth_header(mcp_server, mcp_auth_header) if mcp_auth_header else None
+            return byok, merged, mcp_auth_header
+        return authorization, merged, per_server
+    if isinstance(per_server, str) and per_server:
+        return per_server, forwarded, per_server
+    if mcp_auth_header:
+        return _format_byok_openapi_auth_header(mcp_server, mcp_auth_header), forwarded, mcp_auth_header
+    return None, forwarded, None
+
+
 async def _resolve_byok_mcp_auth_header(
     mcp_server: MCPServer,
     user_api_key_auth: UserAPIKeyAuth | None,
@@ -3193,10 +3241,6 @@ class MCPServerManager:
             # Get server-specific auth header if available
             server_auth_header: str | dict[str, str] | None = None
             if mcp_server_auth_headers:
-                from litellm.proxy._experimental.mcp_server.utils import (
-                    lookup_mcp_server_auth_in_headers,
-                )
-
                 server_auth_header = lookup_mcp_server_auth_in_headers(
                     mcp_server_auth_headers,
                     alias=server.alias,
@@ -5221,11 +5265,6 @@ class MCPServerManager:
         # the exact case of server alias/name (e.g., '1litellmagcgateway' vs '1LiteLLMAGCGateway')
         server_auth_header: dict[str, str] | str | None = None
         if mcp_server_auth_headers:
-            # Normalize keys for case-insensitive lookup
-            from litellm.proxy._experimental.mcp_server.utils import (
-                lookup_mcp_server_auth_in_headers,
-            )
-
             server_auth_header = lookup_mcp_server_auth_in_headers(
                 mcp_server_auth_headers,
                 alias=mcp_server.alias,
@@ -5718,16 +5757,20 @@ class MCPServerManager:
                     server_name,
                 )
 
-            auth_header_value: Final = (
-                _format_byok_openapi_auth_header(mcp_server, mcp_auth_header) if mcp_auth_header else None
+            auth_header_value, openapi_forwarded_headers, upstream_credential = _resolve_openapi_tool_auth(
+                mcp_server=mcp_server,
+                mcp_auth_header=mcp_auth_header,
+                mcp_server_auth_headers=mcp_server_auth_headers,
+                raw_headers=raw_headers,
+                user_api_key_auth=user_api_key_auth,
             )
             resolved_auth_headers, forwarded_headers = await self.resolve_openapi_upstream_auth(
                 mcp_server=mcp_server,
                 oauth2_headers=caller_oauth2_headers,
                 raw_headers=raw_headers,
-                mcp_auth_header=mcp_auth_header,
+                mcp_auth_header=upstream_credential,
                 user_api_key_auth=user_api_key_auth,
-                forwarded_headers=_openapi_forwarded_extra_headers(mcp_server, raw_headers, user_api_key_auth),
+                forwarded_headers=openapi_forwarded_headers,
             )
 
             async def _call_openapi_via_handler():

litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py

@@ -9,10 +9,11 @@ import os
 import re
 from collections.abc import Mapping, Sequence
 from pathlib import PurePosixPath
-from typing import Any, Final, TypeAlias, TypedDict
+from typing import Any, Final, TypedDict
 from urllib.parse import quote
 
 import httpx
+from typing_extensions import ReadOnly, Required
 
 # Tool names emitted from OpenAPI specs must work across all major LLM providers.
 # OpenAI/Anthropic/Bedrock all enforce a character class roughly equivalent to
@@ -47,11 +48,17 @@ from litellm.proxy._experimental.mcp_server.tool_registry import (
     global_mcp_tool_registry,
 )
 
-_OpenAPIParameter: TypeAlias = Mapping[str, Any]
-
 
 class _OpenAPIJSONSchema(TypedDict, total=False):
     properties: Mapping[str, object]
+    type: ReadOnly[str]
+
+
+class _OpenAPIParameter(TypedDict, total=False):
+    name: Required[ReadOnly[str]]
+    description: ReadOnly[str]
+    required: ReadOnly[bool]
+    schema: ReadOnly[_OpenAPIJSONSchema]
 
 
 class _OpenAPIMediaType(TypedDict, total=False):
@@ -241,7 +248,7 @@ def resolve_operation_params(
     operation: _OpenAPIOperation,
     path_item: _OpenAPIPathItem,
     components: _OpenAPIComponents,
-) -> dict[str, Any]:
+) -> _OpenAPIOperation:
     """Return a copy of *operation* with fully-resolved, merged parameters.
 
     Handles two common patterns in real-world OpenAPI specs:
@@ -261,12 +268,11 @@ def resolve_operation_params(
     op_level: Final = _resolve_param_list(operation.get("parameters", []), component_params)
     op_keys: Final = {(p["name"], p.get("in")) for p in op_level}
     merged: Final = [p for p in path_level if (p["name"], p.get("in")) not in op_keys] + op_level
-    result: Final = dict(operation)
-    result["parameters"] = merged
+    result: Final[_OpenAPIOperation] = {**operation, "parameters": merged}
     return result
 
 
-def extract_parameters(operation: Mapping[str, Any]) -> tuple[Sequence[str], Sequence[str], Sequence[str]]:
+def extract_parameters(operation: _OpenAPIOperation) -> tuple[Sequence[str], Sequence[str], Sequence[str]]:
     """Extract parameter names from OpenAPI operation."""
     path_params: Final = []
     query_params: Final = []
@@ -292,7 +298,7 @@ def extract_parameters(operation: Mapping[str, Any]) -> tuple[Sequence[str], Seq
     return path_params, query_params, body_params
 
 
-def build_input_schema(operation: Mapping[str, Any]) -> dict[str, Any]:
+def build_input_schema(operation: _OpenAPIOperation) -> dict[str, object]:
     """Build MCP input schema from OpenAPI operation."""
     properties: Final = {}
     required: Final = []
@@ -389,7 +395,7 @@ def _merge_openapi_tool_request_headers(
 def create_tool_function(
     path: str,
     method: str,
-    operation: Mapping[str, Any],
+    operation: _OpenAPIOperation,
     base_url: str,
     headers: dict[str, str] | None = None,
 ):
@@ -443,7 +449,7 @@ def create_tool_function(
                 url = url.replace("{{" + param_name + "}}", safe_value)
 
         # Build query params using original parameter names
-        params: Final[dict[str, Any]] = {}
+        params: Final[dict[str, object]] = {}
         for param_name in query_params:
             param_value = kwargs.get(param_name, "")
             if param_value:
@@ -451,7 +457,7 @@ def create_tool_function(
                 params[param_name] = param_value
 
         # Build request body
-        json_body: dict[str, Any] | None = None
+        json_body: dict[str, object] | None = None
         if body_params:
             # Try "body" first (most common), then check all body param names
             body_value = kwargs.get("body", {})
@@ -492,7 +498,7 @@ def create_tool_function(
 
 def register_tools_from_openapi(spec: Mapping[str, Any], base_url: str) -> None:
     """Register MCP tools from OpenAPI specification."""
-    paths: Final[Mapping[str, Mapping[str, Any]]] = spec.get("paths", {})
+    paths: Final[Mapping[str, Mapping[str, _OpenAPIOperation]]] = spec.get("paths", {})
     used_names: Final = set()
 
     for path, path_item in paths.items():

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -41,6 +41,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 if TYPE_CHECKING:
     from mcp.types import CallToolResult
 
+    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
     from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload
 from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
 from litellm.types.mcp import MCPAuth
@@ -108,7 +109,7 @@ if MCP_AVAILABLE:
     ########################################################
     ############ MCP Server REST API Routes #################
     async def _safe_fire_mcp_tool_call_logging(
-        logging_obj: Any | None,
+        logging_obj: "LiteLLMLoggingObj | None",
         result: "CallToolResult",
         start_time: datetime,
         end_time: datetime,
@@ -158,7 +159,7 @@ if MCP_AVAILABLE:
         data: dict[str, Any],
         tool_name: str,
         user_api_key_dict: UserAPIKeyAuth,
-    ) -> Any:
+    ) -> "CallToolResult":
         """Handle the virtual ``mcp_tool_search`` / ``mcp_tool_call`` REST tools (gated on
         ``mcp_tool_search_enabled``). Kept out of ``call_tool_rest_api`` so that endpoint stays a single
         dispatch. An upstream 401 raised by the virtual ``mcp_tool_call`` propagates unhandled to the
@@ -298,8 +299,8 @@ if MCP_AVAILABLE:
         """
         if not _is_v1_resolved_oauth2_server(server):
             return None
-        user_id: Final = getattr(user_api_key_dict, "user_id", None)
-        server_id: Final = getattr(server, "server_id", None)
+        user_id: Final[str | None] = getattr(user_api_key_dict, "user_id", None)
+        server_id: Final[str | None] = getattr(server, "server_id", None)
         if not user_id or not server_id:
             return None
         try:
@@ -343,7 +344,7 @@ if MCP_AVAILABLE:
         Returns a dict keyed by server_id. Used to avoid N+1 DB queries when
         iterating over multiple OAuth2 MCP servers.
         """
-        user_id: Final = getattr(user_api_key_dict, "user_id", None)
+        user_id: Final[str | None] = getattr(user_api_key_dict, "user_id", None)
         if not user_id:
             return {}
         try:
@@ -664,7 +665,7 @@ if MCP_AVAILABLE:
             "message": "Successfully retrieved tools",
         }
 
-    def _as_query_str(value: Any) -> str | None:
+    def _as_query_str(value: object) -> str | None:
         """Coerce an Optional[str] Query param to str|None, dropping unresolved FastAPI defaults."""
         return value if isinstance(value, str) else None
 
@@ -935,8 +936,8 @@ if MCP_AVAILABLE:
             user_api_key_dict = await acting_user_auth(user_api_key_dict)
             data = await request.json()
 
-            tool_name: Final = data.get("name")
-            tool_arguments: Final = data.get("arguments") or {}
+            tool_name: Final[str | None] = data.get("name")
+            tool_arguments: Final[dict[str, object]] = data.get("arguments") or {}
 
             from litellm.proxy._experimental.mcp_server.tool_search import (
                 MCP_TOOL_CALL_TOOL_NAME,
@@ -947,7 +948,7 @@ if MCP_AVAILABLE:
                 return await _handle_virtual_mcp_tool(request, data, tool_name, user_api_key_dict)
 
             # Validate required parameters early
-            server_id: Final = data.get("server_id")
+            server_id: Final[str | None] = data.get("server_id")
             if not server_id:
                 raise HTTPException(
                     status_code=400,
@@ -1123,11 +1124,11 @@ if MCP_AVAILABLE:
 
     async def _execute_with_mcp_client(
         request: NewMCPServerRequest,
-        operation: Callable[..., Awaitable[Any]],
+        operation: Callable[..., Awaitable[Mapping[str, object]]],
         mcp_auth_header: str | dict[str, str] | None = None,
         oauth2_headers: dict[str, str] | None = None,
         raw_headers: dict[str, str] | None = None,
-    ) -> dict:
+    ) -> Mapping[str, object]:
         """
         Create a temporary MCP client from *request*, run *operation*, and return the result.
 

litellm/proxy/_experimental/mcp_server/server.py

@@ -430,6 +430,7 @@ if MCP_AVAILABLE:
         MCPServerManager,
         _caller_authorization_fans_out,
         _client_forwarded_authorization_headers,
+        _resolve_openapi_tool_auth,
         _should_strip_caller_authorization,
         _without_authorization,
         global_mcp_server_manager,
@@ -2835,58 +2836,26 @@ if MCP_AVAILABLE:
                 arguments = hook_result["arguments"]
 
             verbose_logger.debug("Executing local registry tool: %s", name)
-            # For BYOK servers the credential must be injected via a ContextVar
-            # because the tool function has headers baked into its closure.
-            # Pre-format the full Authorization header value using the server's
-            # configured auth_type so the generator doesn't need to know the prefix.
-            auth_header_value: str | None = None
-            if mcp_auth_header:
-                server_auth_type: Final = getattr(mcp_server, "auth_type", None) if mcp_server else None
-                if server_auth_type == MCPAuth.api_key:
-                    auth_header_value = f"ApiKey {mcp_auth_header}"
-                elif server_auth_type == MCPAuth.basic:
-                    auth_header_value = f"Basic {mcp_auth_header}"
-                else:
-                    auth_header_value = f"Bearer {mcp_auth_header}"
-
-            # Forward named client headers to OpenAPI tool upstream requests.
-            # MCPServer.extra_headers lists header names to copy from raw_headers.
-            # The strip decision is centralized in _should_strip_caller_authorization so this
-            # OpenAPI/local path agrees with the managed paths: M2M and the resolver-owned modes
-            # (token_exchange's raw subject token, authorization_code's stored token) must never
-            # have the caller's Authorization forwarded verbatim upstream.
-            forwarded_headers: dict[str, str] | None = None
-            if mcp_server and mcp_server.extra_headers and raw_headers:
-                normalized_raw: Final = {str(k).lower(): v for k, v in raw_headers.items() if isinstance(k, str)}
-                skip_caller_authorization: Final = _should_strip_caller_authorization(
-                    mcp_server=mcp_server,
-                    raw_headers=raw_headers,
-                    user_api_key_auth=user_api_key_auth,
-                )
-                for header_name in mcp_server.extra_headers:
-                    if not isinstance(header_name, str):
-                        continue
-                    if skip_caller_authorization and header_name.lower() == "authorization":
-                        continue
-                    value = normalized_raw.get(header_name.lower())
-                    if value is not None:
-                        if forwarded_headers is None:
-                            forwarded_headers = {}
-                        forwarded_headers[header_name] = value
-
-            resolved_auth_headers: dict[str, str] | None = None
-            if mcp_server:
-                (
-                    resolved_auth_headers,
-                    forwarded_headers,
-                ) = await global_mcp_server_manager.resolve_openapi_upstream_auth(
-                    mcp_server=mcp_server,
-                    oauth2_headers=oauth2_headers,
-                    raw_headers=raw_headers,
-                    mcp_auth_header=mcp_auth_header,
-                    user_api_key_auth=user_api_key_auth,
-                    forwarded_headers=forwarded_headers,
-                )
+            # The credential rides ContextVars because the tool function has its
+            # headers baked into the closure at registration time.
+            auth_header_value, openapi_forwarded_headers, upstream_credential = _resolve_openapi_tool_auth(
+                mcp_server=mcp_server,
+                mcp_auth_header=mcp_auth_header,
+                mcp_server_auth_headers=mcp_server_auth_headers,
+                raw_headers=raw_headers,
+                user_api_key_auth=user_api_key_auth,
+            )
+            (
+                resolved_auth_headers,
+                forwarded_headers,
+            ) = await global_mcp_server_manager.resolve_openapi_upstream_auth(
+                mcp_server=mcp_server,
+                oauth2_headers=oauth2_headers,
+                raw_headers=raw_headers,
+                mcp_auth_header=upstream_credential,
+                user_api_key_auth=user_api_key_auth,
+                forwarded_headers=openapi_forwarded_headers,
+            )
 
             _auth_token: Final = _request_auth_header.set(auth_header_value)
             _extra_token: Final = _request_extra_headers.set(forwarded_headers)

litellm/proxy/agent_endpoints/a2a_endpoints.py

@@ -168,14 +168,13 @@ def _jsonrpc_error(
     )
 
 
-def _get_agent(agent_id: str):
+async def _get_agent(agent_id: str) -> "AgentResponse | None":
     """Look up an agent by ID or name. Returns None if not found."""
-    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+    from litellm.proxy.common_utils.registry_read_through import (
+        get_agent_with_read_through,
+    )
 
-    agent = global_agent_registry.get_agent_by_id(agent_id=agent_id)
-    if agent is None:
-        agent = global_agent_registry.get_agent_by_name(agent_name=agent_id)
-    return agent
+    return await get_agent_with_read_through(agent_id)
 
 
 def _enforce_inbound_trace_id(agent: "AgentResponse", request: Request) -> None:
@@ -559,7 +558,7 @@ async def get_agent_card(
     )
 
     try:
-        agent: Final = _get_agent(agent_id)
+        agent: Final = await _get_agent(agent_id)
         if agent is None:
             raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found")
 
@@ -673,7 +672,7 @@ async def invoke_agent_a2a(
                 params.pop(key)
 
         # Find the agent
-        agent: Final = _get_agent(agent_id)
+        agent: Final = await _get_agent(agent_id)
         if agent is None:
             return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found", 404)
 

litellm/proxy/agent_endpoints/a2a_routing.py

@@ -25,10 +25,12 @@ async def route_a2a_agent_request(
     Returns None if not an A2A request (allows normal routing to continue).
     """
     # Import here to avoid circular imports
-    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
     from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
         AgentRequestHandler,
     )
+    from litellm.proxy.common_utils.registry_read_through import (
+        get_agent_with_read_through,
+    )
     from litellm.proxy.route_llm_request import (
         ROUTE_ENDPOINT_MAPPING,
         ProxyModelNotFoundError,
@@ -44,11 +46,11 @@ async def route_a2a_agent_request(
     agent_name: Final = model_name[4:]
 
     # Look up agent in registry
-    agent: Final = global_agent_registry.get_agent_by_name(agent_name)
+    agent: Final = await get_agent_with_read_through(agent_name)
     if agent is None:
         verbose_proxy_logger.error("[A2A] Agent '%s' not found in registry", agent_name)
         route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
-        raise ProxyModelNotFoundError(route=route_name, model_name=model_name)
+        raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False)
 
     # Verify the caller is permitted to use this agent (admins bypass the check)
     is_admin: Final = user_api_key_dict is not None and (
@@ -70,7 +72,7 @@ async def route_a2a_agent_request(
     if not agent.agent_card_params or "url" not in agent.agent_card_params:
         verbose_proxy_logger.error("[A2A] Agent '%s' has no URL configured", agent_name)
         route_name = ROUTE_ENDPOINT_MAPPING.get(route_type, route_type)
-        raise ProxyModelNotFoundError(route=route_name, model_name=model_name)
+        raise ProxyModelNotFoundError(route=route_name, model_name=model_name, retryable_with_model_read_through=False)
 
     # Inject API base and route to litellm
     data["api_base"] = agent.agent_card_params["url"]

litellm/proxy/agent_endpoints/agent_registry.py

@@ -600,3 +600,4 @@ class AgentRegistry:
 
 
 global_agent_registry: Final = AgentRegistry()
+AGENT_RECONCILE_LOCK: Final = asyncio.Lock()

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

@@ -1,5 +1,6 @@
 import builtins
 import json
+from collections.abc import Mapping, Sequence
 from datetime import datetime
 from typing import Any, Final, Literal
 
@@ -7,10 +8,30 @@ import click
 import requests
 import rich
 from rich.table import Table
+from typing_extensions import ReadOnly, TypedDict
 
 from ...keys import KeysManagementClient
 
 
+class _CliContext(TypedDict):
+    """Values the top-level CLI group stores on the click context."""
+
+    base_url: ReadOnly[str]
+    api_key: ReadOnly[str | None]
+
+
+class _CliContextView(TypedDict):
+    obj: ReadOnly[_CliContext]
+
+
+class _KeyRowsView(TypedDict):
+    rows: ReadOnly[Sequence[Mapping[str, object]]]
+
+
+class _JsonBodyView(TypedDict):
+    body: ReadOnly[object]
+
+
 @click.group()
 def keys():
     """Manage API keys for the LiteLLM proxy server"""
@@ -53,7 +74,8 @@ def list(
     return_full_object: bool,
 ):
     """List all API keys"""
-    client: Final = KeysManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
+    context: Final[_CliContextView] = {"obj": ctx.obj}
+    client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
     response: Final = client.list(
         page=page,
         size=size,
@@ -70,14 +92,16 @@ def list(
     if output_format == "json":
         rich.print_json(data=response)
     else:
-        rich.print(f"Showing {len(response.get('keys', []))} keys out of {response.get('total_count', 0)}")
+        listed: Final[_KeyRowsView] = {"rows": response.get("keys", [])}
+        rich.print(f"Showing {len(listed['rows'])} keys out of {response.get('total_count', 0)}")
         table: Final = Table(title="API Keys")
         table.add_column("Key Hash", style="cyan")
         table.add_column("Alias", style="green")
         table.add_column("User ID", style="magenta")
         table.add_column("Team ID", style="yellow")
         table.add_column("Spend", style="red")
-        for key in response.get("keys", []):
+        key_rows: Final[_KeyRowsView] = {"rows": response.get("keys", [])}
+        for key in key_rows["rows"]:
             table.add_row(
                 str(key.get("token", "")),
                 str(key.get("key_alias", "")),
@@ -116,7 +140,8 @@ def generate(
     config: str | None,
 ):
     """Generate a new API key"""
-    client: Final = KeysManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
+    context: Final[_CliContextView] = {"obj": ctx.obj}
+    client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
     try:
         models_list: Final = [m.strip() for m in models.split(",")] if models else None
         aliases_dict: Final = json.loads(aliases) if aliases else None
@@ -139,8 +164,8 @@ def generate(
     except requests.exceptions.HTTPError as e:
         click.echo(f"Error: HTTP {e.response.status_code}", err=True)
         try:
-            error_body: Final = e.response.json()
-            rich.print_json(data=error_body)
+            error_body: Final[_JsonBodyView] = {"body": e.response.json()}
+            rich.print_json(data=error_body["body"])
         except json.JSONDecodeError:
             click.echo(e.response.text, err=True)
         raise click.Abort()
@@ -152,7 +177,8 @@ def generate(
 @click.pass_context
 def delete(ctx: click.Context, keys: str | None, key_aliases: str | None):
     """Delete API keys by key or alias"""
-    client: Final = KeysManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
+    context: Final[_CliContextView] = {"obj": ctx.obj}
+    client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
     keys_list: Final = [k.strip() for k in keys.split(",")] if keys else None
     aliases_list: Final = [a.strip() for a in key_aliases.split(",")] if key_aliases else None
     try:
@@ -161,8 +187,8 @@ def delete(ctx: click.Context, keys: str | None, key_aliases: str | None):
     except requests.exceptions.HTTPError as e:
         click.echo(f"Error: HTTP {e.response.status_code}", err=True)
         try:
-            error_body: Final = e.response.json()
-            rich.print_json(data=error_body)
+            error_body: Final[_JsonBodyView] = {"body": e.response.json()}
+            rich.print_json(data=error_body["body"])
         except json.JSONDecodeError:
             click.echo(e.response.text, err=True)
         raise click.Abort()
@@ -189,10 +215,10 @@ def _parse_created_since_filter(created_since: str | None) -> datetime | None:
 
 def _fetch_all_keys_with_pagination(
     source_client: KeysManagementClient, source_base_url: str
-) -> builtins.list[dict[str, Any]]:
+) -> Sequence[Mapping[str, object]]:
     """Fetch all keys from source instance using pagination."""
     click.echo(f"Fetching keys from source server: {source_base_url}")
-    source_keys: Final = []
+    source_keys: Final[builtins.list[Mapping[str, object]]] = []
     page = 1
     page_size: Final = 100  # Use a larger page size to minimize API calls
 
@@ -200,7 +226,7 @@ def _fetch_all_keys_with_pagination(
         source_response = source_client.list(return_full_object=True, page=page, size=page_size)
         # source_client.list() returns Dict[str, Any] when return_request is False (default)
         assert isinstance(source_response, dict), "Expected dict response from list API"
-        page_keys = source_response.get("keys", [])
+        page_keys: Sequence[Mapping[str, object]] = source_response.get("keys", [])
 
         if not page_keys:
             break
@@ -218,15 +244,15 @@ def _fetch_all_keys_with_pagination(
 
 
 def _filter_keys_by_created_since(
-    source_keys: builtins.list[dict[str, Any]],
+    source_keys: Sequence[Mapping[str, object]],
     created_since_dt: datetime | None,
     created_since: str,
-) -> builtins.list[dict[str, Any]]:
+) -> Sequence[Mapping[str, object]]:
     """Filter keys by created_since date if specified."""
     if not created_since_dt:
         return source_keys
 
-    filtered_keys: Final = []
+    filtered_keys: Final[builtins.list[Mapping[str, object]]] = []
     for key in source_keys:
         key_created_at = key.get("created_at")
         if key_created_at:
@@ -248,7 +274,7 @@ def _filter_keys_by_created_since(
     return filtered_keys
 
 
-def _display_dry_run_table(source_keys: builtins.list[dict[str, Any]]) -> None:
+def _display_dry_run_table(source_keys: Sequence[Mapping[str, object]]) -> None:
     """Display a table of keys that would be imported in dry-run mode."""
     click.echo("\n--- DRY RUN MODE ---")
     table: Final = Table(title="Keys that would be imported")
@@ -271,7 +297,7 @@ def _display_dry_run_table(source_keys: builtins.list[dict[str, Any]]) -> None:
     rich.print(table)
 
 
-def _prepare_key_import_data(key: dict[str, Any]) -> dict[str, Any]:
+def _prepare_key_import_data(key: Mapping[str, object]) -> dict[str, Any]:
     """Prepare key data for import by extracting relevant fields."""
     import_data: Final = {}
 
@@ -293,7 +319,7 @@ def _prepare_key_import_data(key: dict[str, Any]) -> dict[str, Any]:
 
 
 def _import_keys_to_destination(
-    source_keys: builtins.list[dict[str, Any]], dest_client: KeysManagementClient
+    source_keys: Sequence[Mapping[str, object]], dest_client: KeysManagementClient
 ) -> tuple[int, int]:
     """Import each key to the destination instance and return counts."""
     imported_count = 0
@@ -351,7 +377,8 @@ def import_keys(
 
     # Create clients for both source and destination
     source_client: Final = KeysManagementClient(source_base_url, source_api_key)
-    dest_client: Final = KeysManagementClient(ctx.obj["base_url"], ctx.obj["api_key"])
+    context: Final[_CliContextView] = {"obj": ctx.obj}
+    dest_client: Final = KeysManagementClient(context["obj"]["base_url"], context["obj"]["api_key"])
 
     try:
         # Get all keys from source instance with pagination
@@ -383,8 +410,8 @@ def import_keys(
     except requests.exceptions.HTTPError as e:
         click.echo(f"Error: HTTP {e.response.status_code}", err=True)
         try:
-            error_body: Final = e.response.json()
-            rich.print_json(data=error_body)
+            error_body: Final[_JsonBodyView] = {"body": e.response.json()}
+            rich.print_json(data=error_body["body"])
         except json.JSONDecodeError:
             click.echo(e.response.text, err=True)
         raise click.Abort()

litellm/proxy/common_utils/registry_read_through.py

@@ -0,0 +1,224 @@
+"""Read-through recovery for in-memory registries in multi-replica deployments.
+
+A management write (POST /model/new, /guardrails, /v1/agents) lands on one
+replica and reaches Postgres, but sibling replicas only refresh their in-memory
+registries on the periodic config reload, so a request using the new object
+immediately can land on a sibling that has never heard of it and fail 400/404.
+On a registry miss, callers here fetch the missing row from the DB and load it
+into the local registry before giving up. A short negative-result TTL per key
+plus a global resync budget per window bound the DB load from lookups of
+genuinely unknown names.
+"""
+
+import asyncio
+import time
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING, Final
+
+from litellm._logging import verbose_proxy_logger
+from litellm.caching.in_memory_cache import InMemoryCache
+
+if TYPE_CHECKING:
+    from prisma.types import (
+        LiteLLM_AgentsTableInclude,
+        LiteLLM_AgentsTableWhereUniqueInput,
+        LiteLLM_GuardrailsTableWhereInput,
+        LiteLLM_ProxyModelTableWhereInput,
+    )
+
+    from litellm.integrations.custom_guardrail import CustomGuardrail
+    from litellm.types.agents import AgentResponse
+
+READ_THROUGH_MISS_TTL_SECONDS: Final = 2.0
+READ_THROUGH_RESYNC_WINDOW_SECONDS: Final = 5.0
+READ_THROUGH_MAX_RESYNCS_PER_WINDOW: Final = 20
+
+
+class RegistryReadThrough:
+    __slots__ = (
+        "_lock",
+        "_max_resyncs_per_window",
+        "_miss_ttl_seconds",
+        "_recent_misses",
+        "_resync",
+        "_resync_window_seconds",
+        "_window_resyncs",
+        "_window_started_at",
+    )
+
+    def __init__(
+        self,
+        resync: Callable[[str], Awaitable[bool]],
+        miss_ttl_seconds: float = READ_THROUGH_MISS_TTL_SECONDS,
+        max_resyncs_per_window: int = READ_THROUGH_MAX_RESYNCS_PER_WINDOW,
+        resync_window_seconds: float = READ_THROUGH_RESYNC_WINDOW_SECONDS,
+    ) -> None:
+        self._resync = resync
+        self._miss_ttl_seconds = miss_ttl_seconds
+        self._max_resyncs_per_window = max_resyncs_per_window
+        self._resync_window_seconds = resync_window_seconds
+        self._lock = asyncio.Lock()
+        self._recent_misses = InMemoryCache(max_size_in_memory=1000)
+        self._window_started_at = float("-inf")
+        self._window_resyncs = 0
+
+    def _consume_resync_budget(self) -> bool:
+        now: Final = time.monotonic()
+        if now - self._window_started_at >= self._resync_window_seconds:
+            self._window_started_at = now
+            self._window_resyncs = 0
+        if self._window_resyncs >= self._max_resyncs_per_window:
+            return False
+        self._window_resyncs += 1
+        return True
+
+    async def attempt(self, key: str) -> bool:
+        if self._recent_misses.get_cache(key) is not None:
+            return False
+        async with self._lock:
+            if self._recent_misses.get_cache(key) is not None:
+                return False
+            if not self._consume_resync_budget():
+                verbose_proxy_logger.warning(
+                    "registry read-through for %r skipped: resync budget of %s per %ss exhausted",
+                    key,
+                    self._max_resyncs_per_window,
+                    self._resync_window_seconds,
+                )
+                return False
+            try:
+                found: Final = await self._resync(key)
+            except Exception as e:  # noqa: BLE001  # a failed read-through must surface the original miss error, not a 500
+                verbose_proxy_logger.warning("registry read-through for %r failed: %s", key, e)
+                return False
+            if not found:
+                self._recent_misses.set_cache(key, True, ttl=self._miss_ttl_seconds)
+            return found
+
+
+def _db_backed_registries_enabled(object_type: str) -> bool:
+    from litellm.proxy import proxy_server
+
+    if proxy_server.prisma_client is None or proxy_server.store_model_in_db is not True:
+        return False
+    return proxy_server.should_load_db_object(object_type=object_type)
+
+
+async def _resync_model_deployments(model_name: str) -> bool:
+    from litellm.proxy import proxy_server
+    from litellm.repositories.model_repository import ModelRepository
+
+    if not _db_backed_registries_enabled("models"):
+        return False
+    prisma_client: Final = proxy_server.prisma_client
+    assert prisma_client is not None
+    table: Final = ModelRepository(prisma_client).table
+    name_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_name": model_name}
+    id_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_id": model_name}
+    rows: Final = await table.find_many(where=name_filter) or await table.find_many(where=id_filter)
+    if not rows:
+        return False
+    router: Final = proxy_server.llm_router
+    if router is None:
+        await proxy_server.proxy_config.add_deployment(
+            prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj
+        )
+        return proxy_server.llm_router is not None
+    async with proxy_server.MODEL_RECONCILE_LOCK:
+        proxy_server.proxy_config._add_deployment(db_models=rows)
+        proxy_server.llm_model_list = router.get_model_list()
+    return True
+
+
+async def _resync_guardrails(guardrail_name: str) -> bool:
+    from litellm.proxy import proxy_server
+    from litellm.proxy.guardrails.guardrail_registry import (
+        GUARDRAIL_RECONCILE_LOCK,
+        IN_MEMORY_GUARDRAIL_HANDLER,
+    )
+    from litellm.repositories.table_repositories import GuardrailsRepository
+    from litellm.types.guardrails import Guardrail
+
+    if not _db_backed_registries_enabled("guardrails"):
+        return False
+    prisma_client: Final = proxy_server.prisma_client
+    assert prisma_client is not None
+    active_row_filter: Final[LiteLLM_GuardrailsTableWhereInput] = {
+        "guardrail_name": guardrail_name,
+        "status": "active",
+    }
+    row: Final = await GuardrailsRepository(prisma_client).table.find_first(where=active_row_filter)
+    if row is None:
+        return False
+    async with GUARDRAIL_RECONCILE_LOCK:
+        IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(guardrail=Guardrail(**dict(row)))
+    return _initialized_guardrail(guardrail_name) is not None
+
+
+async def _resync_agents(agent_id_or_name: str) -> bool:
+    from litellm.proxy import proxy_server
+    from litellm.proxy.agent_endpoints.agent_registry import (
+        AGENT_RECONCILE_LOCK,
+        agents_table,
+        global_agent_registry,
+    )
+    from litellm.types.agents import AgentResponse
+
+    if not _db_backed_registries_enabled("agents"):
+        return False
+    if _agent_from_registry(agent_id_or_name) is not None:
+        return True
+    prisma_client: Final = proxy_server.prisma_client
+    assert prisma_client is not None
+    table: Final = agents_table(prisma_client)
+    id_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id_or_name}
+    name_filter: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_name": agent_id_or_name}
+    include_permission: Final[LiteLLM_AgentsTableInclude] = {"object_permission": True}
+    async with AGENT_RECONCILE_LOCK:
+        if _agent_from_registry(agent_id_or_name) is not None:
+            return True
+        row: Final = await table.find_unique(where=id_filter, include=include_permission) or await table.find_unique(
+            where=name_filter, include=include_permission
+        )
+        if row is None:
+            return False
+        global_agent_registry.register_agent(agent_config=AgentResponse.model_validate(row.model_dump()))
+        return True
+
+
+model_registry_read_through: Final = RegistryReadThrough(resync=_resync_model_deployments)
+guardrail_registry_read_through: Final = RegistryReadThrough(resync=_resync_guardrails)
+agent_registry_read_through: Final = RegistryReadThrough(resync=_resync_agents)
+
+
+def _agent_from_registry(agent_id_or_name: str) -> "AgentResponse | None":
+    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+
+    by_id: Final = global_agent_registry.get_agent_by_id(agent_id=agent_id_or_name)
+    if by_id is not None:
+        return by_id
+    return global_agent_registry.get_agent_by_name(agent_name=agent_id_or_name)
+
+
+async def get_agent_with_read_through(agent_id_or_name: str) -> "AgentResponse | None":
+    agent: Final = _agent_from_registry(agent_id_or_name)
+    if agent is not None:
+        return agent
+    if not await agent_registry_read_through.attempt(agent_id_or_name):
+        return None
+    return _agent_from_registry(agent_id_or_name)
+
+
+def _initialized_guardrail(guardrail_name: str) -> "CustomGuardrail | None":
+    from litellm.proxy.guardrails import guardrail_endpoints
+
+    return guardrail_endpoints.GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(guardrail_name=guardrail_name)
+
+
+async def get_initialized_guardrail_with_read_through(guardrail_name: str) -> "CustomGuardrail | None":
+    active: Final = _initialized_guardrail(guardrail_name)
+    if active is not None:
+        return active
+    if not await guardrail_registry_read_through.attempt(guardrail_name):
+        return None
+    return _initialized_guardrail(guardrail_name)

litellm/proxy/guardrails/guardrail_endpoints.py

@@ -2305,8 +2305,12 @@ async def apply_guardrail(
     litellm_logging_obj = None
     start_time: Final = datetime.now(timezone.utc)
 
+    from litellm.proxy.common_utils.registry_read_through import (
+        get_initialized_guardrail_with_read_through,
+    )
+
     try:
-        active_guardrail: Final[CustomGuardrail | None] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback(
+        active_guardrail: Final[CustomGuardrail | None] = await get_initialized_guardrail_with_read_through(
             guardrail_name=request.guardrail_name
         )
         if active_guardrail is None:

litellm/proxy/guardrails/guardrail_hooks/aim/aim.py

@@ -7,10 +7,11 @@
 import asyncio
 import json
 import os
-from collections.abc import AsyncGenerator
-from typing import TYPE_CHECKING, Any, Final
+from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence
+from typing import TYPE_CHECKING, Final, TypeAlias
 
 from pydantic import BaseModel
+from typing_extensions import NotRequired, ReadOnly, TypedDict
 from websockets.asyncio.client import ClientConnection, connect
 
 from litellm import DualCache
@@ -31,8 +32,7 @@ from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import (
     CallTypesLiteral,
     Choices,
-    EmbeddingResponse,
-    ImageResponse,
+    LLMResponseTypes,
     ModelResponse,
     ModelResponseStream,
 )
@@ -45,6 +45,58 @@ class AimGuardrailMissingSecrets(Exception):
     pass
 
 
+class AimRequiredAction(TypedDict):
+    """The ``required_action`` block of an Aim ``/fw/v1/analyze`` response."""
+
+    action_type: ReadOnly[NotRequired[str]]
+    detection_message: ReadOnly[str]
+
+
+class AimAnalysisResult(TypedDict):
+    """The ``analysis_result`` block of an Aim ``/fw/v1/analyze`` response."""
+
+    policy_drill_down: ReadOnly[Mapping[str, object]]
+
+
+class AimRedactedMessage(TypedDict):
+    """One entry of Aim's ``redacted_chat.all_redacted_messages``."""
+
+    role: ReadOnly[str]
+    content: ReadOnly[str]
+
+
+class AimRedactedChat(TypedDict):
+    """The ``redacted_chat`` block of an Aim ``/fw/v1/analyze`` response."""
+
+    all_redacted_messages: ReadOnly[Sequence[AimRedactedMessage]]
+
+
+class AimAnalyzeResponse(TypedDict):
+    """Body returned by Aim's ``POST /fw/v1/analyze``."""
+
+    required_action: ReadOnly[AimRequiredAction]
+    analysis_result: ReadOnly[AimAnalysisResult]
+    redacted_chat: ReadOnly[NotRequired[AimRedactedChat]]
+
+
+class AimOutputGuardrailResult(TypedDict, total=False):
+    """Outcome of inspecting one model completion with Aim."""
+
+    detection_message: ReadOnly[str]
+    redacted_output: ReadOnly[str]
+
+
+class AimStreamMessage(TypedDict, total=False):
+    """One frame of Aim's ``/fw/v1/analyze/stream`` websocket protocol."""
+
+    verified_chunk: ReadOnly[Mapping[str, object]]
+    done: ReadOnly[bool]
+    blocking_message: ReadOnly[str]
+
+
+AimStreamChunk: TypeAlias = BaseModel | Mapping[str, object] | str | bytes
+
+
 class AimGuardrail(CustomGuardrail):
     @classmethod
     def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@@ -110,7 +162,7 @@ class AimGuardrail(CustomGuardrail):
             json={"messages": self._build_aim_inspection_messages(data)},
         )
         response.raise_for_status()
-        res: Final = response.json()
+        res: Final[AimAnalyzeResponse] = response.json()
         required_action: Final = res.get("required_action")
         action_type: Final = required_action and required_action.get("action_type", None)
         if action_type is None:
@@ -145,7 +197,7 @@ class AimGuardrail(CustomGuardrail):
             openai_code=openai_code,
         )
 
-    def _handle_block_action(self, analysis_result: Any, required_action: Any) -> None:
+    def _handle_block_action(self, analysis_result: AimAnalysisResult, required_action: AimRequiredAction) -> None:
         detection_message: Final = required_action.get("detection_message", None)
         verbose_proxy_logger.info(
             "Aim: Violation detected enabled policies: {policies}".format(
@@ -154,7 +206,7 @@ class AimGuardrail(CustomGuardrail):
         )
         raise self._rejection(detection_message, openai_code="content_policy_violation")
 
-    def _anonymize_request(self, res: Any, data: dict) -> dict:
+    def _anonymize_request(self, res: AimAnalyzeResponse, data: dict) -> dict:
         verbose_proxy_logger.info("Aim: anonymize action")
         redacted_chat: Final = res.get("redacted_chat")
         if not redacted_chat:
@@ -185,7 +237,7 @@ class AimGuardrail(CustomGuardrail):
 
     async def call_aim_guardrail_on_output(
         self, request_data: dict, output: str, hook: str, key_alias: str | None
-    ) -> dict | None:
+    ) -> AimOutputGuardrailResult | None:
         user_email: Final = request_data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
         call_id: Final = request_data.get("litellm_call_id")
         response: Final = await self.async_handler.post(
@@ -202,7 +254,7 @@ class AimGuardrail(CustomGuardrail):
             },
         )
         response.raise_for_status()
-        res: Final = response.json()
+        res: Final[AimAnalyzeResponse] = response.json()
         required_action: Final = res.get("required_action")
         action_type: Final = required_action and required_action.get("action_type", None)
         if action_type and action_type == "block_action":
@@ -213,7 +265,9 @@ class AimGuardrail(CustomGuardrail):
             return {"redacted_output": redacted_chat["all_redacted_messages"][-1]["content"]}
         return {"redacted_output": output}
 
-    def _handle_block_action_on_output(self, analysis_result: Any, required_action: Any) -> dict | None:
+    def _handle_block_action_on_output(
+        self, analysis_result: AimAnalysisResult, required_action: AimRequiredAction
+    ) -> AimOutputGuardrailResult | None:
         detection_message: Final = required_action.get("detection_message", None)
         verbose_proxy_logger.info(
             "Aim: detected: {detected}, enabled policies: {policies}".format(
@@ -260,8 +314,8 @@ class AimGuardrail(CustomGuardrail):
         self,
         data: dict,
         user_api_key_dict: UserAPIKeyAuth,
-        response: Any | ModelResponse | EmbeddingResponse | ImageResponse,
-    ) -> Any:
+        response: LLMResponseTypes,
+    ) -> LLMResponseTypes:
         if not (isinstance(response, ModelResponse) and response.choices):
             return response
         # Inspect every choice — when ``n>1`` the additional completions
@@ -289,9 +343,11 @@ class AimGuardrail(CustomGuardrail):
         for choice, aim_output_guardrail_result in zip(choices_to_inspect, results):
             if isinstance(aim_output_guardrail_result, BaseException):
                 raise aim_output_guardrail_result
-            if aim_output_guardrail_result and aim_output_guardrail_result.get("detection_message"):
+            if aim_output_guardrail_result and (
+                detection_message := aim_output_guardrail_result.get("detection_message")
+            ):
                 raise self._rejection(
-                    aim_output_guardrail_result.get("detection_message"),
+                    detection_message,
                     openai_code="content_policy_violation",
                 )
             if aim_output_guardrail_result and aim_output_guardrail_result.get("redacted_output"):
@@ -301,7 +357,7 @@ class AimGuardrail(CustomGuardrail):
     async def async_post_call_streaming_iterator_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
-        response,
+        response: AsyncIterator[AimStreamChunk],
         request_data: dict,
     ) -> AsyncGenerator[ModelResponseStream, None]:
         user_email: Final = request_data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
@@ -317,7 +373,7 @@ class AimGuardrail(CustomGuardrail):
         ) as websocket:
             sender: Final = asyncio.create_task(self.forward_the_stream_to_aim(websocket, response))
             while True:
-                result = json.loads(await websocket.recv())
+                result: AimStreamMessage = json.loads(await websocket.recv())
                 if verified_chunk := result.get("verified_chunk"):
                     yield ModelResponseStream.model_validate(verified_chunk)
                 else:
@@ -334,7 +390,7 @@ class AimGuardrail(CustomGuardrail):
     async def forward_the_stream_to_aim(
         self,
         websocket: ClientConnection,
-        response_iter,
+        response_iter: AsyncIterator[AimStreamChunk],
     ) -> None:
         async for chunk in response_iter:
             if isinstance(chunk, BaseModel):

litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py

@@ -7,10 +7,13 @@ and provide safe, sandboxed functionality for common guardrail operations.
 
 import json
 import re
+from collections.abc import Mapping, Sequence
 from typing import Any, Final
 from urllib.parse import urlparse
 
 import httpx
+from pydantic import JsonValue
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
 from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
@@ -21,7 +24,7 @@ from litellm.types.llms.custom_http import httpxSpecialProvider
 # =============================================================================
 
 
-def allow() -> dict[str, Any]:
+def allow() -> dict[str, object]:
     """
     Allow the request/response to proceed unchanged.
 
@@ -31,7 +34,7 @@ def allow() -> dict[str, Any]:
     return {"action": "allow"}
 
 
-def block(reason: str, detection_info: dict[str, Any] | None = None) -> dict[str, Any]:
+def block(reason: str, detection_info: Mapping[str, object] | None = None) -> dict[str, object]:
     """
     Block the request/response with a reason.
 
@@ -42,17 +45,17 @@ def block(reason: str, detection_info: dict[str, Any] | None = None) -> dict[str
     Returns:
         Dict indicating the request should be blocked
     """
-    result: Final[dict[str, Any]] = {"action": "block", "reason": reason}
+    result: Final[dict[str, object]] = {"action": "block", "reason": reason}
     if detection_info:
         result["detection_info"] = detection_info
     return result
 
 
 def modify(
-    texts: list[str] | None = None,
-    images: list[Any] | None = None,
-    tool_calls: list[Any] | None = None,
-) -> dict[str, Any]:
+    texts: Sequence[str] | None = None,
+    images: Sequence[object] | None = None,
+    tool_calls: Sequence[object] | None = None,
+) -> dict[str, object]:
     """
     Modify the request/response content.
 
@@ -64,7 +67,7 @@ def modify(
     Returns:
         Dict indicating the content should be modified
     """
-    result: Final[dict[str, Any]] = {"action": "modify"}
+    result: Final[dict[str, object]] = {"action": "modify"}
     if texts is not None:
         result["texts"] = texts
     if images is not None:
@@ -161,7 +164,15 @@ def regex_find_all(text: str, pattern: str, flags: int = 0) -> list[str]:
 # =============================================================================
 
 
-def json_parse(text: str) -> Any | None:
+class JsonSchemaNode(TypedDict, total=False):
+    """Subset of JSON Schema keywords understood by the built-in validator."""
+
+    type: ReadOnly[str]
+    required: ReadOnly[Sequence[str]]
+    properties: ReadOnly[Mapping[str, "JsonSchemaNode"]]
+
+
+def json_parse(text: str) -> JsonValue:
     """
     Parse a JSON string into a Python object.
 
@@ -178,7 +189,7 @@ def json_parse(text: str) -> Any | None:
         return None
 
 
-def json_stringify(obj: Any) -> str:
+def json_stringify(obj: object) -> str:
     """
     Convert a Python object to a JSON string.
 
@@ -195,7 +206,7 @@ def json_stringify(obj: Any) -> str:
         return ""
 
 
-def json_schema_valid(obj: Any, schema: dict[str, Any]) -> bool:
+def json_schema_valid(obj: JsonValue, schema: JsonSchemaNode) -> bool:
     """
     Validate an object against a JSON schema.
 
@@ -226,7 +237,7 @@ def json_schema_valid(obj: Any, schema: dict[str, Any]) -> bool:
         return False
 
 
-def _basic_json_schema_validate(obj: Any, schema: dict[str, Any], max_depth: int = 50) -> bool:
+def _basic_json_schema_validate(obj: JsonValue, schema: JsonSchemaNode, max_depth: int = 50) -> bool:
     """
     Basic JSON schema validation without external library.
     Handles: type, required, properties
@@ -234,7 +245,7 @@ def _basic_json_schema_validate(obj: Any, schema: dict[str, Any], max_depth: int
     Uses an iterative approach with a stack to avoid recursion limits.
     max_depth limits nesting to prevent infinite loops from circular schemas.
     """
-    type_map: Final[dict[str, type | tuple[type, ...]]] = {
+    type_map: Final[Mapping[str, type | tuple[type, ...]]] = {
         "object": dict,
         "array": list,
         "string": str,
@@ -245,7 +256,7 @@ def _basic_json_schema_validate(obj: Any, schema: dict[str, Any], max_depth: int
     }
 
     # Stack of (obj, schema, depth) tuples to process
-    stack: Final[list[tuple[Any, dict[str, Any], int]]] = [(obj, schema, 0)]
+    stack: Final[list[tuple[JsonValue, JsonSchemaNode, int]]] = [(obj, schema, 0)]
 
     while stack:
         current_obj, current_schema, depth = stack.pop()
@@ -257,19 +268,19 @@ def _basic_json_schema_validate(obj: Any, schema: dict[str, Any], max_depth: int
         # Check type
         schema_type = current_schema.get("type")
         if schema_type:
-            expected_type = type_map.get(schema_type)
+            expected_type: type | tuple[type, ...] | None = type_map.get(schema_type)
             if expected_type is not None and not isinstance(current_obj, expected_type):
                 return False
 
         # Check required fields and properties for dicts
         if isinstance(current_obj, dict):
-            required = current_schema.get("required", [])
+            required: Sequence[str] = current_schema.get("required", [])
             for field in required:
                 if field not in current_obj:
                     return False
 
             # Queue property validations
-            properties = current_schema.get("properties", {})
+            properties: Mapping[str, JsonSchemaNode] = current_schema.get("properties", {})
             for prop_name, prop_schema in properties.items():
                 if prop_name in current_obj:
                     stack.append((current_obj[prop_name], prop_schema, depth + 1))
@@ -358,7 +369,17 @@ _HTTP_DEFAULT_TIMEOUT: Final = 30.0
 _HTTP_MAX_TIMEOUT: Final = 60.0
 
 
-def _http_error_response(error: str) -> dict[str, Any]:
+class HttpResponseResult(TypedDict):
+    """Outcome of an HTTP primitive call, as handed back to custom code."""
+
+    status_code: ReadOnly[int]
+    body: ReadOnly[JsonValue]
+    headers: ReadOnly[Mapping[str, str]]
+    success: ReadOnly[bool]
+    error: ReadOnly[str | None]
+
+
+def _http_error_response(error: str) -> HttpResponseResult:
     """Create a standardized error response for HTTP requests."""
     return {
         "status_code": 0,
@@ -369,9 +390,9 @@ def _http_error_response(error: str) -> dict[str, Any]:
     }
 
 
-def _http_success_response(response: httpx.Response) -> dict[str, Any]:
+def _http_success_response(response: httpx.Response) -> HttpResponseResult:
     """Create a standardized success response from an httpx Response."""
-    parsed_body: Any
+    parsed_body: JsonValue
     try:
         parsed_body = response.json()
     except (json.JSONDecodeError, ValueError):
@@ -387,8 +408,8 @@ def _http_success_response(response: httpx.Response) -> dict[str, Any]:
 
 
 def _prepare_http_body(
-    body: Any | None,
-) -> tuple[dict[str, Any] | None, str | None]:
+    body: JsonValue,
+) -> tuple[dict[str, JsonValue] | None, str | None]:
     """Prepare body arguments for HTTP request - returns (json_body, data_body)."""
     if body is None:
         return None, None
@@ -405,9 +426,9 @@ async def http_request(
     url: str,
     method: str = "GET",
     headers: dict[str, str] | None = None,
-    body: Any | None = None,
+    body: JsonValue = None,
     timeout: float | None = None,
-) -> dict[str, Any]:
+) -> HttpResponseResult:
     """
     Make an async HTTP request to an external service.
 
@@ -491,7 +512,7 @@ async def _execute_http_request(
     method: str,
     url: str,
     headers: dict[str, str] | None,
-    body: Any | None,
+    body: JsonValue,
     timeout: float,
 ) -> httpx.Response:
     """Execute the HTTP request using the appropriate client method."""
@@ -515,7 +536,7 @@ async def http_get(
     url: str,
     headers: dict[str, str] | None = None,
     timeout: float | None = None,
-) -> dict[str, Any]:
+) -> HttpResponseResult:
     """
     Make an async HTTP GET request.
 
@@ -534,10 +555,10 @@ async def http_get(
 
 async def http_post(
     url: str,
-    body: Any | None = None,
+    body: JsonValue = None,
     headers: dict[str, str] | None = None,
     timeout: float | None = None,
-) -> dict[str, Any]:
+) -> HttpResponseResult:
     """
     Make an async HTTP POST request.
 
@@ -755,7 +776,7 @@ def trim(text: str) -> str:
 # =============================================================================
 
 
-def get_custom_code_primitives() -> dict[str, Any]:
+def get_custom_code_primitives() -> dict[str, object]:
     """
     Get all primitives to inject into the custom code environment.
 

litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import os
 from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict
+from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict
 from urllib.parse import urlparse
 from uuid import uuid4
 
@@ -76,6 +76,36 @@ class _HiddenlayerChoice(TypedDict, total=False):
     message: ReadOnly[_HiddenlayerChoiceMessage]
 
 
+class _HiddenlayerV2Output(TypedDict, total=False):
+    messages: ReadOnly[Sequence[_HiddenlayerOutputMessage]]
+    choices: ReadOnly[Sequence[_HiddenlayerChoice]]
+
+
+class _LoggedCallDetails(Protocol):
+    """Logging object view that exposes its untyped call details with the shape this guardrail reads."""
+
+    @property
+    def model_call_details(self) -> Mapping[str, _LoggedCallLitellmParams]: ...
+
+
+class _TokenPayloadSource(Protocol):
+    """Response view that decodes the HiddenLayer OAuth token body as a string mapping."""
+
+    def json(self) -> Mapping[str, str]: ...
+
+
+def _logged_request_headers(logging_obj: _LoggedCallDetails) -> Mapping[str, str]:
+    return logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {})
+
+
+def _token_payload(response: _TokenPayloadSource) -> Mapping[str, str]:
+    return response.json()
+
+
+def _header_value(headers: Mapping[str, str], key: str, default: str) -> str:
+    return headers.get(key, default)
+
+
 def is_saas(host: str) -> bool:
     """Checks whether the connection is to the SaaS platform"""
 
@@ -102,7 +132,7 @@ def _get_jwt(auth_url, api_id, api_key) -> str:
             f"Unable to get authentication credentials for the HiddenLayer API - invalid response: {resp.json()}"
         )
 
-    return resp.json()["access_token"]
+    return _token_payload(resp)["access_token"]
 
 
 class HiddenlayerGuardrail(CustomGuardrail):
@@ -176,10 +206,7 @@ class HiddenlayerGuardrail(CustomGuardrail):
         # from the logger object on the response from the model.
         headers = request_data.get("proxy_server_request", {}).get("headers", {})
         if not headers and logging_obj and logging_obj.model_call_details:
-            logged_litellm_params: Final[_LoggedCallLitellmParams] = logging_obj.model_call_details.get(
-                "litellm_params", {}
-            )
-            headers = logged_litellm_params.get("metadata", {}).get("headers", {})
+            headers = _logged_request_headers(logging_obj)
 
         hl_request_metadata["requester_id"] = headers.get("hl-requester-id") or "LiteLLM"
         project_id: Final = headers.get("hl-project-id")
@@ -418,8 +445,9 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
 
         response: Final = await self._call_hiddenlayer(payload, input_type, hl_headers)
         output: Final = response.json()
+        evaluated_output: Final[_HiddenlayerV2Output] = output
 
-        if response.headers.get("hl-runtime-action", "").lower() == "block":
+        if _header_value(response.headers, "hl-runtime-action", "").lower() == "block":
             raise HTTPException(
                 status_code=400,
                 detail={
@@ -432,7 +460,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
         if input_type == "request":
             inputs["structured_messages"] = output
 
-            modified_messages: Final[Sequence[_HiddenlayerOutputMessage]] = output.get("messages", [])
+            modified_messages: Final[Sequence[_HiddenlayerOutputMessage]] = evaluated_output.get("messages", [])
             for message in modified_messages:
                 content = message.get("content", "")
                 if isinstance(content, list):
@@ -447,7 +475,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail):
             inputs["texts"] = new_texts
 
         elif input_type == "response" and inputs.get("texts"):
-            redacted_choices: Final[Sequence[_HiddenlayerChoice]] = output.get("choices", [{}])
+            redacted_choices: Final[Sequence[_HiddenlayerChoice]] = evaluated_output.get("choices", [{}])
             inputs["texts"] = [redacted_choices[-1].get("message", {}).get("content", "")]
         elif input_type == "response" and inputs.get("tool_calls"):
             inputs["tool_calls"] = output

litellm/proxy/guardrails/guardrail_hooks/microsoft_purview/purview_dlp.py

@@ -10,9 +10,9 @@ Supports three modes:
 import asyncio
 import threading
 import uuid
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, Final, Union, cast
+from typing import TYPE_CHECKING, Any, Final, cast
 
 import httpx
 from fastapi import HTTPException
@@ -36,14 +36,15 @@ from litellm.types.utils import (
 from .base import PurviewGuardrailBase
 
 if TYPE_CHECKING:
+    from litellm.caching.dual_cache import DualCache
     from litellm.proxy._types import UserAPIKeyAuth
+    from litellm.types.llms.openai import AllMessageValues
     from litellm.types.proxy.guardrails.guardrail_hooks.base import (
         GuardrailConfigModel,
     )
     from litellm.types.utils import (
         CallTypesLiteral,
-        EmbeddingResponse,
-        ImageResponse,
+        LLMResponseTypes,
     )
 
 
@@ -63,7 +64,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
         client_secret: str,
         purview_app_name: str = "LiteLLM",
         user_id_field: str = "user_id",
-        **kwargs: Any,
+        **kwargs: object,
     ):
         super().__init__(
             tenant_id=tenant_id,
@@ -104,7 +105,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
         activity: str,
         request_data: dict[str, Any],
         block_on_violation: bool = True,
-    ) -> dict[str, Any]:
+    ) -> dict[str, object]:
         """Evaluate content against Purview DLP policies.
 
         Args:
@@ -119,7 +120,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
         """
         start_time: Final = datetime.now()
         status: GuardrailStatus = "success"
-        response: dict[str, Any] = {}
+        response: dict[str, object] = {}
 
         try:
             etag, _ = await self._compute_protection_scopes(user_id)
@@ -149,7 +150,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
                 upstream_status: Final = exc.response.status_code
                 client_status: Final = 502 if upstream_status in (401, 403) else upstream_status
                 headers: dict[str, str] | None = None
-                retry_after: Final = exc.response.headers.get("retry-after")
+                retry_after: Final[str | None] = exc.response.headers.get("retry-after")
                 if retry_after:
                     headers = {"Retry-After": retry_after}
                 raise HTTPException(
@@ -205,7 +206,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
         return response
 
     @staticmethod
-    def _extract_responses_api_function_call_args(result: Any) -> list[str]:
+    def _extract_responses_api_function_call_args(result: object) -> list[str]:
         """Return tool-call argument strings from a ``ResponsesAPIResponse.output``.
 
         ``ResponsesAPIResponse.output_text`` only aggregates ``output_text``
@@ -215,7 +216,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
         chat (``ModelResponse``) path.
         """
         args: Final[list[str]] = []
-        output: Final = getattr(result, "output", None)
+        output: Final[Sequence[object] | None] = getattr(result, "output", None)
         if not output:
             return args
         for item in output:
@@ -230,7 +231,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
                     args.append(arguments)
         return args
 
-    def _completion_response_text_parts(self, result: Any) -> list[str]:
+    def _completion_response_text_parts(self, result: object) -> list[str]:
         """Collect non-empty text segments from chat, text completions, or responses API.
 
         Includes assistant message content *and* model-generated tool-call
@@ -266,7 +267,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
                 parts.extend(self._extract_tool_call_args_from_message(msg))
         return parts
 
-    def _assemble_responses_api_from_chunks(self, chunks: list[Any]) -> tuple[bool, ResponsesAPIResponse | None]:
+    def _assemble_responses_api_from_chunks(self, chunks: Sequence[object]) -> tuple[bool, ResponsesAPIResponse | None]:
         """Extract the final ``ResponsesAPIResponse`` from a buffered Responses API stream.
 
         Returns a ``(is_responses_api_stream, assembled)`` tuple so the caller
@@ -314,7 +315,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
                 input=input_data if input_data is not None else "",
                 responses_api_request=data,
             )
-            return self.get_prompt_text_for_dlp(cast(list[Any], messages))
+            return self.get_prompt_text_for_dlp(cast(list["AllMessageValues"], messages))
         except Exception:
             verbose_proxy_logger.warning(
                 "Purview DLP: failed to transform responses API input",
@@ -338,8 +339,8 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
 
     def _resolve_user_id_for_blocking(
         self,
-        data: dict[str, Any],
-        user_api_key_dict: Any,
+        data: Mapping[str, object],
+        user_api_key_dict: "UserAPIKeyAuth",
     ) -> str:
         """Resolve user ID for blocking (pre_call / post_call) DLP hooks.
 
@@ -386,10 +387,10 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
     async def async_pre_call_hook(
         self,
         user_api_key_dict: "UserAPIKeyAuth",
-        cache: Any,
+        cache: "DualCache",
         data: dict[str, Any],
         call_type: "CallTypesLiteral",
-    ) -> dict[str, Any] | None:
+    ) -> dict[str, object] | None:
         """Check user prompt against Purview DLP policies before LLM call."""
         user_id: Final = self._resolve_user_id_for_blocking(data, user_api_key_dict)
 
@@ -423,7 +424,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
         else:
             messages: Final[list | None] = data.get("messages")
             if messages:
-                prompt_text = self.get_prompt_text_for_dlp(cast(list[Any], messages))
+                prompt_text = self.get_prompt_text_for_dlp(cast(list["AllMessageValues"], messages))
 
         if not prompt_text:
             return data
@@ -446,8 +447,8 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
         self,
         data: dict,
         user_api_key_dict: "UserAPIKeyAuth",
-        response: Union[Any, ModelResponse, "EmbeddingResponse", "ImageResponse"],
-    ) -> Any:
+        response: "LLMResponseTypes",
+    ) -> "LLMResponseTypes":
         """Check LLM response against Purview DLP policies (non-streaming only).
 
         Streaming responses are handled by ``async_post_call_streaming_iterator_hook``
@@ -472,7 +473,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
     async def async_post_call_streaming_iterator_hook(
         self,
         user_api_key_dict: "UserAPIKeyAuth",
-        response: Any,
+        response: AsyncIterable[ModelResponseStream],
         request_data: dict,
     ) -> AsyncGenerator[ModelResponseStream, None]:
         """Check streaming LLM responses against Purview DLP policies.
@@ -592,7 +593,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
     # Logging-only hook — audit without blocking
     # ------------------------------------------------------------------
 
-    def logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]:
+    def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]:
         """Fire-and-forget async audit logging; returns original (kwargs, result) immediately.
 
         In the proxy's async success path, litellm independently calls both
@@ -640,7 +641,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
 
         return kwargs, result
 
-    async def async_logging_hook(self, kwargs: dict, result: Any, call_type: str) -> tuple[dict, Any]:
+    async def async_logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]:
         """Send both prompt and response to Purview for audit logging.
 
         Errors are logged but never raised — this mode is non-blocking.
@@ -670,7 +671,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
             else:
                 messages: Final = kwargs.get("messages")
                 if messages:
-                    prompt_text = self.get_prompt_text_for_dlp(cast(list[Any], messages))
+                    prompt_text = self.get_prompt_text_for_dlp(cast(list["AllMessageValues"], messages))
 
             if prompt_text:
                 await self._check_content(

litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py

@@ -1,9 +1,11 @@
 import asyncio
 import base64
 import os
-from typing import TYPE_CHECKING, Any, Final, Literal, Optional
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Final, Literal, Optional
 
 from fastapi import HTTPException
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
 from litellm.integrations.custom_guardrail import (
@@ -26,6 +28,41 @@ class PromptSecurityGuardrailMissingSecrets(Exception):
     pass
 
 
+class _ProtectVerdict(TypedDict, total=False):
+    """One side (``prompt`` or ``response``) of an ``/api/protect`` verdict."""
+
+    action: ReadOnly[str]
+    violations: ReadOnly[Sequence[str]]
+    modified_messages: ReadOnly[Sequence[Mapping[str, object]]]
+    modified_text: ReadOnly[str]
+
+
+class _ProtectResult(TypedDict, total=False):
+    prompt: ReadOnly[_ProtectVerdict | None]
+    response: ReadOnly[_ProtectVerdict | None]
+
+
+class _ProtectResponse(TypedDict, total=False):
+    result: ReadOnly[_ProtectResult]
+
+
+class _SanitizeUploadResponse(TypedDict, total=False):
+    jobId: ReadOnly[str]
+
+
+class _SanitizeMetadata(TypedDict, total=False):
+    action: ReadOnly[str]
+    violations: ReadOnly[Sequence[str]]
+
+
+class _SanitizeStatusResponse(TypedDict, total=False):
+    """One poll of ``/api/sanitizeFile``."""
+
+    status: ReadOnly[str]
+    content: ReadOnly[str]
+    metadata: ReadOnly[_SanitizeMetadata]
+
+
 class PromptSecurityGuardrail(CustomGuardrail):
     @classmethod
     def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@@ -199,7 +236,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
             json=payload,
         )
         response.raise_for_status()
-        res: Final = response.json()
+        res: Final[_ProtectResponse] = response.json()
 
         self._log_api_response(
             url=f"{self.api_base}/api/protect",
@@ -261,7 +298,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
             json=payload,
         )
         response.raise_for_status()
-        res: Final = response.json()
+        res: Final[_ProtectResponse] = response.json()
 
         self._log_api_response(
             url=f"{self.api_base}/api/protect",
@@ -290,7 +327,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
 
         return inputs
 
-    def _extract_texts_from_messages(self, messages: list) -> list[str]:
+    def _extract_texts_from_messages(self, messages: Sequence[Mapping[str, object]]) -> list[str]:
         """Extract text content from messages."""
         texts: Final = []
         for message in messages:
@@ -379,7 +416,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
             files=files,
         )
         upload_response.raise_for_status()
-        upload_result: Final = upload_response.json()
+        upload_result: Final[_SanitizeUploadResponse] = upload_response.json()
         job_id: Final = upload_result.get("jobId")
 
         self._log_api_response(
@@ -409,7 +446,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
                 params={"jobId": job_id},
             )
             poll_response.raise_for_status()
-            result = poll_response.json()
+            result: _SanitizeStatusResponse = poll_response.json()
 
             self._log_api_response(
                 url=f"{self.api_base}/api/sanitizeFile",
@@ -656,7 +693,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
         method: str,
         url: str,
         headers: dict,
-        payload: Any,
+        payload: object,
     ) -> None:
         verbose_proxy_logger.debug(
             "Prompt Security request %s %s headers=%s payload=%s",
@@ -670,7 +707,7 @@ class PromptSecurityGuardrail(CustomGuardrail):
         self,
         url: str,
         status_code: int,
-        payload: Any,
+        payload: object,
     ) -> None:
         verbose_proxy_logger.debug(
             "Prompt Security response %s status=%s payload=%s",

litellm/proxy/guardrails/guardrail_hooks/tool_permission.py

@@ -1,6 +1,6 @@
 import json
 import re
-from collections.abc import AsyncGenerator, Sequence
+from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence
 from typing import Any, Final, Literal
 
 from fastapi import HTTPException
@@ -41,6 +41,16 @@ from litellm.types.utils import (
 GUARDRAIL_NAME: Final = "tool_permission"
 
 
+def _object_mapping(value: object) -> Mapping[str, object] | None:
+    """Return ``value`` as an opaque mapping when it is a dict."""
+    return value if isinstance(value, dict) else None
+
+
+def _object_list(value: object) -> Sequence[object] | None:
+    """Return ``value`` as an opaque sequence when it is a list."""
+    return value if isinstance(value, list) else None
+
+
 class ToolPermissionGuardrail(CustomGuardrail):
     def __init__(
         self,
@@ -274,12 +284,12 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
     def _parse_tool_call_arguments(
         self, tool_call: ChatCompletionMessageToolCall
-    ) -> tuple[dict[str, Any] | None, str | None]:
+    ) -> tuple[Mapping[str, object] | None, str | None]:
         arguments: Final = getattr(tool_call.function, "arguments", None)
         if not arguments:
             return None, "missing arguments"
 
-        parsed_arguments: Any = {}
+        parsed_arguments: object = {}
         try:
             if isinstance(arguments, str):
                 parsed_arguments = json.loads(arguments)
@@ -306,9 +316,9 @@ class ToolPermissionGuardrail(CustomGuardrail):
 
     def _collect_argument_paths(
         self,
-        value: Any,
+        value: object,
         current_path: str,
-        collected: dict[str, list[Any]],
+        collected: dict[str, list[object]],
         depth: int = 0,
     ) -> None:
         from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
@@ -316,13 +326,15 @@ class ToolPermissionGuardrail(CustomGuardrail):
         if depth > DEFAULT_MAX_RECURSE_DEPTH:
             return
 
-        if isinstance(value, dict):
-            for key, sub_value in value.items():
+        mapping_value: Final = _object_mapping(value)
+        list_value: Final = _object_list(value)
+        if mapping_value is not None:
+            for key, sub_value in mapping_value.items():
                 next_path = f"{current_path}.{key}" if current_path else key
                 self._collect_argument_paths(sub_value, next_path, collected, depth + 1)
-        elif isinstance(value, list):
+        elif list_value is not None:
             list_path: Final = f"{current_path}[]" if current_path else "[]"
-            for item in value:
+            for item in list_value:
                 self._collect_argument_paths(item, list_path, collected, depth + 1)
         else:
             if not current_path:
@@ -332,7 +344,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
     def _patterns_match_for_rule(
         self,
         *,
-        arguments: dict[str, Any],
+        arguments: Mapping[str, object],
         rule: ToolPermissionRule,
         tool_name: str | None,
     ) -> tuple[bool, str | None]:
@@ -340,7 +352,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
         if not compiled_patterns:
             return True, None
 
-        path_value_map: Final[dict[str, list[Any]]] = {}
+        path_value_map: Final[dict[str, list[object]]] = {}
         self._collect_argument_paths(arguments, "", path_value_map)
 
         for path, compiled_pattern in compiled_patterns.items():
@@ -493,14 +505,14 @@ class ToolPermissionGuardrail(CustomGuardrail):
         )
 
     @staticmethod
-    def _get_anthropic_content_blocks(response: object) -> tuple[Any, ...] | None:
+    def _get_anthropic_content_blocks(response: object) -> tuple[object, ...] | None:
         if not isinstance(response, dict):
             return None
         content: Final[object] = response.get("content")
         return tuple(content) if isinstance(content, list) else None
 
     def _extract_tool_calls_from_anthropic_content(
-        self, content: tuple[Any, ...]
+        self, content: tuple[object, ...]
     ) -> tuple[ChatCompletionMessageToolCall, ...]:
         return tuple(
             tool_call for block in content if (tool_call := self._anthropic_tool_use_to_tool_call(block)) is not None
@@ -852,7 +864,7 @@ class ToolPermissionGuardrail(CustomGuardrail):
     async def async_post_call_streaming_iterator_hook(
         self,
         user_api_key_dict: UserAPIKeyAuth,
-        response: Any,
+        response: AsyncIterable[ModelResponseStream],
         request_data: dict,
     ) -> AsyncGenerator[ModelResponseStream, None]:
         """

litellm/proxy/guardrails/guardrail_hooks/vigil_guard/vigil_guard.py

@@ -1,8 +1,9 @@
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Mapping, Sequence
 from json import JSONDecodeError
-from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
+from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias, cast
 
 import httpx
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
 from litellm.exceptions import GuardrailRaisedException
@@ -50,7 +51,23 @@ _METADATA_ALLOWLIST: Final = (
     "org_id",
 )
 
-_FallbackMode = Literal["fail_closed", "fail_open"]
+_FallbackMode: TypeAlias = Literal["fail_closed", "fail_open"]
+_MetadataValue: TypeAlias = str | int | float | Sequence[str | int | float]
+
+
+class _AnalyzePayload(TypedDict):
+    """Request body posted to the Vigil Guard analyze endpoint."""
+
+    text: ReadOnly[str]
+    source: ReadOnly[str]
+    mode: ReadOnly[str]
+    metadata: ReadOnly[Mapping[str, _MetadataValue]]
+
+
+class _AnalysisView(TypedDict):
+    """Typed read of the analyze endpoint's decoded JSON body."""
+
+    analysis: ReadOnly[Mapping[str, object]]
 
 
 class _AsyncPostHandler(Protocol):
@@ -59,7 +76,7 @@ class _AsyncPostHandler(Protocol):
         *,
         url: str,
         headers: dict[str, str],
-        json: dict[str, Any],
+        json: _AnalyzePayload,
         timeout: httpx.Timeout,
     ) -> Awaitable[httpx.Response]: ...
 
@@ -244,7 +261,7 @@ class VigilGuardGuardrail(CustomGuardrail):
         exc: Exception,
         inputs: GenericGuardrailAPIInputs,
         source: str,
-        final_texts: list[Any],
+        final_texts: list[str],
         final_tool_calls: Any,
     ) -> GenericGuardrailAPIInputs:
         if self.unreachable_fallback == "fail_open":
@@ -271,7 +288,7 @@ class VigilGuardGuardrail(CustomGuardrail):
     @staticmethod
     def _build_output(
         inputs: GenericGuardrailAPIInputs,
-        final_texts: list[Any],
+        final_texts: list[str],
         final_tool_calls: Any,
     ) -> GenericGuardrailAPIInputs:
         # When nothing was changed, return the input shape verbatim so the guardrail
@@ -292,7 +309,7 @@ class VigilGuardGuardrail(CustomGuardrail):
         return guardrailed
 
     @staticmethod
-    def _tool_call_arguments(tool_calls: Any) -> list[tuple[int, str]]:
+    def _tool_call_arguments(tool_calls: Sequence[object] | None) -> list[tuple[int, str]]:
         pairs: Final[list[tuple[int, str]]] = []
         if isinstance(tool_calls, list):
             for index, tool_call in enumerate(tool_calls):
@@ -312,8 +329,8 @@ class VigilGuardGuardrail(CustomGuardrail):
         updated[index] = tool_call
         return updated
 
-    async def _analyze(self, text: str, source: str, metadata: dict[str, Any]) -> dict[str, Any]:
-        payload: Final = {
+    async def _analyze(self, text: str, source: str, metadata: Mapping[str, _MetadataValue]) -> Mapping[str, object]:
+        payload: Final[_AnalyzePayload] = {
             "text": text,
             "source": source,
             "mode": "full",
@@ -325,9 +342,12 @@ class VigilGuardGuardrail(CustomGuardrail):
             "Content-Type": "application/json",
         }
         response: Final = await self._post_with_retry(endpoint, headers, payload)
-        return response.json()
+        decoded: Final[_AnalysisView] = {"analysis": response.json()}
+        return decoded["analysis"]
 
-    async def _post_with_retry(self, endpoint: str, headers: dict[str, str], payload: dict[str, Any]) -> httpx.Response:
+    async def _post_with_retry(
+        self, endpoint: str, headers: dict[str, str], payload: _AnalyzePayload
+    ) -> httpx.Response:
         for attempt in range(2):
             try:
                 response = await self.async_handler.post(
@@ -364,7 +384,7 @@ class VigilGuardGuardrail(CustomGuardrail):
         )
 
     @staticmethod
-    def _build_block_reason(analysis: dict[str, Any]) -> str:
+    def _build_block_reason(analysis: Mapping[str, object]) -> str:
         for key in ("blockMessage", "decisionReason"):
             value = analysis.get(key)
             if isinstance(value, str) and value.strip():
@@ -377,14 +397,16 @@ class VigilGuardGuardrail(CustomGuardrail):
         return "Blocked by policy"
 
     @staticmethod
-    def _resolve_sanitized_text(original: str, analysis: dict[str, Any]) -> str:
+    def _resolve_sanitized_text(original: str, analysis: Mapping[str, object]) -> str:
         for key in ("sanitizedText", "outputText"):
             value = analysis.get(key)
             if isinstance(value, str):
                 return value
         return original
 
-    def _collect_metadata(self, request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"]) -> dict[str, Any]:
+    def _collect_metadata(
+        self, request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"]
+    ) -> Mapping[str, _MetadataValue]:
         sources: Final[list[dict]] = []
         if isinstance(request_data, dict):
             sources.append(request_data)
@@ -393,7 +415,7 @@ class VigilGuardGuardrail(CustomGuardrail):
                 if isinstance(nested, dict):
                     sources.append(nested)
 
-        collected: Final[dict[str, Any]] = {}
+        collected: Final[dict[str, _MetadataValue]] = {}
         for field in _METADATA_ALLOWLIST:
             for source in sources:
                 if field in source and source[field] is not None:
@@ -409,7 +431,7 @@ class VigilGuardGuardrail(CustomGuardrail):
         return collected
 
     @staticmethod
-    def _clamp_metadata_value(value: Any) -> Any:
+    def _clamp_metadata_value(value: Any) -> _MetadataValue | None:
         if isinstance(value, bool):
             return None
         if isinstance(value, str):
@@ -417,7 +439,7 @@ class VigilGuardGuardrail(CustomGuardrail):
         if isinstance(value, (int, float)):
             return value
         if isinstance(value, list):
-            clamped: Final[list[Any]] = []
+            clamped: Final[list[str | int | float]] = []
             for item in value[:_METADATA_ARRAY_MAX_ITEMS]:
                 if isinstance(item, bool):
                     continue

litellm/proxy/guardrails/guardrail_registry.py

@@ -1,13 +1,14 @@
 # litellm/proxy/guardrails/guardrail_registry.py
 
+import asyncio
 import importlib
 import os
-from collections.abc import Callable, Iterator, Mapping
+from collections.abc import Callable, Iterator, Mapping, Sequence
 from datetime import datetime, timezone
 from itertools import chain, count
-from typing import Any, Final, Literal, Optional, Protocol, cast
+from typing import Final, Literal, Optional, Protocol, cast
 
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
 
 import litellm
 from litellm import Router
@@ -67,6 +68,19 @@ class _GuardrailRowLike(Protocol):
     def __iter__(self) -> Iterator[tuple[str, object]]: ...
 
 
+class _GuardrailTableActions(Protocol):
+    async def create(self, *, data: Mapping[str, object]) -> _GuardrailRowLike: ...
+    async def delete(self, *, where: Mapping[str, str]) -> object: ...
+    async def update(self, *, where: Mapping[str, str], data: Mapping[str, object]) -> _GuardrailRowLike: ...
+    async def find_many(self, *, where: Mapping[str, str], order: Mapping[str, str]) -> Sequence[BaseModel]: ...
+    async def find_unique(self, *, where: Mapping[str, str]) -> BaseModel | None: ...
+
+
+def _guardrail_table(prisma_client: PrismaClient) -> _GuardrailTableActions:
+    """Typed view of the guardrails table actions exposed by the Prisma repository."""
+    return GuardrailsRepository(prisma_client).table
+
+
 guardrail_initializer_registry: Final = {
     SupportedGuardrailIntegrations.BEDROCK.value: initialize_bedrock,
     SupportedGuardrailIntegrations.LAKERA.value: initialize_lakera,
@@ -278,7 +292,7 @@ class GuardrailRegistry:
         try:
             guardrail_name: Final = guardrail.get("guardrail_name")
             # Properly serialize LitellmParams Pydantic model to dict
-            litellm_params_obj: Final[Any] = guardrail.get("litellm_params", {})
+            litellm_params_obj: Final = guardrail.get("litellm_params", {})
             if hasattr(litellm_params_obj, "model_dump"):
                 litellm_params_dict = litellm_params_obj.model_dump()
             else:
@@ -287,7 +301,7 @@ class GuardrailRegistry:
             guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {}))
 
             # Create guardrail in DB
-            created_guardrail: Final[_GuardrailRowLike] = await GuardrailsRepository(prisma_client).table.create(
+            created_guardrail: Final[_GuardrailRowLike] = await _guardrail_table(prisma_client).create(
                 data={
                     "guardrail_name": guardrail_name,
                     "litellm_params": litellm_params,
@@ -311,7 +325,7 @@ class GuardrailRegistry:
         """
         try:
             # Delete from DB
-            await GuardrailsRepository(prisma_client).table.delete(where={"guardrail_id": guardrail_id})
+            await _guardrail_table(prisma_client).delete(where={"guardrail_id": guardrail_id})
 
             return {"message": f"Guardrail {guardrail_id} deleted successfully"}
         except Exception as e:
@@ -324,7 +338,7 @@ class GuardrailRegistry:
         try:
             guardrail_name: Final = guardrail.get("guardrail_name")
             # Properly serialize LitellmParams Pydantic model to dict
-            litellm_params_obj: Final[Any] = guardrail.get("litellm_params", {})
+            litellm_params_obj: Final = guardrail.get("litellm_params", {})
             if hasattr(litellm_params_obj, "model_dump"):
                 litellm_params_dict = litellm_params_obj.model_dump()
             else:
@@ -333,7 +347,7 @@ class GuardrailRegistry:
             guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {}))
 
             # Update in DB
-            updated_guardrail: Final[_GuardrailRowLike] = await GuardrailsRepository(prisma_client).table.update(
+            updated_guardrail: Final[_GuardrailRowLike] = await _guardrail_table(prisma_client).update(
                 where={"guardrail_id": guardrail_id},
                 data={
                     "guardrail_name": guardrail_name,
@@ -357,7 +371,7 @@ class GuardrailRegistry:
         Only rows with status == "active" are returned (pending_review and rejected are excluded).
         """
         try:
-            guardrails_from_db: Final = await GuardrailsRepository(prisma_client).table.find_many(
+            guardrails_from_db: Final = await _guardrail_table(prisma_client).find_many(
                 where={"status": "active"},
                 order={"created_at": "desc"},
             )
@@ -375,9 +389,7 @@ class GuardrailRegistry:
         Get a guardrail by its ID from the database
         """
         try:
-            guardrail: Final = await GuardrailsRepository(prisma_client).table.find_unique(
-                where={"guardrail_id": guardrail_id}
-            )
+            guardrail: Final = await _guardrail_table(prisma_client).find_unique(where={"guardrail_id": guardrail_id})
 
             if not guardrail:
                 return None
@@ -391,7 +403,7 @@ class GuardrailRegistry:
         Get a guardrail by its name from the database
         """
         try:
-            guardrail: Final = await GuardrailsRepository(prisma_client).table.find_unique(
+            guardrail: Final = await _guardrail_table(prisma_client).find_unique(
                 where={"guardrail_name": guardrail_name}
             )
 
@@ -813,4 +825,6 @@ class InMemoryGuardrailHandler:
 # In Memory Guardrail Handler for LiteLLM Proxy
 ########################################################
 IN_MEMORY_GUARDRAIL_HANDLER: Final = InMemoryGuardrailHandler()
+
+GUARDRAIL_RECONCILE_LOCK: Final = asyncio.Lock()
 ########################################################

litellm/proxy/hooks/litellm_skills/main.py

@@ -27,7 +27,7 @@ Usage:
 import base64
 import json
 from collections.abc import Mapping, Sequence
-from typing import TYPE_CHECKING, Any, Final
+from typing import TYPE_CHECKING, Any, Final, Protocol
 
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
@@ -43,6 +43,30 @@ if TYPE_CHECKING:
     from litellm.llms.litellm_proxy.skills.sandbox_executor import SkillsSandboxExecutor
 
 
+class _ToolCallFunction(Protocol):
+    @property
+    def name(self) -> str: ...
+
+    @property
+    def arguments(self) -> str: ...
+
+
+class _ChatToolCall(Protocol):
+    @property
+    def id(self) -> str: ...
+
+    @property
+    def function(self) -> _ToolCallFunction: ...
+
+
+class _ChatMessage(Protocol):
+    @property
+    def content(self) -> str | None: ...
+
+    @property
+    def tool_calls(self) -> Sequence[_ChatToolCall] | None: ...
+
+
 class SkillsInjectionHook(CustomLogger):
     """
     Pre/Post-call hook that processes skills from container.skills parameter.
@@ -443,7 +467,7 @@ class SkillsInjectionHook(CustomLogger):
     async def _execute_code_loop_messages_api(
         self,
         data: dict,
-        response: Any,
+        response: object,
         skill_files: dict[str, bytes],
     ) -> LLMResponseTypes | None:
         """
@@ -673,7 +697,7 @@ print('No executable skill module found')
     async def _execute_code_loop(
         self,
         data: dict,
-        response: Any,
+        response: object,
         skill_files: dict[str, bytes],
     ) -> LLMResponseTypes:
         """
@@ -714,8 +738,8 @@ print('No executable skill module found')
 
         for iteration in range(self.max_iterations):
             # OpenAI format response has choices[0].message
-            assistant_message = current_response.choices[0].message
-            stop_reason = current_response.choices[0].finish_reason
+            assistant_message: _ChatMessage = current_response.choices[0].message
+            stop_reason: str | None = current_response.choices[0].finish_reason
 
             # Build assistant message for conversation history
             assistant_msg_dict: dict[str, object] = {
@@ -784,14 +808,14 @@ print('No executable skill module found')
 
     async def _execute_code_tool(
         self,
-        tool_call: Any,
+        tool_call: _ChatToolCall,
         skill_files: dict[str, bytes],
         executor: "SkillsSandboxExecutor",
         generated_files: list[dict[str, object]],
     ) -> str:
         """Execute a litellm_code_execution tool call and return result string."""
         try:
-            args: Final = json.loads(tool_call.function.arguments)
+            args: Final[Mapping[str, str]] = json.loads(tool_call.function.arguments)
             code: Final[str] = args.get("code", "")
 
             verbose_proxy_logger.debug("SkillsInjectionHook: Executing code (%s chars)", len(code))

litellm/proxy/management_endpoints/auto_router_endpoints.py

@@ -7,7 +7,7 @@ POST /auto_router/test_routing - Route one prompt through an unsaved complexity-
 from collections.abc import Mapping, Sequence
 from datetime import datetime, timedelta, timezone
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Annotated, Final
+from typing import TYPE_CHECKING, Annotated, Final, Protocol
 
 from pydantic import BaseModel, TypeAdapter
 
@@ -29,6 +29,7 @@ from litellm.proxy.auth.auth_checks import (
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.db.autorouter_session_rollup import AUTOROUTER_BENCHMARKS_SQL
 from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
+from litellm.repositories.base_repository import SupportsModelDump
 from litellm.repositories.team_repository import TeamRepository
 from litellm.router_strategy.complexity_router import ComplexityRouter
 from litellm.types.management_endpoints.auto_router_endpoints import (
@@ -61,6 +62,77 @@ else:
 router: Final = APIRouter()
 
 
+class _TeamTable(Protocol):
+    async def find_unique(self, *, where: Mapping[str, object]) -> SupportsModelDump | None: ...
+
+
+class _VerificationTokenRow(Protocol):
+    @property
+    def token(self) -> str: ...
+
+    @property
+    def key_alias(self) -> str | None: ...
+
+    @property
+    def key_name(self) -> str | None: ...
+
+
+class _VerificationTokenTable(Protocol):
+    async def find_unique(self, *, where: Mapping[str, object]) -> _VerificationTokenRow | None: ...
+
+    async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_VerificationTokenRow]: ...
+
+
+class _ShadowEvalJobRow(Protocol):
+    @property
+    def id(self) -> str: ...
+
+
+class _ShadowEvalJobTable(Protocol):
+    async def find_unique(self, *, where: Mapping[str, object]) -> _ShadowEvalJobRow | None: ...
+
+    async def find_first(self, *, where: Mapping[str, object]) -> _ShadowEvalJobRow | None: ...
+
+    async def find_many(
+        self, *, where: Mapping[str, object], order: Mapping[str, str], take: int
+    ) -> Sequence[_ShadowEvalJobRow]: ...
+
+    async def create(self, data: Mapping[str, object]) -> _ShadowEvalJobRow: ...
+
+    async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> _ShadowEvalJobRow | None: ...
+
+
+class _ShadowEvalAttemptRow(Protocol):
+    @property
+    def error(self) -> str | None: ...
+
+
+class _ShadowEvalAttemptTable(Protocol):
+    async def find_first(
+        self, *, where: Mapping[str, object], order: Mapping[str, str]
+    ) -> _ShadowEvalAttemptRow | None: ...
+
+
+def _team_table(prisma_client: "PrismaClient") -> _TeamTable:
+    return TeamRepository(prisma_client).table
+
+
+def _verification_tokens(prisma_client: "PrismaClient") -> _VerificationTokenTable:
+    return prisma_client.db.litellm_verificationtoken
+
+
+def _shadow_eval_jobs(prisma_client: "PrismaClient") -> _ShadowEvalJobTable:
+    return prisma_client.db.litellm_shadowevaljob
+
+
+def _shadow_eval_attempts(prisma_client: "PrismaClient") -> _ShadowEvalAttemptTable:
+    return prisma_client.db.litellm_shadowevalattempt
+
+
+async def _query_raw(prisma_client: "PrismaClient", query: str, *args: object) -> Sequence[Mapping[str, object]]:
+    return await prisma_client.db.query_raw(query, *args)
+
+
 async def _authorize_routing_test(user_api_key_dict: UserAPIKeyAuth, team_id: str | None) -> None:
     """Allow exactly the callers who could create this router.
 
@@ -92,7 +164,7 @@ async def _authorize_routing_test(user_api_key_dict: UserAPIKeyAuth, team_id: st
             },
         )
 
-    team_row: Final = await TeamRepository(prisma_client).table.find_unique(
+    team_row: Final = await _team_table(prisma_client).find_unique(
         where={"team_id": team_id},  # mutable-ok: Prisma query filters are dict-shaped
     )
     if team_row is None:
@@ -342,6 +414,26 @@ def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals:
     )
 
 
+def _benchmark_group(row: _SessionAggRow) -> AutoRouterBenchmarkGroup:
+    totals: Final = _benchmark_totals(row)
+    return AutoRouterBenchmarkGroup(
+        router_name=row.router_name,
+        router_type=row.router_type,
+        tier_turns=row.tier_turns,
+        sessions=totals.sessions,
+        turns=totals.turns,
+        avg_turns_per_session=totals.avg_turns_per_session,
+        avg_session_seconds=totals.avg_session_seconds,
+        avg_tokens_per_session=totals.avg_tokens_per_session,
+        spend=totals.spend,
+        saved_spend=totals.saved_spend,
+        baseline_spend=totals.baseline_spend,
+        saved_pct=totals.saved_pct,
+        saved_per_session=totals.saved_per_session,
+        cache=totals.cache,
+    )
+
+
 def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow:
     return _SessionAggRow(
         router_name="",
@@ -407,21 +499,14 @@ async def get_auto_router_benchmarks(
     if end_day < start_day:
         raise HTTPException(status_code=400, detail="end_date must not be earlier than start_date")
 
-    raw_rows: Final = await prisma_client.db.query_raw(
+    raw_rows: Final = await _query_raw(
+        prisma_client,
         AUTOROUTER_BENCHMARKS_SQL,
         start_day.isoformat(),
         (end_day + timedelta(days=1)).isoformat(),
     )
     rows: Final = _SESSION_AGG_ROWS.validate_python(raw_rows or ())
-    groups: Final = tuple(
-        AutoRouterBenchmarkGroup(
-            router_name=row.router_name,
-            router_type=row.router_type,
-            tier_turns=row.tier_turns,
-            **_benchmark_totals(row).model_dump(),
-        )
-        for row in rows
-    )
+    groups: Final = tuple(_benchmark_group(row) for row in rows)
     return AutoRouterBenchmarksResponse(
         start_date=start_day.strftime("%Y-%m-%d"),
         end_date=end_day.strftime("%Y-%m-%d"),
@@ -584,7 +669,7 @@ async def _with_key_labels(
     so the UI can say whose traffic a job shadows. Deleted keys resolve to None."""
     if not responses:
         return ()
-    key_rows: Final = await prisma_client.db.litellm_verificationtoken.find_many(
+    key_rows: Final = await _verification_tokens(prisma_client).find_many(
         where={"token": {"in": sorted({response.api_key_id for response in responses})}}  # mutable-ok: Prisma filter
     )
     labels: Final[Mapping[str, tuple[str | None, str | None]]] = {
@@ -608,12 +693,12 @@ async def _shadow_eval_results(prisma_client: "PrismaClient", job_id: str) -> Sh
     "for the turns the router sent to X, did X beat the baseline" in reverse. Reads are
     bounded by the job's own attempts (<= max_turns) via the job_id index."""
     by_tier: Final = _ATTEMPT_AGG_ROWS.validate_python(
-        await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_TIER_SQL, job_id) or ()
+        await _query_raw(prisma_client, _ATTEMPT_AGG_BY_TIER_SQL, job_id) or ()
     )
     if not by_tier:
         return None
     by_model: Final = _ATTEMPT_AGG_ROWS.validate_python(
-        await prisma_client.db.query_raw(_ATTEMPT_AGG_BY_MODEL_SQL, job_id) or ()
+        await _query_raw(prisma_client, _ATTEMPT_AGG_BY_MODEL_SQL, job_id) or ()
     )
     total_turns: Final = sum(r.turn_count for r in by_tier)
     return ShadowEvalResult(
@@ -661,7 +746,7 @@ async def start_shadow_eval(
     _validate_plain_model(llm_router, data.judge_model, "judge_model")
     if data.baseline_model is not None:
         _validate_plain_model(llm_router, data.baseline_model, "baseline_model")
-    key_row: Final = await prisma_client.db.litellm_verificationtoken.find_unique(
+    key_row: Final = await _verification_tokens(prisma_client).find_unique(
         where={"token": data.api_key_id}  # mutable-ok: Prisma filter
     )
     if key_row is None:
@@ -677,7 +762,7 @@ async def start_shadow_eval(
     # still holds its slot in the per-key, per-direction partial unique index until
     # stamped; free it so a new eval can start. Sweeping both directions is deliberate.
     await prisma_client.db.execute_raw(_SWEEP_FINISHED_JOBS_SQL, data.api_key_id)
-    active: Final = await prisma_client.db.litellm_shadowevaljob.find_first(
+    active: Final = await _shadow_eval_jobs(prisma_client).find_first(
         where={  # mutable-ok: Prisma filter
             "api_key_id": data.api_key_id,
             "direction": data.direction,
@@ -691,7 +776,7 @@ async def start_shadow_eval(
         )
     now: Final = datetime.now(timezone.utc)
     try:
-        job: Final = await prisma_client.db.litellm_shadowevaljob.create(
+        job: Final = await _shadow_eval_jobs(prisma_client).create(
             data={  # mutable-ok: Prisma payload
                 "api_key_id": data.api_key_id,
                 "router_name": data.router_name,
@@ -735,7 +820,7 @@ async def list_shadow_eval_jobs(
     _require_admin_viewer(user_api_key_dict, "view shadow evals")
     if prisma_client is None:
         raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
-    records: Final = await prisma_client.db.litellm_shadowevaljob.find_many(
+    records: Final = await _shadow_eval_jobs(prisma_client).find_many(
         where={"api_key_id": api_key_id} if api_key_id else {},  # mutable-ok: Prisma filter
         order={"created_at": "desc"},  # mutable-ok: Prisma order
         take=limit,
@@ -762,15 +847,15 @@ async def get_shadow_eval_job(
     _require_admin_viewer(user_api_key_dict, "view shadow evals")
     if prisma_client is None:
         raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
-    record: Final = await prisma_client.db.litellm_shadowevaljob.find_unique(
+    record: Final = await _shadow_eval_jobs(prisma_client).find_unique(
         where={"id": job_id}  # mutable-ok: Prisma filter
     )
     if record is None:
         raise HTTPException(status_code=404, detail=f"No shadow eval job {job_id}")
     totals: Final = _ATTEMPT_TOTALS_ROWS.validate_python(
-        await prisma_client.db.query_raw(_ATTEMPT_TOTALS_SQL, job_id) or ()
+        await _query_raw(prisma_client, _ATTEMPT_TOTALS_SQL, job_id) or ()
     )
-    latest_error: Final = await prisma_client.db.litellm_shadowevalattempt.find_first(
+    latest_error: Final = await _shadow_eval_attempts(prisma_client).find_first(
         where={"job_id": job_id, "outcome": "error"},  # mutable-ok: Prisma filter
         order={"created_at": "desc"},  # mutable-ok: Prisma order
     )
@@ -804,7 +889,7 @@ async def stop_shadow_eval_job(
     _require_admin_writer(user_api_key_dict, "stop a shadow eval")
     if prisma_client is None:
         raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)
-    record: Final = await prisma_client.db.litellm_shadowevaljob.find_unique(
+    record: Final = await _shadow_eval_jobs(prisma_client).find_unique(
         where={"id": job_id}  # mutable-ok: Prisma filter
     )
     if record is None:
@@ -812,7 +897,7 @@ async def stop_shadow_eval_job(
     current: Final = ShadowEvalJobResponse.model_validate(record, from_attributes=True)
     if current.status != "running":
         raise HTTPException(status_code=400, detail=f"Job {job_id} is already {current.status}")
-    updated: Final = await prisma_client.db.litellm_shadowevaljob.update(
+    updated: Final = await _shadow_eval_jobs(prisma_client).update(
         where={"id": job_id},  # mutable-ok: Prisma filter
         data={"stopped_at": datetime.now(timezone.utc)},  # mutable-ok: Prisma payload
     )

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -194,6 +194,13 @@ class _PrismaTableActions(Protocol[_PrismaRowT]):
         data: Mapping[str, object],
     ) -> _PrismaRowT | None: ...
 
+    async def upsert(
+        self,
+        *,
+        where: Mapping[str, object],
+        data: Mapping[str, object],
+    ) -> _PrismaRowT: ...
+
 
 class _UserRowLike(Protocol):
     user_id: str | None
@@ -209,24 +216,43 @@ class _TxTables(Protocol):
     litellm_proxymodeltable: _PrismaTableActions[object]
 
 
+class _TableSource(Protocol[_PrismaRowT]):
+    """Repository view that exposes its untyped Prisma ``table`` with a concrete row type."""
+
+    @property
+    def table(self) -> _PrismaTableActions[_PrismaRowT]: ...
+
+
+def _table_of(source: _TableSource[_PrismaRowT]) -> _PrismaTableActions[_PrismaRowT]:
+    return source.table
+
+
 def _prisma_table(
     repository: BaseRepository[_RepositoryModelT],
 ) -> _PrismaTableActions[_RepositoryModelT]:
-    return repository.table
+    return _table_of(repository)
 
 
 def _deleted_verification_token_table(
     prisma_client: PrismaClient,
 ) -> _PrismaTableActions[LiteLLM_DeletedVerificationToken]:
-    return DeletedVerificationTokenRepository(prisma_client).table
+    return _table_of(DeletedVerificationTokenRepository(prisma_client))
+
+
+def _deprecated_verification_token_table(prisma_client: PrismaClient) -> _PrismaTableActions[object]:
+    return _table_of(DeprecatedVerificationTokenRepository(prisma_client))
+
+
+def _user_table(prisma_client: PrismaClient) -> _PrismaTableActions[_UserRowLike]:
+    return _table_of(UserRepository(prisma_client))
 
 
 def _credentials_table(prisma_client: PrismaClient) -> _PrismaTableActions[CredentialItem]:
-    return CredentialsRepository(prisma_client).table
+    return _table_of(CredentialsRepository(prisma_client))
 
 
 def _config_table(prisma_client: PrismaClient) -> _PrismaTableActions[ConfigParam]:
-    return ConfigRepository(prisma_client).table
+    return _table_of(ConfigRepository(prisma_client))
 
 
 async def _check_custom_key_allowed(custom_key_value: str | None) -> None:
@@ -4656,7 +4682,7 @@ async def _insert_deprecated_key(
 
     try:
         revoke_at: Final = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds)
-        await DeprecatedVerificationTokenRepository(prisma_client).table.upsert(
+        await _deprecated_verification_token_table(prisma_client).upsert(
             where={"token": old_token_hash},
             data={
                 "create": {
@@ -6059,13 +6085,13 @@ async def _list_key_helper(
     total_pages: Final = -(-total_count // size)  # Ceiling division
 
     # Fetch user information if expand includes "user"
-    user_map = {}
+    user_map = dict[str | None, _UserRowLike]()
     if expand and "user" in expand:
         user_ids: Final = [key.user_id for key in keys if key.user_id]
         created_by_ids: Final = [key.created_by for key in keys if key.created_by]
         all_ids: Final = list(set(user_ids + created_by_ids))  # Remove duplicates
         if all_ids:
-            users: Final[Sequence[_UserRowLike]] = await UserRepository(prisma_client).table.find_many(
+            users: Final[Sequence[_UserRowLike]] = await _user_table(prisma_client).find_many(
                 where={"user_id": {"in": all_ids}}
             )
             user_map = {user.user_id: user for user in users}

litellm/proxy/management_endpoints/model_access_group_management_endpoints.py

@@ -57,20 +57,37 @@ def _model_table(prisma_client: PrismaClient) -> _ModelTableClient:
     return ModelRepository(prisma_client).table
 
 
-def validate_models_exist(model_names: list[str], llm_router: "Router | None") -> tuple[bool, list[str]]:
+def validate_models_exist(model_names: Sequence[str], llm_router: "Router | None") -> tuple[bool, Sequence[str]]:
     """
     Validate that all requested model names exist in the router.
     Checks only exact model name matches.
 
     Returns:
-        Tuple[bool, List[str]]: (all_valid, missing_models)
+        (all_valid, missing_models)
     """
     if llm_router is None:
         return False, model_names
 
-    router_model_names: Final = set(llm_router.get_model_names())
-    missing: Final = [m for m in model_names if m not in router_model_names]
-    return (len(missing) == 0, missing)
+    router_model_names: Final = frozenset(llm_router.get_model_names())
+    missing: Final = tuple(m for m in model_names if m not in router_model_names)
+    return (not missing, missing)
+
+
+async def _missing_models_after_read_through(
+    model_names: Sequence[str], llm_router: "Router | None"
+) -> tuple[str, ...]:
+    from litellm.proxy import proxy_server
+    from litellm.proxy.common_utils.registry_read_through import (
+        model_registry_read_through,
+    )
+
+    _, missing = validate_models_exist(model_names=model_names, llm_router=llm_router)
+    if not missing:
+        return ()
+    for name in missing:
+        await model_registry_read_through.attempt(name)
+    _, still_missing = validate_models_exist(model_names=model_names, llm_router=proxy_server.llm_router)
+    return tuple(still_missing)
 
 
 def add_access_group_to_deployment(model_info: dict[str, Any], access_group: str) -> tuple[dict[str, Any], bool]:
@@ -101,13 +118,21 @@ def _raise_http_if_reload_degraded_serving(
     before: frozenset[str],
     written_models: Sequence[tuple[str, object]],
     access_group: str,
+    still_desired: frozenset[str] | None,
+    live_after: frozenset[str] | None,
 ) -> None:
     """Same verdict as the model-write endpoints, expressed through this file's
     HTTPException error convention, with the metadata-only obligation: these writes
     change group membership, not the models themselves, so a row that was already not
     serving before the reload is never blamed here; only a model this reload stopped
     serving is reported."""
-    missing, collateral = reload_serving_verdict(before=before, written_models=written_models, written_must_serve=False)
+    missing, collateral = reload_serving_verdict(
+        before=before,
+        written_models=written_models,
+        written_must_serve=False,
+        still_desired=still_desired,
+        live_after=live_after,
+    )
     gone: Final = tuple(dict.fromkeys((*missing, *collateral)))
     if not gone:
         return
@@ -390,12 +415,12 @@ async def create_model_group(
     # Validate model_names exist in router (only if using model_names path)
     if not use_model_ids and has_model_names:
         assert data.model_names is not None
-        all_valid, missing_models = validate_models_exist(
+        missing_models: Final = await _missing_models_after_read_through(
             model_names=data.model_names,
             llm_router=llm_router,
         )
 
-        if not all_valid:
+        if missing_models:
             raise HTTPException(
                 status_code=400,
                 detail={"error": f"Model(s) not found: {', '.join(missing_models)}"},
@@ -439,11 +464,13 @@ async def create_model_group(
 
         live_before_reload: Final = live_model_ids_snapshot()
 
-        await clear_cache()
+        reload_outcome: Final = await clear_cache()
         _raise_http_if_reload_degraded_serving(
             before=live_before_reload,
             written_models=updated_pairs,
             access_group=data.access_group,
+            still_desired=reload_outcome.still_desired,
+            live_after=reload_outcome.live_after,
         )
 
         verbose_proxy_logger.info(
@@ -654,12 +681,12 @@ async def update_access_group(
     # Validation: Check if all new models exist (only if using model_names path)
     if not use_model_ids and has_model_names:
         assert data.model_names is not None
-        all_valid, missing_models = validate_models_exist(
+        missing_models: Final = await _missing_models_after_read_through(
             model_names=data.model_names,
             llm_router=llm_router,
         )
 
-        if not all_valid:
+        if missing_models:
             raise HTTPException(
                 status_code=400,
                 detail={"error": f"Model(s) not found: {', '.join(missing_models)}"},
@@ -699,11 +726,13 @@ async def update_access_group(
 
         # Clear cache and reload models to pick up the access group changes
         live_before_reload: Final = live_model_ids_snapshot()
-        await clear_cache()
+        reload_outcome: Final = await clear_cache()
         _raise_http_if_reload_degraded_serving(
             before=live_before_reload,
             written_models=list({**dict(stripped_pairs), **dict(updated_pairs)}.items()),
             access_group=access_group,
+            still_desired=reload_outcome.still_desired,
+            live_after=reload_outcome.live_after,
         )
 
         verbose_proxy_logger.info(
@@ -801,11 +830,13 @@ async def delete_access_group(
 
         # Clear cache and reload models to pick up the access group changes
         live_before_reload: Final = live_model_ids_snapshot()
-        await clear_cache()
+        reload_outcome: Final = await clear_cache()
         _raise_http_if_reload_degraded_serving(
             before=live_before_reload,
             written_models=removed_pairs,
             access_group=access_group,
+            still_desired=reload_outcome.still_desired,
+            live_after=reload_outcome.live_after,
         )
 
         verbose_proxy_logger.info(

litellm/proxy/management_endpoints/model_management_endpoints.py

@@ -114,13 +114,14 @@ class UpdatePublicModelGroupsRequest(BaseModel):
 class _ProxyModelRow(Protocol):
     model_id: str
     model_name: str
+    litellm_params: Mapping[str, object]
     model_info: Mapping[str, object] | None
 
     def model_dump_json(self, *, exclude_none: bool = False) -> str: ...
 
 
 class _ProxyModelTable(Protocol):
-    def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ...
+    def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[BaseModel | None]: ...
 
     def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ...
 
@@ -182,10 +183,7 @@ def _model_alias_table(prisma_client: PrismaClient) -> _ModelAliasTable:
 
 
 async def get_db_model(model_id: str, prisma_client: PrismaClient) -> Deployment | None:
-    db_model: Final = cast(
-        BaseModel | None,
-        await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_id}),
-    )
+    db_model: Final = await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_id})
 
     if not db_model:
         return None
@@ -1577,7 +1575,7 @@ async def delete_model(
                 },
             )
 
-        model_in_db: Final = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_info.id})
+        model_in_db: Final = await _proxy_model_table(prisma_client).find_unique(where={"model_id": model_info.id})
         if model_in_db is None:
             raise HTTPException(
                 status_code=400,
@@ -1914,7 +1912,7 @@ async def update_model(
             )
 
         _model_id: str | None = None
-        _model_info: Final = getattr(model_params, "model_info", None)
+        _model_info: Final[ModelInfo | None] = getattr(model_params, "model_info", None)
         if _model_info is None:
             raise Exception("model_info not provided")
 

litellm/proxy/management_helpers/utils.py

@@ -1,9 +1,9 @@
 # What is this?
 ## Helper utils for the management endpoints (keys/users/teams)
-from collections.abc import Callable
+from collections.abc import Callable, Mapping, MutableMapping, Sequence
 from datetime import datetime
 from functools import wraps
-from typing import Any, Final
+from typing import Any, Final, Protocol
 
 from fastapi import HTTPException, Request
 from pydantic import BaseModel
@@ -23,6 +23,7 @@ from litellm.proxy._types import (  # key request types; user request types; tea
     LiteLLM_UserTable,
     ManagementEndpointLoggingPayload,
     Member,
+    Span,
     SSOUserDefinedValues,
     UpdateCustomerRequest,
     UpdateKeyRequest,
@@ -39,7 +40,53 @@ from litellm.repositories.table_repositories import TeamMembershipRepository
 from litellm.repositories.user_repository import UserRepository
 
 
-def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict:
+class _PrismaRecord(Protocol):
+    """Row surface the management helpers read back from Prisma."""
+
+    def model_dump(self) -> Mapping[str, object]: ...
+
+
+class _PrismaUserRecord(Protocol):
+    """User row surface the management helpers read back from Prisma."""
+
+    user_id: str
+
+    def model_dump(self) -> Mapping[str, object]: ...
+
+
+class _PrismaBudgetRecord(Protocol):
+    """Budget row surface the management helpers read back from Prisma."""
+
+    budget_id: str
+
+    def model_dump(self) -> Mapping[str, object]: ...
+
+
+class _PrismaBudgetTable(Protocol):
+    """Budget table actions the management helpers issue."""
+
+    async def create(self, *, data: Mapping[str, object]) -> _PrismaBudgetRecord: ...
+
+    async def find_unique(self, *, where: Mapping[str, object]) -> _PrismaBudgetRecord | None: ...
+
+
+class _PrismaUserTable(Protocol):
+    """User table actions the management helpers issue."""
+
+    async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
+
+    async def upsert(
+        self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
+    ) -> _PrismaUserRecord | None: ...
+
+
+class _PrismaTeamMembershipTable(Protocol):
+    """Team membership table actions the management helpers issue."""
+
+    async def create(self, *, data: Mapping[str, object], include: Mapping[str, bool]) -> _PrismaRecord: ...
+
+
+def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict[str, object]:
     user_info: Final = litellm.default_internal_user_params or {}
 
     returned_dict: Final[SSOUserDefinedValues] = {
@@ -95,7 +142,7 @@ async def handle_budget_for_entity(
     _budget_data: Final = {k: v for k, v in _json_data.items() if k in budget_params}
 
     # Check if budget_id is explicitly provided in the data
-    data_budget_id: Final = getattr(data, "budget_id", None)
+    data_budget_id: Final[str | None] = getattr(data, "budget_id", None)
 
     # Case 1: Creating new entity - no existing budget_id
     if existing_budget_id is None:
@@ -107,7 +154,7 @@ async def handle_budget_for_entity(
             budget_row: Final = LiteLLM_BudgetTable(**_budget_data)
             new_budget_data: Final = prisma_client.jsonify_object(budget_row.model_dump(exclude_none=True))
 
-            _budget: Final = await BudgetRepository(prisma_client).table.create(
+            _budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create(
                 data={
                     **new_budget_data,
                     "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
@@ -173,9 +220,8 @@ async def _clone_team_default_budget_for_member(
     member while keeping the default's other limits, so an admin can set a
     member's reset cadence without discarding the team default's max_budget.
     """
-    default_budget: Final = await BudgetRepository(prisma_client).table.find_unique(
-        where={"budget_id": default_team_budget_id}
-    )
+    budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table
+    default_budget: Final = await budget_table.find_unique(where={"budget_id": default_team_budget_id})
     if default_budget is None:
         return None
 
@@ -202,7 +248,7 @@ async def _clone_team_default_budget_for_member(
     if cloned_data.get("budget_duration"):
         cloned_data["budget_reset_at"] = get_budget_reset_time(cloned_data["budget_duration"])
 
-    new_budget: Final = await BudgetRepository(prisma_client).table.create(data=cloned_data)
+    new_budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create(data=cloned_data)
     return new_budget.budget_id
 
 
@@ -238,7 +284,7 @@ async def _resolve_member_budget_id(
     if not has_explicit_limit and budget_duration is None:
         return None
 
-    budget_data: Final[dict] = {
+    budget_data: Final[dict[str, object]] = {
         "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
         "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
     }
@@ -249,7 +295,8 @@ async def _resolve_member_budget_id(
     if budget_duration is not None:
         budget_data["budget_duration"] = budget_duration
         budget_data["budget_reset_at"] = get_budget_reset_time(budget_duration=budget_duration)
-    response: Final = await BudgetRepository(prisma_client).table.create(data=budget_data)
+    budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table
+    response: Final = await budget_table.create(data=budget_data)
     return response.budget_id
 
 
@@ -262,7 +309,8 @@ async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, t
     number of teams a user belongs to). Teams added concurrently for a different
     team id are unaffected, since each update filters on its own team id.
     """
-    await UserRepository(prisma_client).table.update_many(
+    user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
+    await user_table.update_many(
         where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}},
         data={"teams": {"push": [team_id]}},
     )
@@ -300,7 +348,8 @@ async def add_new_member(
         # Prisma only compiles an upsert down to INSERT ... ON CONFLICT when it
         # is non-empty, and falls back to a racy SELECT-then-INSERT when it is
         # not, so this re-states user_id as a no-op rather than being empty.
-        _returned_user = await UserRepository(prisma_client).table.upsert(
+        user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
+        _returned_user: _PrismaUserRecord | None = await user_table.upsert(
             where={"user_id": new_member.user_id},
             data={
                 "create": {"teams": [team_id], **new_user_defaults},
@@ -314,7 +363,7 @@ async def add_new_member(
         new_user_defaults = get_new_internal_user_defaults(user_id=str(uuid.uuid4()), user_email=new_member.user_email)
         ## user email is not unique acc. to prisma schema -> future improvement
         ### for now: check if it exists in db, if not - insert it
-        existing_user_row: Final[list | None] = await prisma_client.get_data(
+        existing_user_row: Final[list[_PrismaUserRecord] | None] = await prisma_client.get_data(
             key_val={"user_email": new_member.user_email},
             table_name="user",
             query_type="find_all",
@@ -346,7 +395,8 @@ async def add_new_member(
     )
 
     if _budget_id and returned_user is not None and returned_user.user_id is not None:
-        _returned_team_membership: Final = await TeamMembershipRepository(prisma_client).table.create(
+        membership_table: Final[_PrismaTeamMembershipTable] = TeamMembershipRepository(prisma_client).table
+        _returned_team_membership: Final = await membership_table.create(
             data={
                 "team_id": team_id,
                 "user_id": returned_user.user_id,
@@ -469,8 +519,18 @@ async def send_management_endpoint_alert(
             )
 
 
-def _redacted_env_var(entry: Any) -> dict:
-    get: Final = entry.get if isinstance(entry, dict) else lambda k: getattr(entry, k, None)
+def _object_mapping(value: object) -> Mapping[str, object] | None:
+    """Return ``value`` as an opaque mapping when it is a dict."""
+    return value if isinstance(value, dict) else None
+
+
+def _object_list(value: object) -> Sequence[object] | None:
+    """Return ``value`` as an opaque sequence when it is a list."""
+    return value if isinstance(value, list) else None
+
+
+def _redacted_env_var(entry: object) -> dict[str, object]:
+    get: Final[Callable[[str], object]] = entry.get if isinstance(entry, dict) else lambda k: getattr(entry, k, None)
     return {
         "name": get("name"),
         "scope": get("scope"),
@@ -479,25 +539,28 @@ def _redacted_env_var(entry: Any) -> dict:
     }
 
 
-def _redact_record_env_vars(record: Any) -> Any:
+def _redact_record_env_vars(record: object) -> object:
     """Return ``record`` with its ``env_vars[].value`` blanked.
 
     Copies rather than mutating, because the record aliases the live response
     object that is also returned to the caller. Records without an ``env_vars``
     list are returned unchanged.
     """
-    env_vars: Final = record.get("env_vars") if isinstance(record, dict) else getattr(record, "env_vars", None)
-    if not isinstance(env_vars, list):
+    record_map: Final = _object_mapping(record)
+    env_vars: Final = _object_list(
+        record_map.get("env_vars") if record_map is not None else getattr(record, "env_vars", None)
+    )
+    if env_vars is None:
         return record
     redacted: Final = [_redacted_env_var(entry) for entry in env_vars]
-    if isinstance(record, dict):
-        return {**record, "env_vars": redacted}
+    if record_map is not None:
+        return {**record_map, "env_vars": redacted}
     if isinstance(record, BaseModel):
         return record.model_copy(update={"env_vars": redacted})
     return record
 
 
-def _redact_env_var_values(response: dict) -> None:
+def _redact_env_var_values(response: MutableMapping[str, object]) -> None:
     """Blank ``env_vars[].value`` in a management response before telemetry.
 
     MCP endpoints return decrypted ``scope="global"`` env var values so the admin
@@ -507,18 +570,19 @@ def _redact_env_var_values(response: dict) -> None:
     create/update) and nested under ``items`` (the submissions queue), so both are
     scrubbed. Names, scopes, and descriptions are kept so traces stay useful.
     """
-    if isinstance(response.get("env_vars"), list):
-        response["env_vars"] = [_redacted_env_var(entry) for entry in response["env_vars"]]
+    env_vars: Final = _object_list(response.get("env_vars"))
+    if env_vars is not None:
+        response["env_vars"] = [_redacted_env_var(entry) for entry in env_vars]
 
-    items: Final = response.get("items")
-    if isinstance(items, list):
+    items: Final = _object_list(response.get("items"))
+    if items is not None:
         response["items"] = [_redact_record_env_vars(item) for item in items]
 
 
 async def _emit_management_endpoint_otel_span(
     func: Callable,
     kwargs: dict,
-    parent_otel_span: Any,
+    parent_otel_span: Span | None,
     start_time: datetime,
     end_time: datetime,
     result: Any = None,
@@ -571,10 +635,10 @@ async def _emit_management_endpoint_otel_span(
         }
     )
 
-    _response: dict | None = None
+    _response: dict[str, object] | None = None
     if exception is None and result is not None:
         try:
-            raw: Final = dict(result)
+            raw: Final[Mapping[str, object]] = dict(result)
             _response = {k: v for k, v in raw.items() if k not in _CREDENTIAL_FIELDS}
             _redact_env_var_values(_response)
         except Exception:
@@ -623,7 +687,7 @@ def management_endpoint_wrapper(func):
                     user_api_key_dict=user_api_key_dict,
                     function_name=func.__name__,
                 )
-                parent_otel_span = getattr(user_api_key_dict, "parent_otel_span", None)
+                parent_otel_span: Span | None = getattr(user_api_key_dict, "parent_otel_span", None)
                 if parent_otel_span is not None:
                     await _emit_management_endpoint_otel_span(
                         func=func,

litellm/proxy/proxy_server.py

@@ -4120,6 +4120,31 @@ def _swap_in_model_cost_map(new_model_cost_map: dict) -> int:
     return fetched_model_count
 
 
+def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool:
+    """
+    Check if an object type should be loaded from the database based on general_settings.supported_db_objects.
+
+    Args:
+        object_type: Type of object to check (e.g., SupportedDBObjectType.MODELS, "models", etc.)
+
+    Returns:
+        True if the object should be loaded, False otherwise
+    """
+    supported_db_objects: Final = general_settings.get("supported_db_objects", None)
+
+    if supported_db_objects is None:
+        return True
+
+    if not isinstance(supported_db_objects, list):
+        verbose_proxy_logger.warning(
+            "supported_db_objects is not a list, got %s. Loading all objects.", type(supported_db_objects)
+        )
+        return True
+
+    object_type_str: Final = str(object_type)
+    return any(str(obj) == object_type_str for obj in supported_db_objects)
+
+
 class ProxyConfig:
     """
     Abstraction class on top of config loading/updating logic. Gives us one place to control all config updating logic.
@@ -6522,36 +6547,7 @@ class ProxyConfig:
         return config
 
     def _should_load_db_object(self, object_type: str | SupportedDBObjectType) -> bool:
-        """
-        Check if an object type should be loaded from the database based on general_settings.supported_db_objects.
-
-        Args:
-            object_type: Type of object to check (e.g., SupportedDBObjectType.MODELS, "models", etc.)
-
-        Returns:
-            True if the object should be loaded, False otherwise
-        """
-        global general_settings
-
-        # Get the supported_db_objects configuration
-        supported_db_objects: Final = general_settings.get("supported_db_objects", None)
-
-        # If supported_db_objects is not set, load all objects (default behavior)
-        if supported_db_objects is None:
-            return True
-
-        # If supported_db_objects is set, only load specified objects
-        if not isinstance(supported_db_objects, list):
-            verbose_proxy_logger.warning(
-                "supported_db_objects is not a list, got %s. Loading all objects.", type(supported_db_objects)
-            )
-            return True
-
-        # Convert object_type to string for comparison (handles both str and enum)
-        object_type_str: Final = str(object_type)
-
-        # Check if the object type is in the list (supports both str and enum values)
-        return any(str(obj) == object_type_str for obj in supported_db_objects)
+        return should_load_db_object(object_type=object_type)
 
     async def _get_models_from_db(self, prisma_client: PrismaClient) -> list | None:
         """
@@ -7094,38 +7090,40 @@ class ProxyConfig:
 
     async def _init_guardrails_in_db(self, prisma_client: PrismaClient):
         from litellm.proxy.guardrails.guardrail_registry import (
+            GUARDRAIL_RECONCILE_LOCK,
             IN_MEMORY_GUARDRAIL_HANDLER,
             Guardrail,
             GuardrailRegistry,
         )
 
         try:
-            guardrails_in_db: Final[list[Guardrail]] = await GuardrailRegistry.get_all_guardrails_from_db(
-                prisma_client=prisma_client
-            )
-            verbose_proxy_logger.debug("guardrails from the DB %s", str(guardrails_in_db))
-            db_guardrail_ids: Final[set] = set()
-            for guardrail in guardrails_in_db:
-                guardrail_id = guardrail.get("guardrail_id")
-                if guardrail_id:
-                    db_guardrail_ids.add(guardrail_id)
-                try:
-                    IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
-                        guardrail=cast(Guardrail, guardrail),
-                    )
-                except Exception as e:  # noqa: BLE001  # one unloadable row must not stop the remaining guardrails
-                    verbose_proxy_logger.error(
-                        "litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - "
-                        "skipping guardrail '%s' (ID: %s): %s: %s",
-                        guardrail.get("guardrail_name"),
-                        guardrail_id,
-                        type(e).__name__,
-                        e,
-                    )
+            async with GUARDRAIL_RECONCILE_LOCK:
+                guardrails_in_db: Final[list[Guardrail]] = await GuardrailRegistry.get_all_guardrails_from_db(
+                    prisma_client=prisma_client
+                )
+                verbose_proxy_logger.debug("guardrails from the DB %s", str(guardrails_in_db))
+                db_guardrail_ids: Final[set] = set()
+                for guardrail in guardrails_in_db:
+                    guardrail_id = guardrail.get("guardrail_id")
+                    if guardrail_id:
+                        db_guardrail_ids.add(guardrail_id)
+                    try:
+                        IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
+                            guardrail=cast(Guardrail, guardrail),
+                        )
+                    except Exception as e:  # noqa: BLE001  # one unloadable row must not stop the remaining guardrails
+                        verbose_proxy_logger.error(
+                            "litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - "
+                            "skipping guardrail '%s' (ID: %s): %s: %s",
+                            guardrail.get("guardrail_name"),
+                            guardrail_id,
+                            type(e).__name__,
+                            e,
+                        )
 
-            # Drop in-memory DB-backed entries whose row was deleted on another
-            # pod. Config-loaded entries are never touched.
-            IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails(db_guardrail_ids=db_guardrail_ids)
+                # Drop in-memory DB-backed entries whose row was deleted on another
+                # pod. Config-loaded entries are never touched.
+                IN_MEMORY_GUARDRAIL_HANDLER.reconcile_db_guardrails(db_guardrail_ids=db_guardrail_ids)
         except Exception as e:
             verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:_init_guardrails_in_db - %s", e)
 
@@ -7278,13 +7276,17 @@ class ProxyConfig:
             )
 
     async def _init_agents_in_db(self, prisma_client: PrismaClient):
+        from litellm.proxy.agent_endpoints.agent_registry import (
+            AGENT_RECONCILE_LOCK,
+        )
         from litellm.proxy.agent_endpoints.agent_registry import (
             global_agent_registry as AGENT_REGISTRY,
         )
 
         try:
-            db_agents: Final = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client)
-            AGENT_REGISTRY.load_agents_from_db_and_config(db_agents=db_agents)
+            async with AGENT_RECONCILE_LOCK:
+                db_agents: Final = await AGENT_REGISTRY.get_all_agents_from_db(prisma_client=prisma_client)
+                AGENT_REGISTRY.load_agents_from_db_and_config(db_agents=db_agents)
         except Exception as e:
             verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:_init_agents_in_db - %s", e)
 
@@ -12990,9 +12992,9 @@ async def _filter_models_by_team_id(
 async def _find_model_by_id(
     model_id: str,
     search: str | None,
-    llm_router,
-    prisma_client,
-    proxy_config,
+    llm_router: Router | None,
+    prisma_client: PrismaClient | None,
+    proxy_config: "ProxyConfig",
 ) -> tuple[list, int | None]:
     """Find a model by its ID and optionally filter by search term."""
     found_model = None

litellm/proxy/route_llm_request.py

@@ -146,7 +146,8 @@ ROUTE_ENDPOINT_MAPPING: Final = {
 
 
 class ProxyModelNotFoundError(HTTPException):
-    def __init__(self, route: str, model_name: str):
+    def __init__(self, route: str, model_name: str, retryable_with_model_read_through: bool = True):
+        self.retryable_with_model_read_through: Final = retryable_with_model_read_through
         detail: Final = {
             "error": f"{route}: Invalid model name passed in model={model_name}. Call `/v1/models` to view available models for your key."
         }
@@ -320,112 +321,150 @@ async def add_shared_session_to_data(data: dict) -> None:
             pass
 
 
+RouteType = Literal[
+    "acompletion",
+    "atext_completion",
+    "aembedding",
+    "aimage_generation",
+    "aspeech",
+    "atranscription",
+    "amoderation",
+    "arerank",
+    "aresponses",
+    "aget_responses",
+    "adelete_responses",
+    "acancel_responses",
+    "acompact_responses",
+    "acreate_response_reply",
+    "alist_input_items",
+    "_arealtime",  # private function for realtime API
+    "acreate_realtime_client_secret",
+    "arealtime_calls",
+    "acreate_realtime_transcription_session",
+    "_aresponses_websocket",  # private function for responses WebSocket mode
+    "aimage_edit",
+    "agenerate_content",
+    "agenerate_content_stream",
+    "allm_passthrough_route",
+    "acreate_batch",
+    "aretrieve_batch",
+    "alist_batches",
+    "afile_content",
+    "afile_retrieve",
+    "acreate_fine_tuning_job",
+    "acancel_fine_tuning_job",
+    "alist_fine_tuning_jobs",
+    "aretrieve_fine_tuning_job",
+    "avector_store_search",
+    "avector_store_create",
+    "avector_store_retrieve",
+    "avector_store_list",
+    "avector_store_update",
+    "avector_store_delete",
+    "avector_store_file_create",
+    "avector_store_file_list",
+    "avector_store_file_retrieve",
+    "avector_store_file_content",
+    "avector_store_file_update",
+    "avector_store_file_delete",
+    "aocr",
+    "asearch",
+    "avideo_generation",
+    "avideo_list",
+    "avideo_status",
+    "avideo_content",
+    "avideo_remix",
+    "avideo_create_character",
+    "avideo_get_character",
+    "avideo_edit",
+    "avideo_extension",
+    "acreate_container",
+    "alist_containers",
+    "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",
+    "adelete_skill",
+    "aingest",
+    "anthropic_messages",
+    "acreate_interaction",
+    "aget_interaction",
+    "adelete_interaction",
+    "acancel_interaction",
+    "acreate_agent",
+    "alist_agents",
+    "aget_agent",
+    "adelete_agent",
+    "alist_agent_versions",
+    "asend_message",
+    "call_mcp_tool",
+    "acancel_batch",
+    "afile_delete",
+    "acreate_eval",
+    "alist_evals",
+    "aget_eval",
+    "aupdate_eval",
+    "adelete_eval",
+    "acancel_eval",
+    "acreate_run",
+    "alist_runs",
+    "aget_run",
+    "acancel_run",
+    "adelete_run",
+]
+
+
 async def route_request(
     data: dict,
     llm_router: LitellmRouter | None,
     user_model: str | None,
-    route_type: Literal[
-        "acompletion",
-        "atext_completion",
-        "aembedding",
-        "aimage_generation",
-        "aspeech",
-        "atranscription",
-        "amoderation",
-        "arerank",
-        "aresponses",
-        "aget_responses",
-        "adelete_responses",
-        "acancel_responses",
-        "acompact_responses",
-        "acreate_response_reply",
-        "alist_input_items",
-        "_arealtime",  # private function for realtime API
-        "acreate_realtime_client_secret",
-        "arealtime_calls",
-        "acreate_realtime_transcription_session",
-        "_aresponses_websocket",  # private function for responses WebSocket mode
-        "aimage_edit",
-        "agenerate_content",
-        "agenerate_content_stream",
-        "allm_passthrough_route",
-        "acreate_batch",
-        "aretrieve_batch",
-        "alist_batches",
-        "afile_content",
-        "afile_retrieve",
-        "acreate_fine_tuning_job",
-        "acancel_fine_tuning_job",
-        "alist_fine_tuning_jobs",
-        "aretrieve_fine_tuning_job",
-        "avector_store_search",
-        "avector_store_create",
-        "avector_store_retrieve",
-        "avector_store_list",
-        "avector_store_update",
-        "avector_store_delete",
-        "avector_store_file_create",
-        "avector_store_file_list",
-        "avector_store_file_retrieve",
-        "avector_store_file_content",
-        "avector_store_file_update",
-        "avector_store_file_delete",
-        "aocr",
-        "asearch",
-        "avideo_generation",
-        "avideo_list",
-        "avideo_status",
-        "avideo_content",
-        "avideo_remix",
-        "avideo_create_character",
-        "avideo_get_character",
-        "avideo_edit",
-        "avideo_extension",
-        "acreate_container",
-        "alist_containers",
-        "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",
-        "adelete_skill",
-        "aingest",
-        "anthropic_messages",
-        "acreate_interaction",
-        "aget_interaction",
-        "adelete_interaction",
-        "acancel_interaction",
-        "acreate_agent",
-        "alist_agents",
-        "aget_agent",
-        "adelete_agent",
-        "alist_agent_versions",
-        "asend_message",
-        "call_mcp_tool",
-        "acancel_batch",
-        "afile_delete",
-        "acreate_eval",
-        "alist_evals",
-        "aget_eval",
-        "aupdate_eval",
-        "adelete_eval",
-        "acancel_eval",
-        "acreate_run",
-        "alist_runs",
-        "aget_run",
-        "acancel_run",
-        "adelete_run",
-    ],
+    route_type: RouteType,
     user_api_key_dict: UserAPIKeyAuth | None = None,
 ):
     """
     Common helper to route the request
     """
+    try:
+        return await _route_request_single_attempt(
+            data=data,
+            llm_router=llm_router,
+            user_model=user_model,
+            route_type=route_type,
+            user_api_key_dict=user_api_key_dict,
+        )
+    except ProxyModelNotFoundError as e:
+        requested_model: Final = data.get("model", "")
+        if not e.retryable_with_model_read_through or not isinstance(requested_model, str) or not requested_model:
+            raise
+        from litellm.proxy import proxy_server
+        from litellm.proxy.common_utils.registry_read_through import (
+            model_registry_read_through,
+        )
+
+        if not await model_registry_read_through.attempt(requested_model):
+            raise
+        return await _route_request_single_attempt(
+            data=data,
+            llm_router=proxy_server.llm_router,
+            user_model=user_model,
+            route_type=route_type,
+            user_api_key_dict=user_api_key_dict,
+        )
+
+
+async def _route_request_single_attempt(  # noqa: ANN202  # returns unawaited provider coroutines; the inferred union keeps route_request's callers typed
+    data: dict,  # mutable-ok: request body is the proxy-wide mutable dict contract shared with route_request
+    llm_router: LitellmRouter | None,
+    user_model: str | None,
+    route_type: RouteType,
+    user_api_key_dict: UserAPIKeyAuth | None = None,
+):
     raise_if_required_body_param_missing(route_type=route_type, data=data)
 
     await add_shared_session_to_data(data)

litellm/proxy/spend_tracking/budget_reservation.py

@@ -105,8 +105,8 @@ def _key_reservation_should_release_for_throttle(counter_key: str, valid_token:
 async def _apply_over_budget_reservation_policy(
     counter: _BudgetCounter,
     valid_token: UserAPIKeyAuth | None,
-    entry: dict[str, Any],
-    applied_entries: list[dict[str, Any]],
+    entry: dict[str, float | str],
+    applied_entries: list[dict[str, float | str]],
     reservation_cost: float,
     current_spend: float,
 ) -> float:
@@ -156,7 +156,7 @@ async def reserve_budget_for_request(
     user_api_key_cache: DualCache,
     proxy_logging_obj: ProxyLogging,
     end_user_id: str | None = None,
-    end_user_object: Any | None = None,
+    end_user_object: object = None,
     apply_user_budget_to_team_keys: bool = False,
     fail_closed_budget_enforcement: bool = False,
 ) -> dict | None:
@@ -194,7 +194,7 @@ async def reserve_budget_for_request(
     if reservation_cost is None or reservation_cost <= 0:
         return None
 
-    applied_entries: Final[list[dict[str, Any]]] = []
+    applied_entries: Final[list[dict[str, float | str]]] = []
     try:
         for counter in counters:
             entry = _counter_to_reservation_entry(
@@ -334,7 +334,7 @@ async def _get_budget_counters(
     user_api_key_cache: DualCache,
     proxy_logging_obj: ProxyLogging,
     end_user_id: str | None = None,
-    end_user_object: Any | None = None,
+    end_user_object: object = None,
     apply_user_budget_to_team_keys: bool = False,
 ) -> list[_BudgetCounter]:
     counters: Final[list[_BudgetCounter]] = []
@@ -443,7 +443,7 @@ async def _get_budget_counters(
 async def _get_end_user_budget_counter(
     valid_token: UserAPIKeyAuth,
     end_user_id: str | None,
-    end_user_object: Any | None,
+    end_user_object: object,
 ) -> _BudgetCounter | None:
     end_user_id = end_user_id or valid_token.end_user_id
     if end_user_id is None:
@@ -608,7 +608,7 @@ def _get_budget_limit_counters(
     entity_prefix: str,
     entity_type: str,
     entity_id: str,
-    budget_limits: Sequence[Any] | None,
+    budget_limits: Sequence[object] | None,
     fallback_spend: float,
 ) -> list[_BudgetCounter]:
     counters: Final[list[_BudgetCounter]] = []
@@ -855,7 +855,7 @@ async def _resize_applied_reservation(
 def _counter_to_reservation_entry(
     counter: _BudgetCounter,
     reserved_cost: float,
-) -> dict[str, Any]:
+) -> dict[str, float | str]:
     return {
         "counter_key": counter.counter_key,
         "entity_type": counter.entity_type,
@@ -983,7 +983,7 @@ def _input_cost_for_cost_info(
     request_body: dict,
     route: str,
     model: str,
-    model_info: dict[str, Any],
+    model_info: Mapping[str, object],
 ) -> float | None:
     input_tokens: Final = _estimate_input_tokens(
         request_body=request_body,
@@ -1027,7 +1027,7 @@ def _max_cost_for_cost_info(
     request_body: dict,
     route: str,
     model: str,
-    model_info: dict[str, Any],
+    model_info: Mapping[str, object],
 ) -> float | None:
     image_cost: Final = _estimate_image_generation_cost(
         request_body=request_body,
@@ -1086,7 +1086,7 @@ def _max_cost_for_cost_info(
 
 def _estimate_image_generation_cost(
     request_body: dict,
-    model_info: dict[str, Any],
+    model_info: Mapping[str, object],
 ) -> float | None:
     """
     Reserve `n × per-image cost` for image-generation requests so concurrent
@@ -1125,7 +1125,7 @@ def _estimate_image_generation_cost(
 def _get_model_cost_info(
     model: str,
     llm_router: Router | None,
-) -> dict[str, Any] | None:
+) -> Mapping[str, object] | None:
     if llm_router is not None:
         model_group_info: Final = llm_router.get_model_group_info(model_group=model)
         if model_group_info is not None:
@@ -1136,7 +1136,7 @@ def _get_model_cost_info(
 def _get_model_cost_infos(
     model: str,
     llm_router: Router | None,
-) -> list[dict[str, Any]]:
+) -> Sequence[Mapping[str, object]]:
     """Cost-info candidates to estimate a request against for one model group.
 
     Reservation runs before routing, so the deployment that will serve the request
@@ -1181,7 +1181,7 @@ def _deployment_tiered_pricing_table(
 def _get_deployment_tiered_pricing_tables(
     model: str,
     llm_router: Router | None,
-) -> list[list[dict]]:
+) -> Sequence[Sequence[Mapping[str, object]]]:
     if llm_router is None:
         return []
     deployments: Final = llm_router.get_model_list(model_name=model) or []
@@ -1196,7 +1196,7 @@ def _estimate_input_tokens(
     request_body: dict,
     route: str,
     model: str,
-    model_info: dict[str, Any],
+    model_info: Mapping[str, object],
 ) -> int | None:
     try:
         if "messages" in request_body:
@@ -1233,7 +1233,7 @@ DEFAULT_MAX_OUTPUT_TOKENS_FALLBACK: Final = 16384
 def _estimate_output_tokens(
     request_body: dict,
     route: str,
-    model_info: dict[str, Any],
+    model_info: Mapping[str, object],
 ) -> int | None:
     if _is_input_only_route(route=route):
         return 0

litellm/repositories/config_repository.py

@@ -10,12 +10,41 @@ import asyncio
 import copy
 import json
 import os
-from typing import Any, Final, Literal, cast
+from collections.abc import Mapping, Sequence
+from typing import Any, Final, Literal, Protocol, cast
 
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
 
 
+class _ConfigRow(Protocol):
+    @property
+    def param_name(self) -> str: ...
+
+    @property
+    def param_value(self) -> object: ...
+
+
+class _ConfigTable(Protocol):
+    async def find_unique(self, *, where: Mapping[str, str]) -> _ConfigRow | None: ...
+
+    async def find_many(self) -> Sequence[_ConfigRow]: ...
+
+    async def upsert(self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> _ConfigRow: ...
+
+    async def delete(self, *, where: Mapping[str, str]) -> _ConfigRow | None: ...
+
+
+class _ConfigDb(Protocol):
+    @property
+    def litellm_config(self) -> _ConfigTable: ...
+
+
+class _PrismaHandle(Protocol):
+    @property
+    def db(self) -> _ConfigDb: ...
+
+
 class ConfigParam:
     """Simple wrapper for config parameter from DB."""
 
@@ -38,18 +67,22 @@ class ConfigRepository:
         self._prisma_client = prisma_client
 
     @property
-    def prisma_client(self) -> Any:
+    def prisma_client(self) -> _PrismaHandle:
         if self._prisma_client is None:
             raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
         return self._prisma_client
 
     @property
-    def table(self) -> Any:
+    def _config_table(self) -> _ConfigTable:
         return self.prisma_client.db.litellm_config
 
+    @property
+    def table(self) -> Any:
+        return self._config_table
+
     async def get_param(self, param_name: str) -> ConfigParam | None:
         """Get a config parameter from the database."""
-        record: Final = await self.table.find_unique(where={"param_name": param_name})
+        record: Final = await self._config_table.find_unique(where={"param_name": param_name})
         if record is None:
             return None
         param_value = record.param_value
@@ -60,7 +93,7 @@ class ConfigRepository:
     async def set_param(self, param_name: str, param_value: Any) -> ConfigParam:
         """Set a config parameter in the database."""
         value_json: Final = json.dumps(param_value) if not isinstance(param_value, str) else param_value
-        await self.table.upsert(
+        await self._config_table.upsert(
             where={"param_name": param_name},
             data={
                 "create": {"param_name": param_name, "param_value": value_json},
@@ -72,15 +105,15 @@ class ConfigRepository:
     async def delete_param(self, param_name: str) -> bool:
         """Delete a config parameter from the database."""
         try:
-            await self.table.delete(where={"param_name": param_name})
+            await self._config_table.delete(where={"param_name": param_name})
             return True
         except Exception:
             return False
 
-    async def get_all_params(self) -> dict[str, Any]:
+    async def get_all_params(self) -> dict[str, object]:
         """Get all config parameters from the database."""
-        records: Final = await self.table.find_many()
-        result: Final = {}
+        records: Final = await self._config_table.find_many()
+        result: Final[dict[str, object]] = {}
         for record in records:
             param_value = record.param_value
             if isinstance(param_value, str):
@@ -107,7 +140,9 @@ class ConfigRepository:
                 else:
                     d[k] = v
 
-    def _decrypt_env_variables(self, env_vars: dict[str, Any], return_original_value: bool = True) -> dict[str, str]:
+    def _decrypt_env_variables(
+        self, env_vars: Mapping[str, object], return_original_value: bool = True
+    ) -> dict[str, str]:
         """Decrypt environment variables from database."""
         decrypted: Final[dict[str, str]] = {}
         for key, value in env_vars.items():

litellm/repositories/model_repository.py

@@ -3,7 +3,8 @@ Model repository for database operations on LiteLLM_ProxyModelTable.
 """
 
 import json
-from typing import Any, Final
+from collections.abc import Awaitable, Mapping, Sequence
+from typing import Any, Final, Protocol
 
 from litellm.models.model import LiteLLM_ProxyModelTable
 from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync
@@ -11,28 +12,51 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
     decrypt_value_helper,
     encrypt_value_helper,
 )
-from litellm.repositories.base_repository import BaseRepository
+from litellm.repositories.base_repository import BaseRepository, DbRecord
+
+
+class _PrismaModelDb(Protocol):
+    litellm_proxymodeltable: object
+
+
+class _PrismaClientView(Protocol):
+    db: _PrismaModelDb
+
+
+class _ProxyModelActions(Protocol):
+    """Prisma table actions used by :class:`ModelRepository`."""
+
+    def find_many(self, *, where: Mapping[str, object] | None = None) -> Awaitable[Sequence[DbRecord]]: ...
+
+    def create(self, *, data: Mapping[str, object]) -> Awaitable[DbRecord]: ...
+
+    def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[DbRecord | None]: ...
 
 
 class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
     """Repository for proxy model database operations with encryption support."""
 
-    def __init__(self, prisma_client: Any, encryption_key: str | None = None):
+    def __init__(self, prisma_client: object, encryption_key: str | None = None):
         super().__init__(prisma_client)
         self._encryption_key = encryption_key
 
     @property
     def table(self) -> Any:
+        client: Final[_PrismaClientView] = self.prisma_client
         return wrap_table_actions_for_config_sync(
-            actions=self.prisma_client.db.litellm_proxymodeltable,
+            actions=client.db.litellm_proxymodeltable,
             table_name="litellm_proxymodeltable",
         )
 
+    @property
+    def _model_table(self) -> _ProxyModelActions:
+        return self.table
+
     @property
     def model_class(self) -> type[LiteLLM_ProxyModelTable]:
         return LiteLLM_ProxyModelTable
 
-    def _encrypt_litellm_params(self, litellm_params: dict[str, Any]) -> dict[str, Any]:
+    def _encrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]:
         """Encrypt sensitive values in litellm_params."""
         encrypted: Final = {}
         for key, value in litellm_params.items():
@@ -42,7 +66,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
                 encrypted[key] = value
         return encrypted
 
-    def _decrypt_litellm_params(self, litellm_params: dict[str, Any]) -> dict[str, Any]:
+    def _decrypt_litellm_params(self, litellm_params: Mapping[str, object]) -> Mapping[str, object]:
         """Decrypt sensitive values in litellm_params."""
         decrypted: Final = {}
         for key, value in litellm_params.items():
@@ -76,17 +100,17 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
 
     async def find_by_name(self, model_name: str) -> list[LiteLLM_ProxyModelTable]:
         """Find models by name."""
-        records: Final = await self.table.find_many(where={"model_name": model_name})
+        records: Final = await self._model_table.find_many(where={"model_name": model_name})
         return self._to_model_list(records)
 
     async def find_all(self) -> list[LiteLLM_ProxyModelTable]:
         """Find all models."""
-        records: Final = await self.table.find_many()
+        records: Final = await self._model_table.find_many()
         return self._to_model_list(records)
 
     async def find_unblocked(self) -> list[LiteLLM_ProxyModelTable]:
         """Find all models that are not blocked."""
-        records: Final = await self.table.find_many(where={"blocked": False})
+        records: Final = await self._model_table.find_many(where={"blocked": False})
         return self._to_model_list(records)
 
     async def find_by_team_id(self, team_id: str) -> list[LiteLLM_ProxyModelTable]:
@@ -102,16 +126,16 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
     async def create_model(
         self,
         model_name: str,
-        litellm_params: dict[str, Any],
+        litellm_params: Mapping[str, object],
         created_by: str,
         model_id: str | None = None,
-        model_info: dict[str, Any] | None = None,
+        model_info: Mapping[str, object] | None = None,
         blocked: bool = False,
     ) -> LiteLLM_ProxyModelTable:
         """Create a new model with encryption."""
         encrypted_params: Final = self._encrypt_litellm_params(litellm_params)
 
-        data: Final[dict[str, Any]] = {
+        data: Final[dict[str, str | bool]] = {
             "model_name": model_name,
             "litellm_params": json.dumps(encrypted_params),
             "created_by": created_by,
@@ -123,7 +147,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
         if model_info is not None:
             data["model_info"] = json.dumps(model_info)
 
-        record: Final = await self.table.create(data=data)
+        record: Final = await self._model_table.create(data=data)
         model: Final = self._to_model(record)
         assert model is not None
         return model
@@ -133,12 +157,12 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
         model_id: str,
         updated_by: str,
         model_name: str | None = None,
-        litellm_params: dict[str, Any] | None = None,
-        model_info: dict[str, Any] | None = None,
+        litellm_params: Mapping[str, object] | None = None,
+        model_info: Mapping[str, object] | None = None,
         blocked: bool | None = None,
     ) -> LiteLLM_ProxyModelTable | None:
         """Update a model with encryption."""
-        data: Final[dict[str, Any]] = {"updated_by": updated_by}
+        data: Final[dict[str, str | bool]] = {"updated_by": updated_by}
         if model_name is not None:
             data["model_name"] = model_name
         if litellm_params is not None:
@@ -149,7 +173,7 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]):
         if blocked is not None:
             data["blocked"] = blocked
 
-        record: Final = await self.table.update(where={"model_id": model_id}, data=data)
+        record: Final = await self._model_table.update(where={"model_id": model_id}, data=data)
         return self._to_model(record)
 
     async def delete_model(self, model_id: str) -> LiteLLM_ProxyModelTable | None:

litellm/responses/main.py

@@ -810,7 +810,7 @@ def _responses_try_dispatch_emulated_file_search(
     extra_body: dict[str, object] | None,
     timeout: float | httpx.Timeout | None,
     custom_llm_provider: str | None,
-    kwargs: dict[str, Any],
+    kwargs: dict[str, object],
     _is_async: bool,
 ) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse] | None:
     """Return a response when emulated file_search handles the call; otherwise None."""

litellm/responses/streaming_iterator.py

@@ -79,6 +79,11 @@ def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]:  # guard-ok: verif
     return _is_json_object(value) and all(isinstance(item, str) for item in value.values())
 
 
+def _load_json_object(payload: str | bytes) -> dict[str, object]:
+    """Parse a JSON payload that the caller consumes as an object."""
+    return json.loads(payload)
+
+
 def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None:
     model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None
     model_id: Final = model_info.get("id") if _is_json_object(model_info) else None
@@ -1474,7 +1479,7 @@ class ResponsesWebSocketStreaming:
             event = event.decode("utf-8")
         if isinstance(event, str):
             try:
-                event_obj = json.loads(event)
+                event_obj = _load_json_object(event)
             except (json.JSONDecodeError, TypeError):
                 return
         else:
@@ -1487,7 +1492,7 @@ class ResponsesWebSocketStreaming:
         """Extract user input content from response.create for logging."""
         try:
             if isinstance(message, str):
-                msg_obj = json.loads(message)
+                msg_obj = _load_json_object(message)
             elif _is_json_object(message):
                 msg_obj = message
             else:
@@ -1557,7 +1562,7 @@ class ResponsesWebSocketStreaming:
                 # masked response.completed.
                 if self.output_guardrail_callbacks:
                     try:
-                        _evt_payload: Mapping[str, object] = json.loads(response_str)
+                        _evt_payload: Mapping[str, object] = _load_json_object(response_str)
                         _evt_type = _evt_payload.get("type")
                     except (json.JSONDecodeError, TypeError):
                         _evt_type = None
@@ -1622,7 +1627,7 @@ class ResponsesWebSocketStreaming:
         Non-``response.create`` messages are returned unchanged.
         """
         try:
-            msg_obj: Final[dict[str, object]] = json.loads(message)
+            msg_obj: Final = _load_json_object(message)
         except (json.JSONDecodeError, TypeError):
             return message
 
@@ -1751,7 +1756,7 @@ class ResponsesWebSocketStreaming:
             return response_str
 
         try:
-            evt_obj: Final[dict[str, object]] = json.loads(response_str)
+            evt_obj: Final = _load_json_object(response_str)
         except (json.JSONDecodeError, TypeError):
             return response_str
 
@@ -1807,7 +1812,7 @@ class ResponsesWebSocketStreaming:
             return response_str
 
         try:
-            evt_obj: Final[Mapping[str, object]] = json.loads(response_str)
+            evt_obj: Final[Mapping[str, object]] = _load_json_object(response_str)
         except (json.JSONDecodeError, TypeError):
             return response_str
 
@@ -2141,7 +2146,7 @@ class ManagedResponsesWebSocketHandler:
     async def _parse_message(self, raw_message: str) -> dict[str, object] | None:
         """Parse raw WS text; return the message dict or None (JSON error / ignored type)."""
         try:
-            msg_obj: Final[dict[str, object]] = json.loads(raw_message)
+            msg_obj: Final = _load_json_object(raw_message)
         except json.JSONDecodeError:
             await self._send_error("Invalid JSON in response.create event", "invalid_request_error")
             return None
@@ -2345,7 +2350,7 @@ class ManagedResponsesWebSocketHandler:
                 continue
             if chunk_type == "response.completed" and completed_event is None:
                 try:
-                    completed_event = json.loads(serialized)
+                    completed_event = _load_json_object(serialized)
                 except Exception:
                     pass
             try:

litellm/router.py

@@ -8573,11 +8573,9 @@ class Router:
         Returns:
         - The added/updated deployment
         """
+        _deployment_model_id: Final = deployment.model_info.id or ""
+        _deployment_on_router: Final[Deployment | None] = self.get_deployment(model_id=_deployment_model_id)
         try:
-            # check if deployment already exists
-            _deployment_model_id: Final = deployment.model_info.id or ""
-
-            _deployment_on_router: Final[Deployment | None] = self.get_deployment(model_id=_deployment_model_id)
             if _deployment_on_router is not None:
                 # deployment with this model_id exists on the router
                 if (
@@ -8628,10 +8626,31 @@ class Router:
                     deployment.model_info.id,
                     e,
                 )
+                self._restore_deployment_after_failed_upsert(
+                    previous_deployment=_deployment_on_router, model_id=_deployment_model_id
+                )
                 return None
             else:
                 raise e
 
+    def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None:
+        if previous_deployment is None or self.has_model_id(model_id):
+            return
+        try:
+            self.add_deployment(deployment=previous_deployment)
+            verbose_router_logger.info(
+                "Restored deployment %s (id=%s); it keeps serving its previous configuration.",
+                previous_deployment.model_name,
+                model_id,
+            )
+        except Exception as restore_error:  # noqa: BLE001  # best-effort restore: a second failure must not abort the reload
+            verbose_router_logger.warning(
+                "Could not restore previously served deployment %s (id=%s) after the failed upsert: %s",
+                previous_deployment.model_name,
+                model_id,
+                restore_error,
+            )
+
     @staticmethod
     def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]:
         """The ``litellm.model_cost`` keys a deployment's shared backend info is registered under."""

litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py

@@ -12,6 +12,7 @@ from __future__ import annotations
 
 import contextlib
 import contextvars
+from collections.abc import Mapping, MutableMapping
 from typing import TYPE_CHECKING, Any, Final
 
 import httpx
@@ -26,13 +27,13 @@ from litellm.utils import get_utc_datetime
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
 
-    Span = _Span | Any
+    Span = _Span
 else:
     Span = Any
 
 RoutingArgsTTL: Final = 60
 
-_io_token_rate_limit_request_kwargs: Final[contextvars.ContextVar[dict[str, Any] | None]] = contextvars.ContextVar(
+_io_token_rate_limit_request_kwargs: Final[contextvars.ContextVar[dict[str, object] | None]] = contextvars.ContextVar(
     "io_token_rate_limit_request_kwargs",
     default=None,
 )
@@ -43,7 +44,7 @@ ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key"
 OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key"
 
 
-def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, Any] | None, store_in_context: bool = True) -> None:
+def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, object] | None, store_in_context: bool = True) -> None:
     # The reservation sentinels are server-only, but `metadata` is caller
     # controlled on proxy requests. Strip any client-supplied copies here (this
     # runs before the router stashes its own reservation) so a forged
@@ -60,7 +61,7 @@ def set_io_token_rate_limit_request_kwargs(kwargs: dict[str, Any] | None, store_
     _io_token_rate_limit_request_kwargs.set(kwargs if store_in_context else None)
 
 
-def get_io_token_rate_limit_request_kwargs() -> dict[str, Any] | None:
+def get_io_token_rate_limit_request_kwargs() -> dict[str, object] | None:
     return _io_token_rate_limit_request_kwargs.get()
 
 
@@ -151,14 +152,14 @@ def _resolve_max_tokens(request_kwargs: dict[str, Any] | None, deployment: dict)
     return 4096
 
 
-def _get_usage_tokens(usage: Any) -> tuple[int, int, int]:
+def _get_usage_tokens(usage: object) -> tuple[int, int, int]:
     if usage is None:
         return 0, 0, 0
     if hasattr(usage, "prompt_tokens") or hasattr(usage, "input_tokens"):
         prompt = int(getattr(usage, "prompt_tokens", None) or getattr(usage, "input_tokens", 0) or 0)
         completion = int(getattr(usage, "completion_tokens", None) or getattr(usage, "output_tokens", 0) or 0)
         cached = 0
-        details = getattr(usage, "prompt_tokens_details", None)
+        details: object = getattr(usage, "prompt_tokens_details", None)
         if details is not None:
             cached = int(getattr(details, "cached_tokens", 0) or 0)
         if not cached:
@@ -175,13 +176,13 @@ def _get_usage_tokens(usage: Any) -> tuple[int, int, int]:
     return 0, 0, 0
 
 
-def _extract_response_usage(response_obj: Any) -> Any:
+def _extract_response_usage(response_obj: object) -> object:
     if isinstance(response_obj, dict):
         return response_obj.get("usage")
     return getattr(response_obj, "usage", None)
 
 
-def _usage_is_present(usage: Any) -> bool:
+def _usage_is_present(usage: object) -> bool:
     """
     True only if usage carries an actual input/output breakdown.
 
@@ -199,8 +200,8 @@ def _usage_is_present(usage: Any) -> bool:
 
 
 def _resolve_reconcile_usage_tokens(
-    kwargs: Any,
-    response_obj: Any,
+    kwargs: Mapping[str, object] | None,
+    response_obj: object,
 ) -> tuple[int, int, bool]:
     """
     Resolve billable input and output tokens for post-call reconcile.
@@ -233,7 +234,7 @@ def _resolve_reconcile_usage_tokens(
 
 
 def _stash_reservation_in_metadata(
-    request_kwargs: dict[str, Any] | None,
+    request_kwargs: dict[str, object] | None,
     *,
     itpm_reserved: int,
     otpm_reserved: int,
@@ -256,7 +257,7 @@ def _stash_reservation_in_metadata(
             request_kwargs[channel] = dict(reservation)
 
 
-def _extract_reservation(reservation: dict[str, Any]) -> tuple[int, int, str | None, str | None]:
+def _extract_reservation(reservation: Mapping[str, int | str | None]) -> tuple[int, int, str | None, str | None]:
     itpm_cache_key: Final = reservation.get(ITPM_CACHE_KEY)
     otpm_cache_key: Final = reservation.get(OTPM_CACHE_KEY)
     return (
@@ -267,7 +268,12 @@ def _extract_reservation(reservation: dict[str, Any]) -> tuple[int, int, str | N
     )
 
 
-def _reservation_channels(kwargs: Any) -> tuple[Any, ...]:
+def _as_mutable_mapping(value: object) -> MutableMapping[str, object] | None:
+    """``value`` when it is a dict, else ``None``."""
+    return value if isinstance(value, dict) else None
+
+
+def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object, ...]:
     """
     Places a reservation may live, in priority order: the top-level metadata
     channels win over litellm_params.metadata (so a top-level stash is never
@@ -275,30 +281,29 @@ def _reservation_channels(kwargs: Any) -> tuple[Any, ...]:
     """
     if not isinstance(kwargs, dict):
         return ()
-    channels: Final = [kwargs.get("metadata"), kwargs.get("litellm_metadata")]
-    litellm_params: Final = kwargs.get("litellm_params")
-    if isinstance(litellm_params, dict):
-        channels.append(litellm_params.get("metadata"))
-    standard_logging_object: Final = kwargs.get("standard_logging_object")
-    if isinstance(standard_logging_object, dict):
-        channels.append(standard_logging_object.get("metadata"))
-    return tuple(channels)
+    top_level: Final = (kwargs.get("metadata"), kwargs.get("litellm_metadata"))
+    litellm_params: Final = _as_mutable_mapping(kwargs.get("litellm_params"))
+    from_params: Final = () if litellm_params is None else (litellm_params.get("metadata"),)
+    standard_logging_object: Final = _as_mutable_mapping(kwargs.get("standard_logging_object"))
+    from_logging_object: Final = () if standard_logging_object is None else (standard_logging_object.get("metadata"),)
+    return top_level + from_params + from_logging_object
 
 
-def _read_reservation_from_kwargs(kwargs: Any) -> tuple[int, int, str | None, str | None]:
+def _read_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> tuple[int, int, str | None, str | None]:
     for channel_dict in _reservation_channels(kwargs):
         if isinstance(channel_dict, dict) and ITPM_RESERVED_KEY in channel_dict:
             return _extract_reservation(channel_dict)
     return 0, 0, None, None
 
 
-def _clear_reservation_from_kwargs(kwargs: Any) -> None:
+def _clear_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> None:
     """
     Remove the stashed reservation so a retry on a different (e.g. non-IO)
     deployment does not re-process the already-reconciled/refunded reservation.
     """
-    for channel_dict in _reservation_channels(kwargs):
-        if isinstance(channel_dict, dict):
+    for channel in _reservation_channels(kwargs):
+        channel_dict = _as_mutable_mapping(channel)
+        if channel_dict is not None:
             for key in (ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY):
                 channel_dict.pop(key, None)
 
@@ -524,11 +529,13 @@ def io_token_reconcile_success(
     kwargs: Any,
     response_obj: Any,
 ) -> None:
-    itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
+    request_kwargs: Final[Mapping[str, object] | None] = kwargs
+    response: Final[object] = response_obj
+    itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
     if itpm_key is None and otpm_key is None:
         return
 
-    billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(kwargs, response_obj)
+    billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response)
 
     try:
         if usage_resolved:
@@ -556,7 +563,7 @@ def io_token_reconcile_success(
                 otpm_reserved,
             )
     finally:
-        _clear_reservation_from_kwargs(kwargs)
+        _clear_reservation_from_kwargs(request_kwargs)
 
     verbose_router_logger.debug(
         "[IO TOKEN LIMIT] reconciled (usage_resolved=%s, itpm_reserved=%s, billable_input=%s, otpm_reserved=%s, output=%s)",
@@ -575,11 +582,13 @@ async def async_io_token_reconcile_success(
     *,
     parent_otel_span: Span | None = None,
 ) -> None:
-    itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
+    request_kwargs: Final[Mapping[str, object] | None] = kwargs
+    response: Final[object] = response_obj
+    itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
     if itpm_key is None and otpm_key is None:
         return
 
-    billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(kwargs, response_obj)
+    billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response)
 
     # Reconcile against the exact key that held the reservation (which encodes
     # the reservation's minute), not a key recomputed at response time. This
@@ -615,7 +624,7 @@ async def async_io_token_reconcile_success(
                 otpm_reserved,
             )
     finally:
-        _clear_reservation_from_kwargs(kwargs)
+        _clear_reservation_from_kwargs(request_kwargs)
 
     verbose_router_logger.debug(
         "[IO TOKEN LIMIT] reconciled (usage_resolved=%s, itpm_reserved=%s, billable_input=%s, otpm_reserved=%s, output=%s)",
@@ -631,7 +640,8 @@ def io_token_refund_failure(
     dual_cache: DualCache,
     kwargs: Any,
 ) -> None:
-    itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
+    request_kwargs: Final[Mapping[str, object] | None] = kwargs
+    itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
     if itpm_key is None and otpm_key is None:
         return
     if itpm_key is not None and itpm_reserved > 0:
@@ -646,11 +656,11 @@ def io_token_refund_failure(
             value=-otpm_reserved,
             ttl=RoutingArgsTTL,
         )
-    _clear_reservation_from_kwargs(kwargs)
+    _clear_reservation_from_kwargs(request_kwargs)
     verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved)
 
 
-def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: dict[str, Any] | None) -> None:
+def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping[str, object] | None) -> None:
     """
     Synchronously refund and clear any reservation a previous deployment
     attempt stashed in ``kwargs``, before it's overwritten for the next
@@ -683,7 +693,8 @@ async def async_io_token_refund_failure(
     *,
     parent_otel_span: Span | None = None,
 ) -> None:
-    itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(kwargs)
+    request_kwargs: Final[Mapping[str, object] | None] = kwargs
+    itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
     if itpm_key is None and otpm_key is None:
         return
     if itpm_key is not None and itpm_reserved > 0:
@@ -700,7 +711,7 @@ async def async_io_token_refund_failure(
             ttl=RoutingArgsTTL,
             parent_otel_span=parent_otel_span,
         )
-    _clear_reservation_from_kwargs(kwargs)
+    _clear_reservation_from_kwargs(request_kwargs)
     verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved)
 
 

ruff-strict-budget.json

@@ -1,6 +1,6 @@
 {
   "ANN001": {
-    "limit": 3026
+    "limit": 3020
   },
   "ANN002": {
     "limit": 71
@@ -12,19 +12,19 @@
     "limit": 2017
   },
   "ANN202": {
-    "limit": 855
+    "limit": 852
   },
   "ANN204": {
     "limit": 711
   },
   "ANN205": {
-    "limit": 114
+    "limit": 113
   },
   "ANN206": {
     "limit": 133
   },
   "ANN401": {
-    "limit": 1290
+    "limit": 1188
   },
   "ASYNC230": {
     "limit": 11
@@ -234,7 +234,7 @@
     "limit": 5
   },
   "TID251": {
-    "limit": 1216
+    "limit": 1212
   },
   "TRY002": {
     "limit": 524

tests/e2e/e2e_http.py

@@ -75,6 +75,8 @@ class NetworkError(BaseModel):
 
 class UnauthorizedError(BaseModel):
     kind: Literal["unauthorized"] = "unauthorized"
+    # litellm 401s for key auth, model access, and tag routing alike, so keep the body to tell them apart.
+    body: str = ""
 
 
 class RateLimitedError(BaseModel):
@@ -289,7 +291,7 @@ def _classify[R: BaseModel](
     resp: requests.Response, response_type: type[R]
 ) -> Result[R]:
     if resp.status_code == 401:
-        return UnauthorizedError()
+        return UnauthorizedError(body=resp.text)
     if resp.status_code == 429:
         return RateLimitedError(body=resp.text)
     if not resp.ok:

tests/e2e/router/test_auto_router_regressions_e2e.py

@@ -69,6 +69,7 @@ PLAIN_MODEL = "anthropic/claude-sonnet-5"
 CHEAP_MODEL = "anthropic/claude-haiku-4-5"
 STRONG_MODEL = "openai/gpt-5.6"
 MAX_TOKENS = 16
+TAG_DENIAL_MESSAGE = "Not allowed to access model due to tags configuration"
 PLAIN_SERVED = frozenset({PLAIN_MODEL, "claude-sonnet-5"})
 CHEAP_SERVED = frozenset({CHEAP_MODEL, "claude-haiku-4-5"})
 EMBEDDING_MODEL = "openai/text-embedding-3-small"
@@ -442,6 +443,9 @@ class TestUntaggedTierDeployments:
         assert isinstance(result, UnauthorizedError), (
             f"expected the tagged direct call to an untagged deployment to be denied with 401, got {result}"
         )
+        assert TAG_DENIAL_MESSAGE in result.body, (
+            f"expected the denial to come from tag routing, got a 401 reading {result.body[:300]}"
+        )
 
 
 class TestResponsesApiTagRouting:

tests/test_litellm/conftest.py

@@ -100,6 +100,12 @@ def isolate_host_aws_config(monkeypatch, isolated_aws_credentials_dir):
     monkeypatch.delenv("AWS_DEFAULT_REGION", raising=False)
 
 
+@pytest.fixture(scope="function", autouse=True)
+def isolate_host_proxy_base_url(monkeypatch):
+    """Prevent a host PROXY_BASE_URL from outranking request-derived URLs during unit tests."""
+    monkeypatch.delenv("PROXY_BASE_URL", raising=False)
+
+
 def _run_coroutine_if_needed(result):
     if not asyncio.iscoroutine(result):
         return

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py

@@ -42,6 +42,7 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
     _deserialize_json_list,
     _normalize_mcp_server_cost_info,
     _obo_retry_applies,
+    _resolve_openapi_tool_auth,
     _should_strip_caller_authorization,
     _without_authorization,
 )
@@ -10664,3 +10665,103 @@ class TestClientForwardedDiscoveryFailureIsNotFatal:
         assert resolved.registration_url == "https://idp.example.com/register"
         assert manager.config_mcp_servers[server.server_id].authorization_url == "https://idp.example.com/authorize"
         assert manager._oauth_discovery_slot(server.server_id) is None
+
+
+class TestResolveOpenapiToolAuth:
+    """The credential matrix for a ``spec_path`` server's two OpenAPI dispatch arms.
+
+    A per-server ``x-mcp-{alias}-authorization`` is already a complete header value and is forwarded
+    verbatim; a BYOK credential is a raw secret that takes the auth-type prefix. Conflating them
+    ships ``Bearer Bearer <token>``, so every cell pins which of the two a value came from.
+    """
+
+    @staticmethod
+    def _server(auth_type: MCPAuthType = MCPAuth.oauth_delegate) -> MCPServer:
+        return MCPServer(
+            server_id="srv-openapi",
+            name="report_api",
+            server_name="report_api",
+            alias="report_api",
+            url="https://api.internal.example.com",
+            transport=MCPTransport.http,
+            auth_type=auth_type,
+            spec_path="https://api.internal.example.com/openapi.json",
+            extra_headers=["X-Tenant"],
+        )
+
+    @pytest.mark.parametrize(
+        "per_server, byok, expected_auth, expected_extra_keys, expected_credential",
+        [
+            (
+                {"report_api": "Bearer caller-token"},
+                None,
+                "Bearer caller-token",
+                {"X-Tenant"},
+                "Bearer caller-token",
+            ),
+            (
+                {"report_api": {"Authorization": "Bearer caller-token", "X-Trace": "abc"}},
+                None,
+                "Bearer caller-token",
+                {"X-Tenant", "X-Trace"},
+                {"Authorization": "Bearer caller-token", "X-Trace": "abc"},
+            ),
+            (
+                {"report_api": {"X-Api-Key": "k1"}},
+                "byok-secret",
+                "Bearer byok-secret",
+                {"X-Tenant", "X-Api-Key"},
+                "byok-secret",
+            ),
+            ({"report_api": {"X-Api-Key": "k1"}}, None, None, {"X-Tenant", "X-Api-Key"}, None),
+            (None, "byok-secret", "Bearer byok-secret", {"X-Tenant"}, "byok-secret"),
+            (None, None, None, {"X-Tenant"}, None),
+            ({"other_server": "Bearer wrong"}, None, None, {"X-Tenant"}, None),
+        ],
+    )
+    def test_credential_matrix(
+        self,
+        per_server: dict | None,
+        byok: str | None,
+        expected_auth: str | None,
+        expected_extra_keys: set,
+        expected_credential: object,
+    ):
+        auth_value, forwarded, credential = _resolve_openapi_tool_auth(
+            mcp_server=self._server(),
+            mcp_auth_header=byok,
+            mcp_server_auth_headers=per_server,
+            raw_headers={"x-tenant": "acme", "authorization": "Bearer admission-key"},
+            user_api_key_auth=None,
+        )
+
+        assert auth_value == expected_auth
+        assert set(forwarded or {}) == expected_extra_keys
+        assert credential == expected_credential
+
+    def test_per_server_value_is_never_re_prefixed(self):
+        """The regression that a naive wiring produces: the caller already sent ``Bearer <token>``."""
+        auth_value, _, credential = _resolve_openapi_tool_auth(
+            mcp_server=self._server(auth_type=MCPAuth.api_key),
+            mcp_auth_header="byok-secret",
+            mcp_server_auth_headers={"report_api": "Bearer caller-token"},
+            raw_headers=None,
+            user_api_key_auth=None,
+        )
+
+        assert auth_value == "Bearer caller-token"
+        assert credential == "Bearer caller-token"
+        assert not auth_value.startswith("ApiKey ")
+
+    def test_per_server_authorization_is_not_also_left_in_forwarded_headers(self):
+        """``resolve_openapi_upstream_auth`` pops Authorization out of the forwarded headers, so a
+        second copy there would give the passthrough arm two sources to reconcile."""
+        _, forwarded, _ = _resolve_openapi_tool_auth(
+            mcp_server=self._server(),
+            mcp_auth_header=None,
+            mcp_server_auth_headers={"report_api": {"Authorization": "Bearer caller-token"}},
+            raw_headers={"x-tenant": "acme"},
+            user_api_key_auth=None,
+        )
+
+        assert "Authorization" not in (forwarded or {})

tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_tool_auth.py

@@ -189,9 +189,11 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable():
 
     pre_call = AsyncMock(return_value={})
     handle_local = AsyncMock(return_value=[])
+    resolve_auth = MagicMock()
 
     # `_get_mcp_server_from_tool_name` returns None — no server context.
     with (
+        patch.object(mcp_module, "_resolve_openapi_tool_auth", new=resolve_auth),
         patch.object(
             mcp_module.global_mcp_server_manager,
             "_get_mcp_server_from_tool_name",
@@ -228,6 +230,9 @@ async def test_openapi_local_tool_denied_when_server_not_resolvable():
     assert exc.value.status_code == 503
     pre_call.assert_not_awaited()
     handle_local.assert_not_awaited()
+    # The credential resolver takes a non-optional server, so the 503 guard above it is what keeps
+    # a missing server from ever reaching it. Pinned here so moving the guard reds this test.
+    resolve_auth.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -549,3 +554,101 @@ async def test_unknown_tool_name_still_reports_not_found():
 
     assert exc.value.status_code == 404
     assert "not found" in str(exc.value.detail)
+
+
+OPENAPI_PER_SERVER_TOKEN = "Bearer per-server-upstream-token"
+
+
+def _spec_path_server() -> MCPServer:
+    return MCPServer(
+        server_id="srv-reports",
+        name="report_api",
+        server_name="report_api",
+        alias="report_api",
+        url="https://api.internal.example.com",
+        transport=MCPTransport.http,
+        auth_type=MCPAuth.oauth_delegate,
+        spec_path="https://api.internal.example.com/openapi.json",
+    )
+
+
+@pytest.mark.parametrize("dispatch_arm", ["local_registry", "call_tool"])
+@pytest.mark.asyncio
+async def test_per_server_auth_header_reaches_both_openapi_dispatch_arms(dispatch_arm: str):
+    """`x-mcp-{alias}-authorization` must survive on BOTH OpenAPI dispatch arms.
+
+    OpenAPI tools live in the local tool registry, so `execute_mcp_tool` serves MCP-protocol and REST
+    tool calls while `MCPServerManager.call_tool` serves the responses-API handler. Both arms sourced
+    the upstream credential only from the deprecated global / BYOK `mcp_auth_header`, so the
+    per-server header was dropped and the upstream API saw no Authorization at all.
+
+    Asserting on the resolver kwarg as well as the ContextVar is deliberate: for the client-forwarded
+    modes the credential has to reach `resolve_openapi_upstream_auth`, whose passthrough arm outranks
+    the ContextVar when it materializes a header.
+    """
+    from litellm.proxy._experimental.mcp_server import server as mcp_module
+    from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
+        _request_auth_header,
+    )
+
+    server = _spec_path_server()
+    auth_headers = {"report_api": {"Authorization": OPENAPI_PER_SERVER_TOKEN}}
+    user = UserAPIKeyAuth(api_key="sk-user", user_id="alice", user_role=LitellmUserRoles.INTERNAL_USER.value)
+    captured: dict = {}
+
+    async def fake_resolver(**kwargs):
+        captured["resolver_credential"] = kwargs["mcp_auth_header"]
+        return None, kwargs["forwarded_headers"]
+
+    async def capture_local(_name, _arguments):
+        captured["injected"] = _request_auth_header.get()
+        return []
+
+    async def capture_openapi_handler(_server, _name, _arguments):
+        captured["injected"] = _request_auth_header.get()
+        return []
+
+    manager = mcp_module.global_mcp_server_manager
+    with (
+        patch.object(manager, "resolve_openapi_upstream_auth", new=fake_resolver),
+        patch.object(manager, "pre_call_tool_check", new=AsyncMock(return_value={})),
+    ):
+        if dispatch_arm == "local_registry":
+            fake_tool = MagicMock()
+            fake_tool.name = "list_reports"
+            with (
+                patch.object(manager, "_get_mcp_server_from_tool_name", return_value=server),
+                patch.object(mcp_module.global_mcp_tool_registry, "get_tool", return_value=fake_tool),
+                patch(
+                    "litellm.proxy._experimental.mcp_server.server._handle_local_mcp_tool",
+                    new=capture_local,
+                ),
+                patch(
+                    "litellm.proxy._experimental.mcp_server.server.MCPRequestHandler.is_tool_allowed",
+                    return_value=True,
+                ),
+            ):
+                await mcp_module.execute_mcp_tool(
+                    name="list_reports",
+                    arguments={},
+                    allowed_mcp_servers=[server],
+                    start_time=datetime.now(timezone.utc),
+                    user_api_key_auth=user,
+                    mcp_server_auth_headers=auth_headers,
+                )
+        else:
+            with (
+                patch.object(manager, "_resolve_mcp_server_for_tool_call", return_value=server),
+                patch.object(manager, "_call_openapi_tool_handler", new=capture_openapi_handler),
+            ):
+                await manager.call_tool(
+                    server_name="report_api",
+                    name="list_reports",
+                    arguments={},
+                    user_api_key_auth=user,
+                    mcp_server_auth_headers=auth_headers,
+                )
+
+    assert captured["resolver_credential"] == {"Authorization": OPENAPI_PER_SERVER_TOKEN}
+    assert captured["injected"] == OPENAPI_PER_SERVER_TOKEN
+    assert _request_auth_header.get() is None

tests/test_litellm/proxy/common_utils/test_registry_read_through.py

@@ -0,0 +1,476 @@
+import asyncio
+from typing import Final
+
+import pytest
+
+from litellm.proxy.common_utils.registry_read_through import RegistryReadThrough
+
+
+class ResyncSpy:
+    def __init__(self, found: bool = True, error: Exception | None = None) -> None:
+        self.found = found
+        self.error = error
+        self.calls: list[str] = []
+
+    async def __call__(self, key: str) -> bool:
+        self.calls.append(key)
+        if self.error is not None:
+            raise self.error
+        return self.found
+
+
+@pytest.mark.asyncio
+async def test_attempt_returns_true_when_resync_finds_object():
+    spy: Final = ResyncSpy(found=True)
+    read_through: Final = RegistryReadThrough(resync=spy)
+
+    assert await read_through.attempt("new-model") is True
+    assert spy.calls == ["new-model"]
+
+
+@pytest.mark.asyncio
+async def test_attempt_found_key_is_not_negative_cached():
+    spy: Final = ResyncSpy(found=True)
+    read_through: Final = RegistryReadThrough(resync=spy)
+
+    assert await read_through.attempt("new-model") is True
+    assert await read_through.attempt("new-model") is True
+    assert spy.calls == ["new-model", "new-model"]
+
+
+@pytest.mark.asyncio
+async def test_missing_key_is_negative_cached_within_ttl():
+    spy: Final = ResyncSpy(found=False)
+    read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0)
+
+    assert await read_through.attempt("ghost-model") is False
+    assert await read_through.attempt("ghost-model") is False
+    assert spy.calls == ["ghost-model"]
+
+
+@pytest.mark.asyncio
+async def test_negative_cache_expires_and_resync_runs_again():
+    spy: Final = ResyncSpy(found=False)
+    read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=0.05)
+
+    assert await read_through.attempt("ghost-model") is False
+    await asyncio.sleep(0.1)
+    assert await read_through.attempt("ghost-model") is False
+    assert spy.calls == ["ghost-model", "ghost-model"]
+
+
+@pytest.mark.asyncio
+async def test_resync_exception_returns_false_without_negative_caching():
+    spy: Final = ResyncSpy(error=RuntimeError("db down"))
+    read_through: Final = RegistryReadThrough(resync=spy)
+
+    assert await read_through.attempt("new-model") is False
+    assert await read_through.attempt("new-model") is False
+    assert spy.calls == ["new-model", "new-model"]
+
+
+@pytest.mark.asyncio
+async def test_concurrent_attempts_for_missing_key_resync_once():
+    class SlowResyncSpy(ResyncSpy):
+        async def __call__(self, key: str) -> bool:
+            await asyncio.sleep(0.05)
+            return await super().__call__(key)
+
+    spy: Final = SlowResyncSpy(found=False)
+    read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0)
+
+    results: Final = await asyncio.gather(*(read_through.attempt("ghost-model") for _ in range(5)))
+    assert results == [False] * 5
+    assert spy.calls == ["ghost-model"]
+
+
+@pytest.mark.asyncio
+async def test_distinct_keys_do_not_share_negative_cache():
+    spy: Final = ResyncSpy(found=False)
+    read_through: Final = RegistryReadThrough(resync=spy, miss_ttl_seconds=60.0)
+
+    assert await read_through.attempt("ghost-a") is False
+    assert await read_through.attempt("ghost-b") is False
+    assert spy.calls == ["ghost-a", "ghost-b"]
+
+
+@pytest.mark.asyncio
+async def test_resync_budget_exhausted_blocks_resync_without_negative_caching():
+    spy: Final = ResyncSpy(found=False)
+    read_through: Final = RegistryReadThrough(
+        resync=spy, miss_ttl_seconds=60.0, max_resyncs_per_window=2, resync_window_seconds=60.0
+    )
+
+    assert await read_through.attempt("ghost-a") is False
+    assert await read_through.attempt("ghost-b") is False
+    assert await read_through.attempt("ghost-c") is False
+    assert spy.calls == ["ghost-a", "ghost-b"]
+    assert read_through._recent_misses.get_cache("ghost-c") is None
+
+
+@pytest.mark.asyncio
+async def test_resync_budget_replenishes_after_window():
+    spy: Final = ResyncSpy(found=True)
+    read_through: Final = RegistryReadThrough(resync=spy, max_resyncs_per_window=1, resync_window_seconds=0.05)
+
+    assert await read_through.attempt("model-a") is True
+    assert await read_through.attempt("model-b") is False
+    await asyncio.sleep(0.1)
+    assert await read_through.attempt("model-b") is True
+    assert spy.calls == ["model-a", "model-b"]
+
+
+class FakeAgentRow:
+    def __init__(self, agent_id: str, agent_name: str) -> None:
+        self.agent_id = agent_id
+        self.agent_name = agent_name
+        self.object_permission = None
+        self.spend = 0.0
+
+    def model_dump(self):
+        return {
+            "agent_id": self.agent_id,
+            "agent_name": self.agent_name,
+            "agent_card_params": {"name": self.agent_name, "url": "http://db-agent"},
+            "litellm_params": {},
+            "object_permission": None,
+            "spend": self.spend,
+        }
+
+
+@pytest.fixture
+def clean_agent_registry():
+    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+
+    original_agents: Final = list(global_agent_registry.agent_list)
+    original_config_agents: Final = getattr(global_agent_registry, "config_agents", ())
+    global_agent_registry.agent_list = []
+    global_agent_registry.config_agents = ()
+    try:
+        yield global_agent_registry
+    finally:
+        global_agent_registry.agent_list = original_agents
+        global_agent_registry.config_agents = original_config_agents
+
+
+@pytest.mark.asyncio
+async def test_get_agent_with_read_through_recovers_agent_created_on_sibling_replica(
+    clean_agent_registry, monkeypatch
+):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
+
+    agent_id: Final = "read-through-db-agent-id"
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
+        return_value=FakeAgentRow(agent_id, "read-through-db-agent")
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+
+    assert clean_agent_registry.get_agent_by_id(agent_id=agent_id) is None
+    agent: Final = await get_agent_with_read_through(agent_id)
+
+    assert agent is not None
+    assert agent.agent_id == agent_id
+    prisma_client.db.litellm_agentstable.find_unique.assert_awaited_once_with(
+        where={"agent_id": agent_id},
+        include={"object_permission": True},
+    )
+
+
+@pytest.mark.asyncio
+async def test_get_agent_with_read_through_recovers_agent_by_name(clean_agent_registry, monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
+
+    agent_name: Final = "read-through-db-agent-by-name"
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
+        side_effect=[None, FakeAgentRow("read-through-name-lookup-id", agent_name)]
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+
+    agent: Final = await get_agent_with_read_through(agent_name)
+
+    assert agent is not None
+    assert agent.agent_name == agent_name
+    prisma_client.db.litellm_agentstable.find_unique.assert_awaited_with(
+        where={"agent_name": agent_name},
+        include={"object_permission": True},
+    )
+
+
+@pytest.mark.asyncio
+async def test_get_agent_with_read_through_returns_none_for_unknown_agent(clean_agent_registry, monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through
+
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+
+    assert await get_agent_with_read_through("agent-nobody-created") is None
+    assert prisma_client.db.litellm_agentstable.find_unique.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_resync_agents_already_registered_skips_db(clean_agent_registry, monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import _resync_agents
+
+    agent_id: Final = "read-through-dedup-agent-id"
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
+        return_value=FakeAgentRow(agent_id, "read-through-dedup-agent")
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+
+    assert await _resync_agents(agent_id) is True
+    assert await _resync_agents(agent_id) is True
+    assert prisma_client.db.litellm_agentstable.find_unique.await_count == 1
+    assert len(clean_agent_registry.agent_list) == 1
+
+
+class FakeGuardrailRow:
+    def __init__(self, guardrail_id: str, guardrail_name: str) -> None:
+        self.guardrail_id = guardrail_id
+        self.guardrail_name = guardrail_name
+
+    def __iter__(self):
+        return iter(
+            {
+                "guardrail_id": self.guardrail_id,
+                "guardrail_name": self.guardrail_name,
+                "litellm_params": {
+                    "guardrail": "litellm_content_filter",
+                    "mode": "pre_call",
+                    "default_on": True,
+                    "blocked_words": [{"keyword": "secret", "action": "BLOCK"}],
+                },
+                "guardrail_info": {},
+                "status": "active",
+            }.items()
+        )
+
+
+@pytest.mark.asyncio
+async def test_get_guardrail_with_read_through_recovers_guardrail_created_on_sibling_replica(monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import (
+        get_initialized_guardrail_with_read_through,
+    )
+    from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
+
+    guardrail_id: Final = "read-through-db-guardrail-id"
+    guardrail_name: Final = "read-through-db-guardrail"
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(
+        return_value=FakeGuardrailRow(guardrail_id, guardrail_name)
+    )
+    prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(
+        side_effect=AssertionError("full-table guardrail scan on read-through miss")
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+
+    try:
+        guardrail: Final = await get_initialized_guardrail_with_read_through(guardrail_name=guardrail_name)
+        assert guardrail is not None
+        assert guardrail.guardrail_name == guardrail_name
+        prisma_client.db.litellm_guardrailstable.find_first.assert_awaited_once_with(
+            where={"guardrail_name": guardrail_name, "status": "active"}
+        )
+    finally:
+        IN_MEMORY_GUARDRAIL_HANDLER.delete_in_memory_guardrail(guardrail_id)
+
+
+@pytest.mark.asyncio
+async def test_get_guardrail_with_read_through_returns_none_for_unknown_guardrail(monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import (
+        get_initialized_guardrail_with_read_through,
+    )
+
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(return_value=None)
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+
+    assert await get_initialized_guardrail_with_read_through(guardrail_name="guardrail-nobody-created") is None
+
+
+@pytest.mark.asyncio
+async def test_resync_guardrails_never_loads_non_active_rows(monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import _resync_guardrails
+
+    pending_name: Final = "pending-review-guardrail"
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(return_value=None)
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+
+    assert await _resync_guardrails(pending_name) is False
+    prisma_client.db.litellm_guardrailstable.find_first.assert_awaited_once_with(
+        where={"guardrail_name": pending_name, "status": "active"}
+    )
+
+
+@pytest.mark.asyncio
+async def test_resync_guardrails_syncs_under_guardrail_reconcile_lock(monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.common_utils.registry_read_through as read_through_module
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import _resync_guardrails
+    from litellm.proxy.guardrails.guardrail_registry import (
+        GUARDRAIL_RECONCILE_LOCK,
+        IN_MEMORY_GUARDRAIL_HANDLER,
+    )
+
+    guardrail_name: Final = "lock-scope-guardrail"
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_guardrailstable.find_first = AsyncMock(
+        return_value=FakeGuardrailRow("lock-scope-guardrail-id", guardrail_name)
+    )
+    lock_states: list[bool] = []
+
+    def record_sync(guardrail) -> None:
+        lock_states.append(GUARDRAIL_RECONCILE_LOCK.locked())
+
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+    monkeypatch.setattr(IN_MEMORY_GUARDRAIL_HANDLER, "sync_guardrail_from_db", record_sync)
+    monkeypatch.setattr(read_through_module, "_initialized_guardrail", lambda guardrail_name: MagicMock())
+
+    assert await _resync_guardrails(guardrail_name) is True
+    assert lock_states == [True]
+    assert not GUARDRAIL_RECONCILE_LOCK.locked()
+
+
+@pytest.mark.asyncio
+async def test_resync_model_deployments_mutates_router_under_model_reconcile_lock(monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments
+
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[MagicMock()])
+    router: Final = MagicMock()
+    router.get_model_list.return_value = []
+    lock_states: list[bool] = []
+
+    def record_add_deployment(db_models) -> None:
+        lock_states.append(proxy_server.MODEL_RECONCILE_LOCK.locked())
+
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+    monkeypatch.setattr(proxy_server, "llm_model_list", None)
+    monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", record_add_deployment)
+
+    assert await _resync_model_deployments("lock-scope-model") is True
+    assert lock_states == [True]
+    assert not proxy_server.MODEL_RECONCILE_LOCK.locked()
+
+
+@pytest.mark.asyncio
+async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments
+
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(
+        side_effect=AssertionError("db hit for an object type this replica does not load")
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+    monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["guardrails"]})
+
+    assert await _resync_model_deployments("gated-out-model") is False
+
+
+@pytest.mark.asyncio
+async def test_resync_guardrails_respects_supported_db_objects(monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import _resync_guardrails
+
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock(
+        side_effect=AssertionError("db hit for an object type this replica does not load")
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+    monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]})
+
+    assert await _resync_guardrails("gated-out-guardrail") is False
+
+
+@pytest.mark.asyncio
+async def test_resync_agents_respects_supported_db_objects(clean_agent_registry, monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.common_utils.registry_read_through import _resync_agents
+
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
+        side_effect=AssertionError("db hit for an object type this replica does not load")
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+    monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]})
+
+    assert await _resync_agents("gated-out-agent") is False
+
+
+@pytest.mark.asyncio
+async def test_resync_agents_waits_for_agent_reload_and_skips_duplicate_registration(clean_agent_registry, monkeypatch):
+    from unittest.mock import AsyncMock, MagicMock
+
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.agent_endpoints.agent_registry import AGENT_RECONCILE_LOCK
+    from litellm.proxy.common_utils.registry_read_through import _resync_agents
+    from litellm.types.agents import AgentResponse
+
+    agent_id: Final = "reload-race-agent-id"
+    prisma_client: Final = MagicMock()
+    prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
+        side_effect=AssertionError("db hit while the agent reload held the reconcile lock")
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+
+    async with AGENT_RECONCILE_LOCK:
+        resync_task: Final = asyncio.ensure_future(_resync_agents(agent_id))
+        await asyncio.sleep(0.05)
+        assert not resync_task.done()
+        clean_agent_registry.register_agent(
+            agent_config=AgentResponse.model_validate(FakeAgentRow(agent_id, "reload-race-agent").model_dump())
+        )
+
+    assert await resync_task is True
+    assert len(clean_agent_registry.agent_list) == 1

tests/test_litellm/proxy/management_endpoints/test_access_group_management.py

@@ -13,6 +13,9 @@ sys.path.insert(
 )  # Adds the parent directory to the system path
 
 from litellm import Router
+from litellm.proxy.management_endpoints.model_management_endpoints import (
+    ReconcileOutcome,
+)
 
 
 @pytest.mark.asyncio
@@ -121,7 +124,7 @@ async def test_create_access_group_with_model_ids_tags_only_specific_deployments
         patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
         patch(
             "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
-            new_callable=AsyncMock,
+            new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
         ),
     ):
         response = await create_model_group(
@@ -186,7 +189,7 @@ async def test_create_access_group_with_model_names_tags_all_deployments():
         patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
         patch(
             "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
-            new_callable=AsyncMock,
+            new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
         ),
     ):
         response = await create_model_group(
@@ -236,7 +239,7 @@ async def test_create_access_group_model_ids_takes_priority_over_model_names():
         patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
         patch(
             "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
-            new_callable=AsyncMock,
+            new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
         ),
     ):
         response = await create_model_group(
@@ -313,7 +316,7 @@ async def test_create_access_group_invalid_model_id_returns_400():
         patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
         patch(
             "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
-            new_callable=AsyncMock,
+            new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
         ),
     ):
         with pytest.raises(HTTPException) as exc_info:
@@ -352,7 +355,7 @@ async def test_create_access_group_surfaces_dropped_models():
         patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
         patch(
             "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
-            new=AsyncMock(return_value=None),
+            new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
         ),
     ):
         with pytest.raises(HTTPException) as exc_info:
@@ -365,6 +368,50 @@ async def test_create_access_group_surfaces_dropped_models():
     assert "deploy-A" in str(exc_info.value.detail)
 
 
+
+@pytest.mark.asyncio
+async def test_create_access_group_trusts_reload_snapshot_over_post_lock_fresh_read():
+    """A concurrent reconcile sampled after the lock is released must not make this
+    write's reload look like it dropped the tagged model: the verdict has to judge from
+    the ReconcileOutcome the reload captured under the lock, not a fresh router read."""
+    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+    from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
+        create_model_group,
+    )
+    from litellm.types.proxy.management_endpoints.model_management_endpoints import (
+        NewModelGroupRequest,
+    )
+
+    deploy_a = MagicMock(model_id="deploy-A", model_name="gpt-4o", model_info={})
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
+    mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a)
+    mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
+
+    concurrently_wiped_router = MagicMock()
+    concurrently_wiped_router.get_model_ids.side_effect = [["deploy-A"], []]
+    with (
+        patch("litellm.proxy.proxy_server.llm_router", concurrently_wiped_router),
+        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+        patch(
+            "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
+            new=AsyncMock(
+                return_value=ReconcileOutcome(
+                    still_desired=frozenset({"deploy-A"}), live_after=frozenset({"deploy-A"})
+                )
+            ),
+        ),
+    ):
+        response = await create_model_group(
+            data=NewModelGroupRequest(access_group="production-models", model_ids=["deploy-A"]),
+            user_api_key_dict=UserAPIKeyAuth(user_id="test_admin", user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
+
+    assert response.models_updated == 1
+    assert concurrently_wiped_router.get_model_ids.call_count == 1
+
+
 @pytest.mark.asyncio
 async def test_tag_deployment_parses_string_model_info_and_refuses_corrupt():
     """The model_info column can arrive as its JSON string; tagging must parse it rather
@@ -420,7 +467,7 @@ async def test_delete_access_group_ignores_models_that_were_already_dead():
         patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
         patch(
             "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
-            new=AsyncMock(return_value=None),
+            new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
         ),
     ):
         response = await delete_access_group(
@@ -430,3 +477,99 @@ async def test_delete_access_group_ignores_models_that_were_already_dead():
 
     assert response.models_updated == 1
     mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_create_access_group_read_through_recovers_model_created_on_sibling_replica():
+    """Regression: an access group referencing a model that another replica just wrote
+    to the DB must be created instead of 400ing until the periodic config reload."""
+    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+    from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
+        create_model_group,
+    )
+    from litellm.types.proxy.management_endpoints.model_management_endpoints import (
+        NewModelGroupRequest,
+    )
+
+    from types import SimpleNamespace
+
+    model_name = "e2e-ag-sibling-replica-model"
+    db_row = SimpleNamespace(
+        model_id=f"{model_name}-id",
+        model_name=model_name,
+        litellm_params={"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "hi"},
+        model_info={},
+        blocked=False,
+    )
+
+    mock_router = Router(
+        model_list=[
+            {
+                "model_name": "some-other-model",
+                "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
+            }
+        ]
+    )
+
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=[[db_row], [], [db_row]])
+    mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
+
+    with (
+        patch("litellm.proxy.proxy_server.llm_router", mock_router),
+        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+        patch("litellm.proxy.proxy_server.store_model_in_db", True),
+        patch(
+            "litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
+            new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)),
+        ),
+    ):
+        response = await create_model_group(
+            data=NewModelGroupRequest(access_group="replica-lag-group", model_names=[model_name]),
+            user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
+
+    assert response.models_updated == 1
+    assert response.model_names == [model_name]
+    assert mock_prisma.db.litellm_proxymodeltable.find_many.await_args_list[0].kwargs["where"] == {
+        "model_name": model_name
+    }
+
+
+@pytest.mark.asyncio
+async def test_create_access_group_model_missing_everywhere_still_400s():
+    from fastapi import HTTPException
+
+    from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+    from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
+        create_model_group,
+    )
+    from litellm.types.proxy.management_endpoints.model_management_endpoints import (
+        NewModelGroupRequest,
+    )
+
+    model_name = "e2e-ag-model-nobody-created"
+    mock_router = Router(
+        model_list=[
+            {
+                "model_name": "some-other-model",
+                "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
+            }
+        ]
+    )
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
+
+    with (
+        patch("litellm.proxy.proxy_server.llm_router", mock_router),
+        patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
+        patch("litellm.proxy.proxy_server.store_model_in_db", True),
+    ):
+        with pytest.raises(HTTPException) as exc_info:
+            await create_model_group(
+                data=NewModelGroupRequest(access_group="ghost-group", model_names=[model_name]),
+                user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
+            )
+
+    assert exc_info.value.status_code == 400
+    assert model_name in str(exc_info.value.detail)

tests/test_litellm/proxy/test_proxy_server.py

@@ -11100,6 +11100,60 @@ async def test_moderations_reraises_proxy_exception_unwrapped():
     mock_logging.post_call_failure_hook.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_init_agents_in_db_rebuilds_registry_under_agent_reconcile_lock(monkeypatch):
+    from litellm.proxy.agent_endpoints.agent_registry import (
+        AGENT_RECONCILE_LOCK,
+        global_agent_registry,
+    )
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    lock_states: list[bool] = []
+
+    async def fake_get_all_agents_from_db(prisma_client) -> list:
+        lock_states.append(AGENT_RECONCILE_LOCK.locked())
+        return []
+
+    def fake_load_agents_from_db_and_config(db_agents) -> None:
+        lock_states.append(AGENT_RECONCILE_LOCK.locked())
+
+    monkeypatch.setattr(global_agent_registry, "get_all_agents_from_db", fake_get_all_agents_from_db)
+    monkeypatch.setattr(global_agent_registry, "load_agents_from_db_and_config", fake_load_agents_from_db_and_config)
+
+    await ProxyConfig()._init_agents_in_db(prisma_client=MagicMock())
+
+    assert lock_states == [True, True]
+    assert not AGENT_RECONCILE_LOCK.locked()
+
+
+@pytest.mark.asyncio
+async def test_init_guardrails_in_db_snapshots_and_reconciles_under_guardrail_reconcile_lock(monkeypatch):
+    from litellm.proxy.guardrails.guardrail_registry import (
+        GUARDRAIL_RECONCILE_LOCK,
+        IN_MEMORY_GUARDRAIL_HANDLER,
+        GuardrailRegistry,
+    )
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    lock_states: list[bool] = []
+
+    async def fake_get_all_guardrails_from_db(prisma_client) -> list:
+        lock_states.append(GUARDRAIL_RECONCILE_LOCK.locked())
+        return []
+
+    def fake_reconcile_db_guardrails(db_guardrail_ids) -> list:
+        lock_states.append(GUARDRAIL_RECONCILE_LOCK.locked())
+        return []
+
+    monkeypatch.setattr(GuardrailRegistry, "get_all_guardrails_from_db", fake_get_all_guardrails_from_db)
+    monkeypatch.setattr(IN_MEMORY_GUARDRAIL_HANDLER, "reconcile_db_guardrails", fake_reconcile_db_guardrails)
+
+    await ProxyConfig()._init_guardrails_in_db(prisma_client=MagicMock())
+
+    assert lock_states == [True, True]
+    assert not GUARDRAIL_RECONCILE_LOCK.locked()
+
+
 class TestEmbeddingsFailureHookRequestData:
     @pytest.mark.asyncio
     async def test_failure_hook_gets_post_setup_data_with_logging_obj(self):

tests/test_litellm/proxy/test_route_a2a_models.py

@@ -50,6 +50,7 @@ async def test_route_a2a_model_bypasses_router():
     )
 
     mock_registry = Mock()
+    mock_registry.get_agent_by_id = Mock(return_value=None)
     mock_registry.get_agent_by_name = Mock(return_value=mock_agent)
 
     # Mock litellm.acompletion to verify it's called
@@ -106,3 +107,79 @@ async def test_route_non_a2a_model_raises_error_if_not_in_router():
             user_model=None,
             route_type="acompletion",
         )
+
+
+class _DbAgentRow:
+    def __init__(self, agent_id: str, agent_name: str) -> None:
+        self.agent_id = agent_id
+        self.agent_name = agent_name
+        self.object_permission = None
+        self.spend = 0.0
+
+    def model_dump(self):
+        return {
+            "agent_id": self.agent_id,
+            "agent_name": self.agent_name,
+            "agent_card_params": {"name": self.agent_name, "url": "http://sibling-db-agent.example.com"},
+            "litellm_params": {},
+            "object_permission": None,
+            "spend": self.spend,
+        }
+
+
+def _router_without_models():
+    mock_router = Mock()
+    mock_router.model_names = []
+    mock_router.deployment_names = []
+    mock_router.has_model_id = Mock(return_value=False)
+    mock_router.model_group_alias = None
+    mock_router.router_general_settings = Mock(pass_through_all_models=False)
+    mock_router.default_deployment = None
+    mock_router.pattern_router = Mock(patterns=[])
+    mock_router.map_team_model = Mock(return_value=None)
+    mock_router.is_recognized_model = Mock(return_value=False)
+    mock_router.team_public_model_names = []
+    return mock_router
+
+
+@pytest.mark.asyncio
+async def test_route_a2a_model_read_through_recovers_agent_created_on_sibling_replica(monkeypatch):
+    import litellm.proxy.proxy_server as proxy_server
+    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
+
+    agent_name = "a2a-sibling-replica-agent"
+    prisma_client = Mock()
+    prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
+        side_effect=[None, _DbAgentRow("a2a-sibling-replica-agent-id", agent_name)]
+    )
+    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+
+    original_agents = list(global_agent_registry.agent_list)
+    original_config_agents = getattr(global_agent_registry, "config_agents", ())
+    global_agent_registry.agent_list = []
+    global_agent_registry.config_agents = ()
+
+    data = {
+        "model": f"a2a/{agent_name}",
+        "messages": [{"role": "user", "content": "Hello"}],
+    }
+    mock_acompletion = AsyncMock(return_value={"id": "read-through-response"})
+
+    try:
+        with patch("litellm.acompletion", mock_acompletion):
+            await route_request(
+                data=data,
+                llm_router=_router_without_models(),
+                user_model=None,
+                route_type="acompletion",
+            )
+    finally:
+        global_agent_registry.agent_list = original_agents
+        global_agent_registry.config_agents = original_config_agents
+
+    mock_acompletion.assert_called_once()
+    call_kwargs = mock_acompletion.call_args.kwargs
+    assert call_kwargs["model"] == f"a2a/{agent_name}"
+    assert call_kwargs["api_base"] == "http://sibling-db-agent.example.com"
+    prisma_client.db.litellm_agentstable.find_unique.assert_awaited()

tests/test_litellm/proxy/test_route_llm_request.py

@@ -1119,6 +1119,126 @@ async def test_route_request_rejects_chat_completion_without_messages():
     llm_router.acompletion.assert_not_called()
 
 
+class FakeProxyModelTable:
+    def __init__(self, rows):
+        self.rows = rows
+        self.find_many_wheres = []
+
+    async def find_many(self, where=None, **kwargs):
+        self.find_many_wheres.append(where)
+        return list(self.rows)
+
+
+def _fake_prisma_client_with_models(rows):
+    from types import SimpleNamespace
+
+    table = FakeProxyModelTable(rows)
+    return SimpleNamespace(db=SimpleNamespace(litellm_proxymodeltable=table)), table
+
+
+def _db_model_row(model_name: str, mock_response: str):
+    from types import SimpleNamespace
+
+    return SimpleNamespace(
+        model_id=f"{model_name}-id",
+        model_name=model_name,
+        litellm_params={"model": "openai/gpt-4o", "api_key": "fake", "mock_response": mock_response},
+        model_info={},
+        blocked=False,
+    )
+
+
+@pytest.mark.asyncio
+async def test_route_request_read_through_recovers_model_created_on_sibling_replica(monkeypatch):
+    """Regression: a model written to the DB by another replica must be served on
+    first request instead of 400ing until the periodic config reload."""
+    import litellm
+    import litellm.proxy.proxy_server as proxy_server
+
+    model_name = "e2e-sibling-replica-model"
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "some-other-model",
+                "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
+            }
+        ]
+    )
+    fake_prisma, table = _fake_prisma_client_with_models([_db_model_row(model_name, "hello-from-db")])
+    monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+
+    llm_call = await route_request(
+        data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]},
+        llm_router=router,
+        user_model=None,
+        route_type="acompletion",
+    )
+    response = await llm_call
+
+    assert response.choices[0].message.content == "hello-from-db"
+    assert len(table.find_many_wheres) == 1
+    assert table.find_many_wheres[0] == {"model_name": model_name}
+
+
+@pytest.mark.asyncio
+async def test_route_request_unknown_model_raises_and_hits_db_once_within_ttl(monkeypatch):
+    import litellm
+    import litellm.proxy.proxy_server as proxy_server
+
+    model_name = "e2e-model-nobody-created"
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "some-other-model",
+                "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
+            }
+        ]
+    )
+    fake_prisma, table = _fake_prisma_client_with_models([])
+    monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+
+    data = {"model": model_name, "messages": [{"role": "user", "content": "hi"}]}
+    with pytest.raises(ProxyModelNotFoundError):
+        await route_request(data=data, llm_router=router, user_model=None, route_type="acompletion")
+    with pytest.raises(ProxyModelNotFoundError):
+        await route_request(data=data, llm_router=router, user_model=None, route_type="acompletion")
+
+    assert table.find_many_wheres == [{"model_name": model_name}, {"model_id": model_name}]
+
+
+@pytest.mark.asyncio
+async def test_route_request_read_through_disabled_without_store_model_in_db(monkeypatch):
+    import litellm
+    import litellm.proxy.proxy_server as proxy_server
+
+    model_name = "e2e-config-only-proxy-model"
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "some-other-model",
+                "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
+            }
+        ]
+    )
+    fake_prisma, table = _fake_prisma_client_with_models([_db_model_row(model_name, "should-not-load")])
+    monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", False)
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+
+    with pytest.raises(ProxyModelNotFoundError):
+        await route_request(
+            data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]},
+            llm_router=router,
+            user_model=None,
+            route_type="acompletion",
+        )
+
+    assert table.find_many_wheres == []
+
 @pytest.mark.asyncio
 async def test_route_request_routing_group_name_passes_model_gate():
     from unittest.mock import AsyncMock, patch
@@ -1141,3 +1261,39 @@ async def test_route_request_routing_group_name_passes_model_gate():
 
     assert response == "group_response"
     spy.assert_called_once_with(**data)
+
+
+@pytest.mark.asyncio
+async def test_route_request_a2a_agent_miss_does_not_consume_model_read_through(monkeypatch):
+    from types import SimpleNamespace
+    from unittest.mock import AsyncMock
+
+    import litellm
+    import litellm.proxy.proxy_server as proxy_server
+
+    model_name = "a2a/agent-nobody-created"
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "some-other-model",
+                "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"},
+            }
+        ]
+    )
+    fake_prisma, model_table = _fake_prisma_client_with_models([])
+    agents_find_unique = AsyncMock(return_value=None)
+    fake_prisma.db.litellm_agentstable = SimpleNamespace(find_unique=agents_find_unique)
+    monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
+    monkeypatch.setattr(proxy_server, "store_model_in_db", True)
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+
+    with pytest.raises(ProxyModelNotFoundError):
+        await route_request(
+            data={"model": model_name, "messages": [{"role": "user", "content": "hi"}]},
+            llm_router=router,
+            user_model=None,
+            route_type="acompletion",
+        )
+
+    assert agents_find_unique.await_count == 2
+    assert model_table.find_many_wheres == []

tests/test_litellm/test_conftest.py

@@ -0,0 +1,43 @@
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Final
+
+REPO_ROOT: Final = Path(__file__).resolve().parents[2]
+
+PROXY_BASE_URL_SENSITIVE_NODE: Final = (
+    "tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py"
+    "::TestTemporaryMCPSessionEndpoints"
+    "::test_mcp_token_opens_sealed_passthrough_code_and_exchanges_with_minted_client"
+)
+
+COVERAGE_SUBPROCESS_VARS: Final = frozenset(
+    {"COV_CORE_SOURCE", "COV_CORE_CONFIG", "COV_CORE_DATAFILE", "COV_CORE_CONTEXT", "COVERAGE_PROCESS_START"}
+)
+
+
+def test_host_proxy_base_url_cannot_reach_request_derived_url_tests():
+    child_env: Final = {
+        key: value for key, value in os.environ.items() if key not in COVERAGE_SUBPROCESS_VARS
+    } | {"PROXY_BASE_URL": "https://leaked-host-origin.example.com"}
+
+    completed: Final = subprocess.run(
+        [
+            sys.executable,
+            "-m",
+            "pytest",
+            PROXY_BASE_URL_SENSITIVE_NODE,
+            "-q",
+            "--no-header",
+            "-p",
+            "no:cacheprovider",
+        ],
+        cwd=REPO_ROOT,
+        env=child_env,
+        capture_output=True,
+        text=True,
+        check=False,
+    )
+
+    assert completed.returncode == 0, completed.stdout + completed.stderr

tests/test_litellm/test_router.py

@@ -7643,6 +7643,95 @@ def test_pre_call_checks_keeps_deployment_when_provider_is_unresolvable(monkeypa
     assert len(result) == 1
 
 
+class TestUpsertDeploymentRollback:
+    """
+    Regression tests: `upsert_deployment` pops the previous deployment before
+    re-adding the edited one. When the re-add raises under
+    `ignore_invalid_deployments=True`, the pop must be rolled back so this pod
+    keeps serving the previous configuration instead of silently dropping a live
+    deployment (the "Error upserting deployment" drop behind the access-group
+    reload 500 in the 2-replica e2e suite).
+    """
+
+    def test_failed_upsert_keeps_previous_deployment_serving(self):
+        from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
+
+        router = litellm.Router(
+            model_list=[
+                {
+                    "model_name": "prod-model",
+                    "litellm_params": {"model": "gpt-4o", "api_key": "sk-old"},
+                    "model_info": {"id": "prod-1", "db_model": True},
+                }
+            ],
+            ignore_invalid_deployments=True,
+        )
+
+        result = router.upsert_deployment(
+            deployment=Deployment(
+                model_name="prod-model",
+                litellm_params=LiteLLM_Params(model="auto_router/broken"),
+                model_info=ModelInfo(id="prod-1", db_model=True),
+            )
+        )
+
+        assert result is None
+        restored = router.get_deployment(model_id="prod-1")
+        assert restored is not None
+        assert restored.litellm_params.model == "gpt-4o"
+        assert [model["model_name"] for model in router.model_list] == ["prod-model"]
+
+    def test_failed_fresh_add_returns_none_without_restore(self):
+        from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
+
+        router = litellm.Router(model_list=[], ignore_invalid_deployments=True)
+
+        result = router.upsert_deployment(
+            deployment=Deployment(
+                model_name="fresh-router",
+                litellm_params=LiteLLM_Params(model="auto_router/broken"),
+                model_info=ModelInfo(id="fresh-1", db_model=True),
+            )
+        )
+
+        assert result is None
+        assert router.get_deployment(model_id="fresh-1") is None
+        assert router.model_list == []
+
+    def test_restore_re_adds_popped_deployment(self):
+        router = litellm.Router(
+            model_list=[
+                {
+                    "model_name": "prod-model",
+                    "litellm_params": {"model": "gpt-4o", "api_key": "sk-old"},
+                    "model_info": {"id": "prod-1", "db_model": True},
+                }
+            ],
+            ignore_invalid_deployments=True,
+        )
+        previous = router.get_deployment(model_id="prod-1")
+        router.delete_deployment(id="prod-1")
+        assert router.has_model_id("prod-1") is False
+
+        router._restore_deployment_after_failed_upsert(
+            previous_deployment=previous, model_id="prod-1"
+        )
+
+        restored = router.get_deployment(model_id="prod-1")
+        assert restored is not None
+        assert restored.litellm_params.model == "gpt-4o"
+
+        router._restore_deployment_after_failed_upsert(
+            previous_deployment=previous, model_id="prod-1"
+        )
+        assert len(router.model_list) == 1
+
+        router._restore_deployment_after_failed_upsert(
+            previous_deployment=None, model_id="prod-1"
+        )
+        assert len(router.model_list) == 1
+
+
 class TestConsumedRequestTagsStamp:
     """Issue #36621: when a request's tags select a tagged pre-routing strategy, those
     tags are consumed by the selection; the hook must stamp the rewritten model group so

type-discipline-budget.json

@@ -1,9 +1,9 @@
 {
   "LIT001": {
-    "limit": 22894
+    "limit": 22809
   },
   "LIT002": {
-    "limit": 26888
+    "limit": 26878
   },
   "LIT003": {
     "limit": 269
@@ -15,7 +15,7 @@
     "limit": 0
   },
   "LIT006": {
-    "limit": 1071
+    "limit": 1069
   },
   "LIT007": {
     "limit": 0
@@ -27,10 +27,10 @@
     "limit": 0
   },
   "LIT010": {
-    "limit": 16700
+    "limit": 16695
   },
   "LIT011": {
-    "limit": 5590
+    "limit": 5588
   },
   "LIT012": {
     "limit": 4519

ui/litellm-dashboard/CLAUDE.md

@@ -16,6 +16,8 @@ Do not trust `eslint --fix` for these two plugins. Fixing the suite in bulk prod
 
 `jest-dom/prefer-to-have-value` stays off because its fixer is wrong here, not merely noisy. It matches any attribute whose name contains "value", so it rewrites `toHaveAttribute("aria-valuenow", n)` into `toHaveValue(n)`, and jest-dom's `toHaveValue` only supports form controls, so the assertion fails on the `role="meter"` elements the dashboard renders. Assert ARIA value attributes with `toHaveAttribute`
 
+Reach for `fireEvent.change` rather than `user.type` when a test only needs a field to hold a value. `user.type` dispatches one event per character and re-renders the whole form each time, which is why a single form test could burn seven seconds. Keep `user.type` where the typing itself is the behaviour under test: an autocomplete that filters per keystroke, a debounce, a key handler, or any Base UI combobox, whose filter state is driven by real keyboard input and does not react to a raw change event
+
 A test may reach for a component library's own CSS class only when that library exposes no role, label, title or ARIA state to query instead, and then the line carries a suppression naming the rule and the reason. Check first: antd icons render as `role="img"` with an `aria-label`, and antd `Form.Item` associates its label with the control, so both are reachable accessibly. When a label does not resolve, suspect the control rather than the test, since a custom wrapper that destructures props without spreading them drops the `id` antd injects and leaves the rendered label pointing at nothing
 
 Rules beyond the enabled set were measured against the whole suite and left off rather than recorded in a budget file, because a ceiling that permits a violation anywhere is worse than an honest gap. `no-node-access` and `no-container` are the ones worth revisiting first, since they catch the DOM archaeology the rules above only discourage. `prefer-implicit-assert` and `prefer-explicit-assert` contradict each other, so neither is enabled

ui/litellm-dashboard/eslint-suppressions.json

@@ -793,16 +793,6 @@
       "count": 1
     }
   },
-  "src/app/(dashboard)/models-and-endpoints/panels/AddModelPanel.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
-  "src/app/(dashboard)/models-and-endpoints/panels/LlmCredentialsPanel.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.ts": {
     "prefer-const": {
       "count": 6
@@ -1740,7 +1730,7 @@
   },
   "src/components/add_model/AddModelForm.test.tsx": {
     "no-restricted-imports": {
-      "count": 2
+      "count": 1
     }
   },
   "src/components/add_model/AddModelForm.tsx": {
@@ -1751,7 +1741,7 @@
       "count": 1
     },
     "no-restricted-imports": {
-      "count": 3
+      "count": 2
     }
   },
   "src/components/add_model/ClassificationMethodConfig.tsx": {
@@ -1803,9 +1793,6 @@
     },
     "no-restricted-imports": {
       "count": 3
-    },
-    "prefer-const": {
-      "count": 2
     }
   },
   "src/components/add_model/auto_router_connection_test.tsx": {
@@ -1818,11 +1805,6 @@
       "count": 1
     }
   },
-  "src/components/add_model/conditional_public_model_name.test.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/components/add_model/conditional_public_model_name.tsx": {
     "local/filename-pascal-case": {
       "count": 1
@@ -1847,11 +1829,6 @@
       "count": 1
     }
   },
-  "src/components/add_model/litellm_model_name.test.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/components/add_model/litellm_model_name.tsx": {
     "local/filename-pascal-case": {
       "count": 1
@@ -1871,11 +1848,6 @@
       "count": 2
     }
   },
-  "src/components/add_model/provider_specific_fields.test.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/components/add_model/provider_specific_fields.tsx": {
     "local/filename-pascal-case": {
       "count": 1
@@ -2332,9 +2304,6 @@
     "local/filename-pascal-case": {
       "count": 1
     },
-    "local/no-complex-jsx-arrow": {
-      "count": 2
-    },
     "max-lines": {
       "count": 1
     },
@@ -3014,4 +2983,4 @@
       "count": 1
     }
   }
-}
+}

ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsModal/AccessGroupEditModal.integration.test.tsx

@@ -1,6 +1,6 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event";
-import { renderWithProviders, screen, waitFor } from "../../../../../../tests/test-utils";
+import { fireEvent, renderWithProviders, screen, waitFor } from "../../../../../../tests/test-utils";
 import { AccessGroupEditModal } from "./AccessGroupEditModal";
 import { AccessGroupResponse } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups";
 
@@ -117,7 +117,7 @@ describe("AccessGroupEditModal submit payload", () => {
     const nameInput = await screen.findByDisplayValue("Engineering");
 
     await user.clear(nameInput);
-    await user.type(nameInput, "  Padded  ");
+    fireEvent.change(nameInput, { target: { value: "  Padded  " } });
     await save(user);
 
     await waitFor(() => expect(mutate).toHaveBeenCalled());

ui/litellm-dashboard/src/app/(dashboard)/access-groups/_components/AccessGroupsPage.test.tsx

@@ -1,4 +1,4 @@
-import { renderWithProviders, screen, within } from "@/../tests/test-utils";
+import { fireEvent, renderWithProviders, screen, within } from "@/../tests/test-utils";
 import { waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
@@ -135,7 +135,9 @@ describe("AccessGroupsPage", () => {
   it("filters by name", async () => {
     const user = userEvent.setup();
     renderWithProviders(<AccessGroupsPage />);
-    await user.type(screen.getByPlaceholderText("Search groups by name, ID, or description..."), "Admin");
+    fireEvent.change(screen.getByPlaceholderText("Search groups by name, ID, or description..."), {
+      target: { value: "Admin" },
+    });
     expect(screen.getByText("Admin Group")).toBeInTheDocument();
     expect(screen.queryByText("Read Only")).not.toBeInTheDocument();
   });
@@ -143,7 +145,9 @@ describe("AccessGroupsPage", () => {
   it("filters by ID", async () => {
     const user = userEvent.setup();
     renderWithProviders(<AccessGroupsPage />);
-    await user.type(screen.getByPlaceholderText("Search groups by name, ID, or description..."), "ag-2");
+    fireEvent.change(screen.getByPlaceholderText("Search groups by name, ID, or description..."), {
+      target: { value: "ag-2" },
+    });
     expect(screen.getByText("Read Only")).toBeInTheDocument();
     expect(screen.queryByText("Admin Group")).not.toBeInTheDocument();
   });
@@ -151,7 +155,9 @@ describe("AccessGroupsPage", () => {
   it("filters by description", async () => {
     const user = userEvent.setup();
     renderWithProviders(<AccessGroupsPage />);
-    await user.type(screen.getByPlaceholderText("Search groups by name, ID, or description..."), "read-only");
+    fireEvent.change(screen.getByPlaceholderText("Search groups by name, ID, or description..."), {
+      target: { value: "read-only" },
+    });
     expect(screen.getByText("Read Only")).toBeInTheDocument();
     expect(screen.queryByText("Admin Group")).not.toBeInTheDocument();
   });
@@ -159,7 +165,9 @@ describe("AccessGroupsPage", () => {
   it("shows the filtered empty state when nothing matches", async () => {
     const user = userEvent.setup();
     renderWithProviders(<AccessGroupsPage />);
-    await user.type(screen.getByPlaceholderText("Search groups by name, ID, or description..."), "no-such-group");
+    fireEvent.change(screen.getByPlaceholderText("Search groups by name, ID, or description..."), {
+      target: { value: "no-such-group" },
+    });
     expect(screen.getByText("No matching access groups")).toBeInTheDocument();
     expect(screen.queryByText("Admin Group")).not.toBeInTheDocument();
   });
@@ -244,7 +252,9 @@ describe("AccessGroupsPage", () => {
     expect(screen.queryByText("ag-01")).not.toBeInTheDocument();
 
     // The only match lives on page 1, so the page index must reset or the table reads as empty.
-    await user.type(screen.getByPlaceholderText("Search groups by name, ID, or description..."), "ag-01");
+    fireEvent.change(screen.getByPlaceholderText("Search groups by name, ID, or description..."), {
+      target: { value: "ag-01" },
+    });
     expect(await screen.findByText("ag-01")).toBeInTheDocument();
     expect(screen.queryByText("No matching access groups")).not.toBeInTheDocument();
   });

ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.test.tsx

@@ -1,4 +1,4 @@
-import { render, screen, waitFor, within } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor, within } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import AdminPanel from "./AdminPanel";
@@ -355,7 +355,7 @@ describe("AdminPanel add allowed IP form", () => {
   it("sends the access token and the typed IP address", async () => {
     const user = userEvent.setup();
 
-    await user.type(ipField(), "192.168.1.50");
+    fireEvent.change(ipField(), { target: { value: "192.168.1.50" } });
     await submitAddIP(user);
 
     await waitFor(() => {
@@ -387,7 +387,7 @@ describe("AdminPanel add allowed IP form", () => {
     const user = userEvent.setup();
     mockGetAllowedIPs.mockResolvedValue(["10.0.0.1", "192.168.1.50"]);
 
-    await user.type(ipField(), "192.168.1.50");
+    fireEvent.change(ipField(), { target: { value: "192.168.1.50" } });
     await submitAddIP(user);
 
     expect(await screen.findByText("192.168.1.50")).toBeInTheDocument();

ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_card_discovery.test.tsx

@@ -61,7 +61,9 @@ describe("AgentCardDiscovery", () => {
 
     expect(screen.getByPlaceholderText("https://upstream-agent.example.com")).toBeInTheDocument();
 
-    await user.type(screen.getByPlaceholderText("https://upstream-agent.example.com"), "https://upstream.example.com");
+    fireEvent.change(screen.getByPlaceholderText("https://upstream-agent.example.com"), {
+      target: { value: "https://upstream.example.com" },
+    });
     await vi.advanceTimersByTimeAsync(500);
 
     await waitFor(() => expect(mockDiscover).toHaveBeenCalled());
@@ -85,7 +87,9 @@ describe("AgentCardDiscovery", () => {
     const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime });
     renderWithProviders(<AgentCardDiscovery accessToken="tok" onApply={vi.fn()} />);
 
-    await user.type(screen.getByPlaceholderText("https://upstream-agent.example.com"), "https://upstream.example.com");
+    fireEvent.change(screen.getByPlaceholderText("https://upstream-agent.example.com"), {
+      target: { value: "https://upstream.example.com" },
+    });
     await vi.advanceTimersByTimeAsync(500);
 
     expect(await screen.findByText("Upstream card loaded")).toBeInTheDocument();
@@ -101,7 +105,9 @@ describe("AgentCardDiscovery", () => {
     const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime });
     renderWithProviders(<AgentCardDiscovery accessToken="tok" onApply={vi.fn()} />);
 
-    await user.type(screen.getByPlaceholderText("https://upstream-agent.example.com"), "https://nope.example");
+    fireEvent.change(screen.getByPlaceholderText("https://upstream-agent.example.com"), {
+      target: { value: "https://nope.example" },
+    });
     await vi.advanceTimersByTimeAsync(500);
 
     expect(await screen.findByText("Discovery failed")).toBeInTheDocument();
@@ -117,7 +123,9 @@ describe("AgentCardDiscovery", () => {
     const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime });
     renderWithProviders(<AgentCardDiscovery accessToken="tok" onApply={onApply} />);
 
-    await user.type(screen.getByPlaceholderText("https://upstream-agent.example.com"), "https://upstream.example.com");
+    fireEvent.change(screen.getByPlaceholderText("https://upstream-agent.example.com"), {
+      target: { value: "https://upstream.example.com" },
+    });
     await vi.advanceTimersByTimeAsync(500);
     await screen.findByText("Upstream card loaded");
 
@@ -290,7 +298,9 @@ describe("AgentCardDiscovery", () => {
     const user = userEvent.setup({ advanceTimers: vi.advanceTimersByTime });
     renderWithProviders(<AgentCardDiscovery accessToken={null} onApply={vi.fn()} />);
 
-    await user.type(screen.getByPlaceholderText("https://upstream-agent.example.com"), "https://upstream.example.com");
+    fireEvent.change(screen.getByPlaceholderText("https://upstream-agent.example.com"), {
+      target: { value: "https://upstream.example.com" },
+    });
     await user.click(screen.getByRole("button", { name: /discover/i }));
 
     expect(await screen.findByText(/No access token available/i)).toBeInTheDocument();

ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx

@@ -1,5 +1,5 @@
 import React from "react";
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event";
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import AgentInfoView from "./agent_info";
@@ -176,7 +176,7 @@ describe("AgentInfoView update payload", () => {
     await openEditor(user);
 
     await user.clear(screen.getByLabelText("TPM Limit"));
-    await user.type(screen.getByLabelText("TPM Limit"), "-5");
+    fireEvent.change(screen.getByLabelText("TPM Limit"), { target: { value: "-5" } });
     await save(user);
 
     expect(patchedPayload().tpm_limit).toBe(0);
@@ -222,7 +222,7 @@ describe("AgentInfoView update payload", () => {
     await openEditor(user);
 
     await user.clear(screen.getByLabelText("API Base"));
-    await user.type(screen.getByLabelText("API Base"), "https://other.example.com");
+    fireEvent.change(screen.getByLabelText("API Base"), { target: { value: "https://other.example.com" } });
 
     await save(user);
 

ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.integration.test.tsx

@@ -1,4 +1,4 @@
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import React from "react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
@@ -39,9 +39,9 @@ describe("BudgetModal", () => {
     const user = userEvent.setup();
     renderModal();
 
-    await user.type(screen.getByLabelText("Budget ID"), "budget-alpha");
-    await user.type(screen.getByLabelText("Max Tokens per minute"), "500.567");
-    await user.type(screen.getByLabelText("Max Requests per minute"), "7");
+    fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } });
+    fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "500.567" } });
+    fireEvent.change(screen.getByLabelText("Max Requests per minute"), { target: { value: "7" } });
     await create(user);
 
     await waitFor(() => expect(createMock).toHaveBeenCalledTimes(1));
@@ -56,12 +56,12 @@ describe("BudgetModal", () => {
     const user = userEvent.setup();
     renderModal();
 
-    await user.type(screen.getByLabelText("Budget ID"), "budget-alpha");
-    await user.type(screen.getByLabelText("Max Tokens per minute"), "500.567");
-    await user.type(screen.getByLabelText("Max Requests per minute"), "7");
+    fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } });
+    fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "500.567" } });
+    fireEvent.change(screen.getByLabelText("Max Requests per minute"), { target: { value: "7" } });
 
     await openOptionalSettings(user);
-    await user.type(screen.getByLabelText("Max Budget (USD)"), "42.567");
+    fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } });
 
     await user.click(screen.getByRole("combobox"));
     await user.click(await screen.findByText("monthly"));
@@ -76,10 +76,10 @@ describe("BudgetModal", () => {
     const user = userEvent.setup();
     renderModal();
 
-    await user.type(screen.getByLabelText("Budget ID"), "budget-alpha");
+    fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } });
 
     await openOptionalSettings(user);
-    await user.type(screen.getByLabelText("Max Budget (USD)"), "42.567");
+    fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } });
     await user.click(screen.getByRole("combobox"));
     await user.click(await screen.findByText("monthly"));
 
@@ -95,8 +95,8 @@ describe("BudgetModal", () => {
     const user = userEvent.setup();
     renderModal();
 
-    await user.type(screen.getByLabelText("Budget ID"), "budget-alpha");
-    await user.type(screen.getByLabelText("Max Tokens per minute"), "5");
+    fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } });
+    fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "5" } });
     await user.clear(screen.getByLabelText("Max Tokens per minute"));
     await create(user);
 
@@ -111,7 +111,7 @@ describe("BudgetModal", () => {
     const user = userEvent.setup();
     renderModal();
 
-    await user.type(screen.getByLabelText("Max Tokens per minute"), "5");
+    fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "5" } });
     await create(user);
 
     await waitFor(() => expect(screen.getByLabelText("Budget ID")).toHaveAttribute("aria-invalid", "true"));
@@ -121,10 +121,10 @@ describe("BudgetModal", () => {
   it("keeps a typed Optional Setting when the section is collapsed and reopened, as antd's store did", async () => {
     const user = userEvent.setup();
     renderModal();
-    await user.type(screen.getByLabelText("Budget ID"), "probe-budget");
+    fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "probe-budget" } });
 
     await openOptionalSettings(user);
-    await user.type(screen.getByLabelText("Max Budget (USD)"), "42.5");
+    fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.5" } });
 
     await user.click(screen.getByText("Optional Settings"));
     await user.click(screen.getByText("Optional Settings"));

ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.test.tsx

@@ -1,5 +1,5 @@
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import React from "react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
@@ -113,7 +113,7 @@ describe("Budget Panel", () => {
     renderPanel();
     await waitFor(() => expect(getMock).toHaveBeenCalled());
 
-    await user.type(screen.getByTestId("datatable-search"), "ecc");
+    fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "ecc" } });
     await waitFor(() => expect(lastQuery().q).toBe("ecc"));
     expect(queries().some((query) => query.q === "e" || query.q === "ec")).toBe(false);
   });
@@ -153,8 +153,8 @@ describe("Budget Panel", () => {
     await waitFor(() => expect(getMock).toHaveBeenCalled());
 
     await openFilters(user);
-    await user.type(screen.getByTestId("budget-filter-max-budget-min"), "10");
-    await user.type(screen.getByTestId("budget-filter-max-budget-max"), "500");
+    fireEvent.change(screen.getByTestId("budget-filter-max-budget-min"), { target: { value: "10" } });
+    fireEvent.change(screen.getByTestId("budget-filter-max-budget-max"), { target: { value: "500" } });
     await user.click(screen.getByTestId("filter-drawer-apply"));
 
     await waitFor(() => expect(lastQuery()["filter[max_budget][gte]"]).toBe("10"));
@@ -171,7 +171,7 @@ describe("Budget Panel", () => {
     await waitFor(() => expect(getMock).toHaveBeenCalled());
 
     await openFilters(user);
-    await user.type(screen.getByTestId("budget-filter-max-budget-min"), "10");
+    fireEvent.change(screen.getByTestId("budget-filter-max-budget-min"), { target: { value: "10" } });
     await user.click(screen.getByTestId("budget-filter-max-budget-unlimited"));
     await user.click(screen.getByTestId("filter-drawer-apply"));
 
@@ -185,8 +185,8 @@ describe("Budget Panel", () => {
     await waitFor(() => expect(getMock).toHaveBeenCalled());
 
     await openFilters(user);
-    await user.type(screen.getByTestId("budget-filter-created-from"), "2026-01-05");
-    await user.type(screen.getByTestId("budget-filter-created-to"), "2026-01-06");
+    fireEvent.change(screen.getByTestId("budget-filter-created-from"), { target: { value: "2026-01-05" } });
+    fireEvent.change(screen.getByTestId("budget-filter-created-to"), { target: { value: "2026-01-06" } });
     await user.click(screen.getByTestId("filter-drawer-apply"));
 
     await waitFor(() =>

ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.integration.test.tsx

@@ -1,4 +1,4 @@
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import React from "react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
@@ -49,7 +49,7 @@ describe("EditBudgetModal", () => {
     renderModal();
 
     await user.clear(screen.getByLabelText("Max Tokens per minute"));
-    await user.type(screen.getByLabelText("Max Tokens per minute"), "500.567");
+    fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "500.567" } });
     await save(user);
 
     await waitFor(() => expect(updateMock).toHaveBeenCalledTimes(1));
@@ -65,13 +65,13 @@ describe("EditBudgetModal", () => {
     renderModal();
 
     await user.clear(screen.getByLabelText("Max Tokens per minute"));
-    await user.type(screen.getByLabelText("Max Tokens per minute"), "500.567");
+    fireEvent.change(screen.getByLabelText("Max Tokens per minute"), { target: { value: "500.567" } });
     await user.clear(screen.getByLabelText("Max Requests per minute"));
-    await user.type(screen.getByLabelText("Max Requests per minute"), "7");
+    fireEvent.change(screen.getByLabelText("Max Requests per minute"), { target: { value: "7" } });
 
     await openOptionalSettings(user);
     await user.clear(screen.getByLabelText("Max Budget (USD)"));
-    await user.type(screen.getByLabelText("Max Budget (USD)"), "42.567");
+    fireEvent.change(screen.getByLabelText("Max Budget (USD)"), { target: { value: "42.567" } });
 
     await user.click(screen.getByRole("combobox"));
     await user.click(await screen.findByText("monthly"));
@@ -97,7 +97,7 @@ describe("EditBudgetModal", () => {
     await openOptionalSettings(user);
     const maxBudget = screen.getByLabelText("Max Budget (USD)");
     await user.clear(maxBudget);
-    await user.type(maxBudget, "99.25");
+    fireEvent.change(maxBudget, { target: { value: "99.25" } });
 
     await user.click(screen.getByText("Optional Settings"));
     await user.click(screen.getByText("Optional Settings"));

ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.integration.test.tsx

@@ -1,5 +1,5 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import CacheSettings from "./index";
 
@@ -99,7 +99,7 @@ describe("CacheSettings advanced settings round-trip", () => {
     await screen.findByText("Connection Settings");
 
     await user.click(screen.getByText("Advanced Settings"));
-    await user.type(await screen.findByLabelText("Namespace"), "typed-ns");
+    fireEvent.change(await screen.findByLabelText("Namespace"), { target: { value: "typed-ns" } });
     await user.click(screen.getByText("Advanced Settings"));
     await waitFor(() => expect(screen.queryByLabelText("Namespace")).not.toBeInTheDocument());
 
@@ -116,7 +116,7 @@ describe("CacheSettings advanced settings round-trip", () => {
     await screen.findByText("Connection Settings");
 
     await user.click(screen.getByText("Advanced Settings"));
-    await user.type(await screen.findByLabelText("Namespace"), "typed-ns");
+    fireEvent.change(await screen.findByLabelText("Namespace"), { target: { value: "typed-ns" } });
     await user.click(screen.getByText("Advanced Settings"));
     await waitFor(() => expect(screen.queryByLabelText("Namespace")).not.toBeInTheDocument());
     await user.click(screen.getByText("Advanced Settings"));
@@ -145,7 +145,7 @@ describe("CacheSettings advanced settings round-trip", () => {
     await screen.findByText("Connection Settings");
 
     await user.click(screen.getByText("Advanced Settings"));
-    await user.type(await screen.findByLabelText("TTL (seconds)"), "not-a-number");
+    fireEvent.change(await screen.findByLabelText("TTL (seconds)"), { target: { value: "not-a-number" } });
     await user.click(screen.getByText("Advanced Settings"));
     await waitFor(() => expect(screen.queryByLabelText("TTL (seconds)")).not.toBeInTheDocument());
 

ui/litellm-dashboard/src/app/(dashboard)/caching/_components/cache_settings/index.test.tsx

@@ -1,5 +1,5 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import CacheSettings from "./index";
 
@@ -79,7 +79,7 @@ describe("CacheSettings", () => {
 
       const port = await screen.findByLabelText("Port");
       await user.clear(port);
-      await user.type(port, "99999");
+      fireEvent.change(port, { target: { value: "99999" } });
       await user.click(screen.getByRole("button", { name: /save changes/i }));
 
       expect(await screen.findByText(/Port must be an integer between 1 and 65535/i)).toBeInTheDocument();
@@ -92,7 +92,7 @@ describe("CacheSettings", () => {
       renderSettings();
 
       const startupNodes = await screen.findByLabelText("Startup Nodes");
-      await user.type(startupNodes, "not json");
+      fireEvent.change(startupNodes, { target: { value: "not json" } });
       await user.click(screen.getByRole("button", { name: /save changes/i }));
 
       expect(await screen.findByText(/Must be a valid JSON array/i)).toBeInTheDocument();
@@ -104,7 +104,7 @@ describe("CacheSettings", () => {
       renderSettings();
 
       const db = await screen.findByLabelText("Database Index");
-      await user.type(db, "redis://host:6379/1");
+      fireEvent.change(db, { target: { value: "redis://host:6379/1" } });
       await user.click(screen.getByRole("button", { name: /save changes/i }));
 
       expect(await screen.findByText(/Must be a non-negative integer/i)).toBeInTheDocument();
@@ -118,7 +118,7 @@ describe("CacheSettings", () => {
       renderSettings();
 
       const host = await screen.findByLabelText("Host");
-      await user.type(host, "localhost");
+      fireEvent.change(host, { target: { value: "localhost" } });
       await user.click(screen.getByRole("button", { name: /save changes/i }));
 
       await waitFor(() =>
@@ -136,8 +136,8 @@ describe("CacheSettings", () => {
       const user = userEvent.setup();
       renderSettings();
 
-      await user.type(await screen.findByLabelText("Redis URL"), "redis://host:6379/1");
-      await user.type(await screen.findByLabelText("Database Index"), "2");
+      fireEvent.change(await screen.findByLabelText("Redis URL"), { target: { value: "redis://host:6379/1" } });
+      fireEvent.change(await screen.findByLabelText("Database Index"), { target: { value: "2" } });
       await user.click(screen.getByRole("button", { name: /save changes/i }));
 
       await waitFor(() => expect(updateCacheSettingsCall).toHaveBeenCalled());

ui/litellm-dashboard/src/app/(dashboard)/caching/_components/coordination_redis_settings/index.integration.test.tsx

@@ -1,5 +1,5 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import userEvent from "@testing-library/user-event";
 import CoordinationRedisSettings from "./index";
@@ -53,7 +53,7 @@ describe("CoordinationRedisSettings value retention across redis types", () => {
     await screen.findByLabelText("Host");
 
     await pickRedisType(user, /sentinel/i);
-    await user.type(await screen.findByLabelText("Service Name"), "mymaster");
+    fireEvent.change(await screen.findByLabelText("Service Name"), { target: { value: "mymaster" } });
 
     await pickRedisType(user, /node/i);
     await waitFor(() => expect(screen.queryByLabelText("Service Name")).not.toBeInTheDocument());
@@ -69,7 +69,7 @@ describe("CoordinationRedisSettings value retention across redis types", () => {
     await screen.findByLabelText("Host");
 
     await pickRedisType(user, /sentinel/i);
-    await user.type(await screen.findByLabelText("Service Name"), "mymaster");
+    fireEvent.change(await screen.findByLabelText("Service Name"), { target: { value: "mymaster" } });
 
     await pickRedisType(user, /node/i);
     await waitFor(() => expect(screen.queryByLabelText("Service Name")).not.toBeInTheDocument());

ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.integration.test.tsx

@@ -1,4 +1,4 @@
-import { render, screen } from "@testing-library/react";
+import { fireEvent, render, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 
@@ -27,8 +27,10 @@ describe("PromptCompressionTab submit payload", () => {
     const user = userEvent.setup();
     render(<PromptCompressionTab accessToken="test-token" />);
 
-    await user.type(screen.getByLabelText("Name"), "  headroom-compression  ");
-    await user.type(screen.getByLabelText("Headroom API base"), "  https://headroom.example.com  ");
+    fireEvent.change(screen.getByLabelText("Name"), { target: { value: "  headroom-compression  " } });
+    fireEvent.change(screen.getByLabelText("Headroom API base"), {
+      target: { value: "  https://headroom.example.com  " },
+    });
     await user.click(screen.getByRole("button", { name: "Add guardrail" }));
 
     await vi.waitFor(() =>
@@ -49,8 +51,8 @@ describe("PromptCompressionTab submit payload", () => {
     const user = userEvent.setup();
     render(<PromptCompressionTab accessToken="test-token" />);
 
-    await user.type(screen.getByLabelText("Name"), "headroom-optin");
-    await user.type(screen.getByLabelText("Headroom API base"), "https://headroom.example.com");
+    fireEvent.change(screen.getByLabelText("Name"), { target: { value: "headroom-optin" } });
+    fireEvent.change(screen.getByLabelText("Headroom API base"), { target: { value: "https://headroom.example.com" } });
     await user.click(screen.getByLabelText("Apply to all requests"));
     await user.click(screen.getByRole("button", { name: "Add guardrail" }));
 
@@ -82,7 +84,7 @@ describe("PromptCompressionTab submit payload", () => {
     const user = userEvent.setup();
     render(<PromptCompressionTab accessToken="test-token" />);
 
-    await user.type(screen.getByLabelText("Name"), "headroom-compression");
+    fireEvent.change(screen.getByLabelText("Name"), { target: { value: "headroom-compression" } });
     await user.type(screen.getByLabelText("Headroom API base"), "https://headroom.example.com{Enter}");
 
     await vi.waitFor(() => expect(createGuardrailCall).toHaveBeenCalledTimes(1));
@@ -92,8 +94,8 @@ describe("PromptCompressionTab submit payload", () => {
     const user = userEvent.setup();
     render(<PromptCompressionTab accessToken="test-token" />);
 
-    await user.type(screen.getByLabelText("Name"), "headroom-compression");
-    await user.type(screen.getByLabelText("Headroom API base"), "https://headroom.example.com");
+    fireEvent.change(screen.getByLabelText("Name"), { target: { value: "headroom-compression" } });
+    fireEvent.change(screen.getByLabelText("Headroom API base"), { target: { value: "https://headroom.example.com" } });
     await user.click(screen.getByLabelText("Apply to all requests"));
     await user.click(screen.getByRole("button", { name: "Add guardrail" }));
 

ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/cost_tracking_settings.integration.test.tsx

@@ -1,6 +1,6 @@
 import React from "react";
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { screen, waitFor } from "@testing-library/react";
+import { fireEvent, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { renderWithProviders } from "../../../../../tests/test-utils";
 import CostTrackingSettings from "./cost_tracking_settings";
@@ -87,7 +87,7 @@ describe("CostTrackingSettings submit paths", () => {
 
     await user.click(screen.getAllByRole("combobox")[0]);
     await user.click((await screen.findAllByRole("option"))[0]);
-    await user.type(screen.getByLabelText(/Discount Percentage/i), "5");
+    fireEvent.change(screen.getByLabelText(/Discount Percentage/i), { target: { value: "5" } });
     await user.click(submitDiscount());
 
     await waitFor(() => expect(stableDiscountCallbacks.handleAddProvider).toHaveBeenCalled());
@@ -105,7 +105,7 @@ describe("CostTrackingSettings submit paths", () => {
 
     await user.click(screen.getAllByRole("combobox")[0]);
     await user.click((await screen.findAllByRole("option"))[0]);
-    await user.type(screen.getByLabelText(/Margin Percentage/i), "10");
+    fireEvent.change(screen.getByLabelText(/Margin Percentage/i), { target: { value: "10" } });
 
     const submit = screen
       .getAllByRole("button")

ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/how_it_works.test.tsx

@@ -1,6 +1,6 @@
 import React from "react";
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { screen } from "@testing-library/react";
+import { fireEvent, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { renderWithProviders } from "../../../../../tests/test-utils";
 import HowItWorks from "./how_it_works";
@@ -52,7 +52,7 @@ describe("HowItWorks", () => {
     renderWithProviders(<HowItWorks />);
 
     const responseCostInput = screen.getByPlaceholderText("0.0171938125");
-    await user.type(responseCostInput, "0.01");
+    fireEvent.change(responseCostInput, { target: { value: "0.01" } });
 
     expect(screen.queryByText("Calculated Results")).not.toBeInTheDocument();
   });
@@ -62,7 +62,7 @@ describe("HowItWorks", () => {
     renderWithProviders(<HowItWorks />);
 
     const discountAmountInput = screen.getByPlaceholderText("0.0009049375");
-    await user.type(discountAmountInput, "0.001");
+    fireEvent.change(discountAmountInput, { target: { value: "0.001" } });
 
     expect(screen.queryByText("Calculated Results")).not.toBeInTheDocument();
   });
@@ -74,8 +74,8 @@ describe("HowItWorks", () => {
     const responseCostInput = screen.getByPlaceholderText("0.0171938125");
     const discountAmountInput = screen.getByPlaceholderText("0.0009049375");
 
-    await user.type(responseCostInput, "0.0171938125");
-    await user.type(discountAmountInput, "0.0009049375");
+    fireEvent.change(responseCostInput, { target: { value: "0.0171938125" } });
+    fireEvent.change(discountAmountInput, { target: { value: "0.0009049375" } });
 
     expect(await screen.findByText("Calculated Results")).toBeInTheDocument();
   });
@@ -84,8 +84,8 @@ describe("HowItWorks", () => {
     const user = userEvent.setup();
     renderWithProviders(<HowItWorks />);
 
-    await user.type(screen.getByPlaceholderText("0.0171938125"), "0.0171938125");
-    await user.type(screen.getByPlaceholderText("0.0009049375"), "0.0009049375");
+    fireEvent.change(screen.getByPlaceholderText("0.0171938125"), { target: { value: "0.0171938125" } });
+    fireEvent.change(screen.getByPlaceholderText("0.0009049375"), { target: { value: "0.0009049375" } });
 
     expect(await screen.findByText("Original Cost:")).toBeInTheDocument();
     expect(screen.getByText("Final Cost:")).toBeInTheDocument();

ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_discount_table.test.tsx

@@ -1,6 +1,6 @@
 import React from "react";
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { screen } from "@testing-library/react";
+import { fireEvent, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { renderWithProviders } from "../../../../../tests/test-utils";
 import ProviderDiscountTable from "./provider_discount_table";
@@ -159,7 +159,7 @@ describe("ProviderDiscountTable", () => {
 
     const input = screen.getByPlaceholderText("5");
     await user.clear(input);
-    await user.type(input, "10");
+    fireEvent.change(input, { target: { value: "10" } });
 
     await user.click(rowAction("save"));
 
@@ -268,7 +268,7 @@ describe("ProviderDiscountTable", () => {
     await user.click(rowAction("edit"));
     const input = screen.getByPlaceholderText("5");
     await user.clear(input);
-    await user.type(input, "150");
+    fireEvent.change(input, { target: { value: "150" } });
     await user.click(rowAction("save"));
 
     expect(onDiscountChange).not.toHaveBeenCalled();

ui/litellm-dashboard/src/app/(dashboard)/cost-tracking/_components/provider_margin_table.test.tsx

@@ -1,6 +1,6 @@
 import React from "react";
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { screen } from "@testing-library/react";
+import { fireEvent, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { renderWithProviders } from "../../../../../tests/test-utils";
 import ProviderMarginTable from "./provider_margin_table";
@@ -188,7 +188,7 @@ describe("ProviderMarginTable", () => {
 
     const percentInput = screen.getByPlaceholderText("10");
     await user.clear(percentInput);
-    await user.type(percentInput, "20");
+    fireEvent.change(percentInput, { target: { value: "20" } });
 
     await user.click(rowAction("save"));
 
@@ -208,7 +208,7 @@ describe("ProviderMarginTable", () => {
     await user.click(rowAction("edit"));
 
     await user.clear(screen.getByPlaceholderText("10"));
-    await user.type(screen.getByPlaceholderText("0.001"), "0.002");
+    fireEvent.change(screen.getByPlaceholderText("0.001"), { target: { value: "0.002" } });
 
     await user.click(rowAction("save"));
 
@@ -316,10 +316,10 @@ describe("ProviderMarginTable", () => {
 
       const percentInput = screen.getByPlaceholderText("10");
       await user.clear(percentInput);
-      await user.type(percentInput, "5");
+      fireEvent.change(percentInput, { target: { value: "5" } });
 
       const fixedInput = screen.getByPlaceholderText("0.001");
-      await user.type(fixedInput, "0.002");
+      fireEvent.change(fixedInput, { target: { value: "0.002" } });
 
       await user.click(rowAction("save"));
 

ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.test.tsx

@@ -1,6 +1,6 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import userEvent from "@testing-library/user-event";
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { EvaluationSettingsModal } from "./EvaluationSettingsModal";
 
 const mockFetchAvailableModels = vi.fn();
@@ -60,7 +60,7 @@ describe("EvaluationSettingsModal", () => {
 
     const promptBox = screen.getByDisplayValue(/Evaluate whether this guardrail's decision was correct/);
     await user.clear(promptBox);
-    await user.type(promptBox, "custom prompt");
+    fireEvent.change(promptBox, { target: { value: "custom prompt" } });
     expect(screen.getByDisplayValue("custom prompt")).toBeInTheDocument();
 
     await user.click(screen.getByText("Reset to default"));

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.test.tsx

@@ -1,5 +1,5 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { render, screen, waitFor } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import CustomPatternModal from "./CustomPatternModal";
 
@@ -42,7 +42,7 @@ describe("CustomPatternModal", () => {
 
     // Find and fill the pattern name input
     const nameInput = screen.getByPlaceholderText("e.g., internal_id, employee_code");
-    await user.type(nameInput, "employee_id");
+    fireEvent.change(nameInput, { target: { value: "employee_id" } });
 
     // Find and fill the regex pattern input - use paste instead of type to avoid special char issues
     const regexInput = screen.getByPlaceholderText("e.g., ID-[0-9]{6}");

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.test.tsx

@@ -1,5 +1,5 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { render, screen } from "@testing-library/react";
+import { fireEvent, render, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import KeywordModal from "./KeywordModal";
 
@@ -35,7 +35,9 @@ describe("KeywordModal", () => {
     const user = userEvent.setup();
     renderModal();
 
-    await user.type(await screen.findByPlaceholderText("Enter sensitive keyword or phrase"), "s");
+    fireEvent.change(await screen.findByPlaceholderText("Enter sensitive keyword or phrase"), {
+      target: { value: "s" },
+    });
 
     expect(handlers.onKeywordChange).toHaveBeenCalledWith("s");
   });
@@ -44,7 +46,9 @@ describe("KeywordModal", () => {
     const user = userEvent.setup();
     renderModal();
 
-    await user.type(await screen.findByPlaceholderText("Explain why this keyword is sensitive"), "x");
+    fireEvent.change(await screen.findByPlaceholderText("Explain why this keyword is sensitive"), {
+      target: { value: "x" },
+    });
 
     expect(handlers.onDescriptionChange).toHaveBeenCalledWith("x");
   });

ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ThresholdInput.test.tsx

@@ -1,5 +1,5 @@
 import { useState } from "react";
-import { render, screen } from "@testing-library/react";
+import { fireEvent, render, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { describe, expect, it, vi } from "vitest";
 
@@ -43,7 +43,7 @@ describe("ThresholdInput", () => {
     await user.clear(input);
     expect(onValueChange).toHaveBeenLastCalledWith(null);
 
-    await user.type(input, "0.55");
+    fireEvent.change(input, { target: { value: "0.55" } });
     expect(onValueChange).toHaveBeenLastCalledWith(0.55);
   });
 
@@ -54,7 +54,7 @@ describe("ThresholdInput", () => {
 
     const input = screen.getByRole("spinbutton");
     await user.clear(input);
-    await user.type(input, "5");
+    fireEvent.change(input, { target: { value: "5" } });
     await user.tab();
 
     expect(onValueChange).toHaveBeenLastCalledWith(1);

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/AwsSigV4Fields.tsx

@@ -1,7 +1,25 @@
 import React from "react";
-import { Form, Input, Tooltip } from "antd";
+import { Input, Tooltip } from "antd";
 import { InfoCircleOutlined } from "@ant-design/icons";
 
+import { MountedFormField } from "@/components/common_components/MountedFormField";
+import { antdRequired } from "@/components/common_components/antdFormRules";
+import { requiredWhenSiblingSet, textControl } from "./mcpFieldRules";
+
+const fieldClassName = "rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500";
+
+const FieldLabel: React.FC<{ label: string; tooltip: string }> = ({ label, tooltip }) => (
+  <span className="text-sm font-medium text-gray-700 flex items-center">
+    {label}
+    <Tooltip title={tooltip}>
+      <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
+    </Tooltip>
+  </span>
+);
+
+const ACCESS_KEY_PATH = ["credentials", "aws_access_key_id"] as const;
+const SECRET_KEY_PATH = ["credentials", "aws_secret_access_key"] as const;
+
 const AwsSigV4Fields: React.FC = () => (
   <>
     <p className="text-sm text-gray-500 mb-2">
@@ -15,140 +33,120 @@ const AwsSigV4Fields: React.FC = () => (
         View docs &rarr;
       </a>
     </p>
-    <Form.Item
-      label={
-        <span className="text-sm font-medium text-gray-700 flex items-center">
-          AWS Region
-          <Tooltip title="AWS region for SigV4 signing (e.g., us-east-1)">
-            <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-          </Tooltip>
-        </span>
-      }
+    <MountedFormField
+      label={<FieldLabel label="AWS Region" tooltip="AWS region for SigV4 signing (e.g., us-east-1)" />}
       name={["credentials", "aws_region_name"]}
-      rules={[{ required: true, message: "AWS region is required for SigV4 auth" }]}
+      required
+      rules={{ validate: { required: antdRequired("AWS region is required for SigV4 auth") } }}
     >
-      <Input placeholder="us-east-1" className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500" />
-    </Form.Item>
-    <Form.Item
+      {(control) => <Input {...textControl(control)} placeholder="us-east-1" className={fieldClassName} />}
+    </MountedFormField>
+    <MountedFormField
       label={
-        <span className="text-sm font-medium text-gray-700 flex items-center">
-          AWS Service Name
-          <Tooltip title="AWS service name for SigV4 signing. Defaults to 'bedrock-agentcore'.">
-            <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-          </Tooltip>
-        </span>
+        <FieldLabel
+          label="AWS Service Name"
+          tooltip="AWS service name for SigV4 signing. Defaults to 'bedrock-agentcore'."
+        />
       }
       name={["credentials", "aws_service_name"]}
     >
-      <Input
-        placeholder="bedrock-agentcore"
-        className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-      />
-    </Form.Item>
-    <Form.Item
+      {(control) => <Input {...textControl(control)} placeholder="bedrock-agentcore" className={fieldClassName} />}
+    </MountedFormField>
+    <MountedFormField
       label={
-        <span className="text-sm font-medium text-gray-700 flex items-center">
-          AWS Access Key ID
-          <Tooltip title="Optional. If not provided, falls back to the boto3 credential chain (IAM role, env vars, etc.).">
-            <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-          </Tooltip>
-        </span>
+        <FieldLabel
+          label="AWS Access Key ID"
+          tooltip="Optional. If not provided, falls back to the boto3 credential chain (IAM role, env vars, etc.)."
+        />
       }
-      name={["credentials", "aws_access_key_id"]}
-      dependencies={[["credentials", "aws_secret_access_key"]]}
-      rules={[
-        ({ getFieldValue }) => ({
-          validator(_, value) {
-            const secretKey = getFieldValue(["credentials", "aws_secret_access_key"]);
-            if (secretKey && !value) {
-              return Promise.reject(new Error("Access Key ID is required when Secret Access Key is provided"));
-            }
-            return Promise.resolve();
-          },
-        }),
-      ]}
+      name={ACCESS_KEY_PATH}
+      rules={{
+        deps: ["credentials.aws_secret_access_key"],
+        validate: {
+          pairedWithSecret: requiredWhenSiblingSet(
+            SECRET_KEY_PATH,
+            "Access Key ID is required when Secret Access Key is provided",
+          ),
+        },
+      }}
     >
-      <Input.Password
-        placeholder="AKIA... (optional — uses IAM role if blank)"
-        className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-      />
-    </Form.Item>
-    <Form.Item
+      {(control) => (
+        <Input.Password
+          {...textControl(control)}
+          placeholder="AKIA... (optional — uses IAM role if blank)"
+          className={fieldClassName}
+        />
+      )}
+    </MountedFormField>
+    <MountedFormField
       label={
-        <span className="text-sm font-medium text-gray-700 flex items-center">
-          AWS Secret Access Key
-          <Tooltip title="Optional. Required if AWS Access Key ID is provided.">
-            <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-          </Tooltip>
-        </span>
+        <FieldLabel label="AWS Secret Access Key" tooltip="Optional. Required if AWS Access Key ID is provided." />
       }
-      name={["credentials", "aws_secret_access_key"]}
-      dependencies={[["credentials", "aws_access_key_id"]]}
-      rules={[
-        ({ getFieldValue }) => ({
-          validator(_, value) {
-            const accessKeyId = getFieldValue(["credentials", "aws_access_key_id"]);
-            if (accessKeyId && !value) {
-              return Promise.reject(new Error("Secret Access Key is required when Access Key ID is provided"));
-            }
-            return Promise.resolve();
-          },
-        }),
-      ]}
+      name={SECRET_KEY_PATH}
+      rules={{
+        deps: ["credentials.aws_access_key_id"],
+        validate: {
+          pairedWithAccessKey: requiredWhenSiblingSet(
+            ACCESS_KEY_PATH,
+            "Secret Access Key is required when Access Key ID is provided",
+          ),
+        },
+      }}
     >
-      <Input.Password
-        placeholder="Enter secret key (optional — uses IAM role if blank)"
-        className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-      />
-    </Form.Item>
-    <Form.Item
-      label={
-        <span className="text-sm font-medium text-gray-700 flex items-center">
-          AWS Session Token
-          <Tooltip title="Optional. Only needed for temporary STS credentials.">
-            <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-          </Tooltip>
-        </span>
-      }
+      {(control) => (
+        <Input.Password
+          {...textControl(control)}
+          placeholder="Enter secret key (optional — uses IAM role if blank)"
+          className={fieldClassName}
+        />
+      )}
+    </MountedFormField>
+    <MountedFormField
+      label={<FieldLabel label="AWS Session Token" tooltip="Optional. Only needed for temporary STS credentials." />}
       name={["credentials", "aws_session_token"]}
     >
-      <Input.Password
-        placeholder="Enter session token (optional)"
-        className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-      />
-    </Form.Item>
-    <Form.Item
+      {(control) => (
+        <Input.Password
+          {...textControl(control)}
+          placeholder="Enter session token (optional)"
+          className={fieldClassName}
+        />
+      )}
+    </MountedFormField>
+    <MountedFormField
       label={
-        <span className="text-sm font-medium text-gray-700 flex items-center">
-          AWS Role ARN
-          <Tooltip title="Optional. IAM role ARN to assume via STS before signing. If set, LiteLLM calls sts:AssumeRole to get temporary credentials. Uses ambient credentials (IAM role, env vars) as the source identity unless explicit keys are also provided.">
-            <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-          </Tooltip>
-        </span>
+        <FieldLabel
+          label="AWS Role ARN"
+          tooltip="Optional. IAM role ARN to assume via STS before signing. If set, LiteLLM calls sts:AssumeRole to get temporary credentials. Uses ambient credentials (IAM role, env vars) as the source identity unless explicit keys are also provided."
+        />
       }
       name={["credentials", "aws_role_name"]}
     >
-      <Input
-        placeholder="arn:aws:iam::123456789012:role/MyRole (optional)"
-        className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-      />
-    </Form.Item>
-    <Form.Item
+      {(control) => (
+        <Input
+          {...textControl(control)}
+          placeholder="arn:aws:iam::123456789012:role/MyRole (optional)"
+          className={fieldClassName}
+        />
+      )}
+    </MountedFormField>
+    <MountedFormField
       label={
-        <span className="text-sm font-medium text-gray-700 flex items-center">
-          AWS Session Name
-          <Tooltip title="Optional. Session name for the AssumeRole call — appears in CloudTrail logs. Auto-generated if omitted.">
-            <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-          </Tooltip>
-        </span>
+        <FieldLabel
+          label="AWS Session Name"
+          tooltip="Optional. Session name for the AssumeRole call — appears in CloudTrail logs. Auto-generated if omitted."
+        />
       }
       name={["credentials", "aws_session_name"]}
     >
-      <Input
-        placeholder="litellm-prod (optional, auto-generated if blank)"
-        className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-      />
-    </Form.Item>
+      {(control) => (
+        <Input
+          {...textControl(control)}
+          placeholder="litellm-prod (optional, auto-generated if blank)"
+          className={fieldClassName}
+        />
+      )}
+    </MountedFormField>
   </>
 );
 

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.integration.test.tsx

@@ -266,11 +266,11 @@ describe("CreateMCPServer", () => {
 
       // Fill in server name (use id to avoid duplicate placeholder)
       const nameInput = getServerNameInput();
-      await user.type(nameInput, "Test_Server");
+      fireEvent.change(nameInput, { target: { value: "Test_Server" } });
 
       // Fill in URL
       const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
-      await user.type(urlInput, "https://example.com/mcp");
+      fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } });
 
       // Select API Key auth type
       await selectAntOption("Authentication", "API Key");
@@ -310,10 +310,10 @@ describe("CreateMCPServer", () => {
       const user = userEvent.setup({ delay: null });
 
       const nameInput = getServerNameInput();
-      await user.type(nameInput, "Test_Server");
+      fireEvent.change(nameInput, { target: { value: "Test_Server" } });
 
       const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
-      await user.type(urlInput, "https://example.com/mcp");
+      fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } });
 
       await selectAntOption("Authentication", "Bearer Token");
 
@@ -351,10 +351,10 @@ describe("CreateMCPServer", () => {
       const user = userEvent.setup({ delay: null });
 
       const nameInput = getServerNameInput();
-      await user.type(nameInput, "My_Server");
+      fireEvent.change(nameInput, { target: { value: "My_Server" } });
 
       const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
-      await user.type(urlInput, "https://example.com/mcp");
+      fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } });
 
       await selectAntOption("Authentication", "API Key");
 
@@ -364,7 +364,7 @@ describe("CreateMCPServer", () => {
 
       // Fill in auth value
       const authInput = screen.getByPlaceholderText("Enter token or secret");
-      await user.type(authInput, "my-secret-key");
+      fireEvent.change(authInput, { target: { value: "my-secret-key" } });
 
       vi.mocked(networking.createMCPServer).mockResolvedValue({
         server_id: "new-server-1",
@@ -397,8 +397,10 @@ describe("CreateMCPServer", () => {
       await selectHttpTransport();
 
       const user = userEvent.setup({ delay: null });
-      await user.type(getServerNameInput(), "PT_Server");
-      await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+      fireEvent.change(getServerNameInput(), { target: { value: "PT_Server" } });
+      fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+        target: { value: "https://example.com/mcp" },
+      });
 
       await selectAntOption("Authentication", "True Passthrough (no LiteLLM auth)");
 
@@ -423,8 +425,10 @@ describe("CreateMCPServer", () => {
       await selectHttpTransport();
 
       const user = userEvent.setup({ delay: null });
-      await user.type(getServerNameInput(), "CF_Server");
-      await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+      fireEvent.change(getServerNameInput(), { target: { value: "CF_Server" } });
+      fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+        target: { value: "https://example.com/mcp" },
+      });
 
       await selectAntOption("Authentication", optionLabel);
 
@@ -487,19 +491,22 @@ describe("CreateMCPServer", () => {
         await selectHttpTransport();
 
         const user = userEvent.setup({ delay: null });
-        await user.type(getServerNameInput(), "CF_App_Server");
-        await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+        fireEvent.change(getServerNameInput(), { target: { value: "CF_App_Server" } });
+        fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+          target: { value: "https://example.com/mcp" },
+        });
 
         await selectAntOption("Authentication", optionLabel);
 
         // Admin declares the org's pre-registered upstream app; unlike the browser-authorized
         // token, this is config and must survive onto the server row so internal users'
         // Tools-page Authorize relays through it (required for non-DCR upstreams like Slack).
-        await user.type(
-          screen.getByPlaceholderText("Leave blank to use dynamic client registration"),
-          "org-app-client-id",
-        );
-        await user.type(screen.getByPlaceholderText("Leave blank for public clients / PKCE"), "org-app-secret");
+        fireEvent.change(screen.getByPlaceholderText("Leave blank to use dynamic client registration"), {
+          target: { value: "org-app-client-id" },
+        });
+        fireEvent.change(screen.getByPlaceholderText("Leave blank for public clients / PKCE"), {
+          target: { value: "org-app-secret" },
+        });
 
         await waitFor(() => expect(oauthHook.onTokenReceived).toBeTruthy());
         await act(async () => {
@@ -548,16 +555,19 @@ describe("CreateMCPServer", () => {
       await selectHttpTransport();
 
       const user = userEvent.setup({ delay: null });
-      await user.type(getServerNameInput(), "CF_Keep_Server");
-      await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+      fireEvent.change(getServerNameInput(), { target: { value: "CF_Keep_Server" } });
+      fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+        target: { value: "https://example.com/mcp" },
+      });
 
       await selectAntOption("Authentication", "True Passthrough (no LiteLLM auth)");
 
-      await user.type(
-        screen.getByPlaceholderText("Leave blank to use dynamic client registration"),
-        "org-app-client-id",
-      );
-      await user.type(screen.getByPlaceholderText("Leave blank for public clients / PKCE"), "org-app-secret");
+      fireEvent.change(screen.getByPlaceholderText("Leave blank to use dynamic client registration"), {
+        target: { value: "org-app-client-id" },
+      });
+      fireEvent.change(screen.getByPlaceholderText("Leave blank for public clients / PKCE"), {
+        target: { value: "org-app-secret" },
+      });
 
       await waitFor(() => expect(oauthHook.onTokenReceived).toBeTruthy());
       await act(async () => {
@@ -605,8 +615,10 @@ describe("CreateMCPServer", () => {
       await selectHttpTransport();
 
       const user = userEvent.setup({ delay: null });
-      await user.type(getServerNameInput(), "Switch_Server");
-      await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+      fireEvent.change(getServerNameInput(), { target: { value: "Switch_Server" } });
+      fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+        target: { value: "https://example.com/mcp" },
+      });
 
       await selectAntOption("Authentication", "OAuth");
 
@@ -653,8 +665,10 @@ describe("CreateMCPServer", () => {
     it("keeps the DCR-minted client out of form.credentials but reuses it via getCredentials", async () => {
       await selectHttpTransport();
       const user = userEvent.setup({ delay: null });
-      await user.type(getServerNameInput(), "DCR_Server");
-      await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+      fireEvent.change(getServerNameInput(), { target: { value: "DCR_Server" } });
+      fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+        target: { value: "https://example.com/mcp" },
+      });
       await selectAntOption("Authentication", "OAuth");
 
       await waitFor(() => expect(oauthHook.onTokenReceived).toBeTruthy());
@@ -678,8 +692,10 @@ describe("CreateMCPServer", () => {
       await selectAntOption("Transport Type", "Streamable HTTP");
       expect(await screen.findByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument();
       const user = userEvent.setup({ delay: null });
-      await user.type(getServerNameInput(), "Leak_Server");
-      await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+      fireEvent.change(getServerNameInput(), { target: { value: "Leak_Server" } });
+      fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+        target: { value: "https://example.com/mcp" },
+      });
       await selectAntOption("Authentication", "OAuth");
 
       await waitFor(() => expect(oauthHook.onTokenReceived).toBeTruthy());
@@ -705,8 +721,10 @@ describe("CreateMCPServer", () => {
     it("persists the DCR client on an oauth2 submit via the ref", async () => {
       await selectHttpTransport();
       const user = userEvent.setup({ delay: null });
-      await user.type(getServerNameInput(), "DCR_Submit_Server");
-      await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+      fireEvent.change(getServerNameInput(), { target: { value: "DCR_Submit_Server" } });
+      fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+        target: { value: "https://example.com/mcp" },
+      });
       await selectAntOption("Authentication", "OAuth");
 
       await waitFor(() => expect(oauthHook.onTokenReceived).toBeTruthy());
@@ -823,10 +841,14 @@ describe("CreateMCPServer", () => {
     it("keeps the typed app but warns when the URL changes after a client-forwarded authorize", async () => {
       await selectHttpTransport();
       const user = userEvent.setup({ delay: null });
-      await user.type(getServerNameInput(), "CF_Warn");
-      await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+      fireEvent.change(getServerNameInput(), { target: { value: "CF_Warn" } });
+      fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+        target: { value: "https://example.com/mcp" },
+      });
       await selectAntOption("Authentication", "True Passthrough (no LiteLLM auth)");
-      await user.type(screen.getByPlaceholderText("Leave blank to use dynamic client registration"), "app-id");
+      fireEvent.change(screen.getByPlaceholderText("Leave blank to use dynamic client registration"), {
+        target: { value: "app-id" },
+      });
 
       await waitFor(() => expect(oauthHook.onTokenReceived).toBeTruthy());
       await act(async () => {
@@ -845,12 +867,17 @@ describe("CreateMCPServer", () => {
 
     it("keeps client_secret when only client_id is edited after a client-forwarded authorize", async () => {
       await selectHttpTransport();
-      const user = userEvent.setup({ delay: null });
-      await user.type(getServerNameInput(), "CF_Keystroke");
-      await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+      fireEvent.change(getServerNameInput(), { target: { value: "CF_Keystroke" } });
+      fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+        target: { value: "https://example.com/mcp" },
+      });
       await selectAntOption("Authentication", "True Passthrough (no LiteLLM auth)");
-      await user.type(screen.getByPlaceholderText("Leave blank to use dynamic client registration"), "app-id");
-      await user.type(screen.getByPlaceholderText("Leave blank for public clients / PKCE"), "app-secret");
+      fireEvent.change(screen.getByPlaceholderText("Leave blank to use dynamic client registration"), {
+        target: { value: "app-id" },
+      });
+      fireEvent.change(screen.getByPlaceholderText("Leave blank for public clients / PKCE"), {
+        target: { value: "app-secret" },
+      });
 
       await waitFor(() => expect(oauthHook.onTokenReceived).toBeTruthy());
       await act(async () => {
@@ -859,7 +886,9 @@ describe("CreateMCPServer", () => {
 
       // Editing only client_id fires an invalidation whose changedValues carries only the client_id
       // sub-field; the preserve + deep-merge re-apply must keep client_secret from being dropped.
-      await user.type(screen.getByPlaceholderText("Leave blank to use dynamic client registration"), "2");
+      fireEvent.change(screen.getByPlaceholderText("Leave blank to use dynamic client registration"), {
+        target: { value: "app-id2" },
+      });
 
       const cfKeystrokeServer = {
         server_id: "cf-keystroke",
@@ -886,8 +915,10 @@ describe("CreateMCPServer", () => {
     it("replaces the token set on re-authorize instead of leaving stale siblings", async () => {
       await selectHttpTransport();
       const user = userEvent.setup({ delay: null });
-      await user.type(getServerNameInput(), "Reauth_Server");
-      await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+      fireEvent.change(getServerNameInput(), { target: { value: "Reauth_Server" } });
+      fireEvent.change(screen.getByPlaceholderText("https://your-mcp-server.com"), {
+        target: { value: "https://example.com/mcp" },
+      });
       await selectAntOption("Authentication", "OAuth");
 
       await waitFor(() => expect(oauthHook.onTokenReceived).toBeTruthy());
@@ -922,10 +953,10 @@ describe("CreateMCPServer", () => {
       const user = userEvent.setup({ delay: null });
 
       const nameInput = getServerNameInput();
-      await user.type(nameInput, "No_Auth_Server");
+      fireEvent.change(nameInput, { target: { value: "No_Auth_Server" } });
 
       const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
-      await user.type(urlInput, "https://example.com/mcp");
+      fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } });
 
       await selectAntOption("Authentication", "None");
 
@@ -1001,15 +1032,15 @@ describe("CreateMCPServer", () => {
       const user = userEvent.setup({ delay: null });
 
       const nameInput = getServerNameInput();
-      await user.type(nameInput, "Limited_Server");
+      fireEvent.change(nameInput, { target: { value: "Limited_Server" } });
 
       const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
-      await user.type(urlInput, "https://example.com/mcp");
+      fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } });
 
       await selectAntOption("Authentication", "None");
 
       const limitInput = screen.getByPlaceholderText("e.g. 10");
-      await user.type(limitInput, "5");
+      fireEvent.change(limitInput, { target: { value: "5" } });
 
       vi.mocked(networking.createMCPServer).mockResolvedValue({
         server_id: "new-server-1",
@@ -1265,10 +1296,10 @@ describe("CreateMCPServer", () => {
       const user = userEvent.setup({ delay: null });
 
       const nameInput = getServerNameInput();
-      await user.type(nameInput, "Locked_Down_Server");
+      fireEvent.change(nameInput, { target: { value: "Locked_Down_Server" } });
 
       const urlInput = screen.getByPlaceholderText("https://your-mcp-server.com");
-      await user.type(urlInput, "https://example.com/mcp");
+      fireEvent.change(urlInput, { target: { value: "https://example.com/mcp" } });
 
       await selectAntOption("Authentication", "None");
 
@@ -2147,7 +2178,7 @@ describe("CreateMCPServer dcr_bridge toggle", () => {
   });
 
   // Forcing dcr_bridge false for every non-client-forwarded auth type is covered in
-  // createServerPayload.test.ts. The two form-state cases below stay: they prove the Form.Item
+  // createServerPayload.test.ts. The two form-state cases below stay: they prove the field
   // unmounts on a switch away, and that the live value survives a client-forwarded swap.
 
   it("forces dcr_bridge: false when the auth type is switched away after toggling", async () => {
@@ -2179,7 +2210,7 @@ describe("CreateMCPServer dcr_bridge toggle", () => {
     });
     expect(getDcrToggle()).toHaveAttribute("aria-checked", "true");
 
-    // The Form.Item is mounted in both client-forwarded modes, so switching between them keeps the
+    // The field is mounted in both client-forwarded modes, so switching between them keeps the
     // live toggle value rather than forcing it back to the default or to false.
     await selectAntOption("Authentication", "OAuth Delegate (client-supplied upstream token)");
     await waitFor(() => {

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.permissions.integration.test.tsx

@@ -0,0 +1,160 @@
+import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import * as networking from "@/components/networking";
+import CreateMCPServer from "./CreateMCPServer";
+import { selectAntOption } from "./testUtils";
+
+vi.mock("@/components/networking", () => ({
+  createMCPServer: vi.fn(),
+  fetchOpenAPIRegistry: vi.fn().mockResolvedValue({ apis: [] }),
+  registerMCPServer: vi.fn(),
+  storeMCPOAuthUserCredential: vi.fn().mockResolvedValue({}),
+  testMCPToolsListRequest: vi.fn().mockResolvedValue({ tools: [], error: null }),
+}));
+
+vi.mock("@/utils/mcpTokenStore", () => ({
+  setToken: vi.fn(),
+}));
+
+vi.mock("./OpenAPIQuickPicker", () => ({
+  default: () => null,
+}));
+
+vi.mock("@/hooks/useMcpOAuthFlow", () => ({
+  useMcpOAuthFlow: () => ({
+    startOAuthFlow: vi.fn(),
+    status: "idle",
+    error: null,
+    tokenResponse: null,
+    reset: vi.fn(),
+  }),
+}));
+
+vi.mock("./mcp_server_cost_config", () => ({
+  default: () => <div data-testid="mcp-cost-config" />,
+}));
+
+vi.mock("./mcp_tool_configuration", () => ({
+  default: () => <div data-testid="mcp-tool-config" />,
+}));
+
+vi.mock("./mcp_connection_status", () => ({
+  default: () => <div data-testid="mcp-connection-status" />,
+}));
+
+vi.mock("./StdioConfiguration", () => ({
+  default: () => <div data-testid="stdio-config" />,
+}));
+
+const defaultProps = {
+  userRole: "Admin",
+  accessToken: "test-token",
+  onCreateSuccess: vi.fn(),
+  isModalVisible: true,
+  setModalVisible: vi.fn(),
+  availableAccessGroups: ["group-a", "group-b"],
+};
+
+const getServerNameInput = () => document.getElementById("server_name") as HTMLInputElement;
+
+const switchFor = (labelText: string): HTMLElement => {
+  const label = screen.getByText(labelText);
+  const row = label.closest(".flex.items-start.justify-between");
+  const control = row?.querySelector("button[role='switch']");
+  if (control === null || control === undefined) {
+    throw new Error(`no switch found for "${labelText}"`);
+  }
+  return control as HTMLElement;
+};
+
+const fillMinimalHttpServer = async (name: string) => {
+  await selectAntOption("Transport Type", "Streamable HTTP");
+  await waitFor(() => {
+    expect(screen.getByPlaceholderText("https://your-mcp-server.com")).toBeInTheDocument();
+  });
+  const user = userEvent.setup({ delay: null });
+  await user.type(getServerNameInput(), name);
+  await user.type(screen.getByPlaceholderText("https://your-mcp-server.com"), "https://example.com/mcp");
+  await selectAntOption("Authentication", "None");
+};
+
+const submitAndReadPayload = async () => {
+  await act(async () => {
+    fireEvent.click(screen.getByRole("button", { name: "Add MCP Server" }));
+  });
+  await waitFor(() => expect(networking.createMCPServer).toHaveBeenCalledTimes(1));
+  return vi.mocked(networking.createMCPServer).mock.calls[0][1];
+};
+
+const createdServer = {
+  server_id: "new-server-1",
+  server_name: "Perm_Server",
+  alias: "Perm_Server",
+  url: "https://example.com/mcp",
+  transport: "http",
+  auth_type: "none",
+  created_at: "2024-01-01T00:00:00Z",
+  created_by: "user-1",
+  updated_at: "2024-01-01T00:00:00Z",
+  updated_by: "user-1",
+};
+
+describe("CreateMCPServer permission toggles reaching the payload", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.mocked(networking.createMCPServer).mockResolvedValue(createdServer);
+  });
+
+  it("sends the panel's untouched defaults rather than dropping the keys the panel owns", async () => {
+    render(<CreateMCPServer {...defaultProps} />);
+    await fillMinimalHttpServer("Perm_Server");
+
+    const payload = await submitAndReadPayload();
+
+    expect(payload.allow_all_keys).toBe(false);
+    expect(payload.available_on_public_internet).toBe(true);
+  });
+
+  it("sends allow_all_keys true once the operator turns the public-to-all-keys switch on", async () => {
+    render(<CreateMCPServer {...defaultProps} />);
+    await fillMinimalHttpServer("Perm_Server");
+
+    await act(async () => {
+      fireEvent.click(switchFor("Allow All LiteLLM Keys"));
+    });
+
+    const payload = await submitAndReadPayload();
+
+    expect(payload.allow_all_keys).toBe(true);
+  });
+
+  it("sends available_on_public_internet false when the operator restricts the server to the internal network", async () => {
+    render(<CreateMCPServer {...defaultProps} />);
+    await fillMinimalHttpServer("Perm_Server");
+
+    const internalOnly = switchFor("Internal network only");
+    expect(internalOnly).toHaveAttribute("aria-checked", "false");
+
+    await act(async () => {
+      fireEvent.click(internalOnly);
+    });
+    expect(internalOnly).toHaveAttribute("aria-checked", "true");
+
+    const payload = await submitAndReadPayload();
+
+    expect(payload.available_on_public_internet).toBe(false);
+  });
+
+  it("omits delegate_auth_to_upstream's true value on a none-auth server, whose gate never mounts that switch", async () => {
+    render(<CreateMCPServer {...defaultProps} />);
+    await fillMinimalHttpServer("Perm_Server");
+
+    expect(screen.getByText("Allow All LiteLLM Keys")).toBeInTheDocument();
+    expect(screen.queryByText("Delegate auth to upstream (PKCE passthrough)")).not.toBeInTheDocument();
+
+    const payload = await submitAndReadPayload();
+
+    expect(payload.delegate_auth_to_upstream).toBe(false);
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/CreateMCPServer.tsx

@@ -1,5 +1,6 @@
 import React, { useState } from "react";
-import { Modal, Tooltip, Form, Select, Input as AntdInput, InputNumber, Collapse } from "antd";
+import { Modal, Tooltip, Select, Input as AntdInput, InputNumber, Collapse } from "antd";
+import { FormProvider, useForm, useWatch } from "react-hook-form";
 import { InfoCircleOutlined } from "@ant-design/icons";
 import { Button } from "@/components/ui/button";
 import { Input } from "@/components/ui/input";
@@ -49,6 +50,16 @@ import { validateMCPServerUrl, validateMCPServerName } from "./utils";
 import { toast } from "@/lib/toast";
 import { useMcpOAuthFlow } from "@/hooks/useMcpOAuthFlow";
 import { useTestMCPConnection } from "@/hooks/useTestMCPConnection";
+import {
+  MountedFormField,
+  MountedFormProvider,
+  projectMountedValues,
+  useMountRegistry,
+  type MountedFormValues,
+} from "@/components/common_components/MountedFormField";
+import { antdRequired, antdRules } from "@/components/common_components/antdFormRules";
+import { allFieldsValue, mountedPaths, resetFields, setFieldsValue, singleBranchChange } from "./mcpFormStore";
+import { numberControl, notOnlyWhitespace, selectControl, textControl } from "./mcpFieldRules";
 import mcpLogo from "../../../../../public/assets/logos/mcp_logo.png";
 
 export const mcpLogoImg = mcpLogo.src;
@@ -76,6 +87,13 @@ const payloadErrorMessage = (result: Exclude<BuildCreatePayloadResult, { kind: "
   }
 };
 
+const CREATE_DEFAULTS: MountedFormValues = {
+  allow_all_keys: false,
+  available_on_public_internet: true,
+  delegate_auth_to_upstream: false,
+  oauth_passthrough: false,
+};
+
 const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
   userID,
   userRole,
@@ -87,7 +105,8 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
   prefillData,
   onBackToDiscovery,
 }) => {
-  const [form] = Form.useForm();
+  const form = useForm<MountedFormValues>({ mode: "onChange", defaultValues: CREATE_DEFAULTS });
+  const registry = useMountRegistry();
   const [isLoading, setIsLoading] = useState(false);
   const [costConfig, setCostConfig] = useState<MCPServerCostInfo>({});
   const [formValues, setFormValues] = useState<Record<string, any>>({});
@@ -136,6 +155,8 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
     enabled: true,
   });
 
+  const authSectionMounted = transportType !== "stdio" && transportType !== "";
+  const watchedAuthType = useWatch({ control: form.control, name: "auth_type" }) as string | undefined;
   const authType = formValues.auth_type as string | undefined;
   const shouldShowAuthValueField = authType ? AUTH_TYPES_REQUIRING_AUTH_VALUE.includes(authType) : false;
   const isOAuthAuthType = authType === AUTH_TYPE.OAUTH2;
@@ -147,7 +168,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
   const persistCreateUiState = () => {
     writeCreateUiSnapshot({
       modalVisible: isModalVisible,
-      formValues: form.getFieldsValue(true),
+      formValues: allFieldsValue(form),
       transportType,
       costConfig,
       allowedTools,
@@ -170,11 +191,11 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
     // Merge the ref-held DCR client so a re-authorize reuses the registered client instead of
     // re-registering; the form store itself never holds the DCR client (see onTokenReceived).
     getCredentials: () => ({
-      ...((form.getFieldValue("credentials") as Record<string, unknown> | undefined) ?? {}),
+      ...((allFieldsValue(form).credentials as Record<string, unknown> | undefined) ?? {}),
       ...(dcrClientRef.current ?? {}),
     }),
     getTemporaryPayload: () => {
-      const values = form.getFieldsValue(true);
+      const values = allFieldsValue(form);
       const transport = values.transport || transportType;
       // For OpenAPI transport the form has spec_path instead of url.
       // We pass the spec_path as url so the temp-session endpoint has something
@@ -218,12 +239,12 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
         return;
       }
 
-      if (isClientForwardedTokenMode(form.getFieldValue("auth_type"))) {
+      if (isClientForwardedTokenMode(allFieldsValue(form).auth_type)) {
         // Browser-only modes: the token is held in local state (oauthAccessToken) for tool preview
         // and committed to sessionStorage on submit; it must never be written into form.credentials,
         // which would persist it as server-level credentials on the created server row. Mirrors the
         // edit form's onTokenReceived early return.
-        setAuthorizedIdentity(getOAuthAuthorizationIdentity(form.getFieldsValue(true)));
+        setAuthorizedIdentity(getOAuthAuthorizationIdentity(allFieldsValue(form)));
         toast.success(
           "Token held for this browser session. Tools can now be previewed and configured; the token is not saved to LiteLLM.",
         );
@@ -240,7 +261,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
           }
         : null;
 
-      const current = (form.getFieldValue("credentials") as Record<string, unknown> | undefined) ?? {};
+      const current = (allFieldsValue(form).credentials as Record<string, unknown> | undefined) ?? {};
       const nextCredentials = {
         ...(preservedAdminCredentials(current) ?? {}),
         ...(current.scopes !== undefined && { scopes: current.scopes }),
@@ -252,10 +273,10 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
       // Path-replace (not deep-merge) so a re-authorize with fewer token fields does not leave stale
       // siblings from the previous token behind; the admin-typed client keys and scopes are carried
       // explicitly above.
-      form.setFieldValue("credentials", nextCredentials);
+      form.setValue("credentials", nextCredentials);
       // Capture the identity AFTER writing the token so the held token is not spuriously invalidated by
       // its own credential write.
-      setAuthorizedIdentity(getOAuthAuthorizationIdentity(form.getFieldsValue(true)));
+      setAuthorizedIdentity(getOAuthAuthorizationIdentity(allFieldsValue(form)));
 
       toast.success("OAuth authorization successful! Please click 'Create MCP Server' to save the configuration.");
     },
@@ -277,10 +298,10 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
     // Capture the admin-typed app before resetFields destroys it, then re-apply it: the app is
     // upstream-scoped config, not minted material, so it survives every invalidation (the token is
     // what gets discarded). Token-shaped keys are excluded by the helper's key filter.
-    const keptAdminCredentials = preservedAdminCredentials(form.getFieldValue("credentials"));
-    form.resetFields([...CLEARED_ON_INVALIDATION]);
+    const keptAdminCredentials = preservedAdminCredentials(allFieldsValue(form).credentials);
+    resetFields(form, [...CLEARED_ON_INVALIDATION]);
     if (keptAdminCredentials) {
-      form.setFieldsValue({ credentials: keptAdminCredentials });
+      setFieldsValue(form, { credentials: keptAdminCredentials });
     }
     // Re-apply the in-flight edit last; rc-field-form deep-merges nested objects, so a changed
     // credentials sub-field composes with the preserved sibling instead of replacing the object.
@@ -288,7 +309,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
       CLEARED_ON_INVALIDATION.filter((key) => key in changedValues).map((key) => [key, changedValues[key]]),
     );
     if (Object.keys(preserved).length > 0) {
-      form.setFieldsValue(preserved);
+      setFieldsValue(form, preserved);
     }
   };
 
@@ -337,7 +358,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
       // wait until transportType state catches up so the URL field is mounted
       return;
     }
-    form.setFieldsValue(pendingRestoredValues.values);
+    setFieldsValue(form, pendingRestoredValues.values);
     setFormValues(pendingRestoredValues.values);
     setPendingRestoredValues(null);
   }, [pendingRestoredValues, form, transportType]);
@@ -381,11 +402,20 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
       prefillValues.url = prefillData.url;
     }
 
-    form.setFieldsValue(prefillValues);
+    setFieldsValue(form, prefillValues);
     setFormValues(prefillValues);
     setAliasManuallyEdited(false);
   }, [isModalVisible, prefillData, form]);
 
+  const handleSubmit = async (event: React.FormEvent) => {
+    event.preventDefault();
+    const isValid = await form.trigger(mountedPaths(registry) as string[]);
+    if (!isValid) {
+      return;
+    }
+    await handleCreate(projectMountedValues(registry, form.getValues));
+  };
+
   const handleCreate = async (values: Record<string, unknown>) => {
     const built = buildCreateServerPayload(values, {
       transportType,
@@ -446,7 +476,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
             description: "Once an admin approves it, the server will appear in your MCP Servers list.",
           });
         }
-        form.resetFields();
+        form.reset(CREATE_DEFAULTS);
         setCostConfig({});
         clearTools();
         setAllowedTools([]);
@@ -466,7 +496,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
 
   // state
   const handleCancel = () => {
-    form.resetFields();
+    form.reset(CREATE_DEFAULTS);
     setCostConfig({});
     clearTools();
     setAllowedTools([]);
@@ -489,11 +519,11 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
           ? { url: undefined, command: undefined, args: undefined, env: undefined }
           : { spec_path: undefined, command: undefined, args: undefined, env: undefined };
 
-    form.setFieldsValue(transportValues);
-    if (isHeldOAuthTokenStale(form.getFieldsValue(true), authorizedIdentity)) {
+    setFieldsValue(form, transportValues);
+    if (isHeldOAuthTokenStale(allFieldsValue(form), authorizedIdentity)) {
       clearHeldOAuthToken();
     }
-    setFormValues(form.getFieldsValue(true));
+    setFormValues(allFieldsValue(form));
   };
 
   // Generate options with existing groups and potential new group
@@ -532,7 +562,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
   React.useEffect(() => {
     if (!aliasManuallyEdited && formValues.server_name) {
       const normalized = formValues.server_name.replace(/\s+/g, "_");
-      form.setFieldsValue({ alias: normalized });
+      setFieldsValue(form, { alias: normalized });
       setFormValues((prev) => ({ ...prev, alias: normalized }));
     }
   }, [formValues.server_name]);
@@ -549,7 +579,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
     const wasVisible = wasModalVisibleRef.current;
     wasModalVisibleRef.current = isModalVisible;
     if (!isModalVisible && wasVisible) {
-      form.resetFields();
+      form.reset(CREATE_DEFAULTS);
       setFormValues({});
       setOauthAccessToken(null);
       clearTools();
@@ -582,19 +612,35 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
       const upstreamChanged = ["url", "spec_path", "issuer", "authorization_url", "token_url", "registration_url"].some(
         (key) => key in changedValues,
       );
-      const hasDeclaredApp = preservedDeclaredAppCredentials(form.getFieldValue("credentials")) !== undefined;
+      const hasDeclaredApp = preservedDeclaredAppCredentials(allFieldsValue(form).credentials) !== undefined;
       if (upstreamChanged && hasDeclaredApp) {
         setAppMayNotMatchUpstream(true);
       }
     }
-    if (isHeldOAuthTokenStale(form.getFieldsValue(true), authorizedIdentity)) {
+    if (isHeldOAuthTokenStale(allFieldsValue(form), authorizedIdentity)) {
       clearHeldOAuthToken(changedValues);
-      setFormValues(form.getFieldsValue(true));
+      setFormValues(allFieldsValue(form));
       return;
     }
     setFormValues(allValues);
   };
 
+  const valuesChangeRef = React.useRef(handleFormValuesChange);
+  valuesChangeRef.current = handleFormValuesChange;
+
+  React.useEffect(() => {
+    const subscription = form.watch((values, { name, type }) => {
+      if (type !== "change" || name === undefined) {
+        return;
+      }
+      valuesChangeRef.current(
+        singleBranchChange(name, values as MountedFormValues),
+        projectMountedValues(registry, form.getValues),
+      );
+    });
+    return () => subscription.unsubscribe();
+  }, [form, registry]);
+
   // rendering
   return (
     <Modal
@@ -636,334 +682,370 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
       }}
     >
       <div className="mt-6">
-        <Form
-          form={form}
-          onFinish={handleCreate}
-          onValuesChange={handleFormValuesChange}
-          layout="vertical"
-          className="space-y-6"
-        >
-          {!isAdmin && (
-            <div className="rounded-md bg-blue-50 border border-blue-200 px-4 py-3 text-sm text-blue-800">
-              Your submission will be sent for admin review. Once approved, the server will appear in your MCP Servers
-              list. The request must be made with a team-scoped API key.
-            </div>
-          )}
-          <div className="grid grid-cols-1 gap-6">
-            <Form.Item
-              label={
-                <span className="text-sm font-medium text-gray-700 flex items-center">
-                  MCP Server Name
-                  <Tooltip title="Best practice: Use a descriptive name that indicates the server's purpose (e.g., 'GitHub_MCP', 'Email_Service'). Cannot contain spaces or hyphens; use underscores instead. Names must comply with SEP-986 and will be rejected if invalid (https://modelcontextprotocol.io/specification/2025-11-25/server/tools#tool-names).">
-                    <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-                  </Tooltip>
-                </span>
-              }
-              name="server_name"
-              rules={[
-                { required: false, message: "Please enter a server name" },
-                { validator: (_, value) => validateMCPServerName(value) },
-              ]}
-            >
-              <Input
-                placeholder="e.g., GitHub_MCP, Zapier_MCP, etc."
-                className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-              />
-            </Form.Item>
+        <FormProvider {...form}>
+          <MountedFormProvider value={{ control: form.control, registry }}>
+            <form onSubmit={handleSubmit} className="space-y-6">
+              {!isAdmin && (
+                <div className="rounded-md bg-blue-50 border border-blue-200 px-4 py-3 text-sm text-blue-800">
+                  Your submission will be sent for admin review. Once approved, the server will appear in your MCP
+                  Servers list. The request must be made with a team-scoped API key.
+                </div>
+              )}
+              <div className="grid grid-cols-1 gap-6">
+                <MountedFormField
+                  label={
+                    <span className="text-sm font-medium text-gray-700 flex items-center">
+                      MCP Server Name
+                      <Tooltip title="Best practice: Use a descriptive name that indicates the server's purpose (e.g., 'GitHub_MCP', 'Email_Service'). Cannot contain spaces or hyphens; use underscores instead. Names must comply with SEP-986 and will be rejected if invalid (https://modelcontextprotocol.io/specification/2025-11-25/server/tools#tool-names).">
+                        <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
+                      </Tooltip>
+                    </span>
+                  }
+                  name="server_name"
+                  rules={{ validate: antdRules({ validator: (_, value) => validateMCPServerName(value) }) }}
+                >
+                  {(control) => (
+                    <Input
+                      {...textControl(control)}
+                      placeholder="e.g., GitHub_MCP, Zapier_MCP, etc."
+                      className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
+                    />
+                  )}
+                </MountedFormField>
 
-            <Form.Item
-              label={
-                <span className="text-sm font-medium text-gray-700 flex items-center">
-                  Alias
-                  <Tooltip title="A short, unique identifier for this server. Defaults to the server name if not provided. Cannot contain spaces or hyphens; use underscores instead.">
-                    <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-                  </Tooltip>
-                </span>
-              }
-              name="alias"
-              rules={[{ required: false }, { validator: (_, value) => validateMCPServerName(value) }]}
-            >
-              <Input
-                placeholder="e.g., GitHub_MCP, Zapier_MCP, etc."
-                className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-                onChange={() => setAliasManuallyEdited(true)}
-              />
-            </Form.Item>
+                <MountedFormField
+                  label={
+                    <span className="text-sm font-medium text-gray-700 flex items-center">
+                      Alias
+                      <Tooltip title="A short, unique identifier for this server. Defaults to the server name if not provided. Cannot contain spaces or hyphens; use underscores instead.">
+                        <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
+                      </Tooltip>
+                    </span>
+                  }
+                  name="alias"
+                  rules={{ validate: antdRules({ validator: (_, value) => validateMCPServerName(value) }) }}
+                >
+                  {(control) => (
+                    <Input
+                      {...textControl(control)}
+                      placeholder="e.g., GitHub_MCP, Zapier_MCP, etc."
+                      className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
+                      onChange={(event) => {
+                        control.onChange(event);
+                        setAliasManuallyEdited(true);
+                      }}
+                    />
+                  )}
+                </MountedFormField>
 
-            <Form.Item
-              label={<span className="text-sm font-medium text-gray-700">Description</span>}
-              name="description"
-              rules={[
-                {
-                  required: false,
-                  message: "Please enter a server description",
-                },
-              ]}
-            >
-              <Input
-                placeholder="Brief description of what this server does"
-                className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-              />
-            </Form.Item>
+                <MountedFormField
+                  label={<span className="text-sm font-medium text-gray-700">Description</span>}
+                  name="description"
+                >
+                  {(control) => (
+                    <Input
+                      {...textControl(control)}
+                      placeholder="Brief description of what this server does"
+                      className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
+                    />
+                  )}
+                </MountedFormField>
 
-            <MCPLogoSelector value={logoUrl} onChange={setLogoUrl} />
+                <MCPLogoSelector value={logoUrl} onChange={setLogoUrl} />
 
-            <Form.Item
-              label={<span className="text-sm font-medium text-gray-700">GitHub / Source URL</span>}
-              name="source_url"
-            >
-              <Input
-                placeholder="https://github.com/org/mcp-server"
-                className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-              />
-            </Form.Item>
+                <MountedFormField
+                  label={<span className="text-sm font-medium text-gray-700">GitHub / Source URL</span>}
+                  name="source_url"
+                >
+                  {(control) => (
+                    <Input
+                      {...textControl(control)}
+                      placeholder="https://github.com/org/mcp-server"
+                      className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
+                    />
+                  )}
+                </MountedFormField>
 
-            <Form.Item
-              label={<span className="text-sm font-medium text-gray-700">Transport Type</span>}
-              name="transport"
-              rules={[{ required: true, message: "Please select a transport type" }]}
-            >
-              <Select
-                placeholder="Select transport"
-                className="rounded-lg"
-                size="large"
-                onChange={handleTransportChange}
-                value={transportType}
-              >
-                <Select.Option value="http">Streamable HTTP (Recommended)</Select.Option>
-                <Select.Option value="sse">Server-Sent Events (SSE)</Select.Option>
-                <Select.Option value="stdio">Standard Input/Output (stdio)</Select.Option>
-                <Select.Option value={TRANSPORT.OPENAPI}>OpenAPI Spec</Select.Option>
-              </Select>
-            </Form.Item>
+                <MountedFormField
+                  label={<span className="text-sm font-medium text-gray-700">Transport Type</span>}
+                  name="transport"
+                  required
+                  rules={{ validate: { required: antdRequired("Please select a transport type") } }}
+                >
+                  {(control) => (
+                    <Select
+                      {...selectControl<string>(control)}
+                      placeholder="Select transport"
+                      className="rounded-lg"
+                      size="large"
+                      onChange={(value: string) => {
+                        control.onChange(value);
+                        handleTransportChange(value);
+                      }}
+                    >
+                      <Select.Option value="http">Streamable HTTP (Recommended)</Select.Option>
+                      <Select.Option value="sse">Server-Sent Events (SSE)</Select.Option>
+                      <Select.Option value="stdio">Standard Input/Output (stdio)</Select.Option>
+                      <Select.Option value={TRANSPORT.OPENAPI}>OpenAPI Spec</Select.Option>
+                    </Select>
+                  )}
+                </MountedFormField>
 
-            {/* URL field - only show for HTTP and SSE */}
-            {(transportType === "http" || transportType === "sse") && (
-              <Form.Item
-                label={<span className="text-sm font-medium text-gray-700">MCP Server URL</span>}
-                name="url"
-                rules={[
-                  { required: true, message: "Please enter a server URL" },
-                  { validator: (_, value) => validateMCPServerUrl(value) },
-                ]}
-              >
-                <AntdInput
-                  placeholder="https://your-mcp-server.com"
-                  className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
-                />
-              </Form.Item>
-            )}
+                {/* URL field - only show for HTTP and SSE */}
+                {(transportType === "http" || transportType === "sse") && (
+                  <MountedFormField
+                    label={<span className="text-sm font-medium text-gray-700">MCP Server URL</span>}
+                    name="url"
+                    required
+                    rules={{
+                      validate: {
+                        required: antdRequired("Please enter a server URL"),
+                        ...antdRules({ validator: (_, value) => validateMCPServerUrl(value) }),
+                      },
+                    }}
+                  >
+                    {(control) => (
+                      <AntdInput
+                        {...textControl(control)}
+                        placeholder="https://your-mcp-server.com"
+                        className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
+                      />
+                    )}
+                  </MountedFormField>
+                )}
 
-            {/* OpenAPI: logo picker + spec URL input */}
-            {transportType === TRANSPORT.OPENAPI && (
-              <OpenAPIFormSection
-                form={form}
-                accessToken={isModalVisible ? accessToken : null}
-                onValuesChange={(updates) =>
-                  handleFormValuesChange(updates, { ...form.getFieldsValue(true), ...updates })
-                }
-                onKeyToolsChange={setKeyTools}
-                onLogoUrlChange={setLogoUrl}
-                onOAuthDocsUrlChange={setOauthDocsUrl}
-              />
-            )}
+                {/* OpenAPI: logo picker + spec URL input */}
+                {transportType === TRANSPORT.OPENAPI && (
+                  <OpenAPIFormSection
+                    form={form}
+                    accessToken={isModalVisible ? accessToken : null}
+                    onValuesChange={(updates) =>
+                      handleFormValuesChange(updates, { ...allFieldsValue(form), ...updates })
+                    }
+                    onKeyToolsChange={setKeyTools}
+                    onLogoUrlChange={setLogoUrl}
+                    onOAuthDocsUrlChange={setOauthDocsUrl}
+                  />
+                )}
 
-            {/* BYOK toggle - only for OpenAPI */}
-            {transportType === TRANSPORT.OPENAPI && <OpenApiByokFields />}
+                {/* BYOK toggle - only for OpenAPI */}
+                {transportType === TRANSPORT.OPENAPI && <OpenApiByokFields />}
 
-            <Form.Item
-              label={
-                <span className="text-sm font-medium text-gray-700 flex items-center">
-                  Max Concurrent Requests (optional)
-                  <Tooltip title="Maximum number of tool calls LiteLLM will run against this server at the same time. Additional calls wait for a free slot. Leave blank for no limit.">
-                    <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-                  </Tooltip>
-                </span>
-              }
-              name="max_concurrent_requests"
-            >
-              <InputNumber
-                min={1}
-                precision={0}
-                placeholder="e.g. 10"
-                style={{ width: "100%" }}
-                className="rounded-lg"
-              />
-            </Form.Item>
+                <MountedFormField
+                  label={
+                    <span className="text-sm font-medium text-gray-700 flex items-center">
+                      Max Concurrent Requests (optional)
+                      <Tooltip title="Maximum number of tool calls LiteLLM will run against this server at the same time. Additional calls wait for a free slot. Leave blank for no limit.">
+                        <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
+                      </Tooltip>
+                    </span>
+                  }
+                  name="max_concurrent_requests"
+                >
+                  {(control) => (
+                    <InputNumber
+                      {...numberControl(control)}
+                      min={1}
+                      precision={0}
+                      placeholder="e.g. 10"
+                      style={{ width: "100%" }}
+                      className="rounded-lg"
+                    />
+                  )}
+                </MountedFormField>
 
-            {/* Authentication - show for HTTP, SSE, and OpenAPI */}
-            {transportType !== "stdio" && transportType !== "" && (
-              <Collapse
-                defaultActiveKey={["auth"]}
-                className="mb-4"
-                items={[
-                  {
-                    key: "auth",
-                    label: <span className="text-sm font-semibold text-gray-700">Authentication</span>,
-                    children: (
-                      <>
-                        <Form.Item name="auth_type" rules={[{ required: true, message: "Please select an auth type" }]}>
-                          <Select placeholder="Select auth type" className="rounded-lg" size="large" virtual={false}>
-                            <Select.Option value="none">None</Select.Option>
-                            <Select.Option value="api_key">API Key</Select.Option>
-                            <Select.Option value="bearer_token">Bearer Token</Select.Option>
-                            <Select.Option value="token">Token</Select.Option>
-                            <Select.Option value="basic">Basic Auth</Select.Option>
-                            <Select.Option value="oauth2">OAuth</Select.Option>
-                            <Select.Option value="oauth2_token_exchange">OAuth Token Exchange (OBO)</Select.Option>
-                            <Select.Option value="oauth2_id_jag">ID-JAG (Okta Cross App Access)</Select.Option>
-                            <Select.Option value="aws_sigv4">AWS SigV4 (Bedrock AgentCore MCPs)</Select.Option>
-                            <Select.Option value="true_passthrough">True Passthrough (no LiteLLM auth)</Select.Option>
-                            <Select.Option value="oauth_delegate">
-                              OAuth Delegate (client-supplied upstream token)
-                            </Select.Option>
-                          </Select>
-                        </Form.Item>
+                {/* Authentication - show for HTTP, SSE, and OpenAPI */}
+                {transportType !== "stdio" && transportType !== "" && (
+                  <Collapse
+                    defaultActiveKey={["auth"]}
+                    className="mb-4"
+                    items={[
+                      {
+                        key: "auth",
+                        label: <span className="text-sm font-semibold text-gray-700">Authentication</span>,
+                        children: (
+                          <>
+                            <MountedFormField
+                              name="auth_type"
+                              required
+                              rules={{ validate: { required: antdRequired("Please select an auth type") } }}
+                            >
+                              {(control) => (
+                                <Select
+                                  {...selectControl<string>(control)}
+                                  placeholder="Select auth type"
+                                  className="rounded-lg"
+                                  size="large"
+                                  virtual={false}
+                                >
+                                  <Select.Option value="none">None</Select.Option>
+                                  <Select.Option value="api_key">API Key</Select.Option>
+                                  <Select.Option value="bearer_token">Bearer Token</Select.Option>
+                                  <Select.Option value="token">Token</Select.Option>
+                                  <Select.Option value="basic">Basic Auth</Select.Option>
+                                  <Select.Option value="oauth2">OAuth</Select.Option>
+                                  <Select.Option value="oauth2_token_exchange">
+                                    OAuth Token Exchange (OBO)
+                                  </Select.Option>
+                                  <Select.Option value="oauth2_id_jag">ID-JAG (Okta Cross App Access)</Select.Option>
+                                  <Select.Option value="aws_sigv4">AWS SigV4 (Bedrock AgentCore MCPs)</Select.Option>
+                                  <Select.Option value="true_passthrough">
+                                    True Passthrough (no LiteLLM auth)
+                                  </Select.Option>
+                                  <Select.Option value="oauth_delegate">
+                                    OAuth Delegate (client-supplied upstream token)
+                                  </Select.Option>
+                                </Select>
+                              )}
+                            </MountedFormField>
 
-                        <TruePassthroughWarning authType={authType} />
+                            <TruePassthroughWarning authType={authType} />
 
-                        <PassthroughAuthorizeSection
-                          authType={authType}
-                          dcrBridgeInitialChecked
-                          oauthFlow={{
-                            startOAuthFlow,
-                            status: oauthStatus,
-                            error: oauthError,
-                            tokenResponse: oauthTokenResponse,
-                          }}
-                          appMayNotMatchUpstream={appMayNotMatchUpstream}
-                        />
-
-                        {shouldShowAuthValueField && (
-                          <Form.Item
-                            label={
-                              <span className="text-sm font-medium text-gray-700 flex items-center">
-                                Authentication Value
-                                <Tooltip title="Token, password, or header value to send with each request for the selected auth type.">
-                                  <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
-                                </Tooltip>
-                              </span>
-                            }
-                            name={["credentials", "auth_value"]}
-                            rules={[
-                              {
-                                validator: (_, value) =>
-                                  value && typeof value === "string" && value.trim() === ""
-                                    ? Promise.reject(new Error("Authentication value cannot be empty whitespace"))
-                                    : Promise.resolve(),
-                              },
-                            ]}
-                          >
-                            <AntdInput.Password
-                              placeholder="Enter token or secret"
-                              className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
+                            <PassthroughAuthorizeSection
+                              authType={authType}
+                              dcrBridgeInitialChecked
+                              oauthFlow={{
+                                startOAuthFlow,
+                                status: oauthStatus,
+                                error: oauthError,
+                                tokenResponse: oauthTokenResponse,
+                              }}
+                              appMayNotMatchUpstream={appMayNotMatchUpstream}
                             />
-                          </Form.Item>
-                        )}
 
-                        {isOAuthAuthType && (
-                          <OAuthFormFields
-                            isM2M={isM2MFlow}
-                            initialFlowType={OAUTH_FLOW.INTERACTIVE}
-                            docsUrl={oauthDocsUrl}
-                            oauthFlow={{
-                              startOAuthFlow,
-                              status: oauthStatus,
-                              error: oauthError,
-                              tokenResponse: oauthTokenResponse,
-                            }}
-                          />
-                        )}
+                            {shouldShowAuthValueField && (
+                              <MountedFormField
+                                label={
+                                  <span className="text-sm font-medium text-gray-700 flex items-center">
+                                    Authentication Value
+                                    <Tooltip title="Token, password, or header value to send with each request for the selected auth type.">
+                                      <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
+                                    </Tooltip>
+                                  </span>
+                                }
+                                name={["credentials", "auth_value"]}
+                                rules={{
+                                  validate: {
+                                    notWhitespace: notOnlyWhitespace("Authentication value cannot be empty whitespace"),
+                                  },
+                                }}
+                              >
+                                {(control) => (
+                                  <AntdInput.Password
+                                    {...textControl(control)}
+                                    placeholder="Enter token or secret"
+                                    className="rounded-lg border-gray-300 focus:border-blue-500 focus:ring-blue-500"
+                                  />
+                                )}
+                              </MountedFormField>
+                            )}
 
-                        {isTokenExchangeAuthType && <TokenExchangeFormFields />}
+                            {isOAuthAuthType && (
+                              <OAuthFormFields
+                                isM2M={isM2MFlow}
+                                initialFlowType={OAUTH_FLOW.INTERACTIVE}
+                                docsUrl={oauthDocsUrl}
+                                oauthFlow={{
+                                  startOAuthFlow,
+                                  status: oauthStatus,
+                                  error: oauthError,
+                                  tokenResponse: oauthTokenResponse,
+                                }}
+                              />
+                            )}
 
-                        {isIdJagAuthType && <IdJagFormFields />}
-                      </>
-                    ),
-                  },
-                ]}
-              />
-            )}
+                            {isTokenExchangeAuthType && <TokenExchangeFormFields />}
 
-            {transportType !== "stdio" && transportType !== "" && isAwsSigV4AuthType && <AwsSigV4Fields />}
+                            {isIdJagAuthType && <IdJagFormFields />}
+                          </>
+                        ),
+                      },
+                    ]}
+                  />
+                )}
 
-            {/* Stdio Configuration - only show for stdio transport */}
-            <StdioConfiguration isVisible={transportType === "stdio"} />
-          </div>
+                {transportType !== "stdio" && transportType !== "" && isAwsSigV4AuthType && <AwsSigV4Fields />}
 
-          {/* Environment Variables Section */}
-          <div className="mt-8">
-            <EnvVarsSection />
-          </div>
+                {/* Stdio Configuration - only show for stdio transport */}
+                <StdioConfiguration isVisible={transportType === "stdio"} />
+              </div>
 
-          {/* Permission Management / Access Control Section */}
-          <div className="mt-8">
-            <MCPPermissionManagement
-              availableAccessGroups={availableAccessGroups}
-              mcpServer={null}
-              searchValue={searchValue}
-              setSearchValue={setSearchValue}
-              getAccessGroupOptions={getAccessGroupOptions}
-            />
-          </div>
+              {/* Environment Variables Section */}
+              <div className="mt-8">
+                <EnvVarsSection />
+              </div>
 
-          {/* Connection Status Section */}
-          <div className="mt-8 pt-6 border-t border-gray-200">
-            <MCPConnectionStatus
-              formValues={formValues}
-              tools={tools}
-              isLoadingTools={isLoadingTools}
-              toolsError={toolsError}
-              toolsErrorStatus={toolsErrorStatus}
-              toolsErrorStackTrace={toolsErrorStackTrace}
-              canFetchTools={canFetchTools}
-              fetchTools={fetchTools}
-            />
-          </div>
+              {/* Permission Management / Access Control Section */}
+              <div className="mt-8">
+                <MCPPermissionManagement
+                  availableAccessGroups={availableAccessGroups}
+                  mcpServer={null}
+                  mountedAuthType={authSectionMounted ? watchedAuthType : undefined}
+                  searchValue={searchValue}
+                  setSearchValue={setSearchValue}
+                  getAccessGroupOptions={getAccessGroupOptions}
+                />
+              </div>
 
-          {/* Tool Configuration Section */}
-          <div className="mt-6">
-            <MCPToolConfiguration
-              accessToken={accessToken}
-              formValues={formValues}
-              allowedTools={allowedTools}
-              existingAllowedTools={null}
-              onAllowedToolsChange={setAllowedTools}
-              hasToolAllowlistInteraction={hasToolAllowlistInteraction}
-              onToolAllowlistInteraction={() => setHasToolAllowlistInteraction(true)}
-              toolNameToDisplayName={toolNameToDisplayName}
-              toolNameToDescription={toolNameToDescription}
-              onToolNameToDisplayNameChange={setToolNameToDisplayName}
-              onToolNameToDescriptionChange={setToolNameToDescription}
-              keyTools={keyTools}
-              externalTools={tools}
-              externalIsLoading={isLoadingTools}
-              externalError={toolsError}
-              externalErrorStatus={toolsErrorStatus}
-              externalCanFetch={canFetchTools}
-            />
-          </div>
+              {/* Connection Status Section */}
+              <div className="mt-8 pt-6 border-t border-gray-200">
+                <MCPConnectionStatus
+                  formValues={formValues}
+                  tools={tools}
+                  isLoadingTools={isLoadingTools}
+                  toolsError={toolsError}
+                  toolsErrorStatus={toolsErrorStatus}
+                  toolsErrorStackTrace={toolsErrorStackTrace}
+                  canFetchTools={canFetchTools}
+                  fetchTools={fetchTools}
+                />
+              </div>
 
-          {/* Cost Configuration Section */}
-          <div className="mt-6">
-            <MCPServerCostConfig
-              value={costConfig}
-              onChange={setCostConfig}
-              tools={tools.filter((tool) => allowedTools.includes(tool.name))}
-              disabled={false}
-            />
-          </div>
+              {/* Tool Configuration Section */}
+              <div className="mt-6">
+                <MCPToolConfiguration
+                  accessToken={accessToken}
+                  formValues={formValues}
+                  allowedTools={allowedTools}
+                  existingAllowedTools={null}
+                  onAllowedToolsChange={setAllowedTools}
+                  hasToolAllowlistInteraction={hasToolAllowlistInteraction}
+                  onToolAllowlistInteraction={() => setHasToolAllowlistInteraction(true)}
+                  toolNameToDisplayName={toolNameToDisplayName}
+                  toolNameToDescription={toolNameToDescription}
+                  onToolNameToDisplayNameChange={setToolNameToDisplayName}
+                  onToolNameToDescriptionChange={setToolNameToDescription}
+                  keyTools={keyTools}
+                  externalTools={tools}
+                  externalIsLoading={isLoadingTools}
+                  externalError={toolsError}
+                  externalErrorStatus={toolsErrorStatus}
+                  externalCanFetch={canFetchTools}
+                />
+              </div>
 
-          <div className="flex items-center justify-end space-x-3 pt-6 border-t border-gray-100">
-            <Button variant="secondary" onClick={handleCancel}>
-              Cancel
-            </Button>
-            <Button type="submit" disabled={isLoading} aria-busy={isLoading}>
-              {isLoading && <UiLoadingSpinner className="size-4" />}
-              {isLoading ? "Creating..." : "Add MCP Server"}
-            </Button>
-          </div>
-        </Form>
+              {/* Cost Configuration Section */}
+              <div className="mt-6">
+                <MCPServerCostConfig
+                  value={costConfig}
+                  onChange={setCostConfig}
+                  tools={tools.filter((tool) => allowedTools.includes(tool.name))}
+                  disabled={false}
+                />
+              </div>
+
+              <div className="flex items-center justify-end space-x-3 pt-6 border-t border-gray-100">
+                <Button variant="secondary" onClick={handleCancel}>
+                  Cancel
+                </Button>
+                <Button type="submit" disabled={isLoading} aria-busy={isLoading}>
+                  {isLoading && <UiLoadingSpinner className="size-4" />}
+                  {isLoading ? "Creating..." : "Add MCP Server"}
+                </Button>
+              </div>
+            </form>
+          </MountedFormProvider>
+        </FormProvider>
       </div>
     </Modal>
   );

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/DcrBridgeToggle.tsx

@@ -1,16 +1,19 @@
 import React from "react";
-import { Form, Switch, Tooltip } from "antd";
+import { Switch, Tooltip } from "antd";
 import { InfoCircleOutlined } from "@ant-design/icons";
+
+import { MountedFormField } from "@/components/common_components/MountedFormField";
 import { isClientForwardedTokenMode } from "@/components/mcp_tools/types";
+import { switchControl } from "./mcpFieldRules";
 
 /**
  * DCR-bridge toggle for the client-forwarded token modes (true_passthrough /
  * oauth_delegate); self-gates to those two auth types and renders nothing
  * otherwise. When on, OAuth-only clients like Claude Desktop can register and
  * sign in through the gateway; when off, the gateway relays the upstream
- * server's own OAuth metadata instead. `initialChecked` seeds the antd
- * Form.Item `initialValue` (not the Switch's DOM defaultChecked): the create
- * form defaults it on, the edit form seeds it from the stored value.
+ * server's own OAuth metadata instead. `initialChecked` seeds the field's
+ * default value (not the Switch's DOM defaultChecked): the create form defaults
+ * it on, the edit form seeds it from the stored value.
  */
 export default function DcrBridgeToggle({
   authType,
@@ -21,7 +24,7 @@ export default function DcrBridgeToggle({
 }) {
   if (!isClientForwardedTokenMode(authType)) return null;
   return (
-    <Form.Item
+    <MountedFormField
       label={
         <span className="text-sm font-medium text-gray-700 flex items-center">
           Gateway-hosted sign-in (DCR bridge)
@@ -31,10 +34,9 @@ export default function DcrBridgeToggle({
         </span>
       }
       name="dcr_bridge"
-      valuePropName="checked"
-      initialValue={initialChecked}
+      defaultValue={initialChecked}
     >
-      <Switch />
-    </Form.Item>
+      {(control) => <Switch {...switchControl(control)} />}
+    </MountedFormField>
   );
 }

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.test.tsx

@@ -0,0 +1,86 @@
+import React from "react";
+import { describe, it, expect, vi } from "vitest";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { FormProvider, useForm } from "react-hook-form";
+
+import {
+  MountedFormProvider,
+  projectMountedValues,
+  useMountRegistry,
+  type MountedFormValues,
+} from "@/components/common_components/MountedFormField";
+import EnvVarsSection from "./EnvVarsSection";
+
+const renderSection = (defaultValues: MountedFormValues) => {
+  const onFinish = vi.fn();
+  const Harness: React.FC = () => {
+    const form = useForm<MountedFormValues>({ mode: "onChange", defaultValues });
+    const registry = useMountRegistry();
+    return (
+      <FormProvider {...form}>
+        <MountedFormProvider value={{ control: form.control, registry }}>
+          <form
+            onSubmit={(event) => {
+              event.preventDefault();
+              onFinish(projectMountedValues(registry, form.getValues));
+            }}
+          >
+            <EnvVarsSection />
+            <button type="submit">Submit</button>
+          </form>
+        </MountedFormProvider>
+      </FormProvider>
+    );
+  };
+  render(<Harness />);
+  return onFinish;
+};
+
+describe("EnvVarsSection", () => {
+  it("submits a per-user row whole, keeping the value key whose input the scope hides", async () => {
+    const onFinish = renderSection({
+      env_vars: [{ name: "DB_USER", value: "admin", scope: "user", description: "Your DB username" }],
+    });
+
+    expect(screen.queryByPlaceholderText("e.g. postgresql")).not.toBeInTheDocument();
+    await userEvent.click(screen.getByText("Submit"));
+
+    expect(onFinish).toHaveBeenCalledWith(
+      expect.objectContaining({
+        env_vars: [{ name: "DB_USER", value: "admin", scope: "user", description: "Your DB username" }],
+      }),
+    );
+  });
+
+  it("submits an empty env_vars key when the list has no rows, rather than dropping the key", async () => {
+    const onFinish = renderSection({ env_vars: [] });
+
+    await userEvent.click(screen.getByText("Submit"));
+
+    expect(onFinish.mock.calls[0][0]).toHaveProperty("env_vars", []);
+  });
+
+  it("carries a row added after mount into the submitted list, scoped global without the user picking one", async () => {
+    const onFinish = renderSection({ env_vars: [] });
+
+    await userEvent.click(screen.getByText("Add Variable"));
+    await userEvent.type(screen.getByPlaceholderText("e.g. DB_PROTOCOL"), "DB_PROTOCOL");
+    await userEvent.type(screen.getByPlaceholderText("e.g. postgresql"), "postgresql");
+    await userEvent.click(screen.getByText("Submit"));
+
+    expect(onFinish).toHaveBeenCalledWith(
+      expect.objectContaining({
+        env_vars: [expect.objectContaining({ name: "DB_PROTOCOL", value: "postgresql", scope: "global" })],
+      }),
+    );
+  });
+
+  it("rejects a variable name that starts with a digit", async () => {
+    renderSection({ env_vars: [{ name: "", value: "", scope: "global", description: "" }] });
+
+    await userEvent.type(screen.getByPlaceholderText("e.g. DB_PROTOCOL"), "9LIVES");
+
+    expect(await screen.findByText("Use letters, digits, underscores; cannot start with a digit.")).toBeInTheDocument();
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/EnvVarsSection.tsx

@@ -1,6 +1,16 @@
 import React from "react";
-import { Form, Input, Select, Button, Tooltip, Typography } from "antd";
+import { Input, Select, Button, Tooltip, Typography } from "antd";
 import { InfoCircleOutlined, MinusCircleOutlined, PlusOutlined } from "@ant-design/icons";
+import { useFieldArray, useFormContext, useWatch } from "react-hook-form";
+
+import {
+  MountedFormField,
+  useMountedName,
+  type MountedFormValues,
+} from "@/components/common_components/MountedFormField";
+import { antdRequired } from "@/components/common_components/antdFormRules";
+import { matchesPattern, selectControl, textControl } from "./mcpFieldRules";
+import { listControl } from "./mcpFormStore";
 
 const { Text } = Typography;
 
@@ -9,6 +19,8 @@ const SCOPE_OPTIONS = [
   { value: "user", label: "Per-user" },
 ];
 
+const VARIABLE_NAME_PATTERN = /^[A-Za-z_][A-Za-z0-9_]*$/;
+
 /**
  * Form section for admin-configured MCP environment variables.
  *
@@ -20,6 +32,10 @@ const SCOPE_OPTIONS = [
  * The parent form reads the ``env_vars`` field from the form values.
  */
 const EnvVarsSection: React.FC = () => {
+  const { control } = useFormContext<MountedFormValues>();
+  const { fields, append, remove } = useFieldArray({ control: listControl(control), name: "env_vars" });
+  useMountedName("env_vars");
+
   return (
     <div className="rounded-lg border border-gray-200 bg-gray-50 p-4">
       <div className="flex items-center gap-2 mb-1">
@@ -48,60 +64,52 @@ const EnvVarsSection: React.FC = () => {
         </code>
       </Text>
 
-      <Form.List name="env_vars">
-        {(fields, { add, remove }) => (
-          <div className="space-y-2">
-            {fields.length > 0 && (
-              <div className="flex gap-3 px-1 text-xs font-medium text-gray-500 uppercase tracking-wide">
-                <div style={{ flex: 1 }}>Variable Name</div>
-                <div style={{ flex: 1 }}>Value / Description</div>
-                <div style={{ width: 160 }}>Scope</div>
-                <div style={{ width: 24 }} />
-              </div>
-            )}
-            {fields.map(({ key, name, ...restField }) => (
-              <div key={key} className="flex gap-3 items-start">
-                <Form.Item
-                  {...restField}
-                  name={[name, "name"]}
-                  className="mb-0"
-                  style={{ flex: 1 }}
-                  rules={[
-                    { required: true, message: "Variable name is required" },
-                    {
-                      pattern: /^[A-Za-z_][A-Za-z0-9_]*$/,
-                      message: "Use letters, digits, underscores; cannot start with a digit.",
-                    },
-                  ]}
-                >
-                  <Input placeholder="e.g. DB_PROTOCOL" className="rounded-md font-mono" />
-                </Form.Item>
-                <div style={{ flex: 1 }}>
-                  <ScopedValueOrDescription name={name} restField={restField} />
-                </div>
-                <Form.Item
-                  {...restField}
-                  name={[name, "scope"]}
-                  className="mb-0"
-                  initialValue="global"
-                  style={{ width: 160 }}
-                >
-                  <Select options={SCOPE_OPTIONS} />
-                </Form.Item>
-                <div style={{ width: 24, height: 32 }} className="flex items-center justify-center">
-                  <MinusCircleOutlined
-                    onClick={() => remove(name)}
-                    className="text-gray-500 hover:text-red-500 cursor-pointer"
-                  />
-                </div>
-              </div>
-            ))}
-            <Button type="dashed" onClick={() => add({ scope: "global" })} icon={<PlusOutlined />} block>
-              Add Variable
-            </Button>
+      <div className="space-y-2">
+        {fields.length > 0 && (
+          <div className="flex gap-3 px-1 text-xs font-medium text-gray-500 uppercase tracking-wide">
+            <div style={{ flex: 1 }}>Variable Name</div>
+            <div style={{ flex: 1 }}>Value / Description</div>
+            <div style={{ width: 160 }}>Scope</div>
+            <div style={{ width: 24 }} />
           </div>
         )}
-      </Form.List>
+        {fields.map((item, index) => (
+          <div key={item.id} className="flex gap-3 items-start">
+            <MountedFormField
+              name={["env_vars", String(index), "name"]}
+              className="mb-0 flex-1"
+              rules={{
+                validate: {
+                  required: antdRequired("Variable name is required"),
+                  pattern: matchesPattern(
+                    VARIABLE_NAME_PATTERN,
+                    "Use letters, digits, underscores; cannot start with a digit.",
+                  ),
+                },
+              }}
+            >
+              {(control) => (
+                <Input {...textControl(control)} placeholder="e.g. DB_PROTOCOL" className="rounded-md font-mono" />
+              )}
+            </MountedFormField>
+            <div style={{ flex: 1 }}>
+              <ScopedValueOrDescription index={index} />
+            </div>
+            <MountedFormField name={["env_vars", String(index), "scope"]} className="mb-0 w-40" defaultValue="global">
+              {(control) => <Select {...selectControl<string>(control)} options={SCOPE_OPTIONS} />}
+            </MountedFormField>
+            <div style={{ width: 24, height: 32 }} className="flex items-center justify-center">
+              <MinusCircleOutlined
+                onClick={() => remove(index)}
+                className="text-gray-500 hover:text-red-500 cursor-pointer"
+              />
+            </div>
+          </div>
+        ))}
+        <Button type="dashed" onClick={() => append({ scope: "global" })} icon={<PlusOutlined />} block>
+          Add Variable
+        </Button>
+      </div>
     </div>
   );
 };
@@ -109,33 +117,33 @@ const EnvVarsSection: React.FC = () => {
 // For instance-scoped vars this column holds the admin value. For per-user
 // vars the value comes from each user later, so the column instead captures an
 // optional description that the per-user fill-in modal shows as a hint.
-const ScopedValueOrDescription: React.FC<{
-  name: number;
-  restField: object;
-}> = ({ name, restField }) => {
-  const isPerUser = Form.useWatch(["env_vars", name, "scope"]) === "user";
+const ScopedValueOrDescription: React.FC<{ index: number }> = ({ index }) => {
+  const isPerUser = useWatch({ name: `env_vars.${index}.scope` }) === "user";
   if (isPerUser) {
     return (
-      <Form.Item {...restField} name={[name, "description"]} className="mb-0">
-        <Input
-          addonBefore={
-            <Tooltip title="Per-user variables have no shared value. This text is only a hint shown to each user when they fill in their own value.">
-              <span className="text-xs text-gray-500 cursor-help whitespace-nowrap">
-                <InfoCircleOutlined className="mr-1" />
-                Hint
-              </span>
-            </Tooltip>
-          }
-          placeholder="e.g. Your DB username"
-          styles={{ input: { color: "#9ca3af" } }}
-        />
-      </Form.Item>
+      <MountedFormField name={["env_vars", String(index), "description"]} className="mb-0">
+        {(control) => (
+          <Input
+            {...textControl(control)}
+            addonBefore={
+              <Tooltip title="Per-user variables have no shared value. This text is only a hint shown to each user when they fill in their own value.">
+                <span className="text-xs text-gray-500 cursor-help whitespace-nowrap">
+                  <InfoCircleOutlined className="mr-1" />
+                  Hint
+                </span>
+              </Tooltip>
+            }
+            placeholder="e.g. Your DB username"
+            styles={{ input: { color: "#9ca3af" } }}
+          />
+        )}
+      </MountedFormField>
     );
   }
   return (
-    <Form.Item {...restField} name={[name, "value"]} className="mb-0">
-      <Input placeholder="e.g. postgresql" className="rounded-md font-mono" />
-    </Form.Item>
+    <MountedFormField name={["env_vars", String(index), "value"]} className="mb-0">
+      {(control) => <Input {...textControl(control)} placeholder="e.g. postgresql" className="rounded-md font-mono" />}
+    </MountedFormField>
   );
 };
 

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/IdJagFormFields.tsx

@@ -1,7 +1,11 @@
 import React from "react";
-import { Form, Input, Select, Tooltip } from "antd";
+import { Input, Select, Tooltip } from "antd";
 import { InfoCircleOutlined } from "@ant-design/icons";
 
+import { MountedFormField } from "@/components/common_components/MountedFormField";
+import { antdRequired } from "@/components/common_components/antdFormRules";
+import { requiredUnlessSiblingSet, selectControl, textControl } from "./mcpFieldRules";
+
 interface IdJagFormFieldsProps {
   isEditing?: boolean;
 }
@@ -17,12 +21,16 @@ const FieldLabel: React.FC<{ label: string; tooltip: string }> = ({ label, toolt
   </span>
 );
 
+const PRIVATE_KEY_PATH = ["credentials", "client_private_key"] as const;
+
 const IdJagFormFields: React.FC<IdJagFormFieldsProps> = ({ isEditing = false }) => {
   const placeholderSuffix = isEditing ? " (leave blank to keep existing)" : "";
+  const requiredWhenCreating = (message: string) =>
+    isEditing ? undefined : { validate: { required: antdRequired(message) } };
 
   return (
     <>
-      <Form.Item
+      <MountedFormField
         label={
           <FieldLabel
             label="Org Token Endpoint (leg 1)"
@@ -30,11 +38,18 @@ const IdJagFormFields: React.FC<IdJagFormFieldsProps> = ({ isEditing = false })
           />
         }
         name="token_exchange_endpoint"
-        rules={[{ required: !isEditing, message: "The org token endpoint is required for ID-JAG" }]}
+        required={!isEditing}
+        rules={requiredWhenCreating("The org token endpoint is required for ID-JAG")}
       >
-        <Input placeholder="https://your-org.okta.com/oauth2/v1/token" className={fieldClassName} />
-      </Form.Item>
-      <Form.Item
+        {(control) => (
+          <Input
+            {...textControl(control)}
+            placeholder="https://your-org.okta.com/oauth2/v1/token"
+            className={fieldClassName}
+          />
+        )}
+      </MountedFormField>
+      <MountedFormField
         label={
           <FieldLabel
             label="Resource Token Endpoint (leg 2)"
@@ -42,18 +57,32 @@ const IdJagFormFields: React.FC<IdJagFormFieldsProps> = ({ isEditing = false })
           />
         }
         name={["credentials", "id_jag_resource_token_endpoint"]}
-        rules={[{ required: !isEditing, message: "The resource token endpoint is required for ID-JAG" }]}
+        required={!isEditing}
+        rules={requiredWhenCreating("The resource token endpoint is required for ID-JAG")}
       >
-        <Input placeholder="https://upstream.example.com/oauth2/token" className={fieldClassName} />
-      </Form.Item>
-      <Form.Item
+        {(control) => (
+          <Input
+            {...textControl(control)}
+            placeholder="https://upstream.example.com/oauth2/token"
+            className={fieldClassName}
+          />
+        )}
+      </MountedFormField>
+      <MountedFormField
         label={<FieldLabel label="Client ID" tooltip="OAuth2 client ID LiteLLM authenticates as on both legs." />}
         name={["credentials", "client_id"]}
-        rules={[{ required: !isEditing, message: "Client ID is required for ID-JAG" }]}
+        required={!isEditing}
+        rules={requiredWhenCreating("Client ID is required for ID-JAG")}
       >
-        <Input.Password placeholder={`Enter OAuth client ID${placeholderSuffix}`} className={fieldClassName} />
-      </Form.Item>
-      <Form.Item
+        {(control) => (
+          <Input.Password
+            {...textControl(control)}
+            placeholder={`Enter OAuth client ID${placeholderSuffix}`}
+            className={fieldClassName}
+          />
+        )}
+      </MountedFormField>
+      <MountedFormField
         label={
           <FieldLabel
             label="Client Secret"
@@ -61,36 +90,47 @@ const IdJagFormFields: React.FC<IdJagFormFieldsProps> = ({ isEditing = false })
           />
         }
         name={["credentials", "client_secret"]}
-        dependencies={[["credentials", "client_private_key"]]}
-        rules={[
-          ({ getFieldValue }) => ({
-            validator: (_, value) => {
-              if (isEditing || value || getFieldValue(["credentials", "client_private_key"])) {
-                return Promise.resolve();
+        rules={
+          isEditing
+            ? undefined
+            : {
+                deps: ["credentials.client_private_key"],
+                validate: {
+                  secretOrPrivateKey: requiredUnlessSiblingSet(
+                    PRIVATE_KEY_PATH,
+                    "Provide either a client secret or a client private key",
+                  ),
+                },
               }
-              return Promise.reject(new Error("Provide either a client secret or a client private key"));
-            },
-          }),
-        ]}
+        }
       >
-        <Input.Password placeholder={`Enter OAuth client secret${placeholderSuffix}`} className={fieldClassName} />
-      </Form.Item>
-      <Form.Item
+        {(control) => (
+          <Input.Password
+            {...textControl(control)}
+            placeholder={`Enter OAuth client secret${placeholderSuffix}`}
+            className={fieldClassName}
+          />
+        )}
+      </MountedFormField>
+      <MountedFormField
         label={
           <FieldLabel
             label="Client Private Key (PEM)"
             tooltip="PEM private key signing the RFC 7523 private_key_jwt client assertion. Okta Cross App Access normally requires this. When set it takes precedence over the client secret."
           />
         }
-        name={["credentials", "client_private_key"]}
+        name={PRIVATE_KEY_PATH}
       >
-        <Input.TextArea
-          rows={3}
-          placeholder={`-----BEGIN PRIVATE KEY-----${placeholderSuffix}`}
-          className={fieldClassName}
-        />
-      </Form.Item>
-      <Form.Item
+        {(control) => (
+          <Input.TextArea
+            {...textControl(control)}
+            rows={3}
+            placeholder={`-----BEGIN PRIVATE KEY-----${placeholderSuffix}`}
+            className={fieldClassName}
+          />
+        )}
+      </MountedFormField>
+      <MountedFormField
         label={
           <FieldLabel
             label="Private Key ID (optional)"
@@ -99,9 +139,9 @@ const IdJagFormFields: React.FC<IdJagFormFieldsProps> = ({ isEditing = false })
         }
         name={["credentials", "client_private_key_id"]}
       >
-        <Input placeholder="my-signing-key-1" className={fieldClassName} />
-      </Form.Item>
-      <Form.Item
+        {(control) => <Input {...textControl(control)} placeholder="my-signing-key-1" className={fieldClassName} />}
+      </MountedFormField>
+      <MountedFormField
         label={
           <FieldLabel
             label="Client Assertion Signing Algorithm (optional)"
@@ -110,9 +150,9 @@ const IdJagFormFields: React.FC<IdJagFormFieldsProps> = ({ isEditing = false })
         }
         name={["credentials", "client_assertion_signing_alg"]}
       >
-        <Input placeholder="RS256" className={fieldClassName} />
-      </Form.Item>
-      <Form.Item
+        {(control) => <Input {...textControl(control)} placeholder="RS256" className={fieldClassName} />}
+      </MountedFormField>
+      <MountedFormField
         label={
           <FieldLabel
             label="Audience (optional)"
@@ -121,9 +161,11 @@ const IdJagFormFields: React.FC<IdJagFormFieldsProps> = ({ isEditing = false })
         }
         name="audience"
       >
-        <Input placeholder="https://upstream.example.com" className={fieldClassName} />
-      </Form.Item>
-      <Form.Item
+        {(control) => (
+          <Input {...textControl(control)} placeholder="https://upstream.example.com" className={fieldClassName} />
+        )}
+      </MountedFormField>
+      <MountedFormField
         label={
           <FieldLabel
             label="Resource Indicator (optional)"
@@ -132,9 +174,11 @@ const IdJagFormFields: React.FC<IdJagFormFieldsProps> = ({ isEditing = false })
         }
         name={["credentials", "id_jag_resource"]}
       >
-        <Input placeholder="https://upstream.example.com/mcp" className={fieldClassName} />
-      </Form.Item>
-      <Form.Item
+        {(control) => (
+          <Input {...textControl(control)} placeholder="https://upstream.example.com/mcp" className={fieldClassName} />
+        )}
+      </MountedFormField>
+      <MountedFormField
         label={
           <FieldLabel
             label="Subject Token Type (optional)"
@@ -143,14 +187,29 @@ const IdJagFormFields: React.FC<IdJagFormFieldsProps> = ({ isEditing = false })
         }
         name="subject_token_type"
       >
-        <Input placeholder="urn:ietf:params:oauth:token-type:id_token" className={fieldClassName} />
-      </Form.Item>
-      <Form.Item
+        {(control) => (
+          <Input
+            {...textControl(control)}
+            placeholder="urn:ietf:params:oauth:token-type:id_token"
+            className={fieldClassName}
+          />
+        )}
+      </MountedFormField>
+      <MountedFormField
         label={<FieldLabel label="Scopes (optional)" tooltip="Scopes requested on leg 1 of the exchange." />}
         name={["credentials", "scopes"]}
       >
-        <Select mode="tags" tokenSeparators={[","]} placeholder="Add scopes" className="rounded-lg" size="large" />
-      </Form.Item>
+        {(control) => (
+          <Select
+            {...selectControl(control)}
+            mode="tags"
+            tokenSeparators={[","]}
+            placeholder="Add scopes"
+            className="rounded-lg"
+            size="large"
+          />
+        )}
+      </MountedFormField>
     </>
   );
 };

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.test.tsx

@@ -1,10 +1,10 @@
 import React from "react";
-import { render, screen } from "@testing-library/react";
+import { screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { describe, it, expect } from "vitest";
-import { Form } from "antd";
 
 import MCPPermissionManagement from "./MCPPermissionManagement";
+import { renderInMcpForm } from "./McpFormTestHarness";
 
 const defaultProps = {
   availableAccessGroups: [],
@@ -12,6 +12,7 @@ const defaultProps = {
   searchValue: "",
   setSearchValue: () => {},
   getAccessGroupOptions: () => [],
+  mountedAuthType: undefined,
 };
 
 describe("MCPPermissionManagement", () => {
@@ -24,22 +25,8 @@ describe("MCPPermissionManagement", () => {
     return user;
   };
 
-  const renderWithForm = (props = {}) => {
-    const Wrapper: React.FC = ({ children }) => {
-      const [form] = Form.useForm();
-      return (
-        <Form form={form} initialValues={{ allow_all_keys: false }}>
-          {children}
-        </Form>
-      );
-    };
-
-    return render(
-      <Wrapper>
-        <MCPPermissionManagement {...defaultProps} {...props} />
-      </Wrapper>,
-    );
-  };
+  const renderWithForm = (props = {}) =>
+    renderInMcpForm(<MCPPermissionManagement {...defaultProps} {...props} />, { allow_all_keys: false });
 
   it("should default allow_all_keys switch to unchecked for new servers", async () => {
     renderWithForm();
@@ -51,27 +38,15 @@ describe("MCPPermissionManagement", () => {
     expect(toggle).not.toBeChecked();
   });
 
-  const renderWithInitialValues = (initialValues: Record<string, unknown>, props = {}) => {
-    const Wrapper: React.FC = ({ children }) => {
-      const [form] = Form.useForm();
-      return (
-        <Form form={form} initialValues={initialValues}>
-          {/* In the real app auth_type is registered by the parent form; the
-              component only watches it. Register a hidden field here so
-              Form.useWatch("auth_type") resolves the initial value. */}
-          <Form.Item name="auth_type" hidden>
-            <input />
-          </Form.Item>
-          {children}
-        </Form>
-      );
-    };
-    return render(
-      <Wrapper>
-        <MCPPermissionManagement {...defaultProps} {...props} />
-      </Wrapper>,
+  const renderWithInitialValues = (initialValues: Record<string, unknown>, props = {}) =>
+    renderInMcpForm(
+      <MCPPermissionManagement
+        {...defaultProps}
+        mountedAuthType={initialValues.auth_type as string | undefined}
+        {...props}
+      />,
+      initialValues,
     );
-  };
 
   it("shows only the oauth2 PKCE-delegation toggle for oauth2 servers", async () => {
     renderWithInitialValues({ allow_all_keys: false, auth_type: "oauth2" });

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPPermissionManagement.tsx

@@ -1,7 +1,17 @@
 import React, { useEffect } from "react";
-import { Alert, Form, Select, Tooltip, Collapse, Input, Space, Button, Switch } from "antd";
+import { Alert, Select, Tooltip, Collapse, Input, Space, Button, Switch } from "antd";
 import { InfoCircleOutlined, MinusCircleOutlined, PlusOutlined } from "@ant-design/icons";
+import { useFieldArray, useFormContext, useWatch } from "react-hook-form";
 import { MCPServer, AUTH_TYPE } from "@/components/mcp_tools/types";
+import {
+  MountedFormField,
+  useMountedName,
+  type MountedFormValues,
+} from "@/components/common_components/MountedFormField";
+import { antdRequired } from "@/components/common_components/antdFormRules";
+import { Field, FieldLabel } from "@/components/shared/form/field";
+import { invertedSwitchControl, selectControl, switchControl, textControl } from "./mcpFieldRules";
+import { listControl } from "./mcpFormStore";
 const { Panel } = Collapse;
 
 interface MCPPermissionManagementProps {
@@ -13,20 +23,79 @@ interface MCPPermissionManagementProps {
     value: string;
     label: React.ReactNode;
   }>;
+  /**
+   * The auth type as seen through the gate that mounts the auth_type field.
+   * Callers pass undefined whenever that field is unmounted, because both
+   * toggles below are mounted from this value and the payload only carries
+   * what is mounted.
+   */
+  mountedAuthType: string | null | undefined;
 }
 
+const StaticHeadersFieldArray: React.FC = () => {
+  const { control } = useFormContext<MountedFormValues>();
+  const { fields, append, remove } = useFieldArray({ control: listControl(control), name: "static_headers" });
+  useMountedName("static_headers");
+
+  return (
+    <div className="space-y-3">
+      {fields.map((item, index) => (
+        <Space key={item.id} className="flex w-full" align="baseline" size="middle">
+          <MountedFormField
+            name={["static_headers", String(index), "header"]}
+            className="flex-1"
+            rules={{ validate: { required: antdRequired("Header name is required") } }}
+          >
+            {(headerControl) => (
+              <Input
+                {...textControl(headerControl)}
+                size="large"
+                allowClear
+                className="rounded-lg"
+                placeholder="Header name (e.g., X-API-Key)"
+              />
+            )}
+          </MountedFormField>
+          <MountedFormField
+            name={["static_headers", String(index), "value"]}
+            className="flex-1"
+            rules={{ validate: { required: antdRequired("Header value is required") } }}
+          >
+            {(valueControl) => (
+              <Input
+                {...textControl(valueControl)}
+                size="large"
+                allowClear
+                className="rounded-lg"
+                placeholder="Header value"
+              />
+            )}
+          </MountedFormField>
+          <MinusCircleOutlined
+            onClick={() => remove(index)}
+            className="text-gray-500 hover:text-red-500 cursor-pointer"
+          />
+        </Space>
+      ))}
+      <Button type="dashed" onClick={() => append({})} icon={<PlusOutlined />} block>
+        Add Static Header
+      </Button>
+    </div>
+  );
+};
+
 const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
   availableAccessGroups,
   mcpServer,
   searchValue,
   setSearchValue,
   getAccessGroupOptions,
+  mountedAuthType,
 }) => {
-  const form = Form.useFormInstance();
-  const watchedAuthType = Form.useWatch("auth_type", form);
-  const isOAuth2 = watchedAuthType === AUTH_TYPE.OAUTH2;
-  const isNoneAuth = watchedAuthType === AUTH_TYPE.NONE || watchedAuthType == null;
-  const watchedExtraHeaders = Form.useWatch("extra_headers", form);
+  const { setValue } = useFormContext<MountedFormValues>();
+  const isOAuth2 = mountedAuthType === AUTH_TYPE.OAUTH2;
+  const isNoneAuth = mountedAuthType === AUTH_TYPE.NONE || mountedAuthType == null;
+  const watchedExtraHeaders = useWatch({ name: "extra_headers" });
   const hasAuthorizationHeader =
     Array.isArray(watchedExtraHeaders) &&
     watchedExtraHeaders.some((h) => typeof h === "string" && h.toLowerCase() === "authorization");
@@ -39,8 +108,8 @@ const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
   // Kept as separate flags so neither silently implies the other and existing
   // oauth2 servers can't regress into pass-through behavior.
   const canEnableOAuthPassthrough = isNoneAuth && hasAuthorizationHeader;
-  const watchedDelegateAuth = Form.useWatch("delegate_auth_to_upstream", form);
-  const watchedPublicInternet = Form.useWatch("available_on_public_internet", form);
+  const watchedDelegateAuth = useWatch({ name: "delegate_auth_to_upstream" });
+  const watchedPublicInternet = useWatch({ name: "available_on_public_internet" });
   const showInternalDelegatePkceWarning = isOAuth2 && watchedDelegateAuth === true && watchedPublicInternet === false;
 
   // Set initial values when mcpServer changes
@@ -51,10 +120,10 @@ const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
           header,
           value: value != null ? String(value) : "",
         }));
-        form.setFieldValue("static_headers", staticHeaders);
+        setValue("static_headers", staticHeaders);
       }
       if (Array.isArray(mcpServer.env_vars) && mcpServer.env_vars.length > 0) {
-        form.setFieldValue(
+        setValue(
           "env_vars",
           mcpServer.env_vars.map((entry) => ({
             name: entry.name,
@@ -65,41 +134,41 @@ const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
         );
       }
       if (typeof mcpServer.allow_all_keys === "boolean") {
-        form.setFieldValue("allow_all_keys", mcpServer.allow_all_keys);
+        setValue("allow_all_keys", mcpServer.allow_all_keys);
       }
       if (typeof mcpServer.available_on_public_internet === "boolean") {
-        form.setFieldValue("available_on_public_internet", mcpServer.available_on_public_internet);
+        setValue("available_on_public_internet", mcpServer.available_on_public_internet);
       }
       if (typeof mcpServer.delegate_auth_to_upstream === "boolean") {
-        form.setFieldValue("delegate_auth_to_upstream", mcpServer.delegate_auth_to_upstream);
+        setValue("delegate_auth_to_upstream", mcpServer.delegate_auth_to_upstream);
       }
       if (typeof mcpServer.oauth_passthrough === "boolean") {
-        form.setFieldValue("oauth_passthrough", mcpServer.oauth_passthrough);
+        setValue("oauth_passthrough", mcpServer.oauth_passthrough);
       }
     } else {
-      form.setFieldValue("allow_all_keys", false);
-      form.setFieldValue("available_on_public_internet", true);
-      form.setFieldValue("delegate_auth_to_upstream", false);
-      form.setFieldValue("oauth_passthrough", false);
+      setValue("allow_all_keys", false);
+      setValue("available_on_public_internet", true);
+      setValue("delegate_auth_to_upstream", false);
+      setValue("oauth_passthrough", false);
     }
-  }, [mcpServer, form]);
+  }, [mcpServer, setValue]);
 
   // delegate_auth_to_upstream is only honored server-side for oauth2 servers.
   // Force it back to false whenever the user switches away from oauth2 so a
   // stale toggle value doesn't get persisted unexpectedly.
   useEffect(() => {
     if (!isOAuth2) {
-      form.setFieldValue("delegate_auth_to_upstream", false);
+      setValue("delegate_auth_to_upstream", false);
     }
-  }, [isOAuth2, form]);
+  }, [isOAuth2, setValue]);
 
   // oauth_passthrough is only honored for auth_type=none servers that forward
   // Authorization upstream. Force it back to false otherwise.
   useEffect(() => {
     if (!canEnableOAuthPassthrough) {
-      form.setFieldValue("oauth_passthrough", false);
+      setValue("oauth_passthrough", false);
     }
-  }, [canEnableOAuthPassthrough, form]);
+  }, [canEnableOAuthPassthrough, setValue]);
 
   return (
     <Collapse className="bg-gray-50 border border-gray-200 rounded-lg" expandIconPosition="end" ghost={false}>
@@ -130,14 +199,9 @@ const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
                 Enable if this server should be &quot;public&quot; to all keys.
               </p>
             </div>
-            <Form.Item
-              name="allow_all_keys"
-              valuePropName="checked"
-              initialValue={mcpServer?.allow_all_keys ?? false}
-              className="mb-0"
-            >
-              <Switch />
-            </Form.Item>
+            <MountedFormField name="allow_all_keys" defaultValue={mcpServer?.allow_all_keys ?? false} className="mb-0">
+              {(control) => <Switch {...switchControl(control)} />}
+            </MountedFormField>
           </div>
 
           <div className="flex items-start justify-between gap-4">
@@ -152,16 +216,9 @@ const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
                 Turn on to restrict access to callers within your internal network only.
               </p>
             </div>
-            <Form.Item
-              name="available_on_public_internet"
-              valuePropName="checked"
-              getValueProps={(value) => ({ checked: !value })}
-              getValueFromEvent={(checked: boolean) => !checked}
-              initialValue={true}
-              className="mb-0"
-            >
-              <Switch />
-            </Form.Item>
+            <MountedFormField name="available_on_public_internet" defaultValue={true} className="mb-0">
+              {(control) => <Switch {...invertedSwitchControl(control)} />}
+            </MountedFormField>
           </div>
 
           {isOAuth2 && (
@@ -177,14 +234,13 @@ const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
                   Bypass LiteLLM auth so clients authenticate directly with the upstream OAuth MCP server.
                 </p>
               </div>
-              <Form.Item
+              <MountedFormField
                 name="delegate_auth_to_upstream"
-                valuePropName="checked"
-                initialValue={mcpServer?.delegate_auth_to_upstream ?? false}
+                defaultValue={mcpServer?.delegate_auth_to_upstream ?? false}
                 className="mb-0"
               >
-                <Switch />
-              </Form.Item>
+                {(control) => <Switch {...switchControl(control)} />}
+              </MountedFormField>
             </div>
           )}
 
@@ -202,14 +258,13 @@ const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
                   upstream MCP server.
                 </p>
               </div>
-              <Form.Item
+              <MountedFormField
                 name="oauth_passthrough"
-                valuePropName="checked"
-                initialValue={mcpServer?.oauth_passthrough ?? false}
+                defaultValue={mcpServer?.oauth_passthrough ?? false}
                 className="mb-0"
               >
-                <Switch />
-              </Form.Item>
+                {(control) => <Switch {...switchControl(control)} />}
+              </MountedFormField>
             </div>
           )}
 
@@ -223,7 +278,7 @@ const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
             />
           )}
 
-          <Form.Item
+          <MountedFormField
             label={
               <span className="text-sm font-medium text-gray-700 flex items-center">
                 MCP Access Groups
@@ -235,21 +290,24 @@ const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
             name="mcp_access_groups"
             className="mb-4"
           >
-            <Select
-              mode="tags"
-              showSearch
-              placeholder="Select existing groups or type to create new ones"
-              optionFilterProp="value"
-              filterOption={(input, option) => (option?.value ?? "").toLowerCase().includes(input.toLowerCase())}
-              onSearch={(value) => setSearchValue(value)}
-              tokenSeparators={[","]}
-              options={getAccessGroupOptions()}
-              maxTagCount="responsive"
-              allowClear
-            />
-          </Form.Item>
+            {(control) => (
+              <Select
+                {...selectControl(control)}
+                mode="tags"
+                showSearch
+                placeholder="Select existing groups or type to create new ones"
+                optionFilterProp="value"
+                filterOption={(input, option) => (option?.value ?? "").toLowerCase().includes(input.toLowerCase())}
+                onSearch={(value) => setSearchValue(value)}
+                tokenSeparators={[","]}
+                options={getAccessGroupOptions()}
+                maxTagCount="responsive"
+                allowClear
+              />
+            )}
+          </MountedFormField>
 
-          <Form.Item
+          <MountedFormField
             label={
               <span className="text-sm font-medium text-gray-700 flex items-center">
                 Extra Headers
@@ -265,70 +323,34 @@ const MCPPermissionManagement: React.FC<MCPPermissionManagementProps> = ({
             }
             name="extra_headers"
           >
-            <Select
-              mode="tags"
-              placeholder={
-                mcpServer?.extra_headers && mcpServer.extra_headers.length > 0
-                  ? `Currently: ${mcpServer.extra_headers.join(", ")}`
-                  : "Enter header names (e.g., Authorization, X-Custom-Header)"
-              }
-              className="rounded-lg"
-              size="large"
-              tokenSeparators={[","]}
-              allowClear
-            />
-          </Form.Item>
+            {(control) => (
+              <Select
+                {...selectControl(control)}
+                mode="tags"
+                placeholder={
+                  mcpServer?.extra_headers && mcpServer.extra_headers.length > 0
+                    ? `Currently: ${mcpServer.extra_headers.join(", ")}`
+                    : "Enter header names (e.g., Authorization, X-Custom-Header)"
+                }
+                className="rounded-lg"
+                size="large"
+                tokenSeparators={[","]}
+                allowClear
+              />
+            )}
+          </MountedFormField>
 
-          <Form.Item
-            label={
+          <Field>
+            <FieldLabel>
               <span className="text-sm font-medium text-gray-700 flex items-center">
                 Static Headers
                 <Tooltip title="Send these key-value headers with every request to this MCP server.">
                   <InfoCircleOutlined className="ml-2 text-blue-400 hover:text-blue-600 cursor-help" />
                 </Tooltip>
               </span>
-            }
-            required={false}
-          >
-            <Form.List name="static_headers">
-              {(fields, { add, remove }) => (
-                <div className="space-y-3">
-                  {fields.map(({ key, name, ...restField }) => (
-                    <Space key={key} className="flex w-full" align="baseline" size="middle">
-                      <Form.Item
-                        {...restField}
-                        name={[name, "header"]}
-                        className="flex-1"
-                        rules={[{ required: true, message: "Header name is required" }]}
-                      >
-                        <Input
-                          size="large"
-                          allowClear
-                          className="rounded-lg"
-                          placeholder="Header name (e.g., X-API-Key)"
-                        />
-                      </Form.Item>
-                      <Form.Item
-                        {...restField}
-                        name={[name, "value"]}
-                        className="flex-1"
-                        rules={[{ required: true, message: "Header value is required" }]}
-                      >
-                        <Input size="large" allowClear className="rounded-lg" placeholder="Header value" />
-                      </Form.Item>
-                      <MinusCircleOutlined
-                        onClick={() => remove(name)}
-                        className="text-gray-500 hover:text-red-500 cursor-pointer"
-                      />
-                    </Space>
-                  ))}
-                  <Button type="dashed" onClick={() => add()} icon={<PlusOutlined />} block>
-                    Add Static Header
-                  </Button>
-                </div>
-              )}
-            </Form.List>
-          </Form.Item>
+            </FieldLabel>
+            <StaticHeadersFieldArray />
+          </Field>
         </div>
       </Panel>
     </Collapse>

ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/MCPToolsetsTab.integration.test.tsx

@@ -1,5 +1,5 @@
 import React from "react";
-import { render, screen, waitFor, within } from "@testing-library/react";
+import { fireEvent, render, screen, waitFor, within } from "@testing-library/react";
 import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event";
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
@@ -68,8 +68,10 @@ describe("MCPToolsetsTab create/edit toolset form", () => {
     renderTab();
 
     const dialog = await openCreate(user);
-    await user.type(dialog.getByPlaceholderText("e.g. github-linear-tools"), "github-linear-tools");
-    await user.type(dialog.getByPlaceholderText("Optional description"), "tools for triage");
+    fireEvent.change(dialog.getByPlaceholderText("e.g. github-linear-tools"), {
+      target: { value: "github-linear-tools" },
+    });
+    fireEvent.change(dialog.getByPlaceholderText("Optional description"), { target: { value: "tools for triage" } });
     await user.click(dialog.getByRole("button", { name: "Create Toolset" }));
 
     await waitFor(() => {
@@ -90,7 +92,7 @@ describe("MCPToolsetsTab create/edit toolset form", () => {
     renderTab();
 
     const dialog = await openCreate(user);
-    await user.type(dialog.getByPlaceholderText("e.g. github-linear-tools"), "solo");
+    fireEvent.change(dialog.getByPlaceholderText("e.g. github-linear-tools"), { target: { value: "solo" } });
     await user.click(dialog.getByRole("button", { name: "Create Toolset" }));
 
     await waitFor(() => {
@@ -121,8 +123,8 @@ describe("MCPToolsetsTab create/edit toolset form", () => {
     renderTab();
 
     const dialog = await openCreate(user);
-    await user.type(dialog.getByPlaceholderText("e.g. github-linear-tools"), "spaced");
-    await user.type(dialog.getByPlaceholderText("Optional description"), "  ");
+    fireEvent.change(dialog.getByPlaceholderText("e.g. github-linear-tools"), { target: { value: "spaced" } });
+    fireEvent.change(dialog.getByPlaceholderText("Optional description"), { target: { value: "  " } });
     await user.click(dialog.getByRole("button", { name: "Create Toolset" }));
 
     await waitFor(() => {
@@ -154,7 +156,7 @@ describe("MCPToolsetsTab create/edit toolset form", () => {
     expect(dialog.getByPlaceholderText("Optional description")).toHaveValue("old description");
 
     await user.clear(name);
-    await user.type(name, "renamed");
+    fireEvent.change(name, { target: { value: "renamed" } });
     await user.click(dialog.getByRole("button", { name: "Save Changes" }));
 
     const expectedUpdate = {