[litellm-agent] Staging → litellm_internal_staging (5/7/2026)

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -12,7 +12,8 @@
 )
 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 @@
 )
 
 
-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)

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 ``<proxy>/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")

litellm/proxy/auth/handle_jwt.py

@@ -224,6 +224,36 @@ def get_team_ids_from_jwt(self, token: dict) -> List[str]:
 
         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]:

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(