feat(team): show whether a member follows the team default budget and allow resetting to it

Adds budget_source (team_default, custom, none) to each membership in /team/info and a
POST /team/{team_id}/member/{user_id}/reset_budget route that relinks a member to the team's
shared team_member_budget row without touching their spend. The Admin UI team members table
shows a Team default or Custom badge next to each member's budget and offers a
Use team default action on customized members

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan 2026-09-19 00:31:55 +00:00
parent 500e880a40
commit b32d1112a6
8 changed files with 679 additions and 13 deletions

View file

@ -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 (
@ -4588,11 +4588,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):

View file

@ -80,11 +80,14 @@ from litellm.proxy._types import (
TeamEditNone,
TeamEditUnrestricted,
TeamInfoMember,
TeamInfoMembership,
TeamInfoResponseObject,
TeamInfoResponseObjectTeamTable,
TeamListResponseObject,
TeamMemberAddRequest,
TeamMemberBudgetSource,
TeamMemberDeleteRequest,
TeamMemberResetBudgetResponse,
TeamMemberUpdateRequest,
TeamMemberUpdateResponse,
TeamModelAddRequest,
@ -3954,6 +3957,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,
@ -4722,9 +4818,7 @@ 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,
@ -4757,7 +4851,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, team_member_budget_id),
}
)
)
for tm in returned_tm
),
)
return response_object

View file

@ -46,6 +46,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,
@ -14432,6 +14433,219 @@ 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):
"""An admin undoing a per-member budget must put the membership back on the team's shared
default row (a connect, not a copy) so later /team/update changes reach the member again,
and must drop the cached membership so the old cap stops being enforced. The shared row and
the member's tracked spend are never written."""
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
):
"""With no shared default to link to, reset leaves the member exactly where a freshly added
member would be: no budget row at all, reported as budget_source='none', rather than
connecting to a budget_id that does not exist or leaving the custom cap in place."""
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()
@pytest.mark.asyncio
async def test_team_info_reports_whether_each_member_follows_the_team_default_budget():
"""/team/info must tell the caller which members still follow the team's shared member budget
and which carry their own row, since budget_id alone only means something to a reader who
also knows the team's team_member_budget_id."""
from fastapi import Request
from litellm.proxy.management_endpoints import team_endpoints
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),
]
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=LiteLLM_BudgetTable(budget_id="team-default-b", max_budget=100.0)
)
mock_prisma.get_data = AsyncMock(return_value=[])
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)),
):
response = await team_endpoints.team_info(
http_request=MagicMock(spec=Request),
team_id="team-1",
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert {tm.user_id: tm.budget_source for tm in response["team_memberships"]} == {
"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():
"""A team that never set team_member_budget has nothing for members to inherit, so an
unlinked member is 'none' rather than 'team_default', while a member with their own row is
still 'custom'."""
from fastapi import Request
from litellm.proxy.management_endpoints import team_endpoints
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),
]
mock_prisma = MagicMock()
mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=LiteLLM_TeamTable(team_id="team-1"))
mock_prisma.get_data = AsyncMock(return_value=[])
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch.object(team_endpoints, "get_all_team_memberships", AsyncMock(return_value=memberships)),
):
response = await team_endpoints.team_info(
http_request=MagicMock(spec=Request),
team_id="team-1",
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert {tm.user_id: tm.budget_source for tm in response["team_memberships"]} == {
"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

View file

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

View file

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

View file

@ -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();
});
});
});

View file

@ -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&apos;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>
</>
);
}

View file

@ -16272,6 +16272,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;
@ -38547,6 +38570,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: {
/**
@ -61730,6 +61769,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;