From 501f8c787294a95e18cecb92af95370bbf3cf3b3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 05:24:59 -0700 Subject: [PATCH] fix(cli): finish lite login for the admin when PROXY_ADMIN_ID is set The admin's browser cookie names default_user_id while its session key belongs to PROXY_ADMIN_ID, so the CLI login page kept sending the admin back to the sign-in form --- .../proxy/mcp_server/byok_oauth_endpoints.py | 27 +++++++-- .../mcp_server/test_byok_oauth_endpoints.py | 55 +++++++++++++++++++ 2 files changed, 78 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/mcp_server/byok_oauth_endpoints.py index 41823f5b799..61fd5f53ff2 100644 --- a/litellm/proxy/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/mcp_server/byok_oauth_endpoints.py @@ -15,6 +15,7 @@ Endpoints implemented here: import base64 import hashlib import html as _html_module +import os import time import uuid from typing import Final, cast @@ -25,6 +26,7 @@ from fastapi import APIRouter, Depends, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse from litellm._logging import verbose_proxy_logger +from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.mcp_server.catalog import public_catalog_operation from litellm.proxy.mcp_server.db import store_user_credential @@ -176,11 +178,28 @@ async def _session_key_is_live(session_key: str | None) -> bool: async def get_authenticated_browser_user_id(request: Request) -> str | None: - session: Final = await get_authenticated_browser_session(request) - return None if session is None else session.user_id + session: Final = await _live_browser_session(request) + if session is None: + return None + cookie_user_id, auth = session + return cookie_user_id if auth.user_id == cookie_user_id else None async def get_authenticated_browser_session(request: Request) -> UserAPIKeyAuth | None: + session: Final = await _live_browser_session(request) + if session is None: + return None + cookie_user_id, auth = session + return auth if auth.user_id in _session_key_owners(cookie_user_id) else None + + +def _session_key_owners(cookie_user_id: str) -> frozenset[str]: + if cookie_user_id != LITELLM_PROXY_ADMIN_NAME: + return frozenset({cookie_user_id}) + return frozenset({cookie_user_id, os.getenv("PROXY_ADMIN_ID", LITELLM_PROXY_ADMIN_NAME)}) + + +async def _live_browser_session(request: Request) -> tuple[str, UserAPIKeyAuth] | None: from datetime import datetime, timezone from pydantic import TypeAdapter, ValidationError @@ -206,14 +225,14 @@ async def get_authenticated_browser_session(request: Request) -> UserAPIKeyAuth ) except Exception: return None - if auth is None or auth.user_id != user_id or auth.blocked or auth.expires is None: + if auth is None or auth.blocked or auth.expires is None: return None try: expiration: Final = TypeAdapter(datetime).validate_python(auth.expires) except ValidationError: return None expires: Final = expiration.replace(tzinfo=timezone.utc) if expiration.tzinfo is None else expiration - return auth if expires > datetime.now(timezone.utc) else None + return (user_id, auth) if expires > datetime.now(timezone.utc) else None async def _byok_session_auth(request: Request) -> UserAPIKeyAuth: 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..21fef847f67 100644 --- a/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py +++ b/tests/unit/proxy/mcp_server/test_byok_oauth_endpoints.py @@ -1994,3 +1994,58 @@ async def test_authenticated_browser_identity_checks_embedded_session( ) request: Final = Request({"type": "http", "headers": [(b"cookie", ("token=" + cookie).encode())]}) assert await get_authenticated_browser_user_id(request) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "cookie_user_id,key_owner,proxy_admin_id,cli_owner,byok_user_id", + [ + ("default_user_id", "litellm-qa-admin-id", "litellm-qa-admin-id", "litellm-qa-admin-id", None), + ("default_user_id", "default_user_id", "litellm-qa-admin-id", "default_user_id", "default_user_id"), + ("default_user_id", "bob", "litellm-qa-admin-id", None, None), + ("default_user_id", "litellm-qa-admin-id", None, None, None), + ("alice", "litellm-qa-admin-id", "litellm-qa-admin-id", None, None), + ("alice", "alice", "litellm-qa-admin-id", "alice", "alice"), + ], +) +async def test_cli_browser_session_accepts_the_proxy_admin_key_that_login_issues( + monkeypatch: pytest.MonkeyPatch, + cookie_user_id: str, + key_owner: str, + proxy_admin_id: str | None, + cli_owner: str | None, + byok_user_id: str | None, +) -> None: + from datetime import datetime, timezone + + import jwt + from starlette.requests import Request + + from litellm.proxy import proxy_server + from litellm.proxy.auth.auth_checks import LITELLM_SESSION_TOKEN_PREFIX + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token + from litellm.proxy.mcp_server.byok_oauth_endpoints import ( + get_authenticated_browser_session, + get_authenticated_browser_user_id, + ) + + monkeypatch.setenv("LITELLM_SALT_KEY", "local-test-encryption-key") + if proxy_admin_id is None: + monkeypatch.delenv("PROXY_ADMIN_ID", raising=False) + else: + monkeypatch.setenv("PROXY_ADMIN_ID", proxy_admin_id) + monkeypatch.setattr(proxy_server, "master_key", "local-test-cookie-signing-key-32-bytes") + monkeypatch.setattr(proxy_server, "prisma_client", object()) + session: Final = UserAPIKeyAuth( + user_id=key_owner, expires=datetime(2099, 1, 1, tzinfo=timezone.utc), is_session_token=True + ) + key: Final = encrypt_bearer_token(session.model_dump_json(exclude_none=True), LITELLM_SESSION_TOKEN_PREFIX) + cookie: Final = jwt.encode( + {"user_id": cookie_user_id, "key": key, "login_method": "username_password", "exp": 9999999999}, + "local-test-cookie-signing-key-32-bytes", + algorithm="HS256", + ) + request: Final = Request({"type": "http", "headers": [(b"cookie", ("token=" + cookie).encode())]}) + cli_session: Final = await get_authenticated_browser_session(request) + assert (None if cli_session is None else cli_session.user_id) == cli_owner + assert await get_authenticated_browser_user_id(request) == byok_user_id