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>

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

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

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)