mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
feat(mcp): gateway DCR session admission at the aggregate /mcp endpoint (LIT-3637)
Admits a keyless SSO user (no virtual key) at the aggregate /mcp endpoint from a gateway DCR session bearer, resolving team/org/SCIM/budget authorization fresh on every call. - Aggregate DCR front door: stateless /register (sealed llm_dcrc_ client ids), SSO-backed /authorize + /authorize/complete, and /token minting identity-only session tokens with PKCE, single-use codes/flows, and rotating refresh tokens. - Admission: a session-shaped Authorization at the aggregate scope opens via _admit_gateway_session, reloads the live user, and runs the centralized policy gate; failures return the RFC 9728 invalid_token challenge. Gated on the un-forgeable, server-only mcp_admitted_user_subject marker, so virtual-key and JWT auth are unchanged. - Authorization model: an admitted subject is resolved as one plain UserAPIKeyAuth per grant source (its own grants, plus each team it is a live roster member of), each answered by the SAME resolver virtual keys use, then unioned. That branch is the FIRST statement of BOTH public resolvers, so no single-credential prelude runs for it and a fault in a lookup it never uses cannot deny its grants. A source team counts only while it is a live grantor: roster membership, not blocked, and neither the team nor its owning org over budget (enforced through the SAME _team_max_budget_check / _organization_max_budget_check owners common_checks uses for keys). Each team source carries that team's own org, so the existing org ceiling caps it; for a keyless source the org list only ever intersects (a ceiling must not become a grant) and an unresolvable ceiling denies rather than silently uncapping, on both the server and tool axes. _roster_team_object is the single owner of "which teams count": a team whose roster no longer lists the user neither grants servers nor throttles, in one place. - Rate limits: the subject is bounded by its user rpm/tpm AND by the per-server mcp_rpm_limit of the team a call is ATTRIBUTED to — the same single source billing charges, from the same owner. A key charges its one pinned team's bucket; a keyless subject has no team_id, so admission stamps each granting team's limit map onto the auth (server-only field, stripped from validated input like the marker) and the limiter emits that team's mcp_per_team descriptor. Charging every granting team instead would let one cross-team user drain several teams' SHARED buckets on a single call and block their other members; and a server the user's OWN grant reaches charges no team bucket at all, because no team provided it. Per-KEY MCP limits do not apply because there is no key. - Wrapper channels: the manager-level union treats the admitted subject by the same grant model. The admin-role short-circuit and the absolute no_mcp_servers early-return are key-credential rules and never apply to it (a session bearer is a third-party client credential, not the dashboard, and the subject's opt-out silences only its own source). Operator-open channels (allow_all_keys, the user's own BYOM submissions) are owned by one operator_open_server_ids helper that BOTH the server union and the admitted tool resolution consult (suppress-BYOM-when- explicitly-scoped is a key-credential rule and never applies to the subject, whose user row carries the DB-default empty mcp_servers), so an open-channel server is default-open for tools instead of listable but uninvokable. - Redirect URIs: one owner, validate_redirect_uri_shape, decides redirect-URI hygiene (bad scheme, fragment, missing host, userinfo, backslash host) and resolves allowlisted native callbacks, shared by DCR registration and the OAuth endpoints. Registration keeps a deliberately wider trust policy than validate_trusted_redirect_uri: public dynamic registration accepts any https client, and its controls are mandatory S256 PKCE plus the consent screen. - Egress leak-defense: a gateway admission credential (session bearer / bridge envelope) is scrubbed from EVERY egress header context, anchored to the credential shape, so it can never be forwarded upstream and replayed. - Single-use guard: auth-code, refresh and connect-flow claims resolve the proxy's cross-worker redis cache themselves rather than trusting the cache passed in, and fail CLOSED on a Redis fault instead of falling back to a per-worker count that a captured id could replay through another worker. - Sign-in return_to: one shared, never-raising helper persists a safe return_to for every sign-in branch (SSO/Okta/generic and username/password), and every branch RESUMES through the same _sso_return_to_redirect the SSO callback uses, so however a deployment signs in the stored value is honored identically (same-origin path directly; control_plane_url via the one-time login-code handoff). A stale cookie is ignored rather than failing a completed sign-in. - Budgets, both halves: ENFORCEMENT (an already over-budget team or its owning org stops being a grantor, in the source gate) and ACCOUNTING (a team-derived tool call is billed to the granting team and ITS org, so that budget accumulates and the right organization is charged). A server the user's own grant reaches bills the user; when several teams grant one server the pick is the lowest team_id, stable and auditable. Billing rides a COPY, so authorization still sees the full union, and it is inert when the target server cannot be resolved from the tool name. Deferred (tracked): client-selected server scoping of the session token (LIT-4680). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Tin Chi Lo
parent
86eee8bd05
commit
a78130461f
28 changed files with 4545 additions and 393 deletions
File diff suppressed because it is too large
Load diff
|
|
@ -33,6 +33,7 @@ from litellm.proxy._experimental.mcp_server.bridge_token_flow import (
|
|||
_finish_bridge_mint,
|
||||
_prepare_bridge_mint,
|
||||
_prepare_bridge_refresh,
|
||||
_reload_active_user_by_id,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.faults import (
|
||||
CallerRejected,
|
||||
|
|
@ -43,6 +44,14 @@ from litellm.proxy._experimental.mcp_server.faults import (
|
|||
dcr_fault_detail,
|
||||
render_token_fault,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
|
||||
aggregate_authorize,
|
||||
aggregate_token,
|
||||
complete_connect_flow,
|
||||
is_gateway_dcr_client_id,
|
||||
register_aggregate_client,
|
||||
relative_request_url,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import (
|
||||
TOKEN_NO_CACHE_HEADERS,
|
||||
get_request_base_url,
|
||||
|
|
@ -324,14 +333,25 @@ def redeem_passthrough_authorization_code(
|
|||
return sealed
|
||||
|
||||
|
||||
def _session_cookie_user_id(request: Request) -> str | None:
|
||||
"""The signed-in litellm user for a browser request, or ``None``. Thin wrapper so the
|
||||
aggregate DCR flow's verbs receive the identity as a plain value instead of parsing
|
||||
cookies themselves."""
|
||||
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( # noqa: PLC0415 # circular import at module load
|
||||
_user_id_from_session_cookie,
|
||||
)
|
||||
|
||||
return _user_id_from_session_cookie(request)
|
||||
|
||||
|
||||
def _redirect_to_litellm_login(request: Request) -> RedirectResponse:
|
||||
"""Send an unauthenticated browser through litellm login before the interactive bridge authorize
|
||||
can capture its identity. The bridge oauth_delegate flow seals the SSO user into the gateway code,
|
||||
so a session is required; without one there is nothing to bind. After login the user re-initiates
|
||||
the connection, which then finds the session cookie (the seamless return-to round-trip, which is
|
||||
origin-validated against the control-plane URL, is a follow-up)."""
|
||||
so a session is required; without one there is nothing to bind. A same-origin relative
|
||||
``return_to`` (honored by the SSO callback) brings the browser straight back to this authorize
|
||||
request after login instead of stranding it on the dashboard."""
|
||||
base_url = get_request_base_url(request)
|
||||
return RedirectResponse(f"{base_url}/sso/key/generate")
|
||||
return RedirectResponse(f"{base_url}/sso/key/generate?{urlencode({'return_to': relative_request_url(request)})}")
|
||||
|
||||
|
||||
# LIT-4197: some upstream authorization servers reject an over-long ``state``
|
||||
|
|
@ -1601,6 +1621,18 @@ async def authorize(
|
|||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
if mcp_server_name is None and client_id and is_gateway_dcr_client_id(client_id):
|
||||
return aggregate_authorize(
|
||||
request=request,
|
||||
client_id=client_id,
|
||||
redirect_uri=redirect_uri,
|
||||
state=state,
|
||||
code_challenge=code_challenge,
|
||||
code_challenge_method=code_challenge_method,
|
||||
response_type=response_type,
|
||||
session_user_id=_session_cookie_user_id(request),
|
||||
)
|
||||
|
||||
lookup_name: Optional[str] = mcp_server_name or client_id
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
mcp_server = (
|
||||
|
|
@ -1664,6 +1696,25 @@ async def token_endpoint(
|
|||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
if mcp_server_name is None and is_gateway_dcr_client_id(client_id):
|
||||
from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load
|
||||
master_key,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
return await aggregate_token(
|
||||
request=request,
|
||||
grant_type=grant_type,
|
||||
code=code,
|
||||
redirect_uri=redirect_uri,
|
||||
client_id=client_id,
|
||||
code_verifier=code_verifier,
|
||||
refresh_token=refresh_token,
|
||||
master_key=master_key,
|
||||
reload_user=_reload_active_user_by_id,
|
||||
cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
lookup_name = mcp_server_name or client_id
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(lookup_name, client_ip=client_ip)
|
||||
|
|
@ -1685,6 +1736,21 @@ async def token_endpoint(
|
|||
)
|
||||
|
||||
|
||||
@router.post("/authorize/complete")
|
||||
async def authorize_complete(request: Request, flow: str = Form(...)):
|
||||
"""Finish an aggregate connect flow: mint the gateway authorization code for the
|
||||
signed-in user and redirect back to the DCR client. POST plus the per-flow HttpOnly
|
||||
cookie set at /authorize; an anonymous or bad-flow request just 400s."""
|
||||
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 # circular import at module load
|
||||
|
||||
return await complete_connect_flow(
|
||||
request=request,
|
||||
flow_handle=flow,
|
||||
session_user_id=_session_cookie_user_id(request),
|
||||
cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
|
||||
# Per RFC 6749 §4.1.2.1, an IdP that rejects an OAuth authorization request
|
||||
# redirects back to the configured redirect URI with ``error`` /
|
||||
# ``error_description`` / ``error_uri`` query params and no ``code``. The MCP
|
||||
|
|
@ -2422,6 +2488,13 @@ async def register_client(request: Request, mcp_server_name: Optional[str] = Non
|
|||
}
|
||||
client_ip = IPAddressUtils.get_mcp_client_ip(request)
|
||||
if not mcp_server_name:
|
||||
# A real DCR request carries redirect_uris (RFC 7591): route it to the aggregate DCR
|
||||
# endpoint the aggregate authorization-server metadata advertises. A single-server
|
||||
# deployment registers at /{server}/register instead (its bare-origin discovery
|
||||
# advertises that), so this does not affect it. A request without redirect_uris is not
|
||||
# a DCR request, so the legacy single-server-or-dummy fallback is kept for it.
|
||||
if data.get("redirect_uris"):
|
||||
return await register_aggregate_client(request=request, request_body=data)
|
||||
resolved = _resolve_oauth2_server_for_root_endpoints(client_ip=client_ip)
|
||||
if resolved:
|
||||
return await register_client_with_server(
|
||||
|
|
|
|||
637
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
Normal file
637
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
Normal file
|
|
@ -0,0 +1,637 @@
|
|||
"""The gateway-level DCR flow for the aggregate ``/mcp`` endpoint (``mcp_gateway_dcr``).
|
||||
|
||||
An OAuth-only DCR client (Claude Desktop, Claude Code, MCP Inspector) pointed at the
|
||||
aggregate ``/mcp`` endpoint discovers the gateway as its authorization server (PR 1 of
|
||||
this track) and then walks the flow implemented here:
|
||||
|
||||
1. ``POST /register``: stateless dynamic client registration. The ``client_id`` IS the
|
||||
registration: the client's redirect URIs are sealed into it with the repo's
|
||||
authenticated symmetric helper, so nothing is persisted and a forged or tampered
|
||||
client_id simply fails to open. Clients are always public (``token_endpoint_auth_method
|
||||
"none"``); PKCE S256 is what protects the code.
|
||||
2. ``GET /authorize``: validates the client and redirect URI, requires S256 PKCE, and
|
||||
interposes LiteLLM sign-in. Without a session cookie the browser is sent through
|
||||
``/sso/key/generate`` with a same-origin ``return_to`` so it lands back here after
|
||||
login. With a session, the flow parameters and the SSO user are sealed into a per-flow
|
||||
HttpOnly cookie (the same pattern as the upstream OAuth state relay) and the browser is
|
||||
sent to the connect page, where the user authorizes individual servers (vaulting those
|
||||
tokens server-side) before finishing.
|
||||
3. ``POST /authorize/complete``: the deliberate finish step. A POST (not GET) bound to the
|
||||
SameSite=Lax flow cookie, so a cross-site link cannot silently mint a code with the
|
||||
victim's session, and the signed-in user must match the user sealed into the flow.
|
||||
Mints a short-lived, single-use, gateway-sealed authorization code and redirects to the
|
||||
client's registered redirect URI.
|
||||
4. ``POST /token``: exchanges the code (PKCE-verified, client- and redirect-bound,
|
||||
single-use) for the identity-only session tokens of
|
||||
:mod:`.outbound_credentials.session_token`, re-validating that the litellm user is
|
||||
still active first; the ``refresh_token`` grant rotates the pair the same way.
|
||||
|
||||
Nothing here stores state server-side except the single-use code guard (a TTL cache
|
||||
entry). Every sealed value is authenticated encryption over the proxy salt/master key
|
||||
family, opened totally (bad input maps to an OAuth error, never a raise), and every
|
||||
identity is a stable reference re-validated live at mint, refresh, and (in the admission
|
||||
PR) tool-call time. Upstream server credentials never appear anywhere in this flow; they
|
||||
are vaulted per user by the existing ``/v1/mcp`` authorize endpoints and resolved at
|
||||
egress by user id.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import secrets
|
||||
from base64 import urlsafe_b64encode
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import Awaitable, Callable, Literal, TypeVar
|
||||
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
from fastapi.responses import JSONResponse, RedirectResponse, Response
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import (
|
||||
TOKEN_NO_CACHE_HEADERS,
|
||||
get_request_base_url,
|
||||
is_loopback_redirect_host,
|
||||
validate_redirect_uri_shape,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
|
||||
SessionRefreshOpened,
|
||||
open_session_refresh_bearer,
|
||||
session_keys_from_master_key,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import (
|
||||
SESSION_REFRESH_TTL_SECONDS,
|
||||
MintedSessionToken,
|
||||
SessionKeys,
|
||||
SessionPrincipal,
|
||||
mint_session_refresh_token,
|
||||
mint_session_token,
|
||||
)
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
decrypt_value_helper,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
|
||||
GATEWAY_DCR_CLIENT_ID_PREFIX = "llm_dcrc_"
|
||||
"""Marker prefix on every gateway-issued DCR client_id so the root authorize/token
|
||||
endpoints can route an aggregate-flow request without decrypting, and existing per-server
|
||||
flows (whose client_ids are upstream-issued) are never captured by the aggregate arm."""
|
||||
|
||||
GATEWAY_AUTH_CODE_PREFIX = "llm_gcode_"
|
||||
"""Marker prefix on the gateway-sealed authorization code, distinct from the bridge
|
||||
``llm_bcode_`` so neither flow can consume the other's codes."""
|
||||
|
||||
CONNECT_FLOW_COOKIE_PREFIX = "mcp_connect_flow_"
|
||||
"""Per-flow HttpOnly cookie holding the sealed connect flow, keyed by a short random
|
||||
handle carried in the connect-page URL (the same handle-plus-cookie pattern as the
|
||||
``mcp_oauth_state_`` upstream relay, for the same reasons: replica-safe with no
|
||||
server-side session store, and the sealed value never appears in a URL)."""
|
||||
|
||||
CONNECT_FLOW_TTL_SECONDS = 600
|
||||
GATEWAY_AUTH_CODE_TTL_SECONDS = 120
|
||||
_CLAIM_TTL_BUFFER_SECONDS = 60
|
||||
_USED_CODE_CACHE_PREFIX = "mcp_gateway_dcr_code_used:"
|
||||
_USED_FLOW_CACHE_PREFIX = "mcp_gateway_dcr_flow_used:"
|
||||
_USED_REFRESH_CACHE_PREFIX = "mcp_gateway_dcr_refresh_used:"
|
||||
|
||||
MAX_REDIRECT_URIS = 3
|
||||
MAX_REDIRECT_URI_LENGTH = 256
|
||||
MAX_CLIENT_ID_LENGTH = 2048
|
||||
"""Registration bounds. They exist to bound the sealed client_id, which rides inside
|
||||
every session-token claim set: 3 URIs of 256 bytes seal to roughly 1.2KB, comfortably
|
||||
under this cap and under the session token's own 4KB ceiling. Claude Desktop and MCP
|
||||
Inspector register one or two redirect URIs."""
|
||||
|
||||
MAX_STATE_LENGTH = 1024
|
||||
"""Bound on the client ``state`` sealed into the flow cookie and echoed on the auth-code
|
||||
redirect. An unbounded ``state`` can push the sealed cookie past the browser's ~4KB cap
|
||||
(silently dropped, breaking the flow); spec clients send a short opaque value."""
|
||||
|
||||
MIN_CODE_VERIFIER_LENGTH = 43
|
||||
MAX_CODE_VERIFIER_LENGTH = 128
|
||||
"""RFC 7636 section 4.1 bounds for the PKCE ``code_verifier``. Enforced so an out-of-range
|
||||
verifier gets a clean ``invalid_request`` instead of an opaque PKCE-mismatch."""
|
||||
|
||||
_UNPREFIXED = ""
|
||||
"""Prefix for a sealed value that carries no wire marker because it is never routed by
|
||||
prefix (the connect flow lives only in its own per-handle cookie, opened by that one
|
||||
handle). Named so the empty-string argument to ``_seal`` / ``_open_sealed`` reads as
|
||||
deliberate rather than a typo."""
|
||||
|
||||
_CLIENT_RECORD_DEBUG_KEY = "gateway_dcr_client"
|
||||
_CONNECT_FLOW_DEBUG_KEY = "gateway_connect_flow"
|
||||
_AUTH_CODE_DEBUG_KEY = "gateway_authorization_code"
|
||||
|
||||
ReloadUserFailure = Literal["unresolvable", "unavailable", "no_active_key"]
|
||||
ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]]
|
||||
"""Injected live-user revalidation (the token endpoint's mirror of admission):
|
||||
``None`` means the user is active; ``unavailable`` is a retryable DB outage; anything
|
||||
else fails the grant closed."""
|
||||
|
||||
|
||||
class GatewayDcrClient(BaseModel):
|
||||
"""The registration record sealed into a gateway DCR ``client_id``.
|
||||
|
||||
``extra="forbid"`` so a sealed value of another type (an auth code, a connect flow)
|
||||
that happened to decrypt under the shared key can never validate as a client record:
|
||||
cross-type confusion is rejected at the model boundary, not left to differing required
|
||||
fields."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
redirect_uris: tuple[str, ...] = Field(min_length=1, max_length=MAX_REDIRECT_URIS)
|
||||
iat: int
|
||||
|
||||
|
||||
class _ConnectFlow(BaseModel):
|
||||
"""One in-flight authorize: the SSO user it belongs to and the client parameters
|
||||
needed to mint the code at the finish step. Sealed into the per-flow cookie. ``jti``
|
||||
makes the flow single-use at complete; ``extra="forbid"`` rejects cross-type
|
||||
confusion."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
user_id: str = Field(min_length=1)
|
||||
client_id: str = Field(min_length=1)
|
||||
redirect_uri: str = Field(min_length=1)
|
||||
state: str
|
||||
code_challenge: str = Field(min_length=1)
|
||||
jti: str = Field(min_length=1)
|
||||
exp: int
|
||||
|
||||
|
||||
class _GatewayAuthCode(BaseModel):
|
||||
"""The gateway-sealed authorization code: the user consent it represents and the
|
||||
bindings the token endpoint must verify (client, redirect URI, PKCE challenge),
|
||||
plus a ``jti`` for the single-use guard. ``extra="forbid"`` rejects cross-type
|
||||
confusion."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
user_id: str = Field(min_length=1)
|
||||
client_id: str = Field(min_length=1)
|
||||
redirect_uri: str = Field(min_length=1)
|
||||
code_challenge: str = Field(min_length=1)
|
||||
jti: str = Field(min_length=1)
|
||||
iat: int
|
||||
exp: int
|
||||
|
||||
|
||||
def is_gateway_dcr_client_id(client_id: str | None) -> bool:
|
||||
"""Cheap prefix routing test so the root endpoints only enter the aggregate arm for
|
||||
clients this flow registered; every other client_id keeps today's behavior."""
|
||||
return client_id is not None and client_id.startswith(GATEWAY_DCR_CLIENT_ID_PREFIX)
|
||||
|
||||
|
||||
def _oauth_error(status_code: int, error: str, description: str) -> JSONResponse:
|
||||
"""RFC 6749 section 5.2 / RFC 7591 section 3.2.2 error body. Descriptions carry no
|
||||
token, code, or URL material so they are safe to relay to any client."""
|
||||
return JSONResponse(
|
||||
status_code=status_code,
|
||||
content={"error": error, "error_description": description},
|
||||
headers=TOKEN_NO_CACHE_HEADERS,
|
||||
)
|
||||
|
||||
|
||||
def _seal(prefix: str, payload: BaseModel) -> str:
|
||||
return prefix + encrypt_value_helper(payload.model_dump_json())
|
||||
|
||||
|
||||
_SealedModelT = TypeVar("_SealedModelT", bound=BaseModel)
|
||||
|
||||
|
||||
def _open_sealed(value: str, prefix: str, model: type[_SealedModelT], debug_key: str) -> _SealedModelT | None:
|
||||
"""Open a sealed value totally: anything that is not prefix-shaped, does not decrypt,
|
||||
or does not validate returns ``None`` for the caller to map onto an OAuth error."""
|
||||
if not value.startswith(prefix):
|
||||
return None
|
||||
decrypted = decrypt_value_helper(value[len(prefix) :], debug_key, return_original_value=False)
|
||||
if not isinstance(decrypted, str):
|
||||
return None
|
||||
try:
|
||||
return model.model_validate_json(decrypted)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def open_gateway_dcr_client(client_id: str) -> GatewayDcrClient | None:
|
||||
return _open_sealed(client_id, GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient, _CLIENT_RECORD_DEBUG_KEY)
|
||||
|
||||
|
||||
async def register_aggregate_client(request: Request, request_body: Mapping[str, object]) -> Response:
|
||||
"""RFC 7591 dynamic registration against the gateway itself, statelessly.
|
||||
|
||||
Only ``redirect_uris`` is authoritative; every client is registered as a public
|
||||
``token_endpoint_auth_method "none"`` client regardless of what it asked for (RFC
|
||||
7591 lets the server override metadata), because the gateway never issues client
|
||||
secrets: possession of a secret would add nothing over the mandatory S256 PKCE, and a
|
||||
stateless registration has nowhere to keep one. Nothing is persisted, so open
|
||||
registration cannot be used to fill storage.
|
||||
|
||||
Redirect-URI *hygiene* is not decided here: :func:`validate_redirect_uri_shape` is
|
||||
the single owner of that rule across the MCP OAuth surface, so allowlisted native
|
||||
callbacks (``cursor://``) are accepted and fragments, missing hosts, userinfo
|
||||
(``https://claude.ai@attacker.example/cb``) and backslash hosts are rejected exactly
|
||||
as they are on /authorize and /callback.
|
||||
|
||||
What this endpoint does decide is its own trust policy, which is deliberately wider
|
||||
than :func:`validate_trusted_redirect_uri`'s: registration is *public*, so any https
|
||||
client may register (that is what lets a hosted MCP client register at all), and the
|
||||
controls are mandatory S256 PKCE plus the consent screen showing the client origin.
|
||||
http is confined to loopback per RFC 8252 section 7.3.
|
||||
"""
|
||||
raw_uris = request_body.get("redirect_uris")
|
||||
if not isinstance(raw_uris, list) or not raw_uris or len(raw_uris) > MAX_REDIRECT_URIS:
|
||||
return _oauth_error(
|
||||
400,
|
||||
"invalid_redirect_uri",
|
||||
f"redirect_uris must be a list of 1 to {MAX_REDIRECT_URIS} URIs",
|
||||
)
|
||||
if not all(isinstance(uri, str) and len(uri) <= MAX_REDIRECT_URI_LENGTH for uri in raw_uris):
|
||||
return _oauth_error(
|
||||
400,
|
||||
"invalid_redirect_uri",
|
||||
f"each redirect URI must be a string of at most {MAX_REDIRECT_URI_LENGTH} characters",
|
||||
)
|
||||
for uri in raw_uris:
|
||||
parsed = urlparse(uri)
|
||||
try:
|
||||
if validate_redirect_uri_shape(parsed):
|
||||
continue # allowlisted native callback, e.g. cursor://
|
||||
except HTTPException as exc:
|
||||
# The shared validator speaks HTTP; RFC 7591 registration answers with an OAuth
|
||||
# error object, so translate the shape without re-deciding the rule.
|
||||
return _oauth_error(400, "invalid_redirect_uri", str(exc.detail))
|
||||
if parsed.scheme == "https" or (parsed.scheme == "http" and is_loopback_redirect_host(parsed)):
|
||||
continue
|
||||
return _oauth_error(
|
||||
400,
|
||||
"invalid_redirect_uri",
|
||||
"each redirect URI must be https, http on a loopback host, or a registered native callback",
|
||||
)
|
||||
now = datetime.now(timezone.utc)
|
||||
client_id = _seal(
|
||||
GATEWAY_DCR_CLIENT_ID_PREFIX, GatewayDcrClient(redirect_uris=tuple(raw_uris), iat=int(now.timestamp()))
|
||||
)
|
||||
if len(client_id) > MAX_CLIENT_ID_LENGTH:
|
||||
return _oauth_error(400, "invalid_client_metadata", "registered metadata is too large")
|
||||
return JSONResponse(
|
||||
status_code=201,
|
||||
content={
|
||||
"client_id": client_id,
|
||||
"client_id_issued_at": int(now.timestamp()),
|
||||
"redirect_uris": list(raw_uris),
|
||||
"token_endpoint_auth_method": "none",
|
||||
"grant_types": ["authorization_code", "refresh_token"],
|
||||
"response_types": ["code"],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _flow_cookie_name(handle: str) -> str:
|
||||
return f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}"
|
||||
|
||||
|
||||
def _cookie_path_and_secure(request: Request) -> tuple[str, bool]:
|
||||
parsed = urlparse(get_request_base_url(request))
|
||||
return parsed.path or "/", parsed.scheme == "https"
|
||||
|
||||
|
||||
def _append_query_params(url: str, params: dict[str, str]) -> str:
|
||||
parsed = urlparse(url)
|
||||
query = parse_qsl(parsed.query, keep_blank_values=True) + list(params.items())
|
||||
return urlunparse(parsed._replace(query=urlencode(query)))
|
||||
|
||||
|
||||
def relative_request_url(request: Request) -> str:
|
||||
"""The request's own path and query as a same-origin ``return_to`` target for the
|
||||
login round-trip; relative by construction, so it can never leave the gateway."""
|
||||
path = request.url.path
|
||||
return f"{path}?{request.url.query}" if request.url.query else path
|
||||
|
||||
|
||||
def aggregate_authorize(
|
||||
request: Request,
|
||||
client_id: str,
|
||||
redirect_uri: str,
|
||||
state: str,
|
||||
code_challenge: str | None,
|
||||
code_challenge_method: str | None,
|
||||
response_type: str | None,
|
||||
session_user_id: str | None,
|
||||
) -> Response:
|
||||
"""The aggregate authorize verb: validate the client, require S256 PKCE, interpose
|
||||
LiteLLM sign-in, and hand the browser to the connect page with the flow sealed into a
|
||||
per-flow cookie.
|
||||
|
||||
Validation failures respond directly with 400 and never redirect: per RFC 6749
|
||||
section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and
|
||||
once the client is at fault there is no trusted place to send the browser.
|
||||
"""
|
||||
client = open_gateway_dcr_client(client_id)
|
||||
if client is None:
|
||||
return _oauth_error(400, "invalid_client", "unknown or malformed client_id")
|
||||
if redirect_uri not in client.redirect_uris:
|
||||
return _oauth_error(400, "invalid_request", "redirect_uri is not registered for this client")
|
||||
if response_type != "code":
|
||||
return _oauth_error(400, "unsupported_response_type", "response_type must be 'code'")
|
||||
if not code_challenge or code_challenge_method != "S256":
|
||||
return _oauth_error(
|
||||
400,
|
||||
"invalid_request",
|
||||
"PKCE is required: send code_challenge with code_challenge_method=S256",
|
||||
)
|
||||
if len(state) > MAX_STATE_LENGTH:
|
||||
return _oauth_error(400, "invalid_request", f"state must be at most {MAX_STATE_LENGTH} characters")
|
||||
base_url = get_request_base_url(request)
|
||||
if session_user_id is None:
|
||||
login_url = f"{base_url}/sso/key/generate?{urlencode({'return_to': relative_request_url(request)})}"
|
||||
return RedirectResponse(login_url, status_code=303)
|
||||
now = datetime.now(timezone.utc)
|
||||
handle = secrets.token_urlsafe(24)
|
||||
flow = _ConnectFlow(
|
||||
user_id=session_user_id,
|
||||
client_id=client_id,
|
||||
redirect_uri=redirect_uri,
|
||||
state=state,
|
||||
code_challenge=code_challenge,
|
||||
jti=secrets.token_urlsafe(24),
|
||||
exp=int(now.timestamp()) + CONNECT_FLOW_TTL_SECONDS,
|
||||
)
|
||||
connect_url = _append_query_params(
|
||||
f"{base_url}/ui/chat/integrations",
|
||||
{"connect_flow": handle, "connect_client": _origin_only(redirect_uri)},
|
||||
)
|
||||
response = RedirectResponse(connect_url, status_code=303)
|
||||
path, secure = _cookie_path_and_secure(request)
|
||||
response.set_cookie(
|
||||
key=_flow_cookie_name(handle),
|
||||
value=_seal(_UNPREFIXED, flow),
|
||||
max_age=CONNECT_FLOW_TTL_SECONDS,
|
||||
path=path,
|
||||
secure=secure,
|
||||
httponly=True,
|
||||
samesite="lax",
|
||||
)
|
||||
return response
|
||||
|
||||
|
||||
def _origin_only(url: str) -> str:
|
||||
"""Scheme+host for display on the connect page; never the full redirect URI, whose
|
||||
path or query could carry values that do not belong in a page URL or logs."""
|
||||
parsed = urlparse(url)
|
||||
return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else ""
|
||||
|
||||
|
||||
async def complete_connect_flow(
|
||||
request: Request,
|
||||
flow_handle: str,
|
||||
session_user_id: str | None,
|
||||
cache: DualCache,
|
||||
) -> Response:
|
||||
"""The deliberate finish step of the connect flow: mint the gateway authorization
|
||||
code and send the browser back to the client.
|
||||
|
||||
Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly
|
||||
per-flow cookie plus an exact match between the signed-in user and the user sealed
|
||||
into the flow: a link crafted by another party dies here with ``access_denied``
|
||||
instead of minting a code for the victim's identity. The flow is single-use (an atomic
|
||||
claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in.
|
||||
"""
|
||||
sealed_flow = request.cookies.get(_flow_cookie_name(flow_handle))
|
||||
if sealed_flow is None:
|
||||
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
|
||||
flow = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY)
|
||||
if flow is None:
|
||||
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
|
||||
now = datetime.now(timezone.utc)
|
||||
if now.timestamp() >= flow.exp:
|
||||
return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection")
|
||||
if session_user_id is None:
|
||||
return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting")
|
||||
if session_user_id != flow.user_id:
|
||||
return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow")
|
||||
if not await _SingleUseGuard(cache).claim(
|
||||
f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
|
||||
):
|
||||
return _oauth_error(400, "invalid_request", "this connect flow was already completed; restart the connection")
|
||||
code = _seal(
|
||||
GATEWAY_AUTH_CODE_PREFIX,
|
||||
_GatewayAuthCode(
|
||||
user_id=flow.user_id,
|
||||
client_id=flow.client_id,
|
||||
redirect_uri=flow.redirect_uri,
|
||||
code_challenge=flow.code_challenge,
|
||||
jti=secrets.token_urlsafe(24),
|
||||
iat=int(now.timestamp()),
|
||||
exp=int(now.timestamp()) + GATEWAY_AUTH_CODE_TTL_SECONDS,
|
||||
),
|
||||
)
|
||||
params = {"code": code, **({"state": flow.state} if flow.state else {})}
|
||||
response = RedirectResponse(_append_query_params(flow.redirect_uri, params), status_code=303)
|
||||
path, secure = _cookie_path_and_secure(request)
|
||||
response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax")
|
||||
return response
|
||||
|
||||
|
||||
def _pkce_verifier_matches(code_verifier: str, code_challenge: str) -> bool:
|
||||
"""RFC 7636 S256 verification, total over hostile input. The comparison is over bytes
|
||||
so a non-ASCII ``code_challenge`` (which reaches here unvalidated from the client's
|
||||
authorize request) simply fails to match instead of raising ``TypeError`` the way
|
||||
``hmac.compare_digest`` does on two ``str`` with non-ASCII content. The verifier is
|
||||
ASCII per spec; a compliant client's challenge is base64url and matches."""
|
||||
digest = hashlib.sha256(code_verifier.encode("ascii", "replace")).digest()
|
||||
computed = urlsafe_b64encode(digest).rstrip(b"=")
|
||||
return hmac.compare_digest(computed, code_challenge.encode("utf-8"))
|
||||
|
||||
|
||||
class _SingleUseGuard:
|
||||
"""Atomic single-use claim for a one-time id (an auth-code, connect-flow ``jti``, or refresh-token
|
||||
``jti``) over the injected proxy cache.
|
||||
|
||||
Uses an atomic increment rather than a get-then-set: two concurrent redemptions of the same id
|
||||
cannot both observe "unused", because exactly one increment returns 1. The claim IS the gate, so it
|
||||
fails closed. Crucially, the increment must be recorded in a backend SHARED across replicas, or the
|
||||
single-use property is per-worker only (each replica's in-memory counter returns 1, so a captured
|
||||
id replays through a different worker):
|
||||
|
||||
- When a Redis backend is configured it is the SOLE authority: the claim goes straight to Redis
|
||||
(``INCR`` is atomic across replicas), and any Redis fault fails the claim CLOSED — it never falls
|
||||
back to the per-worker in-memory count (``DualCache.async_increment_cache`` does fall back, which
|
||||
is exactly the replay window this avoids).
|
||||
- With no Redis configured (single-replica) the in-memory increment is authoritative within the one
|
||||
process. A multi-worker deployment must run Redis for the guarantee to hold across workers.
|
||||
|
||||
The id's own TTL is the outer bound. For the auth code, PKCE binding is the primary defense against
|
||||
interception; this makes the RFC 6749 4.1.2 single-use property reliable on top of it."""
|
||||
|
||||
def __init__(self, cache: DualCache) -> None:
|
||||
self._cache = cache
|
||||
|
||||
async def claim(self, key: str, ttl_seconds: int) -> bool:
|
||||
"""Atomically claim ``key``. ``True`` iff this caller is the first (increment to 1); ``False``
|
||||
on a replay (>1) or when the claim could not be recorded in the shared backend (fail closed)."""
|
||||
from litellm.proxy.proxy_server import redis_usage_cache # noqa: PLC0415 # circular import at module load
|
||||
|
||||
# Resolve the shared authority HERE rather than trusting the injected cache: callers pass
|
||||
# user_api_key_cache, which only carries a redis_cache when enable_redis_auth_cache is set
|
||||
# (off by default), so a guard that read its injected cache silently degraded every claim to
|
||||
# a per-worker count on a stock multi-worker deployment. redis_usage_cache is the store the
|
||||
# proxy already treats as cross-worker, so no call site can wire the guarantee away.
|
||||
redis_cache = redis_usage_cache or getattr(self._cache, "redis_cache", None)
|
||||
if redis_cache is not None:
|
||||
# Shared, atomic authority for multi-replica deployments. Claim ONLY against Redis and fail
|
||||
# CLOSED on any Redis fault (async_increment re-raises) rather than fall back to the
|
||||
# per-worker in-memory count, which would let each replica observe count==1 and replay the id.
|
||||
try:
|
||||
count = await redis_cache.async_increment(key, 1, ttl=ttl_seconds)
|
||||
except Exception as e: # noqa: BLE001 # ANY Redis fault fails the single-use claim closed
|
||||
verbose_logger.warning(
|
||||
"mcp gateway single-use claim: shared cache backend unavailable, failing closed: %s", e
|
||||
)
|
||||
return False
|
||||
return count == 1
|
||||
# No shared backend configured (single-replica): the in-memory increment is authoritative.
|
||||
count = await self._cache.async_increment_cache(key, 1, ttl=ttl_seconds, local_only=True)
|
||||
return count == 1
|
||||
|
||||
|
||||
def _session_token_pair(principal: SessionPrincipal, keys: SessionKeys, now: datetime) -> Response:
|
||||
access = mint_session_token(principal, keys, now)
|
||||
refresh = mint_session_refresh_token(principal, keys, now)
|
||||
if not isinstance(access, MintedSessionToken) or not isinstance(refresh, MintedSessionToken):
|
||||
return _oauth_error(500, "server_error", "failed to mint the session credential")
|
||||
return JSONResponse(
|
||||
status_code=200,
|
||||
content={
|
||||
"access_token": access.token.get_secret_value(),
|
||||
"token_type": "Bearer",
|
||||
"expires_in": int((access.expires_at - now).total_seconds()),
|
||||
"refresh_token": refresh.token.get_secret_value(),
|
||||
},
|
||||
headers=TOKEN_NO_CACHE_HEADERS,
|
||||
)
|
||||
|
||||
|
||||
def _reload_failure_response(failure: ReloadUserFailure) -> Response:
|
||||
"""Map the live-user revalidation failure onto its OAuth error, exhaustively, so a new
|
||||
``ReloadUserFailure`` member is a type error here rather than silently 400ing."""
|
||||
match failure:
|
||||
case "unavailable":
|
||||
return _oauth_error(503, "temporarily_unavailable", "the gateway database is unavailable; retry")
|
||||
case "unresolvable":
|
||||
return _oauth_error(500, "server_error", "the gateway is not configured to resolve users")
|
||||
case "no_active_key":
|
||||
return _oauth_error(400, "invalid_grant", "the user for this grant is no longer active")
|
||||
case _:
|
||||
assert_never(failure)
|
||||
|
||||
|
||||
async def aggregate_token(
|
||||
request: Request,
|
||||
grant_type: str,
|
||||
code: str | None,
|
||||
redirect_uri: str | None,
|
||||
client_id: str,
|
||||
code_verifier: str | None,
|
||||
refresh_token: str | None,
|
||||
master_key: str | None,
|
||||
reload_user: ReloadUser,
|
||||
cache: DualCache,
|
||||
) -> Response:
|
||||
"""The aggregate token verb: authorization_code and refresh_token grants for the
|
||||
identity-only session pair. Every path re-validates the litellm user live before
|
||||
minting, so a deactivated user cannot obtain or renew a session."""
|
||||
if master_key is None:
|
||||
verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured")
|
||||
return _oauth_error(500, "server_error", "the gateway has no master key configured")
|
||||
keys = session_keys_from_master_key(master_key)
|
||||
now = datetime.now(timezone.utc)
|
||||
if grant_type == "authorization_code":
|
||||
return await _authorization_code_grant(
|
||||
code=code,
|
||||
redirect_uri=redirect_uri,
|
||||
client_id=client_id,
|
||||
code_verifier=code_verifier,
|
||||
keys=keys,
|
||||
now=now,
|
||||
reload_user=reload_user,
|
||||
guard=_SingleUseGuard(cache),
|
||||
)
|
||||
if grant_type == "refresh_token":
|
||||
return await _refresh_token_grant(
|
||||
refresh_token=refresh_token,
|
||||
client_id=client_id,
|
||||
keys=keys,
|
||||
now=now,
|
||||
reload_user=reload_user,
|
||||
guard=_SingleUseGuard(cache),
|
||||
)
|
||||
return _oauth_error(400, "unsupported_grant_type", "grant_type must be authorization_code or refresh_token")
|
||||
|
||||
|
||||
async def _authorization_code_grant(
|
||||
code: str | None,
|
||||
redirect_uri: str | None,
|
||||
client_id: str,
|
||||
code_verifier: str | None,
|
||||
keys: SessionKeys,
|
||||
now: datetime,
|
||||
reload_user: ReloadUser,
|
||||
guard: _SingleUseGuard,
|
||||
) -> Response:
|
||||
if not code or not redirect_uri or not code_verifier:
|
||||
return _oauth_error(400, "invalid_request", "code, redirect_uri, and code_verifier are required")
|
||||
if not MIN_CODE_VERIFIER_LENGTH <= len(code_verifier) <= MAX_CODE_VERIFIER_LENGTH:
|
||||
return _oauth_error(400, "invalid_request", "code_verifier must be 43 to 128 characters (RFC 7636)")
|
||||
parsed = _open_sealed(code, GATEWAY_AUTH_CODE_PREFIX, _GatewayAuthCode, _AUTH_CODE_DEBUG_KEY)
|
||||
if parsed is None:
|
||||
return _oauth_error(400, "invalid_grant", "the authorization code is invalid")
|
||||
if now.timestamp() >= parsed.exp:
|
||||
return _oauth_error(400, "invalid_grant", "the authorization code has expired")
|
||||
if client_id != parsed.client_id or redirect_uri != parsed.redirect_uri:
|
||||
return _oauth_error(400, "invalid_grant", "the authorization code was issued to a different client")
|
||||
if not _pkce_verifier_matches(code_verifier, parsed.code_challenge):
|
||||
return _oauth_error(400, "invalid_grant", "PKCE verification failed")
|
||||
# Revalidate the user BEFORE claiming the code, so a transient DB outage (a retryable
|
||||
# 503) does not consume a still-valid code and force the client to restart sign-in.
|
||||
failure = await reload_user(parsed.user_id)
|
||||
if failure is not None:
|
||||
return _reload_failure_response(failure)
|
||||
# Atomic single-use claim is the gate: on a concurrent double-redeem exactly one caller
|
||||
# wins, and a claim that cannot be recorded fails closed.
|
||||
if not await guard.claim(
|
||||
f"{_USED_CODE_CACHE_PREFIX}{parsed.jti}", GATEWAY_AUTH_CODE_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
|
||||
):
|
||||
return _oauth_error(400, "invalid_grant", "the authorization code was already used")
|
||||
return _session_token_pair(SessionPrincipal(user_id=parsed.user_id, client_id=client_id), keys, now)
|
||||
|
||||
|
||||
async def _refresh_token_grant(
|
||||
refresh_token: str | None,
|
||||
client_id: str,
|
||||
keys: SessionKeys,
|
||||
now: datetime,
|
||||
reload_user: ReloadUser,
|
||||
guard: _SingleUseGuard,
|
||||
) -> Response:
|
||||
if not refresh_token:
|
||||
return _oauth_error(400, "invalid_request", "refresh_token is required")
|
||||
opened = open_session_refresh_bearer(refresh_token, keys, now, expected_client_id=client_id)
|
||||
if not isinstance(opened, SessionRefreshOpened):
|
||||
return _oauth_error(400, "invalid_grant", "the refresh token is invalid for this client")
|
||||
failure = await reload_user(opened.principal.user_id)
|
||||
if failure is not None:
|
||||
return _reload_failure_response(failure)
|
||||
# Refresh-token rotation (OAuth 2.0 Security BCP section 4.13): the presented refresh token is
|
||||
# single-use. Claim its jti before issuing the replacement pair, so a captured or replayed
|
||||
# refresh token cannot mint a second pair after the legitimate holder rotated. Claimed AFTER
|
||||
# user revalidation so a transient DB 503 does not burn a still-valid token; a claim that
|
||||
# cannot be recorded fails closed, exactly like the authorization-code path.
|
||||
if not await guard.claim(
|
||||
f"{_USED_REFRESH_CACHE_PREFIX}{opened.jti}", SESSION_REFRESH_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
|
||||
):
|
||||
return _oauth_error(400, "invalid_grant", "the refresh token was already used")
|
||||
return _session_token_pair(opened.principal, keys, now)
|
||||
|
|
@ -49,6 +49,7 @@ from litellm.litellm_core_utils.url_utils import SSRFError, async_safe_get
|
|||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
|
||||
MCPRequestHandler,
|
||||
_is_mcp_admitted_user_subject,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.exceptions import (
|
||||
MCPServerListError,
|
||||
|
|
@ -2197,6 +2198,56 @@ class MCPServerManager:
|
|||
|
||||
return [server_id for server_id in submitted_server_ids if self.get_mcp_server_by_id(server_id) is not None]
|
||||
|
||||
async def operator_open_server_ids(
|
||||
self,
|
||||
user_api_key_auth: UserAPIKeyAuth | None = None,
|
||||
*,
|
||||
allow_all_server_ids: list[str] | None = None,
|
||||
submitted_server_ids: list[str] | None = None,
|
||||
) -> set:
|
||||
"""Servers reachable through OPEN channels rather than a grant: operator-opened
|
||||
``allow_all_keys`` servers, plus the caller's own active BYOM submissions when the caller
|
||||
carries no explicit ``mcp_servers`` scope.
|
||||
|
||||
The single owner of that question for BOTH axes. The server union in
|
||||
``get_allowed_mcp_servers`` adds these ids, and the admitted subject's tool resolution asks
|
||||
the same question to treat an open-channel server as default-open for tools — exactly how a
|
||||
virtual key experiences it. Encoding the channel membership twice is how a server ends up
|
||||
listable but uninvokable.
|
||||
|
||||
Empty inside a toolset scope: toolset_mcp_route / dynamic_mcp_route set
|
||||
``_mcp_active_toolset_id`` before calling the handler, pinning the request to the toolset's
|
||||
own servers (checking op.mcp_toolsets==[] instead would false-positive on DB-default rows
|
||||
where Postgres initialises the column to ARRAY[]::TEXT[]).
|
||||
|
||||
``allow_all_server_ids`` / ``submitted_server_ids`` are injectable so the server union,
|
||||
which precomputes both for its fallback path, does not compute them twice."""
|
||||
from litellm.proxy._experimental.mcp_server.mcp_context import ( # noqa: PLC0415
|
||||
_mcp_active_toolset_id,
|
||||
)
|
||||
|
||||
if _mcp_active_toolset_id.get() is not None:
|
||||
return set()
|
||||
if allow_all_server_ids is None:
|
||||
allow_all_server_ids = self.get_allow_all_keys_server_ids()
|
||||
open_ids = set(allow_all_server_ids)
|
||||
key_object_permission = user_api_key_auth.object_permission if user_api_key_auth else None
|
||||
# "Explicitly scoped, so do not widen with BYOM" is a rule about a CREDENTIAL that carries
|
||||
# its own mcp_servers list. It does not describe a keyless admitted subject: its
|
||||
# object_permission is the user's own row, whose mcp_servers column is [] by DB default, so
|
||||
# applying this rule would hide almost every admitted user's OWN submitted servers. Their
|
||||
# submissions are theirs by authorship, and their scope comes from the per-source union.
|
||||
has_explicit_object_permission = (
|
||||
not _is_mcp_admitted_user_subject(user_api_key_auth)
|
||||
and key_object_permission is not None
|
||||
and (key_object_permission.mcp_servers is not None)
|
||||
)
|
||||
if not has_explicit_object_permission:
|
||||
if submitted_server_ids is None:
|
||||
submitted_server_ids = await self._get_active_submitted_mcp_server_ids_for_user(user_api_key_auth)
|
||||
open_ids.update(submitted_server_ids)
|
||||
return open_ids
|
||||
|
||||
async def get_allowed_mcp_servers(self, user_api_key_auth: Optional[UserAPIKeyAuth] = None) -> list[str]:
|
||||
"""
|
||||
Get the allowed MCP Servers for the user.
|
||||
|
|
@ -2210,11 +2261,22 @@ class MCPServerManager:
|
|||
|
||||
allow_all_server_ids = self.get_allow_all_keys_server_ids()
|
||||
|
||||
# A keyless admitted subject is resolved per grant source, and channel decisions that are
|
||||
# absolute for a scoped KEY credential are not absolute for it: its own opt-out silences its
|
||||
# own source (handled per source in the resolver), never its teams' grants, and its admin
|
||||
# role does not swallow the grant model — a session bearer is a third-party client
|
||||
# credential, not the dashboard, so an admin signing in through the connect flow gets their
|
||||
# grants like anyone else rather than handing the client the full registry ahead of every
|
||||
# per-team org ceiling.
|
||||
is_admitted_subject = _is_mcp_admitted_user_subject(user_api_key_auth)
|
||||
|
||||
# The key explicitly opted out of every MCP server. Return zero before
|
||||
# layering on allow_all_keys or submitted servers so the opt-out is absolute.
|
||||
key_object_permission = user_api_key_auth.object_permission if user_api_key_auth else None
|
||||
if key_object_permission is not None and (
|
||||
SpecialMCPServerNames.no_mcp_servers.value in (key_object_permission.mcp_servers or [])
|
||||
if (
|
||||
not is_admitted_subject
|
||||
and key_object_permission is not None
|
||||
and (SpecialMCPServerNames.no_mcp_servers.value in (key_object_permission.mcp_servers or []))
|
||||
):
|
||||
return []
|
||||
|
||||
|
|
@ -2234,8 +2296,14 @@ class MCPServerManager:
|
|||
)
|
||||
|
||||
try:
|
||||
# If admin but NO explicit object permission, get all servers
|
||||
if user_api_key_auth and _user_has_admin_view(user_api_key_auth) and not has_explicit_object_permission:
|
||||
# If admin but NO explicit object permission, get all servers (never for an admitted
|
||||
# subject — see is_admitted_subject above)
|
||||
if (
|
||||
user_api_key_auth
|
||||
and not is_admitted_subject
|
||||
and _user_has_admin_view(user_api_key_auth)
|
||||
and not has_explicit_object_permission
|
||||
):
|
||||
verbose_logger.debug("Admin user without explicit object_permission - returning all servers")
|
||||
return list(self.get_registry().keys())
|
||||
|
||||
|
|
@ -2243,20 +2311,14 @@ class MCPServerManager:
|
|||
allowed_mcp_servers = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
|
||||
verbose_logger.debug(f"Allowed MCP Servers for user api key auth: {allowed_mcp_servers}")
|
||||
combined_servers = set(allowed_mcp_servers)
|
||||
# Only skip allow_all_keys servers when the request is inside a toolset
|
||||
# scope. toolset_mcp_route / dynamic_mcp_route set _mcp_active_toolset_id
|
||||
# before calling the handler — that ContextVar is the reliable signal.
|
||||
# Using op.mcp_toolsets==[] would false-positive on DB-default rows where
|
||||
# Postgres initialises the column to ARRAY[]::TEXT[].
|
||||
from litellm.proxy._experimental.mcp_server.mcp_context import ( # noqa: PLC0415
|
||||
_mcp_active_toolset_id,
|
||||
combined_servers.update(
|
||||
await self.operator_open_server_ids(
|
||||
user_api_key_auth,
|
||||
allow_all_server_ids=allow_all_server_ids,
|
||||
submitted_server_ids=submitted_server_ids,
|
||||
)
|
||||
)
|
||||
|
||||
in_toolset_scope = _mcp_active_toolset_id.get() is not None
|
||||
if not in_toolset_scope:
|
||||
combined_servers.update(allow_all_server_ids)
|
||||
combined_servers.update(submitted_server_ids)
|
||||
|
||||
# For anonymous callers (no user_id, no role), also surface any
|
||||
# servers the operator has opted into upstream-delegated auth.
|
||||
# These servers handle their own auth at the upstream level, so
|
||||
|
|
|
|||
|
|
@ -343,8 +343,36 @@ def _parse_redirect_uri_for_validation(redirect_uri: str) -> ParseResult:
|
|||
)
|
||||
|
||||
|
||||
def _validate_trusted_http_redirect_shape(parsed: ParseResult) -> bool:
|
||||
"""Return True when ``parsed`` is an allowlisted native callback (caller may return)."""
|
||||
def is_loopback_redirect_host(parsed: ParseResult) -> bool:
|
||||
"""True when the redirect host is loopback (RFC 8252 section 7.3).
|
||||
|
||||
Shared by every redirect-URI policy in the MCP OAuth surface so that none of them
|
||||
hand-rolls its own host list: a literal ``("localhost", "127.0.0.1", "::1")`` tuple
|
||||
silently misses the rest of 127.0.0.0/8 and IPv6-mapped forms.
|
||||
"""
|
||||
host = (parsed.hostname or "").lower()
|
||||
if host == "localhost":
|
||||
return True
|
||||
try:
|
||||
return ip_address(host).is_loopback
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def validate_redirect_uri_shape(parsed: ParseResult) -> bool:
|
||||
"""Validate redirect-URI *hygiene* and resolve allowlisted native callbacks.
|
||||
|
||||
Returns True when ``parsed`` is an allowlisted native callback (the caller may accept
|
||||
it outright); returns False for http/https, leaving the trust decision to the caller;
|
||||
raises for a URI that no policy should ever accept (bad scheme, fragment, missing
|
||||
host, userinfo, backslash in the host).
|
||||
|
||||
This is deliberately separate from :func:`validate_trusted_redirect_uri`, which adds
|
||||
the *first-party* trust policy (same-origin, loopback, ops allowlist) appropriate to
|
||||
the proxy's own OAuth endpoints. Public dynamic-client registration accepts any https
|
||||
client and relies on PKCE plus the consent screen instead, so it shares this hygiene
|
||||
rule but not that trust policy.
|
||||
"""
|
||||
if parsed.scheme not in ("http", "https"):
|
||||
if _matches_trusted_native_redirect_uri(parsed):
|
||||
return True
|
||||
|
|
@ -396,14 +424,8 @@ def _trusted_redirect_uri_is_allowed(
|
|||
):
|
||||
return True
|
||||
|
||||
host = (parsed.hostname or "").lower()
|
||||
if host == "localhost":
|
||||
if is_loopback_redirect_host(parsed):
|
||||
return True
|
||||
try:
|
||||
if ip_address(host).is_loopback:
|
||||
return True
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
if parsed.scheme == "https":
|
||||
for entry in _parse_trusted_redirect_origins():
|
||||
|
|
@ -522,7 +544,7 @@ def validate_trusted_redirect_uri(request: Request, redirect_uri: str) -> None:
|
|||
:func:`validate_loopback_redirect_uri`.
|
||||
"""
|
||||
parsed = _parse_redirect_uri_for_validation(redirect_uri)
|
||||
if _validate_trusted_http_redirect_shape(parsed):
|
||||
if validate_redirect_uri_shape(parsed):
|
||||
return
|
||||
redirect_netloc = _strip_default_port(parsed.scheme, parsed.netloc)
|
||||
proxy_base = _resolve_proxy_base_for_redirect(request)
|
||||
|
|
|
|||
|
|
@ -149,6 +149,7 @@ class SessionRefreshOpened(BaseModel):
|
|||
model_config = ConfigDict(frozen=True)
|
||||
tag: Literal["opened"] = "opened"
|
||||
principal: SessionPrincipal
|
||||
jti: str
|
||||
|
||||
|
||||
class SessionRefreshInvalid(BaseModel):
|
||||
|
|
@ -187,4 +188,4 @@ def open_session_refresh_bearer(
|
|||
return SessionRefreshInvalid()
|
||||
if opened.principal.client_id != expected_client_id:
|
||||
return SessionRefreshInvalid()
|
||||
return SessionRefreshOpened(principal=opened.principal)
|
||||
return SessionRefreshOpened(principal=opened.principal, jti=opened.jti)
|
||||
|
|
|
|||
|
|
@ -113,10 +113,12 @@ class MintedSessionToken(BaseModel):
|
|||
|
||||
|
||||
class OpenedSessionToken(BaseModel):
|
||||
"""A validated session token of either kind: the principal it was minted for."""
|
||||
"""A validated session token of either kind: the principal it was minted for, plus the
|
||||
``jti`` so the token endpoint can enforce single-use rotation on a refresh token."""
|
||||
|
||||
model_config = ConfigDict(frozen=True)
|
||||
principal: SessionPrincipal
|
||||
jti: str
|
||||
|
||||
|
||||
class SessionTokenTooLarge(BaseModel):
|
||||
|
|
@ -320,7 +322,9 @@ def _open(
|
|||
return SessionMalformed()
|
||||
if now.timestamp() >= claims.exp:
|
||||
return SessionExpired()
|
||||
return OpenedSessionToken(principal=SessionPrincipal(user_id=claims.user_id, client_id=claims.client_id))
|
||||
return OpenedSessionToken(
|
||||
principal=SessionPrincipal(user_id=claims.user_id, client_id=claims.client_id), jti=claims.jti
|
||||
)
|
||||
|
||||
|
||||
def _decode_claims(
|
||||
|
|
|
|||
|
|
@ -977,7 +977,17 @@ if MCP_AVAILABLE:
|
|||
data = await add_litellm_data_to_request(
|
||||
data=body_data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_auth,
|
||||
# Bill a team-derived call to the team that granted it. A keyless admitted
|
||||
# subject carries no team_id, so spend skipped team updates entirely and
|
||||
# charged the user's PRIMARY org — the granting team's budget never
|
||||
# accumulated (so it could never begin to block) and, cross-org, the wrong
|
||||
# organization was charged. This is the ACCOUNTING half; the enforcement
|
||||
# half (an already-over-budget team stops granting) lives in the source gate.
|
||||
# Authorization is unaffected: it ran before this, and the union is resolved
|
||||
# from the untouched auth object passed to call_mcp_tool below.
|
||||
user_api_key_dict=await MCPRequestHandler.billing_auth_for_tool_call(
|
||||
user_api_key_auth, tool_name=name
|
||||
),
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -2605,6 +2605,17 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
user_max_budget: Optional[float] = None
|
||||
request_route: Optional[str] = None
|
||||
is_session_token: bool = False
|
||||
# Server-only marker set exclusively by the MCP gateway admission path
|
||||
# (_reload_admitted_user) for a keyless user-subject admitted via a gateway DCR session
|
||||
# bearer or bridge envelope. Not a DB column and never populated from caller-controlled key
|
||||
# metadata or JWT claims, so it cannot be forged to gain the team-inherited MCP grant union
|
||||
# or to escape the caller-Authorization egress scrub. exclude=True keeps it out of serialization.
|
||||
mcp_admitted_user_subject: bool = Field(default=False, exclude=True)
|
||||
# team_id -> that team's mcp_rpm_limit map, for a keyless admitted subject that reaches MCP
|
||||
# servers through several teams at once and therefore has no single team_id for the limiter to
|
||||
# key off. Server-only and stripped from validated input for the same reason as the marker
|
||||
# above: a forged entry would let a caller pick which team's rpm bucket it is charged against.
|
||||
mcp_source_team_rpm_limits: dict[str, dict[str, int]] | None = Field(default=None, exclude=True)
|
||||
budget_reservation: Optional[Dict[str, Any]] = Field(default=None, exclude=True)
|
||||
budget_throttle_pct: Optional[float] = Field(default=None, exclude=True)
|
||||
user: Optional[Any] = None # Expanded user object when expand=user is used
|
||||
|
|
@ -2625,6 +2636,11 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
|
|||
# If values is already an instance (not a dict), return it as-is
|
||||
if not isinstance(values, dict):
|
||||
return values
|
||||
# mcp_admitted_user_subject is a server-only marker, set ONLY by the MCP gateway admission
|
||||
# path via post-construction assignment. Strip it from any validated input (constructor
|
||||
# kwargs, model_validate, a JWT/key claim splat) so it can never be forged from caller data.
|
||||
values.pop("mcp_admitted_user_subject", None)
|
||||
values.pop("mcp_source_team_rpm_limits", None)
|
||||
if values.get("api_key") is not None:
|
||||
values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))})
|
||||
if isinstance(values.get("api_key"), str):
|
||||
|
|
|
|||
|
|
@ -2661,6 +2661,15 @@ async def get_managed_vector_store_rows_by_uuids(
|
|||
return result
|
||||
|
||||
|
||||
class OrganizationNotFoundError(Exception):
|
||||
"""The organization row is CONFIRMED absent, as opposed to a lookup that failed.
|
||||
|
||||
Subclasses Exception so every existing except Exception caller keeps its current
|
||||
behavior; it exists so a caller that wants to treat "no such org" as "no restriction" can do
|
||||
that WITHOUT also swallowing an outage and silently dropping a real org ceiling.
|
||||
"""
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def get_org_object(
|
||||
org_id: str,
|
||||
|
|
@ -2707,25 +2716,30 @@ async def get_org_object(
|
|||
query_kwargs["include"] = {"litellm_budget_table": True}
|
||||
|
||||
response = await OrganizationRepository(prisma_client).table.find_unique(**query_kwargs)
|
||||
|
||||
if response is None:
|
||||
raise Exception
|
||||
|
||||
_org_obj = LiteLLM_OrganizationTable(**response.model_dump())
|
||||
# Cache the result
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=_org_obj,
|
||||
model_type=LiteLLM_OrganizationTable,
|
||||
ttl=DEFAULT_IN_MEMORY_TTL,
|
||||
)
|
||||
|
||||
return _org_obj
|
||||
except Exception:
|
||||
raise Exception(
|
||||
# An operational failure (DB down, timeout, cache fault) is NOT the same fact as a confirmed
|
||||
# missing row, and relabelling it as "doesn't exist" made every caller unable to tell them
|
||||
# apart — a caller that treats absence as "this org places no restriction" then drops a real
|
||||
# org ceiling during an outage. Propagate the real error; callers that already catch
|
||||
# Exception are unaffected.
|
||||
raise
|
||||
|
||||
if response is None:
|
||||
raise OrganizationNotFoundError(
|
||||
f"Organization doesn't exist in db. Organization={org_id}. Create organization via `/organization/new` call."
|
||||
)
|
||||
|
||||
_org_obj = LiteLLM_OrganizationTable(**response.model_dump())
|
||||
# Cache the result
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=cache_key,
|
||||
value=_org_obj,
|
||||
model_type=LiteLLM_OrganizationTable,
|
||||
ttl=DEFAULT_IN_MEMORY_TTL,
|
||||
)
|
||||
|
||||
return _org_obj
|
||||
|
||||
|
||||
async def _get_resources_from_access_groups(
|
||||
access_group_ids: List[str],
|
||||
|
|
|
|||
|
|
@ -7,12 +7,15 @@ login endpoints (e.g., /login and /v2/login).
|
|||
|
||||
import os
|
||||
import secrets
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Literal, Optional, cast
|
||||
|
||||
import jwt
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm.constants import LITELLM_PROXY_ADMIN_NAME, LITELLM_UI_SESSION_DURATION
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_UserTable,
|
||||
LitellmUserRoles,
|
||||
|
|
@ -313,6 +316,29 @@ async def authenticate_user(
|
|||
)
|
||||
|
||||
|
||||
def _ui_session_exp_timestamp() -> int:
|
||||
"""The ``exp`` claim (unix seconds) for a UI session cookie, ``LITELLM_UI_SESSION_DURATION``
|
||||
from now. The virtual key sealed inside the cookie already expires after this same
|
||||
duration; stamping the JWT itself gives the cookie the bounded lifetime the dashboard's
|
||||
client-side expiry check and the server-side session-cookie readers both assume, instead
|
||||
of a token that stays signature-valid until the master key rotates."""
|
||||
ttl_seconds = duration_in_seconds(LITELLM_UI_SESSION_DURATION)
|
||||
return int((datetime.now(timezone.utc) + timedelta(seconds=ttl_seconds)).timestamp())
|
||||
|
||||
|
||||
def encode_ui_session_jwt(returned_ui_token_object: ReturnedUITokenObject, master_key: str) -> str:
|
||||
"""Encode a UI session cookie JWT with a bounded ``exp``.
|
||||
|
||||
The single choke point every UI login path (SSO and username/password /login, /v2,
|
||||
/v3) uses to mint the ``token`` cookie, so the cookie's lifetime is set in exactly one
|
||||
place and cannot drift between paths. Without the ``exp`` the cookie is valid until the
|
||||
master key rotates, and the session-cookie readers that require a bounded lifetime
|
||||
(the MCP interactive sign-in) reject it.
|
||||
"""
|
||||
claims = {**cast(dict, returned_ui_token_object), "exp": _ui_session_exp_timestamp()}
|
||||
return jwt.encode(claims, master_key, algorithm="HS256")
|
||||
|
||||
|
||||
def create_ui_token_object(
|
||||
login_result: LoginResult,
|
||||
general_settings: dict,
|
||||
|
|
|
|||
|
|
@ -1781,28 +1781,38 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
"""
|
||||
from litellm.proxy.auth.auth_utils import get_team_mcp_rpm_limit
|
||||
|
||||
if not mcp_server_name or not user_api_key_dict.team_id:
|
||||
if not mcp_server_name:
|
||||
return
|
||||
|
||||
mcp_rpm_limit = get_team_mcp_rpm_limit(user_api_key_dict)
|
||||
if not mcp_rpm_limit:
|
||||
return
|
||||
# Which teams' buckets does this call charge? A key is pinned to exactly one team. A keyless
|
||||
# MCP-admitted subject reaches servers through SEVERAL teams at once and has no team_id, so
|
||||
# without the second source below its calls charged no team bucket at all and it outran every
|
||||
# team's mcp_rpm_limit. Every applicable team is charged rather than one being picked: the
|
||||
# limiter enforces all descriptors, so each team's own ceiling binds on a call made through
|
||||
# its grant, and there is no arbitrary attribution when several teams grant the same server.
|
||||
team_limits: list[tuple[str | None, dict[str, int] | None]] = []
|
||||
if user_api_key_dict.team_id:
|
||||
team_limits.append((user_api_key_dict.team_id, get_team_mcp_rpm_limit(user_api_key_dict)))
|
||||
for source_team_id, source_limit in (user_api_key_dict.mcp_source_team_rpm_limits or {}).items():
|
||||
team_limits.append((source_team_id, source_limit))
|
||||
|
||||
server_rpm_limit = mcp_rpm_limit.get(mcp_server_name)
|
||||
if server_rpm_limit is None:
|
||||
return
|
||||
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="mcp_per_team",
|
||||
value=f"{user_api_key_dict.team_id}:{mcp_server_name}",
|
||||
rate_limit={
|
||||
"requests_per_unit": server_rpm_limit,
|
||||
"tokens_per_unit": None,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
for team_id, mcp_rpm_limit in team_limits:
|
||||
if not team_id or not mcp_rpm_limit:
|
||||
continue
|
||||
server_rpm_limit = mcp_rpm_limit.get(mcp_server_name)
|
||||
if server_rpm_limit is None:
|
||||
continue
|
||||
descriptors.append(
|
||||
RateLimitDescriptor(
|
||||
key="mcp_per_team",
|
||||
value=f"{team_id}:{mcp_server_name}",
|
||||
rate_limit={
|
||||
"requests_per_unit": server_rpm_limit,
|
||||
"tokens_per_unit": None,
|
||||
"window_size": self.window_size,
|
||||
},
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
def _should_enforce_rate_limit(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -36,7 +36,7 @@ if TYPE_CHECKING:
|
|||
import httpx
|
||||
|
||||
import jwt
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
|
||||
from fastapi.responses import RedirectResponse
|
||||
|
||||
import litellm
|
||||
|
|
@ -965,15 +965,8 @@ async def google_login(
|
|||
state=cli_state,
|
||||
request=request,
|
||||
)
|
||||
if return_to is not None and sso_redirect is not None:
|
||||
if SSOAuthenticationHandler._validate_return_to(return_to):
|
||||
sso_redirect.set_cookie(
|
||||
key="litellm_cp_return_to",
|
||||
value=return_to,
|
||||
max_age=600,
|
||||
httponly=True,
|
||||
samesite="lax",
|
||||
)
|
||||
if sso_redirect is not None:
|
||||
_persist_return_to_cookie(sso_redirect, return_to)
|
||||
return sso_redirect
|
||||
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
|
@ -982,13 +975,19 @@ async def google_login(
|
|||
os.getenv("LITELLM_HIDE_DEFAULT_CREDENTIALS_HINT", "false").lower() == "true"
|
||||
or general_settings.get("hide_default_credentials_hint", False) is True
|
||||
)
|
||||
return HTMLResponse(
|
||||
form_response = HTMLResponse(
|
||||
content=build_ui_login_form(
|
||||
show_deprecation_banner=True,
|
||||
hide_default_credentials_hint=hide_default_credentials_hint,
|
||||
),
|
||||
status_code=200,
|
||||
)
|
||||
# Preserve return_to across the username/password sign-in too, via the SAME shared, never-raising
|
||||
# helper the SSO branch uses, so /login can resume the connect flow instead of dead-ending at the
|
||||
# dashboard. One implementation → the two sign-in branches cannot diverge (and the login form always
|
||||
# renders, since the helper never raises on a bad return_to).
|
||||
_persist_return_to_cookie(form_response, return_to)
|
||||
return form_response
|
||||
|
||||
|
||||
def generic_response_convertor(
|
||||
|
|
@ -2418,6 +2417,92 @@ async def sso_readiness():
|
|||
)
|
||||
|
||||
|
||||
def _is_same_origin_return_path(return_to: str) -> bool:
|
||||
"""True for a strictly relative return path that stays on the gateway's own origin by
|
||||
construction, and is therefore safe to honor without a configured ``control_plane_url``.
|
||||
Used by the MCP gateway DCR authorize round-trip so a browser sent through login lands
|
||||
back on the authorize request.
|
||||
|
||||
Requires a single leading ``/`` (not protocol-relative ``//``), no backslash (browsers
|
||||
fold ``\\`` to ``/``, so ``/\\evil.com`` would escape the origin), and no control or
|
||||
whitespace characters. Rejecting control chars keeps a ``\\r\\n``/tab-bearing value out
|
||||
of the redirect ``Location`` and the ``litellm_cp_return_to`` cookie entirely, rather
|
||||
than relying on downstream header encoding to neutralize it."""
|
||||
if not return_to.startswith("/") or return_to.startswith("//") or "\\" in return_to:
|
||||
return False
|
||||
return not any(ord(ch) < 0x20 or ch in (" ", "\x7f") for ch in return_to)
|
||||
|
||||
|
||||
async def _sso_return_to_redirect(
|
||||
return_to: str | None,
|
||||
jwt_token: str,
|
||||
redis_usage_cache,
|
||||
user_api_key_cache,
|
||||
) -> RedirectResponse | None:
|
||||
"""Resolve the post-SSO redirect for a ``return_to``, or None to fall through to the dashboard.
|
||||
|
||||
Two arms, both clearing the one-shot ``litellm_cp_return_to`` cookie:
|
||||
- **Same-origin relative path** (the MCP gateway DCR authorize round-trip): set the session cookie
|
||||
exactly like the dashboard path, then send the browser back where it came from.
|
||||
- **Control-plane cross-origin** (``control_plane_url``): stash the JWT behind a single-use opaque
|
||||
code (60s TTL) so the token never lands in browser history/logs; the control plane redeems it via
|
||||
``POST /v3/login/exchange``.
|
||||
|
||||
Extracted from ``get_redirect_response_from_openid`` to keep that method inside the complexity
|
||||
budget; behavior is identical to the inline arms it replaces (including letting
|
||||
``_validate_return_to`` raise for a mismatched absolute return_to, as before)."""
|
||||
if return_to is None:
|
||||
return None
|
||||
|
||||
if _is_same_origin_return_path(return_to):
|
||||
redirect_response = RedirectResponse(url=return_to, status_code=303)
|
||||
redirect_response.set_cookie(key="token", value=jwt_token)
|
||||
redirect_response.delete_cookie("litellm_cp_return_to")
|
||||
return redirect_response
|
||||
|
||||
if SSOAuthenticationHandler._validate_return_to(return_to):
|
||||
code = secrets.token_urlsafe(32)
|
||||
cache_key = f"login_code:{code}"
|
||||
cache_value = {"token": jwt_token, "redirect_url": return_to}
|
||||
if redis_usage_cache is not None:
|
||||
await redis_usage_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60)
|
||||
else:
|
||||
await user_api_key_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60)
|
||||
|
||||
separator = "&" if "?" in return_to else "?"
|
||||
redirect_url = return_to + separator + urlencode({"login": "success", "code": code})
|
||||
verbose_proxy_logger.info("Cross-origin SSO: redirecting to control plane with login code")
|
||||
redirect_response = RedirectResponse(url=redirect_url, status_code=303)
|
||||
redirect_response.delete_cookie("litellm_cp_return_to")
|
||||
return redirect_response
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _persist_return_to_cookie(response: Response, return_to: str | None) -> None:
|
||||
"""Best-effort: persist a SAFE ``return_to`` on ``response`` as the one-shot ``litellm_cp_return_to``
|
||||
cookie so ANY sign-in path — SSO / Okta / generic OR the username/password form — can resume there
|
||||
afterwards. THIS is the single source of truth, called by every sign-in branch so they cannot
|
||||
diverge (a per-branch reimplementation is exactly how the two drifted before). Honors a strictly
|
||||
relative same-origin path, and (when ``control_plane_url`` is configured) a return_to matching that
|
||||
origin. It NEVER raises: a mismatched or invalid ``return_to`` is simply not stored, so it can never
|
||||
block sign-in — the login entrypoint must always render."""
|
||||
if return_to is None:
|
||||
return
|
||||
try:
|
||||
safe = _is_same_origin_return_path(return_to) or SSOAuthenticationHandler._validate_return_to(return_to)
|
||||
except HTTPException:
|
||||
return # a non-matching absolute return_to is ignored, never blocks sign-in
|
||||
if safe:
|
||||
response.set_cookie(
|
||||
key="litellm_cp_return_to",
|
||||
value=return_to,
|
||||
max_age=600,
|
||||
httponly=True,
|
||||
samesite="lax",
|
||||
)
|
||||
|
||||
|
||||
class SSOAuthenticationHandler:
|
||||
"""
|
||||
Handler for SSO Authentication across all SSO providers
|
||||
|
|
@ -3055,7 +3140,6 @@ class SSOAuthenticationHandler:
|
|||
return_to: Optional[str] = None,
|
||||
sso_assertion: SSOIdentityAssertion | None = None,
|
||||
) -> RedirectResponse:
|
||||
import jwt
|
||||
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings,
|
||||
|
|
@ -3219,30 +3303,21 @@ class SSOAuthenticationHandler:
|
|||
server_root_path=get_server_root_path(),
|
||||
)
|
||||
|
||||
jwt_token = jwt.encode(
|
||||
cast(dict, returned_ui_token_object),
|
||||
master_key or "",
|
||||
algorithm="HS256",
|
||||
from litellm.proxy.auth.login_utils import encode_ui_session_jwt
|
||||
|
||||
jwt_token = encode_ui_session_jwt(returned_ui_token_object, master_key or "")
|
||||
|
||||
# Post-SSO return_to handling (the same-origin DCR round-trip and the control-plane
|
||||
# cross-origin code exchange) lives in one shared helper so this method stays inside the
|
||||
# complexity budget. None falls through to the dashboard redirect below.
|
||||
return_to_redirect = await _sso_return_to_redirect(
|
||||
return_to=return_to,
|
||||
jwt_token=jwt_token,
|
||||
redis_usage_cache=redis_usage_cache,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
# Control-plane cross-origin: store JWT behind a single-use opaque
|
||||
# code (60s TTL) so the token never appears in browser history / logs.
|
||||
# The control plane redeems it via POST /v3/login/exchange.
|
||||
if return_to is not None and SSOAuthenticationHandler._validate_return_to(return_to):
|
||||
code = secrets.token_urlsafe(32)
|
||||
cache_key = f"login_code:{code}"
|
||||
cache_value = {"token": jwt_token, "redirect_url": return_to}
|
||||
if redis_usage_cache is not None:
|
||||
await redis_usage_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60)
|
||||
else:
|
||||
await user_api_key_cache.async_set_cache(key=cache_key, value=cache_value, ttl=60)
|
||||
|
||||
separator = "&" if "?" in return_to else "?"
|
||||
redirect_url = return_to + separator + urlencode({"login": "success", "code": code})
|
||||
verbose_proxy_logger.info("Cross-origin SSO: redirecting to control plane with login code")
|
||||
redirect_response = RedirectResponse(url=redirect_url, status_code=303)
|
||||
redirect_response.delete_cookie("litellm_cp_return_to")
|
||||
return redirect_response
|
||||
if return_to_redirect is not None:
|
||||
return return_to_redirect
|
||||
|
||||
if user_id is not None and isinstance(user_id, str):
|
||||
litellm_dashboard_ui += "?login=success"
|
||||
|
|
|
|||
|
|
@ -13472,7 +13472,7 @@ async def fallback_login(request: Request):
|
|||
@router.post("/login", include_in_schema=False) # hidden since this is a helper for UI sso login
|
||||
async def login(request: Request):
|
||||
global premium_user, general_settings, master_key
|
||||
from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object
|
||||
from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt
|
||||
from litellm.proxy.utils import get_custom_url
|
||||
|
||||
form = await request.form()
|
||||
|
|
@ -13495,13 +13495,7 @@ async def login(request: Request):
|
|||
)
|
||||
|
||||
# Generate JWT token
|
||||
import jwt
|
||||
|
||||
jwt_token = jwt.encode(
|
||||
cast(dict, returned_ui_token_object),
|
||||
cast(str, master_key),
|
||||
algorithm="HS256",
|
||||
)
|
||||
jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key))
|
||||
|
||||
# Build redirect URL
|
||||
litellm_dashboard_ui = get_custom_url(str(request.base_url))
|
||||
|
|
@ -13511,16 +13505,51 @@ async def login(request: Request):
|
|||
litellm_dashboard_ui += "/ui/"
|
||||
litellm_dashboard_ui += "?login=success"
|
||||
|
||||
# Honor a same-origin return_to preserved by the sign-in page (e.g. the aggregate DCR connect flow's
|
||||
# authorize round-trip), mirroring the SSO callback; otherwise land on the dashboard. Gated by
|
||||
# _is_same_origin_return_path (strictly relative path) so it can never be an open redirect, and the
|
||||
# one-shot cookie is cleared after use.
|
||||
from litellm.proxy.management_endpoints.ui_sso import _sso_return_to_redirect
|
||||
|
||||
# Resume through the SAME resumer the SSO callback uses, rather than a second, narrower arm.
|
||||
# _persist_return_to_cookie stores both shapes it accepts (a relative same-origin path AND a
|
||||
# control_plane_url-matching absolute URL); honoring only the relative one here silently dropped
|
||||
# the control-plane case, landing the user on the dashboard. One function decides how a stored
|
||||
# return_to is honored for EVERY sign-in branch, so the write and read sets cannot diverge: it
|
||||
# sets the token cookie on the same-origin arm and hands off via a one-time login code on the
|
||||
# cross-origin arm, and clears the one-shot cookie in both.
|
||||
cp_return_to = request.cookies.get("litellm_cp_return_to")
|
||||
if cp_return_to:
|
||||
try:
|
||||
resumed = await _sso_return_to_redirect(
|
||||
return_to=cp_return_to,
|
||||
jwt_token=jwt_token,
|
||||
redis_usage_cache=redis_usage_cache,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # resuming must NEVER block a completed sign-in
|
||||
# The symmetric half of _persist_return_to_cookie's "never raises" contract. The resumer
|
||||
# rejects a return_to that no longer matches control_plane_url (a config change between
|
||||
# the cookie's write and this read), and the user has ALREADY authenticated here —
|
||||
# failing their login over a stale one-shot cookie is the worst possible outcome. Land
|
||||
# on the dashboard instead; the cookie is cleared below either way.
|
||||
verbose_proxy_logger.info("Ignoring stale litellm_cp_return_to cookie; landing on dashboard")
|
||||
resumed = None
|
||||
if resumed is not None:
|
||||
return resumed
|
||||
|
||||
# Create redirect response with cookie
|
||||
redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303)
|
||||
redirect_response.set_cookie(key="token", value=jwt_token)
|
||||
if cp_return_to:
|
||||
redirect_response.delete_cookie(key="litellm_cp_return_to")
|
||||
return redirect_response
|
||||
|
||||
|
||||
@router.post("/v2/login", include_in_schema=False) # hidden helper for UI logins via API
|
||||
async def login_v2(request: Request):
|
||||
global premium_user, general_settings, master_key
|
||||
from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object
|
||||
from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt
|
||||
from litellm.proxy.utils import get_custom_url
|
||||
|
||||
try:
|
||||
|
|
@ -13541,13 +13570,7 @@ async def login_v2(request: Request):
|
|||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
import jwt
|
||||
|
||||
jwt_token = jwt.encode(
|
||||
cast(dict, returned_ui_token_object),
|
||||
cast(str, master_key),
|
||||
algorithm="HS256",
|
||||
)
|
||||
jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key))
|
||||
|
||||
litellm_dashboard_ui = get_custom_url(str(request.base_url))
|
||||
if litellm_dashboard_ui.endswith("/"):
|
||||
|
|
@ -13591,7 +13614,7 @@ async def login_v2(request: Request):
|
|||
) # control-plane login — always returns token in body for cross-origin use
|
||||
async def login_v3(request: Request):
|
||||
global premium_user, general_settings, master_key
|
||||
from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object
|
||||
from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object, encode_ui_session_jwt
|
||||
from litellm.proxy.utils import get_custom_url
|
||||
|
||||
try:
|
||||
|
|
@ -13620,13 +13643,7 @@ async def login_v3(request: Request):
|
|||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
import jwt
|
||||
|
||||
jwt_token = jwt.encode(
|
||||
cast(dict, returned_ui_token_object),
|
||||
cast(str, master_key),
|
||||
algorithm="HS256",
|
||||
)
|
||||
jwt_token = encode_ui_session_jwt(returned_ui_token_object, cast(str, master_key))
|
||||
|
||||
litellm_dashboard_ui = get_custom_url(str(request.base_url))
|
||||
if litellm_dashboard_ui.endswith("/"):
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -3302,19 +3302,20 @@ async def test_token_root_does_not_resolve_private_server_for_external_client():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_root_resolves_single_oauth2_server():
|
||||
"""When /register is hit without server name and exactly 1 OAuth2 server exists, resolve it."""
|
||||
try:
|
||||
from fastapi import Request
|
||||
async def test_register_root_does_aggregate_dcr_not_single_server_resolution():
|
||||
"""Root /register is the aggregate DCR endpoint: it mints a stateless llm_dcrc_ client
|
||||
from the request's redirect_uris and does NOT resolve a single configured oauth2 server
|
||||
(a single-server deployment registers at /{server}/register instead)."""
|
||||
import json
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP discoverable endpoints not available")
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
global_mcp_server_manager.registry.clear()
|
||||
oauth2_server = _create_oauth2_server()
|
||||
|
|
@ -3325,33 +3326,37 @@ async def test_register_root_resolves_single_oauth2_server():
|
|||
mock_request.headers = {}
|
||||
|
||||
try:
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body",
|
||||
new=AsyncMock(return_value={}),
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body",
|
||||
new=AsyncMock(return_value={"redirect_uris": ["https://claude.ai/cb"]}),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit3637"),
|
||||
):
|
||||
result = await register_client(request=mock_request, mcp_server_name=None)
|
||||
response = await register_client(request=mock_request, mcp_server_name=None)
|
||||
|
||||
# Should resolve to the single server and return its name as client_id
|
||||
assert result["client_id"] == "test_oauth"
|
||||
assert "redirect_uris" in result
|
||||
body = json.loads(response.body)
|
||||
assert body["client_id"].startswith("llm_dcrc_")
|
||||
assert body["client_id"] != "test_oauth"
|
||||
assert body["token_endpoint_auth_method"] == "none"
|
||||
finally:
|
||||
global_mcp_server_manager.registry.clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_root_does_not_resolve_private_server_for_external_client():
|
||||
"""Root /register must not reveal or use a hidden MCP server."""
|
||||
try:
|
||||
from fastapi import Request
|
||||
async def test_register_root_does_not_leak_a_private_server():
|
||||
"""Root /register never resolves or reveals a configured server, so a private one cannot
|
||||
leak to an external caller: it always mints the aggregate DCR client instead."""
|
||||
import json
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP discoverable endpoints not available")
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
|
||||
register_client,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
global_mcp_server_manager.registry.clear()
|
||||
oauth2_server = _create_oauth2_server(available_on_public_internet=False)
|
||||
|
|
@ -3365,17 +3370,19 @@ async def test_register_root_does_not_resolve_private_server_for_external_client
|
|||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints._read_request_body",
|
||||
new=AsyncMock(return_value={}),
|
||||
new=AsyncMock(return_value={"redirect_uris": ["https://claude.ai/cb"]}),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.discoverable_endpoints.IPAddressUtils.get_mcp_client_ip",
|
||||
return_value="198.51.100.10",
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit3637"),
|
||||
):
|
||||
result = await register_client(request=mock_request, mcp_server_name=None)
|
||||
response = await register_client(request=mock_request, mcp_server_name=None)
|
||||
|
||||
assert result["client_id"] == "dummy_client"
|
||||
assert result["redirect_uris"] == ["https://llm.example.com/callback"]
|
||||
body = json.loads(response.body)
|
||||
assert body["client_id"].startswith("llm_dcrc_")
|
||||
assert "test_oauth" not in body["client_id"]
|
||||
finally:
|
||||
global_mcp_server_manager.registry.clear()
|
||||
|
||||
|
|
@ -5155,7 +5162,10 @@ async def test_bridge_refresh_grant_with_non_envelope_is_invalid_grant_before_up
|
|||
|
||||
|
||||
def _mint_test_refresh_envelope(
|
||||
server_id="bridge_srv", key_hash="hashed-litellm-key-77", upstream_refresh="UPSTREAM-REFRESH", identity=None,
|
||||
server_id="bridge_srv",
|
||||
key_hash="hashed-litellm-key-77",
|
||||
upstream_refresh="UPSTREAM-REFRESH",
|
||||
identity=None,
|
||||
scope=None,
|
||||
):
|
||||
"""Mint a refresh envelope the way the producer does, for driving the refresh_token grant in tests.
|
||||
|
|
@ -5178,7 +5188,9 @@ def _mint_test_refresh_envelope(
|
|||
keys = envelope_keys_from_master_key(_BRIDGE_MASTER_KEY)
|
||||
identity = identity if identity is not None else key_hash_identity(server_id=server_id, key_hash=key_hash)
|
||||
sealed = build_bridge_refresh_token_response(
|
||||
identity, RefreshCredential(refresh_token=SecretStr(upstream_refresh), scope=scope), keys,
|
||||
identity,
|
||||
RefreshCredential(refresh_token=SecretStr(upstream_refresh), scope=scope),
|
||||
keys,
|
||||
datetime.now(timezone.utc),
|
||||
)
|
||||
assert isinstance(sealed, SealedEnvelope)
|
||||
|
|
@ -5413,7 +5425,10 @@ async def test_bridge_refresh_re_requests_the_sealed_scope_when_client_omits_it(
|
|||
)
|
||||
captured: dict = {}
|
||||
response = await _refresh_for_bridge_server(
|
||||
server, refresh_env, {"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600}, None,
|
||||
server,
|
||||
refresh_env,
|
||||
{"access_token": "NEW-ACCESS", "token_type": "Bearer", "expires_in": 3600},
|
||||
None,
|
||||
fake_client_out=captured,
|
||||
)
|
||||
|
||||
|
|
@ -5553,7 +5568,9 @@ async def test_bridge_refresh_upstream_invalid_grant_maps_to_invalid_grant():
|
|||
error_response = MagicMock()
|
||||
error_response.status_code = 400
|
||||
error_response.text = '{"error": "invalid_grant", "error_description": "refresh token expired"}'
|
||||
error_response.json = MagicMock(return_value={"error": "invalid_grant", "error_description": "refresh token expired"})
|
||||
error_response.json = MagicMock(
|
||||
return_value={"error": "invalid_grant", "error_description": "refresh token expired"}
|
||||
)
|
||||
error_response.raise_for_status = MagicMock(
|
||||
side_effect=httpx.HTTPStatusError("bad", request=MagicMock(), response=error_response)
|
||||
)
|
||||
|
|
@ -7183,7 +7200,9 @@ def _upstream_token_response(status_code: int, *, json_body: object = None, text
|
|||
return httpx.Response(status_code, text=text_body, request=request)
|
||||
|
||||
|
||||
async def _exchange_with_upstream_response(upstream_response, *, server_client_id="web-client.apps.googleusercontent.com"):
|
||||
async def _exchange_with_upstream_response(
|
||||
upstream_response, *, server_client_id="web-client.apps.googleusercontent.com"
|
||||
):
|
||||
"""Run the raw (non-bridge) authorization_code exchange against a canned upstream token-endpoint
|
||||
response and return what the gateway would hand the client. ``server_client_id=None`` models the
|
||||
caller-supplied-credentials flow (no stored client on the server)."""
|
||||
|
|
@ -7334,9 +7353,7 @@ async def test_token_exchange_bounds_relayed_error_fields():
|
|||
async def test_token_exchange_200_without_access_token_is_502_not_keyerror():
|
||||
"""A 200 whose body has no usable access_token used to KeyError into a 500; the raw arm now
|
||||
answers 502 with the same wording as the bridge arm's no_upstream_token rejection."""
|
||||
response = await _exchange_with_upstream_response(
|
||||
_upstream_token_response(200, json_body={"token_type": "Bearer"})
|
||||
)
|
||||
response = await _exchange_with_upstream_response(_upstream_token_response(200, json_body={"token_type": "Bearer"}))
|
||||
|
||||
assert response.status_code == 502
|
||||
body = json.loads(response.body)
|
||||
|
|
@ -7357,7 +7374,9 @@ async def test_token_exchange_relays_rejection_when_http_client_raises():
|
|||
)
|
||||
raising_client = MagicMock()
|
||||
raising_client.post = AsyncMock(
|
||||
side_effect=httpx.HTTPStatusError("Client error '401 Unauthorized'", request=rejection.request, response=rejection)
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
"Client error '401 Unauthorized'", request=rejection.request, response=rejection
|
||||
)
|
||||
)
|
||||
|
||||
from fastapi import Request
|
||||
|
|
@ -7422,7 +7441,9 @@ async def test_register_relays_rejection_when_http_client_raises():
|
|||
)
|
||||
raising_client = MagicMock()
|
||||
raising_client.post = AsyncMock(
|
||||
side_effect=httpx.HTTPStatusError("Client error '400 Bad Request'", request=rejection.request, response=rejection)
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
"Client error '400 Bad Request'", request=rejection.request, response=rejection
|
||||
)
|
||||
)
|
||||
|
||||
oauth2_server = _bridge_server(auth_type=MCPAuth.oauth2, dcr_bridge=None)
|
||||
|
|
@ -7800,9 +7821,7 @@ async def test_hydrate_does_not_overwrite_explicit_config_client_id():
|
|||
auth_type=MCPAuth.oauth2,
|
||||
client_id="explicit-from-config",
|
||||
)
|
||||
store_read = AsyncMock(
|
||||
return_value={"client_id": "stale-store-client", "client_secret": "x", "redirect_uris": []}
|
||||
)
|
||||
store_read = AsyncMock(return_value={"client_id": "stale-store-client", "client_secret": "x", "redirect_uris": []})
|
||||
with (
|
||||
patch("litellm.proxy.utils.get_prisma_client_or_throw", return_value=MagicMock()),
|
||||
patch(
|
||||
|
|
@ -8019,9 +8038,7 @@ async def test_bare_origin_discovery_resolves_single_server_not_aggregate():
|
|||
mock_request.headers = {}
|
||||
|
||||
try:
|
||||
authorization_response = _build_oauth_authorization_server_response(
|
||||
request=mock_request, mcp_server_name=None
|
||||
)
|
||||
authorization_response = _build_oauth_authorization_server_response(request=mock_request, mcp_server_name=None)
|
||||
resource_response = await _build_oauth_protected_resource_response(
|
||||
request=mock_request, mcp_server_name=None, use_standard_pattern=True
|
||||
)
|
||||
|
|
@ -8033,6 +8050,65 @@ async def test_bare_origin_discovery_resolves_single_server_not_aggregate():
|
|||
global_mcp_server_manager.registry.clear()
|
||||
|
||||
|
||||
def test_gateway_dcr_flow_routing_engages_only_for_llm_dcrc_clients(monkeypatch):
|
||||
"""The aggregate DCR arms engage for llm_dcrc_ client_ids (register always mints one,
|
||||
authorize/token route into the aggregate flow); a non-gateway client_id keeps the
|
||||
per-server behavior, and /authorize/complete exists but 400s without a valid flow."""
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
||||
global_mcp_server_manager,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-for-lit3637")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-salt-for-lit3637", raising=False)
|
||||
global_mcp_server_manager.registry.clear()
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
client = TestClient(app)
|
||||
|
||||
registered = client.post("/register", json={"redirect_uris": ["https://claude.ai/cb"]})
|
||||
assert registered.status_code == 201
|
||||
assert registered.json()["client_id"].startswith("llm_dcrc_")
|
||||
assert registered.json()["token_endpoint_auth_method"] == "none"
|
||||
|
||||
authorize_params = {
|
||||
"client_id": "llm_dcrc_bogus",
|
||||
"redirect_uri": "https://claude.ai/cb",
|
||||
"response_type": "code",
|
||||
"code_challenge": "c" * 43,
|
||||
"code_challenge_method": "S256",
|
||||
}
|
||||
bogus_client = client.get("/authorize", params=authorize_params)
|
||||
assert bogus_client.status_code == 400
|
||||
assert bogus_client.json()["error"] == "invalid_client"
|
||||
|
||||
no_cookie = client.post("/authorize/complete", data={"flow": "h"})
|
||||
assert no_cookie.status_code == 400
|
||||
assert no_cookie.json()["error"] == "invalid_request"
|
||||
|
||||
token_response = client.post(
|
||||
"/token",
|
||||
data={
|
||||
"grant_type": "authorization_code",
|
||||
"client_id": "llm_dcrc_bogus",
|
||||
"code": "x",
|
||||
"redirect_uri": "https://claude.ai/cb",
|
||||
"code_verifier": "v" * 43,
|
||||
},
|
||||
)
|
||||
assert token_response.status_code == 400
|
||||
assert token_response.json()["error"] == "invalid_grant"
|
||||
|
||||
upstream_shaped = client.post(
|
||||
"/token",
|
||||
data={"grant_type": "authorization_code", "client_id": "regular-upstream-client", "code": "x"},
|
||||
)
|
||||
assert upstream_shaped.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorize_wall_names_the_fix_for_urlless_servers():
|
||||
"""LIT-4629: the authorize wall previously said only "authorization url is not set" with no
|
||||
|
|
|
|||
|
|
@ -0,0 +1,590 @@
|
|||
"""Tests for the aggregate gateway DCR flow (register, authorize, complete, token)."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from base64 import urlsafe_b64encode
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from http.cookies import SimpleCookie
|
||||
from urllib.parse import parse_qs, urlparse
|
||||
|
||||
import pytest
|
||||
from starlette.requests import Request
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
|
||||
CONNECT_FLOW_COOKIE_PREFIX,
|
||||
GATEWAY_AUTH_CODE_PREFIX,
|
||||
GATEWAY_AUTH_CODE_TTL_SECONDS,
|
||||
GATEWAY_DCR_CLIENT_ID_PREFIX,
|
||||
_GatewayAuthCode,
|
||||
_seal,
|
||||
aggregate_authorize,
|
||||
aggregate_token,
|
||||
complete_connect_flow,
|
||||
is_gateway_dcr_client_id,
|
||||
open_gateway_dcr_client,
|
||||
register_aggregate_client,
|
||||
)
|
||||
from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import (
|
||||
resolve_session_bearer,
|
||||
session_keys_from_master_key,
|
||||
SessionBearerAdmitted,
|
||||
)
|
||||
|
||||
MASTER_KEY = "sk-gateway-dcr-flow-tests"
|
||||
REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback"
|
||||
CODE_VERIFIER = "verifier-" + "v" * 43
|
||||
CODE_CHALLENGE = urlsafe_b64encode(hashlib.sha256(CODE_VERIFIER.encode("ascii")).digest()).rstrip(b"=").decode("ascii")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _salt_key(monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", MASTER_KEY)
|
||||
|
||||
|
||||
def _request(path="/authorize", query="", cookies=None, method="GET"):
|
||||
cookie_header = []
|
||||
if cookies:
|
||||
cookie = SimpleCookie()
|
||||
for name, value in cookies.items():
|
||||
cookie[name] = value
|
||||
cookie_header = [(b"cookie", cookie.output(header="", sep="; ").strip().encode())]
|
||||
return Request(
|
||||
{
|
||||
"type": "http",
|
||||
"method": method,
|
||||
"scheme": "https",
|
||||
"path": path,
|
||||
"query_string": query.encode(),
|
||||
"headers": [(b"host", b"llm.example.com"), *cookie_header],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _register(redirect_uris) -> dict:
|
||||
response = await register_aggregate_client(
|
||||
request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris}
|
||||
)
|
||||
return json.loads(response.body)
|
||||
|
||||
|
||||
async def _reload_user_active(user_id: str):
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_mints_stateless_public_client():
|
||||
body = await _register([REDIRECT_URI])
|
||||
assert body["token_endpoint_auth_method"] == "none"
|
||||
assert "client_secret" not in body
|
||||
assert body["redirect_uris"] == [REDIRECT_URI]
|
||||
assert is_gateway_dcr_client_id(body["client_id"])
|
||||
record = open_gateway_dcr_client(body["client_id"])
|
||||
assert record is not None
|
||||
assert record.redirect_uris == (REDIRECT_URI,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_allows_loopback_http_for_dev_clients():
|
||||
body = await _register(["http://localhost:6274/oauth/callback"])
|
||||
assert is_gateway_dcr_client_id(body["client_id"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"code_challenge",
|
||||
["short", "", "p" * 300, "ünïcode-challenge", "AAAA" * 20],
|
||||
)
|
||||
def test_pkce_mismatched_challenge_returns_false_never_raises(code_challenge):
|
||||
"""A wrong-length or non-ASCII code_challenge must VERIFY FALSE, not raise.
|
||||
|
||||
Pins the reason this compares bytes rather than str: hmac.compare_digest raises TypeError on
|
||||
two str with non-ASCII content, but on bytes of unequal length it simply returns False. A
|
||||
review flagged this as an unhandled 500 on length mismatch; encoding both sides to bytes is
|
||||
exactly what makes that impossible, so the claim is pinned here rather than in a comment."""
|
||||
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import _pkce_verifier_matches
|
||||
|
||||
assert _pkce_verifier_matches("a" * 43, code_challenge) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_allows_allowlisted_native_callback():
|
||||
"""Native MCP clients register a private-use scheme, not https. Registration shares
|
||||
the one redirect-URI shape owner with /authorize, so the callback the allowlist
|
||||
already trusts there is registrable here rather than rejected as non-https."""
|
||||
body = await _register(["cursor://anysphere.cursor-mcp/oauth/callback"])
|
||||
assert is_gateway_dcr_client_id(body["client_id"])
|
||||
record = open_gateway_dcr_client(body["client_id"])
|
||||
assert record is not None
|
||||
assert record.redirect_uris == ("cursor://anysphere.cursor-mcp/oauth/callback",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_register_rejects_userinfo_spoofed_origin():
|
||||
"""``https://claude.ai@attacker.example/cb`` parses with netloc
|
||||
``claude.ai@attacker.example``, so a naive origin display on the consent screen reads
|
||||
as claude.ai while the code would be delivered to attacker.example. Rejected at
|
||||
registration, which is the only way such a URI could enter a sealed client."""
|
||||
response = await register_aggregate_client(
|
||||
request=_request(path="/register", method="POST"),
|
||||
request_body={"redirect_uris": ["https://claude.ai@attacker.example/callback"]},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] == "invalid_redirect_uri"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"redirect_uris",
|
||||
[
|
||||
[],
|
||||
"not-a-list",
|
||||
["http://evil.example.com/callback"],
|
||||
["https://claude.ai/cb#fragment"],
|
||||
["ftp://claude.ai/cb"],
|
||||
["https://a.example.com/" + "p" * 300],
|
||||
["https://a.example.com/1", "https://a.example.com/2", "https://a.example.com/3", "https://a.example.com/4"],
|
||||
[12345],
|
||||
],
|
||||
)
|
||||
async def test_register_rejects_bad_redirect_uris(redirect_uris):
|
||||
response = await register_aggregate_client(
|
||||
request=_request(path="/register", method="POST"), request_body={"redirect_uris": redirect_uris}
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] in ("invalid_redirect_uri", "invalid_client_metadata")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tampered_client_id_does_not_open():
|
||||
body = await _register([REDIRECT_URI])
|
||||
tampered = body["client_id"][:-4] + "AAAA"
|
||||
assert open_gateway_dcr_client(tampered) is None
|
||||
assert open_gateway_dcr_client("llm_dcrc_garbage") is None
|
||||
assert open_gateway_dcr_client("other_prefix") is None
|
||||
|
||||
|
||||
def _authorize(
|
||||
client_id, session_user_id, redirect_uri=REDIRECT_URI, challenge=CODE_CHALLENGE, method="S256", response_type="code"
|
||||
):
|
||||
return aggregate_authorize(
|
||||
request=_request(query=f"client_id={client_id}"),
|
||||
client_id=client_id,
|
||||
redirect_uri=redirect_uri,
|
||||
state="client-state-123",
|
||||
code_challenge=challenge,
|
||||
code_challenge_method=method,
|
||||
response_type=response_type,
|
||||
session_user_id=session_user_id,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorize_validation_failures_never_redirect_to_client():
|
||||
client_id = (await _register([REDIRECT_URI]))["client_id"]
|
||||
for response, expected_error in (
|
||||
(_authorize("llm_dcrc_bogus", "u1"), "invalid_client"),
|
||||
(_authorize(client_id, "u1", redirect_uri="https://attacker.example.com/cb"), "invalid_request"),
|
||||
(_authorize(client_id, "u1", response_type="token"), "unsupported_response_type"),
|
||||
(_authorize(client_id, "u1", challenge=None), "invalid_request"),
|
||||
(_authorize(client_id, "u1", method="plain"), "invalid_request"),
|
||||
):
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] == expected_error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorize_without_session_redirects_to_login_with_return_to():
|
||||
client_id = (await _register([REDIRECT_URI]))["client_id"]
|
||||
response = _authorize(client_id, session_user_id=None)
|
||||
assert response.status_code == 303
|
||||
location = response.headers["location"]
|
||||
assert location.startswith("https://llm.example.com/sso/key/generate?return_to=")
|
||||
assert "return_to=%2Fauthorize" in location
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorize_with_session_hands_browser_to_connect_page_with_flow_cookie():
|
||||
client_id = (await _register([REDIRECT_URI]))["client_id"]
|
||||
response = _authorize(client_id, session_user_id="u1")
|
||||
assert response.status_code == 303
|
||||
location = urlparse(response.headers["location"])
|
||||
assert location.path == "/ui/chat/integrations"
|
||||
params = parse_qs(location.query)
|
||||
handle = params["connect_flow"][0]
|
||||
assert params["connect_client"] == ["https://claude.ai"]
|
||||
set_cookie = response.headers["set-cookie"]
|
||||
assert f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" in set_cookie
|
||||
assert "HttpOnly" in set_cookie
|
||||
return handle, set_cookie
|
||||
|
||||
|
||||
def _flow_cookie_from(response) -> tuple:
|
||||
location = urlparse(response.headers["location"])
|
||||
handle = parse_qs(location.query)["connect_flow"][0]
|
||||
cookie = SimpleCookie()
|
||||
cookie.load(response.headers["set-cookie"])
|
||||
name = f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}"
|
||||
return handle, {name: cookie[name].value}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_walk_register_authorize_complete_token_and_replay():
|
||||
"""The whole front door on one deterministic walk: register -> authorize ->
|
||||
complete -> token, then the security edges on the same artifacts (user mismatch,
|
||||
PKCE mismatch, single-use replay, refresh rotation, cross-client refresh)."""
|
||||
client_id = (await _register([REDIRECT_URI]))["client_id"]
|
||||
authorize_response = _authorize(client_id, session_user_id="u1")
|
||||
handle, cookies = _flow_cookie_from(authorize_response)
|
||||
|
||||
denied = await complete_connect_flow(
|
||||
request=_request("/authorize/complete", cookies=cookies, method="POST"),
|
||||
flow_handle=handle,
|
||||
session_user_id="attacker",
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert denied.status_code == 403
|
||||
|
||||
anonymous = await complete_connect_flow(
|
||||
request=_request("/authorize/complete", cookies=cookies, method="POST"),
|
||||
flow_handle=handle,
|
||||
session_user_id=None,
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert anonymous.status_code == 401
|
||||
|
||||
completed = await complete_connect_flow(
|
||||
request=_request("/authorize/complete", cookies=cookies, method="POST"),
|
||||
flow_handle=handle,
|
||||
session_user_id="u1",
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert completed.status_code == 303
|
||||
redirect = urlparse(completed.headers["location"])
|
||||
assert f"{redirect.scheme}://{redirect.netloc}{redirect.path}" == REDIRECT_URI
|
||||
params = parse_qs(redirect.query)
|
||||
assert params["state"] == ["client-state-123"]
|
||||
code = params["code"][0]
|
||||
assert code.startswith(GATEWAY_AUTH_CODE_PREFIX)
|
||||
|
||||
cache = DualCache()
|
||||
|
||||
async def _token(**overrides):
|
||||
arguments = {
|
||||
"request": _request("/token", method="POST"),
|
||||
"grant_type": "authorization_code",
|
||||
"code": code,
|
||||
"redirect_uri": REDIRECT_URI,
|
||||
"client_id": client_id,
|
||||
"code_verifier": CODE_VERIFIER,
|
||||
"refresh_token": None,
|
||||
"master_key": MASTER_KEY,
|
||||
"reload_user": _reload_user_active,
|
||||
"cache": cache,
|
||||
}
|
||||
return await aggregate_token(**{**arguments, **overrides})
|
||||
|
||||
wrong_verifier = await _token(code_verifier="wrong-" + "w" * 43)
|
||||
assert json.loads(wrong_verifier.body)["error"] == "invalid_grant"
|
||||
|
||||
wrong_client = await _token(client_id=(await _register([REDIRECT_URI]))["client_id"])
|
||||
assert json.loads(wrong_client.body)["error"] == "invalid_grant"
|
||||
|
||||
token_response = await _token()
|
||||
assert token_response.status_code == 200
|
||||
payload = json.loads(token_response.body)
|
||||
assert payload["token_type"] == "Bearer"
|
||||
assert 0 < payload["expires_in"] <= 3600
|
||||
|
||||
keys = session_keys_from_master_key(MASTER_KEY)
|
||||
admitted = resolve_session_bearer(f"Bearer {payload['access_token']}", keys, datetime.now(timezone.utc))
|
||||
assert isinstance(admitted, SessionBearerAdmitted)
|
||||
assert admitted.principal.user_id == "u1"
|
||||
assert admitted.principal.client_id == client_id
|
||||
|
||||
replay = await _token()
|
||||
assert json.loads(replay.body)["error"] == "invalid_grant"
|
||||
|
||||
refreshed = await _token(grant_type="refresh_token", code=None, refresh_token=payload["refresh_token"])
|
||||
assert refreshed.status_code == 200
|
||||
rotated = json.loads(refreshed.body)
|
||||
assert rotated["refresh_token"] != payload["refresh_token"]
|
||||
|
||||
# Rotation is single-use: replaying the now-consumed refresh token cannot mint a second pair
|
||||
# (a captured token is dead once the legitimate holder has rotated).
|
||||
replayed = await _token(grant_type="refresh_token", code=None, refresh_token=payload["refresh_token"])
|
||||
assert json.loads(replayed.body)["error"] == "invalid_grant"
|
||||
assert "already used" in json.loads(replayed.body).get("error_description", "")
|
||||
|
||||
cross_client = await _token(
|
||||
grant_type="refresh_token",
|
||||
code=None,
|
||||
refresh_token=payload["refresh_token"],
|
||||
client_id=(await _register([REDIRECT_URI]))["client_id"],
|
||||
)
|
||||
assert json.loads(cross_client.body)["error"] == "invalid_grant"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_complete_rejects_missing_tampered_and_expired_flows():
|
||||
missing = await complete_connect_flow(
|
||||
request=_request("/authorize/complete", method="POST"),
|
||||
flow_handle="nope",
|
||||
session_user_id="u1",
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert missing.status_code == 400
|
||||
|
||||
tampered = await complete_connect_flow(
|
||||
request=_request("/authorize/complete", cookies={f"{CONNECT_FLOW_COOKIE_PREFIX}h1": "garbage"}, method="POST"),
|
||||
flow_handle="h1",
|
||||
session_user_id="u1",
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert tampered.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_rejects_expired_code_and_missing_configuration():
|
||||
expired_code = _seal(
|
||||
GATEWAY_AUTH_CODE_PREFIX,
|
||||
_GatewayAuthCode(
|
||||
user_id="u1",
|
||||
client_id="llm_dcrc_x",
|
||||
redirect_uri=REDIRECT_URI,
|
||||
code_challenge=CODE_CHALLENGE,
|
||||
jti="jti-1",
|
||||
iat=int((datetime.now(timezone.utc) - timedelta(seconds=500)).timestamp()),
|
||||
exp=int((datetime.now(timezone.utc) - timedelta(seconds=500 - GATEWAY_AUTH_CODE_TTL_SECONDS)).timestamp()),
|
||||
),
|
||||
)
|
||||
response = await aggregate_token(
|
||||
request=_request("/token", method="POST"),
|
||||
grant_type="authorization_code",
|
||||
code=expired_code,
|
||||
redirect_uri=REDIRECT_URI,
|
||||
client_id="llm_dcrc_x",
|
||||
code_verifier=CODE_VERIFIER,
|
||||
refresh_token=None,
|
||||
master_key=MASTER_KEY,
|
||||
reload_user=_reload_user_active,
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert json.loads(response.body)["error"] == "invalid_grant"
|
||||
|
||||
no_master_key = await aggregate_token(
|
||||
request=_request("/token", method="POST"),
|
||||
grant_type="authorization_code",
|
||||
code="llm_gcode_x",
|
||||
redirect_uri=REDIRECT_URI,
|
||||
client_id="llm_dcrc_x",
|
||||
code_verifier=CODE_VERIFIER,
|
||||
refresh_token=None,
|
||||
master_key=None,
|
||||
reload_user=_reload_user_active,
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert no_master_key.status_code == 500
|
||||
assert json.loads(no_master_key.body)["error"] == "server_error"
|
||||
|
||||
unsupported = await aggregate_token(
|
||||
request=_request("/token", method="POST"),
|
||||
grant_type="password",
|
||||
code=None,
|
||||
redirect_uri=None,
|
||||
client_id="llm_dcrc_x",
|
||||
code_verifier=None,
|
||||
refresh_token=None,
|
||||
master_key=MASTER_KEY,
|
||||
reload_user=_reload_user_active,
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert json.loads(unsupported.body)["error"] == "unsupported_grant_type"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"failure,expected_status,expected_error",
|
||||
[
|
||||
("no_active_key", 400, "invalid_grant"),
|
||||
("unavailable", 503, "temporarily_unavailable"),
|
||||
("unresolvable", 500, "server_error"),
|
||||
],
|
||||
)
|
||||
async def test_token_gates_on_live_user_revalidation(failure, expected_status, expected_error):
|
||||
client_id = (await _register([REDIRECT_URI]))["client_id"]
|
||||
authorize_response = _authorize(client_id, session_user_id="deactivated-user")
|
||||
handle, cookies = _flow_cookie_from(authorize_response)
|
||||
completed = await complete_connect_flow(
|
||||
request=_request("/authorize/complete", cookies=cookies, method="POST"),
|
||||
flow_handle=handle,
|
||||
session_user_id="deactivated-user",
|
||||
cache=DualCache(),
|
||||
)
|
||||
code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0]
|
||||
|
||||
async def _reload_user_failing(user_id: str):
|
||||
return failure
|
||||
|
||||
response = await aggregate_token(
|
||||
request=_request("/token", method="POST"),
|
||||
grant_type="authorization_code",
|
||||
code=code,
|
||||
redirect_uri=REDIRECT_URI,
|
||||
client_id=client_id,
|
||||
code_verifier=CODE_VERIFIER,
|
||||
refresh_token=None,
|
||||
master_key=MASTER_KEY,
|
||||
reload_user=_reload_user_failing,
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert response.status_code == expected_status
|
||||
assert json.loads(response.body)["error"] == expected_error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flow_is_single_use_shared_cache_rejects_second_complete():
|
||||
"""A double-submit of the finish step mints only ONE code: the second complete over the
|
||||
same cache fails invalid_request (atomic flow claim), so one sign-in cannot yield two codes."""
|
||||
cache = DualCache()
|
||||
client_id = (await _register([REDIRECT_URI]))["client_id"]
|
||||
handle, cookies = _flow_cookie_from(_authorize(client_id, session_user_id="u1"))
|
||||
|
||||
first = await complete_connect_flow(
|
||||
request=_request("/authorize/complete", cookies=cookies, method="POST"),
|
||||
flow_handle=handle,
|
||||
session_user_id="u1",
|
||||
cache=cache,
|
||||
)
|
||||
assert first.status_code == 303
|
||||
second = await complete_connect_flow(
|
||||
request=_request("/authorize/complete", cookies=cookies, method="POST"),
|
||||
flow_handle=handle,
|
||||
session_user_id="u1",
|
||||
cache=cache,
|
||||
)
|
||||
assert second.status_code == 400
|
||||
assert json.loads(second.body)["error"] == "invalid_request"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_rejects_out_of_range_code_verifier():
|
||||
"""RFC 7636: a code_verifier outside 43-128 chars is invalid_request, not a confusing
|
||||
invalid_grant PKCE-mismatch."""
|
||||
for bad in ["short", "x" * 200]:
|
||||
response = await aggregate_token(
|
||||
request=_request("/token", method="POST"),
|
||||
grant_type="authorization_code",
|
||||
code="llm_gcode_whatever",
|
||||
redirect_uri=REDIRECT_URI,
|
||||
client_id="llm_dcrc_x",
|
||||
code_verifier=bad,
|
||||
refresh_token=None,
|
||||
master_key=MASTER_KEY,
|
||||
reload_user=_reload_user_active,
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] == "invalid_request"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authorize_rejects_over_long_state():
|
||||
client_id = (await _register([REDIRECT_URI]))["client_id"]
|
||||
response = aggregate_authorize(
|
||||
request=_request(query=f"client_id={client_id}"),
|
||||
client_id=client_id,
|
||||
redirect_uri=REDIRECT_URI,
|
||||
state="s" * 2000,
|
||||
code_challenge=CODE_CHALLENGE,
|
||||
code_challenge_method="S256",
|
||||
response_type="code",
|
||||
session_user_id="u1",
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] == "invalid_request"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_ascii_code_challenge_fails_grant_not_500():
|
||||
"""A non-ASCII code_challenge (unvalidated from the client) must yield a clean
|
||||
invalid_grant, never a TypeError-driven 500 (bytes comparison, not str)."""
|
||||
client_id = (await _register([REDIRECT_URI]))["client_id"]
|
||||
# Seal a code carrying a non-ASCII challenge directly (authorize requires S256 shape,
|
||||
# but the challenge charset is not validated there, so this state is reachable).
|
||||
from datetime import datetime, timezone
|
||||
|
||||
code = _seal(
|
||||
GATEWAY_AUTH_CODE_PREFIX,
|
||||
_GatewayAuthCode(
|
||||
user_id="u1",
|
||||
client_id=client_id,
|
||||
redirect_uri=REDIRECT_URI,
|
||||
code_challenge="challenge-with-€-non-ascii",
|
||||
jti="jti-x",
|
||||
iat=int(datetime.now(timezone.utc).timestamp()),
|
||||
exp=int(datetime.now(timezone.utc).timestamp()) + 120,
|
||||
),
|
||||
)
|
||||
response = await aggregate_token(
|
||||
request=_request("/token", method="POST"),
|
||||
grant_type="authorization_code",
|
||||
code=code,
|
||||
redirect_uri=REDIRECT_URI,
|
||||
client_id=client_id,
|
||||
code_verifier=CODE_VERIFIER,
|
||||
refresh_token=None,
|
||||
master_key=MASTER_KEY,
|
||||
reload_user=_reload_user_active,
|
||||
cache=DualCache(),
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert json.loads(response.body)["error"] == "invalid_grant"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_use_guard_in_memory_is_single_use_within_process():
|
||||
"""No Redis configured (single-replica): the in-memory increment is authoritative — the first claim
|
||||
wins, a replay of the same id loses."""
|
||||
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import _SingleUseGuard
|
||||
|
||||
guard = _SingleUseGuard(DualCache()) # redis_cache is None
|
||||
assert await guard.claim("jti-inmem", 60) is True
|
||||
assert await guard.claim("jti-inmem", 60) is False # replay of the same id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_use_guard_uses_redis_as_sole_authority_when_configured():
|
||||
"""With Redis configured it is the SOLE authority: the shared INCR result decides the claim (1 →
|
||||
first caller, >1 → replay), and the per-worker in-memory count is never consulted."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import _SingleUseGuard
|
||||
|
||||
cache = DualCache()
|
||||
cache.redis_cache = MagicMock()
|
||||
cache.redis_cache.async_increment = AsyncMock(return_value=1)
|
||||
# in-memory must NOT be consulted when Redis is configured — poison it so any fallback is visible.
|
||||
cache.async_increment_cache = AsyncMock(side_effect=AssertionError("must not fall back to in-memory"))
|
||||
|
||||
guard = _SingleUseGuard(cache)
|
||||
assert await guard.claim("jti-redis", 60) is True
|
||||
cache.redis_cache.async_increment = AsyncMock(return_value=2)
|
||||
assert await guard.claim("jti-redis", 60) is False # Redis says 2 → replay
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_use_guard_fails_closed_when_redis_errors():
|
||||
"""A Redis fault must fail the claim CLOSED (refuse the id) rather than fall back to the per-worker
|
||||
in-memory count — which would let each replica observe count==1 and replay the one-time id (the
|
||||
Cursor/Veria replay-across-workers finding)."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import _SingleUseGuard
|
||||
|
||||
cache = DualCache()
|
||||
cache.redis_cache = MagicMock()
|
||||
cache.redis_cache.async_increment = AsyncMock(side_effect=ConnectionError("redis down"))
|
||||
cache.async_increment_cache = AsyncMock(return_value=1) # would fail OPEN if the guard fell back
|
||||
|
||||
guard = _SingleUseGuard(cache)
|
||||
assert await guard.claim("jti-fault", 60) is False # fail closed, not a fallback count of 1
|
||||
|
|
@ -176,9 +176,7 @@ async def test_authenticate_user_invalid_credentials():
|
|||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
|
||||
|
||||
with patch.dict(
|
||||
os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": "correct-password"}
|
||||
):
|
||||
with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": "correct-password"}):
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await authenticate_user(
|
||||
username=ui_username,
|
||||
|
|
@ -227,9 +225,7 @@ async def test_authenticate_user_wrong_password():
|
|||
)
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(
|
||||
return_value=mock_user
|
||||
)
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=mock_user)
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
|
|
@ -279,9 +275,7 @@ async def test_authenticate_user_email_case_insensitive_login():
|
|||
return None
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(
|
||||
side_effect=mock_find_first
|
||||
)
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(side_effect=mock_find_first)
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
|
|
@ -334,9 +328,7 @@ async def test_authenticate_user_database_required_for_admin():
|
|||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
|
||||
|
||||
with patch.dict(
|
||||
os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": ui_password}
|
||||
):
|
||||
with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": ui_password}):
|
||||
with patch(
|
||||
"litellm.proxy.auth.login_utils.user_update",
|
||||
new_callable=AsyncMock,
|
||||
|
|
@ -429,9 +421,7 @@ def test_authenticate_user_non_ascii_direct_comparison():
|
|||
assert result is True
|
||||
|
||||
# And correctly returns False for different passwords
|
||||
result = secrets.compare_digest(
|
||||
password.encode("utf-8"), "different£pass".encode("utf-8")
|
||||
)
|
||||
result = secrets.compare_digest(password.encode("utf-8"), "different£pass".encode("utf-8"))
|
||||
assert result is False
|
||||
|
||||
|
||||
|
|
@ -531,9 +521,7 @@ async def test_authenticate_user_database_login_with_non_ascii_password():
|
|||
return None
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(
|
||||
side_effect=mock_find_first
|
||||
)
|
||||
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(side_effect=mock_find_first)
|
||||
|
||||
with patch.dict(
|
||||
os.environ,
|
||||
|
|
@ -559,3 +547,58 @@ async def test_authenticate_user_database_login_with_non_ascii_password():
|
|||
assert isinstance(result, LoginResult)
|
||||
assert result.user_id == "test-user-123"
|
||||
assert result.user_email == user_email
|
||||
|
||||
|
||||
class TestEncodeUiSessionJwt:
|
||||
"""The UI session cookie must carry a bounded exp so it does not stay
|
||||
signature-valid until the master key rotates, and so the session-cookie readers
|
||||
that require a bounded lifetime (the MCP interactive sign-in) accept it."""
|
||||
|
||||
def _decode(self, token: str) -> dict:
|
||||
import jwt
|
||||
|
||||
return jwt.decode(token, "sk-master-for-tests", algorithms=["HS256"])
|
||||
|
||||
def test_encoded_cookie_carries_bounded_exp(self):
|
||||
import time
|
||||
|
||||
from litellm.proxy.auth.login_utils import encode_ui_session_jwt
|
||||
|
||||
token_object = {"user_id": "u1", "key": "sk-abc", "login_method": "username_password"}
|
||||
with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "24h"):
|
||||
token = encode_ui_session_jwt(token_object, "sk-master-for-tests")
|
||||
claims = self._decode(token)
|
||||
assert claims["user_id"] == "u1"
|
||||
assert claims["login_method"] == "username_password"
|
||||
remaining = claims["exp"] - int(time.time())
|
||||
assert 23 * 3600 < remaining <= 24 * 3600
|
||||
|
||||
def test_duration_is_honored_from_env(self):
|
||||
import time
|
||||
|
||||
from litellm.proxy.auth.login_utils import encode_ui_session_jwt
|
||||
|
||||
with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "1h"):
|
||||
token = encode_ui_session_jwt({"user_id": "u1"}, "sk-master-for-tests")
|
||||
remaining = self._decode(token)["exp"] - int(time.time())
|
||||
assert 0 < remaining <= 3600
|
||||
|
||||
def test_cookie_is_accepted_by_the_exp_requiring_session_reader(self):
|
||||
"""The regression this change exists for: before it, the UI cookie carried no
|
||||
exp and _user_id_from_session_cookie (require=["exp"]) rejected every real login,
|
||||
so the MCP interactive sign-in could never capture identity. A cookie minted by
|
||||
this helper must now be accepted."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
|
||||
_user_id_from_session_cookie,
|
||||
)
|
||||
from litellm.proxy.auth.login_utils import encode_ui_session_jwt
|
||||
|
||||
token_object = {"user_id": "cornell-user", "key": "sk-abc", "login_method": "sso"}
|
||||
with patch("litellm.proxy.auth.login_utils.LITELLM_UI_SESSION_DURATION", "24h"):
|
||||
token = encode_ui_session_jwt(token_object, "sk-master-for-tests")
|
||||
request = MagicMock()
|
||||
request.cookies = {"token": token}
|
||||
with patch("litellm.proxy.proxy_server.master_key", "sk-master-for-tests"):
|
||||
assert _user_id_from_session_cookie(request) == "cornell-user"
|
||||
|
|
|
|||
|
|
@ -7763,3 +7763,80 @@ async def test_cli_completion_persists_assertion_under_db_user_id():
|
|||
|
||||
retain_mock.assert_awaited_once_with(user_id="cli-user-id", assertion=assertion)
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
class TestSameOriginReturnPath:
|
||||
"""The same-origin relative return_to arm added for the MCP gateway DCR authorize
|
||||
round-trip: only strictly relative paths qualify, so login can never redirect the
|
||||
browser off the gateway origin."""
|
||||
|
||||
def test_accepts_relative_paths(self):
|
||||
from litellm.proxy.management_endpoints.ui_sso import _is_same_origin_return_path
|
||||
|
||||
assert _is_same_origin_return_path("/authorize?client_id=llm_dcrc_x&state=s") is True
|
||||
assert _is_same_origin_return_path("/some_server/authorize") is True
|
||||
|
||||
def test_rejects_absolute_protocol_relative_and_backslash_paths(self):
|
||||
from litellm.proxy.management_endpoints.ui_sso import _is_same_origin_return_path
|
||||
|
||||
assert _is_same_origin_return_path("https://evil.example.com/authorize") is False
|
||||
assert _is_same_origin_return_path("//evil.example.com/authorize") is False
|
||||
assert _is_same_origin_return_path("/\\evil.example.com") is False
|
||||
assert _is_same_origin_return_path("javascript:alert(1)") is False
|
||||
assert _is_same_origin_return_path("") is False
|
||||
|
||||
|
||||
class TestPersistReturnToCookieSharedHelper:
|
||||
"""The single shared return_to helper used by EVERY sign-in branch (SSO / Okta / generic AND the
|
||||
username/password form). It must be best-effort and NEVER raise — a bad return_to can never block
|
||||
sign-in. Regression: the password form previously 400'd because it called _validate_return_to
|
||||
directly (which raises for a non-matching absolute return_to when control_plane_url is set)."""
|
||||
|
||||
@staticmethod
|
||||
def _cookie(resp) -> str:
|
||||
return resp.headers.get("set-cookie", "")
|
||||
|
||||
def test_sets_cookie_for_same_origin_relative_path(self, monkeypatch):
|
||||
from fastapi import Response
|
||||
|
||||
from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
|
||||
resp = Response()
|
||||
_persist_return_to_cookie(resp, "/mcp/authorize?client_id=llm_dcrc_abc")
|
||||
assert "litellm_cp_return_to=" in self._cookie(resp)
|
||||
|
||||
def test_bad_absolute_with_control_plane_configured_does_not_raise_and_is_not_stored(self, monkeypatch):
|
||||
"""THE regression: a non-matching absolute return_to with control_plane_url set must NOT raise
|
||||
(it did, blocking the login form) and must NOT be stored — sign-in proceeds."""
|
||||
from fastapi import Response
|
||||
|
||||
from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"}
|
||||
)
|
||||
resp = Response()
|
||||
_persist_return_to_cookie(resp, "https://evil.example.com/steal") # must not raise
|
||||
assert "litellm_cp_return_to=" not in self._cookie(resp)
|
||||
|
||||
def test_none_return_to_is_a_noop(self):
|
||||
from fastapi import Response
|
||||
|
||||
from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie
|
||||
|
||||
resp = Response()
|
||||
_persist_return_to_cookie(resp, None)
|
||||
assert "litellm_cp_return_to=" not in self._cookie(resp)
|
||||
|
||||
def test_control_plane_matching_absolute_is_stored(self, monkeypatch):
|
||||
from fastapi import Response
|
||||
|
||||
from litellm.proxy.management_endpoints.ui_sso import _persist_return_to_cookie
|
||||
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.general_settings", {"control_plane_url": "https://cp.example.com"}
|
||||
)
|
||||
resp = Response()
|
||||
_persist_return_to_cookie(resp, "https://cp.example.com/ui?page=models")
|
||||
assert "litellm_cp_return_to=" in self._cookie(resp)
|
||||
|
|
|
|||
|
|
@ -49,9 +49,7 @@ def _install_login_mocks(monkeypatch, raise_on_auth: bool = False) -> None:
|
|||
}
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.auth.login_utils.authenticate_user", _fake_auth)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.auth.login_utils.create_ui_token_object", _fake_token_object
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.auth.login_utils.create_ui_token_object", _fake_token_object)
|
||||
monkeypatch.setattr(ps, "master_key", "sk-test-master")
|
||||
monkeypatch.setattr(ps, "general_settings", {})
|
||||
monkeypatch.setattr(ps, "premium_user", False)
|
||||
|
|
@ -69,9 +67,7 @@ def test_fallback_login_returns_html_form(client, monkeypatch):
|
|||
body_lower = response.text.lower()
|
||||
shape = {
|
||||
"status": response.status_code,
|
||||
"content_type_html": response.headers.get("content-type", "").startswith(
|
||||
"text/html"
|
||||
),
|
||||
"content_type_html": response.headers.get("content-type", "").startswith("text/html"),
|
||||
"has_form": "<form" in body_lower or "username" in body_lower,
|
||||
}
|
||||
assert shape == {
|
||||
|
|
@ -88,9 +84,7 @@ def test_fallback_login_returns_html_form_with_ui_username_set(client, monkeypat
|
|||
body_lower = response.text.lower()
|
||||
shape = {
|
||||
"status": response.status_code,
|
||||
"content_type_html": response.headers.get("content-type", "").startswith(
|
||||
"text/html"
|
||||
),
|
||||
"content_type_html": response.headers.get("content-type", "").startswith("text/html"),
|
||||
"has_form_or_username": "<form" in body_lower or "username" in body_lower,
|
||||
}
|
||||
assert shape == {
|
||||
|
|
@ -126,11 +120,7 @@ def test_fallback_login_invalid_method_405(client):
|
|||
"""POST against the GET-only /fallback/login is rejected (error path)."""
|
||||
response = client.post("/fallback/login")
|
||||
assert response.status_code == 405
|
||||
body = (
|
||||
response.json()
|
||||
if response.headers.get("content-type", "").startswith("application/json")
|
||||
else {}
|
||||
)
|
||||
body = response.json() if response.headers.get("content-type", "").startswith("application/json") else {}
|
||||
assert isinstance(body, dict)
|
||||
|
||||
|
||||
|
|
@ -193,15 +183,15 @@ def test_v2_login_success_returns_token_and_redirect(client, monkeypatch):
|
|||
json={"username": "admin", "password": "password"},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
assert normalize(
|
||||
response.json(), volatile=frozenset({"token", "redirect_url"})
|
||||
) == {"redirect_url": "<VOLATILE>", "token": "<VOLATILE>"}
|
||||
assert normalize(response.json(), volatile=frozenset({"token", "redirect_url"})) == {
|
||||
"redirect_url": "<VOLATILE>",
|
||||
"token": "<VOLATILE>",
|
||||
}
|
||||
body = response.json()
|
||||
set_cookie = response.headers.get("set-cookie", "")
|
||||
shape = {
|
||||
"redirect_url_has_ui": "/ui/" in body.get("redirect_url", ""),
|
||||
"redirect_url_has_login_success": "login=success"
|
||||
in body.get("redirect_url", ""),
|
||||
"redirect_url_has_login_success": "login=success" in body.get("redirect_url", ""),
|
||||
"token_in_body": bool(body.get("token")),
|
||||
"token_cookie_set": "token=" in set_cookie,
|
||||
}
|
||||
|
|
@ -264,9 +254,7 @@ def test_v3_login_success_returns_code(client, monkeypatch):
|
|||
from litellm.proxy import proxy_server as ps
|
||||
|
||||
_install_login_mocks(monkeypatch)
|
||||
monkeypatch.setattr(
|
||||
ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}
|
||||
)
|
||||
monkeypatch.setattr(ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"})
|
||||
# Force the local (non-redis) cache path
|
||||
monkeypatch.setattr(ps, "redis_usage_cache", None)
|
||||
fake_cache = MagicMock()
|
||||
|
|
@ -301,9 +289,7 @@ def test_v3_login_authenticate_failure_500(client, monkeypatch):
|
|||
from litellm.proxy import proxy_server as ps
|
||||
|
||||
_install_login_mocks(monkeypatch, raise_on_auth=True)
|
||||
monkeypatch.setattr(
|
||||
ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}
|
||||
)
|
||||
monkeypatch.setattr(ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"})
|
||||
|
||||
response = client.post(
|
||||
"/v3/login",
|
||||
|
|
@ -337,9 +323,7 @@ def test_v3_login_exchange_missing_code_400(client, monkeypatch):
|
|||
"""Error path: missing 'code' in body -> 400 with 'Missing' message."""
|
||||
from litellm.proxy import proxy_server as ps
|
||||
|
||||
monkeypatch.setattr(
|
||||
ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}
|
||||
)
|
||||
monkeypatch.setattr(ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"})
|
||||
|
||||
response = client.post("/v3/login/exchange", json={})
|
||||
assert response.status_code == 400
|
||||
|
|
@ -352,9 +336,7 @@ def test_v3_login_exchange_invalid_code_401(client, monkeypatch):
|
|||
"""Error path: code that isn't in cache -> 401 'Invalid or expired'."""
|
||||
from litellm.proxy import proxy_server as ps
|
||||
|
||||
monkeypatch.setattr(
|
||||
ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}
|
||||
)
|
||||
monkeypatch.setattr(ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"})
|
||||
monkeypatch.setattr(ps, "redis_usage_cache", None)
|
||||
fake_cache = MagicMock()
|
||||
fake_cache.async_get_cache = AsyncMock(return_value=None)
|
||||
|
|
@ -372,9 +354,7 @@ def test_v3_login_exchange_success_returns_token_and_redirect(client, monkeypatc
|
|||
"""Pin: valid code -> JSON {token, redirect_url} + token cookie + cache deleted (single-use)."""
|
||||
from litellm.proxy import proxy_server as ps
|
||||
|
||||
monkeypatch.setattr(
|
||||
ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"}
|
||||
)
|
||||
monkeypatch.setattr(ps, "general_settings", {"control_plane_url": "https://cp.example.invalid"})
|
||||
monkeypatch.setattr(ps, "redis_usage_cache", None)
|
||||
|
||||
cached_payload = {
|
||||
|
|
@ -388,9 +368,10 @@ def test_v3_login_exchange_success_returns_token_and_redirect(client, monkeypatc
|
|||
|
||||
response = client.post("/v3/login/exchange", json={"code": "valid-code"})
|
||||
assert response.status_code == 200
|
||||
assert normalize(
|
||||
response.json(), volatile=frozenset({"token", "redirect_url"})
|
||||
) == {"token": "<VOLATILE>", "redirect_url": "<VOLATILE>"}
|
||||
assert normalize(response.json(), volatile=frozenset({"token", "redirect_url"})) == {
|
||||
"token": "<VOLATILE>",
|
||||
"redirect_url": "<VOLATILE>",
|
||||
}
|
||||
body = response.json()
|
||||
set_cookie = response.headers.get("set-cookie", "")
|
||||
shape = {
|
||||
|
|
@ -405,3 +386,77 @@ def test_v3_login_exchange_success_returns_token_and_redirect(client, monkeypatc
|
|||
"token_cookie_set": True,
|
||||
"cache_deleted_once": True,
|
||||
}
|
||||
|
||||
|
||||
def test_login_form_honors_same_origin_return_to_cookie(client, monkeypatch):
|
||||
"""The aggregate DCR connect flow preserves a same-origin return_to in the litellm_cp_return_to
|
||||
cookie; /login must RESUME there after password sign-in instead of dead-ending at the dashboard."""
|
||||
_install_login_mocks(monkeypatch)
|
||||
return_to = "/mcp/authorize?client_id=llm_dcrc_abc&response_type=code"
|
||||
response = client.post(
|
||||
"/login",
|
||||
data={"username": "admin", "password": "password"},
|
||||
cookies={"litellm_cp_return_to": return_to},
|
||||
follow_redirects=False,
|
||||
)
|
||||
assert response.status_code == 303
|
||||
assert response.headers.get("location", "") == return_to # resumed the connect flow, not the dashboard
|
||||
assert "token=" in response.headers.get("set-cookie", "")
|
||||
|
||||
|
||||
def test_login_form_honors_control_plane_return_to_cookie(client, monkeypatch):
|
||||
"""/login resumes through the SAME resumer the SSO callback uses, so it honors BOTH shapes
|
||||
_persist_return_to_cookie is willing to store. Honoring only the relative one silently dropped
|
||||
a control-plane return_to and landed the user on the dashboard."""
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
_install_login_mocks(monkeypatch)
|
||||
monkeypatch.setitem(ps.general_settings, "control_plane_url", "https://cp.example.com")
|
||||
response = client.post(
|
||||
"/login",
|
||||
data={"username": "admin", "password": "password"},
|
||||
cookies={"litellm_cp_return_to": "https://cp.example.com/console"},
|
||||
follow_redirects=False,
|
||||
)
|
||||
location = response.headers.get("location", "")
|
||||
assert response.status_code == 303
|
||||
assert location.startswith("https://cp.example.com/console")
|
||||
# Cross-origin arm hands the JWT off via a one-time code rather than a cookie.
|
||||
assert "code=" in location and "login=success" in location
|
||||
assert "token=" not in response.headers.get("set-cookie", "")
|
||||
|
||||
|
||||
def test_login_form_survives_stale_control_plane_return_to(client, monkeypatch):
|
||||
"""A stale one-shot cookie must NEVER fail a completed sign-in. The resumer rejects a return_to
|
||||
that no longer matches control_plane_url (a config change between the cookie's write and this
|
||||
read); the user has already authenticated, so land on the dashboard instead of erroring."""
|
||||
import litellm.proxy.proxy_server as ps
|
||||
|
||||
_install_login_mocks(monkeypatch)
|
||||
monkeypatch.setitem(ps.general_settings, "control_plane_url", "https://new-cp.example.com")
|
||||
response = client.post(
|
||||
"/login",
|
||||
data={"username": "admin", "password": "password"},
|
||||
cookies={"litellm_cp_return_to": "https://old-cp.example.com/console"},
|
||||
follow_redirects=False,
|
||||
)
|
||||
assert response.status_code == 303, "login must not break on a stale return_to cookie"
|
||||
location = response.headers.get("location", "")
|
||||
assert "old-cp.example.com" not in location
|
||||
assert "/ui/" in location
|
||||
|
||||
|
||||
def test_login_form_ignores_open_redirect_return_to(client, monkeypatch):
|
||||
"""A non-same-origin return_to (open-redirect attempt) is rejected — /login falls back to the
|
||||
dashboard rather than honoring an absolute/foreign URL."""
|
||||
_install_login_mocks(monkeypatch)
|
||||
response = client.post(
|
||||
"/login",
|
||||
data={"username": "admin", "password": "password"},
|
||||
cookies={"litellm_cp_return_to": "https://evil.example.com/steal"},
|
||||
follow_redirects=False,
|
||||
)
|
||||
assert response.status_code == 303
|
||||
location = response.headers.get("location", "")
|
||||
assert "evil.example.com" not in location
|
||||
assert "/ui/" in location # dashboard fallback
|
||||
|
|
|
|||
|
|
@ -127,11 +127,15 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch):
|
|||
general_settings={},
|
||||
premium_user=False,
|
||||
)
|
||||
mock_jwt_encode.assert_called_once_with(
|
||||
{"user_id": "test-user"},
|
||||
"test-master-key",
|
||||
algorithm="HS256",
|
||||
)
|
||||
mock_jwt_encode.assert_called_once()
|
||||
payload, secret = mock_jwt_encode.call_args.args
|
||||
# The UI session token carries a bounded-lifetime `exp` claim (dynamic timestamp), alongside
|
||||
# the user_id; assert its presence rather than an exact expiry value.
|
||||
assert payload["user_id"] == "test-user"
|
||||
assert isinstance(payload.get("exp"), int) and payload["exp"] > 0
|
||||
assert set(payload.keys()) == {"user_id", "exp"}
|
||||
assert secret == "test-master-key"
|
||||
assert mock_jwt_encode.call_args.kwargs == {"algorithm": "HS256"}
|
||||
|
||||
|
||||
def test_login_v2_returns_json_on_proxy_exception(monkeypatch):
|
||||
|
|
|
|||
|
|
@ -2972,7 +2972,7 @@
|
|||
"count": 1
|
||||
},
|
||||
"no-nested-ternary": {
|
||||
"count": 7
|
||||
"count": 6
|
||||
}
|
||||
},
|
||||
"src/components/chat/MCPConnectPicker.tsx": {
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import { Suspense, useEffect } from "react";
|
|||
import { useRouter, useSearchParams } from "next/navigation";
|
||||
import { useChatShell } from "@/contexts/ChatShellContext";
|
||||
import MCPAppsPanel from "@/components/chat/MCPAppsPanel";
|
||||
import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner";
|
||||
|
||||
// useSearchParams() requires a Suspense boundary for static export.
|
||||
function IntegrationsPageContent() {
|
||||
|
|
@ -11,6 +12,13 @@ function IntegrationsPageContent() {
|
|||
const router = useRouter();
|
||||
const searchParams = useSearchParams();
|
||||
const oauthReturn = searchParams.get("mcpOauthReturn");
|
||||
// Set by the gateway DCR authorize when a DCR client sends the user here to
|
||||
// authorize servers before finishing sign-in (see gateway_dcr_flow.py). The
|
||||
// handle keys the sealed per-flow cookie; connect_client is the client origin
|
||||
// for display only. connect_flow is NOT cleaned from the URL: the finish form
|
||||
// needs it, and the sealed cookie (not the URL) is the security boundary.
|
||||
const connectFlow = searchParams.get("connect_flow");
|
||||
const connectClient = searchParams.get("connect_client");
|
||||
|
||||
// Clean up the OAuth return param after it's been consumed — real routing means
|
||||
// we no longer need it to pick a tab, but it should not linger in the address bar.
|
||||
|
|
@ -24,7 +32,13 @@ function IntegrationsPageContent() {
|
|||
|
||||
return (
|
||||
<div className="flex-1 min-h-0 overflow-auto w-full py-8 px-8">
|
||||
<MCPAppsPanel accessToken={accessToken} selectedServers={selectedMCPServers} onChange={setSelectedMCPServers} />
|
||||
{connectFlow && <ConnectFlowBanner flowHandle={connectFlow} clientOrigin={connectClient} />}
|
||||
<MCPAppsPanel
|
||||
accessToken={accessToken}
|
||||
selectedServers={selectedMCPServers}
|
||||
onChange={setSelectedMCPServers}
|
||||
connectMode={!!connectFlow}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,51 @@
|
|||
import { afterEach, describe, expect, it, vi } from "vitest";
|
||||
import { render, screen } from "@testing-library/react";
|
||||
import ConnectFlowBanner from "./ConnectFlowBanner";
|
||||
|
||||
vi.mock("@/components/networking", () => ({
|
||||
getProxyBaseUrl: () => "https://gateway.example.com",
|
||||
}));
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
sessionStorage.clear();
|
||||
});
|
||||
|
||||
describe("ConnectFlowBanner", () => {
|
||||
it("posts the flow handle to the proxy /authorize/complete as a full-page form", () => {
|
||||
const { container } = render(<ConnectFlowBanner flowHandle="flow-handle-123" clientOrigin="https://claude.ai" />);
|
||||
|
||||
const form = container.querySelector("form")!;
|
||||
expect(form.getAttribute("method")).toBe("POST");
|
||||
expect(form.getAttribute("action")).toBe("https://gateway.example.com/authorize/complete");
|
||||
|
||||
const hidden = form.querySelector('input[name="flow"]') as HTMLInputElement;
|
||||
expect(hidden.value).toBe("flow-handle-123");
|
||||
// No token, code, or secret is ever placed in the form; the sealed cookie carries them.
|
||||
expect(form.innerHTML).not.toContain("token");
|
||||
});
|
||||
|
||||
it("shows the client origin so the user knows what they are connecting to", () => {
|
||||
render(<ConnectFlowBanner flowHandle="h" clientOrigin="https://claude.ai" />);
|
||||
expect(screen.getAllByText(/claude\.ai/).length).toBeGreaterThan(0);
|
||||
expect(screen.getByRole("button", { name: /finish connecting/i })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("falls back to a generic label when the client origin is unknown", () => {
|
||||
render(<ConnectFlowBanner flowHandle="h" clientOrigin={null} />);
|
||||
expect(screen.getAllByText(/the application/).length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
it("does NOT complete the flow on pagehide (completion requires the explicit button)", () => {
|
||||
// Security regression: an attacker could lure a signed-in victim to their own client's
|
||||
// authorize URL; the victim merely closing the tab must NOT deliver a victim-bound code.
|
||||
// Completion is a deliberate button press, never a side effect of leaving the page.
|
||||
const beaconMock = vi.fn(() => true);
|
||||
vi.stubGlobal("navigator", { ...navigator, sendBeacon: beaconMock });
|
||||
render(<ConnectFlowBanner flowHandle="flow-xyz" clientOrigin="https://claude.ai" />);
|
||||
|
||||
window.dispatchEvent(new Event("pagehide"));
|
||||
|
||||
expect(beaconMock).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,59 @@
|
|||
"use client";
|
||||
|
||||
import React from "react";
|
||||
import { CheckCircle } from "lucide-react";
|
||||
import { getProxyBaseUrl } from "@/components/networking";
|
||||
|
||||
interface Props {
|
||||
flowHandle: string;
|
||||
clientOrigin: string | null;
|
||||
}
|
||||
|
||||
/**
|
||||
* The interlude shown when a DCR client (Claude Desktop, MCP Inspector) sends the user
|
||||
* through the gateway sign-in and lands them on the apps grid to authorize servers. The
|
||||
* grid below authorizes individual servers into the per-user vault; this banner is the
|
||||
* finish step that returns the user to the client.
|
||||
*
|
||||
* Finishing requires the explicit "Finish connecting" button: a native form POST to the proxy's
|
||||
* /authorize/complete, which mints the gateway authorization code and 303-redirects to the DCR
|
||||
* client's own redirect URI (the full-page navigation carries the HttpOnly per-flow cookie and
|
||||
* follows the cross-origin redirect to the client's loopback).
|
||||
*
|
||||
* The button press IS the consent gate and must not be bypassed. An earlier version auto-finished
|
||||
* on tab close via navigator.sendBeacon; that let an attacker who lured a signed-in victim to their
|
||||
* own client's authorize URL harvest a victim-bound code the moment the victim closed the tab
|
||||
* (no click). Merely visiting the authorize URL is attacker-inducible, so completion has to be a
|
||||
* deliberate user action, not a side effect of leaving the page.
|
||||
*/
|
||||
const ConnectFlowBanner: React.FC<Props> = ({ flowHandle, clientOrigin }) => {
|
||||
const action = `${getProxyBaseUrl()}/authorize/complete`;
|
||||
const clientLabel = clientOrigin ?? "the application";
|
||||
|
||||
return (
|
||||
<div className="mb-6 rounded-lg border border-primary/30 bg-primary/5 px-5 py-4">
|
||||
<div className="flex items-start justify-between gap-4 flex-wrap">
|
||||
<div className="flex items-start gap-3 min-w-0">
|
||||
<CheckCircle className="h-5 w-5 text-primary shrink-0 mt-0.5" />
|
||||
<div className="min-w-0">
|
||||
<p className="text-sm font-semibold text-foreground">Connect your MCP servers to {clientLabel}</p>
|
||||
<p className="text-[13px] text-muted-foreground mt-0.5">
|
||||
Authorize the servers you want to use below, then click Finish connecting to return to {clientLabel}.
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
<form method="POST" action={action} className="shrink-0">
|
||||
<input type="hidden" name="flow" value={flowHandle} />
|
||||
<button
|
||||
type="submit"
|
||||
className="h-[38px] rounded-md bg-primary px-4 text-sm font-semibold text-primary-foreground hover:bg-primary/90"
|
||||
>
|
||||
Finish connecting
|
||||
</button>
|
||||
</form>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default ConnectFlowBanner;
|
||||
|
|
@ -13,7 +13,7 @@ import {
|
|||
getMCPOAuthUserCredentialStatus,
|
||||
listMCPTools,
|
||||
} from "../networking";
|
||||
import { AUTH_TYPE, MCPServer, MCPTool, handleTransport } from "../mcp_tools/types";
|
||||
import { AUTH_TYPE, MCPServer, MCPTool, handleTransport, isUnsupportedOnGatewayConnect } from "../mcp_tools/types";
|
||||
import { Logo } from "@/components/molecules/logo/Logo";
|
||||
import MessageManager from "@/components/molecules/message_manager";
|
||||
import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow";
|
||||
|
|
@ -71,6 +71,7 @@ interface Props {
|
|||
accessToken: string;
|
||||
selectedServers: string[];
|
||||
onChange: (servers: string[]) => void;
|
||||
connectMode?: boolean;
|
||||
}
|
||||
|
||||
const AVATAR_COLORS = [
|
||||
|
|
@ -96,7 +97,7 @@ type TabKey = "all" | "connected";
|
|||
|
||||
const TOOLS_FETCH_CONCURRENCY = 5;
|
||||
|
||||
const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange }) => {
|
||||
const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange, connectMode }) => {
|
||||
const [servers, setServers] = useState<MCPServer[]>([]);
|
||||
const [loading, setLoading] = useState(true);
|
||||
const [query, setQuery] = useState("");
|
||||
|
|
@ -106,6 +107,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
const [toolCounts, setToolCounts] = useState<Record<string, number>>({});
|
||||
const [loadingCounts, setLoadingCounts] = useState(false);
|
||||
const [oauthConnected, setOauthConnected] = useState<Set<string>>(new Set());
|
||||
const [oauthChecking, setOauthChecking] = useState<Set<string>>(new Set());
|
||||
|
||||
const serversRef = useRef<MCPServer[]>([]);
|
||||
useEffect(() => {
|
||||
|
|
@ -148,6 +150,14 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
}
|
||||
} catch {
|
||||
// ignore
|
||||
} finally {
|
||||
if (!fetchLoadCancelledRef.current) {
|
||||
setOauthChecking((prev) => {
|
||||
const next = new Set(prev);
|
||||
next.delete(server.server_id);
|
||||
return next;
|
||||
});
|
||||
}
|
||||
}
|
||||
},
|
||||
[accessToken],
|
||||
|
|
@ -160,9 +170,13 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
.then(async (serverData) => {
|
||||
if (fetchLoadCancelledRef.current) return;
|
||||
const list: MCPServer[] = Array.isArray(serverData) ? serverData : serverData?.data ?? [];
|
||||
const oauthServers = list.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2);
|
||||
setServers(list);
|
||||
setOauthChecking(new Set(oauthServers.map((s) => s.server_id)));
|
||||
setLoading(false);
|
||||
|
||||
oauthServers.forEach((s) => checkOauthCredential(s));
|
||||
|
||||
setLoadingCounts(true);
|
||||
const chunks = Array.from({ length: Math.ceil(list.length / TOOLS_FETCH_CONCURRENCY) }, (_, i) =>
|
||||
list.slice(i * TOOLS_FETCH_CONCURRENCY, (i + 1) * TOOLS_FETCH_CONCURRENCY),
|
||||
|
|
@ -172,9 +186,6 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
await Promise.allSettled(chunk.map((s) => fetchToolCount(s)));
|
||||
}
|
||||
if (!fetchLoadCancelledRef.current) setLoadingCounts(false);
|
||||
|
||||
const oauthServers = list.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2);
|
||||
oauthServers.forEach((s) => checkOauthCredential(s));
|
||||
})
|
||||
.catch(() => {
|
||||
if (!fetchLoadCancelledRef.current) {
|
||||
|
|
@ -231,6 +242,36 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
}
|
||||
};
|
||||
|
||||
const renderConnectionIndicator = (server: MCPServer) => {
|
||||
if (connectMode && isUnsupportedOnGatewayConnect(server.auth_type)) {
|
||||
return (
|
||||
<span className="text-[11px] text-muted-foreground shrink-0 whitespace-nowrap">
|
||||
Not supported on this connection
|
||||
</span>
|
||||
);
|
||||
}
|
||||
if (server.auth_type === AUTH_TYPE.OAUTH2) {
|
||||
if (oauthConnected.has(server.server_id)) {
|
||||
return <CheckCircle className="h-3.5 w-3.5 text-emerald-600 shrink-0" />;
|
||||
}
|
||||
if (oauthChecking.has(server.server_id)) {
|
||||
return <Skeleton className="h-6 w-16 shrink-0 rounded-md" />;
|
||||
}
|
||||
return (
|
||||
<OAuth2ConnectButton
|
||||
server={server}
|
||||
accessToken={accessToken}
|
||||
onConnect={(id) => setOauthConnected((prev) => new Set(prev).add(id))}
|
||||
variant="badge"
|
||||
/>
|
||||
);
|
||||
}
|
||||
if (selectedServers.includes(nameOf(server))) {
|
||||
return <span className="w-[7px] h-[7px] rounded-full bg-emerald-600 dark:bg-emerald-400 shrink-0" />;
|
||||
}
|
||||
return null;
|
||||
};
|
||||
|
||||
const { data: detailToolsResult, isLoading: loadingTools } = useQuery({
|
||||
queryKey: ["mcp-apps-panel-detail-tools", detailServer?.server_id],
|
||||
queryFn: () => listMCPTools(accessToken, detailServer!.server_id),
|
||||
|
|
@ -390,24 +431,30 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
<div>
|
||||
<div className="flex items-center gap-2 mb-1">
|
||||
<h2 className="m-0 text-lg font-semibold text-foreground">MCP Servers</h2>
|
||||
<span className="text-[10px] font-semibold text-primary bg-primary/10 rounded px-1.5 py-0.5 uppercase tracking-wider">
|
||||
Beta
|
||||
</span>
|
||||
</div>
|
||||
<div className="flex items-center gap-3">
|
||||
<p className="m-0 text-[13px] text-muted-foreground">Browse tools, authenticate once, use in chat</p>
|
||||
{loadingCounts ? (
|
||||
<span className="flex items-center gap-1.5 text-xs text-muted-foreground">
|
||||
<Loader2 className="h-3 w-3 animate-spin" />
|
||||
Loading tools...
|
||||
{!connectMode && (
|
||||
<span className="text-[10px] font-semibold text-primary bg-primary/10 rounded px-1.5 py-0.5 uppercase tracking-wider">
|
||||
Beta
|
||||
</span>
|
||||
) : totalTools > 0 ? (
|
||||
<span className="flex items-center gap-1 text-xs text-muted-foreground">
|
||||
<Wrench className="h-3 w-3" />
|
||||
{totalTools} tool{totalTools !== 1 ? "s" : ""} available
|
||||
</span>
|
||||
) : null}
|
||||
)}
|
||||
</div>
|
||||
{connectMode ? (
|
||||
<p className="m-0 text-[13px] text-muted-foreground">Click a server to see its tools and connect</p>
|
||||
) : (
|
||||
<div className="flex items-center gap-3">
|
||||
<p className="m-0 text-[13px] text-muted-foreground">Browse tools, authenticate once, use in chat</p>
|
||||
{loadingCounts ? (
|
||||
<span className="flex items-center gap-1.5 text-xs text-muted-foreground">
|
||||
<Loader2 className="h-3 w-3 animate-spin" />
|
||||
Loading tools...
|
||||
</span>
|
||||
) : totalTools > 0 ? (
|
||||
<span className="flex items-center gap-1 text-xs text-muted-foreground">
|
||||
<Wrench className="h-3 w-3" />
|
||||
{totalTools} tool{totalTools !== 1 ? "s" : ""} available
|
||||
</span>
|
||||
) : null}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
<div className="relative w-[220px]">
|
||||
<Search className="absolute left-3 top-1/2 -translate-y-1/2 h-3.5 w-3.5 text-muted-foreground" />
|
||||
|
|
@ -458,10 +505,10 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
<div className="grid grid-cols-2 border rounded-lg overflow-hidden">
|
||||
{filtered.map((server, idx) => {
|
||||
const name = nameOf(server);
|
||||
const isConnected = selectedServers.includes(name);
|
||||
const color = getAvatarColor(name);
|
||||
const isLeftCol = idx % 2 === 0;
|
||||
const count = toolCounts[name];
|
||||
const unsupported = !!connectMode && isUnsupportedOnGatewayConnect(server.auth_type);
|
||||
|
||||
return (
|
||||
<div
|
||||
|
|
@ -469,7 +516,9 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
onClick={() => setDetailServer(server)}
|
||||
className={`flex items-center gap-3 p-4 bg-card cursor-pointer transition-colors hover:bg-accent/30 min-w-0 ${
|
||||
isLeftCol ? "border-r" : ""
|
||||
} ${Math.floor(idx / 2) < Math.floor((filtered.length - 1) / 2) ? "border-b" : ""}`}
|
||||
} ${Math.floor(idx / 2) < Math.floor((filtered.length - 1) / 2) ? "border-b" : ""} ${
|
||||
unsupported ? "opacity-50" : ""
|
||||
}`}
|
||||
>
|
||||
{server.mcp_info?.logo_url ? (
|
||||
<Logo
|
||||
|
|
@ -500,22 +549,7 @@ const MCPAppsPanel: React.FC<Props> = ({ accessToken, selectedServers, onChange
|
|||
) : null}
|
||||
</div>
|
||||
</div>
|
||||
{server.auth_type === AUTH_TYPE.OAUTH2 ? (
|
||||
oauthConnected.has(server.server_id) ? (
|
||||
<CheckCircle className="h-3.5 w-3.5 text-emerald-600 shrink-0" />
|
||||
) : (
|
||||
<OAuth2ConnectButton
|
||||
server={server}
|
||||
accessToken={accessToken}
|
||||
onConnect={(id) => {
|
||||
setOauthConnected((prev) => new Set(prev).add(id));
|
||||
}}
|
||||
variant="badge"
|
||||
/>
|
||||
)
|
||||
) : isConnected ? (
|
||||
<span className="w-[7px] h-[7px] rounded-full bg-emerald-600 dark:bg-emerald-400 shrink-0" />
|
||||
) : null}
|
||||
{renderConnectionIndicator(server)}
|
||||
<ChevronRight className="h-3 w-3 text-muted-foreground/40 shrink-0" />
|
||||
</div>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import {
|
|||
preservedDeclaredAppCredentials,
|
||||
withoutMintedTokenCredentials,
|
||||
credentialAuthClass,
|
||||
isUnsupportedOnGatewayConnect,
|
||||
} from "./types";
|
||||
|
||||
describe("getOAuthAuthorizationIdentity", () => {
|
||||
|
|
@ -267,3 +268,23 @@ describe("credentialAuthClass", () => {
|
|||
expect(credentialAuthClass(null)).toBeNull();
|
||||
});
|
||||
});
|
||||
|
||||
describe("isUnsupportedOnGatewayConnect", () => {
|
||||
it("flags the modes that need a caller-supplied upstream token or subject", () => {
|
||||
// client-forwarded: caller presents the upstream Authorization per call
|
||||
expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.TRUE_PASSTHROUGH)).toBe(true);
|
||||
expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.OAUTH_DELEGATE)).toBe(true);
|
||||
// OBO: caller's own IdP token is the exchange subject, which the session bearer is not
|
||||
expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE)).toBe(true);
|
||||
});
|
||||
|
||||
it("does not flag modes the gateway can serve from server-side state or interactive vaulting", () => {
|
||||
// interactive authorization_code is the one mode the connect grid vaults per user
|
||||
expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.OAUTH2)).toBe(false);
|
||||
// server-configured credentials need no per-user connect
|
||||
expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.API_KEY)).toBe(false);
|
||||
expect(isUnsupportedOnGatewayConnect(AUTH_TYPE.NONE)).toBe(false);
|
||||
expect(isUnsupportedOnGatewayConnect(null)).toBe(false);
|
||||
expect(isUnsupportedOnGatewayConnect(undefined)).toBe(false);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -65,6 +65,15 @@ export const gatewayMintsClientFor = (server: { auth_type?: string | null; dcr_b
|
|||
server.auth_type === AUTH_TYPE.TRUE_PASSTHROUGH ||
|
||||
(server.auth_type === AUTH_TYPE.OAUTH_DELEGATE && !server.dcr_bridge);
|
||||
|
||||
// Auth modes that cannot be used through the gateway aggregate connect flow, where the client holds
|
||||
// only an identity-only session bearer and upstream credentials are resolved server-side from the
|
||||
// per-user vault. The vault is only populated by interactive authorization_code (oauth2). The
|
||||
// client-forwarded modes need the caller to present the upstream Authorization per call, and
|
||||
// oauth2_token_exchange (OBO) needs the caller's own IdP token as the subject to exchange; the
|
||||
// session bearer is neither, so none of these can complete a tool call on this connection.
|
||||
export const isUnsupportedOnGatewayConnect = (authType?: string | null): boolean =>
|
||||
isClientForwardedTokenMode(authType) || authType === AUTH_TYPE.OAUTH2_TOKEN_EXCHANGE;
|
||||
|
||||
export const OAUTH_FLOW = {
|
||||
INTERACTIVE: "interactive",
|
||||
M2M: "m2m",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue