From 5e6a14f42c5d7adb3db04e9b07886bad8233ae61 Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 12:44:15 -0700 Subject: [PATCH 1/2] fix(proxy): gate key budget_duration and budget_id updates like max_budget A key owner or team member could shorten budget_duration or swap budget_id on /key/update without the budget authorization max_budget gets, so the reset job cleared spend on a capped key or a bigger tier lifted its cap Fixes #46916 --- litellm/litellm_core_utils/duration_parser.py | 10 +- .../management_endpoints/key_budget_change.py | 44 +++++- .../test_key_budget_owner_updates.py | 132 ++++++++++++++++++ .../test_key_budget_change.py | 109 +++++++++++++++ 4 files changed, 286 insertions(+), 9 deletions(-) create mode 100644 tests/integration/authorization/test_key_budget_owner_updates.py diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index d52716951df..95bc77e3eba 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/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: diff --git a/litellm/proxy/management_endpoints/key_budget_change.py b/litellm/proxy/management_endpoints/key_budget_change.py index 3143ef53d76..197a13cdb13 100644 --- a/litellm/proxy/management_endpoints/key_budget_change.py +++ b/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: diff --git a/tests/integration/authorization/test_key_budget_owner_updates.py b/tests/integration/authorization/test_key_budget_owner_updates.py new file mode 100644 index 00000000000..f6e85aac5cc --- /dev/null +++ b/tests/integration/authorization/test_key_budget_owner_updates.py @@ -0,0 +1,132 @@ +from collections.abc import Iterator +from hashlib import sha256 +from types import MappingProxyType +from typing import Final + +import pytest +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 _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_budget_id_to_a_higher_tier(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 + + denied: Final = gateway.request("POST", "/key/update", {"key": target, "budget_id": high}, key=owner) + + assert denied.status_code == 403, denied.text + assert _key_row(gateway, target)["budget_id"] == low + 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" diff --git a/tests/unit/proxy/management_endpoints/test_key_budget_change.py b/tests/unit/proxy/management_endpoints/test_key_budget_change.py index baf8d870117..484f8638513 100644 --- a/tests/unit/proxy/management_endpoints/test_key_budget_change.py +++ b/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 From 675d900a7f5c0184248964fbb404d0db6741d3ce Mon Sep 17 00:00:00 2001 From: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:19:16 -0700 Subject: [PATCH 2/2] test(proxy): cover tier keys and self-serve budget policies for key budget updates --- .../test_key_budget_owner_updates.py | 71 +++++++++++++++++-- 1 file changed, 67 insertions(+), 4 deletions(-) diff --git a/tests/integration/authorization/test_key_budget_owner_updates.py b/tests/integration/authorization/test_key_budget_owner_updates.py index f6e85aac5cc..8415c044b6f 100644 --- a/tests/integration/authorization/test_key_budget_owner_updates.py +++ b/tests/integration/authorization/test_key_budget_owner_updates.py @@ -1,9 +1,11 @@ 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, @@ -33,6 +35,30 @@ def resetting(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: 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")) @@ -86,7 +112,7 @@ def test_key_owner_cannot_shorten_budget_duration_to_restart_an_exhausted_budget assert _chat_status(resetting, model, target) == BUDGET_EXCEEDED -def test_key_owner_cannot_swap_budget_id_to_a_higher_tier(gateway: Gateway) -> None: +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) @@ -95,11 +121,14 @@ def test_key_owner_cannot_swap_budget_id_to_a_higher_tier(gateway: Gateway) -> N 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) - denied: Final = gateway.request("POST", "/key/update", {"key": target, "budget_id": high}, key=owner) + 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 denied.status_code == 403, denied.text - assert _key_row(gateway, target)["budget_id"] == low + 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 @@ -130,3 +159,37 @@ def test_key_owner_keeps_edits_that_cannot_restart_a_budget(gateway: Gateway) -> 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