From 9bb292f478dab11c0cd8088fde353fc349fb8e93 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 3 Dec 2025 12:41:45 -0800 Subject: [PATCH] V2 login route --- litellm/proxy/proxy_server.py | 47 +++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 52 +++++++++++++++++++ 2 files changed, 99 insertions(+) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9789f495508..4d971e8ce42 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8318,6 +8318,53 @@ async def login(request: Request): # noqa: PLR0915 return redirect_response +@router.post( + "/v2/login", include_in_schema=False +) # hidden helper for UI logins via API +async def login_v2(request: Request): # noqa: PLR0915 + global premium_user, general_settings, master_key + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.utils import get_custom_url + + body = await request.json() + username = str(body.get("username")) + password = str(body.get("password")) + + login_result = await authenticate_user( + username=username, + password=password, + master_key=master_key, + prisma_client=prisma_client, + ) + + returned_ui_token_object = create_ui_token_object( + login_result=login_result, + general_settings=general_settings, + premium_user=premium_user, + ) + + import jwt + + jwt_token = jwt.encode( # type: ignore + cast(dict, returned_ui_token_object), + master_key, + algorithm="HS256", + ) + + litellm_dashboard_ui = get_custom_url(str(request.base_url)) + if litellm_dashboard_ui.endswith("/"): + litellm_dashboard_ui += "ui/" + else: + litellm_dashboard_ui += "/ui/" + litellm_dashboard_ui += "?login=success" + + json_response = JSONResponse( + content={"redirect_url": litellm_dashboard_ui}, + status_code=status.HTTP_200_OK, + ) + json_response.set_cookie(key="token", value=jwt_token) + return json_response + @app.get("/onboarding/get_token", include_in_schema=False) async def onboarding(invite_link: str, request: Request): """ diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8e87b67933f..7b3157f8ee1 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -71,6 +71,58 @@ def client_no_auth(): return TestClient(app) +def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): + mock_login_result = {"user_id": "test-user"} + mock_prisma_client = MagicMock() + mock_authenticate_user = AsyncMock(return_value=mock_login_result) + mock_create_ui_token_object = MagicMock(return_value={"user_id": "test-user"}) + mock_jwt_encode = MagicMock(return_value="signed-token") + + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + mock_authenticate_user, + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + mock_create_ui_token_object, + ) + monkeypatch.setattr("jwt.encode", mock_jwt_encode) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + client = TestClient(app) + response = client.post( + "/v2/login", + json={"username": "alice", "password": "secret"}, + ) + + assert response.status_code == 200 + assert ( + response.json() + == {"redirect_url": "http://testserver/ui/?login=success"} + ) + assert response.cookies.get("token") == "signed-token" + + mock_authenticate_user.assert_awaited_once_with( + username="alice", + password="secret", + master_key="test-master-key", + prisma_client=mock_prisma_client, + ) + mock_create_ui_token_object.assert_called_once_with( + login_result=mock_login_result, + general_settings={}, + premium_user=False, + ) + mock_jwt_encode.assert_called_once_with( + {"user_id": "test-user"}, + "test-master-key", + algorithm="HS256", + ) + + @pytest.mark.asyncio async def test_initialize_scheduled_jobs_credentials(monkeypatch): """