Merge e9fc0d6c2c into dddbaa803d

litellm/litellm_core_utils/duration_parser.py

@@ -22,7 +22,7 @@ _BUDGET_DURATION_WORD_ALIASES: Final[dict[str, str]] = {
 }
 
 
-def _normalize_duration(duration: str) -> str:
+def normalize_duration(duration: str) -> str:
     return _BUDGET_DURATION_WORD_ALIASES.get(duration.strip().lower(), duration)
 
 
@@ -61,7 +61,7 @@ def duration_in_seconds(duration: str) -> int:
 
     Returns time in seconds till when budget needs to be reset
     """
-    value, unit = _extract_from_regex(duration=_normalize_duration(duration))
+    value, unit = _extract_from_regex(duration=normalize_duration(duration))
 
     if unit == "s":
         return value
@@ -89,7 +89,7 @@ def duration_in_seconds_from(start: datetime, duration: str) -> int:
     month of a different length needs its own start to know when it ends. Every other
     unit is a fixed length and ignores `start`.
     """
-    value, unit = _extract_from_regex(duration=_normalize_duration(duration))
+    value, unit = _extract_from_regex(duration=normalize_duration(duration))
     if unit == "mo":
         return _months_in_seconds(start=start, months=value)
     return duration_in_seconds(duration)
@@ -129,7 +129,7 @@ def get_next_standardized_reset_time(
     current_time, _ = _setup_timezone(current_time, timezone_str)
 
     # Parse duration
-    value, unit = _parse_duration(_normalize_duration(duration))
+    value, unit = _parse_duration(normalize_duration(duration))
     if value is None:
         verbose_logger.warning(
             "Unrecognized budget_duration %r; falling back to a next-midnight reset. "
@@ -171,7 +171,7 @@ def get_budget_window_start(duration: str, reset_at: datetime) -> datetime:
     """Start of the budget period that ends at `reset_at`, under the same rules
     `get_next_standardized_reset_time` used to pick it: `30d` and `Nmo` span calendar
     months and a duration it does not recognize resets at the next midnight."""
-    value, unit = _parse_duration(_normalize_duration(duration))
+    value, unit = _parse_duration(normalize_duration(duration))
     if value is None:
         return reset_at - timedelta(days=1)
     match unit:

litellm/proxy/management_endpoints/key_budget_change.py

@@ -9,6 +9,7 @@ from typing import Final, Literal, TypeAlias
 from pydantic import TypeAdapter, ValidationError
 
 from litellm._logging import verbose_proxy_logger
+from litellm.litellm_core_utils.duration_parser import normalize_duration
 from litellm.models.team import BudgetLimitEntry
 from litellm.models.verification_token import LiteLLM_VerificationToken
 from litellm.proxy._types import UpdateKeyRequest
@@ -99,6 +100,30 @@ def _budget_limits_change(data: UpdateKeyRequest, existing: LiteLLM_Verification
     return "tightened" if kept_or_lowered else "loosened"
 
 
+def _reset_schedule(budget_duration: str | None) -> str | None:
+    return None if budget_duration is None else normalize_duration(budget_duration)
+
+
+def _changes_budget_duration(data: UpdateKeyRequest, existing: LiteLLM_VerificationToken) -> bool:
+    return "budget_duration" in data.model_fields_set and _reset_schedule(data.budget_duration) != _reset_schedule(
+        existing.budget_duration
+    )
+
+
+def _changes_budget_id(data: UpdateKeyRequest, existing: LiteLLM_VerificationToken) -> bool:
+    return "budget_id" in data.model_fields_set and data.budget_id != existing.budget_id
+
+
+def _budget_duration_change(data: UpdateKeyRequest, existing: LiteLLM_VerificationToken) -> _FieldChange:
+    """
+    Any new `budget_duration` restarts the reset schedule from now, so only clearing it tightens the
+    stored `max_budget`. A key with no `max_budget` has no cap for the schedule to reset.
+    """
+    if existing.max_budget is None or not _changes_budget_duration(data, existing):
+        return "unchanged"
+    return "tightened" if data.budget_duration is None else "loosened"
+
+
 @dataclass(frozen=True, slots=True)
 class EffectiveKeyBudget:
     max_budget: float | None
@@ -122,12 +147,23 @@ def effective_key_budget(data: UpdateKeyRequest, existing: LiteLLM_VerificationT
 
 
 def classify_key_budget_change(data: UpdateKeyRequest, existing: LiteLLM_VerificationToken) -> KeyBudgetChange:
-    overrides_linked_budget: Final = (
-        existing.budget_id is not None and _max_budget_change(data, existing) != "unchanged"
+    overrides_linked_budget: Final = existing.budget_id is not None and (
+        _max_budget_change(data, existing) != "unchanged" or _changes_budget_duration(data, existing)
     )
-    if data.spend is not None or "soft_budget" in data.model_fields_set or overrides_linked_budget:
+    if (
+        data.spend is not None
+        or "soft_budget" in data.model_fields_set
+        or overrides_linked_budget
+        or _changes_budget_id(data, existing)
+    ):
         return KeyBudgetAdminOnly()
-    changes: Final = frozenset((_max_budget_change(data, existing), _budget_limits_change(data, existing)))
+    changes: Final = frozenset(
+        (
+            _max_budget_change(data, existing),
+            _budget_limits_change(data, existing),
+            _budget_duration_change(data, existing),
+        )
+    )
     if "loosened" in changes:
         return KeyBudgetLoosened()
     if "tightened" in changes:

tests/integration/authorization/test_key_budget_owner_updates.py

@@ -0,0 +1,195 @@
+from collections.abc import Iterator
+from hashlib import sha256
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import pytest
+import yaml
+from integration._support.client import (
+    JSON_OBJECT,
+    Gateway,
+    Scenario,
+    eventually,
+    gateway_from_environment,
+    string_value,
+)
+from integration._support.process import owned_proxy
+from pydantic import JsonValue
+
+pytestmark: Final = pytest.mark.shares_worker_state
+
+CAP: Final = 0.1
+BUDGET_EXCEEDED: Final = 422
+FAST_RESET_TICK: Final = MappingProxyType(
+    {"PROXY_BUDGET_RESCHEDULER_MIN_TIME": "2", "PROXY_BUDGET_RESCHEDULER_MAX_TIME": "2"}
+)
+
+
+@pytest.fixture(scope="module")
+def resetting(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
+    with (
+        gateway_from_environment() as shared,
+        owned_proxy(shared, tmp_path_factory.mktemp("fast-reset"), FAST_RESET_TICK) as replica,
+    ):
+        yield replica
+
+
+def _policy_config(directory: Path, policy: str) -> Path:
+    shared: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+    general_settings: Final = {
+        **JSON_OBJECT.validate_python(shared["general_settings"]),
+        "self_serve_budget_policy": policy,
+    }
+    path: Final = directory / f"{policy}.yaml"
+    path.write_text(yaml.safe_dump({**shared, "general_settings": general_settings}))
+    return path
+
+
+@pytest.fixture(scope="module", params=("lower_only", "ceiling"))
+def self_serve(
+    request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory
+) -> Iterator[tuple[str, Gateway]]:
+    policy: Final = string_value(request.param)
+    directory: Final = tmp_path_factory.mktemp(policy)
+    with (
+        gateway_from_environment() as shared,
+        owned_proxy(shared, directory, {}, config=_policy_config(directory, policy)) as replica,
+    ):
+        yield policy, replica
+
+
+def _owner_credential(scenario: Scenario) -> str:
+    return scenario.key(user_id=scenario.user(user_role="internal_user"))
+
+
+def _owned_key(gateway: Gateway, scenario: Scenario, owner: str, **fields: JsonValue) -> str:
+    response: Final = gateway.request("POST", "/key/generate", fields, key=owner)
+    assert response.status_code == 200, response.text
+    token: Final = string_value(JSON_OBJECT.validate_json(response.content)["key"])
+    scenario.cleanups.callback(scenario.delete_key, token)
+    return token
+
+
+def _key_row(gateway: Gateway, key: str) -> dict[str, JsonValue]:
+    rows: Final = gateway.rows(
+        "SELECT spend, max_budget, budget_duration, budget_id, budget_reset_at::text AS reset_at"
+        ' FROM "LiteLLM_VerificationToken" WHERE token = %s',
+        (sha256(key.encode()).hexdigest(),),
+    )
+    assert len(rows) == 1, rows
+    return rows[0]
+
+
+def _exhaust(gateway: Gateway, key: str) -> None:
+    gateway.post("/key/update", {"key": key, "spend": CAP})
+    assert _key_row(gateway, key)["spend"] == CAP
+
+
+def _chat_status(gateway: Gateway, model: str, key: str) -> int:
+    return gateway.request(
+        "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "budget"}]}, key=key
+    ).status_code
+
+
+def test_key_owner_cannot_shorten_budget_duration_to_restart_an_exhausted_budget(resetting: Gateway) -> None:
+    with resetting.scenario() as scenario:
+        model: Final = scenario.model()
+        owner: Final = _owner_credential(scenario)
+        target: Final = _owned_key(resetting, scenario, owner, models=[model], max_budget=CAP, budget_duration="30d")
+        control: Final = _owned_key(resetting, scenario, owner, models=[model], max_budget=CAP, budget_duration="30d")
+        _exhaust(resetting, target)
+        _exhaust(resetting, control)
+        assert _chat_status(resetting, model, target) == BUDGET_EXCEEDED
+        before: Final = _key_row(resetting, target)
+
+        denied: Final = resetting.request("POST", "/key/update", {"key": target, "budget_duration": "1s"}, key=owner)
+
+        resetting.post("/key/update", {"key": control, "budget_duration": "1s"})
+        eventually(lambda: _key_row(resetting, control)["spend"], lambda spend: spend == 0.0, seconds=60)
+        assert denied.status_code == 403, denied.text
+        assert _key_row(resetting, target) == before
+        assert _chat_status(resetting, model, target) == BUDGET_EXCEEDED
+
+
+def test_key_owner_cannot_swap_or_reschedule_a_tier_key(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        model: Final = scenario.model()
+        low: Final = scenario.budget(max_budget=CAP)
+        high: Final = scenario.budget(max_budget=100.0)
+        owner: Final = _owner_credential(scenario)
+        target: Final = _owned_key(gateway, scenario, owner, models=[model], budget_id=low)
+        _exhaust(gateway, target)
+        assert _chat_status(gateway, model, target) == BUDGET_EXCEEDED
+        before: Final = _key_row(gateway, target)
+
+        swapped: Final = gateway.request("POST", "/key/update", {"key": target, "budget_id": high}, key=owner)
+        rescheduled: Final = gateway.request("POST", "/key/update", {"key": target, "budget_duration": "1s"}, key=owner)
+
+        assert swapped.status_code == 403, swapped.text
+        assert rescheduled.status_code == 403, rescheduled.text
+        assert _key_row(gateway, target) == before
+        assert _chat_status(gateway, model, target) == BUDGET_EXCEEDED
+
+
+def test_team_member_cannot_shorten_budget_duration_on_a_team_key(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        team: Final = scenario.team(team_member_permissions=["/key/update", "/key/info"])
+        member: Final = scenario.key(user_id=scenario.member(team), team_id=team)
+        target: Final = scenario.key(team_id=team, max_budget=CAP, budget_duration="30d")
+        before: Final = _key_row(gateway, target)
+
+        denied: Final = gateway.request("POST", "/key/update", {"key": target, "budget_duration": "1s"}, key=member)
+
+        assert denied.status_code == 403, denied.text
+        assert _key_row(gateway, target) == before
+
+
+def test_key_owner_keeps_edits_that_cannot_restart_a_budget(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        owner: Final = _owner_credential(scenario)
+        capped: Final = _owned_key(gateway, scenario, owner, max_budget=CAP, budget_duration="monthly")
+        uncapped: Final = _owned_key(gateway, scenario, owner)
+
+        renamed: Final = gateway.request("POST", "/key/update", {"key": capped, "key_alias": "renamed"}, key=owner)
+        resent: Final = gateway.request("POST", "/key/update", {"key": capped, "budget_duration": "30d"}, key=owner)
+        scheduled: Final = gateway.request("POST", "/key/update", {"key": uncapped, "budget_duration": "1d"}, key=owner)
+
+        assert renamed.status_code == 200, renamed.text
+        assert resent.status_code == 200, resent.text
+        assert scheduled.status_code == 200, scheduled.text
+        assert _key_row(gateway, uncapped)["budget_duration"] == "1d"
+
+
+def test_self_serve_policy_keeps_tier_keys_admin_only(self_serve: tuple[str, Gateway]) -> None:
+    _, proxy = self_serve
+    with proxy.scenario() as scenario:
+        low: Final = scenario.budget(max_budget=CAP)
+        high: Final = scenario.budget(max_budget=100.0)
+        owner: Final = _owner_credential(scenario)
+        target: Final = _owned_key(proxy, scenario, owner, budget_id=low)
+        before: Final = _key_row(proxy, target)
+
+        swapped: Final = proxy.request("POST", "/key/update", {"key": target, "budget_id": high}, key=owner)
+        rescheduled: Final = proxy.request("POST", "/key/update", {"key": target, "budget_duration": "1d"}, key=owner)
+
+        assert swapped.status_code == 403, swapped.text
+        assert rescheduled.status_code == 403, rescheduled.text
+        assert _key_row(proxy, target) == before
+
+
+def test_self_serve_policy_gates_a_new_reset_period_like_a_raised_cap(self_serve: tuple[str, Gateway]) -> None:
+    policy, proxy = self_serve
+    with proxy.scenario() as scenario:
+        owner: Final = _owner_credential(scenario)
+        shortened: Final = _owned_key(proxy, scenario, owner, max_budget=CAP, budget_duration="30d")
+        cleared: Final = _owned_key(proxy, scenario, owner, max_budget=CAP, budget_duration="30d")
+
+        shorten: Final = proxy.request("POST", "/key/update", {"key": shortened, "budget_duration": "1s"}, key=owner)
+        clear: Final = proxy.request("POST", "/key/update", {"key": cleared, "budget_duration": None}, key=owner)
+
+        assert (shorten.status_code, _key_row(proxy, shortened)["budget_duration"]) == (
+            (403, "30d") if policy == "lower_only" else (200, "1s")
+        ), shorten.text
+        assert clear.status_code == 200, clear.text
+        assert _key_row(proxy, cleared)["budget_duration"] is None

tests/unit/proxy/management_endpoints/test_key_budget_change.py

@@ -45,6 +45,7 @@ def _key(
     max_budget: float | None = None,
     budget_limits: Sequence[Mapping[str, object]] | None = None,
     budget_id: str | None = None,
+    budget_duration: str | None = None,
 ) -> LiteLLM_VerificationToken:
     return LiteLLM_VerificationToken(
         token="hashed-owner-key",
@@ -54,6 +55,7 @@ def _key(
         max_budget=max_budget,
         budget_limits=None if budget_limits is None else [dict(window) for window in budget_limits],
         budget_id=budget_id,
+        budget_duration=budget_duration,
         metadata={},
     )
 
@@ -127,6 +129,56 @@ def _window(duration: str, cap: float) -> Mapping[str, object]:
         (UpdateKeyRequest(key="sk-1", spend=0.0), _key(), KeyBudgetAdminOnly()),
         (UpdateKeyRequest(key="sk-1", soft_budget=1.0), _key(), KeyBudgetAdminOnly()),
         (UpdateKeyRequest(key="sk-1", soft_budget=None, max_budget=5.0), _key(max_budget=10.0), KeyBudgetAdminOnly()),
+        (
+            UpdateKeyRequest(key="sk-1", budget_duration="1s"),
+            _key(max_budget=10.0, budget_duration="30d"),
+            KeyBudgetLoosened(),
+        ),
+        (
+            UpdateKeyRequest(key="sk-1", budget_duration="60d"),
+            _key(max_budget=10.0, budget_duration="30d"),
+            KeyBudgetLoosened(),
+        ),
+        (UpdateKeyRequest(key="sk-1", budget_duration="1d"), _key(max_budget=10.0), KeyBudgetLoosened()),
+        (
+            UpdateKeyRequest(key="sk-1", budget_duration=None),
+            _key(max_budget=10.0, budget_duration="30d"),
+            KeyBudgetTightened(),
+        ),
+        (
+            UpdateKeyRequest(key="sk-1", budget_duration="30d"),
+            _key(max_budget=10.0, budget_duration="30d"),
+            KeyBudgetUnchanged(),
+        ),
+        (
+            UpdateKeyRequest(key="sk-1", budget_duration="30d"),
+            _key(max_budget=10.0, budget_duration="Monthly"),
+            KeyBudgetUnchanged(),
+        ),
+        (UpdateKeyRequest(key="sk-1", budget_duration="1s"), _key(budget_duration="30d"), KeyBudgetUnchanged()),
+        (UpdateKeyRequest(key="sk-1", max_budget=5.0, budget_duration="1s"), _key(), KeyBudgetTightened()),
+        (
+            UpdateKeyRequest(key="sk-1", budget_duration="1s"),
+            _key(budget_id="shared-tier-10", budget_duration="30d"),
+            KeyBudgetAdminOnly(),
+        ),
+        (
+            UpdateKeyRequest(key="sk-1", budget_duration=None),
+            _key(budget_id="shared-tier-10", budget_duration="30d"),
+            KeyBudgetAdminOnly(),
+        ),
+        (
+            UpdateKeyRequest(key="sk-1", budget_id="shared-tier-100"),
+            _key(budget_id="shared-tier-10"),
+            KeyBudgetAdminOnly(),
+        ),
+        (UpdateKeyRequest(key="sk-1", budget_id=None), _key(budget_id="shared-tier-10"), KeyBudgetAdminOnly()),
+        (UpdateKeyRequest(key="sk-1", budget_id="shared-tier-10"), _key(), KeyBudgetAdminOnly()),
+        (
+            UpdateKeyRequest(key="sk-1", budget_id="shared-tier-10"),
+            _key(budget_id="shared-tier-10"),
+            KeyBudgetUnchanged(),
+        ),
     ],
 )
 def test_classify_key_budget_change(
@@ -465,3 +517,60 @@ async def test_owner_cannot_clear_max_budget_by_default(monkeypatch: pytest.Monk
             existing=_key(max_budget=10.0),
         )
     assert exc.value.status_code == 403
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("policy", "data", "existing"),
+    [
+        (None, UpdateKeyRequest(key="sk-1", budget_duration="1s"), _key(max_budget=10.0, budget_duration="30d")),
+        (
+            "lower_only",
+            UpdateKeyRequest(key="sk-1", budget_duration="1s"),
+            _key(max_budget=10.0, budget_duration="30d"),
+        ),
+        ("ceiling", UpdateKeyRequest(key="sk-1", budget_id="shared-tier-100"), _key(budget_id="shared-tier-10")),
+        ("ceiling", UpdateKeyRequest(key="sk-1", budget_duration="1s"), _key(budget_id="shared-tier-10")),
+    ],
+)
+async def test_owner_cannot_restart_or_relink_a_key_budget_without_admin(
+    monkeypatch: pytest.MonkeyPatch, policy: str | None, data: UpdateKeyRequest, existing: LiteLLM_VerificationToken
+) -> None:
+    with pytest.raises(HTTPException) as exc:
+        await _update(monkeypatch, policy=policy, data=data, existing=existing)
+    assert exc.value.status_code == 403
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("policy", "data", "existing"),
+    [
+        (None, UpdateKeyRequest(key="sk-1", budget_duration="1d"), _key(budget_duration="30d")),
+        (None, UpdateKeyRequest(key="sk-1", budget_duration="24h"), _key(max_budget=10.0, budget_duration="daily")),
+        (
+            "lower_only",
+            UpdateKeyRequest(key="sk-1", budget_duration=None),
+            _key(max_budget=10.0, budget_duration="30d"),
+        ),
+        ("ceiling", UpdateKeyRequest(key="sk-1", budget_duration="1s"), _key(max_budget=10.0, budget_duration="30d")),
+    ],
+)
+async def test_owner_changes_budget_duration_when_it_cannot_restart_a_cap_or_policy_allows(
+    monkeypatch: pytest.MonkeyPatch, policy: str | None, data: UpdateKeyRequest, existing: LiteLLM_VerificationToken
+) -> None:
+    assert await _denial(monkeypatch, policy=policy, data=data, existing=existing) is None
+
+
+@pytest.mark.asyncio
+async def test_team_member_cannot_restart_a_team_key_budget_created_by_someone_else(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    with pytest.raises(HTTPException) as exc:
+        await _update(
+            monkeypatch,
+            policy=None,
+            data=UpdateKeyRequest(key="sk-1", budget_duration="1s"),
+            existing=_key(created_by=_TEAM_ADMIN, team_id=_TEAM, max_budget=10.0, budget_duration="30d"),
+            team=_team(["/key/update"]),
+        )
+    assert exc.value.status_code == 403