diff --git a/docs/my-website/docs/proxy/admin_ui_sso.md b/docs/my-website/docs/proxy/admin_ui_sso.md index a0dde80e9cf..b8aa152ed8e 100644 --- a/docs/my-website/docs/proxy/admin_ui_sso.md +++ b/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 diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 24b13c4ecda..36ab13c693e 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/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: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 710c181b08a..bca368f3482 100644 --- a/litellm/proxy/utils.py +++ b/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: diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 75bf17cfd62..47684e64d2e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/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() diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py new file mode 100644 index 00000000000..b6e9f6f6ee6 --- /dev/null +++ b/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/"