Connect {server_name}

""" BYOK (Bring Your Own Key) OAuth 2.1 Authorization Server endpoints for MCP servers. When an MCP client connects to a BYOK-enabled server and no stored credential exists, LiteLLM runs a minimal OAuth 2.1 authorization code flow. The "authorization page" is just a form that asks the user for their API key — not a full identity-provider OAuth. Endpoints implemented here: GET /.well-known/oauth-authorization-server — OAuth authorization server metadata GET /.well-known/oauth-protected-resource — OAuth protected resource metadata GET /v1/mcp/oauth/authorize — Shows HTML form to collect the API key POST /v1/mcp/oauth/authorize — Stores temp auth code and redirects POST /v1/mcp/oauth/token — Exchanges code for a bearer JWT token """ import base64 import hashlib import html as _html_module import time import uuid from typing import Dict, Optional, cast from urllib.parse import urlencode, urlparse import jwt from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.db import store_user_credential from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( get_request_base_url, ) # --------------------------------------------------------------------------- # In-memory store for pending authorization codes. # Each entry: {code: {api_key, server_id, code_challenge, redirect_uri, user_id, expires_at}} # --------------------------------------------------------------------------- _byok_auth_codes: Dict[str, dict] = {} # Authorization codes expire after 5 minutes. _AUTH_CODE_TTL_SECONDS = 300 # Hard cap to prevent memory exhaustion from incomplete OAuth flows. _AUTH_CODES_MAX_SIZE = 1000 router = APIRouter(tags=["mcp"]) # --------------------------------------------------------------------------- # PKCE helper # --------------------------------------------------------------------------- def _verify_pkce(code_verifier: str, code_challenge: str) -> bool: """Return True iff SHA-256(code_verifier) == code_challenge (base64url, no padding).""" digest = hashlib.sha256(code_verifier.encode()).digest() computed = base64.urlsafe_b64encode(digest).rstrip(b"=").decode() return computed == code_challenge # --------------------------------------------------------------------------- # Cleanup of expired auth codes (called lazily on each request) # --------------------------------------------------------------------------- def _purge_expired_codes() -> None: now = time.time() expired = [k for k, v in _byok_auth_codes.items() if v["expires_at"] < now] for k in expired: del _byok_auth_codes[k] def _build_authorize_html( server_name: str, server_initial: str, client_id: str, redirect_uri: str, code_challenge: str, code_challenge_method: str, state: str, server_id: str, access_items: list, help_url: str, ) -> str: """Build the 2-step BYOK OAuth authorization page HTML.""" # Escape all user-supplied / externally-derived values before interpolation e = _html_module.escape server_name = e(server_name) server_initial = e(server_initial) client_id = e(client_id) redirect_uri = e(redirect_uri) code_challenge = e(code_challenge) code_challenge_method = e(code_challenge_method) state = e(state) server_id = e(server_id) # Build access checklist rows access_rows = "".join( f'

✓{e(item)}

' for item in access_items ) access_section = "" if access_rows: access_section = f"""

▮ Requested Access

{access_rows}

""" # Help link for step 2 help_link_html = "" if help_url: help_link_html = f'Where do I find my API key? ↗' return f""" Connect {server_name} — LiteLLM """ # --------------------------------------------------------------------------- # OAuth metadata discovery endpoints # --------------------------------------------------------------------------- @router.get("/.well-known/oauth-authorization-server", include_in_schema=False) async def oauth_authorization_server_metadata(request: Request) -> JSONResponse: """RFC 8414 Authorization Server Metadata for the BYOK OAuth flow.""" base_url = get_request_base_url(request) return JSONResponse( { "issuer": base_url, "authorization_endpoint": f"{base_url}/v1/mcp/oauth/authorize", "token_endpoint": f"{base_url}/v1/mcp/oauth/token", "response_types_supported": ["code"], "grant_types_supported": ["authorization_code"], "code_challenge_methods_supported": ["S256"], } ) @router.get("/.well-known/oauth-protected-resource", include_in_schema=False) async def oauth_protected_resource_metadata(request: Request) -> JSONResponse: """RFC 9728 Protected Resource Metadata pointing back at this server.""" base_url = get_request_base_url(request) return JSONResponse( { "resource": base_url, "authorization_servers": [base_url], } ) # --------------------------------------------------------------------------- # Authorization endpoint — GET (show form) and POST (process form) # --------------------------------------------------------------------------- @router.get("/v1/mcp/oauth/authorize", include_in_schema=False) async def byok_authorize_get( request: Request, client_id: Optional[str] = None, redirect_uri: Optional[str] = None, response_type: Optional[str] = None, code_challenge: Optional[str] = None, code_challenge_method: Optional[str] = None, state: Optional[str] = None, server_id: Optional[str] = None, ) -> HTMLResponse: """ Show the BYOK API-key entry form. The MCP client navigates the user here; the user types their API key and clicks "Connect & Authorize", which POSTs back to this same path. """ if response_type != "code": raise HTTPException(status_code=400, detail="response_type must be 'code'") if not redirect_uri: raise HTTPException(status_code=400, detail="redirect_uri is required") if not code_challenge: raise HTTPException(status_code=400, detail="code_challenge is required") # Resolve server metadata (name, description items, help URL). server_name = "MCP Server" access_items: list = [] help_url = "" if server_id: try: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) registry = global_mcp_server_manager.get_registry() if server_id in registry: srv = registry[server_id] server_name = srv.server_name or srv.name access_items = list(srv.byok_description or []) help_url = srv.byok_api_key_help_url or "" except Exception: pass server_initial = (server_name[0].upper()) if server_name else "S" html = _build_authorize_html( server_name=server_name, server_initial=server_initial, client_id=client_id or "", redirect_uri=redirect_uri, code_challenge=code_challenge, code_challenge_method=code_challenge_method or "S256", state=state or "", server_id=server_id or "", access_items=access_items, help_url=help_url, ) return HTMLResponse(content=html) @router.post("/v1/mcp/oauth/authorize", include_in_schema=False) async def byok_authorize_post( request: Request, client_id: str = Form(default=""), redirect_uri: str = Form(...), code_challenge: str = Form(...), code_challenge_method: str = Form(default="S256"), state: str = Form(default=""), server_id: str = Form(default=""), api_key: str = Form(...), ) -> RedirectResponse: """ Process the BYOK API-key form submission. Stores a short-lived authorization code and redirects the client back to redirect_uri with ?code=...&state=... query parameters. """ _purge_expired_codes() # Validate redirect_uri scheme to prevent open redirect parsed_uri = urlparse(redirect_uri) if parsed_uri.scheme not in ("http", "https"): raise HTTPException(status_code=400, detail="Invalid redirect_uri scheme") # Reject new codes if the store is at capacity (prevents memory exhaustion # from a burst of abandoned OAuth flows). if len(_byok_auth_codes) >= _AUTH_CODES_MAX_SIZE: raise HTTPException(status_code=503, detail="Too many pending authorization flows") if code_challenge_method != "S256": raise HTTPException( status_code=400, detail="Only S256 code_challenge_method is supported" ) auth_code = str(uuid.uuid4()) _byok_auth_codes[auth_code] = { "api_key": api_key, "server_id": server_id, "code_challenge": code_challenge, "redirect_uri": redirect_uri, "user_id": client_id, # external client passes LiteLLM user-id as client_id "expires_at": time.time() + _AUTH_CODE_TTL_SECONDS, } params = urlencode({"code": auth_code, "state": state}) separator = "&" if "?" in redirect_uri else "?" location = f"{redirect_uri}{separator}{params}" return RedirectResponse(url=location, status_code=302) # --------------------------------------------------------------------------- # Token endpoint # --------------------------------------------------------------------------- @router.post("/v1/mcp/oauth/token", include_in_schema=False) async def byok_token( request: Request, grant_type: str = Form(...), code: str = Form(...), redirect_uri: str = Form(default=""), code_verifier: str = Form(...), client_id: str = Form(default=""), ) -> JSONResponse: """ Exchange an authorization code for a short-lived BYOK session JWT. 1. Validates the authorization code and PKCE challenge. 2. Stores the API key via store_user_credential(). 3. Issues a signed JWT with type="byok_session". """ from litellm.proxy.proxy_server import master_key, prisma_client _purge_expired_codes() if grant_type != "authorization_code": raise HTTPException(status_code=400, detail="unsupported_grant_type") record = _byok_auth_codes.get(code) if record is None: raise HTTPException(status_code=400, detail="invalid_grant") if time.time() > record["expires_at"]: del _byok_auth_codes[code] raise HTTPException(status_code=400, detail="invalid_grant") # PKCE verification if not _verify_pkce(code_verifier, record["code_challenge"]): raise HTTPException(status_code=400, detail="invalid_grant") # Consume the code (one-time use) del _byok_auth_codes[code] server_id: str = record["server_id"] api_key_value: str = record["api_key"] # Prefer the user_id that was stored when the code was issued; fall back to # whatever client_id the token request supplies (they should match). user_id: str = record.get("user_id") or client_id if not user_id: raise HTTPException( status_code=400, detail="Cannot determine user_id; pass LiteLLM user id as client_id", ) # Persist the BYOK credential if prisma_client is not None: try: await store_user_credential( prisma_client=prisma_client, user_id=user_id, server_id=server_id, credential=api_key_value, ) # Invalidate any cached negative result so the user isn't blocked # for up to the TTL period after completing the OAuth flow. from litellm.proxy._experimental.mcp_server.server import ( _invalidate_byok_cred_cache, ) _invalidate_byok_cred_cache(user_id, server_id) except Exception as exc: verbose_proxy_logger.error( "byok_token: failed to store user credential for user=%s server=%s: %s", user_id, server_id, exc, ) raise HTTPException(status_code=500, detail="Failed to store credential") else: verbose_proxy_logger.warning( "byok_token: prisma_client is None — credential not persisted" ) if master_key is None: raise HTTPException( status_code=500, detail="Master key not configured; cannot issue token" ) now = int(time.time()) payload = { "user_id": user_id, "server_id": server_id, # "type" distinguishes this from regular proxy auth tokens. # The proxy's SSO JWT path uses asymmetric keys (RS256/ES256), so an # HS256 token signed with master_key cannot be accepted there. "type": "byok_session", "iat": now, "exp": now + 3600, } access_token = jwt.encode(payload, cast(str, master_key), algorithm="HS256") return JSONResponse( { "access_token": access_token, "token_type": "bearer", "expires_in": 3600, } )

Connect {server_name} MCP

How it works

Provide API Key