From 1c599cadc8d93e93dcc861ce1353ee2b37c421f1 Mon Sep 17 00:00:00 2001 From: Darien Kindlund Date: Fri, 10 Apr 2026 23:26:00 -0400 Subject: [PATCH] fix(proxy): use model_group for model_max_budget spend tracking cache key The model_max_budget limiter tracks spend in one code path (async_log_success_event) and enforces budget limits in another (is_key_within_model_budget via user_api_key_auth). These two paths used different model name formats to build cache keys: - Tracking used standard_logging_payload["model"], which is the deployment-level model name (e.g. "vertex_ai/claude-opus-4-6@default") - Enforcement used request_data["model"], which is the model group alias (e.g. "claude-opus-4-6") Because the cache keys never matched, the enforcement path always read None for current spend, silently allowing all requests through even after the budget was exceeded. This affected any provider that decorates model names with provider prefixes or version suffixes (Vertex AI, Bedrock, etc.). Fix: use model_group (the user-facing alias) from StandardLoggingPayload for spend tracking, falling back to model when model_group is None. This aligns the tracking cache key with the enforcement cache key. Fixes the same root cause reported in #15223 and #10052. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../proxy/hooks/model_max_budget_limiter.py | 11 +- ...test_unit_test_max_model_budget_limiter.py | 149 ++++++++++++++++++ 2 files changed, 159 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 5e48ef2879ee..95ffafb7bad4 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -255,7 +255,16 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti return response_cost: float = standard_logging_payload.get("response_cost", 0) - model = standard_logging_payload.get("model") + # Use model_group (the user-facing model alias, e.g. "gpt-4o") when + # available. The enforcement path (is_key_within_model_budget) receives + # the model name from request_data["model"] which is the model group + # alias, so the spend tracking cache key must use the same name. + # Falling back to the deployment-level "model" field preserves + # behaviour for non-proxy or non-router deployments where model_group + # is None. + model = standard_logging_payload.get( + "model_group" + ) or standard_logging_payload.get("model") virtual_key = standard_logging_payload.get("metadata", {}).get( "user_api_key_hash" ) diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 030d452e55fc..b4aac113f578 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -219,6 +219,155 @@ async def test_get_end_user_spend_for_model(budget_limiter): assert spend == 50.0 +@pytest.mark.asyncio +async def test_async_log_success_event_uses_model_group_for_cache_key(budget_limiter): + """ + When model_group is present in StandardLoggingPayload (proxy/router + deployments), spend must be tracked under the model_group name — not the + deployment-level model name — so the cache key matches the one used by + is_key_within_model_budget (which receives request_data["model"], the + model group alias). + + Without this, providers that decorate model names (e.g. Vertex AI + "vertex_ai/claude-opus-4-6@default") track spend under a different cache + key than enforcement reads, silently disabling budget limits. + """ + from litellm.proxy.hooks.model_max_budget_limiter import ( + VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, + ) + + virtual_key = "test-key-hash" + model_group = "claude-opus-4-6" + deployment_model = "vertex_ai/claude-opus-4-6@default" + budget_duration = "1d" + user_api_key_model_max_budget = { + model_group: {"budget_limit": 50.0, "time_period": budget_duration}, + } + kwargs = { + "standard_logging_object": { + "response_cost": 0.10, + "model": deployment_model, + "model_group": model_group, + "metadata": {"user_api_key_hash": virtual_key}, + }, + "litellm_params": { + "metadata": { + "user_api_key_model_max_budget": user_api_key_model_max_budget, + }, + }, + } + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + mock_increment.assert_awaited_once() + call_kwargs = mock_increment.call_args.kwargs + spend_key = call_kwargs["spend_key"] + # The cache key must use the model_group name, NOT the deployment name + assert spend_key == ( + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model_group}:{budget_duration}" + ) + assert call_kwargs["response_cost"] == 0.10 + + +@pytest.mark.asyncio +async def test_async_log_success_event_falls_back_to_model_when_no_model_group( + budget_limiter, +): + """ + When model_group is None (non-proxy / non-router usage), spend tracking + must fall back to using the model field so existing behaviour is preserved. + """ + from litellm.proxy.hooks.model_max_budget_limiter import ( + VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, + ) + + virtual_key = "test-key-hash" + model = "gpt-4" + budget_duration = "1d" + user_api_key_model_max_budget = { + model: {"budget_limit": 100.0, "time_period": budget_duration}, + } + kwargs = { + "standard_logging_object": { + "response_cost": 0.05, + "model": model, + "model_group": None, + "metadata": {"user_api_key_hash": virtual_key}, + }, + "litellm_params": { + "metadata": { + "user_api_key_model_max_budget": user_api_key_model_max_budget, + }, + }, + } + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + mock_increment.assert_awaited_once() + call_kwargs = mock_increment.call_args.kwargs + spend_key = call_kwargs["spend_key"] + assert spend_key == ( + f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{virtual_key}:{model}:{budget_duration}" + ) + + +@pytest.mark.asyncio +async def test_async_log_success_event_end_user_uses_model_group(budget_limiter): + """ + End-user model budget tracking must also use model_group when available, + matching the enforcement path in is_end_user_within_model_budget. + """ + from litellm.proxy.hooks.model_max_budget_limiter import ( + END_USER_SPEND_CACHE_KEY_PREFIX, + ) + + end_user_id = "test-user" + model_group = "claude-sonnet-4-6" + deployment_model = "vertex_ai/claude-sonnet-4-6@default" + budget_duration = "1d" + user_api_key_end_user_model_max_budget = { + model_group: {"budget_limit": 25.0, "time_period": budget_duration}, + } + kwargs = { + "standard_logging_object": { + "response_cost": 0.03, + "model": deployment_model, + "model_group": model_group, + "end_user": end_user_id, + "metadata": {"user_api_key_end_user_id": end_user_id}, + }, + "litellm_params": { + "metadata": { + "user_api_key_end_user_model_max_budget": user_api_key_end_user_model_max_budget, + }, + }, + } + with patch.object( + budget_limiter, + "_increment_spend_for_key", + new_callable=AsyncMock, + ) as mock_increment: + await budget_limiter.async_log_success_event( + kwargs, response_obj=None, start_time=None, end_time=None + ) + mock_increment.assert_awaited_once() + call_kwargs = mock_increment.call_args.kwargs + spend_key = call_kwargs["spend_key"] + assert spend_key == ( + f"{END_USER_SPEND_CACHE_KEY_PREFIX}:{end_user_id}:{model_group}:{budget_duration}" + ) + + @pytest.mark.asyncio async def test_async_log_success_event_uses_end_user_model_budget_duration( budget_limiter,