mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge pull request #41906 from BerriAI/litellm_team_member_budget_source_reset
feat(team): show whether a member follows the team default budget and allow resetting to it
This commit is contained in:
ryan-crabbe-berri
GitHub
commit
5573265013
10 changed files with 883 additions and 13 deletions
|
|
@ -4,7 +4,7 @@ import os
|
|||
from collections.abc import Callable, Mapping
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, TypeAlias
|
||||
|
||||
import httpx
|
||||
from pydantic import (
|
||||
|
|
@ -872,6 +872,7 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/management/v1/teams/{team_id}/members/bulk_update",
|
||||
"/team/member_update",
|
||||
"/team/{team_id}/member/{user_id}/reset_spend",
|
||||
"/team/{team_id}/member/{user_id}/reset_budget",
|
||||
"/team/permissions_list",
|
||||
"/team/permissions_update",
|
||||
"/team/daily/activity",
|
||||
|
|
@ -4627,11 +4628,26 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable):
|
|||
caller_edit_access: TeamEditAccess = Field(default_factory=TeamEditNone)
|
||||
|
||||
|
||||
TeamMemberBudgetSource: TypeAlias = Literal["team_default", "custom", "none"]
|
||||
|
||||
|
||||
class TeamInfoMembership(LiteLLM_TeamMembership):
|
||||
budget_source: TeamMemberBudgetSource
|
||||
|
||||
|
||||
class TeamInfoResponseObject(TypedDict):
|
||||
team_id: str
|
||||
team_info: TeamInfoResponseObjectTeamTable
|
||||
keys: list
|
||||
team_memberships: list[LiteLLM_TeamMembership]
|
||||
team_memberships: ReadOnly[tuple[TeamInfoMembership, ...]]
|
||||
|
||||
|
||||
class TeamMemberResetBudgetResponse(BaseModel):
|
||||
team_id: str
|
||||
user_id: str
|
||||
budget_id: str | None
|
||||
previous_budget_id: str | None
|
||||
budget_source: TeamMemberBudgetSource
|
||||
|
||||
|
||||
class TeamListResponseObject(LiteLLM_TeamTable):
|
||||
|
|
|
|||
|
|
@ -80,11 +80,14 @@ from litellm.proxy._types import (
|
|||
TeamEditNone,
|
||||
TeamEditUnrestricted,
|
||||
TeamInfoMember,
|
||||
TeamInfoMembership,
|
||||
TeamInfoResponseObject,
|
||||
TeamInfoResponseObjectTeamTable,
|
||||
TeamListResponseObject,
|
||||
TeamMemberAddRequest,
|
||||
TeamMemberBudgetSource,
|
||||
TeamMemberDeleteRequest,
|
||||
TeamMemberResetBudgetResponse,
|
||||
TeamMemberUpdateRequest,
|
||||
TeamMemberUpdateResponse,
|
||||
TeamModelAddRequest,
|
||||
|
|
@ -4058,6 +4061,99 @@ async def reset_team_member_spend_fn(
|
|||
}
|
||||
|
||||
|
||||
class _TeamMetadataView(BaseModel):
|
||||
metadata: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
def _team_default_budget_id(team: LiteLLM_TeamTable) -> str | None:
|
||||
view: Final = _TeamMetadataView.model_validate(team, from_attributes=True)
|
||||
raw: Final = view.metadata.get("team_member_budget_id") if view.metadata is not None else None
|
||||
return raw if isinstance(raw, str) else None
|
||||
|
||||
|
||||
async def _existing_team_default_budget_id(team: LiteLLM_TeamTable, prisma_client: PrismaClient) -> str | None:
|
||||
budget_id: Final = _team_default_budget_id(team)
|
||||
if budget_id is None:
|
||||
return None
|
||||
row: Final = await _budget_db(prisma_client).find_unique(
|
||||
where={"budget_id": budget_id}, # mutable-ok: prisma client requires a plain dict where= argument
|
||||
)
|
||||
return budget_id if row is not None else None
|
||||
|
||||
|
||||
def _member_budget_source(budget_id: str | None, team_default_budget_id: str | None) -> TeamMemberBudgetSource:
|
||||
if budget_id is not None and budget_id != team_default_budget_id:
|
||||
return "custom"
|
||||
return "team_default" if team_default_budget_id is not None else "none"
|
||||
|
||||
|
||||
@router.post(
|
||||
"/team/{team_id}/member/{user_id}/reset_budget",
|
||||
tags=["team management"], # mutable-ok: FastAPI's `tags` param is typed as list[str], not Sequence
|
||||
dependencies=(Depends(user_api_key_auth),),
|
||||
response_model=TeamMemberResetBudgetResponse,
|
||||
)
|
||||
@management_endpoint_wrapper
|
||||
async def reset_team_member_budget_fn(
|
||||
team_id: str,
|
||||
user_id: str,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> TeamMemberResetBudgetResponse:
|
||||
"""
|
||||
Put a team member back on the team's shared default member budget (`team_member_budget`).
|
||||
|
||||
Drops the member's own budget row link so team-wide changes made through /team/update
|
||||
reach them again. Leaves the member with no budget when the team has no default. Spend is untouched.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
|
||||
if prisma_client is None:
|
||||
_raise_reset_spend_error(status.HTTP_500_INTERNAL_SERVER_ERROR, "DB not connected. prisma_client is None")
|
||||
|
||||
team_obj: Final = await get_team_object(
|
||||
team_id=team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
check_db_only=True,
|
||||
)
|
||||
await _verify_team_access(team_obj=team_obj, user_api_key_dict=user_api_key_dict)
|
||||
|
||||
membership_where: Final = { # mutable-ok: prisma client requires a plain dict where= argument
|
||||
"user_id_team_id": {"user_id": user_id, "team_id": team_id} # mutable-ok: same prisma where= argument
|
||||
}
|
||||
membership_row: Final = await _team_membership_db(prisma_client).find_unique(where=membership_where)
|
||||
if membership_row is None:
|
||||
_raise_reset_spend_error(status.HTTP_404_NOT_FOUND, f"User {user_id} is not a member of team {team_id}.")
|
||||
|
||||
team_default_budget_id: Final = await _existing_team_default_budget_id(team_obj, prisma_client)
|
||||
budget_link: Final = (
|
||||
{
|
||||
"connect": {"budget_id": team_default_budget_id}
|
||||
} # mutable-ok: prisma client requires a plain dict data= argument
|
||||
if team_default_budget_id is not None
|
||||
else {"disconnect": True} # mutable-ok: same prisma data= argument
|
||||
)
|
||||
await _team_membership_db(prisma_client).update(
|
||||
where=membership_where,
|
||||
data={"litellm_budget_table": budget_link}, # mutable-ok: prisma client requires a plain dict data= argument
|
||||
)
|
||||
await invalidate_team_member_spend_state(
|
||||
user_id=user_id,
|
||||
team_id=team_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
return TeamMemberResetBudgetResponse(
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
budget_id=team_default_budget_id,
|
||||
previous_budget_id=membership_row.budget_id,
|
||||
budget_source=_member_budget_source(team_default_budget_id, team_default_budget_id),
|
||||
)
|
||||
|
||||
|
||||
def _create_results_from_response(
|
||||
members: list[Member],
|
||||
response: TeamAddMemberResponse,
|
||||
|
|
@ -4826,15 +4922,16 @@ async def team_info(
|
|||
_team_info = TeamInfoResponseObjectTeamTable()
|
||||
|
||||
## GET TEAM BUDGET (if exists) ##
|
||||
team_member_budget_id: Final = (
|
||||
_team_info.metadata.get("team_member_budget_id") if _team_info.metadata is not None else None
|
||||
)
|
||||
team_member_budget_id: Final = _team_default_budget_id(_team_info)
|
||||
if team_member_budget_id is not None:
|
||||
_team_info = await _add_team_member_budget_table(
|
||||
team_member_budget_id=team_member_budget_id,
|
||||
prisma_client=prisma_client,
|
||||
team_info_response_object=_team_info,
|
||||
)
|
||||
active_default_budget_id: Final = (
|
||||
team_member_budget_id if _team_info.team_member_budget_table is not None else None
|
||||
)
|
||||
|
||||
# Resolve resources inherited from access groups
|
||||
resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info)
|
||||
|
|
@ -4861,7 +4958,17 @@ async def team_info(
|
|||
team_id=team_id,
|
||||
team_info=hydrated_team_info,
|
||||
keys=keys,
|
||||
team_memberships=returned_tm,
|
||||
team_memberships=tuple(
|
||||
TeamInfoMembership.model_validate(
|
||||
MappingProxyType(
|
||||
{
|
||||
**tm.model_dump(),
|
||||
"budget_source": _member_budget_source(tm.budget_id, active_default_budget_id),
|
||||
}
|
||||
)
|
||||
)
|
||||
for tm in returned_tm
|
||||
),
|
||||
)
|
||||
return response_object
|
||||
|
||||
|
|
|
|||
|
|
@ -81,6 +81,7 @@ POST /prompts/test
|
|||
POST /search_tools/test_connection
|
||||
POST /team/bulk_member_add
|
||||
POST /team/{team_id}/member/{user_id}/reset_spend
|
||||
POST /team/{team_id}/member/{user_id}/reset_budget
|
||||
POST /team/key/bulk_update
|
||||
POST /team/permissions_bulk_update
|
||||
POST /team/{team_id}/disable_logging
|
||||
|
|
|
|||
197
tests/proxy_behavior/management/test_team_member_reset_budget.py
Normal file
197
tests/proxy_behavior/management/test_team_member_reset_budget.py
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
from .actors import Actor
|
||||
from .conftest import create_scratch_team
|
||||
|
||||
pytestmark = pytest.mark.asyncio(loop_scope="session")
|
||||
|
||||
_SEED_SPEND = 5.0
|
||||
_TEAM_DEFAULT_MAX_BUDGET = 100.0
|
||||
_CUSTOM_MAX_BUDGET = 50.0
|
||||
|
||||
_MATRIX = [
|
||||
("alpha/proxy_admin", Actor.PROXY_ADMIN, "alpha", 200),
|
||||
("alpha/org_admin", Actor.ORG_ADMIN, "alpha", 200),
|
||||
("alpha/team_admin", Actor.TEAM_ADMIN, "alpha", 200),
|
||||
("alpha/internal_user", Actor.INTERNAL_USER, "alpha", 403),
|
||||
("alpha/owner", Actor.OWNER, "alpha", 403),
|
||||
("alpha/unrelated_same_org", Actor.UNRELATED_SAME_ORG, "alpha", 403),
|
||||
("alpha/cross_org_user", Actor.CROSS_ORG_USER, "alpha", 403),
|
||||
("alpha/service_account", Actor.SERVICE_ACCOUNT, "alpha", 403),
|
||||
("alpha/org_b_admin", Actor.ORG_B_ADMIN, "alpha", 403),
|
||||
("beta/proxy_admin", Actor.PROXY_ADMIN, "beta", 200),
|
||||
("beta/org_admin", Actor.ORG_ADMIN, "beta", 403),
|
||||
("beta/team_admin", Actor.TEAM_ADMIN, "beta", 403),
|
||||
("beta/org_b_admin", Actor.ORG_B_ADMIN, "beta", 200),
|
||||
]
|
||||
|
||||
|
||||
async def _seed_budget(prisma, budget_id: str, max_budget: float) -> str:
|
||||
await prisma.db.litellm_budgettable.create(
|
||||
data={
|
||||
"budget_id": budget_id,
|
||||
"max_budget": max_budget,
|
||||
"created_by": "phase4-scratch",
|
||||
"updated_by": "phase4-scratch",
|
||||
}
|
||||
)
|
||||
return budget_id
|
||||
|
||||
|
||||
async def _seed_team_with_default_budget(prisma, world, shape: str, team_id: str, scratch) -> str:
|
||||
default_budget_id = await _seed_budget(prisma, scratch.tag("team-default-budget"), _TEAM_DEFAULT_MAX_BUDGET)
|
||||
metadata = {"team_member_budget_id": default_budget_id}
|
||||
if shape == "alpha":
|
||||
await create_scratch_team(
|
||||
prisma,
|
||||
team_id,
|
||||
organization_id=world.org_a_id,
|
||||
admin_user_ids=[world.keys[Actor.TEAM_ADMIN].user_id],
|
||||
metadata=metadata,
|
||||
)
|
||||
elif shape == "beta":
|
||||
await create_scratch_team(prisma, team_id, organization_id=world.org_b_id, metadata=metadata)
|
||||
else: # pragma: no cover - guard
|
||||
pytest.fail(f"unknown shape={shape}")
|
||||
return default_budget_id
|
||||
|
||||
|
||||
async def _seed_custom_member(prisma, team_id: str, member_id: str, scratch) -> str:
|
||||
custom_budget_id = await _seed_budget(prisma, scratch.tag("custom-budget"), _CUSTOM_MAX_BUDGET)
|
||||
await prisma.db.litellm_teammembership.create(
|
||||
data={
|
||||
"user_id": member_id,
|
||||
"team_id": team_id,
|
||||
"spend": _SEED_SPEND,
|
||||
"litellm_budget_table": {"connect": {"budget_id": custom_budget_id}},
|
||||
}
|
||||
)
|
||||
return custom_budget_id
|
||||
|
||||
|
||||
async def _membership(prisma, team_id: str, member_id: str):
|
||||
row = await prisma.db.litellm_teammembership.find_unique(
|
||||
where={"user_id_team_id": {"user_id": member_id, "team_id": team_id}}
|
||||
)
|
||||
assert row is not None
|
||||
return row
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"actor,shape,expected_status",
|
||||
[(a, sh, s) for (_id, a, sh, s) in _MATRIX],
|
||||
ids=[s[0] for s in _MATRIX],
|
||||
)
|
||||
async def test_team_member_reset_budget_authz_matrix(
|
||||
actor: Actor,
|
||||
shape: str,
|
||||
expected_status: int,
|
||||
proxy_client,
|
||||
prisma,
|
||||
scratch,
|
||||
world,
|
||||
):
|
||||
member_id = scratch.tag("member")
|
||||
default_budget_id = await _seed_team_with_default_budget(prisma, world, shape, scratch.prefix, scratch)
|
||||
custom_budget_id = await _seed_custom_member(prisma, scratch.prefix, member_id, scratch)
|
||||
caller = world.keys[actor]
|
||||
|
||||
resp = await proxy_client.post(
|
||||
f"/team/{scratch.prefix}/member/{member_id}/reset_budget",
|
||||
headers={"Authorization": f"Bearer {caller.cleartext}"},
|
||||
)
|
||||
assert resp.status_code == expected_status, f"{actor.value} {shape}: {resp.status_code} {resp.text}"
|
||||
|
||||
row = await _membership(prisma, scratch.prefix, member_id)
|
||||
assert row.spend == _SEED_SPEND, "reset_budget must never touch spend"
|
||||
if expected_status == 200:
|
||||
assert row.budget_id == default_budget_id
|
||||
body = resp.json()
|
||||
assert body["budget_id"] == default_budget_id
|
||||
assert body["previous_budget_id"] == custom_budget_id
|
||||
assert body["budget_source"] == "team_default"
|
||||
else:
|
||||
assert row.budget_id == custom_budget_id, "denied but budget relinked"
|
||||
|
||||
|
||||
async def test_team_member_reset_budget_leaves_shared_default_row_untouched(proxy_client, prisma, scratch, world):
|
||||
member_id = scratch.tag("member")
|
||||
default_budget_id = await _seed_team_with_default_budget(prisma, world, "alpha", scratch.prefix, scratch)
|
||||
await _seed_custom_member(prisma, scratch.prefix, member_id, scratch)
|
||||
|
||||
resp = await proxy_client.post(
|
||||
f"/team/{scratch.prefix}/member/{member_id}/reset_budget",
|
||||
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
|
||||
default_row = await prisma.db.litellm_budgettable.find_unique(where={"budget_id": default_budget_id})
|
||||
assert default_row is not None and default_row.max_budget == _TEAM_DEFAULT_MAX_BUDGET
|
||||
|
||||
info = await proxy_client.get(
|
||||
f"/team/info?team_id={scratch.prefix}",
|
||||
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
|
||||
)
|
||||
assert info.status_code == 200, info.text
|
||||
memberships = {tm["user_id"]: tm for tm in info.json()["team_memberships"]}
|
||||
assert memberships[member_id]["budget_source"] == "team_default"
|
||||
assert memberships[member_id]["litellm_budget_table"]["max_budget"] == _TEAM_DEFAULT_MAX_BUDGET
|
||||
|
||||
|
||||
async def test_team_member_reset_budget_without_team_default_detaches_member(proxy_client, prisma, scratch, world):
|
||||
member_id = scratch.tag("member")
|
||||
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
|
||||
await _seed_custom_member(prisma, scratch.prefix, member_id, scratch)
|
||||
|
||||
resp = await proxy_client.post(
|
||||
f"/team/{scratch.prefix}/member/{member_id}/reset_budget",
|
||||
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["budget_id"] is None
|
||||
assert resp.json()["budget_source"] == "none"
|
||||
|
||||
row = await _membership(prisma, scratch.prefix, member_id)
|
||||
assert row.budget_id is None
|
||||
assert row.spend == _SEED_SPEND
|
||||
|
||||
|
||||
async def test_team_member_reset_budget_with_deleted_team_default_detaches_member(proxy_client, prisma, scratch, world):
|
||||
member_id = scratch.tag("member")
|
||||
await create_scratch_team(
|
||||
prisma,
|
||||
scratch.prefix,
|
||||
organization_id=world.org_a_id,
|
||||
metadata={"team_member_budget_id": scratch.tag("deleted-budget")},
|
||||
)
|
||||
await _seed_custom_member(prisma, scratch.prefix, member_id, scratch)
|
||||
|
||||
resp = await proxy_client.post(
|
||||
f"/team/{scratch.prefix}/member/{member_id}/reset_budget",
|
||||
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
|
||||
)
|
||||
assert resp.status_code == 200, resp.text
|
||||
assert resp.json()["budget_id"] is None
|
||||
assert resp.json()["budget_source"] == "none"
|
||||
|
||||
row = await _membership(prisma, scratch.prefix, member_id)
|
||||
assert row.budget_id is None
|
||||
|
||||
|
||||
async def test_team_member_reset_budget_missing_team_is_404(proxy_client, world):
|
||||
resp = await proxy_client.post(
|
||||
f"/team/behavior-pin-no-such-team/member/{uuid.uuid4().hex}/reset_budget",
|
||||
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
|
||||
)
|
||||
assert resp.status_code == 404, resp.text
|
||||
|
||||
|
||||
async def test_team_member_reset_budget_missing_membership_is_404(proxy_client, prisma, scratch, world):
|
||||
await create_scratch_team(prisma, scratch.prefix, organization_id=world.org_a_id)
|
||||
resp = await proxy_client.post(
|
||||
f"/team/{scratch.prefix}/member/{uuid.uuid4().hex}/reset_budget",
|
||||
headers={"Authorization": f"Bearer {world.keys[Actor.PROXY_ADMIN].cleartext}"},
|
||||
)
|
||||
assert resp.status_code == 404, resp.text
|
||||
|
|
@ -51,6 +51,7 @@ from litellm.proxy.management_endpoints.team_endpoints import (
|
|||
_verify_team_access,
|
||||
delete_team,
|
||||
list_available_teams,
|
||||
reset_team_member_budget_fn,
|
||||
reset_team_member_spend_fn,
|
||||
router,
|
||||
team_member_add_duplication_check,
|
||||
|
|
@ -15109,6 +15110,221 @@ async def test_reset_team_member_spend_fn_proxy_admin_can_reset_own_spend(monkey
|
|||
assert response["spend"] == 0.0
|
||||
|
||||
|
||||
def _reset_budget_admin() -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user")
|
||||
|
||||
|
||||
def _team_with_default_budget(team_id: str, budget_id: str) -> LiteLLM_TeamTable:
|
||||
return LiteLLM_TeamTable(team_id=team_id, metadata={"team_member_budget_id": budget_id})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_team_member_budget_fn_relinks_custom_member_to_team_default(monkeypatch):
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
real_cache = UserApiKeyCache()
|
||||
await real_cache.async_set_cache(key="team-1_member-1", value="stale-membership")
|
||||
await real_cache.async_set_cache(key="team_membership:member-1:team-1", value="stale-membership")
|
||||
|
||||
membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", spend=10.0, budget_id="custom-b1")
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
|
||||
return_value=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0)
|
||||
)
|
||||
mock_prisma_client.db.litellm_budgettable.update = AsyncMock()
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", real_cache)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
|
||||
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
|
||||
AsyncMock(return_value=_team_with_default_budget("team-1", "team-default-b")),
|
||||
):
|
||||
response = await reset_team_member_budget_fn(
|
||||
team_id="team-1", user_id="member-1", user_api_key_dict=_reset_budget_admin()
|
||||
)
|
||||
|
||||
assert response.budget_id == "team-default-b"
|
||||
assert response.previous_budget_id == "custom-b1"
|
||||
assert response.budget_source == "team_default"
|
||||
mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with(
|
||||
where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}},
|
||||
data={"litellm_budget_table": {"connect": {"budget_id": "team-default-b"}}},
|
||||
)
|
||||
mock_prisma_client.db.litellm_budgettable.update.assert_not_awaited()
|
||||
assert await real_cache.async_get_cache(key="team-1_member-1") is None
|
||||
assert await real_cache.async_get_cache(key="team_membership:member-1:team-1") is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"team_obj, default_row",
|
||||
[
|
||||
(LiteLLM_TeamTable(team_id="team-1"), None),
|
||||
(_team_with_default_budget("team-1", "gone-b"), None),
|
||||
],
|
||||
ids=["no_default_configured", "configured_default_row_missing"],
|
||||
)
|
||||
async def test_reset_team_member_budget_fn_detaches_member_when_team_has_no_usable_default(
|
||||
monkeypatch, team_obj, default_row
|
||||
):
|
||||
mock_prisma_client = MagicMock()
|
||||
membership_row = LiteLLM_TeamMembership(user_id="member-1", team_id="team-1", budget_id="custom-b1")
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.litellm_teammembership.update = AsyncMock(return_value=membership_row)
|
||||
mock_prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=default_row)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
|
||||
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
|
||||
AsyncMock(return_value=team_obj),
|
||||
):
|
||||
response = await reset_team_member_budget_fn(
|
||||
team_id="team-1", user_id="member-1", user_api_key_dict=_reset_budget_admin()
|
||||
)
|
||||
|
||||
assert response.budget_id is None
|
||||
assert response.previous_budget_id == "custom-b1"
|
||||
assert response.budget_source == "none"
|
||||
mock_prisma_client.db.litellm_teammembership.update.assert_awaited_once_with(
|
||||
where={"user_id_team_id": {"user_id": "member-1", "team_id": "team-1"}},
|
||||
data={"litellm_budget_table": {"disconnect": True}},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_team_member_budget_fn_membership_not_found(monkeypatch):
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None)
|
||||
mock_prisma_client.db.litellm_teammembership.update = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
|
||||
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
|
||||
AsyncMock(return_value=_team_with_default_budget("team-1", "team-default-b")),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await reset_team_member_budget_fn(
|
||||
team_id="team-1", user_id="ghost-user", user_api_key_dict=_reset_budget_admin()
|
||||
)
|
||||
assert exc.value.status_code == 404
|
||||
mock_prisma_client.db.litellm_teammembership.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reset_team_member_budget_fn_forbidden_for_non_admin(monkeypatch):
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teammembership.update = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock())
|
||||
|
||||
with patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.management_endpoints.team_endpoints.get_team_object",
|
||||
AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1", members_with_roles=[])),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await reset_team_member_budget_fn(
|
||||
team_id="team-1",
|
||||
user_id="member-1",
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-user", user_id="plain-user"
|
||||
),
|
||||
)
|
||||
assert exc.value.status_code == 403
|
||||
mock_prisma_client.db.litellm_teammembership.update.assert_not_awaited()
|
||||
|
||||
|
||||
async def _team_info_budget_sources(
|
||||
team_row: LiteLLM_TeamTable,
|
||||
memberships: list[LiteLLM_TeamMembership],
|
||||
default_budget_row: LiteLLM_BudgetTable | None,
|
||||
) -> dict[str, str]:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy.management_endpoints import team_endpoints
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
|
||||
mock_prisma.db.litellm_budgettable.find_unique = AsyncMock(return_value=default_budget_row)
|
||||
mock_prisma.get_data = AsyncMock(return_value=[])
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: no live DB here; matches this file's established convention for endpoint-logic unit tests
|
||||
"litellm.proxy.proxy_server.prisma_client", mock_prisma
|
||||
),
|
||||
patch.object( # test-quality-ok: membership lookup is a module-level DB query with no injection point
|
||||
team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)
|
||||
),
|
||||
):
|
||||
response = await team_endpoints.team_info(
|
||||
http_request=MagicMock(spec=Request),
|
||||
team_id=team_row.team_id,
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
|
||||
)
|
||||
return {tm.user_id: tm.budget_source for tm in response["team_memberships"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_info_reports_whether_each_member_follows_the_team_default_budget():
|
||||
sources = await _team_info_budget_sources(
|
||||
team_row=_team_with_default_budget("team-1", "team-default-b"),
|
||||
memberships=[
|
||||
LiteLLM_TeamMembership(user_id="inherits", team_id="team-1", budget_id="team-default-b"),
|
||||
LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"),
|
||||
LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None),
|
||||
],
|
||||
default_budget_row=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0),
|
||||
)
|
||||
|
||||
assert sources == {
|
||||
"inherits": "team_default",
|
||||
"customized": "custom",
|
||||
"unlinked": "team_default",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_info_reports_no_budget_source_when_team_has_no_default():
|
||||
sources = await _team_info_budget_sources(
|
||||
team_row=LiteLLM_TeamTable(team_id="team-1"),
|
||||
memberships=[
|
||||
LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"),
|
||||
LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None),
|
||||
],
|
||||
default_budget_row=None,
|
||||
)
|
||||
|
||||
assert sources == {
|
||||
"customized": "custom",
|
||||
"unlinked": "none",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_info_reports_no_budget_source_when_team_default_row_was_deleted():
|
||||
sources = await _team_info_budget_sources(
|
||||
team_row=_team_with_default_budget("team-1", "deleted-b"),
|
||||
memberships=[
|
||||
LiteLLM_TeamMembership(user_id="customized", team_id="team-1", budget_id="own-b"),
|
||||
LiteLLM_TeamMembership(user_id="unlinked", team_id="team-1", budget_id=None),
|
||||
],
|
||||
default_budget_row=None,
|
||||
)
|
||||
|
||||
assert sources == {
|
||||
"customized": "custom",
|
||||
"unlinked": "none",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_update_invalidates_team_member_spend_state_when_budget_patch_applied(monkeypatch):
|
||||
"""Raising a stuck member's max_budget_in_team via the documented /team/member_update
|
||||
|
|
|
|||
|
|
@ -0,0 +1,16 @@
|
|||
import { useMutation } from "@tanstack/react-query";
|
||||
import { fetchClient } from "@/lib/http/api";
|
||||
|
||||
export interface ResetTeamMemberBudgetParams {
|
||||
teamId: string;
|
||||
userId: string;
|
||||
}
|
||||
|
||||
export const resetTeamMemberBudget = async ({ teamId, userId }: ResetTeamMemberBudgetParams): Promise<void> => {
|
||||
await fetchClient.POST("/team/{team_id}/member/{user_id}/reset_budget", {
|
||||
params: { path: { team_id: teamId, user_id: userId } },
|
||||
});
|
||||
};
|
||||
|
||||
export const useResetTeamMemberBudget = () =>
|
||||
useMutation<void, Error, ResetTeamMemberBudgetParams>({ mutationFn: resetTeamMemberBudget });
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import type { components } from "@/lib/http/schema";
|
||||
import useCan from "@/app/(dashboard)/hooks/useCan";
|
||||
import { organizationKeys, useOrganizations } from "@/app/(dashboard)/hooks/organizations/useOrganizations";
|
||||
import { useQueryClient } from "@tanstack/react-query";
|
||||
|
|
@ -247,10 +248,13 @@ export const retainedMcpToolPermissions = (
|
|||
export const mcpUnresolvableSaveError = (reason: string): string =>
|
||||
`Cannot save MCP tool permissions because ${reason}. Retry once the page has finished loading`;
|
||||
|
||||
export type TeamMemberBudgetSource = components["schemas"]["TeamMemberResetBudgetResponse"]["budget_source"];
|
||||
|
||||
export interface TeamMembership {
|
||||
user_id: string;
|
||||
team_id: string;
|
||||
budget_id: string;
|
||||
budget_id: string | null;
|
||||
budget_source: TeamMemberBudgetSource;
|
||||
spend: number;
|
||||
total_spend: number | null;
|
||||
litellm_budget_table: {
|
||||
|
|
@ -1361,6 +1365,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
|
|||
canEditTeam={canEditTeam}
|
||||
handleMemberDelete={handleMemberDelete}
|
||||
onMemberSpendReset={refreshTeamData}
|
||||
onMemberBudgetReset={refreshTeamData}
|
||||
setSelectedEditMember={setSelectedEditMember}
|
||||
setIsEditMemberModalVisible={setIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={setIsAddMemberModalVisible}
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ const mockSetSelectedEditMember = vi.fn();
|
|||
const mockSetIsEditMemberModalVisible = vi.fn();
|
||||
const mockSetIsAddMemberModalVisible = vi.fn();
|
||||
const mockOnMemberSpendReset = vi.fn();
|
||||
const mockOnMemberBudgetReset = vi.fn();
|
||||
|
||||
const budgetResetIso = new Date(2026, 6, 15, 12, 0, 0).toISOString();
|
||||
|
||||
|
|
@ -74,6 +75,7 @@ const createMockTeamData = (overrides: Partial<TeamData> = {}): TeamData => ({
|
|||
user_id: "user1@test.com",
|
||||
team_id: "team-123",
|
||||
budget_id: "budget1",
|
||||
budget_source: "custom",
|
||||
spend: 100.5,
|
||||
total_spend: 1538.2608,
|
||||
litellm_budget_table: {
|
||||
|
|
@ -126,6 +128,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -142,6 +145,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -161,6 +165,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -180,6 +185,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam: false,
|
||||
handleMemberDelete: mockHandleMemberDelete,
|
||||
onMemberSpendReset: mockOnMemberSpendReset,
|
||||
onMemberBudgetReset: mockOnMemberBudgetReset,
|
||||
setSelectedEditMember: mockSetSelectedEditMember,
|
||||
setIsEditMemberModalVisible: mockSetIsEditMemberModalVisible,
|
||||
setIsAddMemberModalVisible: mockSetIsAddMemberModalVisible,
|
||||
|
|
@ -204,6 +210,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={true}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -231,6 +238,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -258,6 +266,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -274,6 +283,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -293,6 +303,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -309,6 +320,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -326,6 +338,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -346,6 +359,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={true}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -381,6 +395,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={true}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -435,6 +450,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={true}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -466,6 +482,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={true}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -486,6 +503,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={true}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -503,6 +521,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={false}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -521,6 +540,7 @@ describe("TeamMembersComponent", () => {
|
|||
canEditTeam={true}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
|
|
@ -603,4 +623,134 @@ describe("TeamMembersComponent", () => {
|
|||
expect(screen.getByTestId("reset-member-spend")).toBeVisible();
|
||||
});
|
||||
});
|
||||
|
||||
describe("budget source", () => {
|
||||
const teamDataWithDefault = () => {
|
||||
const base = createMockTeamData();
|
||||
return createMockTeamData({
|
||||
team_info: {
|
||||
...base.team_info,
|
||||
team_member_budget_table: { max_budget: 25, budget_duration: null, tpm_limit: null, rpm_limit: null },
|
||||
},
|
||||
team_memberships: [
|
||||
base.team_memberships[0],
|
||||
{
|
||||
user_id: "user2@test.com",
|
||||
team_id: "team-123",
|
||||
budget_id: "team-default-budget",
|
||||
budget_source: "team_default",
|
||||
spend: 0,
|
||||
total_spend: null,
|
||||
litellm_budget_table: {
|
||||
budget_id: "team-default-budget",
|
||||
soft_budget: null,
|
||||
max_budget: 25,
|
||||
max_parallel_requests: null,
|
||||
tpm_limit: null,
|
||||
rpm_limit: null,
|
||||
model_max_budget: null,
|
||||
budget_duration: null,
|
||||
budget_reset_at: null,
|
||||
},
|
||||
},
|
||||
],
|
||||
});
|
||||
};
|
||||
|
||||
const renderTab = (teamData: TeamData, canEditTeam = true) =>
|
||||
renderWithProviders(
|
||||
<TeamMembersComponent
|
||||
teamData={teamData}
|
||||
canEditTeam={canEditTeam}
|
||||
handleMemberDelete={mockHandleMemberDelete}
|
||||
onMemberSpendReset={mockOnMemberSpendReset}
|
||||
onMemberBudgetReset={mockOnMemberBudgetReset}
|
||||
setSelectedEditMember={mockSetSelectedEditMember}
|
||||
setIsEditMemberModalVisible={mockSetIsEditMemberModalVisible}
|
||||
setIsAddMemberModalVisible={mockSetIsAddMemberModalVisible}
|
||||
/>,
|
||||
);
|
||||
|
||||
it("labels each member's budget as Custom or Team default and shows the team amount for inherited members", () => {
|
||||
renderTab(teamDataWithDefault());
|
||||
|
||||
const customRow = screen.getByRole("row", { name: /user1@test\.com/ });
|
||||
const inheritedRow = screen.getByRole("row", { name: /user2@test\.com/ });
|
||||
expect(within(customRow).getByTestId("member-budget-source")).toHaveTextContent("Custom");
|
||||
expect(customRow).toHaveTextContent("$1,000.00");
|
||||
expect(within(inheritedRow).getByTestId("member-budget-source")).toHaveTextContent("Team default");
|
||||
expect(inheritedRow).toHaveTextContent("$25.00");
|
||||
});
|
||||
|
||||
it("shows no source label for a member with neither a custom nor a team budget", () => {
|
||||
renderTab(createMockTeamData({ team_memberships: [] }));
|
||||
|
||||
expect(screen.queryByTestId("member-budget-source")).not.toBeInTheDocument();
|
||||
expect(screen.queryByTestId("reset-member-budget")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("only offers Use team default on customized members, and only to editors", () => {
|
||||
const { unmount } = renderTab(teamDataWithDefault());
|
||||
|
||||
expect(
|
||||
within(screen.getByRole("row", { name: /user1@test\.com/ })).getByTestId("reset-member-budget"),
|
||||
).toBeVisible();
|
||||
expect(
|
||||
within(screen.getByRole("row", { name: /user2@test\.com/ })).queryByTestId("reset-member-budget"),
|
||||
).not.toBeInTheDocument();
|
||||
|
||||
unmount();
|
||||
renderTab(teamDataWithDefault(), false);
|
||||
expect(screen.queryByTestId("reset-member-budget")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("puts the member back on the team default after confirming, then refreshes the team", async () => {
|
||||
const user = userEvent.setup();
|
||||
POST.mockResolvedValue({ data: {} });
|
||||
renderTab(teamDataWithDefault());
|
||||
|
||||
await user.click(screen.getByTestId("reset-member-budget"));
|
||||
|
||||
const dialog = await screen.findByRole("dialog", { name: "Reset Team Member Budget" });
|
||||
expect(dialog).toHaveTextContent("user1@test.com");
|
||||
expect(dialog).toHaveTextContent("team default of $25.00");
|
||||
expect(dialog).toHaveTextContent("Custom budget: $1,000.00");
|
||||
expect(POST).not.toHaveBeenCalled();
|
||||
|
||||
await user.click(within(dialog).getByRole("button", { name: "Use team default" }));
|
||||
|
||||
await waitFor(() => expect(mockOnMemberBudgetReset).toHaveBeenCalledTimes(1));
|
||||
expect(POST).toHaveBeenCalledExactlyOnceWith("/team/{team_id}/member/{user_id}/reset_budget", {
|
||||
params: { path: { team_id: "team-123", user_id: "user1@test.com" } },
|
||||
});
|
||||
expect(mockOnMemberSpendReset).not.toHaveBeenCalled();
|
||||
expect(screen.queryByRole("dialog")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("keeps the dialog open and does not refresh the team when the reset fails", async () => {
|
||||
const user = userEvent.setup();
|
||||
POST.mockRejectedValue(new Error("Team admin cannot reset budgets"));
|
||||
renderTab(teamDataWithDefault());
|
||||
|
||||
await user.click(screen.getByTestId("reset-member-budget"));
|
||||
const dialog = await screen.findByRole("dialog", { name: "Reset Team Member Budget" });
|
||||
await user.click(within(dialog).getByRole("button", { name: "Use team default" }));
|
||||
|
||||
await waitFor(() => expect(POST).toHaveBeenCalledTimes(1));
|
||||
expect(mockOnMemberBudgetReset).not.toHaveBeenCalled();
|
||||
expect(screen.getByRole("dialog", { name: "Reset Team Member Budget" })).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("does not call the API when the dialog is cancelled", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderTab(teamDataWithDefault());
|
||||
|
||||
await user.click(screen.getByTestId("reset-member-budget"));
|
||||
const dialog = await screen.findByRole("dialog", { name: "Reset Team Member Budget" });
|
||||
await user.click(within(dialog).getByRole("button", { name: "Cancel" }));
|
||||
|
||||
await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument());
|
||||
expect(POST).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
import { useResetTeamMemberBudget } from "@/app/(dashboard)/hooks/teams/useResetTeamMemberBudget";
|
||||
import { useResetTeamMemberSpend } from "@/app/(dashboard)/hooks/teams/useResetTeamMemberSpend";
|
||||
import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
|
||||
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
|
||||
import { SimpleTooltip } from "@/components/ui/tooltip";
|
||||
|
|
@ -13,7 +15,15 @@ import { formatNumberWithCommas } from "@/utils/dataUtils";
|
|||
import { isProxyAdminRole, isUserTeamAdminForSingleTeam } from "@/utils/roles";
|
||||
import { CircleHelp } from "lucide-react";
|
||||
import { useState, type ComponentProps } from "react";
|
||||
import { TeamData, TeamMembership } from "./TeamInfo";
|
||||
import { TeamData, TeamMemberBudgetSource, TeamMembership } from "./TeamInfo";
|
||||
|
||||
const BUDGET_SOURCE_LABELS: Record<Exclude<TeamMemberBudgetSource, "none">, string> = {
|
||||
team_default: "Team default",
|
||||
custom: "Custom",
|
||||
};
|
||||
|
||||
const formatBudget = (value: number | null): string =>
|
||||
value === null ? "Unlimited" : `$${formatNumberWithCommas(value, 2)}`;
|
||||
|
||||
export const seedMemberBudgetFields = (
|
||||
record: Member,
|
||||
|
|
@ -37,6 +47,7 @@ interface TeamMemberTabProps {
|
|||
setIsEditMemberModalVisible: (visible: boolean) => void;
|
||||
setIsAddMemberModalVisible: (visible: boolean) => void;
|
||||
onMemberSpendReset: () => void;
|
||||
onMemberBudgetReset: () => void;
|
||||
}
|
||||
|
||||
export default function TeamMemberTab({
|
||||
|
|
@ -47,9 +58,13 @@ export default function TeamMemberTab({
|
|||
setIsEditMemberModalVisible,
|
||||
setIsAddMemberModalVisible,
|
||||
onMemberSpendReset,
|
||||
onMemberBudgetReset,
|
||||
}: TeamMemberTabProps) {
|
||||
const [memberToResetSpend, setMemberToResetSpend] = useState<Member | null>(null);
|
||||
const [memberToResetBudget, setMemberToResetBudget] = useState<Member | null>(null);
|
||||
const { mutate: resetMemberSpend, isPending: isResettingSpend } = useResetTeamMemberSpend();
|
||||
const { mutate: resetMemberBudget, isPending: isResettingBudget } = useResetTeamMemberBudget();
|
||||
const teamDefaultBudget = teamData.team_info.team_member_budget_table?.max_budget ?? null;
|
||||
|
||||
const formatNumber = (value: number | null): string => {
|
||||
if (value === null || value === undefined) return "0";
|
||||
|
|
@ -82,10 +97,19 @@ export default function TeamMemberTab({
|
|||
return membership?.total_spend ?? 0;
|
||||
};
|
||||
|
||||
const getUserBudgetSource = (userId: string | null): TeamMemberBudgetSource => {
|
||||
if (!userId) return "none";
|
||||
const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
|
||||
return membership?.budget_source ?? "none";
|
||||
};
|
||||
|
||||
const getUserBudget = (userId: string | null): number | null => {
|
||||
if (!userId) return null;
|
||||
const membership = teamData.team_memberships.find((tm) => tm.user_id === userId);
|
||||
return membership?.litellm_budget_table?.max_budget ?? null;
|
||||
return (
|
||||
membership?.litellm_budget_table?.max_budget ??
|
||||
(membership?.budget_source === "team_default" ? teamDefaultBudget : null)
|
||||
);
|
||||
};
|
||||
|
||||
// Helper function to get rate limits for a user
|
||||
|
|
@ -182,12 +206,40 @@ export default function TeamMemberTab({
|
|||
render: (record: Member) => <MoneyCell value={getUserTotalSpend(record.user_id)} decimals={2} />,
|
||||
},
|
||||
{
|
||||
title: "Team Member Budget (USD)",
|
||||
title: (
|
||||
<span className="flex items-center gap-1">
|
||||
Team Member Budget (USD)
|
||||
<SimpleTooltip content="Team default follows the team's member budget, so changing it in team settings updates this member too. Custom is set on this member only and ignores later team changes.">
|
||||
<CircleHelp className="size-4" aria-label="Team member budget information" />
|
||||
</SimpleTooltip>
|
||||
</span>
|
||||
),
|
||||
key: "budget",
|
||||
sortValue: (record: Member) => getUserBudget(record.user_id),
|
||||
render: (record: Member) => (
|
||||
<MoneyCell value={getUserBudget(record.user_id)} decimals={2} emptyText="Unlimited" showZero />
|
||||
),
|
||||
render: (record: Member) => {
|
||||
const source = getUserBudgetSource(record.user_id);
|
||||
return (
|
||||
<span className="flex items-center gap-2">
|
||||
<MoneyCell value={getUserBudget(record.user_id)} decimals={2} emptyText="Unlimited" showZero />
|
||||
{source !== "none" && (
|
||||
<Badge variant={source === "custom" ? "outline" : "secondary"} data-testid="member-budget-source">
|
||||
{BUDGET_SOURCE_LABELS[source]}
|
||||
</Badge>
|
||||
)}
|
||||
{source === "custom" && canEditTeam && (
|
||||
<Button
|
||||
variant="link"
|
||||
size="xs"
|
||||
className="h-auto p-0"
|
||||
data-testid="reset-member-budget"
|
||||
onClick={() => setMemberToResetBudget(record)}
|
||||
>
|
||||
Use team default
|
||||
</Button>
|
||||
)}
|
||||
</span>
|
||||
);
|
||||
},
|
||||
},
|
||||
{
|
||||
title: "Budget Reset",
|
||||
|
|
@ -224,6 +276,21 @@ export default function TeamMemberTab({
|
|||
);
|
||||
};
|
||||
|
||||
const handleResetBudget = () => {
|
||||
if (!memberToResetBudget?.user_id) return;
|
||||
resetMemberBudget(
|
||||
{ teamId: teamData.team_id, userId: memberToResetBudget.user_id },
|
||||
{
|
||||
onSuccess: () => {
|
||||
toast.success("Team member budget reset to the team default");
|
||||
setMemberToResetBudget(null);
|
||||
onMemberBudgetReset();
|
||||
},
|
||||
onError: (error) => toast.fromError(parseErrorMessage(error)),
|
||||
},
|
||||
);
|
||||
};
|
||||
|
||||
return (
|
||||
<>
|
||||
<MemberTable
|
||||
|
|
@ -273,6 +340,30 @@ export default function TeamMemberTab({
|
|||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
<Dialog open={memberToResetBudget !== null} onOpenChange={(open) => !open && setMemberToResetBudget(null)}>
|
||||
<DialogContent>
|
||||
<DialogHeader>
|
||||
<DialogTitle>Reset Team Member Budget</DialogTitle>
|
||||
</DialogHeader>
|
||||
<p>
|
||||
Remove the custom budget for{" "}
|
||||
<strong>{memberToResetBudget?.user_email || memberToResetBudget?.user_id}</strong> and put them back on the
|
||||
team default of <strong>{formatBudget(teamDefaultBudget)}</strong>?
|
||||
</p>
|
||||
<p className="text-sm text-muted-foreground">
|
||||
Custom budget: <strong>{formatBudget(getUserBudget(memberToResetBudget?.user_id ?? null))}</strong>. Their
|
||||
spend is kept. Future changes to the team's member budget will apply to them again.
|
||||
</p>
|
||||
<DialogFooter>
|
||||
<Button variant="outline" onClick={() => setMemberToResetBudget(null)}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button onClick={handleResetBudget} disabled={isResettingBudget}>
|
||||
Use team default
|
||||
</Button>
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
71
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
71
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -16316,6 +16316,29 @@ export interface paths {
|
|||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/team/{team_id}/member/{user_id}/reset_budget": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
get?: never;
|
||||
put?: never;
|
||||
/**
|
||||
* Reset Team Member Budget Fn
|
||||
* @description Put a team member back on the team's shared default member budget (`team_member_budget`).
|
||||
*
|
||||
* Drops the member's own budget row link so team-wide changes made through /team/update
|
||||
* reach them again. Leaves the member with no budget when the team has no default. Spend is untouched.
|
||||
*/
|
||||
post: operations["reset_team_member_budget_fn_team__team_id__member__user_id__reset_budget_post"];
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/team/{team_id}/member/{user_id}/reset_spend": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -38777,6 +38800,22 @@ export interface components {
|
|||
/** User Id */
|
||||
user_id?: string | null;
|
||||
};
|
||||
/** TeamMemberResetBudgetResponse */
|
||||
TeamMemberResetBudgetResponse: {
|
||||
/** Budget Id */
|
||||
budget_id: string | null;
|
||||
/**
|
||||
* Budget Source
|
||||
* @enum {string}
|
||||
*/
|
||||
budget_source: "team_default" | "custom" | "none";
|
||||
/** Previous Budget Id */
|
||||
previous_budget_id: string | null;
|
||||
/** Team Id */
|
||||
team_id: string;
|
||||
/** User Id */
|
||||
user_id: string;
|
||||
};
|
||||
/** TeamMemberUpdateRequest */
|
||||
TeamMemberUpdateRequest: {
|
||||
/**
|
||||
|
|
@ -62053,6 +62092,38 @@ export interface operations {
|
|||
};
|
||||
};
|
||||
};
|
||||
reset_team_member_budget_fn_team__team_id__member__user_id__reset_budget_post: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path: {
|
||||
team_id: string;
|
||||
user_id: string;
|
||||
};
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody?: never;
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["TeamMemberResetBudgetResponse"];
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
422: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["HTTPValidationError"];
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
reset_team_member_spend_fn_team__team_id__member__user_id__reset_spend_post: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue