scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(budgets): invalidate customer hierarchy caches through stored team IDs (#51654)

Browse source 
* fix(budgets): invalidate customer hierarchy caches through stored team IDs

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: restore customer budget alerts after hierarchy schema changes (#51645)

* fix(proxy): upsert customer orgs and customer members safely (#51607)

* fix(proxy): upsert customer orgs and customer members safely

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

* fix(proxy): lock customer org snapshots during upsert

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

* fix(proxy): reject empty customer org upsert identifiers

Co-Authored-By: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>

* fix(proxy): reuse the customer org when its alias is already registered

A request for an alias that another worker registered a moment earlier got
alias_already_exists back from the create and failed. The registration now
treats that the same as an existing id and returns the registered customer org.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
Co-authored-by: kerry <kerry@berri.ai>

* test(spend): synchronize concurrent member resets before validation completes (#51655)

Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(budgets): invalidate customer hierarchy caches through stored team IDs

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(budgets): supply async customer team lookups in caller fixtures

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: ishaan-j-berri <ishaan.j@berri.ai>
Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com>
Co-authored-by: kerry <kerry@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-11-13 18:43:38 +00:00 • committed by GitHub
parent 331a673114
commit 1defaa6cd1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 66 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

32
litellm/proxy/management_endpoints/budget_management_endpoints.py
View file

@ -37,7 +37,12 @@ from litellm.proxy.management_helpers.customer_hierarchy import budget_limit_jso
from litellm.proxy.management_helpers.utils import validate_budget_windows_update
from litellm.proxy.utils import jsonify_object
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.table_repositories import CustomerMemberRepository, CustomerOrganizationRepository
from litellm.repositories.chunked_in import IN_LIST_CHUNK_SIZE
from litellm.repositories.table_repositories import (
CustomerMemberRepository,
CustomerOrganizationRepository,
CustomerTeamRepository,
)
router: Final = APIRouter()
_WRITE_VALUES: Final = TypeAdapter(Mapping[str, object])
@ -239,20 +244,25 @@ async def update_budget(
"OR": (
{"budget_id": budget_obj.budget_id},
{"budget_override_id": budget_obj.budget_id},
{
"parent_team": {
"is": {
"OR": (
{"budget_id": budget_obj.budget_id},
{"parent_org": {"is": {"budget_id": budget_obj.budget_id}}},
)
}
}
},
{"parent_org": {"is": {"budget_id": budget_obj.budget_id}}},
)
}
)
await forget_customer_hierarchies(tuple(member.customer_member_id for member in members), user_api_key_cache)
teams: Final = await CustomerTeamRepository(prisma_client).table.find_many(
where={"OR": ({"budget_id": budget_obj.budget_id}, {"budget_override_id": budget_obj.budget_id})}
)
team_ids: Final = tuple(team.customer_team_id for team in teams)
for offset in range(0, len(team_ids), IN_LIST_CHUNK_SIZE):
await forget_customer_hierarchies(
tuple(
member.customer_member_id
for member in await CustomerMemberRepository(prisma_client).table.find_many(
where={"parent_team_ids": {"has_some": team_ids[offset : offset + IN_LIST_CHUNK_SIZE]}}
)
),
user_api_key_cache,
)
organizations: Final = await CustomerOrganizationRepository(prisma_client).table.find_many(
where={"budget_id": budget_obj.budget_id}
)

14
tests/integration/spend/test_tag_budget_enforcement.py
View file

@ -16,6 +16,20 @@ def _untracked_tag(tag: str) -> str:
return f"{tag}-untracked"
def test_updating_an_unattached_budget_persists_its_limits(gateway: Gateway) -> None:
with gateway.scenario() as scenario:
budget: Final = gateway.post("/budget/new", {"max_budget": 10, "rpm_limit": 20})
budget_id: Final = string_value(budget["budget_id"])
scenario.cleanups.callback(gateway.post, "/budget/delete", {"id": budget_id})
updated: Final = gateway.post(
"/budget/update", {"budget_id": budget_id, "max_budget": 15, "rpm_limit": 30}
)
assert (updated["budget_id"], updated["max_budget"], updated["rpm_limit"]) == (budget_id, 15, 30)
assert read_rows(
'SELECT max_budget, rpm_limit FROM "LiteLLM_BudgetTable" WHERE budget_id=%s', (budget_id,)
) == [{"max_budget": 15.0, "rpm_limit": 30}]
def _tagged_request(
gateway: Gateway,
model: str,

1
tests/unit/proxy/management_endpoints/test_access_group_management.py
View file

@ -845,6 +845,7 @@ class _FakePrismaClient:
self.db = SimpleNamespace(
litellm_budgettable=self.budget_table,
litellm_customermembertable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_customerteamtable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_customerorganizationtable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_modelaccessgroupbudgettable=self.access_group_budget_table,
litellm_proxymodeltable=self.model_table,

21
tests/unit/proxy/management_endpoints/test_budget_endpoints.py
View file

@ -36,6 +36,7 @@ def client_and_mocks(monkeypatch):
mock_prisma.db = types.SimpleNamespace(
litellm_budgettable=mock_table,
litellm_customermembertable=types.SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_customerteamtable=types.SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_customerorganizationtable=types.SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_dailyspend=mock_table,
query_raw=AsyncMock(return_value=[{"referenced": False}]),
@ -143,6 +144,26 @@ def test_budget_edit_evicts_members_and_organization_aliases(
assert cache.get_cache(alias_key, local_only=True) is None
def test_budget_edit_evicts_a_members_inherited_team_budget(
client_and_mocks: tuple[TestClient, MagicMock, MagicMock], monkeypatch: pytest.MonkeyPatch,
) -> None:
client, database, _ = client_and_mocks
cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache())
member_key: Final = "customer_member_id:through-team"
other_key: Final = "customer_member_id:unrelated"
cache.set_cache(member_key, {"max_budget": 10}, local_only=True)
cache.set_cache(other_key, {"max_budget": 20}, local_only=True)
monkeypatch.setattr(ps, "user_api_key_cache", cache)
database.db.litellm_customerteamtable.find_many.return_value = [types.SimpleNamespace(customer_team_id="team")]
database.db.litellm_customermembertable.find_many.side_effect = [
[], [types.SimpleNamespace(customer_member_id="through-team")],
]
response: Final = client.post("/budget/update", json={"budget_id": "team-budget", "max_budget": 3})
assert response.status_code == 200, response.text
assert cache.get_cache(member_key, local_only=True) is None
assert cache.get_cache(other_key, local_only=True) == {"max_budget": 20}
@pytest.mark.asyncio
async def test_new_budget_success(client_and_mocks):
client, _, mock_table = client_and_mocks

2
tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py
View file

@ -516,6 +516,7 @@ async def test_update_tag_clears_or_sets_only_provided_budget_fields(
litellm_proxymodeltable=SimpleNamespace(find_many=find_models),
litellm_budgettable=SimpleNamespace(update=update_budget),
litellm_customermembertable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_customerteamtable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_customerorganizationtable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
)
),
@ -575,6 +576,7 @@ async def test_update_tag_explicit_null_clears_budget_duration():
)
mock_db = Mock()
mock_db.litellm_customermembertable.find_many = AsyncMock(return_value=[])
mock_db.litellm_customerteamtable.find_many = AsyncMock(return_value=[])
mock_db.litellm_customerorganizationtable.find_many = AsyncMock(return_value=[])
mock_prisma = SimpleNamespace(db=mock_db)
mock_db.litellm_tagtable.find_unique = AsyncMock(return_value=existing_tag)

4
tests/unit/proxy/management_endpoints/test_team_endpoints.py
View file

@ -3287,6 +3287,7 @@ async def test_upsert_team_member_budget_table_clears_duration_kept_budget(mock_
team_table.metadata = {"team_member_budget_id": "existing_budget_123"}
mock_db_client.db.litellm_customermembertable.find_many = AsyncMock(return_value=[])
mock_db_client.db.litellm_customerteamtable.find_many = AsyncMock(return_value=[])
mock_db_client.db.litellm_customerorganizationtable.find_many = AsyncMock(return_value=[])
mock_db_client.db.litellm_budgettable.update = AsyncMock(
@ -12006,6 +12007,7 @@ async def test_upsert_team_member_model_max_budget_updates_default_budget_and_ev
budget_cache.async_delete_cache = AsyncMock()
prisma: Final = MagicMock()
prisma.db.litellm_customermembertable.find_many = AsyncMock(return_value=[])
prisma.db.litellm_customerteamtable.find_many = AsyncMock(return_value=[])
prisma.db.litellm_customerorganizationtable.find_many = AsyncMock(return_value=[])
prisma.db.litellm_budgettable.update = AsyncMock(return_value=MagicMock(budget_id="member-budget-1"))
@ -12058,6 +12060,7 @@ async def test_update_team_clears_explicitly_null_member_model_budget_and_evicts
update_budget_response: Final = MagicMock(budget_id="member-budget-1")
prisma: Final = MagicMock()
prisma.db.litellm_customermembertable.find_many = AsyncMock(return_value=[])
prisma.db.litellm_customerteamtable.find_many = AsyncMock(return_value=[])
prisma.db.litellm_customerorganizationtable.find_many = AsyncMock(return_value=[])
with (
@ -12116,6 +12119,7 @@ async def test_update_team_persists_member_model_budget_to_default_budget_row(
update_budget_response: Final = MagicMock(budget_id="member-budget-1")
prisma: Final = MagicMock()
prisma.db.litellm_customermembertable.find_many = AsyncMock(return_value=[])
prisma.db.litellm_customerteamtable.find_many = AsyncMock(return_value=[])
prisma.db.litellm_customerorganizationtable.find_many = AsyncMock(return_value=[])
with (

3
tests/unit/proxy/management_helpers/test_management_helpers_utils.py
View file

@ -308,6 +308,7 @@ async def test_handle_budget_for_entity_updates_with_existing_budget_id_in_reque
db: Final = SimpleNamespace(
litellm_budgettable=SimpleNamespace(update=budget_update, find_unique=AsyncMock(return_value=None)),
litellm_customermembertable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_customerteamtable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_customerorganizationtable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
)
prisma_client: Final = cast(PrismaClient, SimpleNamespace(db=db, writer_db=db))
@ -357,6 +358,7 @@ async def test_handle_budget_for_entity_updates_only_cleared_or_supplied_fields(
db: Final = SimpleNamespace(
litellm_budgettable=SimpleNamespace(update=budget_update, find_unique=AsyncMock(return_value=None)),
litellm_customermembertable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_customerteamtable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
litellm_customerorganizationtable=SimpleNamespace(find_many=AsyncMock(return_value=[])),
)
prisma_client: Final = cast(PrismaClient, SimpleNamespace(db=db, writer_db=db))
@ -983,6 +985,7 @@ class _FakeDb:
self.litellm_teammembership: Final = _FakeMembershipTable(self.litellm_budgettable)
self.litellm_usertable: Final = _FakeUserTable()
self.litellm_customermembertable: Final = SimpleNamespace(find_many=AsyncMock(return_value=[]))
self.litellm_customerteamtable: Final = SimpleNamespace(find_many=AsyncMock(return_value=[]))
self.litellm_customerorganizationtable: Final = SimpleNamespace(find_many=AsyncMock(return_value=[]))