UI / SSO - Update proxy admin id role in DB + Handle SSO redirects with custom root path (#11384)

* fix(ui_sso.py): update user as proxy admin in db table, when checking for proxy_admin_id

Fixes issue where existing internal user, unable to make calls when set as proxy admin id

* fix(utils.py): fix custom base path

docs/my-website/docs/proxy/admin_ui_sso.md

@@ -186,6 +186,10 @@ Set a Proxy Admin when SSO is enabled. Once SSO is enabled, the `user_id` for us
 export PROXY_ADMIN_ID="116544810872468347480"
 ```
 
+This will update the user role in the `LiteLLM_UserTable` to `proxy_admin`. 
+
+If you plan to change this ID, please update the user role via API `/user/update` or UI (Internal Users page). 
+
 #### Step 3: See all proxy keys
 
 <Image img={require('../../img/litellm_ui_admin.png')} />

litellm/proxy/management_endpoints/ui_sso.py

@@ -433,6 +433,29 @@ def apply_user_info_values_to_sso_user_defined_values(
     return user_defined_values
 
 
+async def check_and_update_if_proxy_admin_id(
+    user_role: str, user_id: str, prisma_client: Optional[PrismaClient]
+):
+    """
+    - Check if user role in DB is admin
+    - If not, update user role in DB to admin role
+    """
+    proxy_admin_id = os.getenv("PROXY_ADMIN_ID")
+    if proxy_admin_id is not None and proxy_admin_id == user_id:
+        if user_role and user_role == LitellmUserRoles.PROXY_ADMIN.value:
+            return user_role
+
+        if prisma_client:
+            await prisma_client.db.litellm_usertable.update(
+                where={"user_id": user_id},
+                data={"user_role": LitellmUserRoles.PROXY_ADMIN.value},
+            )
+
+        user_role = LitellmUserRoles.PROXY_ADMIN.value
+
+    return user_role
+
+
 @router.get("/sso/callback", tags=["experimental"], include_in_schema=False)
 async def auth_callback(request: Request):  # noqa: PLR0915
     """Verify login"""
@@ -451,6 +474,7 @@ async def auth_callback(request: Request):  # noqa: PLR0915
         user_api_key_cache,
         user_custom_sso,
     )
+    from litellm.proxy.utils import get_custom_url
     from litellm.types.proxy.ui_sso import ReturnedUITokenObject
 
     if prisma_client is None:
@@ -469,12 +493,11 @@ async def auth_callback(request: Request):  # noqa: PLR0915
             param="master_key",
             code=status.HTTP_500_INTERNAL_SERVER_ERROR,
         )
-    redirect_url = os.getenv("PROXY_BASE_URL", str(request.base_url))
-    if redirect_url.endswith("/"):
-        redirect_url += "sso/callback"
-    else:
-        redirect_url += "/sso/callback"
+    redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso(
+        request=request, sso_callback_route="sso/callback"
+    )
 
+    verbose_proxy_logger.info(f"Redirecting to {redirect_url}")
     result = None
     if google_client_id is not None:
         result = await GoogleSSOHandler.get_google_callback_response(
@@ -602,17 +625,17 @@ async def auth_callback(request: Request):  # noqa: PLR0915
     key = response["token"]  # type: ignore
     user_id = response["user_id"]  # type: ignore
 
-    litellm_dashboard_ui = "/ui/"
+    litellm_dashboard_ui = get_custom_url(
+        request_base_url=str(request.base_url), route="ui/"
+    )
     user_role = (
         user_defined_values["user_role"]
         or LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value
     )
-    if (
-        os.getenv("PROXY_ADMIN_ID", None) is not None
-        and os.environ["PROXY_ADMIN_ID"] == user_id
-    ):
-        # checks if user is admin
-        user_role = LitellmUserRoles.PROXY_ADMIN.value
+    if user_id and isinstance(user_id, str):
+        user_role = await check_and_update_if_proxy_admin_id(
+            user_role=user_role, user_id=user_id, prisma_client=prisma_client
+        )
 
     verbose_proxy_logger.debug(
         f"user_role: {user_role}; ui_access_mode: {ui_access_mode}"
@@ -947,7 +970,9 @@ class SSOAuthenticationHandler:
         """
         Get the redirect URL for SSO
         """
-        redirect_url = os.getenv("PROXY_BASE_URL", str(request.base_url))
+        from litellm.proxy.utils import get_custom_url
+
+        redirect_url = get_custom_url(request_base_url=str(request.base_url))
         if redirect_url.endswith("/"):
             redirect_url += sso_callback_route
         else:

litellm/proxy/utils.py

@@ -2864,7 +2864,16 @@ def is_known_model(model: Optional[str], llm_router: Optional[Router]) -> bool:
     return is_in_list
 
 
-def get_custom_url(request_base_url: str) -> str:
+def join_paths(base_path: str, route: str) -> str:
+    # Remove trailing/leading slashes
+    base_path = base_path.rstrip("/")
+    route = route.lstrip("/")
+
+    # Join with a single slash
+    return f"{base_path}/{route}"
+
+
+def get_custom_url(request_base_url: str, route: Optional[str] = None) -> str:
     """
     Use proxy base url, if set.
 
@@ -2874,6 +2883,8 @@ def get_custom_url(request_base_url: str) -> str:
 
     proxy_base_url = os.getenv("PROXY_BASE_URL")
     server_root_path = os.getenv("SERVER_ROOT_PATH") or ""
+    if route is not None:
+        server_root_path = join_paths(base_path=server_root_path, route=route)
     if proxy_base_url:
         ui_link = str(URL(proxy_base_url).join(server_root_path))
     else:

tests/test_litellm/proxy/management_endpoints/test_ui_sso.py

@@ -629,3 +629,64 @@ async def test_get_user_info_from_db_alternate_user_id():
         user_info = await get_user_info_from_db(**args)
         mock_get_user_object.assert_called_once()
         mock_get_user_object.call_args.kwargs["user_id"] = "krrishd-email1234"
+
+
+@pytest.mark.asyncio
+async def test_check_and_update_if_proxy_admin_id():
+    """
+    Test that a user with matching PROXY_ADMIN_ID gets their role updated to admin
+    """
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.management_endpoints.ui_sso import (
+        check_and_update_if_proxy_admin_id,
+    )
+
+    # Mock Prisma client
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_usertable.update = AsyncMock()
+
+    # Set up test data
+    test_user_id = "test_admin_123"
+    test_user_role = "user"
+
+    with patch.dict(os.environ, {"PROXY_ADMIN_ID": test_user_id}):
+        # Act
+        updated_role = await check_and_update_if_proxy_admin_id(
+            user_role=test_user_role, user_id=test_user_id, prisma_client=mock_prisma
+        )
+
+        # Assert
+        assert updated_role == LitellmUserRoles.PROXY_ADMIN.value
+        mock_prisma.db.litellm_usertable.update.assert_called_once_with(
+            where={"user_id": test_user_id},
+            data={"user_role": LitellmUserRoles.PROXY_ADMIN.value},
+        )
+
+
+@pytest.mark.asyncio
+async def test_check_and_update_if_proxy_admin_id_already_admin():
+    """
+    Test that a user who is already an admin doesn't get their role updated
+    """
+    from litellm.proxy._types import LitellmUserRoles
+    from litellm.proxy.management_endpoints.ui_sso import (
+        check_and_update_if_proxy_admin_id,
+    )
+
+    # Mock Prisma client
+    mock_prisma = MagicMock()
+    mock_prisma.db.litellm_usertable.update = AsyncMock()
+
+    # Set up test data
+    test_user_id = "test_admin_123"
+    test_user_role = LitellmUserRoles.PROXY_ADMIN.value
+
+    with patch.dict(os.environ, {"PROXY_ADMIN_ID": test_user_id}):
+        # Act
+        updated_role = await check_and_update_if_proxy_admin_id(
+            user_role=test_user_role, user_id=test_user_id, prisma_client=mock_prisma
+        )
+
+        # Assert
+        assert updated_role == LitellmUserRoles.PROXY_ADMIN.value
+        mock_prisma.db.litellm_usertable.update.assert_not_called()

tests/test_litellm/proxy/test_proxy_utils.py

@@ -0,0 +1,21 @@
+import json
+import os
+import sys
+
+import pytest
+from fastapi.testclient import TestClient
+
+sys.path.insert(
+    0, os.path.abspath("../../..")
+)  # Adds the parent directory to the system path
+
+
+from unittest.mock import MagicMock
+
+from litellm.proxy.utils import get_custom_url
+
+
+def test_get_custom_url(monkeypatch):
+    monkeypatch.setenv("SERVER_ROOT_PATH", "/litellm")
+    custom_url = get_custom_url(request_base_url="http://0.0.0.0:4000", route="ui/")
+    assert custom_url == "http://0.0.0.0:4000/litellm/ui/"