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

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:

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