diff --git a/litellm/constants.py b/litellm/constants.py index 28bb7747224..d0596bed684 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1060,6 +1060,9 @@ WANDB_MODELS: set = set( "Qwen/Qwen3-235B-A22B-Thinking-2507", # moonshotai "moonshotai/Kimi-K2-Instruct", + "moonshotai/Kimi-K2.5", + # MiniMaxAI + "MiniMaxAI/MiniMax-M2.5", # meta models "meta-llama/Llama-3.1-8B-Instruct", "meta-llama/Llama-3.3-70B-Instruct", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b2c964c57ed..f7189a60a31 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -32414,6 +32414,34 @@ "litellm_provider": "wandb", "mode": "chat" }, + "wandb/moonshotai/Kimi-K2.5": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 3e-06, + "litellm_provider": "wandb", + "mode": "chat", + "source": "https://wandb.ai/inference/coreweave/cw_moonshotai_Kimi-K2.5", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, + "wandb/MiniMaxAI/MiniMax-M2.5": { + "max_tokens": 197000, + "max_input_tokens": 197000, + "max_output_tokens": 197000, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "wandb", + "mode": "chat", + "source": "https://wandb.ai/inference/coreweave/cw_MiniMaxAI_MiniMax-M2.5", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true + }, "wandb/meta-llama/Llama-3.1-8B-Instruct": { "max_tokens": 128000, "max_input_tokens": 128000, diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bcfaed24398..16243038b78 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1,7 +1,7 @@ import asyncio import json import time -from datetime import datetime, timedelta, timezone +from datetime import datetime, timezone from typing import List, Literal, Optional, Union from litellm._logging import verbose_proxy_logger @@ -652,24 +652,13 @@ class ResetBudgetJob: ) -> LiteLLM_BudgetTableFull: try: if budget.budget_duration is not None: - from litellm.litellm_core_utils.duration_parser import ( - duration_in_seconds, + from litellm.proxy.common_utils.timezone_utils import ( + get_budget_reset_time, ) - duration_s = duration_in_seconds(duration=budget.budget_duration) - - # Fallback for existing budgets that do not have a budget_reset_at date set, ensuring the duration is taken into account - if ( - budget.budget_reset_at is None - and budget.created_at + timedelta(seconds=duration_s) > current_time - ): - budget.budget_reset_at = budget.created_at + timedelta( - seconds=duration_s - ) - else: - budget.budget_reset_at = current_time + timedelta( - seconds=duration_s - ) + budget.budget_reset_at = get_budget_reset_time( + budget_duration=budget.budget_duration + ) except Exception as e: verbose_proxy_logger.exception( "Error resetting budget_reset_at for budget: %s. Item: %s", e, budget diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 5e48ef2879e..95ffafb7bad 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -255,7 +255,16 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): 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/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 41a98fa4ad9..90c0d02d1e0 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -12,11 +12,9 @@ All /budget management endpoints """ #### BUDGET TABLE MANAGEMENT #### -from datetime import timedelta - from fastapi import APIRouter, Depends, HTTPException -from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.utils import jsonify_object @@ -86,8 +84,8 @@ async def new_budget( # if no budget_reset_at date is set, but a budget_duration is given, then set budget_reset_at initially to the first completed duration interval in future if budget_obj.budget_reset_at is None and budget_obj.budget_duration is not None: - budget_obj.budget_reset_at = datetime.utcnow() + timedelta( - seconds=duration_in_seconds(duration=budget_obj.budget_duration) + budget_obj.budget_reset_at = get_budget_reset_time( + budget_duration=budget_obj.budget_duration ) budget_obj_json = budget_obj.model_dump(exclude_none=True) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 8e800c4572b..f69d9d2f8d4 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1849,7 +1849,7 @@ async def _process_single_key_update( # Delete cache await _delete_cache_key_object( - hashed_token=hash_token(key_update_item.key), + hashed_token=_hash_token_if_needed(key_update_item.key), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) @@ -3726,7 +3726,7 @@ async def _execute_virtual_key_regeneration( if hashed_api_key or key: await _delete_cache_key_object( - hashed_token=hash_token(key), + hashed_token=_hash_token_if_needed(key), user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cc24770c367..3ffc0ec7c58 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -32399,6 +32399,34 @@ "litellm_provider": "wandb", "mode": "chat" }, + "wandb/moonshotai/Kimi-K2.5": { + "max_tokens": 262144, + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "cache_read_input_token_cost": 1e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 3e-06, + "litellm_provider": "wandb", + "mode": "chat", + "source": "https://wandb.ai/inference/coreweave/cw_moonshotai_Kimi-K2.5", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true + }, + "wandb/MiniMaxAI/MiniMax-M2.5": { + "max_tokens": 197000, + "max_input_tokens": 197000, + "max_output_tokens": 197000, + "input_cost_per_token": 3e-07, + "output_cost_per_token": 1.2e-06, + "litellm_provider": "wandb", + "mode": "chat", + "source": "https://wandb.ai/inference/coreweave/cw_MiniMaxAI_MiniMax-M2.5", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true + }, "wandb/meta-llama/Llama-3.1-8B-Instruct": { "max_tokens": 128000, "max_input_tokens": 128000, 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 030d452e55f..b4aac113f57 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, diff --git a/tests/test_budget_management.py b/tests/test_budget_management.py index ad1e07ce995..09759763559 100644 --- a/tests/test_budget_management.py +++ b/tests/test_budget_management.py @@ -1,12 +1,25 @@ # What is this? ## Unit tests for the /budget/* endpoints from litellm._uuid import uuid -from datetime import datetime, timedelta +from datetime import datetime, timezone import aiohttp import pytest import pytest_asyncio +from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time +from litellm.proxy.common_utils.timezone_utils import get_budget_reset_timezone + + +def _parse_budget_api_datetime(value: str) -> datetime: + """Parse ISO timestamps returned by the proxy JSON API.""" + if value.endswith("Z"): + value = value[:-1] + "+00:00" + dt = datetime.fromisoformat(value) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return dt + async def delete_budget(session, budget_id): url = "http://0.0.0.0:4000/budget/delete" @@ -61,32 +74,30 @@ async def budget_setup(): @pytest.mark.asyncio async def test_create_budget_with_duration(budget_setup): """ - Test creating a budget with a specified duration and verify that the 'budget_reset_at' - timestamp is correctly calculated as 'created_at' plus the budget duration (one day). - - This test uses the budget_setup fixture, which handles both the creation and cleanup of the budget. + Test creating a budget with a specified duration and verify that 'budget_reset_at' + matches the next standardized reset (see get_budget_reset_time / new_budget), not + necessarily created_at + wall-clock duration. """ - # Verify that the response includes a 'budget_reset_at' timestamp. assert ( budget_setup["budget_reset_at"] is not None ), "The budget_reset_at field should not be None" - # Calculate the expected reset time: created_at + 1 day. - # Replace trailing 'Z' with '+00:00' for Python 3.9 compat (fromisoformat - # only learned to accept 'Z' in Python 3.11). - created_at_str = budget_setup["created_at"].replace("Z", "+00:00") - expected_reset_at_date = datetime.fromisoformat(created_at_str) + timedelta(days=1) + created_at = _parse_budget_api_datetime(budget_setup["created_at"]) + expected_reset_at = get_next_standardized_reset_time( + duration=budget_setup["budget_duration"], + current_time=created_at, + timezone_str=get_budget_reset_timezone(), + ) + + actual_reset_at = _parse_budget_api_datetime(budget_setup["budget_reset_at"]) - # Allow for a small tolerance in seconds for the timestamp calculation. tolerance_seconds = 3 - reset_at_str = budget_setup["budget_reset_at"].replace("Z", "+00:00") - actual_reset_at_date = datetime.fromisoformat(reset_at_str) time_difference = abs( - (actual_reset_at_date - expected_reset_at_date).total_seconds() + (actual_reset_at - expected_reset_at).total_seconds() ) assert time_difference <= tolerance_seconds, ( - f"Expected budget_reset_at to be within {tolerance_seconds} seconds of {expected_reset_at_date}, " + f"Expected budget_reset_at to be within {tolerance_seconds} seconds of {expected_reset_at}, " f"but the difference was {time_difference} seconds." ) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index f975460836a..1f2d4f4905f 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -444,6 +444,140 @@ def test_reset_budget_for_keys_linked_to_budgets_empty( assert len(calls) == 0 +@pytest.mark.parametrize( + "budget_duration, expected_day, expected_month", + [ + ("30d", 1, 7), # 30d → 1st of next month + ("1mo", 1, 7), # 1mo → 1st of next month + ("1d", 16, 6), # 1d → next midnight (same month) + ], + ids=["30d-calendar-month", "1mo-calendar-month", "1d-next-midnight"], +) +def test_reset_budget_reset_at_date_calendar_aligned( + budget_duration, expected_day, expected_month +): + """ + Verify that _reset_budget_reset_at_date produces calendar-aligned reset + times (matching get_budget_reset_time), not sliding-window offsets. + """ + from unittest.mock import patch + + # Fix "now" to 2023-06-15 10:30:00 UTC for deterministic results + fixed_now = datetime(2023, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "budget_duration": budget_duration, + "budget_reset_at": fixed_now - timedelta(hours=1), + "budget_id": "test-budget", + "created_at": fixed_now - timedelta(days=30), + }, + ) + + with patch( + "litellm.proxy.common_utils.timezone_utils.datetime" + ) as mock_dt: + mock_dt.now.return_value = fixed_now + mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) + asyncio.run( + ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now) + ) + + assert test_budget.budget_reset_at.day == expected_day + assert test_budget.budget_reset_at.month == expected_month + assert test_budget.budget_reset_at.hour == 0 + assert test_budget.budget_reset_at.minute == 0 + assert test_budget.budget_reset_at.second == 0 + + +def test_reset_budget_reset_at_date_7d_next_monday(): + """Verify 7d budget duration resets to next Monday at midnight.""" + from unittest.mock import patch + + # 2023-06-14 is a Wednesday + fixed_now = datetime(2023, 6, 14, 10, 30, 0, tzinfo=timezone.utc) + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "budget_duration": "7d", + "budget_reset_at": fixed_now - timedelta(hours=1), + "budget_id": "test-budget", + "created_at": fixed_now - timedelta(days=7), + }, + ) + + with patch( + "litellm.proxy.common_utils.timezone_utils.datetime" + ) as mock_dt: + mock_dt.now.return_value = fixed_now + mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) + asyncio.run( + ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now) + ) + + # Next Monday after Wednesday June 14 is June 19 + assert test_budget.budget_reset_at.day == 19 + assert test_budget.budget_reset_at.month == 6 + assert test_budget.budget_reset_at.weekday() == 0 # Monday + assert test_budget.budget_reset_at.hour == 0 + + +def test_reset_budget_reset_at_date_none_duration(): + """Verify that budget_reset_at is unchanged when budget_duration is None.""" + original_reset_at = datetime(2023, 6, 20, 0, 0, 0, tzinfo=timezone.utc) + now = datetime(2023, 6, 15, 10, 0, 0, tzinfo=timezone.utc) + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "budget_duration": None, + "budget_reset_at": original_reset_at, + "budget_id": "test-budget", + "created_at": now - timedelta(days=30), + }, + ) + + asyncio.run(ResetBudgetJob._reset_budget_reset_at_date(test_budget, now)) + assert test_budget.budget_reset_at == original_reset_at + + +def test_reset_budget_reset_at_date_none_reset_at(): + """Verify that budget_reset_at is set correctly even when previously None.""" + from unittest.mock import patch + + fixed_now = datetime(2023, 6, 15, 10, 30, 0, tzinfo=timezone.utc) + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "budget_duration": "30d", + "budget_reset_at": None, + "budget_id": "test-budget", + "created_at": fixed_now - timedelta(days=5), + }, + ) + + with patch( + "litellm.proxy.common_utils.timezone_utils.datetime" + ) as mock_dt: + mock_dt.now.return_value = fixed_now + mock_dt.side_effect = lambda *args, **kwargs: datetime(*args, **kwargs) + asyncio.run( + ResetBudgetJob._reset_budget_reset_at_date(test_budget, fixed_now) + ) + + # Should be set to 1st of next month (July 1) + assert test_budget.budget_reset_at is not None + assert test_budget.budget_reset_at.day == 1 + assert test_budget.budget_reset_at.month == 7 + + def test_budget_table_reset_also_resets_linked_keys( reset_budget_job, mock_prisma_client ): diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 096e0b2bc41..479defbff5c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -5385,9 +5385,15 @@ async def test_bulk_update_keys_success(monkeypatch): ) as mock_hash: mock_hash.side_effect = ["hashed-key-1", "hashed-key-2"] + def _hash_for_bulk_success(token: str) -> str: + return { + "test-key-1": "hashed-key-1", + "test-key-2": "hashed-key-2", + }[token] + with patch( "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", - side_effect=["hashed-key-1", "hashed-key-2"], + side_effect=_hash_for_bulk_success, ): with patch( "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook" @@ -5511,9 +5517,15 @@ async def test_bulk_update_keys_partial_failures(monkeypatch): ) as mock_hash: mock_hash.return_value = "hashed-key-1" + def _hash_for_bulk_partial(token: str) -> str: + return { + "test-key-1": "hashed-key-1", + "non-existent-key": "hashed-non-existent-key", + }[token] + with patch( "litellm.proxy.management_endpoints.key_management_endpoints._hash_token_if_needed", - side_effect=["hashed-key-1", "hashed-non-existent-key"], + side_effect=_hash_for_bulk_partial, ): with patch( "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook" @@ -8746,3 +8758,167 @@ def test_validate_public_image_url_accepts_http_and_noop_empty(): _validate_public_image_url(None, "logo_url") _validate_public_image_url("", "logo_url") _validate_public_image_url(" ", "logo_url") + + +@pytest.mark.asyncio +async def test_process_single_key_update_cache_invalidation_with_token_hash(): + """ + _process_single_key_update must pass the token hash as-is (not + double-hashed) to _delete_cache_key_object when the key is already a + pre-hashed token ID rather than an sk- prefixed key. + + Without this, cache invalidation silently fails: the wrong cache entry + is deleted while the stale entry (with outdated fields) persists and + gets refreshed indefinitely by update_cache on every successful request. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _process_single_key_update, + ) + from litellm.types.proxy.management_endpoints.key_management_endpoints import ( + BulkUpdateKeyRequestItem, + ) + + token_hash = "abc123def456" + + existing_key = LiteLLM_VerificationToken( + token=token_hash, + user_id="user-1", + models=["gpt-4"], + team_id=None, + max_budget=None, + tags=None, + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=existing_key + ) + mock_updated = MagicMock() + mock_updated.model_dump.return_value = {"max_budget": 100.0} + mock_prisma_client.update_data = AsyncMock(return_value={"data": mock_updated}) + + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + mock_llm_router = MagicMock() + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data", + return_value={"max_budget": 100.0}, + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint", + return_value=None, + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ) as mock_delete_cache, patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_updated_hook", + new_callable=AsyncMock, + ): + key_update_item = BulkUpdateKeyRequestItem( + key=token_hash, + max_budget=100.0, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + await _process_single_key_update( + key_update_item=key_update_item, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + prisma_client=mock_prisma_client, + user_api_key_cache=mock_user_api_key_cache, + proxy_logging_obj=mock_proxy_logging_obj, + llm_router=mock_llm_router, + ) + + mock_delete_cache.assert_called_once() + call_kwargs = mock_delete_cache.call_args.kwargs + # The token hash should be passed as-is, NOT double-hashed + assert call_kwargs["hashed_token"] == token_hash + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_cache_invalidation_with_token_hash(): + """ + _execute_virtual_key_regeneration must pass the token hash as-is (not + double-hashed) to _delete_cache_key_object when the key is a + pre-hashed token ID. + """ + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + token_hash = "abc123def456" + + existing_key = LiteLLM_VerificationToken( + token=token_hash, + user_id="user-1", + models=["gpt-4"], + team_id=None, + max_budget=None, + tags=None, + ) + + mock_prisma_client = AsyncMock() + # _execute_virtual_key_regeneration calls dict(updated_token) which + # needs the return value to be iterable as key-value pairs. + class DictLikeResult: + def __init__(self, data): + self._data = data + def __iter__(self): + return iter(self._data.items()) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=DictLikeResult({"token": "new-hashed-token", "key_name": "sk-...ab12", "user_id": "user-1"}) + ) + mock_prisma_client.db.litellm_verificationtoken.create = AsyncMock( + return_value=None + ) + mock_prisma_client.jsonify_object = MagicMock(side_effect=lambda data: data) + + mock_user_api_key_cache = MagicMock() + mock_proxy_logging_obj = MagicMock() + + user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin-user", + ) + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token", + new_callable=AsyncMock, + return_value="sk-newtoken1234ab12", + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key", + new_callable=AsyncMock, + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new_callable=AsyncMock, + ) as mock_delete_cache, patch( + "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook", + new_callable=AsyncMock, + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints.prepare_key_update_data", + new_callable=AsyncMock, + return_value={}, + ): + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=existing_key, + hashed_api_key=token_hash, + key=token_hash, + data=None, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=None, + user_api_key_cache=mock_user_api_key_cache, + proxy_logging_obj=mock_proxy_logging_obj, + ) + + mock_delete_cache.assert_called_once() + call_kwargs = mock_delete_cache.call_args.kwargs + # The token hash should be passed as-is, NOT double-hashed + assert call_kwargs["hashed_token"] == token_hash diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 0258eaabe33..446316a02dc 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -93,6 +93,23 @@ def test_baseten_model_api_pricing_entries(): assert model_info["output_cost_per_token"] == output_cost +def test_wandb_model_api_pricing_entries(): + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + expected_pricing = { + "wandb/moonshotai/Kimi-K2.5": (6e-07, 3e-06), + "wandb/MiniMaxAI/MiniMax-M2.5": (3e-07, 1.2e-06), + } + + for model_name, (input_cost, output_cost) in expected_pricing.items(): + model_info = litellm.model_cost.get(model_name) + assert model_info is not None, f"Missing model pricing entry: {model_name}" + assert model_info["litellm_provider"] == "wandb" + assert model_info["input_cost_per_token"] == input_cost + assert model_info["output_cost_per_token"] == output_cost + + def test_cost_calculator_with_usage(monkeypatch): os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="")