feat(proxy): email alerts at configured percentages of a team member budget

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan 2026-09-23 02:37:00 +00:00
parent 5dfaa8d620
commit 30e1ce2819
16 changed files with 941 additions and 16 deletions

View file

@ -32,6 +32,7 @@ from litellm.integrations.email_templates.key_rotated_email import (
from litellm.integrations.email_templates.templates import (
MAX_BUDGET_ALERT_EMAIL_TEMPLATE,
SOFT_BUDGET_ALERT_EMAIL_TEMPLATE,
TEAM_MEMBER_MAX_BUDGET_ALERT_EMAIL_TEMPLATE,
TEAM_SOFT_BUDGET_ALERT_EMAIL_TEMPLATE,
)
from litellm.integrations.email_templates.user_invitation_email import (
@ -48,6 +49,12 @@ from litellm.secret_managers.main import get_secret_bool
from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL
def _max_budget_alert_id(user_info: CallInfo) -> str:
if user_info.event_group == Litellm_EntityType.TEAM_MEMBER:
return f"team_member:{user_info.user_id}:{user_info.team_id}"
return user_info.token or user_info.user_id or "default_id"
def _parse_email_list(raw) -> List[str]:
"""Parse emails from a list or comma-separated string."""
if isinstance(raw, list):
@ -373,17 +380,31 @@ class BaseEmailLogger(CustomLogger):
greeting = html.escape(
event.user_email or event.key_alias or event.token or ""
)
email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
recipient_email=greeting,
percentage=percentage,
spend=spend_str,
max_budget=max_budget_str,
alert_threshold=alert_threshold_str,
base_url=email_params.base_url,
email_support_contact=email_params.support_contact,
email_footer=email_params.signature,
)
if event.event_group == Litellm_EntityType.TEAM_MEMBER:
email_html_content = TEAM_MEMBER_MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
member=html.escape(event.user_email or event.user_id or ""),
team_alias=html.escape(event.team_alias or event.team_id or ""),
percentage=percentage,
spend=spend_str,
max_budget=max_budget_str,
alert_threshold=alert_threshold_str,
base_url=email_params.base_url,
email_support_contact=email_params.support_contact,
email_footer=email_params.signature,
)
else:
email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
recipient_email=greeting,
percentage=percentage,
spend=spend_str,
max_budget=max_budget_str,
alert_threshold=alert_threshold_str,
base_url=email_params.base_url,
email_support_contact=email_params.support_contact,
email_footer=email_params.signature,
)
await self.send_email(
from_email=self.DEFAULT_LITELLM_EMAIL,
to_email=recipient_emails,
@ -607,7 +628,7 @@ class BaseEmailLogger(CustomLogger):
if user_info.spend < threshold_amount:
continue
_id = user_info.token or user_info.user_id or "default_id"
_id = _max_budget_alert_id(user_info)
_cache_key = (
f"email_budget_alerts:max_budget_alert:{threshold_pct}:{_id}"
)
@ -618,7 +639,7 @@ class BaseEmailLogger(CustomLogger):
emails.append(user_info.user_email)
if not emails:
verbose_proxy_logger.warning(
"No recipients for %d%% threshold on key %s, skipping alert",
"No recipients for %d%% threshold on %s, skipping alert",
threshold_pct,
_id,
)
@ -633,7 +654,11 @@ class BaseEmailLogger(CustomLogger):
if send_count is not None and send_count > 1:
continue
event_message = f"Max Budget Alert - {threshold_pct}% of Maximum Budget Reached"
event_message = (
f"Team Member Budget Alert - {threshold_pct}% of Team Member Budget Reached"
if user_info.event_group == Litellm_EntityType.TEAM_MEMBER
else f"Max Budget Alert - {threshold_pct}% of Maximum Budget Reached"
)
webhook_event = WebhookEvent(
event="max_budget_alert",
event_message=event_message,

View file

@ -63,6 +63,8 @@ class TokenBudgetAlert(BaseBudgetAlertType):
return "Key Budget: "
def get_id(self, user_info: CallInfo) -> str:
if user_info.event_group == Litellm_EntityType.TEAM_MEMBER:
return f"team_member:{user_info.user_id}:{user_info.team_id}"
return user_info.token or "default_id"

View file

@ -131,3 +131,25 @@ MAX_BUDGET_ALERT_EMAIL_TEMPLATE: Final = """
{email_footer}
"""
TEAM_MEMBER_MAX_BUDGET_ALERT_EMAIL_TEMPLATE: Final = """
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" />
<p> Hi, <br/>
Team member <b>{member}</b> has reached <b>{percentage}% of their team member budget</b> in team <b>{team_alias}</b>. <br /> <br />
<b>Current Spend:</b> {spend} <br />
<b>Team Member Budget:</b> {max_budget} <br />
<b>Alert Threshold:</b> {alert_threshold} ({percentage}%) <br />
<p style="color: #dc2626; font-weight: 500;">
⚠️ Warning: Once this member reaches their team member budget of {max_budget}, their requests in this team will be rejected.
</p>
You can view usage and manage team member budgets in the <a href="{base_url}">LiteLLM Dashboard</a>. <br /> <br />
If you have any questions, please send an email to {email_support_contact} <br /> <br />
{email_footer}
"""

View file

@ -5524,6 +5524,44 @@ async def _virtual_key_max_budget_alert_check(
)
TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY: Final = "team_member_max_budget_alert_emails"
def _team_member_max_budget_alert_check(
team_id: str,
team_alias: str | None,
team_metadata: Mapping[str, object] | None,
organization_id: str | None,
user_id: str,
user_email: str | None,
proxy_logging_obj: ProxyLogging,
spend: float,
max_budget: float,
) -> None:
raw_config: Final = (team_metadata or {}).get(TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY)
alert_email_config: Final = _merge_budget_alert_email_configs(
global_cfg=None,
per_key_cfg=raw_config if isinstance(raw_config, Mapping) else None,
)
if not alert_email_config or spend <= 0:
return
min_pct: Final = min((int(k) for k in alert_email_config if k.isdigit()), default=None)
if min_pct is None or spend < max_budget * (min_pct / 100.0):
return
call_info: Final = CallInfo(
spend=spend,
max_budget=max_budget,
user_id=user_id,
team_id=team_id,
team_alias=team_alias,
organization_id=organization_id,
user_email=user_email,
event_group=Litellm_EntityType.TEAM_MEMBER,
max_budget_alert_emails=alert_email_config,
)
asyncio.create_task(proxy_logging_obj.budget_alerts(type="max_budget_alert", user_info=call_info))
async def _check_team_member_budget(
team_object: LiteLLM_TeamTable | None,
user_object: LiteLLM_UserTable | None,
@ -5589,7 +5627,22 @@ async def _check_team_member_budget(
max_budget=team_member_budget,
)
if math.isfinite(team_member_budget) and team_member_spend >= team_member_budget:
if not math.isfinite(team_member_budget):
return
_team_member_max_budget_alert_check(
team_id=team_object.team_id,
team_alias=team_object.team_alias,
team_metadata=team_object.metadata,
organization_id=team_object.organization_id,
user_id=valid_token.user_id,
user_email=user_object.user_email if user_object is not None else None,
proxy_logging_obj=proxy_logging_obj,
spend=team_member_spend,
max_budget=team_member_budget,
)
if team_member_spend >= team_member_budget:
raise litellm.BudgetExceededError(
current_cost=team_member_spend,
max_budget=team_member_budget,

View file

@ -50,6 +50,7 @@ from litellm.proxy.auth.auth_checks import (
_get_user_role,
_is_model_cost_zero,
_is_user_proxy_admin,
_team_member_max_budget_alert_check,
_virtual_key_max_budget_alert_check,
_virtual_key_max_budget_check,
_virtual_key_soft_budget_check,
@ -2282,6 +2283,17 @@ async def _user_api_key_auth_builder(
fallback_spend=team_member_spend,
max_budget=team_member_budget,
)
_team_member_max_budget_alert_check(
team_id=_team_id,
team_alias=valid_token.team_alias,
team_metadata=valid_token.team_metadata,
organization_id=valid_token.org_id,
user_id=_user_id,
user_email=user_obj.user_email if user_obj is not None else None,
proxy_logging_obj=proxy_logging_obj,
spend=team_member_spend,
max_budget=team_member_budget,
)
if team_member_spend >= team_member_budget:
_entity_id: Final = f"{valid_token.user_id}:{valid_token.team_id}"
raise litellm.BudgetExceededError(

View file

@ -0,0 +1,120 @@
from __future__ import annotations
import socketserver
import threading
from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from email import message_from_bytes
from email.message import Message
from queue import SimpleQueue
from typing import Final
@dataclass(frozen=True, slots=True)
class Delivery:
sender: str
recipients: tuple[str, ...]
message: Message
@property
def subject(self) -> str:
return str(self.message["Subject"])
@property
def html(self) -> str:
for part in self.message.walk():
if part.get_content_type() == "text/html":
return part.get_payload(decode=True).decode()
return ""
@dataclass(frozen=True, slots=True)
class Mailbox:
host: str
port: int
received: SimpleQueue[Delivery]
def drain(self) -> tuple[Delivery, ...]:
return tuple(self.received.get_nowait() for _ in range(self.received.qsize()))
def _address(argument: str) -> str:
return argument.split(":", 1)[1].strip().strip("<>")
@contextmanager
def smtp_sink() -> Generator[Mailbox, None, None]:
"""Owned plaintext SMTP peer; deliveries traverse the proxy's real smtplib client."""
received: Final[SimpleQueue[Delivery]] = SimpleQueue()
errors: Final[SimpleQueue[Exception]] = SimpleQueue()
class Handler(socketserver.StreamRequestHandler):
timeout = 5
def handle(self) -> None:
try:
self._session()
except Exception as error:
errors.put(error)
def _reply(self, line: str) -> None:
self.wfile.write(f"{line}\r\n".encode())
self.wfile.flush()
def _session(self) -> None:
self._reply("220 integration-smtp ready")
sender = ""
recipients: tuple[str, ...] = ()
while True:
raw: Final = self.rfile.readline()
if not raw:
return
line = raw.decode().rstrip("\r\n")
verb = line.split(" ", 1)[0].upper()
if verb in {"EHLO", "HELO"}:
self._reply("250 integration-smtp")
elif verb == "MAIL":
sender = _address(line)
self._reply("250 OK")
elif verb == "RCPT":
recipients = (*recipients, _address(line))
self._reply("250 OK")
elif verb == "DATA":
self._reply("354 End data with <CR><LF>.<CR><LF>")
body = bytearray()
while True:
chunk: Final = self.rfile.readline()
if not chunk or chunk == b".\r\n":
break
body.extend(chunk[1:] if chunk.startswith(b"..") else chunk)
received.put(Delivery(sender, recipients, message_from_bytes(bytes(body))))
sender, recipients = "", ()
self._reply("250 OK queued")
elif verb == "RSET":
sender, recipients = "", ()
self._reply("250 OK")
elif verb == "NOOP":
self._reply("250 OK")
elif verb == "QUIT":
self._reply("221 Bye")
return
else:
self._reply("502 Command not implemented")
class OwnedServer(socketserver.ThreadingTCPServer):
allow_reuse_address = True
daemon_threads = False
with OwnedServer(("127.0.0.1", 0), Handler) as server:
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
thread.start()
try:
yield Mailbox("127.0.0.1", server.server_address[1], received)
finally:
server.shutdown()
thread.join(timeout=6)
assert not thread.is_alive(), "Owned SMTP server survived cleanup"
server.server_close()
failure: Final = None if errors.empty() else errors.get_nowait()
assert failure is None, f"Owned SMTP peer failed: {failure!r}"

View file

@ -1794,6 +1794,9 @@
],
"tests/integration/observability/test_callback_delivery.py::test_response_survives_raising_success_deployment_hook[videos]": [
"other.observability.callbacks.raising_success_deployment_hook_keeps_response"
],
"tests/integration/spend/test_team_member_budget_alerts.py::test_team_member_budget_thresholds_email_member_and_configured_recipients": [
"quota_management.budget.team_member.alerts_at_configured_thresholds"
]
},
"browser": {

View file

@ -0,0 +1,104 @@
import uuid
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, eventually
from integration._support.database import read_rows
from integration._support.mail import Delivery, Mailbox, smtp_sink
from integration._support.process import owned_proxy
MEMBER_BUDGET: Final = 0.10
CALL_COST: Final = 20 * 0.001 + 20 * 0.002
def _membership_spend(user_id: str, team_id: str) -> float:
rows: Final = read_rows(
'SELECT spend FROM "LiteLLM_TeamMembership" WHERE user_id = %s AND team_id = %s', (user_id, team_id)
)
return float(str(rows[0]["spend"])) if rows else 0.0
def _deliveries_for(mailbox: Mailbox, seen: list[Delivery]) -> list[Delivery]:
seen.extend(mailbox.drain())
return seen
@pytest.mark.covers("quota_management.budget.team_member.alerts_at_configured_thresholds")
def test_team_member_budget_thresholds_email_member_and_configured_recipients(gateway: Gateway, tmp_path: Path) -> None:
member_email: Final = f"member-{uuid.uuid4().hex}@integration.test"
finance_email: Final = f"finance-{uuid.uuid4().hex}@integration.test"
configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration["general_settings"]["alerting"] = ["email"]
path: Final = tmp_path / "email-alerting.yaml"
path.write_text(yaml.safe_dump(configuration))
with smtp_sink() as mailbox:
overrides: Final = {
"SMTP_HOST": mailbox.host,
"SMTP_PORT": str(mailbox.port),
"SMTP_TLS": "False",
"SMTP_SENDER_EMAIL": "alerts@integration.test",
}
with owned_proxy(gateway, tmp_path, overrides, config=path) as candidate, candidate.scenario() as scenario:
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
user_id: Final = scenario.user(user_email=member_email)
team_id: Final = scenario.team(
models=[model],
team_member_budget=MEMBER_BUDGET,
metadata={"team_member_max_budget_alert_emails": {"50": [], "100": [finance_email]}},
)
candidate.post("/team/member_add", {"team_id": team_id, "member": {"user_id": user_id, "role": "user"}})
key: Final = scenario.key(team_id=team_id, user_id=user_id)
received: Final[list[Delivery]] = []
first: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "first call"}]},
key=key,
)
assert first.status_code == 200, first.text
assert float(first.headers["x-litellm-response-cost"]) == pytest.approx(CALL_COST)
eventually(
lambda: _membership_spend(user_id, team_id), lambda spend: spend == pytest.approx(CALL_COST), seconds=70
)
assert _deliveries_for(mailbox, received) == [], "no threshold is reached before the first call is recorded"
second: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "second call"}]},
key=key,
)
assert second.status_code == 200, second.text
halfway: Final = eventually(
lambda: _deliveries_for(mailbox, received), lambda values: len(values) >= 1, seconds=30
)
assert [delivery.recipients for delivery in halfway] == [(member_email,)], halfway
assert "50%" in halfway[0].subject, halfway[0].subject
assert f"${MEMBER_BUDGET}" in halfway[0].html, halfway[0].html
eventually(
lambda: _membership_spend(user_id, team_id),
lambda spend: spend == pytest.approx(2 * CALL_COST),
seconds=70,
)
third: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "third call"}]},
key=key,
)
assert third.status_code == 422 and third.json()["error"]["type"] == "budget_exceeded", third.text
capped: Final = eventually(
lambda: _deliveries_for(mailbox, received), lambda values: len(values) >= 3, seconds=30
)
hundred: Final = capped[1:]
assert all("100%" in delivery.subject for delivery in hundred), capped
assert {recipient for delivery in hundred for recipient in delivery.recipients} == {
member_email,
finance_email,
}, capped
assert all(member_email in delivery.html and f"${MEMBER_BUDGET}" in delivery.html for delivery in hundred)
assert len(capped) == 3, capped

View file

@ -1090,6 +1090,47 @@ async def test_multi_threshold_empty_emails_only_owner(
assert to_emails == ["owner@co.com"]
@pytest.mark.asyncio
async def test_multi_threshold_team_member_alert_renders_member_template_per_team(
base_email_logger, mock_send_email
):
"""A team member budget alert is keyed per member and team, names the member and team,
and goes to the member plus the threshold's configured recipients"""
user_info = CallInfo(
user_id="member_1",
user_email="member@co.com",
team_id="team_a",
team_alias="Platform",
spend=0.10,
max_budget=0.10,
event_group=Litellm_EntityType.TEAM_MEMBER,
max_budget_alert_emails={"50": [], "100": ["finance@co.com"]},
)
mock_cache = mock.AsyncMock()
mock_cache.async_increment_cache = mock.AsyncMock(return_value=1)
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
cache_keys = sorted(c[1]["key"] for c in mock_cache.async_increment_cache.call_args_list)
assert cache_keys == [
"email_budget_alerts:max_budget_alert:100:team_member:member_1:team_a",
"email_budget_alerts:max_budget_alert:50:team_member:member_1:team_a",
]
assert mock_send_email.call_count == 2
hundred = next(
c.kwargs for c in mock_send_email.call_args_list if "100%" in c.kwargs["subject"]
)
assert hundred["subject"] == "LiteLLM: Team Member Budget Alert - 100% of Team Member Budget Reached"
assert sorted(hundred["to_email"]) == ["finance@co.com", "member@co.com"]
assert "member@co.com" in hundred["html_body"] and "Platform" in hundred["html_body"]
assert "team member budget" in hundred["html_body"] and "$0.1" in hundred["html_body"]
fifty = next(c.kwargs for c in mock_send_email.call_args_list if "50%" in c.kwargs["subject"])
assert fifty["to_email"] == ["member@co.com"]
@pytest.mark.asyncio
async def test_no_map_preserves_old_single_threshold(
base_email_logger, mock_send_email

View file

@ -1,4 +1,7 @@
from litellm.integrations.SlackAlerting.budget_alert_types import SoftBudgetAlert
from litellm.integrations.SlackAlerting.budget_alert_types import (
SoftBudgetAlert,
TokenBudgetAlert,
)
from litellm.proxy._types import CallInfo, Litellm_EntityType
@ -64,3 +67,30 @@ class TestSoftBudgetAlert:
result = alert.get_id(user_info)
assert result == "default_id"
class TestTokenBudgetAlert:
def test_get_id_dedupes_team_member_alerts_per_member_and_team(self):
alert = TokenBudgetAlert()
team_a = CallInfo(
spend=8.0, max_budget=10.0, user_id="member_1", team_id="team_a", event_group=Litellm_EntityType.TEAM_MEMBER
)
team_b = CallInfo(
spend=8.0, max_budget=10.0, user_id="member_1", team_id="team_b", event_group=Litellm_EntityType.TEAM_MEMBER
)
assert alert.get_id(team_a) == "team_member:member_1:team_a"
assert alert.get_id(team_b) == "team_member:member_1:team_b"
def test_get_id_uses_token_for_key_alerts(self):
alert = TokenBudgetAlert()
user_info = CallInfo(
spend=8.0,
max_budget=10.0,
token="hashed_key",
user_id="member_1",
team_id="team_a",
event_group=Litellm_EntityType.KEY,
)
assert alert.get_id(user_info) == "hashed_key"

View file

@ -52,6 +52,7 @@ from litellm.proxy.auth.auth_checks import (
_log_budget_lookup_failure,
_tag_max_budget_check,
_team_max_budget_check,
_team_member_max_budget_alert_check,
_virtual_key_max_budget_alert_check,
_check_agent_caller_model_access,
_virtual_key_max_budget_check,
@ -3336,6 +3337,55 @@ async def test_virtual_key_max_budget_alert_check_without_user_obj():
assert captured_call_info.user_email is None
@pytest.mark.parametrize(
"spend, team_metadata, expect_alert",
[
(0.05, {"team_member_max_budget_alert_emails": {"50": [], "100": ["finance@co.com"]}}, True),
(0.10, {"team_member_max_budget_alert_emails": {"50": [], "100": ["finance@co.com"]}}, True),
(0.049, {"team_member_max_budget_alert_emails": {"50": [], "100": ["finance@co.com"]}}, False),
(0.0, {"team_member_max_budget_alert_emails": {"50": []}}, False),
(0.10, {"team_member_max_budget_alert_emails": {"abc": []}}, False),
(0.10, {"team_member_max_budget_alert_emails": "50"}, False),
(0.10, {"soft_budget_alerting_emails": ["finance@co.com"]}, False),
(0.10, None, False),
],
)
@pytest.mark.asyncio
async def test_team_member_max_budget_alert_check_dispatches_only_at_configured_thresholds(
spend, team_metadata, expect_alert
):
captured: list[tuple[str, CallInfo]] = []
class RecordingProxyLogging:
async def budget_alerts(self, type, user_info):
captured.append((type, user_info))
_team_member_max_budget_alert_check(
team_id="team-1",
team_alias="platform",
team_metadata=team_metadata,
organization_id="org-1",
user_id="user-1",
user_email="member@co.com",
proxy_logging_obj=RecordingProxyLogging(),
spend=spend,
max_budget=0.10,
)
await asyncio.sleep(0)
if not expect_alert:
assert captured == [], captured
return
assert [type for type, _ in captured] == ["max_budget_alert"], captured
call_info = captured[0][1]
assert call_info.event_group == Litellm_EntityType.TEAM_MEMBER
assert (call_info.spend, call_info.max_budget) == (spend, 0.10)
assert (call_info.user_id, call_info.user_email) == ("user-1", "member@co.com")
assert (call_info.team_id, call_info.team_alias, call_info.organization_id) == ("team-1", "platform", "org-1")
assert call_info.max_budget_alert_emails == {"50": [], "100": ["finance@co.com"]}
assert call_info.token is None
@pytest.mark.parametrize(
"spend, max_budget, expect_alert",
[

View file

@ -28,6 +28,7 @@ from litellm.proxy._types import (
LiteLLM_OrganizationTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
Litellm_EntityType,
LitellmUserRoles,
ProxyErrorTypes,
ProxyException,
@ -7957,6 +7958,125 @@ async def test_cached_key_team_member_budget_blocks_at_exact_cap(team_member_spe
assert f"TeamMember={user_id}:{team_id}" in exc_info.value.message
@pytest.mark.asyncio
@pytest.mark.parametrize(
"team_member_spend, expect_blocked, expected_alerts",
[
(1.1, False, 0),
(1.2, False, 1),
(2.4, True, 1),
],
)
async def test_cached_key_team_member_budget_emails_configured_thresholds(
team_member_spend, expect_blocked, expected_alerts
):
"""The team's team_member_max_budget_alert_emails thresholds fire from the cached-key auth path,
including on the request that trips the hard cap, and stay silent below the lowest threshold."""
from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_TeamTableCachedObj
from litellm.proxy.common_utils.user_api_key_cache import team_membership_auth_cache_key
from litellm.proxy.utils import hash_token
api_key = "sk-team-member-alert-thresholds"
hashed_token = hash_token(api_key)
team_id = "team-alert-thresholds"
user_id = "user-alert-thresholds"
alert_emails = {"50": [], "100": ["finance@example.com"]}
user_api_key_cache = DualCache()
await _cache_key_object(
hashed_token=hashed_token,
user_api_key_obj=UserAPIKeyAuth(
token=hashed_token,
team_id=team_id,
team_alias="platform",
team_metadata={"team_member_max_budget_alert_emails": alert_emails},
user_id=user_id,
team_member_spend=team_member_spend,
),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=None,
)
await user_api_key_cache.async_set_cache(
key=f"team_id:{team_id}",
value=LiteLLM_TeamTableCachedObj(team_id=team_id),
)
await user_api_key_cache.async_set_cache(
key=user_id,
value=LiteLLM_UserTable(
user_id=user_id, user_email="member@example.com", user_role=LitellmUserRoles.INTERNAL_USER
),
)
await user_api_key_cache.async_set_cache(
key=team_membership_auth_cache_key(team_id=team_id, user_id=user_id),
value=LiteLLM_TeamMembership(
user_id=user_id,
team_id=team_id,
spend=team_member_spend,
budget_id="budget-alert-thresholds",
litellm_budget_table=LiteLLM_BudgetTable(max_budget=2.4),
),
)
mock_request = MagicMock()
mock_request.url.path = "/v1/messages"
mock_request.method = "POST"
mock_request.headers = {"authorization": f"Bearer {api_key}"}
mock_request.query_params = {}
mock_request.state = SimpleNamespace()
proxy_logging_obj = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock(return_value=None)
async def _auth():
return await _user_api_key_auth_builder(
request=mock_request,
api_key=f"Bearer {api_key}",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data={"model": "claude-sonnet-5", "messages": [{"role": "user", "content": "hi"}]},
)
with (
patch( # test-quality-ok: the builder reads proxy settings from module globals, no injection seam
"litellm.proxy.proxy_server.general_settings", {"disable_budget_reservation": True}
),
patch("litellm.proxy.proxy_server.master_key", "sk-master"), # test-quality-ok: module-global proxy state
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: module-global proxy state
patch( # test-quality-ok: seed the cached key, team and membership without a DB
"litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache
),
patch( # test-quality-ok: module-global proxy state
"litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj
),
patch( # test-quality-ok: the live counter needs Redis or a DB; pin the spend the check compares
"litellm.proxy.proxy_server.get_current_spend",
new=AsyncMock(return_value=team_member_spend),
),
):
if expect_blocked:
with pytest.raises(ProxyException) as exc_info:
await _auth()
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
else:
await _auth()
await asyncio.sleep(0)
assert proxy_logging_obj.budget_alerts.await_count == expected_alerts
if expected_alerts == 0:
return
call_info = proxy_logging_obj.budget_alerts.await_args.kwargs["user_info"]
assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "max_budget_alert"
assert call_info.event_group == Litellm_EntityType.TEAM_MEMBER
assert (call_info.spend, call_info.max_budget) == (team_member_spend, 2.4)
assert (call_info.user_id, call_info.user_email) == (user_id, "member@example.com")
assert (call_info.team_id, call_info.team_alias) == (team_id, "platform")
assert call_info.max_budget_alert_emails == alert_emails
@pytest.mark.asyncio
@pytest.mark.parametrize(
"expiry_offset, expect_blocked",

View file

@ -2332,6 +2332,103 @@ describe("TeamInfoView - the exact bytes the update call sends", () => {
expect(wireBody(payload)).toStrictEqual(expected);
});
const openEditorWithMemberBudgetAlerts = async (user: ReturnType<typeof userEvent.setup>) => {
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({
models: ["gpt-4"],
team_member_budget_table: { max_budget: 42 },
metadata: { team_member_max_budget_alert_emails: { "50": [], "100": ["finance@test.com"] } },
}),
);
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any);
renderWithProviders(<TeamInfoView {...props} />);
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
await user.click(screen.getByRole("tab", { name: "Settings" }));
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
await screen.findByLabelText("Team Name");
};
const memberBudgetAlertEmails = (payload: Record<string, unknown>) =>
(wireBody(payload).metadata as Record<string, unknown>).team_member_max_budget_alert_emails;
it("resends the stored team member budget alert thresholds when the section stays closed", async () => {
const user = userEvent.setup({ delay: null });
await openEditorWithMemberBudgetAlerts(user);
const payload = await save(user);
expect(memberBudgetAlertEmails(payload)).toStrictEqual({ "50": [], "100": ["finance@test.com"] });
});
it("sends the edited team member budget alert thresholds as a percent to recipients map", async () => {
const user = userEvent.setup({ delay: null });
await openEditorWithMemberBudgetAlerts(user);
await user.click(screen.getByText("Team Member Settings"));
await screen.findByLabelText("Default Budget (USD)");
const thresholds = screen.getAllByPlaceholderText("% of budget");
const recipients = screen.getAllByPlaceholderText(/Additional recipients/);
expect(thresholds.map((input) => (input as HTMLInputElement).value)).toStrictEqual(["50", "100"]);
expect(recipients.map((input) => (input as HTMLInputElement).value)).toStrictEqual(["", "finance@test.com"]);
fireEvent.change(thresholds[0], { target: { value: "75" } });
fireEvent.change(recipients[0], { target: { value: " lead@test.com, finance@test.com " } });
await user.click(screen.getByRole("button", { name: "Add Budget Alert Threshold" }));
fireEvent.change(screen.getAllByPlaceholderText("% of budget")[2], { target: { value: "90" } });
const payload = await save(user);
expect(memberBudgetAlertEmails(payload)).toStrictEqual({
"75": ["lead@test.com", "finance@test.com"],
"100": ["finance@test.com"],
"90": [],
});
});
it("drops the team member budget alert thresholds key once every row is removed", async () => {
const user = userEvent.setup({ delay: null });
await openEditorWithMemberBudgetAlerts(user);
await user.click(screen.getByText("Team Member Settings"));
await screen.findByLabelText("Default Budget (USD)");
const removeButtons = screen.getAllByRole("button", { name: "Remove budget alert threshold" });
await user.click(removeButtons[1]);
await user.click(removeButtons[0]);
const payload = await save(user);
expect(memberBudgetAlertEmails(payload)).toBeUndefined();
});
it("blocks the save when a team member budget alert threshold is above 100", async () => {
const user = userEvent.setup({ delay: null });
await openEditorWithMemberBudgetAlerts(user);
await user.click(screen.getByText("Team Member Settings"));
await screen.findByLabelText("Default Budget (USD)");
const threshold = screen.getAllByPlaceholderText("% of budget")[0] as HTMLInputElement;
fireEvent.change(threshold, { target: { value: "150" } });
expect(threshold.validity.rangeOverflow).toBe(true);
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(networking.teamUpdateCall).not.toHaveBeenCalled());
});
it("refuses to save a team member budget alert row with no threshold", async () => {
const user = userEvent.setup({ delay: null });
await openEditorWithMemberBudgetAlerts(user);
await user.click(screen.getByText("Team Member Settings"));
await screen.findByLabelText("Default Budget (USD)");
await user.click(screen.getByRole("button", { name: "Add Budget Alert Threshold" }));
await user.click(screen.getByRole("button", { name: /save changes/i }));
await screen.findByText("Enter a whole number from 1 to 100");
expect(networking.teamUpdateCall).not.toHaveBeenCalled();
});
it("carries every typed value to the update payload at the type and shape antd sends today", async () => {
const user = userEvent.setup({ delay: null });
await openEditor(user);

View file

@ -118,12 +118,20 @@ import {
TEAM_INFO_TAB_LABELS,
} from "./tabVisibilityUtils";
import TeamMembersComponent from "./TeamMemberTab";
import {
isValidThreshold,
TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY,
teamMemberBudgetAlertEmailsFromRows,
teamMemberBudgetAlertRowsFromMetadata,
teamMemberBudgetAlertSummary,
} from "./teamMemberBudgetAlertEmails";
import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable";
const UI_MANAGED_METADATA_KEYS: ReadonlySet<string> = new Set([
"logging",
"secret_manager_settings",
"soft_budget_alerting_emails",
TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY,
"model_tpm_limit",
"model_rpm_limit",
"default_estimated_output_tokens",
@ -351,6 +359,18 @@ const teamUpdateFieldsSchema = z.object({
team_member_key_duration: z.string().optional(),
team_member_tpm_limit: numericInputSchema,
team_member_rpm_limit: numericInputSchema,
team_member_max_budget_alert_emails: z
.array(z.object({ threshold: z.number().nullable(), emails: z.string() }))
.superRefine((rows, ctx) => {
rows.forEach((row, index) => {
if (!isValidThreshold(row.threshold)) {
ctx.addIssue({ code: "custom", message: "Enter a whole number from 1 to 100", path: [index, "threshold"] });
} else if (rows.filter((other) => other.threshold === row.threshold).length > 1) {
ctx.addIssue({ code: "custom", message: "Duplicate threshold", path: [index, "threshold"] });
}
});
})
.optional(),
budget_duration: z.string().nullish(),
tpm_limit: numericInputSchema,
rpm_limit: numericInputSchema,
@ -418,6 +438,7 @@ const TEAM_MEMBER_SETTINGS_FIELDS = [
"team_member_key_duration",
"team_member_tpm_limit",
"team_member_rpm_limit",
"team_member_max_budget_alert_emails",
] as const;
const SEARCH_TOOL_SETTINGS_FIELDS = ["object_permission_search_tools"] as const;
@ -433,6 +454,7 @@ const EMPTY_TEAM_UPDATE_VALUES: TeamUpdateFormValues = {
team_member_key_duration: undefined,
team_member_tpm_limit: undefined,
team_member_rpm_limit: undefined,
team_member_max_budget_alert_emails: [],
budget_duration: undefined,
tpm_limit: undefined,
rpm_limit: undefined,
@ -483,6 +505,7 @@ const toTeamFormValues = (info: TeamInfoRecord, effectiveGuardrails: string[]):
team_member_key_duration: info.metadata?.team_member_key_duration,
team_member_tpm_limit: info.team_member_budget_table?.tpm_limit,
team_member_rpm_limit: info.team_member_budget_table?.rpm_limit,
team_member_max_budget_alert_emails: [...teamMemberBudgetAlertRowsFromMetadata(info.metadata)],
budget_duration: info.budget_duration,
tpm_limit: info.tpm_limit,
rpm_limit: info.rpm_limit,
@ -568,6 +591,11 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
append: appendModelLimit,
remove: removeModelLimit,
} = useFieldArray({ control: form.control, name: "modelLimits" });
const {
fields: memberBudgetAlertRows,
append: appendMemberBudgetAlertRow,
remove: removeMemberBudgetAlertRow,
} = useFieldArray({ control: form.control, name: "team_member_max_budget_alert_emails" });
const [teamMemberSettingsOpen, setTeamMemberSettingsOpen] = useState(false);
const [searchToolSettingsOpen, setSearchToolSettingsOpen] = useState(false);
const [isEditMemberModalVisible, setIsEditMemberModalVisible] = useState(false);
@ -966,6 +994,15 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
? { allowed_passthrough_routes: info.metadata.allowed_passthrough_routes }
: {};
const memberBudgetAlertEmails =
values.team_member_max_budget_alert_emails !== undefined
? teamMemberBudgetAlertEmailsFromRows(values.team_member_max_budget_alert_emails)
: info.metadata?.[TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY];
const memberBudgetAlertEmailsMetadata =
memberBudgetAlertEmails !== undefined && Object.keys(memberBudgetAlertEmails).length > 0
? { [TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY]: memberBudgetAlertEmails }
: {};
const updateData: any = {
team_id: teamId,
team_alias: values.team_alias,
@ -997,6 +1034,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
.filter((email: string) => email.length > 0)
: values.soft_budget_alerting_emails || [],
...(secretManagerSettings !== undefined ? { secret_manager_settings: secretManagerSettings } : {}),
...memberBudgetAlertEmailsMetadata,
},
...(values.policies?.length > 0 ? { policies: values.policies } : {}),
...(values.organization_id !== info.organization_id ? { organization_id: values.organization_id ?? null } : {}),
@ -1581,6 +1619,71 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<NumericalInput {...field} ref={ref} value={value ?? ""} step={1} placeholder="e.g., 100" />
)}
</FormField>
<Field>
<FieldLabel>
{labelWithHint(
"Budget Alert Thresholds",
"Email each member when their spend reaches a percentage of their team member budget. The member is always notified; add comma-separated addresses to notify others as well. Requires email alerting to be configured on the proxy.",
)}
</FieldLabel>
{memberBudgetAlertRows.map((row, index) => (
<div key={row.id} className="mb-2 flex items-start gap-2">
<FormField
control={form.control}
name={`team_member_max_budget_alert_emails.${index}.threshold`}
className="w-32"
>
{({ ref, value, onChange, ...field }) => (
<NumericalInput
{...field}
ref={ref}
value={value ?? ""}
onChange={(event: React.ChangeEvent<HTMLInputElement>) =>
onChange(event.target.value === "" ? null : Number(event.target.value))
}
placeholder="% of budget"
min={1}
max={100}
step={1}
/>
)}
</FormField>
<FormField
control={form.control}
name={`team_member_max_budget_alert_emails.${index}.emails`}
className="flex-1"
>
{({ ref, value, ...field }) => (
<UIInput
{...field}
ref={ref}
value={value ?? ""}
placeholder="Additional recipients, e.g. finance@example.com"
/>
)}
</FormField>
<Button
type="button"
variant="ghost"
size="icon"
aria-label="Remove budget alert threshold"
className="mt-1 text-destructive"
onClick={() => removeMemberBudgetAlertRow(index)}
>
<CircleMinus className="size-4" />
</Button>
</div>
))}
<Button
type="button"
variant="outline"
className="w-full border-dashed"
onClick={() => appendMemberBudgetAlertRow({ threshold: null, emails: "" })}
>
<Plus className="size-4" />
Add Budget Alert Threshold
</Button>
</Field>
</FieldGroup>
</CollapsibleContent>
</Collapsible>
@ -2149,6 +2252,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<div>Key Duration: {info.metadata?.team_member_key_duration || "No Limit"}</div>
<div>TPM Limit: {info.team_member_budget_table?.tpm_limit ?? "No Limit"}</div>
<div>RPM Limit: {info.team_member_budget_table?.rpm_limit ?? "No Limit"}</div>
<div>Budget Alert Thresholds: {teamMemberBudgetAlertSummary(info.metadata).join("; ") || "None"}</div>
</div>
<div>
<p className="font-medium">Router Settings</p>

View file

@ -0,0 +1,87 @@
import { describe, expect, it } from "vitest";
import {
isValidThreshold,
teamMemberBudgetAlertEmailsFromRows,
teamMemberBudgetAlertRowsFromMetadata,
teamMemberBudgetAlertSummary,
} from "./teamMemberBudgetAlertEmails";
describe("teamMemberBudgetAlertRowsFromMetadata", () => {
it("turns the stored threshold map into rows sorted by threshold", () => {
const metadata = {
team_member_max_budget_alert_emails: { "100": ["finance@example.com", "cto@example.com"], "50": [] },
};
expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([
{ threshold: 50, emails: "" },
{ threshold: 100, emails: "finance@example.com, cto@example.com" },
]);
});
it("drops non-numeric thresholds and non-list recipients instead of crashing", () => {
const metadata = {
team_member_max_budget_alert_emails: { fifty: [], "75": "finance@example.com", "90": [1], "100": ["a@b.c"] },
};
expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([{ threshold: 100, emails: "a@b.c" }]);
});
it.each([undefined, null, "50", { team_member_max_budget_alert_emails: "50" }, { soft_budget_alerting_emails: [] }])(
"returns no rows for unrelated or malformed metadata %j",
(metadata) => {
expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([]);
},
);
});
describe("teamMemberBudgetAlertEmailsFromRows", () => {
it("builds the threshold map, splitting, trimming and deduplicating recipients", () => {
expect(
teamMemberBudgetAlertEmailsFromRows([
{ threshold: 50, emails: "" },
{ threshold: 100, emails: " finance@example.com,cto@example.com , finance@example.com, " },
]),
).toEqual({ "50": [], "100": ["finance@example.com", "cto@example.com"] });
});
it("skips rows without a valid threshold", () => {
expect(
teamMemberBudgetAlertEmailsFromRows([
{ threshold: null, emails: "finance@example.com" },
{ threshold: 0, emails: "" },
{ threshold: 101, emails: "" },
{ threshold: 12.5, emails: "" },
{ threshold: 80, emails: "" },
]),
).toEqual({ "80": [] });
});
it("round-trips the stored config", () => {
const stored = { team_member_max_budget_alert_emails: { "50": [], "100": ["finance@example.com"] } };
expect(teamMemberBudgetAlertEmailsFromRows(teamMemberBudgetAlertRowsFromMetadata(stored))).toEqual(
stored.team_member_max_budget_alert_emails,
);
});
});
describe("isValidThreshold", () => {
it.each([
[1, true],
[50, true],
[100, true],
[0, false],
[101, false],
[33.3, false],
[null, false],
])("treats %s as valid=%s", (threshold, valid) => {
expect(isValidThreshold(threshold)).toBe(valid);
});
});
describe("teamMemberBudgetAlertSummary", () => {
it("states that the member is always notified and lists extra recipients", () => {
expect(
teamMemberBudgetAlertSummary({
team_member_max_budget_alert_emails: { "100": ["finance@example.com"], "50": [] },
}),
).toEqual(["50%: member", "100%: member, finance@example.com"]);
});
});

View file

@ -0,0 +1,55 @@
export const TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY = "team_member_max_budget_alert_emails" as const;
export interface TeamMemberBudgetAlertRow {
readonly threshold: number | null;
readonly emails: string;
}
export type TeamMemberBudgetAlertEmails = Readonly<Record<string, readonly string[]>>;
const isEmailList = (value: unknown): value is readonly string[] =>
Array.isArray(value) && value.every((email) => typeof email === "string");
const splitEmails = (emails: string): readonly string[] =>
Array.from(
new Set(
emails
.split(",")
.map((email) => email.trim())
.filter((email) => email.length > 0),
),
);
const THRESHOLD_MIN = 1;
const THRESHOLD_MAX = 100;
export const isValidThreshold = (threshold: number | null): threshold is number => {
const isWholeNumber = threshold !== null && Number.isInteger(threshold);
return isWholeNumber && threshold >= THRESHOLD_MIN && threshold <= THRESHOLD_MAX;
};
export const teamMemberBudgetAlertRowsFromMetadata = (metadata: unknown): readonly TeamMemberBudgetAlertRow[] => {
if (typeof metadata !== "object" || metadata === null) return [];
const config: unknown = (metadata as Record<string, unknown>)[TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY];
if (typeof config !== "object" || config === null || Array.isArray(config)) return [];
return Object.entries(config as Record<string, unknown>)
.flatMap(([key, emails]) => {
const threshold = Number(key);
return /^\d+$/.test(key) && isEmailList(emails) ? [{ threshold, emails: emails.join(", ") }] : [];
})
.sort((a, b) => (a.threshold ?? 0) - (b.threshold ?? 0));
};
export const teamMemberBudgetAlertEmailsFromRows = (
rows: readonly TeamMemberBudgetAlertRow[],
): TeamMemberBudgetAlertEmails =>
Object.fromEntries(
rows
.filter((row) => isValidThreshold(row.threshold))
.map((row) => [String(row.threshold), splitEmails(row.emails)]),
);
export const teamMemberBudgetAlertSummary = (metadata: unknown): readonly string[] =>
teamMemberBudgetAlertRowsFromMetadata(metadata).map((row) =>
row.emails.length > 0 ? `${row.threshold}%: member, ${row.emails}` : `${row.threshold}%: member`,
);