diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index e93c650ed97..17f27e48931 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -37,6 +37,7 @@ from litellm.proxy._types import ( VirtualKeyEvent, WebhookEvent, ) +from litellm.proxy.management_helpers.user_invitation import get_user_invitation_link from litellm.repositories.team_repository import TeamRepository from litellm.repositories.user_repository import UserRepository from litellm.types.integrations.slack_alerting import * @@ -1177,7 +1178,10 @@ Model Info: email_logo_url=email_logo_url, recipient_email=recipient_email, team_name=team_name, - base_url=base_url, + base_url=await get_user_invitation_link( + user_id=recipient_user_id, + base_url=base_url, + ), email_support_contact=email_support_contact, ) else: diff --git a/litellm/proxy/management_helpers/user_invitation.py b/litellm/proxy/management_helpers/user_invitation.py index babc920189a..82859ae7c02 100644 --- a/litellm/proxy/management_helpers/user_invitation.py +++ b/litellm/proxy/management_helpers/user_invitation.py @@ -1,8 +1,11 @@ +from __future__ import annotations + from datetime import timedelta from fastapi import HTTPException import litellm +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, InvitationNew, UserAPIKeyAuth from litellm.repositories.table_repositories import InvitationLinkRepository @@ -46,3 +49,45 @@ async def create_invitation_for_user( }, ) raise HTTPException(status_code=500, detail={"error": str(e)}) + + +def construct_invitation_link(invitation_id: str, base_url: str) -> str: + """ + e.g. http://localhost:4000/ui/onboarding?invitation_id=7a096b3a-37c6-440f-9dd1-ba22e8043f6b + """ + return f"{base_url.rstrip('/')}/ui/onboarding?invitation_id={invitation_id}" + + +async def get_user_invitation_link(user_id: str | None, base_url: str) -> str: + """ + Return the onboarding link for `user_id`, reusing the user's most recent invitation + or creating one if none exists. + + Falls back to `base_url` when the link cannot be built. + """ + from litellm.proxy.proxy_server import prisma_client + + if user_id is None or prisma_client is None: + return base_url + + try: + existing_invitations = await InvitationLinkRepository(prisma_client).table.find_many( + where={"user_id": user_id}, + order={"created_at": "desc"}, + ) + invitation = ( + existing_invitations[0] + if existing_invitations + else await create_invitation_for_user( + data=InvitationNew(user_id=user_id), + user_api_key_dict=UserAPIKeyAuth(user_id=user_id), + ) + ) + except Exception as e: + verbose_proxy_logger.error("Unable to get/create invitation for user_id %s - %s", user_id, str(e)) + return base_url + + if invitation is None: + return base_url + + return construct_invitation_link(invitation_id=invitation.id, base_url=base_url) diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 1ea4795207d..bb8e5e76df3 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -3,9 +3,12 @@ import json import os import sys import unittest +from types import SimpleNamespace from typing import List, Optional, Tuple from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch +import pytest + sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system-path @@ -245,3 +248,112 @@ class TestSlackAlerting(unittest.TestCase): ) self.assertEqual(parsed_data["alerts"], [408]) self.assertEqual(parsed_data["provider_region_id"], "vertex_aius-east1") + + +@pytest.mark.asyncio +async def test_user_invited_email_links_to_existing_invitation(): + """Legacy SMTP invite email must link to /ui/onboarding?invitation_id=... (issue #34555)""" + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import WebhookEvent + + prisma_client = MagicMock() + prisma_client.db.litellm_invitationlink.find_many = AsyncMock( + return_value=[SimpleNamespace(id="1c1e7bfa-0f7e-4d1b-9d1c-6b1b6a5a7f3a")] + ) + prisma_client.db.litellm_invitationlink.create = AsyncMock() + + sent_email = AsyncMock() + slack_alerting = SlackAlerting(alerting=["email"]) + with patch.object(proxy_server, "prisma_client", prisma_client), patch.dict( + os.environ, {"PROXY_BASE_URL": "https://proxy.example.com"} + ), patch("litellm.proxy.utils.send_email", sent_email): + result = await slack_alerting.send_key_created_or_user_invited_email( + webhook_event=WebhookEvent( + event="internal_user_created", + event_group=Litellm_EntityType.USER, + event_message="Welcome to LiteLLM Proxy", + spend=0.0, + user_id="new-user-id", + user_email="new-user@example.com", + ) + ) + + assert result is True + prisma_client.db.litellm_invitationlink.create.assert_not_called() + html = sent_email.call_args.kwargs["html"] + assert ( + 'href="https://proxy.example.com/ui/onboarding?invitation_id=1c1e7bfa-0f7e-4d1b-9d1c-6b1b6a5a7f3a"' + in html + ) + + +@pytest.mark.asyncio +async def test_user_invited_email_creates_invitation_when_missing(): + """No invitation row yet -> one is created and its id is used in the email link""" + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import WebhookEvent + + prisma_client = MagicMock() + prisma_client.db.litellm_invitationlink.find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_invitationlink.create = AsyncMock( + return_value=SimpleNamespace(id="fresh-invite-id") + ) + + sent_email = AsyncMock() + slack_alerting = SlackAlerting(alerting=["email"]) + with patch.object(proxy_server, "prisma_client", prisma_client), patch.dict( + os.environ, {"PROXY_BASE_URL": "https://proxy.example.com/"} + ), patch("litellm.proxy.utils.send_email", sent_email): + await slack_alerting.send_key_created_or_user_invited_email( + webhook_event=WebhookEvent( + event="internal_user_created", + event_group=Litellm_EntityType.USER, + event_message="Welcome to LiteLLM Proxy", + spend=0.0, + user_id="new-user-id", + user_email="new-user@example.com", + ) + ) + + assert ( + prisma_client.db.litellm_invitationlink.create.call_args.kwargs["data"][ + "user_id" + ] + == "new-user-id" + ) + html = sent_email.call_args.kwargs["html"] + assert ( + 'href="https://proxy.example.com/ui/onboarding?invitation_id=fresh-invite-id"' + in html + ) + + +@pytest.mark.asyncio +async def test_user_invited_email_falls_back_to_base_url_on_db_error(): + """If the invitation lookup blows up, the email still goes out with the base url""" + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import WebhookEvent + + prisma_client = MagicMock() + prisma_client.db.litellm_invitationlink.find_many = AsyncMock( + side_effect=Exception("db down") + ) + + sent_email = AsyncMock() + slack_alerting = SlackAlerting(alerting=["email"]) + with patch.object(proxy_server, "prisma_client", prisma_client), patch.dict( + os.environ, {"PROXY_BASE_URL": "https://proxy.example.com"} + ), patch("litellm.proxy.utils.send_email", sent_email): + result = await slack_alerting.send_key_created_or_user_invited_email( + webhook_event=WebhookEvent( + event="internal_user_created", + event_group=Litellm_EntityType.USER, + event_message="Welcome to LiteLLM Proxy", + spend=0.0, + user_id="new-user-id", + user_email="new-user@example.com", + ) + ) + + assert result is True + assert 'href="https://proxy.example.com"' in sent_email.call_args.kwargs["html"]