diff --git a/litellm/proxy/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/mcp_server/byok_oauth_endpoints.py index 868b73e39f1..fda132f610d 100644 --- a/litellm/proxy/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/mcp_server/byok_oauth_endpoints.py @@ -22,7 +22,7 @@ from urllib.parse import urlencode, urlparse import jwt from fastapi import APIRouter, Depends, Form, HTTPException, Request -from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth @@ -224,6 +224,28 @@ async def _byok_session_auth(request: Request) -> UserAPIKeyAuth: return UserAPIKeyAuth(api_key="byok_session_cookie", user_id=user_id) +async def _byok_session_user_id_or_none(request: Request) -> str | None: + """Same check as ``_byok_session_auth``, but returns ``None`` instead of + raising, so the authorize GET can send a signed-out browser to the login + page before the user types an API key the POST would then reject.""" + user_id, session_key = _session_identity_from_cookie(request) + if not user_id or not await _session_key_is_live(session_key): + return None + return user_id + + +def _redirect_to_ui_login(request: Request) -> RedirectResponse: + """Send the browser through the UI login and back to this authorize request. + + ``return_to`` is the request's own relative path and query, so it is + same-origin by construction; the login endpoints only honor relative + same-origin paths (or the configured control-plane origin).""" + from litellm.proxy.mcp_server.gateway_dcr_flow import relative_request_url + + return_to: Final = urlencode({"return_to": relative_request_url(request)}) + return RedirectResponse(f"{get_request_base_url(request)}/sso/key/generate?{return_to}", status_code=303) + + def _build_authorize_html( server_name: str, server_initial: str, @@ -752,16 +774,18 @@ async def byok_authorize_get( code_challenge_method: str | None = None, state: str | None = None, server_id: str | None = None, -) -> HTMLResponse: + session_user_id: str | None = Depends(_byok_session_user_id_or_none), +) -> Response: """ 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. - This GET is intentionally unauthenticated: it only renders HTML with no - state change. The POST handler enforces ``user_api_key_auth`` and pins - the stored credential to the authenticated session. + This GET makes no state change. The POST handler requires the UI session + and pins the stored credential to it, so a browser without a valid UI + session is first redirected to the UI login, which returns here after + sign-in instead of letting the user type a key the POST would reject. """ if response_type != "code": raise HTTPException(status_code=400, detail="response_type must be 'code'") @@ -772,6 +796,8 @@ async def byok_authorize_get( validate_loopback_redirect_uri(redirect_uri) if not code_challenge: raise HTTPException(status_code=400, detail="code_challenge is required") + if session_user_id is None: + return _redirect_to_ui_login(request) # Resolve server metadata (name, description items, help URL). server_name = "MCP Server" diff --git a/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py b/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py index 37da00ab2ba..4cbe5d3f9e7 100644 --- a/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py @@ -20,6 +20,7 @@ import uuid from typing import Any, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException from fastapi.testclient import TestClient @@ -68,6 +69,7 @@ from fastapi import FastAPI from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.mcp_server.byok_oauth_endpoints import ( _byok_session_auth, + _byok_session_user_id_or_none, ) _test_app = FastAPI() @@ -81,10 +83,12 @@ def client(): _test_app.dependency_overrides[_byok_session_auth] = lambda: UserAPIKeyAuth( api_key="hashed", user_id="user-123" ) + _test_app.dependency_overrides[_byok_session_user_id_or_none] = lambda: "user-123" try: yield TestClient(_test_app, raise_server_exceptions=False) finally: _test_app.dependency_overrides.pop(_byok_session_auth, None) + _test_app.dependency_overrides.pop(_byok_session_user_id_or_none, None) @pytest.fixture @@ -285,6 +289,110 @@ def test_authorize_get_rejects_non_loopback_redirect_uri(client): assert resp.status_code == 400 +_AUTHORIZE_GET_PARAMS: Final = { + "client_id": "test-client", + "redirect_uri": "http://127.0.0.1:3000/callback", + "response_type": "code", + "code_challenge": "abc123", + "code_challenge_method": "S256", + "state": "xyz", + "server_id": "my-server", +} + + +# Fixed expiry (2100-01-01) so the cookie fixtures do not depend on the host clock. +_SESSION_COOKIE_EXP: Final = 4102444800 + + +def _assert_redirects_to_login_and_back(resp: httpx.Response) -> None: + from urllib.parse import parse_qs, urlparse + + assert resp.status_code == 303 + location: Final = urlparse(resp.headers["location"]) + assert location.path == "/sso/key/generate" + return_to: Final = parse_qs(location.query)["return_to"][0] + assert return_to.startswith("/v1/mcp/oauth/authorize?") + assert parse_qs(urlparse(return_to).query) == {k: [v] for k, v in _AUTHORIZE_GET_PARAMS.items()} + + +def test_authorize_get_without_session_redirects_to_login(unauthenticated_client: TestClient) -> None: + """Without a UI session the POST would reject the form with 401 + ``login_required`` after the user typed the key, so the GET sends the + browser to the UI login first, with a return path back to this page.""" + with patch("litellm.proxy.proxy_server.master_key", "test-master-key"): + resp: Final = unauthenticated_client.get( + "/v1/mcp/oauth/authorize", params=_AUTHORIZE_GET_PARAMS, follow_redirects=False + ) + _assert_redirects_to_login_and_back(resp) + assert "apiKey" not in resp.text + + +def test_authorize_get_with_revoked_session_redirects_to_login(unauthenticated_client: TestClient) -> None: + """A signature-valid cookie whose session key no longer resolves is not a session.""" + import jwt as _jwt + + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + # Stand-in for the database: the revoked session key has no row. MCP + # servers stay config-only so the catalog does not query it. + fake_prisma_client: Final = MagicMock() + fake_prisma_client.get_data = AsyncMock(return_value=None) + + with ( + patch("litellm.proxy.proxy_server.master_key", "test-master-key"), + patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), + patch("litellm.proxy.proxy_server.user_api_key_cache", UserApiKeyCache()), + patch.dict("litellm.proxy.proxy_server.general_settings", {"supported_db_objects": ["models"]}), + ): + cookie_jwt: Final = _jwt.encode( + { + "user_id": "browser-user-42", + "key": "sk-revoked-session-key", + "login_method": "sso", + "exp": _SESSION_COOKIE_EXP, + }, + "test-master-key", + algorithm="HS256", + ) + resp: Final = unauthenticated_client.get( + "/v1/mcp/oauth/authorize", + params=_AUTHORIZE_GET_PARAMS, + cookies={"token": cookie_jwt}, + follow_redirects=False, + ) + _assert_redirects_to_login_and_back(resp) + fake_prisma_client.get_data.assert_awaited() + + +def test_authorize_get_with_session_cookie_renders_form(unauthenticated_client: TestClient) -> None: + import jwt as _jwt + + with patch("litellm.proxy.proxy_server.master_key", "test-master-key"): + cookie_jwt: Final = _jwt.encode( + {"user_id": "browser-user-42", "login_method": "sso", "exp": _SESSION_COOKIE_EXP}, + "test-master-key", + algorithm="HS256", + ) + resp: Final = unauthenticated_client.get( + "/v1/mcp/oauth/authorize", + params=_AUTHORIZE_GET_PARAMS, + cookies={"token": cookie_jwt}, + follow_redirects=False, + ) + assert resp.status_code == 200 + assert "Connect & Authorize" in resp.text + + +def test_authorize_get_validates_request_before_login_redirect(unauthenticated_client: TestClient) -> None: + """A malformed request is rejected outright rather than bounced through login.""" + resp: Final = unauthenticated_client.get( + "/v1/mcp/oauth/authorize", + params={**_AUTHORIZE_GET_PARAMS, "redirect_uri": "https://attacker.example.com/cb"}, + follow_redirects=False, + ) + assert resp.status_code == 400 + + # --------------------------------------------------------------------------- # Authorization POST endpoint # ---------------------------------------------------------------------------