diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1222995529..fc9a3ff4a8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -91,6 +91,7 @@ from litellm.proxy._types import ( TeamDefaultSettings, TokenCountRequest, TransformRequestBody, + UI_TEAM_ID, UserAPIKeyAuth, ) from litellm.proxy.common_utils.callback_utils import ( @@ -12032,7 +12033,7 @@ async def onboarding(invite_link: str, request: Request): """ - Get the invite link - Validate it's still 'valid' - - Invalidate the link (prevents abuse) + - Return a short-lived onboarding token - Get user from db - Pass in user_email if set """ @@ -12070,7 +12071,7 @@ async def onboarding(invite_link: str, request: Request): ) #### CHECK IF ALREADY USED - if invite_obj.is_accepted is True: + if invite_obj.is_accepted is True or invite_obj.accepted_at is not None: raise HTTPException( status_code=401, detail={"error": "Invitation link has already been used."}, @@ -12086,24 +12087,6 @@ async def onboarding(invite_link: str, request: Request): status_code=401, detail={"error": "User does not exist in db."} ) - user_email = user_obj.user_email - - response = await generate_key_helper_fn( - request_type="key", - **{ - "user_role": user_obj.user_role, - "duration": LITELLM_UI_SESSION_DURATION, - "key_max_budget": litellm.max_ui_session_budget, - "models": [], - "aliases": {}, - "config": {}, - "spend": 0, - "user_id": user_obj.user_id, - "team_id": "litellm-dashboard", - }, # type: ignore - ) - key = response["token"] # type: ignore - litellm_dashboard_ui = get_custom_url(str(request.base_url)) if litellm_dashboard_ui.endswith("/"): litellm_dashboard_ui += "ui/onboarding" @@ -12111,13 +12094,24 @@ async def onboarding(invite_link: str, request: Request): litellm_dashboard_ui += "/ui/onboarding" import jwt + user_email = user_obj.user_email + onboarding_token = jwt.encode( # type: ignore + { + "token_type": "litellm_onboarding", + "invitation_link": invite_link, + "user_id": user_obj.user_id, + "exp": litellm.utils.get_utc_datetime() + timedelta(minutes=15), + }, + master_key, + algorithm="HS256", + ) disabled_non_admin_personal_key_creation = ( get_disabled_non_admin_personal_key_creation() ) returned_ui_token_object = ReturnedUITokenObject( user_id=user_obj.user_id, - key=key, + key=onboarding_token, user_email=user_obj.user_email, user_role=user_obj.user_role, login_method="username_password", @@ -12142,8 +12136,117 @@ async def onboarding(invite_link: str, request: Request): } +def _get_onboarding_claims_from_request(request: Request) -> dict: + global master_key, general_settings + + if master_key is None: + raise ProxyException( + message="Master Key not set for Proxy. Please set Master Key to use Admin UI. Set `LITELLM_MASTER_KEY` in .env or set general_settings:master_key in config.yaml. https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.", + type=ProxyErrorTypes.auth_error, + param="master_key", + code=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + auth_header_name = general_settings.get("litellm_key_header_name", "Authorization") + onboarding_auth_header = request.headers.get(auth_header_name) + if onboarding_auth_header is None: + raise HTTPException( + status_code=401, + detail={"error": "Missing onboarding session for invitation link."}, + ) + onboarding_token = onboarding_auth_header + if onboarding_token.lower().startswith("bearer "): + onboarding_token = onboarding_token.split(" ", 1)[1] + + import jwt + + try: + return jwt.decode( + onboarding_token, + master_key, + algorithms=["HS256"], + ) + except Exception: + raise HTTPException( + status_code=401, + detail={"error": "Invalid onboarding session for invitation link."}, + ) + + +async def _rollback_onboarding_invite_claim( + invitation_link: str, + user_id: str, +) -> None: + global prisma_client + + if prisma_client is None: + return + + try: + await prisma_client.db.litellm_invitationlink.update_many( + where={"id": invitation_link, "is_accepted": True}, + data={ + "accepted_at": None, + "is_accepted": False, + "updated_at": litellm.utils.get_utc_datetime(), + "updated_by": user_id, + }, + ) + except Exception: + verbose_proxy_logger.exception( + "Failed to roll back onboarding invitation after session key mint failed." + ) + + +async def _generate_onboarding_ui_session_token(user_obj: Any) -> str: + global master_key, general_settings + + response = await generate_key_helper_fn( + request_type="key", + **{ + "user_role": user_obj.user_role, + "duration": LITELLM_UI_SESSION_DURATION, + "key_max_budget": litellm.max_ui_session_budget, + "models": [], + "aliases": {}, + "config": {}, + "spend": 0, + "user_id": user_obj.user_id, + "team_id": UI_TEAM_ID, + }, # type: ignore + ) + key = response["token"] # type: ignore + + from litellm.types.proxy.ui_sso import ReturnedUITokenObject + + import jwt + + disabled_non_admin_personal_key_creation = ( + get_disabled_non_admin_personal_key_creation() + ) + returned_ui_token_object = ReturnedUITokenObject( + user_id=user_obj.user_id, + key=key, + user_email=user_obj.user_email, + user_role=user_obj.user_role, + login_method="username_password", + premium_user=premium_user, + auth_header_name=general_settings.get( + "litellm_key_header_name", "Authorization" + ), + disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, + server_root_path=get_server_root_path(), + ) + assert master_key is not None + return jwt.encode( # type: ignore + cast(dict, returned_ui_token_object), + master_key, + algorithm="HS256", + ) + + @app.post("/onboarding/claim_token", include_in_schema=False) -async def claim_onboarding_link(data: InvitationClaim): +async def claim_onboarding_link(data: InvitationClaim, request: Request): """ Special route. Allows UI link share user to update their password. @@ -12155,7 +12258,7 @@ async def claim_onboarding_link(data: InvitationClaim): This route can only update user password. """ - global prisma_client + global prisma_client, master_key, general_settings ### VALIDATE INVITE LINK ### if prisma_client is None: raise HTTPException( @@ -12180,7 +12283,7 @@ async def claim_onboarding_link(data: InvitationClaim): ) #### CHECK IF ALREADY USED - if invite_obj.is_accepted is True: + if invite_obj.is_accepted is True or invite_obj.accepted_at is not None: raise HTTPException( status_code=401, detail={"error": "Invitation link has already been used."}, @@ -12196,32 +12299,87 @@ async def claim_onboarding_link(data: InvitationClaim): ) }, ) - ### UPDATE USER OBJECT ### - hashed_pw = hash_password(data.password) - user_obj = await prisma_client.db.litellm_usertable.update( - where={"user_id": invite_obj.user_id}, data={"password": hashed_pw} - ) - if user_obj is None: + onboarding_claims = _get_onboarding_claims_from_request(request=request) + if ( + onboarding_claims.get("token_type") != "litellm_onboarding" + or onboarding_claims.get("invitation_link") != data.invitation_link + or onboarding_claims.get("user_id") != data.user_id + ): raise HTTPException( - status_code=401, detail={"error": "User does not exist in db."} + status_code=401, + detail={"error": "Invalid onboarding session for invitation link."}, ) - #### MARK LINK AS USED + hashed_pw = hash_password(data.password) current_time = litellm.utils.get_utc_datetime() - await prisma_client.db.litellm_invitationlink.update( - where={"id": data.invitation_link}, - data={ - "accepted_at": current_time, - "updated_at": current_time, - "is_accepted": True, - "updated_by": invite_obj.user_id, # type: ignore - }, - ) + async with prisma_client.db.tx() as tx: + updated_count = await tx.litellm_invitationlink.update_many( + where={"id": data.invitation_link, "is_accepted": False}, + data={ + "is_accepted": True, + "updated_at": current_time, + "updated_by": invite_obj.user_id, # type: ignore + }, + ) + if updated_count == 0: + raise HTTPException( + status_code=401, + detail={"error": "Invitation link has already been used."}, + ) + + ### UPDATE USER OBJECT ### + user_obj = await tx.litellm_usertable.update( + where={"user_id": invite_obj.user_id}, data={"password": hashed_pw} + ) + + if user_obj is None: + raise HTTPException( + status_code=401, detail={"error": "User does not exist in db."} + ) + + #### MARK LINK AS USED + current_time = litellm.utils.get_utc_datetime() + await tx.litellm_invitationlink.update( + where={"id": data.invitation_link}, + data={ + "accepted_at": current_time, + "updated_at": current_time, + "updated_by": invite_obj.user_id, # type: ignore + }, + ) if user_obj and hasattr(user_obj, "__dict__"): user_obj.__dict__.pop("password", None) - return user_obj + + try: + jwt_token = await _generate_onboarding_ui_session_token(user_obj=user_obj) + except Exception as e: + await _rollback_onboarding_invite_claim( + invitation_link=data.invitation_link, + user_id=data.user_id, + ) + if isinstance(e, HTTPException): + raise e + raise HTTPException( + status_code=500, + detail={ + "error": "Failed to create onboarding session. Please retry the invitation link." + }, + ) from e + + 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" + return { + "login_url": litellm_dashboard_ui, + "token": jwt_token, + "user_email": user_obj.user_email, + "user": user_obj, + } @app.get("/get_logo_url", include_in_schema=False) diff --git a/tests/test_litellm/proxy/auth/test_onboarding.py b/tests/test_litellm/proxy/auth/test_onboarding.py index 57fa871a35..c81f4cb7d6 100644 --- a/tests/test_litellm/proxy/auth/test_onboarding.py +++ b/tests/test_litellm/proxy/auth/test_onboarding.py @@ -2,14 +2,16 @@ Tests for the invite-link onboarding endpoints. Covers the security behavior of: - GET /onboarding/get_token – rejects already-used links before showing any user data - POST /onboarding/claim_token – rejects already-used links; marks is_accepted=True only - after the password is successfully written + GET /onboarding/get_token – rejects already-used links and returns only a + short-lived onboarding token, not a UI session key + POST /onboarding/claim_token – requires that onboarding token; mints the UI + session key only after the password is written """ from datetime import timedelta from unittest.mock import AsyncMock, MagicMock, patch +import jwt import pytest from fastapi import HTTPException @@ -22,14 +24,27 @@ from litellm.proxy._types import InvitationClaim # --------------------------------------------------------------------------- -def _make_invite(*, is_accepted: bool, expired: bool = False) -> MagicMock: +class _AsyncTx: + def __init__(self, db: MagicMock): + self.db = db + + async def __aenter__(self) -> MagicMock: + return self.db + + async def __aexit__(self, exc_type, exc, tb): + return False + + +def _make_invite( + *, is_accepted: bool, expired: bool = False, claimed: bool = False +) -> MagicMock: now = litellm.utils.get_utc_datetime() invite = MagicMock() invite.id = "invite-abc" invite.user_id = "user-123" invite.is_accepted = is_accepted invite.expires_at = now - timedelta(days=1) if expired else now + timedelta(days=6) - invite.accepted_at = None + invite.accepted_at = now if claimed else None return invite @@ -45,11 +60,39 @@ def _make_prisma(invite: MagicMock, user: MagicMock | None = None) -> MagicMock: prisma = MagicMock() prisma.db.litellm_invitationlink.find_unique = AsyncMock(return_value=invite) prisma.db.litellm_invitationlink.update = AsyncMock() + prisma.db.litellm_invitationlink.update_many = AsyncMock(return_value=1) prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user) prisma.db.litellm_usertable.update = AsyncMock(return_value=user) + prisma.db.tx = MagicMock(return_value=_AsyncTx(prisma.db)) return prisma +def _make_onboarding_token( + *, + invitation_link: str = "invite-abc", + user_id: str = "user-123", + token_type: str = "litellm_onboarding", + master_key: str = "sk-test", +) -> str: + return jwt.encode( + { + "token_type": token_type, + "invitation_link": invitation_link, + "user_id": user_id, + "exp": litellm.utils.get_utc_datetime() + timedelta(minutes=15), + }, + master_key, + algorithm="HS256", + ) + + +def _make_claim_request(token: str | None = None) -> MagicMock: + request = MagicMock() + request.headers = {"Authorization": f"Bearer {token}"} if token is not None else {} + request.base_url = "http://localhost:4000/" + return request + + # --------------------------------------------------------------------------- # GET /onboarding/get_token # --------------------------------------------------------------------------- @@ -120,10 +163,10 @@ async def test_get_token_rejects_missing_link(): @pytest.mark.asyncio -async def test_get_token_does_not_set_is_accepted(): +async def test_get_token_returns_onboarding_token_without_minting_ui_key(): """ - A valid, unused link should succeed and must NOT flip is_accepted to True. - That flag is only written after the password is claimed. + A valid, unused link should return a short-lived onboarding token, but + must not reserve the invite or mint a usable UI/API key on GET. """ from litellm.proxy.proxy_server import onboarding @@ -133,6 +176,252 @@ async def test_get_token_does_not_set_is_accepted(): request = MagicMock() request.base_url = "http://localhost:4000/" + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.premium_user", False), + patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key, + patch( + "litellm.proxy.proxy_server.get_custom_url", + return_value="http://localhost:4000/", + ), + patch( + "litellm.proxy.proxy_server.get_disabled_non_admin_personal_key_creation", + return_value=False, + ), + patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""), + ): + result = await onboarding(invite_link="invite-abc", request=request) + + # Endpoint succeeded + assert "token" in result + assert "login_url" in result + + outer_claims = jwt.decode(result["token"], "sk-test", algorithms=["HS256"]) + onboarding_token = outer_claims["key"] + onboarding_claims = jwt.decode(onboarding_token, "sk-test", algorithms=["HS256"]) + assert onboarding_claims["token_type"] == "litellm_onboarding" + assert onboarding_claims["invitation_link"] == "invite-abc" + assert onboarding_claims["user_id"] == "user-123" + assert not onboarding_token.startswith("sk-") + + mock_generate_key.assert_not_called() + prisma.db.litellm_invitationlink.update_many.assert_not_called() + prisma.db.litellm_invitationlink.update.assert_not_called() + + +# --------------------------------------------------------------------------- +# POST /onboarding/claim_token +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_claim_token_rejects_already_used_link(): + """ + If is_accepted is True, the password has already been set. + A second claim attempt must be rejected with 401. + """ + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=True, claimed=True) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=_make_claim_request()) + + assert exc_info.value.status_code == 401 + assert "already been used" in exc_info.value.detail["error"] + # Password must never have been written + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_claim_token_rejects_expired_link(): + """An expired link must be rejected even if is_accepted is False.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False, expired=True) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=_make_claim_request()) + + assert exc_info.value.status_code == 401 + assert "expired" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_claim_token_rejects_mismatched_user_id(): + """The user_id in the request must match the one on the invite.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="wrong-user", + password="NewP@ssw0rd", + ) + + with patch("litellm.proxy.proxy_server.prisma_client", prisma): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=_make_claim_request()) + + assert exc_info.value.status_code == 401 + assert "does not match" in exc_info.value.detail["error"] + + +@pytest.mark.asyncio +async def test_claim_token_rejects_missing_onboarding_token(): + """The password endpoint must require the onboarding token returned by get_token.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=_make_claim_request()) + + assert exc_info.value.status_code == 401 + assert "Missing onboarding session" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_claim_token_rejects_wrong_onboarding_session(): + """The onboarding token must be bound to the invite and user being claimed.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + request = _make_claim_request( + _make_onboarding_token(invitation_link="other-invite") + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.status_code == 401 + assert "Invalid onboarding session" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_claim_token_rejects_invalid_bearer_token(): + """A regular API key must not be accepted as an onboarding token.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + request = _make_claim_request("sk-regular-key") + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + ): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.status_code == 401 + assert "Invalid onboarding session" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + + +@pytest.mark.asyncio +async def test_claim_token_rejects_concurrent_reuse_before_password_write(): + """Only the first valid claim may reserve the invitation.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + prisma = _make_prisma(invite) + prisma.db.litellm_invitationlink.update_many = AsyncMock(return_value=0) + request = _make_claim_request(_make_onboarding_token()) + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key, + ): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.status_code == 401 + assert "already been used" in exc_info.value.detail["error"] + prisma.db.litellm_usertable.update.assert_not_called() + mock_generate_key.assert_not_called() + + +@pytest.mark.asyncio +async def test_claim_token_sets_accepted_at_after_password_written(): + """ + A valid first-time claim must: + 1. Write the hashed password to the user table. + 2. Set accepted_at on the invitation link after the password write succeeds. + """ + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + user = _make_user() + prisma = _make_prisma(invite, user) + request = _make_claim_request(_make_onboarding_token()) + + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + mock_token_response = {"token": "sk-generated-key", "user_id": "user-123"} with ( @@ -155,113 +444,13 @@ async def test_get_token_does_not_set_is_accepted(): ), patch("litellm.proxy.proxy_server.get_server_root_path", return_value=""), ): - result = await onboarding(invite_link="invite-abc", request=request) - - # Endpoint succeeded - assert "token" in result - assert "login_url" in result - - # is_accepted must NOT have been updated here - prisma.db.litellm_invitationlink.update.assert_not_called() - - -# --------------------------------------------------------------------------- -# POST /onboarding/claim_token -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_claim_token_rejects_already_used_link(): - """ - If is_accepted is True, the password has already been set. - A second claim attempt must be rejected with 401. - """ - from litellm.proxy.proxy_server import claim_onboarding_link - - invite = _make_invite(is_accepted=True) - prisma = _make_prisma(invite) - data = InvitationClaim( - invitation_link="invite-abc", - user_id="user-123", - password="NewP@ssw0rd", - ) - - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - with pytest.raises(HTTPException) as exc_info: - await claim_onboarding_link(data=data) - - assert exc_info.value.status_code == 401 - assert "already been used" in exc_info.value.detail["error"] - # Password must never have been written - prisma.db.litellm_usertable.update.assert_not_called() - - -@pytest.mark.asyncio -async def test_claim_token_rejects_expired_link(): - """An expired link must be rejected even if is_accepted is False.""" - from litellm.proxy.proxy_server import claim_onboarding_link - - invite = _make_invite(is_accepted=False, expired=True) - prisma = _make_prisma(invite) - data = InvitationClaim( - invitation_link="invite-abc", - user_id="user-123", - password="NewP@ssw0rd", - ) - - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - with pytest.raises(HTTPException) as exc_info: - await claim_onboarding_link(data=data) - - assert exc_info.value.status_code == 401 - assert "expired" in exc_info.value.detail["error"] - - -@pytest.mark.asyncio -async def test_claim_token_rejects_mismatched_user_id(): - """The user_id in the request must match the one on the invite.""" - from litellm.proxy.proxy_server import claim_onboarding_link - - invite = _make_invite(is_accepted=False) - prisma = _make_prisma(invite) - data = InvitationClaim( - invitation_link="invite-abc", - user_id="wrong-user", - password="NewP@ssw0rd", - ) - - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - with pytest.raises(HTTPException) as exc_info: - await claim_onboarding_link(data=data) - - assert exc_info.value.status_code == 401 - assert "does not match" in exc_info.value.detail["error"] - - -@pytest.mark.asyncio -async def test_claim_token_sets_is_accepted_after_password_written(): - """ - A valid first-time claim must: - 1. Write the hashed password to the user table. - 2. Flip is_accepted to True on the invitation link — and only after the - password write succeeds. - """ - from litellm.proxy.proxy_server import claim_onboarding_link - - invite = _make_invite(is_accepted=False) - user = _make_user() - prisma = _make_prisma(invite, user) - - data = InvitationClaim( - invitation_link="invite-abc", - user_id="user-123", - password="NewP@ssw0rd", - ) - - with patch("litellm.proxy.proxy_server.prisma_client", prisma): - result = await claim_onboarding_link(data=data) + result = await claim_onboarding_link(data=data, request=request) # Password was written + prisma.db.litellm_invitationlink.update_many.assert_called_once() + reserve_kwargs = prisma.db.litellm_invitationlink.update_many.call_args.kwargs + assert reserve_kwargs["where"] == {"id": "invite-abc", "is_accepted": False} + assert reserve_kwargs["data"]["is_accepted"] is True prisma.db.litellm_usertable.update.assert_called_once() call_kwargs = prisma.db.litellm_usertable.update.call_args assert call_kwargs.kwargs["where"] == {"user_id": "user-123"} @@ -270,5 +459,50 @@ async def test_claim_token_sets_is_accepted_after_password_written(): # is_accepted was flipped to True on the invitation link prisma.db.litellm_invitationlink.update.assert_called_once() link_update_data = prisma.db.litellm_invitationlink.update.call_args.kwargs["data"] - assert link_update_data["is_accepted"] is True + assert "is_accepted" not in link_update_data assert link_update_data["accepted_at"] is not None + outer_claims = jwt.decode(result["token"], "sk-test", algorithms=["HS256"]) + assert outer_claims["key"] == "sk-generated-key" + + +@pytest.mark.asyncio +async def test_claim_token_rolls_back_invite_when_session_key_mint_fails(): + """A session key failure must not leave the invite permanently consumed.""" + from litellm.proxy.proxy_server import claim_onboarding_link + + invite = _make_invite(is_accepted=False) + user = _make_user() + prisma = _make_prisma(invite, user) + request = _make_claim_request(_make_onboarding_token()) + + data = InvitationClaim( + invitation_link="invite-abc", + user_id="user-123", + password="NewP@ssw0rd", + ) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.master_key", "sk-test"), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch( + "litellm.proxy.proxy_server.generate_key_helper_fn", + new_callable=AsyncMock, + side_effect=Exception("key mint failed"), + ), + ): + with pytest.raises(HTTPException) as exc_info: + await claim_onboarding_link(data=data, request=request) + + assert exc_info.value.status_code == 500 + assert "Failed to create onboarding session" in exc_info.value.detail["error"] + assert prisma.db.litellm_invitationlink.update_many.call_count == 2 + rollback_kwargs = prisma.db.litellm_invitationlink.update_many.call_args_list[ + 1 + ].kwargs + assert rollback_kwargs["where"] == { + "id": "invite-abc", + "is_accepted": True, + } + assert rollback_kwargs["data"]["accepted_at"] is None + assert rollback_kwargs["data"]["is_accepted"] is False diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx index d9e75388d1..bf73e802cd 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx @@ -1,6 +1,6 @@ -import { render, screen } from "@testing-library/react"; +import { act, fireEvent, render, screen } from "@testing-library/react"; import React from "react"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { OnboardingForm } from "./OnboardingForm"; const mockUseOnboardingCredentials = vi.fn(); @@ -36,14 +36,33 @@ vi.mock("./OnboardingErrorView", () => ({ })); vi.mock("./OnboardingFormBody", () => ({ - OnboardingFormBody: ({ variant, userEmail }: { variant: string; userEmail: string }) => ( + OnboardingFormBody: ({ + variant, + userEmail, + claimError, + onSubmit, + }: { + variant: string; + userEmail: string; + claimError: string | null; + onSubmit: (formValues: { password: string }) => void; + }) => (
Form Body + + {claimError ?
{claimError}
: null}
), })); describe("OnboardingForm", () => { + beforeEach(() => { + vi.clearAllMocks(); + document.cookie = "token=; expires=Thu, 01 Jan 1970 00:00:00 GMT; path=/"; + }); + it("should render loading view when credentials are loading", () => { mockUseOnboardingCredentials.mockReturnValue({ data: undefined, @@ -92,4 +111,24 @@ describe("OnboardingForm", () => { expect(screen.getByTestId("form-body")).toHaveAttribute("data-variant", "reset_password"); }); + + it("should show claim error when claim response is missing final token", async () => { + mockUseOnboardingCredentials.mockReturnValue({ + data: { token: "fake-jwt-token" }, + isLoading: false, + isError: false, + }); + mockClaimToken.mockImplementation((_params, options) => { + options.onSuccess({}); + }); + + render(); + + await act(async () => { + fireEvent.click(screen.getByRole("button", { name: "Submit" })); + }); + + expect(screen.getByTestId("claim-error")).toHaveTextContent("Failed to start session"); + expect(document.cookie).not.toContain("fake-jwt-token"); + }); }); diff --git a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx index 9a9d9d7e72..1e8afe7645 100644 --- a/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx +++ b/ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx @@ -31,18 +31,21 @@ export function OnboardingForm({ variant }: OnboardingFormProps) { const userEmail: string = decoded?.user_email ?? ""; const userId: string | null = decoded?.user_id ?? null; const accessToken: string | null = decoded?.key ?? null; - const jwtToken: string | null = credentialsData?.token ?? null; const handleSubmit = (formValues: { password: string }) => { - if (!accessToken || !jwtToken || !userId || !inviteId) return; + if (!accessToken || !userId || !inviteId) return; setClaimError(null); claimToken( { accessToken, inviteId, userId, password: formValues.password }, { - onSuccess: () => { - document.cookie = `token=${jwtToken}; path=/; SameSite=Lax`; + onSuccess: (data: { token?: string }) => { + if (!data?.token) { + setClaimError("Failed to start session. Please try again."); + return; + } + document.cookie = `token=${data.token}; path=/; SameSite=Lax`; const proxyBaseUrl = getProxyBaseUrl(); window.location.href = proxyBaseUrl ? `${proxyBaseUrl}/ui/?login=success`