From 158b0c28c02556a80bf5992ad750aacb9d40ba6f Mon Sep 17 00:00:00 2001 From: "oss-pr-review-agent-shin[bot]" <281797381+oss-pr-review-agent-shin[bot]@users.noreply.github.com> Date: Thu, 7 May 2026 21:29:47 +0000 Subject: [PATCH 1/3] =?UTF-8?q?[litellm-agent]=20Staging=20=E2=86=92=20lit?= =?UTF-8?q?ellm=5Finternal=5Fstaging=20(5/7/2026)=20(#27375)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Squash-merged by litellm-agent from oss-pr-review-agent-shin[bot]'s PR. --- .../mcp_server/discoverable_endpoints.py | 79 ++---- .../_experimental/mcp_server/oauth_utils.py | 109 ++++++++- litellm/proxy/auth/handle_jwt.py | 30 +++ .../proxy/guardrails/guardrail_endpoints.py | 20 +- .../proxy/guardrails/guardrail_registry.py | 65 ++++- litellm/proxy/guardrails/init_guardrails.py | 1 + .../health_endpoints/_health_endpoints.py | 67 +++-- .../tag_management_endpoints.py | 51 +++- litellm/proxy/management_endpoints/ui_sso.py | 4 +- litellm/proxy/proxy_server.py | 10 + .../mcp_server/test_byok_oauth_endpoints.py | 66 +++++ .../mcp_server/test_discoverable_endpoints.py | 69 +++++- .../proxy/auth/test_handle_jwt.py | 71 +++++- .../guardrails/test_guardrail_endpoints.py | 64 +++++ .../guardrails/test_guardrail_registry.py | 120 +++++++++ .../health_endpoints/test_health_endpoints.py | 230 ++++++++++++++++++ .../test_tag_management_endpoints.py | 111 +++++++++ .../UsagePage/components/UsagePageView.tsx | 41 ++-- .../src/components/model_info_view.test.tsx | 27 ++ .../src/components/model_info_view.tsx | 6 + .../src/components/networking.tsx | 21 +- .../components/templates/key_info_view.tsx | 4 +- 22 files changed, 1145 insertions(+), 121 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 1794cd14381..62691641234 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -12,7 +12,8 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._experimental.mcp_server.oauth_utils import ( TOKEN_NO_CACHE_HEADERS, - validate_loopback_redirect_uri, + get_request_base_url, + validate_trusted_redirect_uri, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -29,51 +30,6 @@ router = APIRouter( ) -def get_request_base_url(request: Request) -> str: - """ - Get the base URL for the request, considering X-Forwarded-* headers. - - X-Forwarded-Proto / X-Forwarded-Host / X-Forwarded-Port are only honoured - when the request comes from a configured trusted proxy - (``use_x_forwarded_for`` enabled AND caller in ``mcp_trusted_proxy_ranges``). - Otherwise the request's literal ``base_url`` is returned, so an - untrusted caller cannot poison OAuth-discovery / redirect_uri values - by injecting headers. - - Args: - request: FastAPI Request object - - Returns: - The reconstructed base URL (e.g., "https://proxy.example.com") - """ - base_url = str(request.base_url).rstrip("/") - parsed = urlparse(base_url) - - if not IPAddressUtils.is_request_from_trusted_proxy(request): - return base_url - - x_forwarded_proto = request.headers.get("X-Forwarded-Proto") - x_forwarded_host = request.headers.get("X-Forwarded-Host") - x_forwarded_port = request.headers.get("X-Forwarded-Port") - - scheme = x_forwarded_proto if x_forwarded_proto else parsed.scheme - - if x_forwarded_host: - # X-Forwarded-Host may already include port (e.g., "example.com:8080") - if ":" in x_forwarded_host and not x_forwarded_host.startswith("["): - netloc = x_forwarded_host - elif x_forwarded_port: - netloc = f"{x_forwarded_host}:{x_forwarded_port}" - else: - netloc = x_forwarded_host - else: - netloc = parsed.netloc - if x_forwarded_port and ":" not in netloc: - netloc = f"{netloc}:{x_forwarded_port}" - - return urlunparse((scheme, netloc, parsed.path, "", "", "")) - - def encode_state_with_base_url( base_url: str, original_state: str, @@ -127,12 +83,14 @@ def decode_state_hash(encrypted_state: str) -> dict: return state_data -def _get_validated_client_redirect_uri(state_data: Dict[str, Any]) -> str: - """Return a loopback client redirect URI from OAuth state.""" +def _get_validated_client_redirect_uri( + request: Request, state_data: Dict[str, Any] +) -> str: + """Return a trusted (same-origin or loopback) client redirect URI from OAuth state.""" redirect_uri = state_data.get("client_redirect_uri") or state_data.get("base_url") if not redirect_uri or not isinstance(redirect_uri, str): raise HTTPException(status_code=400, detail="Invalid redirect URI") - validate_loopback_redirect_uri(redirect_uri) + validate_trusted_redirect_uri(request, redirect_uri) return redirect_uri @@ -338,12 +296,12 @@ async def authorize_with_server( status_code=400, detail="MCP server authorization url is not set" ) - # Loopback-only redirect_uri. The URI is encrypted into the OAuth - # state and decoded on /callback to redirect the user back; a non- - # loopback URI would be an open-redirect + code-theft primitive - # (VERIA-57 root cause B). MCP clients are native apps — loopback is - # the spec-compliant callback pattern. - validate_loopback_redirect_uri(redirect_uri) + # Loopback OR same-origin redirect_uri. The URI is encrypted into the + # OAuth state and decoded on /callback to redirect the user back; + # restricting to trusted origins blocks the open-redirect + + # code-theft primitive (VERIA-57 root cause B). Loopback supports + # native MCP clients; same-origin supports the proxy's own UI callback. + validate_trusted_redirect_uri(request, redirect_uri) parsed = urlparse(redirect_uri) base_url = urlunparse(parsed._replace(query="")) request_base_url = get_request_base_url(request) @@ -660,17 +618,18 @@ async def token_endpoint( @router.get("/callback") -async def callback(code: str, state: str): +async def callback(request: Request, code: str, state: str): try: state_data = decode_state_hash(state) original_state = state_data["original_state"] - # Re-validate loopback at the sink. /authorize rejects non-loopback + # Re-validate at the sink. /authorize rejects untrusted # redirect_uri before encoding into state, but encrypted states # minted before that check was added have no expiry and remain - # valid indefinitely. Validating here blocks the open-redirect + - # code-theft primitive even for pre-fix states. - redirect_uri = _get_validated_client_redirect_uri(state_data) + # valid indefinitely. Validating here (same-origin OR loopback) + # blocks the open-redirect + code-theft primitive even for pre-fix + # states while allowing the UI's same-origin callback to work. + redirect_uri = _get_validated_client_redirect_uri(request, state_data) params = {"code": code, "state": original_state} complete_returned_url = _append_query_params(redirect_uri, params) diff --git a/litellm/proxy/_experimental/mcp_server/oauth_utils.py b/litellm/proxy/_experimental/mcp_server/oauth_utils.py index b13cf83058c..343d1bee613 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth_utils.py +++ b/litellm/proxy/_experimental/mcp_server/oauth_utils.py @@ -2,15 +2,63 @@ (BYOK + discoverable / pass-through OAuth proxy).""" from ipaddress import ip_address -from urllib.parse import urlparse +from urllib.parse import urlparse, urlunparse -from fastapi import HTTPException +from fastapi import HTTPException, Request + +from litellm._logging import verbose_logger +from litellm.proxy.auth.ip_address_utils import IPAddressUtils # RFC 6749 §5.1 / OAuth 2.1 draft-15 §4.1.3: token-endpoint responses # must not be cached — both success and error bodies may reveal secrets. TOKEN_NO_CACHE_HEADERS = {"Cache-Control": "no-store", "Pragma": "no-cache"} +def get_request_base_url(request: Request) -> str: + """ + Get the base URL for the request, considering X-Forwarded-* headers. + + X-Forwarded-Proto / X-Forwarded-Host / X-Forwarded-Port are only honoured + when the request comes from a configured trusted proxy + (``use_x_forwarded_for`` enabled AND caller in ``mcp_trusted_proxy_ranges``). + Otherwise the request's literal ``base_url`` is returned, so an + untrusted caller cannot poison OAuth-discovery / redirect_uri values + by injecting headers. + + Args: + request: FastAPI Request object + + Returns: + The reconstructed base URL (e.g., "https://proxy.example.com") + """ + base_url = str(request.base_url).rstrip("/") + parsed = urlparse(base_url) + + if not IPAddressUtils.is_request_from_trusted_proxy(request): + return base_url + + x_forwarded_proto = request.headers.get("X-Forwarded-Proto") + x_forwarded_host = request.headers.get("X-Forwarded-Host") + x_forwarded_port = request.headers.get("X-Forwarded-Port") + + scheme = x_forwarded_proto if x_forwarded_proto else parsed.scheme + + if x_forwarded_host: + # X-Forwarded-Host may already include port (e.g., "example.com:8080") + if ":" in x_forwarded_host and not x_forwarded_host.startswith("["): + netloc = x_forwarded_host + elif x_forwarded_port: + netloc = f"{x_forwarded_host}:{x_forwarded_port}" + else: + netloc = x_forwarded_host + else: + netloc = parsed.netloc + if x_forwarded_port and ":" not in netloc: + netloc = f"{netloc}:{x_forwarded_port}" + + return urlunparse((scheme, netloc, parsed.path, "", "", "")) + + def validate_loopback_redirect_uri(redirect_uri: str) -> None: """Require a loopback ``redirect_uri`` (OAuth 2.1 §4.1.2.1 + RFC 8252 §7.3 native-app pattern). MCP clients are native apps that listen on @@ -46,3 +94,60 @@ def validate_loopback_redirect_uri(redirect_uri: str) -> None: # don't let it bubble up as a 500. pass raise HTTPException(status_code=400, detail="invalid_request") + + +def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None: + """Accept same-origin (proxy's own origin) OR loopback ``redirect_uri``. + + Same-origin is required for the LiteLLM UI's OAuth flow: the UI + redirects to ``/ui/mcp/oauth/callback`` which is not loopback + but is on the proxy's own trusted HTTPS origin. An attacker cannot + host content on the proxy's own origin without already owning the + proxy, so the open-redirect / code-theft primitive that motivated + :func:`validate_loopback_redirect_uri` does not apply here. + + Loopback continues to be accepted for native MCP clients (per + OAuth 2.1 §4.1.2.1 + RFC 8252 §7.3). + + Use this in the discoverable OAuth proxy endpoints that serve both + native clients and the proxy's own UI. BYOK endpoints that only + support native clients should keep + :func:`validate_loopback_redirect_uri`. + """ + try: + parsed = urlparse(redirect_uri) + except ValueError: + raise HTTPException(status_code=400, detail="invalid_request") + if parsed.scheme not in ("http", "https"): + raise HTTPException(status_code=400, detail="invalid_request") + if parsed.fragment: + raise HTTPException(status_code=400, detail="invalid_request") + + # Same-origin: scheme + netloc (host[:port]) must match the proxy's + # own base URL at this request (honouring trusted X-Forwarded-*). + try: + proxy_base = urlparse(get_request_base_url(request)) + if ( + parsed.netloc + and parsed.scheme == proxy_base.scheme + and parsed.netloc.lower() == proxy_base.netloc.lower() + ): + return + except Exception as exc: + # If we can't determine the proxy's origin, fall through to + # loopback. Log so the failure is diagnosable in production. + verbose_logger.warning( + "validate_trusted_redirect_uri: could not determine proxy origin, " + "falling back to loopback-only check. error=%s", + exc, + ) + + host = (parsed.hostname or "").lower() + if host == "localhost": + return + try: + if ip_address(host).is_loopback: + return + except ValueError: + pass + raise HTTPException(status_code=400, detail="invalid_request") diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index d1fd5818f35..a270031c33c 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -224,6 +224,36 @@ class JWTHandler: return [] + def get_all_jwt_team_ids(self, token: dict) -> List[str]: + """ + Return team IDs from both the plural ``team_ids_jwt_field`` and the + singular ``team_id_jwt_field`` claim, as a deduplicated list preserving + plural-first order. + + Membership-reconciliation paths (SSO callback, JWT-bearer sync) need + to consider both claim shapes. Reading only the plural field — as + callers historically did — silently dropped users whose IdP populates + the singular field, which is what Okta and Auth0 default to when a + user has a single primary team. + + This intentionally does NOT consult ``team_id_default``: that fallback + is a property of how the JWT-bearer auth flow resolves a single + request-bound team, not of the token's claims. Callers that want the + default-team behavior should still go through ``get_team_id``. + """ + team_ids: List[str] = list(self.get_team_ids_from_jwt(token)) + if self.litellm_jwtauth.team_id_jwt_field is not None: + singular = get_nested_value( + data=token, + key_path=self.litellm_jwtauth.team_id_jwt_field, + default=None, + ) + if isinstance(singular, list): + singular = singular[0] if singular else None + if singular and singular not in team_ids: + team_ids.append(singular) + return team_ids + def get_end_user_id( self, token: dict, default_value: Optional[str] ) -> Optional[str]: diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 5351391e5e1..e55f3b6e16b 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -242,6 +242,10 @@ async def list_guardrails_v2( gid = guardrail.get("guardrail_id") if gid in seen_guardrail_ids: continue + # Skip stale DB-backed entries — the DB row was deleted (likely by + # another pod) and reconciliation hasn't fired yet on this pod. + if gid is not None and IN_MEMORY_GUARDRAIL_HANDLER.get_source(gid) == "db": + continue if not is_admin: g_team_id = guardrail.get("team_id") if g_team_id is not None and g_team_id not in caller_team_ids: @@ -360,7 +364,7 @@ async def create_guardrail( try: IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail( - guardrail=cast(Guardrail, result) + guardrail=cast(Guardrail, result), source="db" ) verbose_proxy_logger.info( f"Immediate sync: Successfully initialized guardrail '{guardrail_name}' (ID: {guardrail_id})" @@ -1017,7 +1021,7 @@ async def approve_guardrail_submission( } try: IN_MEMORY_GUARDRAIL_HANDLER.initialize_guardrail( - guardrail=cast(Guardrail, guardrail_dict) + guardrail=cast(Guardrail, guardrail_dict), source="db" ) verbose_proxy_logger.info( "Approved guardrail %s (ID: %s) and initialized in memory", @@ -1295,10 +1299,18 @@ async def get_guardrail_info(guardrail_id: str): guardrail_id=guardrail_id, prisma_client=prisma_client ) if result is None: - result = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id( + in_memory = IN_MEMORY_GUARDRAIL_HANDLER.get_guardrail_by_id( guardrail_id=guardrail_id ) - guardrail_definition_location = GUARDRAIL_DEFINITION_LOCATION.CONFIG + # Only return config-loaded entries here. A DB-backed entry that's + # missing from the DB is stale (deleted on another pod, awaiting + # reconciliation on this one) and must surface as 404. + if ( + in_memory is not None + and IN_MEMORY_GUARDRAIL_HANDLER.get_source(guardrail_id) == "config" + ): + result = in_memory + guardrail_definition_location = GUARDRAIL_DEFINITION_LOCATION.CONFIG if result is None: raise HTTPException( diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 868b23756d2..3aa612c761e 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -3,7 +3,7 @@ import importlib import os from datetime import datetime, timezone -from typing import Any, Dict, List, Optional, Type, cast +from typing import Any, Dict, List, Literal, Optional, Set, Type, cast import litellm from litellm import Router @@ -403,11 +403,19 @@ class InMemoryGuardrailHandler: Guardrail id to CustomGuardrail object mapping """ + self._sources: Dict[str, Literal["db", "config"]] = {} + """ + Guardrail id to provenance marker. "db" entries are reconciled against + the DB on each polling tick; "config" entries are owned by proxy_config.yaml + and never deleted by reconciliation. + """ + def initialize_guardrail( self, guardrail: Guardrail, config_file_path: Optional[str] = None, llm_router: Optional["Router"] = None, + source: Literal["db", "config"] = "config", ) -> Optional[Guardrail]: """ Initialize a guardrail from a dictionary and add it to the litellm callback manager @@ -420,6 +428,10 @@ class InMemoryGuardrailHandler: verbose_proxy_logger.debug( "guardrail_id already exists in IN_MEMORY_GUARDRAILS" ) + # Honor the caller's source even on the early-return path so a + # racing polling tick or a hot-reload of config can correct an + # entry's provenance. + self._sources[guardrail_id] = source return self.IN_MEMORY_GUARDRAILS[guardrail_id] custom_guardrail_callback: Optional[CustomGuardrail] = None @@ -492,6 +504,7 @@ class InMemoryGuardrailHandler: # store references to the guardrail in memory self.IN_MEMORY_GUARDRAILS[guardrail_id] = parsed_guardrail self.guardrail_id_to_custom_guardrail[guardrail_id] = custom_guardrail_callback + self._sources[guardrail_id] = source return parsed_guardrail @@ -552,7 +565,10 @@ class InMemoryGuardrailHandler: return _guardrail_callback def update_in_memory_guardrail( - self, guardrail_id: str, guardrail: Guardrail + self, + guardrail_id: str, + guardrail: Guardrail, + source: Literal["db", "config"] = "db", ) -> None: """ Update a guardrail in memory @@ -561,6 +577,7 @@ class InMemoryGuardrailHandler: - updates the guardrail params in litellm.callback_manager """ self.IN_MEMORY_GUARDRAILS[guardrail_id] = guardrail + self._sources[guardrail_id] = source custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.get( guardrail_id @@ -579,6 +596,7 @@ class InMemoryGuardrailHandler: """ # Remove from in-memory storage self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None) + self._sources.pop(guardrail_id, None) # Remove the callback from litellm.callbacks custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.pop( @@ -603,6 +621,34 @@ class InMemoryGuardrailHandler: """ return self.IN_MEMORY_GUARDRAILS.get(guardrail_id) + def get_source(self, guardrail_id: str) -> Optional[Literal["db", "config"]]: + """ + Return the provenance of an in-memory guardrail. + """ + return self._sources.get(guardrail_id) + + def reconcile_db_guardrails(self, db_guardrail_ids: Set[str]) -> List[str]: + """ + Drop in-memory entries that originated from the DB but are no longer + present in db_guardrail_ids. Config-loaded guardrails are never touched. + + Called by the periodic DB polling tick so that a guardrail deleted + on another pod is eventually purged from this pod's memory + callbacks. + """ + stale_ids = [ + guardrail_id + for guardrail_id, source in self._sources.items() + if source == "db" and guardrail_id not in db_guardrail_ids + ] + for guardrail_id in stale_ids: + verbose_proxy_logger.info( + "Reconcile: removing stale DB-backed guardrail '%s' from memory " + "(deleted in DB by another pod)", + guardrail_id, + ) + self.delete_in_memory_guardrail(guardrail_id) + return stale_ids + def _has_guardrail_params_changed( self, guardrail_id: str, new_guardrail: Guardrail ) -> bool: @@ -656,7 +702,10 @@ class InMemoryGuardrailHandler: return len(changed_fields) > 0 def reinitialize_guardrail( - self, guardrail: Guardrail, config_file_path: Optional[str] = None + self, + guardrail: Guardrail, + config_file_path: Optional[str] = None, + source: Literal["db", "config"] = "config", ) -> Optional[Guardrail]: """ Force re-initialization of a guardrail even if it exists in memory. @@ -675,7 +724,7 @@ class InMemoryGuardrailHandler: # Initialize fresh (will add new callback to litellm.callbacks) return self.initialize_guardrail( - guardrail=guardrail, config_file_path=config_file_path + guardrail=guardrail, config_file_path=config_file_path, source=source ) def sync_guardrail_from_db( @@ -696,9 +745,15 @@ class InMemoryGuardrailHandler: f"Guardrail '{guardrail_name}' (ID: {guardrail_id}) params changed, re-initializing..." ) return self.reinitialize_guardrail( - guardrail=guardrail, config_file_path=config_file_path + guardrail=guardrail, + config_file_path=config_file_path, + source="db", ) + # Params unchanged but the entry is still DB-backed; make sure the + # source marker reflects that even if it was previously set differently + # (e.g. a config entry whose UUID later collided with a DB row). + self._sources[guardrail_id] = "db" return self.IN_MEMORY_GUARDRAILS.get(guardrail_id) diff --git a/litellm/proxy/guardrails/init_guardrails.py b/litellm/proxy/guardrails/init_guardrails.py index d742cc223b4..83f1281dc02 100644 --- a/litellm/proxy/guardrails/init_guardrails.py +++ b/litellm/proxy/guardrails/init_guardrails.py @@ -30,6 +30,7 @@ def init_guardrails_v2( guardrail=cast(Guardrail, guardrail), config_file_path=config_file_path, llm_router=llm_router, + source="config", ) if initialized_guardrail: guardrail_list.append(initialized_guardrail) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 096e23e673d..2c857c0fde8 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1742,29 +1742,54 @@ async def test_model_connection( # Look up model configuration from router if model name is provided # This gets the litellm_params from proxy config (with resolved env vars) config_litellm_params: dict = {} - if model_name and llm_router is not None: + if llm_router is not None: + # Prefer disambiguation by deployment id (`model_info.id`) when + # the caller supplies it. This is required when multiple + # deployments share a `model_name` (e.g. wildcard `openai/*` + # with multiple `api_base` values for failover): the UI's + # "Test Connection" button targets a specific row, and that + # row's id is the only thing that uniquely identifies which + # deployment to probe. Without this, all duplicates collapse + # onto `deployments[0]`. + request_model_info = model_info or {} + request_model_id = request_model_info.get("id") try: - # First try to find by proxy model_name (e.g., "gpt-4o") - deployments = llm_router.get_model_list(model_name=model_name) - - # If not found, try to find by litellm model name (e.g., "azure/gpt-4o") - if not deployments or len(deployments) == 0: - all_deployments = llm_router.get_model_list(model_name=None) - if all_deployments: - for deployment in all_deployments: - if ( - deployment.get("litellm_params", {}).get("model") - == model_name - ): - deployments = [deployment] - break - - if deployments and len(deployments) > 0: - # Use the first deployment's litellm_params as base config - # These already have resolved environment variables from proxy config - config_litellm_params = dict( - deployments[0].get("litellm_params", {}) + deployment_by_id = None + if request_model_id: + deployment_by_id = llm_router.get_deployment( + model_id=request_model_id ) + + if deployment_by_id is not None: + config_litellm_params = deployment_by_id.litellm_params.model_dump( + exclude_none=True + ) + elif model_name: + # Fall back to model_name lookup for callers (e.g. the + # "Add Model" wizard, or curl) that don't supply an id. + # First try to find by proxy model_name (e.g., "gpt-4o") + deployments = llm_router.get_model_list(model_name=model_name) + + # If not found, try to find by litellm model name + # (e.g., "azure/gpt-4o") + if not deployments or len(deployments) == 0: + all_deployments = llm_router.get_model_list(model_name=None) + if all_deployments: + for deployment in all_deployments: + if ( + deployment.get("litellm_params", {}).get("model") + == model_name + ): + deployments = [deployment] + break + + if deployments and len(deployments) > 0: + # Use the first deployment's litellm_params as base + # config. These already have resolved environment + # variables from proxy config. + config_litellm_params = dict( + deployments[0].get("litellm_params", {}) + ) except Exception as e: verbose_proxy_logger.debug( f"Could not find model {model_name} in router: {e}. " diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 0e60820aab1..2a4895d0299 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -12,9 +12,10 @@ All /tag management endpoints import asyncio import json -from typing import TYPE_CHECKING, Dict, List, Optional +from datetime import datetime +from typing import TYPE_CHECKING, Any, Dict, List, Optional -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Query from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth @@ -395,6 +396,32 @@ async def info_tag( raise HTTPException(status_code=500, detail=str(e)) +def _validate_tag_list_date_range( + start_date: Optional[str], end_date: Optional[str] +) -> None: + """Require both dates together, and enforce YYYY-MM-DD format with start <= end.""" + if (start_date is None) != (end_date is None): + raise HTTPException( + status_code=400, + detail="start_date and end_date must be provided together", + ) + if start_date is None: + return + try: + start = datetime.strptime(start_date, "%Y-%m-%d") + end = datetime.strptime(end_date, "%Y-%m-%d") # type: ignore[arg-type] + except ValueError as e: + raise HTTPException( + status_code=400, + detail=f"Invalid date format, expected YYYY-MM-DD: {e}", + ) + if start > end: + raise HTTPException( + status_code=400, + detail="start_date must be on or before end_date", + ) + + @router.get( "/tag/list", tags=["tag management"], @@ -402,6 +429,18 @@ async def info_tag( ) async def list_tags( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + start_date: Optional[str] = Query( + None, + description=( + "Optional start date (YYYY-MM-DD). When provided together with " + "end_date, dynamic tags are limited to those active in the window. " + "Stored tags are always returned." + ), + ), + end_date: Optional[str] = Query( + None, + description="Optional end date (YYYY-MM-DD). Must be given with start_date.", + ), ): """ List all available tags with their budget information. @@ -411,6 +450,8 @@ async def list_tags( if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") + _validate_tag_list_date_range(start_date, end_date) + try: ## QUERY STORED TAGS ## tag_records = await prisma_client.db.litellm_tagtable.find_many( @@ -453,9 +494,13 @@ async def list_tags( # Prisma's distinct fetches all columns for all rows and deduplicates # in application code, which is extremely slow on large tables. # See: https://www.prisma.io/docs/orm/prisma-client/queries/aggregation-grouping-summarizing#distinct-under-the-hood + dynamic_tag_where: Dict[str, Any] = {"tag": {"not": None}} + if start_date is not None and end_date is not None: + dynamic_tag_where["date"] = {"gte": start_date, "lte": end_date} + dynamic_tag_rows = await prisma_client.db.litellm_dailytagspend.group_by( by=["tag"], - where={"tag": {"not": None}}, + where=dynamic_tag_where, min={"created_at": True}, max={"updated_at": True}, ) diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 74ee7c7220d..ea629f24485 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -740,7 +740,7 @@ def generic_response_convertor( all_teams = [] if sso_jwt_handler is not None: - team_ids = sso_jwt_handler.get_team_ids_from_jwt(cast(dict, response)) + team_ids = sso_jwt_handler.get_all_jwt_team_ids(cast(dict, response)) all_teams.extend(team_ids) if team_mappings is not None and team_mappings.team_ids_jwt_field is not None: @@ -755,7 +755,7 @@ def generic_response_convertor( f"Loaded team_ids from DB team_mappings.team_ids_jwt_field='{team_mappings.team_ids_jwt_field}': {team_ids_from_db_mapping}" ) else: - team_ids = jwt_handler.get_team_ids_from_jwt(cast(dict, response)) + team_ids = jwt_handler.get_all_jwt_team_ids(cast(dict, response)) all_teams.extend(team_ids) # Determine user role based on role_mappings if available diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c96d0acb008..f0b3faa7267 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5936,10 +5936,20 @@ class ProxyConfig: verbose_proxy_logger.debug( "guardrails from the DB %s", str(guardrails_in_db) ) + db_guardrail_ids: set = set() for guardrail in guardrails_in_db: + guardrail_id = guardrail.get("guardrail_id") + if guardrail_id: + db_guardrail_ids.add(guardrail_id) IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db( guardrail=cast(Guardrail, guardrail), ) + + # 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 - {}".format( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py index 89992c510f5..9f004318488 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py @@ -1135,3 +1135,69 @@ def test_validate_loopback_redirect_uri_rejects_malformed_cleanly(): with pytest.raises(HTTPException) as exc: validate_loopback_redirect_uri("http://[not-an-ip]/cb") assert exc.value.status_code == 400 + + +def _mock_request_with_base_url(base_url: str): + req = MagicMock() + req.base_url = base_url + req.headers = {} + return req + + +def test_validate_trusted_redirect_uri_accepts_same_origin(): + """UI OAuth flow: redirect_uri on the proxy's own origin is allowed.""" + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + req = _mock_request_with_base_url("https://proxy.example.com/") + # Should not raise. + validate_trusted_redirect_uri( + req, "https://proxy.example.com/ui/mcp/oauth/callback" + ) + + +def test_validate_trusted_redirect_uri_accepts_loopback(): + """Native MCP client flow: loopback is still allowed.""" + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + req = _mock_request_with_base_url("https://proxy.example.com/") + validate_trusted_redirect_uri(req, "http://127.0.0.1:3000/cb") + validate_trusted_redirect_uri(req, "http://localhost:3000/cb") + + +def test_validate_trusted_redirect_uri_rejects_external_origin(): + """An attacker-controlled origin must still be rejected.""" + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + req = _mock_request_with_base_url("https://proxy.example.com/") + with pytest.raises(HTTPException) as exc: + validate_trusted_redirect_uri(req, "https://attacker.example.com/cb") + assert exc.value.status_code == 400 + + +def test_validate_trusted_redirect_uri_rejects_scheme_mismatch(): + """https→http (or vice versa) on the same host is not same-origin.""" + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + req = _mock_request_with_base_url("https://proxy.example.com/") + with pytest.raises(HTTPException) as exc: + validate_trusted_redirect_uri(req, "http://proxy.example.com/ui/callback") + assert exc.value.status_code == 400 + + +def test_validate_trusted_redirect_uri_rejects_fragment(): + from litellm.proxy._experimental.mcp_server.oauth_utils import ( + validate_trusted_redirect_uri, + ) + + req = _mock_request_with_base_url("https://proxy.example.com/") + with pytest.raises(HTTPException) as exc: + validate_trusted_redirect_uri(req, "https://proxy.example.com/ui/cb#code=1") + assert exc.value.status_code == 400 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 85d5d6ba466..581324d47d2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -23,6 +23,20 @@ def mock_mcp_client_ip(): yield +def _mock_callback_request(base_url: str = "http://localhost:3000/"): + """Return a MagicMock Request for callback/authorize same-origin tests. + + The callback handler only uses ``request`` to compute the proxy's own + base URL via ``get_request_base_url`` (which reads ``request.base_url`` + and trusted ``X-Forwarded-*`` headers). A simple MagicMock with the + right attributes is sufficient. + """ + req = MagicMock() + req.base_url = base_url + req.headers = {} + return req + + @pytest.fixture def trust_xff(): """Force ``IPAddressUtils.is_request_from_trusted_proxy`` to True. @@ -1844,6 +1858,7 @@ async def test_oauth_callback_redirects_with_state(): # Call callback endpoint with code and state response = await callback( + request=_mock_callback_request(), code="test_authorization_code_12345", state="encrypted_state_value", ) @@ -1887,6 +1902,7 @@ async def test_oauth_callback_preserves_client_redirect_uri_query(): } response = await callback( + request=_mock_callback_request(), code="test_authorization_code_12345", state="encrypted_state_value", ) @@ -1917,6 +1933,7 @@ async def test_oauth_callback_handles_invalid_state(): # Call callback endpoint with invalid state response = await callback( + request=_mock_callback_request(), code="test_code", state="invalid_encrypted_state", ) @@ -1926,6 +1943,40 @@ async def test_oauth_callback_handles_invalid_state(): assert "Authentication incomplete" in response.body.decode() +@pytest.mark.asyncio +async def test_oauth_callback_accepts_same_origin_ui_redirect(): + """UI OAuth flow: the callback should redirect to the proxy's own UI + origin when the encrypted state carries a same-origin client_redirect_uri.""" + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( + callback, + ) + + with patch( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints.decode_state_hash" + ) as mock_decode: + mock_decode.return_value = { + "base_url": "https://proxy.example.com/ui/mcp/oauth/callback", + "original_state": "state-123", + "code_challenge": None, + "code_challenge_method": None, + "client_redirect_uri": "https://proxy.example.com/ui/mcp/oauth/callback", + } + + response = await callback( + request=_mock_callback_request(base_url="https://proxy.example.com/"), + code="auth-code-123", + state="encrypted_state", + ) + + assert response.status_code == 302 + assert ( + "https://proxy.example.com/ui/mcp/oauth/callback" + in response.headers["location"] + ) + assert "code=auth-code-123" in response.headers["location"] + assert "state=state-123" in response.headers["location"] + + @pytest.mark.asyncio async def test_oauth_authorize_includes_scopes_from_server_config(): """Test that authorize endpoint includes scopes from server configuration.""" @@ -2307,7 +2358,11 @@ async def test_callback_revalidates_loopback_on_decoded_base_url(): "client_redirect_uri": "https://attacker.example.com/cb", } with pytest.raises(HTTPException) as exc_info: - await callback(code="stolen_code", state="encrypted_stale_state") + await callback( + request=_mock_callback_request(), + code="stolen_code", + state="encrypted_stale_state", + ) assert exc_info.value.status_code == 400 @@ -2329,7 +2384,11 @@ async def test_callback_revalidates_loopback_on_decoded_client_redirect_uri(): "client_redirect_uri": "https://attacker.example.com/cb", } with pytest.raises(HTTPException) as exc_info: - await callback(code="stolen_code", state="encrypted_stale_state") + await callback( + request=_mock_callback_request(), + code="stolen_code", + state="encrypted_stale_state", + ) assert exc_info.value.status_code == 400 @@ -2349,7 +2408,11 @@ async def test_callback_rejects_state_missing_redirect_uri(): "code_challenge_method": None, } with pytest.raises(HTTPException) as exc_info: - await callback(code="code", state="encrypted_malformed_state") + await callback( + request=_mock_callback_request(), + code="code", + state="encrypted_malformed_state", + ) assert exc_info.value.status_code == 400 diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index b7dba9c1d16..c09c303deee 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1,5 +1,5 @@ from typing import Optional -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -494,6 +494,75 @@ async def test_sync_user_role_and_teams_no_cache_write_when_nothing_changes(): mock_cache.async_set_cache.assert_not_called() +def test_get_all_jwt_team_ids_unions_singular_and_plural(): + """get_all_jwt_team_ids must include the singular team_id_jwt_field claim + in addition to the plural team_ids_jwt_field, deduplicated.""" + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=MagicMock(), + litellm_jwtauth=LiteLLM_JWTAuth( + team_id_jwt_field="team_id", + team_ids_jwt_field="teams", + ), + ) + + # singular only — Okta/Auth0 default shape + assert jwt_handler.get_all_jwt_team_ids({"team_id": "team-low"}) == ["team-low"] + + # plural only — pre-fix shape + assert jwt_handler.get_all_jwt_team_ids({"teams": ["a", "b"]}) == ["a", "b"] + + # both populated, no overlap + assert jwt_handler.get_all_jwt_team_ids( + {"team_id": "primary", "teams": ["a", "b"]} + ) == ["a", "b", "primary"] + + # both populated with overlap — singular dedup'd + assert jwt_handler.get_all_jwt_team_ids({"team_id": "a", "teams": ["a", "b"]}) == [ + "a", + "b", + ] + + # neither populated + assert jwt_handler.get_all_jwt_team_ids({}) == [] + + +def test_get_all_jwt_team_ids_does_not_use_team_id_default(): + """team_id_default is a JWT-bearer-flow auth-builder fallback, not a token + claim. It must NOT leak into get_all_jwt_team_ids — otherwise SSO logins + would silently start adding users to the default team for any tenant that + has team_id_default configured.""" + jwt_handler = JWTHandler() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=MagicMock(), + litellm_jwtauth=LiteLLM_JWTAuth( + team_id_jwt_field="team_id", + team_ids_jwt_field="teams", + team_id_default="default-team", + ), + ) + + # team_id claim missing — must not fall back to default-team + assert jwt_handler.get_all_jwt_team_ids({"teams": []}) == [] + assert jwt_handler.get_all_jwt_team_ids({}) == [] + + # only the plural is populated — default still must not be added + assert jwt_handler.get_all_jwt_team_ids({"teams": ["a"]}) == ["a"] + + # team_id_jwt_field unset entirely + only default configured: still no default + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=MagicMock(), + litellm_jwtauth=LiteLLM_JWTAuth( + team_ids_jwt_field="teams", + team_id_default="default-team", + ), + ) + assert jwt_handler.get_all_jwt_team_ids({"teams": []}) == [] + + @pytest.mark.asyncio async def test_map_jwt_role_to_litellm_role(): """Test JWT role mapping to LiteLLM roles with various patterns""" diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 033deb3ff42..0d7becd3e2e 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -106,9 +106,11 @@ def mock_in_memory_handler(mocker): mock_handler = mocker.Mock(spec=InMemoryGuardrailHandler) mock_handler.list_in_memory_guardrails.return_value = [MOCK_CONFIG_GUARDRAIL] mock_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL + mock_handler.get_source.return_value = "config" mock_handler.initialize_guardrail = mocker.Mock() mock_handler.update_in_memory_guardrail = mocker.Mock() mock_handler.delete_in_memory_guardrail = mocker.Mock() + mock_handler.reconcile_db_guardrails = mocker.Mock(return_value=[]) return mock_handler @@ -162,6 +164,67 @@ async def test_list_guardrails_v2_with_db_and_config( assert isinstance(config_guardrail.litellm_params, BaseLitellmParams) +@pytest.mark.asyncio +async def test_list_guardrails_v2_skips_stale_db_backed_in_memory_entries(mocker): + """ + A guardrail that's still in this pod's memory tagged source='db' but is no + longer in the DB result (deleted on another pod, awaiting reconcile) must + NOT surface in the list response — pre-fix it leaked as 'config'. + """ + stale_guardrail = { + "guardrail_id": "stale-db-id", + "guardrail_name": "Stale DB Guardrail", + "litellm_params": {"guardrail": "bedrock", "mode": "pre_call"}, + "guardrail_info": {}, + } + mock_prisma_client = mocker.Mock() + mock_prisma_client.db = mocker.Mock() + mock_prisma_client.db.litellm_guardrailstable = mocker.Mock() + mock_prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(return_value=[]) + + mock_in_memory_handler = mocker.Mock() + mock_in_memory_handler.list_in_memory_guardrails.return_value = [stale_guardrail] + mock_in_memory_handler.get_source.return_value = "db" + + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + + admin_auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) + response = await list_guardrails_v2(user_api_key_dict=admin_auth) + + assert response.guardrails == [] + mock_in_memory_handler.get_source.assert_called_with("stale-db-id") + + +@pytest.mark.asyncio +async def test_get_guardrail_info_404s_stale_db_backed_entry( + mocker, mock_prisma_client, mock_in_memory_handler +): + """ + Stale DB-backed entry (in-memory but not in DB) must 404 instead of being + returned as if it were a config-loaded guardrail. + """ + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch( + "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", + mock_in_memory_handler, + ) + mock_prisma_client.db.litellm_guardrailstable.find_unique = AsyncMock( + return_value=None + ) + # In-memory still has it, but it's tagged as 'db' (stale, awaiting reconcile) + mock_in_memory_handler.get_source.return_value = "db" + + with pytest.raises(HTTPException) as exc_info: + await get_guardrail_info("stale-db-id") + + assert exc_info.value.status_code == 404 + assert "not found" in str(exc_info.value.detail) + + @pytest.mark.asyncio async def test_list_guardrails_v2_masks_sensitive_data_in_db_guardrails(mocker): """Test that sensitive litellm_params are masked for DB guardrails in list response""" @@ -1160,6 +1223,7 @@ async def test_get_guardrail_info_endpoint_config_guardrail(mocker): # Mock IN_MEMORY_GUARDRAIL_HANDLER at its source to return config guardrail mock_in_memory_handler = mocker.Mock() mock_in_memory_handler.get_guardrail_by_id.return_value = MOCK_CONFIG_GUARDRAIL + mock_in_memory_handler.get_source.return_value = "config" mocker.patch( "litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", mock_in_memory_handler, diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py index 1d70126681d..9f7173383b0 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_registry.py @@ -60,3 +60,123 @@ def test_update_in_memory_guardrail(): handler.guardrail_id_to_custom_guardrail["123"].event_hook is GuardrailEventHooks.pre_call ) + + +def _make_guardrail(guardrail_id: str, name: str = "g") -> Guardrail: + return Guardrail( + guardrail_id=guardrail_id, + guardrail_name=name, + litellm_params=LitellmParams(guardrail=name, mode="pre_call", default_on=False), + ) + + +def test_reconcile_db_guardrails_drops_stale_db_entries_only(): + """ + The reconcile pass must drop in-memory entries marked source='db' that are + missing from the DB result, and never touch source='config' entries. + Models the multi-pod case where another pod deleted a DB-backed guardrail. + """ + handler = InMemoryGuardrailHandler() + + # Two DB-backed entries on this pod (synced from earlier polling cycles) + handler.IN_MEMORY_GUARDRAILS["db-keep"] = _make_guardrail("db-keep") + handler.IN_MEMORY_GUARDRAILS["db-stale"] = _make_guardrail("db-stale") + handler._sources["db-keep"] = "db" + handler._sources["db-stale"] = "db" + + # One config-loaded entry that must survive reconciliation + handler.IN_MEMORY_GUARDRAILS["cfg"] = _make_guardrail("cfg") + handler._sources["cfg"] = "config" + + # The DB now only contains db-keep — db-stale was deleted on another pod. + removed = handler.reconcile_db_guardrails(db_guardrail_ids={"db-keep"}) + + assert removed == ["db-stale"] + assert "db-stale" not in handler.IN_MEMORY_GUARDRAILS + assert "db-stale" not in handler._sources + assert "db-keep" in handler.IN_MEMORY_GUARDRAILS + assert "cfg" in handler.IN_MEMORY_GUARDRAILS + assert handler._sources["cfg"] == "config" + + +def test_reconcile_does_not_drop_config_entries_missing_from_db(): + """A config-only guardrail (no DB row) must never be reconciled away.""" + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["cfg-only"] = _make_guardrail("cfg-only") + handler._sources["cfg-only"] = "config" + + removed = handler.reconcile_db_guardrails(db_guardrail_ids=set()) + + assert removed == [] + assert "cfg-only" in handler.IN_MEMORY_GUARDRAILS + + +def test_get_source_returns_marker_set_at_insert(): + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["a"] = _make_guardrail("a") + handler._sources["a"] = "db" + handler.IN_MEMORY_GUARDRAILS["b"] = _make_guardrail("b") + handler._sources["b"] = "config" + + assert handler.get_source("a") == "db" + assert handler.get_source("b") == "config" + assert handler.get_source("missing") is None + + +def test_delete_in_memory_guardrail_clears_source_marker(): + handler = InMemoryGuardrailHandler() + handler.IN_MEMORY_GUARDRAILS["a"] = _make_guardrail("a") + handler._sources["a"] = "db" + + handler.delete_in_memory_guardrail("a") + + assert "a" not in handler.IN_MEMORY_GUARDRAILS + assert "a" not in handler._sources + assert handler.get_source("a") is None + + +def test_initialize_guardrail_early_return_updates_source_marker(): + """ + When initialize_guardrail is called for a guardrail that already exists + in memory, the early-return path must still honor the caller's source. + Otherwise a racing polling tick that placed a DB entry in memory first + would leave a later config-init call wrongly marked as 'db' (or vice + versa), and the entry would be reconciled with the wrong classification. + """ + handler = InMemoryGuardrailHandler() + # Simulate a polling tick already placing the entry as DB-backed. + handler.IN_MEMORY_GUARDRAILS["collide"] = _make_guardrail("collide", name="bedrock") + handler._sources["collide"] = "db" + + # Config init re-visits the same id (e.g., hot-reload, or UUID collision). + g = Guardrail( + guardrail_id="collide", + guardrail_name="bedrock", + litellm_params=LitellmParams( + guardrail="bedrock", mode="pre_call", default_on=False + ), + ) + handler.initialize_guardrail(guardrail=g, source="config") + + assert handler.get_source("collide") == "config" + + # And the symmetric direction: db sync should override an entry left + # marked as 'config' from a stale init path. + handler.initialize_guardrail(guardrail=g, source="db") + assert handler.get_source("collide") == "db" + + +def test_sync_guardrail_from_db_marks_source_db_when_unchanged(): + """ + sync_guardrail_from_db must enforce source='db' even when params are + unchanged, so a config entry whose UUID happens to collide with a later + DB row gets re-tagged correctly. + """ + handler = InMemoryGuardrailHandler() + g = _make_guardrail("collide") + handler.IN_MEMORY_GUARDRAILS["collide"] = g + handler._sources["collide"] = "config" + + handler.sync_guardrail_from_db(g) + + assert handler.get_source("collide") == "db" diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 2edcb00c967..bcd7fcb37b3 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -466,6 +466,236 @@ async def test_test_model_connection_loads_config_from_router(): assert "result" in result +@pytest.mark.asyncio +async def test_test_model_connection_uses_model_info_id_to_disambiguate_duplicate_model_names(): + """ + When two deployments share the same `model_name` (e.g. wildcard + `openai/*`) but have different `api_base` values, clicking "Test + Connection" on a specific row in the UI must probe THAT row's + `api_base` — not whichever happens to be `deployments[0]`. + + The UI passes `model_info.id` to identify the deployment the user + actually clicked on. The backend must use that id to look up the + specific deployment rather than always grabbing the first match. + + Regression test for: silent fallback to deployments[0] when + multiple deployments share a wildcard model_name. + """ + from litellm.types.router import Deployment, LiteLLM_Params + + mock_request = MagicMock() + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.token = "test-token" + + mock_prisma_client = MagicMock() + + deployment_a = { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_base": "https://deployment-A-base.invalid/v1", + "api_key": "fake-key-A", + }, + "model_info": {"id": "deployment-A-id"}, + } + deployment_b = { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_base": "https://deployment-B-base.invalid/v1", + "api_key": "fake-key-B", + }, + "model_info": {"id": "deployment-B-id"}, + } + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [deployment_a, deployment_b] + + # Backend uses get_deployment(model_id=...) for O(1) lookup by id. + def _get_deployment_by_id(model_id): + if model_id == "deployment-A-id": + return Deployment( + model_name="openai/*", + litellm_params=LiteLLM_Params(**deployment_a["litellm_params"]), + model_info=deployment_a["model_info"], + ) + if model_id == "deployment-B-id": + return Deployment( + model_name="openai/*", + litellm_params=LiteLLM_Params(**deployment_b["litellm_params"]), + model_info=deployment_b["model_info"], + ) + return None + + mock_router.get_deployment.side_effect = _get_deployment_by_id + + mock_can_user_make_model_call = AsyncMock() + + mock_health_check_result = {"status": "healthy", "response_time_ms": 50} + mock_ahealth_check = AsyncMock(return_value=mock_health_check_result) + mock_run_with_timeout = AsyncMock(return_value=mock_health_check_result) + + def mock_update_params(model_info, litellm_params): + params = litellm_params.copy() + params["messages"] = [{"role": "user", "content": "test"}] + return params + + def mock_reject_os_environ(params): + return None + + with ( + patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ), + patch( + "litellm.proxy.proxy_server.llm_router", + mock_router, + ), + patch( + "litellm.proxy.proxy_server.premium_user", + False, + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + mock_can_user_make_model_call, + ), + patch( + "litellm.proxy.health_endpoints._health_endpoints.litellm.ahealth_check", + mock_ahealth_check, + ), + patch( + "litellm.proxy.health_endpoints._health_endpoints.run_with_timeout", + mock_run_with_timeout, + ), + patch( + "litellm.proxy.health_endpoints._health_endpoints._update_litellm_params_for_health_check", + mock_update_params, + ), + patch( + "litellm.proxy.health_endpoints._health_endpoints._reject_os_environ_references", + mock_reject_os_environ, + ), + ): + # Click "Test Connection" on deployment B (NOT the first one). + # The UI sends only `model` + `model_info.id` — it does NOT + # send `api_base`/`api_key`, so the backend must resolve them + # from the right deployment. + await health_test_model_connection( + request=mock_request, + mode="chat", + litellm_params={"model": "openai/*"}, + model_info={"id": "deployment-B-id"}, + user_api_key_dict=mock_user_api_key_dict, + ) + + # The outbound health check must hit deployment B's api_base. + ahealth_check_call_args = mock_ahealth_check.call_args + assert ahealth_check_call_args is not None + model_params = ahealth_check_call_args.kwargs.get("model_params", {}) + + assert model_params.get("api_base") == ( + "https://deployment-B-base.invalid/v1" + ), ( + "Expected /health/test_connection to probe deployment B's " + "api_base when model_info.id='deployment-B-id' was provided. " + f"Got: {model_params.get('api_base')!r}. This means the " + "backend silently fell back to deployments[0] (A) instead " + "of disambiguating by model_info.id." + ) + assert model_params.get("api_key") == "fake-key-B" + + +@pytest.mark.asyncio +async def test_test_model_connection_falls_back_to_deployments_zero_without_id(): + """ + Backwards-compat: when the request body does NOT include + `model_info.id`, the legacy behavior of using `deployments[0]` + is preserved (single-deployment case, or callers that haven't + been updated to pass an id). + """ + mock_request = MagicMock() + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.token = "test-token" + + mock_prisma_client = MagicMock() + + deployment_a = { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_base": "https://deployment-A-base.invalid/v1", + "api_key": "fake-key-A", + }, + "model_info": {"id": "deployment-A-id"}, + } + deployment_b = { + "model_name": "openai/*", + "litellm_params": { + "model": "openai/*", + "api_base": "https://deployment-B-base.invalid/v1", + "api_key": "fake-key-B", + }, + "model_info": {"id": "deployment-B-id"}, + } + + mock_router = MagicMock() + mock_router.get_model_list.return_value = [deployment_a, deployment_b] + + mock_can_user_make_model_call = AsyncMock() + mock_health_check_result = {"status": "healthy"} + mock_ahealth_check = AsyncMock(return_value=mock_health_check_result) + mock_run_with_timeout = AsyncMock(return_value=mock_health_check_result) + + def mock_update_params(model_info, litellm_params): + params = litellm_params.copy() + params["messages"] = [{"role": "user", "content": "test"}] + return params + + def mock_reject_os_environ(params): + return None + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch("litellm.proxy.proxy_server.premium_user", False), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + mock_can_user_make_model_call, + ), + patch( + "litellm.proxy.health_endpoints._health_endpoints.litellm.ahealth_check", + mock_ahealth_check, + ), + patch( + "litellm.proxy.health_endpoints._health_endpoints.run_with_timeout", + mock_run_with_timeout, + ), + patch( + "litellm.proxy.health_endpoints._health_endpoints._update_litellm_params_for_health_check", + mock_update_params, + ), + patch( + "litellm.proxy.health_endpoints._health_endpoints._reject_os_environ_references", + mock_reject_os_environ, + ), + ): + await health_test_model_connection( + request=mock_request, + mode="chat", + litellm_params={"model": "openai/*"}, + model_info={}, # no id provided + user_api_key_dict=mock_user_api_key_dict, + ) + + # Without id, deployments[0] (A) should be used (legacy behavior). + model_params = mock_ahealth_check.call_args.kwargs.get("model_params", {}) + assert model_params.get("api_base") == "https://deployment-A-base.invalid/v1" + assert model_params.get("api_key") == "fake-key-A" + + @pytest.mark.asyncio async def test_health_services_endpoint_datadog_llm_observability(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py index 39ec6f075d7..ee2d72d2dd4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_tag_management_endpoints.py @@ -380,6 +380,117 @@ async def test_list_tags_no_dynamic_tags(): app.dependency_overrides.clear() +@pytest.mark.asyncio +async def test_list_tags_with_date_range_filters_dynamic_tags(): + """ + /tag/list?start_date=...&end_date=... should push the date window into + the dailytagspend group_by WHERE clause so large tables don't get scanned. + """ + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_db = Mock() + mock_prisma.db = mock_db + mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[]) + group_by_mock = AsyncMock(return_value=[]) + mock_db.litellm_dailytagspend.group_by = group_by_mock + + headers = {"Authorization": "Bearer sk-1234"} + response = client.get( + "/tag/list?start_date=2026-04-01&end_date=2026-04-29", + headers=headers, + ) + + assert response.status_code == 200 + group_by_mock.assert_awaited_once() + where = group_by_mock.await_args.kwargs["where"] + assert where["tag"] == {"not": None} + assert where["date"] == {"gte": "2026-04-01", "lte": "2026-04-29"} + + finally: + app.dependency_overrides.clear() + + +@pytest.mark.asyncio +async def test_list_tags_without_date_range_omits_date_filter(): + """When no date range is passed, the WHERE clause must not carry a date key.""" + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_db = Mock() + mock_prisma.db = mock_db + mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[]) + group_by_mock = AsyncMock(return_value=[]) + mock_db.litellm_dailytagspend.group_by = group_by_mock + + headers = {"Authorization": "Bearer sk-1234"} + response = client.get("/tag/list", headers=headers) + + assert response.status_code == 200 + group_by_mock.assert_awaited_once() + where = group_by_mock.await_args.kwargs["where"] + assert "date" not in where + + finally: + app.dependency_overrides.clear() + + +@pytest.mark.parametrize( + "query, expected_detail_fragment", + [ + ("?start_date=2026-04-01", "must be provided together"), + ("?end_date=2026-04-29", "must be provided together"), + ("?start_date=2026-04-29&end_date=2026-04-01", "on or before end_date"), + ("?start_date=not-a-date&end_date=2026-04-29", "YYYY-MM-DD"), + ], +) +@pytest.mark.asyncio +async def test_list_tags_rejects_invalid_date_range(query, expected_detail_fragment): + from unittest.mock import AsyncMock, Mock + + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + mock_user_auth = UserAPIKeyAuth( + user_id="test-user-123", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + app.dependency_overrides[user_api_key_auth] = lambda: mock_user_auth + + try: + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + mock_db = Mock() + mock_prisma.db = mock_db + mock_db.litellm_tagtable.find_many = AsyncMock(return_value=[]) + mock_db.litellm_dailytagspend.group_by = AsyncMock(return_value=[]) + + headers = {"Authorization": "Bearer sk-1234"} + response = client.get(f"/tag/list{query}", headers=headers) + + assert response.status_code == 400 + assert expected_detail_fragment in response.json()["detail"] + + finally: + app.dependency_overrides.clear() + + @pytest.mark.asyncio async def test_get_deployments_by_model_id(): """ diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 809f1d4e17b..efc1166a398 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -145,23 +145,6 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const [topKeysLimit, setTopKeysLimit] = useState(5); const [topModelsLimit, setTopModelsLimit] = useState(5); const [showTokenBreakdown, setShowTokenBreakdown] = useState(false); - const getAllTags = async () => { - if (!accessToken) { - return; - } - const tags = await tagListCall(accessToken); - setAllTags( - Object.values(tags).map((tag: Tag) => ({ - label: tag.name, - value: tag.name, - })), - ); - }; - - useEffect(() => { - getAllTags(); - }, [accessToken]); - // Sync selectedUserId when auth state settles (isAdmin/userID may be null on initial render) useEffect(() => { if (!isAdmin && userID) { @@ -175,6 +158,30 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const startTime = useMemo(() => (dateValue.from ? new Date(dateValue.from) : null), [dateValue.from]); const endTime = useMemo(() => (dateValue.to ? new Date(dateValue.to) : null), [dateValue.to]); + useEffect(() => { + if (!accessToken) return; + let cancelled = false; + (async () => { + try { + const tags = await tagListCall(accessToken, startTime, endTime); + if (cancelled) return; + setAllTags( + Object.values(tags).map((tag: Tag) => ({ + label: tag.name, + value: tag.name, + })), + ); + } catch (e) { + if (!cancelled) { + console.error("Failed to fetch tag list", e); + } + } + })(); + return () => { + cancelled = true; + }; + }, [accessToken, startTime, endTime]); + // Try aggregated endpoint first, fall back to paginated on failure const aggregatedFetchIdRef = useRef(0); useEffect(() => { diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx index 29eb8e0019b..54144d02eda 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx @@ -250,6 +250,33 @@ describe("ModelInfoView", () => { }); }); + it("should pass model_info.id to disambiguate duplicate model_name deployments", async () => { + // Regression test: when two deployments share `model_name` (e.g. + // wildcard `openai/*` with different `api_base` values), the UI + // must forward the clicked row's `model_info.id` to the backend. + // Otherwise /health/test_connection silently probes deployments[0] + // instead of the deployment the user actually selected. + const user = userEvent.setup(); + render(, { wrapper }); + + await waitFor(() => { + expect(screen.getByText("Model Settings")).toBeInTheDocument(); + }); + + const testButton = screen.getByRole("button", { name: /test connection/i }); + await user.click(testButton); + + await waitFor(() => { + expect(mockTestConnectionRequest).toHaveBeenCalled(); + }); + + const callArgs = mockTestConnectionRequest.mock.calls[0]; + // Signature: (accessToken, litellm_params, model_info, mode) + const modelInfoArg = callArgs[2] as Record; + expect(modelInfoArg).toBeDefined(); + expect(modelInfoArg.id).toBe("123"); + }); + it("should display error notification when connection test fails", async () => { const user = userEvent.setup(); mockTestConnectionRequest.mockRejectedValue(new Error("Connection failed")); diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx index ea8af2fcd62..95a43862de5 100644 --- a/ui/litellm-dashboard/src/components/model_info_view.tsx +++ b/ui/litellm-dashboard/src/components/model_info_view.tsx @@ -379,6 +379,12 @@ export default function ModelInfoView({ model: localModelData.litellm_model_name, }, { + // `id` is required to disambiguate when multiple deployments + // share the same model_name (e.g. wildcard `openai/*` with two + // different `api_base` values for failover). Without it the + // backend silently falls back to deployments[0] and probes + // the wrong endpoint. + id: localModelData.model_info?.id, mode: localModelData.model_info?.mode, }, localModelData.model_info?.mode, diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 44208904a70..abe17616bef 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -7288,10 +7288,29 @@ export const tagInfoCall = async (accessToken: string, tagNames: string[]): Prom } }; -export const tagListCall = async (accessToken: string): Promise => { +const formatYmd = (value: Date): string => { + const year = value.getFullYear(); + const month = String(value.getMonth() + 1).padStart(2, "0"); + const day = String(value.getDate()).padStart(2, "0"); + return `${year}-${month}-${day}`; +}; + +export const tagListCall = async ( + accessToken: string, + startTime?: Date | null, + endTime?: Date | null, +): Promise => { try { let url = proxyBaseUrl ? `${proxyBaseUrl}/tag/list` : `/tag/list`; + if (startTime && endTime) { + const params = new URLSearchParams({ + start_date: formatYmd(startTime), + end_date: formatYmd(endTime), + }); + url = `${url}?${params.toString()}`; + } + const response = await fetch(url, { method: "GET", headers: { diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 492e43cbc81..65bd9d9eb95 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -396,7 +396,7 @@ export default function KeyInfoView({ }; return ( -
+
- +
Key Settings {!isEditing && canModifyKey && ( From bac03ac3f162c752457990430bf8461606aae03e Mon Sep 17 00:00:00 2001 From: milan-berri Date: Fri, 8 May 2026 00:34:49 +0300 Subject: [PATCH 2/3] feat(auth): add scope and wildcard support for JWT routing overrides (#26325) Squash-merged by litellm-agent from milan-berri's PR. --- litellm/proxy/_types.py | 6 + litellm/proxy/auth/user_api_key_auth.py | 41 ++- .../proxy/auth/test_user_api_key_auth.py | 335 +++++++++++++++++- 3 files changed, 377 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ed20fe86cdc..89b422bd5d0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4342,10 +4342,16 @@ class JWTRoutingOverride(BaseModel): A rule matches when all provided selectors match token claims. If matched, request is routed to the configured auth path. + + Wildcard selectors use shell-style patterns (* and ?) and are matched with + case-sensitive semantics; use the same casing your IdP emits in JWT claims. + Space-delimited tokenization applies only to the ``scope`` claim (OAuth/OIDC + scope strings), not to ``iss``, ``aud``, or ``client_id``. """ iss: Union[str, List[str]] client_id: Optional[Union[str, List[str]]] = None + scope: Optional[Union[str, List[str]]] = None aud: Optional[Union[str, List[str]]] = None path: Literal["oauth2"] = "oauth2" diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9d3c06e641f..2f61a40bc9a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -8,6 +8,7 @@ Returns a UserAPIKeyAuth object if the API key is valid """ import asyncio +import fnmatch import re import secrets from datetime import datetime, timezone @@ -183,22 +184,49 @@ def _get_bearer_token_or_received_api_key(api_key: str) -> str: def _routing_selector_matches_claim( - selector_value: Optional[Any], claim_value: Optional[Any] + selector_value: Optional[Any], + claim_value: Optional[Any], + *, + split_space_delimited: bool = False, ) -> bool: if selector_value is None: return True - selector_list = ( + selector_list: List[str] = ( [str(v) for v in selector_value] if isinstance(selector_value, list) else [str(selector_value)] ) + if claim_value is None: + return False + if isinstance(claim_value, list): claim_list = [str(v) for v in claim_value] - return any(v in claim_list for v in selector_list) + elif ( + split_space_delimited + and isinstance(claim_value, str) + and " " in claim_value.strip() + ): + # OAuth/OIDC often sends scope as a single space-delimited string. Only split + # for the scope selector: iss/aud/client_id must stay exact full-string match + # on unverified claims (see routing override security review). The elif guard + # (`" " in claim_value.strip()`) ensures at least two non-empty tokens survive. + claim_list = [v for v in claim_value.strip().split(" ") if v] + else: + claim_list = [str(claim_value)] - return str(claim_value) in selector_list if claim_value is not None else False + def _selector_matches_claim(selector: str, claim: str) -> bool: + # NOTE: wildcard matching is case-sensitive (fnmatch.fnmatchcase). + if "*" in selector or "?" in selector: + return fnmatch.fnmatchcase(claim, selector) + return selector == claim + + return any( + _selector_matches_claim(selector=s, claim=c) + for s in selector_list + for c in claim_list + ) def _matches_routing_override( @@ -209,6 +237,11 @@ def _matches_routing_override( and _routing_selector_matches_claim( override.client_id, token_claims.get("client_id") ) + and _routing_selector_matches_claim( + override.scope, + token_claims.get("scope"), + split_space_delimited=True, + ) and _routing_selector_matches_claim(override.aud, token_claims.get("aud")) ) diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 95b3d746c66..4b0f6c0ea2f 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -29,8 +29,10 @@ from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.auth_checks import get_key_object, _cache_key_object from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import ( - _route_requires_auth_despite_public, + _matches_routing_override, _reserve_budget_after_common_checks, + _route_requires_auth_despite_public, + _routing_selector_matches_claim, _run_centralized_common_checks, _run_post_custom_auth_checks, get_api_key, @@ -572,6 +574,140 @@ def _assert_get_api_key_with_custom_litellm_key_header( ) == (api_key, passed_in_key) +@pytest.mark.parametrize( + "selector_value, claim_value, expected, split_space_delimited", + [ + (None, "any-value", True, False), + ("issuer.example.com", "issuer.example.com", True, False), + ("issuer.example.com", "other-issuer.example.com", False, False), + # iss (and other non-scope claims) must not match via space-split injection + ( + "trusted.example.com", + "trusted.example.com attacker.example.com", + False, + False, + ), + ( + ["issuer-a.example.com", "issuer-b.example.com"], + "issuer-b.example.com", + True, + False, + ), + ("*MID_LITELLM", "STREAM_MID_LITELLM", True, False), + ("*MID_LITELLM", "REDIS_LITELLM", False, False), + ("machine-??", "machine-01", True, False), + ("machine-??", "machine-001", False, False), + # Wildcard matching is case-sensitive (fnmatch.fnmatchcase) + ("*litellm", "BATCH_LITELLM", False, False), + ("*LITELLM", "BATCH_LITELLM", True, False), + ("App:LiteLLM", "App:LiteLLM openid", True, True), + ("App:*", "App:LiteLLM openid", True, True), + (["openid", "App:LiteLLM"], "openid profile", True, True), + (["service-*", "batch-*"], "batch-123", True, False), + (["service-*", "batch-*"], "other-123", False, False), + ("App:LiteLLM", ["openid", "App:LiteLLM"], True, False), + ("App:LiteLLM", None, False, False), + ], +) +def test_routing_selector_matches_claim_parametrized( + selector_value, claim_value, expected, split_space_delimited +): + assert ( + _routing_selector_matches_claim( + selector_value=selector_value, + claim_value=claim_value, + split_space_delimited=split_space_delimited, + ) + is expected + ) + + +@pytest.mark.parametrize( + "override, token_claims, expected", + [ + # Only iss selector is required and should match. + ( + JWTRoutingOverride(iss="oauth-issuer.example.com", path="oauth2"), + {"iss": "oauth-issuer.example.com"}, + True, + ), + # Scope selector narrows the match. + ( + JWTRoutingOverride( + iss="oauth-issuer.example.com", + scope="App:LiteLLM", + path="oauth2", + ), + {"iss": "oauth-issuer.example.com", "scope": "App:LiteLLM openid"}, + True, + ), + # client_id wildcard selector narrows the match. + ( + JWTRoutingOverride( + iss="oauth-issuer.example.com", + client_id="*MID_LITELLM", + path="oauth2", + ), + {"iss": "oauth-issuer.example.com", "client_id": "BATCH_MID_LITELLM"}, + True, + ), + ( + JWTRoutingOverride( + iss="oauth-issuer.example.com", + client_id="*MID_LITELLM", + path="oauth2", + ), + {"iss": "oauth-issuer.example.com", "client_id": "BATCH_PORTAL"}, + False, + ), + # aud selector still works with list claims. + ( + JWTRoutingOverride( + iss="oauth-issuer.example.com", + aud=["api://litellm", "api://fallback"], + path="oauth2", + ), + {"iss": "oauth-issuer.example.com", "aud": ["api://other", "api://litellm"]}, + True, + ), + # All provided selectors are AND-ed. + ( + JWTRoutingOverride( + iss="oauth-issuer.example.com", + scope="App:LiteLLM", + client_id="*MID_LITELLM", + path="oauth2", + ), + { + "iss": "oauth-issuer.example.com", + "scope": "App:LiteLLM openid", + "client_id": "BATCH_MID_LITELLM", + }, + True, + ), + ( + JWTRoutingOverride( + iss="oauth-issuer.example.com", + scope="App:LiteLLM", + client_id="*MID_LITELLM", + path="oauth2", + ), + { + "iss": "oauth-issuer.example.com", + "scope": "App:Other openid", + "client_id": "BATCH_MID_LITELLM", + }, + False, + ), + ], +) +def test_matches_routing_override_parametrized(override, token_claims, expected): + assert ( + _matches_routing_override(token_claims=token_claims, override=override) + is expected + ) + + def test_get_api_key_with_custom_litellm_key_header_bearer_prefix(): token = "sk-" + "1" * 8 header = f"Bearer {token}" @@ -1578,6 +1714,203 @@ class TestJWTOAuth2Coexistence: mock_jwt_auth.assert_not_called() assert result.user_id == "machine-client-aud-list" + @pytest.mark.asyncio + async def test_routing_override_matches_scope_claim(self): + """ + Match routing override when scope selector is configured and scope claim matches. + """ + jwt_token = ( + "eyJhbGciOiJSUzI1NiJ9." + "eyJpc3MiOiJvYXV0aC1pc3N1ZXIuZXhhbXBsZS5jb20iLCJzY29wZSI6IkFwcDpMaXRlTExNIiwiY2xpZW50X2lkIjoiTUFDSElORV9NSURfTElURUxMTSJ9." + "c2ln" + ) + general_settings = { + "enable_oauth2_auth": False, + "enable_jwt_auth": True, + } + mock_oauth2_response = UserAPIKeyAuth( + api_key=jwt_token, + user_id="machine-client-scope-match", + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + + with patch( + "litellm.proxy.proxy_server.general_settings", general_settings + ), patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.proxy.proxy_server.master_key", "sk-master" + ), patch( + "litellm.proxy.proxy_server.prisma_client", None + ), patch( + "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", + new_callable=AsyncMock, + return_value=mock_oauth2_response, + ) as mock_oauth2, patch( + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + ) as mock_jwt_auth: + litellm.proxy.proxy_server.jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth( + routing_overrides=[ + JWTRoutingOverride( + iss="oauth-issuer.example.com", + scope="App:LiteLLM", + path="oauth2", + ) + ] + ), + ) + + result = await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {jwt_token}", + ) + + mock_oauth2.assert_called_once_with(token=jwt_token) + mock_jwt_auth.assert_not_called() + assert result.user_id == "machine-client-scope-match" + + @pytest.mark.asyncio + async def test_routing_override_scope_mismatch_falls_back_to_jwt(self): + """ + If scope selector does not match, continue default JWT flow. + """ + jwt_token = ( + "eyJhbGciOiJSUzI1NiJ9." + "eyJpc3MiOiJvYXV0aC1pc3N1ZXIuZXhhbXBsZS5jb20iLCJzY29wZSI6IkFwcDpPdGhlciIsImNsaWVudF9pZCI6IlBPUlRBTF9NSURfTElURUxMTSJ9." + "c2ln" + ) + general_settings = { + "enable_oauth2_auth": False, + "enable_jwt_auth": True, + } + mock_jwt_result = { + "is_proxy_admin": True, + "team_object": None, + "user_object": None, + "end_user_object": None, + "org_object": None, + "token": jwt_token, + "team_id": "jwt-team", + "user_id": "jwt-user-scope-mismatch", + "end_user_id": None, + "org_id": None, + "team_membership": None, + "jwt_claims": {"sub": "user1"}, + } + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + + with patch( + "litellm.proxy.proxy_server.general_settings", general_settings + ), patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.proxy.proxy_server.master_key", "sk-master" + ), patch( + "litellm.proxy.proxy_server.prisma_client", None + ), patch( + "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", + new_callable=AsyncMock, + ) as mock_oauth2, patch( + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + return_value=mock_jwt_result, + ) as mock_jwt_auth: + litellm.proxy.proxy_server.jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth( + routing_overrides=[ + JWTRoutingOverride( + iss="oauth-issuer.example.com", + scope="App:LiteLLM", + path="oauth2", + ) + ] + ), + ) + + result = await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {jwt_token}", + ) + + mock_oauth2.assert_not_called() + mock_jwt_auth.assert_called_once() + assert result.user_id == "jwt-user-scope-mismatch" + + @pytest.mark.asyncio + async def test_routing_override_matches_scope_and_client_wildcard_when_scope_claim_is_space_delimited( + self, + ): + """ + Integration check: combined scope + wildcard selectors match on OAuth2 path + when scope claim is a space-delimited string. + """ + jwt_token = ( + "eyJhbGciOiJSUzI1NiJ9." + "eyJpc3MiOiJvYXV0aC1pc3N1ZXIuZXhhbXBsZS5jb20iLCJzY29wZSI6IkFwcDpMaXRlTExNIG9wZW5pZCIsImNsaWVudF9pZCI6IkJBVENIX01JRF9MSVRFTExNIn0." + "c2ln" + ) + general_settings = { + "enable_oauth2_auth": False, + "enable_jwt_auth": True, + } + mock_oauth2_response = UserAPIKeyAuth( + api_key=jwt_token, + user_id="machine-client-space-delimited-scope-match", + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + + with patch( + "litellm.proxy.proxy_server.general_settings", general_settings + ), patch("litellm.proxy.proxy_server.premium_user", True), patch( + "litellm.proxy.proxy_server.master_key", "sk-master" + ), patch( + "litellm.proxy.proxy_server.prisma_client", None + ), patch( + "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", + new_callable=AsyncMock, + return_value=mock_oauth2_response, + ) as mock_oauth2, patch( + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + ) as mock_jwt_auth: + litellm.proxy.proxy_server.jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=DualCache(), + litellm_jwtauth=LiteLLM_JWTAuth( + routing_overrides=[ + JWTRoutingOverride( + iss="oauth-issuer.example.com", + scope="App:LiteLLM", + client_id="*MID_LITELLM", + path="oauth2", + ) + ] + ), + ) + + result = await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {jwt_token}", + ) + + mock_oauth2.assert_called_once_with(token=jwt_token) + mock_jwt_auth.assert_not_called() + assert result.user_id == "machine-client-space-delimited-scope-match" + @pytest.mark.asyncio async def test_routing_override_routes_jwt_to_oauth2_when_oauth2_globally_disabled( self, From 055bdc35077c05314680c7c4444b8e6666852940 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 11 May 2026 11:35:56 +0530 Subject: [PATCH 3/3] fix(auth): harden JWT routing wildcard iss and merge list team_id claims Reject fnmatch wildcards on non-scope claims when the claim string contains whitespace so malformed iss values cannot match patterns like trusted.*. Merge every entry when team_id_jwt_field resolves to a list instead of keeping only the first element. Co-authored-by: Cursor --- litellm/proxy/auth/handle_jwt.py | 15 ++- litellm/proxy/auth/user_api_key_auth.py | 5 + .../proxy/auth/test_handle_jwt.py | 5 + .../proxy/auth/test_user_api_key_auth.py | 100 ++++++++++-------- 4 files changed, 77 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index a270031c33c..5b116c142b2 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -227,8 +227,8 @@ class JWTHandler: def get_all_jwt_team_ids(self, token: dict) -> List[str]: """ Return team IDs from both the plural ``team_ids_jwt_field`` and the - singular ``team_id_jwt_field`` claim, as a deduplicated list preserving - plural-first order. + singular ``team_id_jwt_field`` claim (string or list of strings), as a + deduplicated list preserving plural-first order. Membership-reconciliation paths (SSO callback, JWT-bearer sync) need to consider both claim shapes. Reading only the plural field — as @@ -249,9 +249,14 @@ class JWTHandler: default=None, ) if isinstance(singular, list): - singular = singular[0] if singular else None - if singular and singular not in team_ids: - team_ids.append(singular) + for item in singular: + if item is None: + continue + sid = str(item) + if sid and sid not in team_ids: + team_ids.append(sid) + elif singular and str(singular) not in team_ids: + team_ids.append(str(singular)) return team_ids def get_end_user_id( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 2f61a40bc9a..213c68f1db2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -219,6 +219,11 @@ def _routing_selector_matches_claim( def _selector_matches_claim(selector: str, claim: str) -> bool: # NOTE: wildcard matching is case-sensitive (fnmatch.fnmatchcase). if "*" in selector or "?" in selector: + # Without scope splitting, do not let `*` span whitespace: a malformed + # iss like "trusted.example.com evil.com" must not match "trusted.*". + # Scope uses split_space_delimited so each claim token is checked separately. + if not split_space_delimited and any(ch.isspace() for ch in claim): + return False return fnmatch.fnmatchcase(claim, selector) return selector == claim diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index c09c303deee..90c5d4f4fc5 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -524,6 +524,11 @@ def test_get_all_jwt_team_ids_unions_singular_and_plural(): "b", ] + # singular field as multi-element list (some IdPs) — merge all, preserve plural-first order + assert jwt_handler.get_all_jwt_team_ids( + {"team_id": ["primary", "secondary"], "teams": ["a"]} + ) == ["a", "primary", "secondary"] + # neither populated assert jwt_handler.get_all_jwt_team_ids({}) == [] diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 4b0f6c0ea2f..a6483403c27 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -587,6 +587,14 @@ def _assert_get_api_key_with_custom_litellm_key_header( False, False, ), + # Wildcard iss must not match space-containing claim strings (fnmatch * spans spaces) + ( + "trusted.*", + "trusted.example.com attacker.example.com", + False, + False, + ), + ("trusted.*", "trusted.example.com", True, False), ( ["issuer-a.example.com", "issuer-b.example.com"], "issuer-b.example.com", @@ -667,7 +675,10 @@ def test_routing_selector_matches_claim_parametrized( aud=["api://litellm", "api://fallback"], path="oauth2", ), - {"iss": "oauth-issuer.example.com", "aud": ["api://other", "api://litellm"]}, + { + "iss": "oauth-issuer.example.com", + "aud": ["api://other", "api://litellm"], + }, True, ), # All provided selectors are AND-ed. @@ -1738,20 +1749,21 @@ class TestJWTOAuth2Coexistence: mock_request.headers = {"authorization": f"Bearer {jwt_token}"} mock_request.query_params = {} - with patch( - "litellm.proxy.proxy_server.general_settings", general_settings - ), patch("litellm.proxy.proxy_server.premium_user", True), patch( - "litellm.proxy.proxy_server.master_key", "sk-master" - ), patch( - "litellm.proxy.proxy_server.prisma_client", None - ), patch( - "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", - new_callable=AsyncMock, - return_value=mock_oauth2_response, - ) as mock_oauth2, patch( - "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", - new_callable=AsyncMock, - ) as mock_jwt_auth: + with ( + patch("litellm.proxy.proxy_server.general_settings", general_settings), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch( + "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", + new_callable=AsyncMock, + return_value=mock_oauth2_response, + ) as mock_oauth2, + patch( + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + ) as mock_jwt_auth, + ): litellm.proxy.proxy_server.jwt_handler.update_environment( prisma_client=None, user_api_key_cache=DualCache(), @@ -1809,20 +1821,21 @@ class TestJWTOAuth2Coexistence: mock_request.headers = {"authorization": f"Bearer {jwt_token}"} mock_request.query_params = {} - with patch( - "litellm.proxy.proxy_server.general_settings", general_settings - ), patch("litellm.proxy.proxy_server.premium_user", True), patch( - "litellm.proxy.proxy_server.master_key", "sk-master" - ), patch( - "litellm.proxy.proxy_server.prisma_client", None - ), patch( - "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", - new_callable=AsyncMock, - ) as mock_oauth2, patch( - "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", - new_callable=AsyncMock, - return_value=mock_jwt_result, - ) as mock_jwt_auth: + with ( + patch("litellm.proxy.proxy_server.general_settings", general_settings), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch( + "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", + new_callable=AsyncMock, + ) as mock_oauth2, + patch( + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + return_value=mock_jwt_result, + ) as mock_jwt_auth, + ): litellm.proxy.proxy_server.jwt_handler.update_environment( prisma_client=None, user_api_key_cache=DualCache(), @@ -1873,20 +1886,21 @@ class TestJWTOAuth2Coexistence: mock_request.headers = {"authorization": f"Bearer {jwt_token}"} mock_request.query_params = {} - with patch( - "litellm.proxy.proxy_server.general_settings", general_settings - ), patch("litellm.proxy.proxy_server.premium_user", True), patch( - "litellm.proxy.proxy_server.master_key", "sk-master" - ), patch( - "litellm.proxy.proxy_server.prisma_client", None - ), patch( - "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", - new_callable=AsyncMock, - return_value=mock_oauth2_response, - ) as mock_oauth2, patch( - "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", - new_callable=AsyncMock, - ) as mock_jwt_auth: + with ( + patch("litellm.proxy.proxy_server.general_settings", general_settings), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch( + "litellm.proxy.auth.user_api_key_auth.Oauth2Handler.check_oauth2_token", + new_callable=AsyncMock, + return_value=mock_oauth2_response, + ) as mock_oauth2, + patch( + "litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder", + new_callable=AsyncMock, + ) as mock_jwt_auth, + ): litellm.proxy.proxy_server.jwt_handler.update_environment( prisma_client=None, user_api_key_cache=DualCache(),