scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

fix(proxy): narrow project_id without a cast and fold the unbudgeted cases into the budget matrix test
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
LiteLLM Rust / rust-wheel (push) Has been cancelled
Details
Terraform Modules / fmt, validate, test (aws) (push) Has been cancelled
Details
Terraform Modules / fmt, validate, test (gcp) (push) Has been cancelled
Details

Browse source 
The lint job failed on one new typing.cast (LIT006) in the cost callback, and the test-quality gate behind it would have failed next on a test whose only assertion inspected a mock (TQ002). project_id is now narrowed with isinstance, and the zero and negative max_budget cases run through the existing parametrized budget test, which asserts the raised error or a clean admit with no alert
This commit is contained in:
ryan-crabbe-berri 2026-09-18 14:33:41 -07:00
parent a5f6ce7bb1
commit 595768b54b
2 changed files with 19 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -319,7 +319,11 @@ class _ProxyDBLogger(CustomLogger):
user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None))
team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None))
org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None))
project_id: Final = cast(str | None, metadata.get("user_api_key_project_id", None))
project_id: Final = (
project_id_value
if isinstance(project_id_value := metadata.get("user_api_key_project_id"), str)
else None
)
key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None))
end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None)
sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)

45
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -7559,15 +7559,19 @@ def _project_with_budget(spend: float, max_budget: float):
@pytest.mark.asyncio
@pytest.mark.parametrize(
"counter_spend, db_spend, blocks",
"counter_spend, db_spend, max_budget, blocks",
[
pytest.param(5.0, 0.0, True, id="counter-at-budget-blocks-despite-stale-db-row"),
pytest.param(4.99, 0.0, False, id="counter-under-budget-admits"),
pytest.param(None, 5.0, True, id="no-counter-falls-back-to-persisted-spend"),
pytest.param(None, 0.0, False, id="no-counter-and-no-persisted-spend-admits"),
pytest.param(5.0, 0.0, 5.0, True, id="counter-at-budget-blocks-despite-stale-db-row"),
pytest.param(4.99, 0.0, 5.0, False, id="counter-under-budget-admits"),
pytest.param(None, 5.0, 5.0, True, id="no-counter-falls-back-to-persisted-spend"),
pytest.param(None, 0.0, 5.0, False, id="no-counter-and-no-persisted-spend-admits"),
pytest.param(12.5, 12.5, 0.0, False, id="zero-budget-is-unbudgeted"),
pytest.param(12.5, 12.5, -1.0, False, id="negative-budget-is-unbudgeted"),
],
)
async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks):
async def test_project_max_budget_check_blocks_only_when_live_spend_reaches_a_positive_budget(
counter_spend, db_spend, max_budget, blocks
):
from litellm.caching.dual_cache import DualCache
from litellm.proxy.auth.auth_checks import _project_max_budget_check
@ -7583,14 +7587,16 @@ async def test_project_max_budget_check_reads_live_spend_counter(counter_spend,
):
if not blocks:
await _project_max_budget_check(
project_object=_project_with_budget(spend=db_spend, max_budget=5.0),
project_object=_project_with_budget(spend=db_spend, max_budget=max_budget),
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
await asyncio.sleep(0)
proxy_logging_obj.budget_alerts.assert_not_awaited()
return
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _project_max_budget_check(
project_object=_project_with_budget(spend=db_spend, max_budget=5.0),
project_object=_project_with_budget(spend=db_spend, max_budget=max_budget),
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
@ -7603,29 +7609,6 @@ async def test_project_max_budget_check_reads_live_spend_counter(counter_spend,
assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget"
@pytest.mark.asyncio
@pytest.mark.parametrize("max_budget", [0.0, -1.0])
async def test_project_max_budget_check_treats_non_positive_budget_as_unbudgeted(max_budget):
from litellm.caching.dual_cache import DualCache
from litellm.proxy.auth.auth_checks import _project_max_budget_check
real_spend_counter_cache = DualCache()
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=12.5)
proxy_logging_obj = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
):
await _project_max_budget_check(
project_object=_project_with_budget(spend=12.5, max_budget=max_budget),
valid_token=UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget"),
proxy_logging_obj=proxy_logging_obj,
)
proxy_logging_obj.budget_alerts.assert_not_awaited()
def test_is_user_proxy_admin_rejects_view_only_admin():
"""This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an
Admin Viewer answering True here would gain every write route. Read parity for