From 7bca988c78fd125d9e14234c98d2472d0389e559 Mon Sep 17 00:00:00 2001 From: Artur Sidorenko Date: Mon, 2 Nov 2026 12:32:19 +0100 Subject: [PATCH 1/2] feat(mcp): send signed-out browsers to the UI login before the BYOK key form GET /v1/mcp/oauth/authorize rendered the API-key form without a UI session, while the POST rejects such a submission with 401 login_required, so the user typed the key only to lose it. The GET now redirects a browser without a live UI session to /sso/key/generate with a same-origin relative return_to, the same round-trip the MCP gateway authorize flows use, and the login brings the browser back to the authorize request. --- .../proxy/mcp_server/byok_oauth_endpoints.py | 36 ++++++- .../mcp_server/test_byok_oauth_endpoints.py | 97 +++++++++++++++++++ 2 files changed, 128 insertions(+), 5 deletions(-) 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..c89bdf15998 100644 --- a/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py @@ -68,6 +68,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 +82,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 +288,100 @@ 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", +} + + +def _assert_redirects_to_login_and_back(resp) -> None: + from urllib.parse import parse_qs, urlparse + + assert resp.status_code == 303 + location = urlparse(resp.headers["location"]) + assert location.path == "/sso/key/generate" + return_to = parse_qs(location.query)["return_to"][0] + # Relative, same-origin path back to the exact authorize request. + 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): + """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 = 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): + """A signature-valid cookie whose session key no longer resolves is not a session.""" + import jwt as _jwt + + with ( + patch("litellm.proxy.proxy_server.master_key", "test-master-key"), + patch( + "litellm.proxy.mcp_server.byok_oauth_endpoints._session_key_is_live", + new=AsyncMock(return_value=False), + ), + ): + cookie_jwt = _jwt.encode( + { + "user_id": "browser-user-42", + "key": "sk-revoked-session-key", + "login_method": "sso", + "exp": int(time.time()) + 3600, + }, + "test-master-key", + algorithm="HS256", + ) + resp = 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) + + +def test_authorize_get_with_session_cookie_renders_form(unauthenticated_client): + import jwt as _jwt + + with patch("litellm.proxy.proxy_server.master_key", "test-master-key"): + cookie_jwt = _jwt.encode( + {"user_id": "browser-user-42", "login_method": "sso", "exp": int(time.time()) + 3600}, + "test-master-key", + algorithm="HS256", + ) + resp = 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): + """A malformed request is rejected outright rather than bounced through login.""" + resp = 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 # --------------------------------------------------------------------------- From efa948379a2812209e776808433334c88c2d655e Mon Sep 17 00:00:00 2001 From: Artur Sidorenko Date: Fri, 6 Nov 2026 08:19:22 +0100 Subject: [PATCH 2/2] test(mcp): exercise the real session-key lookup and fix cookie expiry in BYOK authorize tests --- .../mcp_server/test_byok_oauth_endpoints.py | 51 +++++++++++-------- 1 file changed, 31 insertions(+), 20 deletions(-) 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 c89bdf15998..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 @@ -299,70 +300,80 @@ _AUTHORIZE_GET_PARAMS: Final = { } -def _assert_redirects_to_login_and_back(resp) -> None: +# 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 = urlparse(resp.headers["location"]) + location: Final = urlparse(resp.headers["location"]) assert location.path == "/sso/key/generate" - return_to = parse_qs(location.query)["return_to"][0] - # Relative, same-origin path back to the exact authorize request. + 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): +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 = unauthenticated_client.get( + 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): +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.mcp_server.byok_oauth_endpoints._session_key_is_live", - new=AsyncMock(return_value=False), - ), + 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 = _jwt.encode( + cookie_jwt: Final = _jwt.encode( { "user_id": "browser-user-42", "key": "sk-revoked-session-key", "login_method": "sso", - "exp": int(time.time()) + 3600, + "exp": _SESSION_COOKIE_EXP, }, "test-master-key", algorithm="HS256", ) - resp = unauthenticated_client.get( + 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): +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 = _jwt.encode( - {"user_id": "browser-user-42", "login_method": "sso", "exp": int(time.time()) + 3600}, + cookie_jwt: Final = _jwt.encode( + {"user_id": "browser-user-42", "login_method": "sso", "exp": _SESSION_COOKIE_EXP}, "test-master-key", algorithm="HS256", ) - resp = unauthenticated_client.get( + resp: Final = unauthenticated_client.get( "/v1/mcp/oauth/authorize", params=_AUTHORIZE_GET_PARAMS, cookies={"token": cookie_jwt}, @@ -372,9 +383,9 @@ def test_authorize_get_with_session_cookie_renders_form(unauthenticated_client): assert "Connect & Authorize" in resp.text -def test_authorize_get_validates_request_before_login_redirect(unauthenticated_client): +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 = unauthenticated_client.get( + resp: Final = unauthenticated_client.get( "/v1/mcp/oauth/authorize", params={**_AUTHORIZE_GET_PARAMS, "redirect_uri": "https://attacker.example.com/cb"}, follow_redirects=False,