diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3927380c8ae..0a1cfda1502 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8382,7 +8382,7 @@ async def async_queue_request( ): global general_settings, user_debug, proxy_logging_obj """ - v2 attempt at a background worker to handle queuing. + v2 attempt at a background worker to handle queuing Just supports /chat/completion calls currently. @@ -8575,44 +8575,68 @@ async def login_v2(request: Request): # noqa: PLR0915 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")) + try: + 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, - ) + 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, - ) + returned_ui_token_object = create_ui_token_object( + login_result=login_result, + general_settings=general_settings, + premium_user=premium_user, + ) - import jwt + import jwt - jwt_token = jwt.encode( - cast(dict, returned_ui_token_object), - cast(str, master_key), - algorithm="HS256", - ) + jwt_token = jwt.encode( + cast(dict, returned_ui_token_object), + cast(str, 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" + 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 + 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 + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.login_v2(): Exception occurred - {}".format( + str(e) + ) + ) + if isinstance(e, ProxyException): + raise e + elif isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", str(e)), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), + ) + else: + error_msg = f"{str(e)}" + raise ProxyException( + message=error_msg, + type=ProxyErrorTypes.auth_error, + param="None", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) @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 05362e50d6a..1f81026b537 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -126,6 +126,114 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): ) +def test_login_v2_returns_json_on_proxy_exception(monkeypatch): + """Test that /v2/login returns JSON error when ProxyException is raised""" + from litellm.proxy._types import ProxyException, ProxyErrorTypes + + mock_prisma_client = MagicMock() + mock_authenticate_user = AsyncMock( + side_effect=ProxyException( + message="Invalid credentials", + type=ProxyErrorTypes.auth_error, + param="password", + code=401, + ) + ) + + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + mock_authenticate_user, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + client = TestClient(app) + response = client.post( + "/v2/login", + json={"username": "alice", "password": "wrong"}, + ) + + assert response.status_code == 401 + assert response.headers["content-type"] == "application/json" + data = response.json() + assert "error" in data + assert data["error"]["message"] == "Invalid credentials" + assert data["error"]["type"] == "auth_error" + + +def test_login_v2_returns_json_on_http_exception(monkeypatch): + """Test that /v2/login converts HTTPException to JSON error response""" + from fastapi import HTTPException + + mock_prisma_client = MagicMock() + mock_authenticate_user = AsyncMock( + side_effect=HTTPException(status_code=401, detail="Unauthorized") + ) + + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + mock_authenticate_user, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + 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 == 401 + assert response.headers["content-type"] == "application/json" + data = response.json() + assert "error" in data + assert isinstance(data["error"], dict) + + +def test_login_v2_returns_json_on_unexpected_exception(monkeypatch): + """Test that /v2/login returns JSON error when unexpected exception occurs""" + mock_prisma_client = MagicMock() + mock_authenticate_user = AsyncMock(side_effect=ValueError("Unexpected error")) + + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + mock_authenticate_user, + ) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + 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 == 500 + assert response.headers["content-type"] == "application/json" + data = response.json() + assert "error" in data + assert isinstance(data["error"], dict) + assert "Unexpected error" in data["error"]["message"] + + +def test_login_v2_returns_json_on_invalid_json_body(monkeypatch): + """Test that /v2/login returns JSON error when request body is invalid JSON""" + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + + client = TestClient(app) + response = client.post( + "/v2/login", + content="invalid json", + headers={"Content-Type": "application/json"}, + ) + + assert response.status_code == 500 + assert response.headers["content-type"] == "application/json" + data = response.json() + assert "error" in data + assert isinstance(data["error"], dict) + + def test_fallback_login_has_no_deprecation_banner(client_no_auth): response = client_no_auth.get("/fallback/login")