ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #26843 from stuxf/codex/fix-onboarding-invite-token

Browse Source 
chore(auth): harden invite-link onboarding token flow
This commit is contained in:
yuneng-jiang 2026-04-30 11:56:26 -07:00 committed by GitHub
commit a9db887bdd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 597 additions and 163 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

242
litellm/proxy/proxy_server.py
View File

@ -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)

462
tests/test_litellm/proxy/auth/test_onboarding.py
View File

@ -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

45
ui/litellm-dashboard/src/app/onboarding/OnboardingForm.test.tsx
View File

@ -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;
}) => (
<div data-testid="form-body" data-variant={variant} data-email={userEmail}>
Form Body
<button type="button" onClick={() => onSubmit({ password: "NewP@ssw0rd" })}>
Submit
</button>
{claimError ? <div data-testid="claim-error">{claimError}</div> : null}
</div>
),
}));
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(<OnboardingForm variant="signup" />);
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");
});
});

11
ui/litellm-dashboard/src/app/onboarding/OnboardingForm.tsx
View File

@ -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`