Merge efa948379a into 08f2a98250

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"

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
 # ---------------------------------------------------------------------------