From 1defaa6cd19abfbdbf98e0a59fc4e092a0a2056d Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 13 Nov 2026 18:43:38 +0000 Subject: [PATCH] fix(budgets): invalidate customer hierarchy caches through stored team IDs (#51654) * 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 * 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 Co-authored-by: Ishaan Jaffer <155045088+ishaan-berri@users.noreply.github.com> Co-authored-by: kerry --- .../budget_management_endpoints.py | 32 ++++++++++++------- .../spend/test_tag_budget_enforcement.py | 14 ++++++++ .../test_access_group_management.py | 1 + .../test_budget_endpoints.py | 21 ++++++++++++ .../test_tag_management_endpoints.py | 2 ++ .../test_team_endpoints.py | 4 +++ .../test_management_helpers_utils.py | 3 ++ 7 files changed, 66 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 9d84dde95d2..091b88bb046 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -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} ) diff --git a/tests/integration/spend/test_tag_budget_enforcement.py b/tests/integration/spend/test_tag_budget_enforcement.py index 5430f77a3ec..2b5790359ba 100644 --- a/tests/integration/spend/test_tag_budget_enforcement.py +++ b/tests/integration/spend/test_tag_budget_enforcement.py @@ -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, diff --git a/tests/unit/proxy/management_endpoints/test_access_group_management.py b/tests/unit/proxy/management_endpoints/test_access_group_management.py index 88669827bf8..c1f8f1a8776 100644 --- a/tests/unit/proxy/management_endpoints/test_access_group_management.py +++ b/tests/unit/proxy/management_endpoints/test_access_group_management.py @@ -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, diff --git a/tests/unit/proxy/management_endpoints/test_budget_endpoints.py b/tests/unit/proxy/management_endpoints/test_budget_endpoints.py index a5a50e19fb9..2a5658e463f 100644 --- a/tests/unit/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_budget_endpoints.py @@ -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 diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index 129c224c8eb..109a06e5b84 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -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) diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index 2a96694847a..ed0dca79602 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -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 ( diff --git a/tests/unit/proxy/management_helpers/test_management_helpers_utils.py b/tests/unit/proxy/management_helpers/test_management_helpers_utils.py index 03c023d8d1f..24c9d8b6af1 100644 --- a/tests/unit/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/unit/proxy/management_helpers/test_management_helpers_utils.py @@ -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=[]))