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,