Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py: 21%
198 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2BYOK (Bring Your Own Key) OAuth 2.1 Authorization Server endpoints for MCP servers.
4When an MCP client connects to a BYOK-enabled server and no stored credential exists,
5LiteLLM runs a minimal OAuth 2.1 authorization code flow. The "authorization page" is
6just a form that asks the user for their API key — not a full identity-provider OAuth.
8Endpoints implemented here:
9 GET /.well-known/oauth-authorization-server — OAuth authorization server metadata
10 GET /v1/mcp/oauth/authorize — Shows HTML form to collect the API key
11 POST /v1/mcp/oauth/authorize — Stores temp auth code and redirects
12 POST /v1/mcp/oauth/token — Exchanges code for a bearer JWT token
13"""
15import base64
16import hashlib
17import html as _html_module
18import time
19import uuid
20from typing import Final, cast
21from urllib.parse import urlencode, urlparse
23import jwt
24from fastapi import APIRouter, Depends, Form, HTTPException, Request
25from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
27from litellm._logging import verbose_proxy_logger
28from litellm.proxy._experimental.mcp_server.db import store_user_credential
29from litellm.proxy._experimental.mcp_server.oauth_utils import (
30 BYOK_RESOURCE_METADATA_PATH,
31 TOKEN_NO_CACHE_HEADERS,
32 get_request_base_url,
33 validate_loopback_redirect_uri,
34 well_known_root_suffix,
35)
36from litellm.proxy._types import UserAPIKeyAuth
37from litellm.proxy.middleware.per_request_root_path_middleware import get_server_root_paths
39# ---------------------------------------------------------------------------
40# In-memory store for pending authorization codes.
41# Each entry: {code: {api_key, server_id, code_challenge, redirect_uri, user_id, expires_at}}
42# ---------------------------------------------------------------------------
43_byok_auth_codes: Final[dict[str, dict]] = {}
45# Authorization codes expire after 5 minutes.
46_AUTH_CODE_TTL_SECONDS: Final = 300
47# Hard cap to prevent memory exhaustion from incomplete OAuth flows.
48_AUTH_CODES_MAX_SIZE: Final = 1000
50router: Final = APIRouter(tags=["mcp"])
53# ---------------------------------------------------------------------------
54# PKCE helper
55# ---------------------------------------------------------------------------
58def _verify_pkce(code_verifier: str, code_challenge: str) -> bool:
59 """Return True iff SHA-256(code_verifier) == code_challenge (base64url, no padding)."""
60 digest: Final = hashlib.sha256(code_verifier.encode()).digest()
61 computed: Final = base64.urlsafe_b64encode(digest).rstrip(b"=").decode()
62 return computed == code_challenge
65# ---------------------------------------------------------------------------
66# Cleanup of expired auth codes (called lazily on each request)
67# ---------------------------------------------------------------------------
70def _purge_expired_codes() -> None:
71 now: Final = time.time()
72 expired: Final = [k for k, v in _byok_auth_codes.items() if v["expires_at"] < now]
73 for k in expired:
74 del _byok_auth_codes[k]
77def _oauth_token_error(code: str, status: int = 400) -> JSONResponse:
78 """RFC 6749 §5.2 token-endpoint error body: ``{"error": "<code>"}``.
79 FastAPI's default ``HTTPException`` renders ``{"detail": ...}`` which
80 spec-compliant OAuth clients parsing the ``error`` field won't recognize.
81 """
82 return JSONResponse(status_code=status, content={"error": code}, headers=TOKEN_NO_CACHE_HEADERS)
85def _user_id_from_session_cookie(request: Request) -> str | None:
86 """Return user_id from the UI ``token`` cookie, or None if missing/invalid."""
87 user_id, _ = _session_identity_from_cookie(request)
88 return user_id
91def _session_identity_from_cookie(request: Request) -> tuple[str | None, str | None]:
92 """Return ``(user_id, session_key)`` from the UI ``token`` cookie
93 (HS256-signed with ``master_key``), or ``(None, None)`` if missing/invalid.
95 The /token endpoint in this file ALSO issues master-key-signed JWTs
96 (type="byok_session") for MCP-client-side use. They must not be
97 accepted here as UI sessions — otherwise a leaked byok_session token
98 could be replayed as a cookie to re-authorize BYOK writes. Distinguish
99 by requiring a ``login_method`` claim (UI tokens set ``"sso"`` or
100 ``"username_password"``; byok_session tokens never set it) and
101 rejecting any token whose ``type`` identifies it as non-UI.
102 """
103 # Inline import avoids a circular dep (proxy_server -> mcp_server router).
104 from litellm.proxy.proxy_server import master_key
106 if not master_key: 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true
107 return None, None
108 token: Final = request.cookies.get("token")
109 if not token: 109 ↛ 111line 109 didn't jump to line 111 because the condition on line 109 was always true
110 return None, None
111 try:
112 payload: Final = jwt.decode(
113 token,
114 master_key,
115 algorithms=["HS256"],
116 # Require an expiry claim so a leaked UI session cookie has a
117 # bounded lifetime. PyJWT verifies exp by default when present;
118 # require=["exp"] additionally rejects tokens that omit it.
119 options={"require": ["exp"]},
120 )
121 except jwt.InvalidTokenError:
122 return None, None
123 if payload.get("type") == "byok_session":
124 return None, None
125 if payload.get("login_method") not in ("sso", "username_password"):
126 return None, None
127 user_id: Final = payload.get("user_id")
128 if not isinstance(user_id, str) or not user_id:
129 return None, None
130 session_key: Final = payload.get("key")
131 return user_id, session_key if isinstance(session_key, str) and session_key else None
134async def _session_key_is_live(session_key: str | None) -> bool:
135 """Whether the session key embedded in the UI cookie still resolves.
137 The cookie JWT stays signature-valid until ``exp``; the DB-backed session
138 key inside it is what ``POST /session/logout`` and password-change
139 revocation actually kill. Trusting the signature alone would let a
140 logged-out cookie keep authorizing BYOK credential writes, so re-resolve
141 the key here.
143 EXPERIMENTAL_UI_LOGIN blob tokens (non-``sk-``) have no DB row and are
144 unrevocable by construction (scoped out of revocation); they pass through
145 on their bounded 10-minute lifetime, as before.
146 """
147 from litellm.proxy._types import hash_token
148 from litellm.proxy.auth.auth_checks import get_key_object
149 from litellm.proxy.proxy_server import (
150 prisma_client,
151 proxy_logging_obj,
152 user_api_key_cache,
153 )
155 if session_key is None:
156 # Older cookies predating the ``key`` claim: nothing to resolve.
157 return True
158 if not session_key.startswith("sk-"):
159 return True
160 if prisma_client is None:
161 return True
162 try:
163 await get_key_object(
164 hashed_token=hash_token(session_key),
165 prisma_client=prisma_client,
166 user_api_key_cache=user_api_key_cache,
167 proxy_logging_obj=proxy_logging_obj,
168 )
169 except Exception:
170 return False
171 return True
174async def _byok_session_auth(request: Request) -> UserAPIKeyAuth:
175 """Require the UI session cookie, with the embedded session key
176 re-resolved against the DB so a revoked (logged-out) session cannot
177 authorize BYOK writes. Programmatic BYOK management uses
178 ``POST /v1/mcp/server/{id}/user-credential`` instead."""
179 user_id, session_key = _session_identity_from_cookie(request)
180 if not user_id:
181 raise HTTPException(status_code=401, detail="login_required")
182 if not await _session_key_is_live(session_key):
183 raise HTTPException(status_code=401, detail="login_required")
184 return UserAPIKeyAuth(api_key="byok_session_cookie", user_id=user_id)
187def _build_authorize_html(
188 server_name: str,
189 server_initial: str,
190 client_id: str,
191 redirect_uri: str,
192 code_challenge: str,
193 code_challenge_method: str,
194 state: str,
195 server_id: str,
196 access_items: list,
197 help_url: str,
198) -> str:
199 """Build the 2-step BYOK OAuth authorization page HTML."""
201 # Escape all user-supplied / externally-derived values before interpolation
202 e: Final = _html_module.escape
203 server_name = e(server_name)
204 server_initial = e(server_initial)
205 client_id = e(client_id)
206 redirect_uri = e(redirect_uri)
207 code_challenge = e(code_challenge)
208 code_challenge_method = e(code_challenge_method)
209 state = e(state)
210 server_id = e(server_id)
212 # Build access checklist rows
213 access_rows: Final = "".join(
214 f'<div class="access-item"><span class="check">✓</span>{e(item)}</div>' for item in access_items
215 )
216 access_section = ""
217 if access_rows:
218 access_section = f"""
219 <div class="access-box">
220 <div class="access-header">
221 <span class="shield">▮</span>
222 <span>Requested Access</span>
223 </div>
224 {access_rows}
225 </div>"""
227 # Help link for step 2
228 help_link_html = ""
229 if help_url:
230 help_link_html = (
231 f'<a class="help-link" href="{e(help_url)}" target="_blank">Where do I find my API key? ↗</a>'
232 )
234 return f"""<!DOCTYPE html>
235<html lang="en">
236<head>
237<meta charset="UTF-8">
238<meta name="viewport" content="width=device-width, initial-scale=1.0">
239<title>Connect {server_name} — LiteLLM</title>
240<style>
241 *, *::before, *::after {{ box-sizing: border-box; margin: 0; padding: 0; }}
242 body {{
243 font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
244 background: #0f172a;
245 min-height: 100vh;
246 display: flex;
247 align-items: center;
248 justify-content: center;
249 padding: 24px;
250 }}
251 .modal {{
252 background: #ffffff;
253 border-radius: 20px;
254 padding: 36px 32px 32px;
255 width: 440px;
256 max-width: 100%;
257 position: relative;
258 box-shadow: 0 25px 60px rgba(0,0,0,0.35);
259 }}
260 /* Progress dots */
261 .dots {{
262 display: flex;
263 justify-content: center;
264 gap: 7px;
265 margin-bottom: 28px;
266 }}
267 .dot {{
268 width: 8px; height: 8px;
269 border-radius: 50%;
270 background: #e2e8f0;
271 }}
272 .dot.active {{ background: #38bdf8; }}
273 /* Close button */
274 .close-btn {{
275 position: absolute;
276 top: 16px; right: 16px;
277 background: none; border: none;
278 font-size: 16px; color: #94a3b8;
279 cursor: pointer; line-height: 1;
280 width: 28px; height: 28px;
281 border-radius: 6px;
282 display: flex; align-items: center; justify-content: center;
283 }}
284 .close-btn:hover {{ background: #f1f5f9; color: #475569; }}
285 /* Logo pair */
286 .logos {{
287 display: flex; align-items: center; justify-content: center;
288 gap: 12px; margin-bottom: 20px;
289 }}
290 .logo {{
291 width: 52px; height: 52px;
292 border-radius: 14px;
293 display: flex; align-items: center; justify-content: center;
294 font-size: 22px; font-weight: 800; color: white;
295 }}
296 .logo-img {{
297 width: 52px; height: 52px;
298 border-radius: 14px;
299 object-fit: cover;
300 border: 1.5px solid #e2e8f0;
301 }}
302 .logo-s {{ background: linear-gradient(135deg, #818cf8 0%, #4f46e5 100%); }}
303 .logo-arrow {{ color: #cbd5e1; font-size: 20px; font-weight: 300; }}
304 /* Headings */
305 .step-title {{
306 text-align: center;
307 font-size: 21px; font-weight: 700;
308 color: #0f172a; margin-bottom: 8px;
309 }}
310 .step-subtitle {{
311 text-align: center;
312 font-size: 14px; color: #64748b;
313 line-height: 1.55; margin-bottom: 22px;
314 }}
315 /* Info box */
316 .info-box {{
317 background: #f8fafc;
318 border-radius: 12px;
319 padding: 14px 16px;
320 display: flex; gap: 12px;
321 margin-bottom: 14px;
322 }}
323 .info-icon {{ font-size: 17px; flex-shrink: 0; margin-top: 1px; color: #38bdf8; }}
324 .info-box h4 {{ font-size: 13px; font-weight: 600; color: #1e293b; margin-bottom: 4px; }}
325 .info-box p {{ font-size: 13px; color: #64748b; line-height: 1.5; }}
326 /* Access checklist */
327 .access-box {{
328 background: #f8fafc;
329 border-radius: 12px;
330 padding: 14px 16px;
331 margin-bottom: 22px;
332 }}
333 .access-header {{
334 display: flex; align-items: center; gap: 8px;
335 margin-bottom: 10px;
336 }}
337 .shield {{ color: #22c55e; font-size: 15px; }}
338 .access-header > span:last-child {{
339 font-size: 11px; font-weight: 700;
340 letter-spacing: 0.07em;
341 text-transform: uppercase;
342 color: #475569;
343 }}
344 .access-item {{
345 display: flex; align-items: center; gap: 9px;
346 font-size: 13.5px; color: #374151;
347 padding: 3px 0;
348 }}
349 .check {{ color: #22c55e; font-weight: 700; font-size: 13px; }}
350 /* Primary CTA */
351 .btn-primary {{
352 width: 100%; padding: 15px;
353 background: #0f172a; color: white;
354 border: none; border-radius: 12px;
355 font-size: 15px; font-weight: 600;
356 cursor: pointer; margin-bottom: 10px;
357 }}
358 .btn-primary:hover {{ background: #1e293b; }}
359 .btn-cancel {{
360 width: 100%; padding: 8px;
361 background: none; border: none;
362 font-size: 13.5px; color: #94a3b8;
363 cursor: pointer;
364 }}
365 .btn-cancel:hover {{ color: #64748b; }}
366 /* Step 2 nav */
367 .step2-nav {{
368 display: flex; align-items: center;
369 justify-content: space-between;
370 margin-bottom: 24px;
371 }}
372 .back-btn {{
373 background: none; border: none;
374 font-size: 13.5px; color: #64748b;
375 cursor: pointer; display: flex; align-items: center; gap: 4px;
376 }}
377 .back-btn:hover {{ color: #374151; }}
378 /* Key icon */
379 .key-icon-wrap {{
380 width: 46px; height: 46px;
381 background: #e0f2fe;
382 border-radius: 12px;
383 display: flex; align-items: center; justify-content: center;
384 margin-bottom: 14px;
385 }}
386 .key-icon-wrap svg {{ width: 22px; height: 22px; color: #0284c7; }}
387 /* Form elements */
388 .field-label {{
389 font-size: 13.5px; font-weight: 600;
390 color: #1e293b; display: block;
391 margin-bottom: 7px;
392 }}
393 .key-input {{
394 width: 100%; padding: 11px 13px;
395 border: 1.5px solid #e2e8f0;
396 border-radius: 10px;
397 font-size: 14px; color: #0f172a;
398 outline: none; transition: border-color 0.15s, box-shadow 0.15s;
399 }}
400 .key-input:focus {{
401 border-color: #38bdf8;
402 box-shadow: 0 0 0 3px rgba(56,189,248,0.12);
403 }}
404 .help-link {{
405 display: inline-flex; align-items: center; gap: 4px;
406 color: #0ea5e9; font-size: 13px;
407 text-decoration: none; margin: 8px 0 16px;
408 }}
409 .help-link:hover {{ text-decoration: underline; }}
410 /* Save toggle card */
411 .save-card {{
412 border: 1.5px solid #e2e8f0;
413 border-radius: 12px;
414 padding: 13px 15px;
415 margin-bottom: 6px;
416 }}
417 .save-row {{
418 display: flex; align-items: center; gap: 10px;
419 }}
420 .save-icon {{ font-size: 16px; }}
421 .save-label {{
422 flex: 1;
423 font-size: 14px; font-weight: 500; color: #1e293b;
424 }}
425 /* Toggle switch */
426 .toggle {{ position: relative; width: 44px; height: 24px; flex-shrink: 0; }}
427 .toggle input {{ opacity: 0; width: 0; height: 0; }}
428 .slider {{
429 position: absolute; inset: 0;
430 background: #e2e8f0;
431 border-radius: 24px; cursor: pointer;
432 transition: background 0.18s;
433 }}
434 .slider::before {{
435 content: '';
436 position: absolute;
437 width: 18px; height: 18px;
438 left: 3px; bottom: 3px;
439 background: white;
440 border-radius: 50%;
441 transition: transform 0.18s;
442 box-shadow: 0 1px 3px rgba(0,0,0,0.18);
443 }}
444 input:checked + .slider {{ background: #38bdf8; }}
445 input:checked + .slider::before {{ transform: translateX(20px); }}
446 /* Duration pills */
447 .duration-section {{ margin-top: 14px; }}
448 .duration-label {{
449 font-size: 12px; font-weight: 600;
450 color: #64748b; margin-bottom: 8px;
451 text-transform: uppercase; letter-spacing: 0.05em;
452 }}
453 .pills {{ display: flex; flex-wrap: wrap; gap: 7px; }}
454 .pill {{
455 padding: 6px 13px;
456 border: 1.5px solid #e2e8f0;
457 border-radius: 20px;
458 font-size: 13px; color: #475569;
459 cursor: pointer; background: white;
460 transition: all 0.13s;
461 user-select: none;
462 }}
463 .pill:hover {{ border-color: #94a3b8; }}
464 .pill.sel {{
465 border-color: #38bdf8;
466 color: #0284c7;
467 background: #e0f2fe;
468 }}
469 /* Security note */
470 .sec-note {{
471 background: #f8fafc;
472 border-radius: 10px;
473 padding: 11px 14px;
474 display: flex; gap: 9px; align-items: flex-start;
475 margin: 16px 0;
476 }}
477 .sec-icon {{ font-size: 13px; color: #94a3b8; margin-top: 1px; flex-shrink: 0; }}
478 .sec-note p {{ font-size: 12.5px; color: #64748b; line-height: 1.5; }}
479 /* Connect button */
480 .btn-connect {{
481 width: 100%; padding: 15px;
482 border: none; border-radius: 12px;
483 font-size: 15px; font-weight: 600;
484 cursor: pointer;
485 background: #bae6fd; color: #0369a1;
486 transition: background 0.15s, color 0.15s;
487 }}
488 .btn-connect.ready {{
489 background: #0ea5e9; color: white;
490 }}
491 .btn-connect.ready:hover {{ background: #0284c7; }}
492 /* Step visibility */
493 .step {{ display: none; }}
494 .step.show {{ display: block; }}
495</style>
496</head>
497<body>
498<div class="modal">
500 <!-- ── STEP 1: Connect ─────────────────────────────────────── -->
501 <div id="s1" class="step show">
502 <div class="dots">
503 <div class="dot active"></div>
504 <div class="dot"></div>
505 </div>
506 <button class="close-btn" type="button" onclick="doCancel()" title="Close">×</button>
508 <div class="logos">
509 <img src="/ui/assets/logos/litellm_logo.jpg" class="logo-img" alt="LiteLLM">
510 <span class="logo-arrow">→</span>
511 <div class="logo logo-s">{server_initial}</div>
512 </div>
514 <h2 class="step-title">Connect {server_name} MCP</h2>
515 <p class="step-subtitle">LiteLLM needs access to {server_name} to complete your request.</p>
517 <div class="info-box">
518 <span class="info-icon">
519 <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><circle cx="12" cy="12" r="10"/><line x1="12" y1="8" x2="12" y2="12"/><line x1="12" y1="16" x2="12.01" y2="16"/></svg>
520 </span>
521 <div>
522 <h4>How it works</h4>
523 <p>LiteLLM acts as a secure bridge. Your requests are routed through our MCP client directly to {server_name}’s API.</p>
524 </div>
525 </div>
527 {access_section}
529 <button class="btn-primary" type="button" onclick="goStep2()">
530 Continue to Authentication →
531 </button>
532 <button class="btn-cancel" type="button" onclick="doCancel()">Cancel</button>
533 </div>
535 <!-- ── STEP 2: Provide API Key ──────────────────────────────── -->
536 <div id="s2" class="step">
537 <div class="step2-nav">
538 <button class="back-btn" type="button" onclick="goStep1()">← Back</button>
539 <div class="dots">
540 <div class="dot active"></div>
541 <div class="dot active"></div>
542 </div>
543 <button class="close-btn" style="position:static;" type="button" onclick="doCancel()" title="Close">×</button>
544 </div>
546 <div class="key-icon-wrap">
547 <svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" fill="none" stroke="#0284c7" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 2l-2 2m-7.61 7.61a5.5 5.5 0 1 1-7.778 7.778 5.5 5.5 0 0 1 7.777-7.777zm0 0L15.5 7.5m0 0l3 3L22 7l-3-3m-3.5 3.5L19 4"/></svg>
548 </div>
549 <h2 class="step-title" style="text-align:left;">Provide API Key</h2>
550 <p class="step-subtitle" style="text-align:left;">Enter your {server_name} API key to authorize this connection.</p>
552 <form method="POST" id="authForm" onsubmit="prepareSubmit()">
553 <input type="hidden" name="client_id" value="{client_id}">
554 <input type="hidden" name="redirect_uri" value="{redirect_uri}">
555 <input type="hidden" name="code_challenge" value="{code_challenge}">
556 <input type="hidden" name="code_challenge_method" value="{code_challenge_method}">
557 <input type="hidden" name="state" value="{state}">
558 <input type="hidden" name="server_id" value="{server_id}">
559 <input type="hidden" name="duration" id="durInput" value="until_revoked">
561 <label class="field-label">{server_name} API Key</label>
562 <input
563 type="password"
564 name="api_key"
565 id="apiKey"
566 class="key-input"
567 placeholder="Enter your API key"
568 required
569 autofocus
570 oninput="syncBtn()"
571 >
573 {help_link_html}
575 <div class="save-card">
576 <div class="save-row">
577 <span class="save-label">Save key for future use</span>
578 <label class="toggle">
579 <input type="checkbox" id="saveToggle" onchange="toggleDur()">
580 <span class="slider"></span>
581 </label>
582 </div>
583 <div id="durSection" class="duration-section" style="display:none;">
584 <div class="duration-label">Duration</div>
585 <div class="pills">
586 <div class="pill" onclick="selDur('1h',this)">1 hour</div>
587 <div class="pill sel" onclick="selDur('24h',this)">24 hours</div>
588 <div class="pill" onclick="selDur('7d',this)">7 days</div>
589 <div class="pill" onclick="selDur('30d',this)">30 days</div>
590 <div class="pill" onclick="selDur('until_revoked',this)">Until I revoke</div>
591 </div>
592 </div>
593 </div>
595 <div class="sec-note">
596 <span class="sec-icon">
597 <svg xmlns="http://www.w3.org/2000/svg" width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><rect x="3" y="11" width="18" height="11" rx="2" ry="2"/><path d="M7 11V7a5 5 0 0 1 10 0v4"/></svg>
598 </span>
599 <p>Your key is stored securely and transmitted over HTTPS. It is never shared with third parties.</p>
600 </div>
602 <button type="submit" class="btn-connect" id="connectBtn">
603 Connect & Authorize
604 </button>
605 </form>
606 </div>
608</div>
609<script>
610 function goStep2() {{
611 document.getElementById('s1').classList.remove('show');
612 document.getElementById('s2').classList.add('show');
613 }}
614 function goStep1() {{
615 document.getElementById('s2').classList.remove('show');
616 document.getElementById('s1').classList.add('show');
617 }}
618 function doCancel() {{
619 if (window.opener) window.close();
620 else window.history.back();
621 }}
622 function toggleDur() {{
623 const on = document.getElementById('saveToggle').checked;
624 document.getElementById('durSection').style.display = on ? 'block' : 'none';
625 }}
626 function selDur(val, el) {{
627 document.querySelectorAll('.pill').forEach(p => p.classList.remove('sel'));
628 el.classList.add('sel');
629 document.getElementById('durInput').value = val;
630 }}
631 function syncBtn() {{
632 const btn = document.getElementById('connectBtn');
633 if (document.getElementById('apiKey').value.length > 0) {{
634 btn.classList.add('ready');
635 }} else {{
636 btn.classList.remove('ready');
637 }}
638 }}
639 function prepareSubmit() {{
640 // nothing extra needed — duration is already in the hidden input
641 }}
642</script>
643</body>
644</html>"""
647# ---------------------------------------------------------------------------
648# OAuth metadata discovery endpoints
649# ---------------------------------------------------------------------------
652def _byok_authorization_server_response(base_url: str, issuer: str) -> JSONResponse:
653 return JSONResponse(
654 {
655 "issuer": issuer,
656 "authorization_endpoint": f"{base_url}/v1/mcp/oauth/authorize",
657 "token_endpoint": f"{base_url}/v1/mcp/oauth/token",
658 "response_types_supported": ["code"],
659 "grant_types_supported": ["authorization_code"],
660 "code_challenge_methods_supported": ["S256"],
661 }
662 )
665@router.get("/.well-known/oauth-authorization-server", include_in_schema=False)
666async def oauth_authorization_server_metadata(request: Request) -> JSONResponse:
667 base_url: Final = get_request_base_url(request)
668 return _byok_authorization_server_response(base_url, base_url)
671@router.get(f"/.well-known/oauth-authorization-server{well_known_root_suffix()}/v1/mcp/oauth", include_in_schema=False)
672async def byok_authorization_server_metadata(request: Request) -> JSONResponse:
673 base_url: Final = get_request_base_url(request)
674 return _byok_authorization_server_response(base_url, f"{base_url}/v1/mcp/oauth")
677@router.get("/.well-known/oauth-authorization-server/{root_path:path}/v1/mcp/oauth", include_in_schema=False)
678async def byok_prefixed_authorization_server_metadata(request: Request, root_path: str) -> JSONResponse:
679 prefix: Final = f"/{root_path}"
680 if prefix not in get_server_root_paths():
681 raise HTTPException(status_code=404, detail="Unknown proxy root path")
682 parsed: Final = urlparse(get_request_base_url(request))
683 base_url: Final = f"{parsed.scheme}://{parsed.netloc}{prefix}"
684 return _byok_authorization_server_response(base_url, f"{base_url}/v1/mcp/oauth")
687@router.get(BYOK_RESOURCE_METADATA_PATH, include_in_schema=False)
688async def byok_protected_resource_metadata(request: Request) -> JSONResponse:
689 base_url: Final = get_request_base_url(request)
690 parsed: Final = urlparse(base_url)
691 return JSONResponse(
692 {
693 "resource": f"{parsed.scheme}://{parsed.netloc}",
694 "authorization_servers": (f"{base_url}/v1/mcp/oauth",),
695 }
696 )
699# ---------------------------------------------------------------------------
700# Authorization endpoint — GET (show form) and POST (process form)
701# ---------------------------------------------------------------------------
704@router.get("/v1/mcp/oauth/authorize", include_in_schema=False)
705async def byok_authorize_get(
706 request: Request,
707 client_id: str | None = None,
708 redirect_uri: str | None = None,
709 response_type: str | None = None,
710 code_challenge: str | None = None,
711 code_challenge_method: str | None = None,
712 state: str | None = None,
713 server_id: str | None = None,
714) -> HTMLResponse:
715 """
716 Show the BYOK API-key entry form.
718 The MCP client navigates the user here; the user types their API key and
719 clicks "Connect & Authorize", which POSTs back to this same path.
721 This GET is intentionally unauthenticated: it only renders HTML with no
722 state change. The POST handler enforces ``user_api_key_auth`` and pins
723 the stored credential to the authenticated session.
724 """
725 if response_type != "code":
726 raise HTTPException(status_code=400, detail="response_type must be 'code'")
727 if not redirect_uri:
728 raise HTTPException(status_code=400, detail="redirect_uri is required")
729 # Validate here too so the user sees the rejection before typing their
730 # API key into the HTML form (the POST handler also validates).
731 validate_loopback_redirect_uri(redirect_uri)
732 if not code_challenge:
733 raise HTTPException(status_code=400, detail="code_challenge is required")
735 # Resolve server metadata (name, description items, help URL).
736 server_name = "MCP Server"
737 access_items: list = []
738 help_url = ""
739 if server_id:
740 try:
741 from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
742 global_mcp_server_manager,
743 )
745 registry: Final = global_mcp_server_manager.get_registry()
746 if server_id in registry:
747 srv: Final = registry[server_id]
748 server_name = srv.server_name or srv.name
749 access_items = list(srv.byok_description or [])
750 help_url = srv.byok_api_key_help_url or ""
751 except Exception:
752 pass
754 server_initial: Final = (server_name[0].upper()) if server_name else "S"
756 html: Final = _build_authorize_html(
757 server_name=server_name,
758 server_initial=server_initial,
759 client_id=client_id or "",
760 redirect_uri=redirect_uri,
761 code_challenge=code_challenge,
762 code_challenge_method=code_challenge_method or "S256",
763 state=state or "",
764 server_id=server_id or "",
765 access_items=access_items,
766 help_url=help_url,
767 )
768 return HTMLResponse(content=html)
771@router.post("/v1/mcp/oauth/authorize", include_in_schema=False)
772async def byok_authorize_post(
773 request: Request,
774 client_id: str = Form(default=""),
775 redirect_uri: str = Form(...),
776 code_challenge: str = Form(...),
777 code_challenge_method: str = Form(default="S256"),
778 state: str = Form(default=""),
779 server_id: str = Form(default=""),
780 api_key: str = Form(...),
781 user_api_key_dict: UserAPIKeyAuth = Depends(_byok_session_auth),
782) -> RedirectResponse:
783 """
784 Process the BYOK API-key form submission.
786 Stores a short-lived authorization code and redirects the client back to
787 redirect_uri with ?code=...&state=... query parameters.
788 """
789 _purge_expired_codes()
791 validate_loopback_redirect_uri(redirect_uri)
793 # Reject new codes if the store is at capacity (prevents memory exhaustion
794 # from a burst of abandoned OAuth flows).
795 if len(_byok_auth_codes) >= _AUTH_CODES_MAX_SIZE:
796 raise HTTPException(status_code=503, detail="Too many pending authorization flows")
798 if code_challenge_method != "S256":
799 raise HTTPException(status_code=400, detail="Only S256 code_challenge_method is supported")
801 # Identity comes from the authenticated session, not the OAuth client_id
802 # form field (RFC 6749 §2.2: client_id identifies the client application,
803 # not the user). We do bind the code to the submitted client_id so the
804 # /token call must present the same value (RFC 6749 §4.1.3).
805 user_id: Final = user_api_key_dict.user_id
806 if not user_id:
807 raise HTTPException(status_code=401, detail="login_required")
809 auth_code: Final = str(uuid.uuid4())
810 _byok_auth_codes[auth_code] = {
811 "api_key": api_key,
812 "server_id": server_id,
813 "code_challenge": code_challenge,
814 "redirect_uri": redirect_uri,
815 # RFC 6749 §4.1.3 defense-in-depth: if the authorization request
816 # declared a client_id, the token request must submit the same
817 # value. Stored even though we don't pre-register clients.
818 "client_id": client_id,
819 "user_id": user_id,
820 "expires_at": time.time() + _AUTH_CODE_TTL_SECONDS,
821 }
823 params: Final = urlencode({"code": auth_code, "state": state})
824 separator: Final = "&" if "?" in redirect_uri else "?"
825 location: Final = f"{redirect_uri}{separator}{params}"
826 return RedirectResponse(url=location, status_code=302)
829# ---------------------------------------------------------------------------
830# Token endpoint
831# ---------------------------------------------------------------------------
834@router.post("/v1/mcp/oauth/token", include_in_schema=False)
835async def byok_token(
836 request: Request,
837 grant_type: str = Form(...),
838 code: str = Form(...),
839 redirect_uri: str = Form(default=""),
840 code_verifier: str = Form(...),
841 client_id: str = Form(default=""),
842) -> JSONResponse:
843 """
844 Exchange an authorization code for a short-lived BYOK session JWT.
846 1. Validates the authorization code and PKCE challenge.
847 2. Stores the API key via store_user_credential().
848 3. Issues a signed JWT with type="byok_session".
849 """
850 from litellm.proxy.proxy_server import master_key, prisma_client
852 _purge_expired_codes()
854 if grant_type != "authorization_code":
855 return _oauth_token_error("unsupported_grant_type")
857 record: Final = _byok_auth_codes.get(code)
858 if record is None:
859 return _oauth_token_error("invalid_grant")
861 if time.time() > record["expires_at"]:
862 del _byok_auth_codes[code]
863 return _oauth_token_error("invalid_grant")
865 # PKCE verification
866 if not _verify_pkce(code_verifier, record["code_challenge"]):
867 return _oauth_token_error("invalid_grant")
869 # RFC 6749 §4.1.3: if redirect_uri was sent with the authorization
870 # request, the token request MUST include the identical value.
871 # OAuth 2.1 draft-15 §4.1.3 drops this requirement — strict OAuth 2.1
872 # clients will omit it. Enforce equality ONLY when the client
873 # actually submitted a value, so we stay RFC 6749-backward-compatible
874 # without breaking OAuth 2.1 clients. PKCE + client_id binding
875 # (checked below) cover the security role redirect_uri played.
876 if record.get("redirect_uri") and redirect_uri and redirect_uri != record["redirect_uri"]:
877 return _oauth_token_error("invalid_grant")
879 # RFC 6749 §4.1.3: if the client was identified at /authorize, the
880 # /token request MUST authenticate as the same client. We don't
881 # pre-register clients, so an empty stored client_id skips the check.
882 if record.get("client_id") and client_id != record["client_id"]:
883 return _oauth_token_error("invalid_grant")
885 server_id: Final[str] = record["server_id"]
886 api_key_value: Final[str] = record["api_key"]
887 # user_id is stamped by the authenticated /authorize POST. No client_id
888 # fallback — that fallback was the credential-hijack primitive. The
889 # token-endpoint client_id is informational per RFC 6749 and is not
890 # cross-checked against user_id (which identifies the resource owner,
891 # not the client application).
892 user_id: Final[str] = record.get("user_id") or ""
893 if not user_id:
894 return _oauth_token_error("invalid_grant")
896 # Verify preconditions that would fail token issuance BEFORE consuming
897 # the code or writing to the DB — otherwise a misconfigured proxy
898 # (missing master_key) silently persists the user's credential without
899 # ever returning an access token, and the user has no way to recover.
900 if master_key is None:
901 return _oauth_token_error("server_error", status=500)
903 # Consume the code (one-time use)
904 del _byok_auth_codes[code]
906 # Persist the BYOK credential
907 if prisma_client is not None:
908 try:
909 await store_user_credential(
910 prisma_client=prisma_client,
911 user_id=user_id,
912 server_id=server_id,
913 credential=api_key_value,
914 )
915 # Invalidate any cached negative result so the user isn't blocked
916 # for up to the TTL period after completing the OAuth flow.
917 from litellm.proxy._experimental.mcp_server.server import (
918 _invalidate_byok_cred_cache,
919 )
921 await _invalidate_byok_cred_cache(user_id, server_id)
922 except Exception as exc:
923 verbose_proxy_logger.error(
924 "byok_token: failed to store user credential for user=%s server=%s: %s",
925 user_id,
926 server_id,
927 exc,
928 )
929 return _oauth_token_error("server_error", status=500)
930 else:
931 verbose_proxy_logger.warning("byok_token: prisma_client is None — credential not persisted")
933 now: Final = int(time.time())
934 payload: Final = {
935 "user_id": user_id,
936 "server_id": server_id,
937 # "type" distinguishes this from regular proxy auth tokens.
938 # The proxy's SSO JWT path uses asymmetric keys (RS256/ES256), so an
939 # HS256 token signed with master_key cannot be accepted there.
940 "type": "byok_session",
941 "iat": now,
942 "exp": now + 3600,
943 }
944 access_token: Final = jwt.encode(payload, cast(str, master_key), algorithm="HS256")
946 return JSONResponse(
947 {
948 "access_token": access_token,
949 "token_type": "bearer",
950 "expires_in": 3600,
951 },
952 headers=TOKEN_NO_CACHE_HEADERS,
953 )