From bad8b1dae853bb1811ae19f73bbae729bf78ed2c Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 13 Nov 2026 11:24:05 -0800 Subject: [PATCH] fix(proxy): drop a customer override expiry sent without an override (#51667) * fix(proxy): drop a customer override expiry sent without an override * test(proxy): clear the override before deleting the row so the integration test leaves no budget behind --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management_helpers/customer_hierarchy.py | 16 +++--- .../test_customer_override_expiry.py | 47 +++++++++++++++++ .../management_v1/test_customers.py | 50 +++++++++++++++++++ 3 files changed, 105 insertions(+), 8 deletions(-) create mode 100644 tests/integration/management/test_customer_override_expiry.py diff --git a/litellm/proxy/management_helpers/customer_hierarchy.py b/litellm/proxy/management_helpers/customer_hierarchy.py index 08398cf9a8e..d3ef6ffb6ac 100644 --- a/litellm/proxy/management_helpers/customer_hierarchy.py +++ b/litellm/proxy/management_helpers/customer_hierarchy.py @@ -316,11 +316,13 @@ def team_response(row: PrismaRecord, now: datetime) -> Team: def override_expiry( - current: datetime | None, payload: TeamPatch | MemberPatch, override_id: str | None + row: Team | CustomerMember | None, payload: TeamPatch | MemberPatch, override_id: str | None ) -> datetime | None: + if override_id is None: + return None if "budget_override_expires_at" in payload.model_fields_set: return payload.budget_override_expires_at - return current if override_id is not None else None + return row.budget_override_expires_at if row is not None and row.budget_override_id is not None else None def utc_now() -> datetime: @@ -727,7 +729,7 @@ class CustomerHierarchyStore: "parent_org_id": payload.parent_org_id, "budget_id": budget_id, "budget_override_id": override_id, - "budget_override_expires_at": payload.budget_override_expires_at, + "budget_override_expires_at": override_expiry(None, payload, override_id), "created_by": actor, "updated_by": actor, } @@ -746,7 +748,7 @@ class CustomerHierarchyStore: **alias_patch(payload), "budget_id": budget_id, "budget_override_id": override_id, - "budget_override_expires_at": override_expiry(team.budget_override_expires_at, payload, override_id), + "budget_override_expires_at": override_expiry(team, payload, override_id), "updated_by": actor, }, ) @@ -788,7 +790,7 @@ class CustomerHierarchyStore: "parent_team_ids": tuple(dict.fromkeys(payload.parent_team_ids)), "budget_id": template_id, "budget_override_id": override_id, - "budget_override_expires_at": payload.budget_override_expires_at, + "budget_override_expires_at": override_expiry(None, payload, override_id), "created_by": actor, "updated_by": actor, } @@ -831,9 +833,7 @@ class CustomerHierarchyStore: **( { "budget_override_id": override_id, - "budget_override_expires_at": override_expiry( - member.budget_override_expires_at, payload, override_id - ), + "budget_override_expires_at": override_expiry(member, payload, override_id), } if payload.model_fields_set & {"budget_override", "budget_override_expires_at"} else {} diff --git a/tests/integration/management/test_customer_override_expiry.py b/tests/integration/management/test_customer_override_expiry.py new file mode 100644 index 00000000000..1e08db467dc --- /dev/null +++ b/tests/integration/management/test_customer_override_expiry.py @@ -0,0 +1,47 @@ +import uuid +from typing import Final + +import pytest + +from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value + +pytestmark: Final = pytest.mark.shares_worker_state + +_ROUTE: Final = "/management/v1" + + +def _created(scenario: Scenario, collection: str, body: dict[str, object], identity: str) -> None: + response: Final = scenario.gateway.request("POST", f"{_ROUTE}/{collection}", body) # pyright: ignore[reportArgumentType] # JsonValue body + assert response.status_code == 201, response.text + scenario.cleanups.callback(scenario.gateway.request, "DELETE", f"{_ROUTE}/{collection}/{identity}") + + +@pytest.mark.parametrize("collection", ("customer_members", "customer_teams")) +def test_an_override_applies_after_a_past_expiry_was_sent_without_one(gateway: Gateway, collection: str) -> None: + suffix: Final = uuid.uuid4().hex + org, team, row = f"override-org-{suffix}", f"override-team-{suffix}", f"override-row-{suffix}" + past_expiry: Final = {"budget_override_expires_at": "2020-01-01T00:00:00Z"} + with gateway.scenario() as scenario: + _created( + scenario, + "customer_organizations", + {"customer_org_id": org, "default_budgets": [{"name": "base", "max_budget": 10}]}, + org, + ) + _created(scenario, "customer_teams", {"customer_team_id": team, "parent_org_id": org}, team) + _created( + scenario, + collection, + {"customer_member_id": row, "parent_team_ids": [team], **past_expiry} + if collection == "customer_members" + else {"customer_team_id": row, "parent_org_id": org, **past_expiry}, + row, + ) + scenario.cleanups.callback(gateway.request, "PATCH", f"{_ROUTE}/{collection}/{row}", {"budget_override": None}) + overridden: Final = gateway.request( + "PATCH", f"{_ROUTE}/{collection}/{row}", {"budget_override": {"max_budget": 50}} + ) + assert overridden.status_code == 200, overridden.text + body: Final = JSON_OBJECT.validate_json(overridden.content) + assert body["budget_override_expires_at"] is None + assert object_value(body["effective_budget"])["max_budget"] == 50 diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py index f7213be2eca..84aadd7c018 100644 --- a/tests/unit/proxy/management_endpoints/management_v1/test_customers.py +++ b/tests/unit/proxy/management_endpoints/management_v1/test_customers.py @@ -352,6 +352,56 @@ def test_clearing_an_override_also_clears_its_expiry() -> None: assert renewed.json()["effective_budget"]["max_budget"] == 7 +_PAST_EXPIRY: Final = "2020-01-01T00:00:00Z" + + +@pytest.mark.parametrize( + ("route", "created_without_override"), + ( + ("/management/v1/customer_members", {"customer_member_id": "row", "parent_team_ids": ["team"]}), + ("/management/v1/customer_teams", {"customer_team_id": "row", "parent_org_id": "org"}), + ), + ids=("member", "team"), +) +def test_an_expiry_sent_without_an_override_is_not_stored_and_a_later_override_applies( + route: str, created_without_override: dict[str, object] +) -> None: + http: Final = client() + family(http) + created: Final = http.post(route, json={**created_without_override, "budget_override_expires_at": _PAST_EXPIRY}) + assert created.status_code == 201, created.text + assert created.json()["budget_override_expires_at"] is None + expiry_only: Final = http.patch(f"{route}/row", json={"budget_override_expires_at": _PAST_EXPIRY}) + assert expiry_only.status_code == 200, expiry_only.text + assert expiry_only.json()["budget_override_expires_at"] is None + overridden: Final = http.patch(f"{route}/row", json={"budget_override": {"max_budget": 7}}) + assert overridden.status_code == 200, overridden.text + assert overridden.json()["budget_override_expires_at"] is None + assert overridden.json()["effective_budget"]["max_budget"] == 7 + + +def test_an_override_added_to_a_row_holding_a_stale_expiry_does_not_inherit_it() -> None: + database: Final = Database() + http: Final = client(database=database) + family(http) + created: Final = http.post( + "/management/v1/customer_members", json={"customer_member_id": "member", "parent_team_ids": ["team"]} + ) + assert created.status_code == 201, created.text + asyncio.run( + database.litellm_customermembertable.update( + where={"customer_member_id": "member"}, + data={"budget_override_expires_at": NOW - timedelta(days=365)}, + ) + ) + overridden: Final = http.patch( + "/management/v1/customer_members/member", json={"budget_override": {"max_budget": 7}} + ) + assert overridden.status_code == 200, overridden.text + assert overridden.json()["budget_override_expires_at"] is None + assert overridden.json()["effective_budget"]["max_budget"] == 7 + + def test_null_budget_clears_the_overall_budget_of_an_organization_and_a_team() -> None: http: Final = client() family(http)