From f54e4e664bd4c0bdf2677a162da616db67eb8ea7 Mon Sep 17 00:00:00 2001 From: Darien Kindlund Date: Sat, 11 Apr 2026 22:36:40 -0400 Subject: [PATCH 1/9] fix(proxy): use _hash_token_if_needed for cache invalidation in bulk update and key rotation (#25552) Two code paths in key_management_endpoints.py call hash_token() unconditionally when invalidating the user_api_key_cache after a key update. When the caller passes a pre-hashed token ID (not an sk- prefixed key), hash_token() double-hashes it, producing a cache key that does not match the actual cached entry. Cache invalidation silently fails. This is compounded by update_cache() which writes the stale cached key object back with a fresh 60s TTL after every successful request, preventing natural TTL expiry. The stale entry (with outdated fields like max_budget=None) persists indefinitely under load. PR #24969 fixed this in update_key_fn but missed two other call sites: - _process_single_key_update (bulk update path) - _execute_virtual_key_regeneration (key rotation path) Fix: replace hash_token() with _hash_token_if_needed() in both locations, matching the pattern already used elsewhere in the file. Co-authored-by: Claude Opus 4.6 (1M context) --- .../key_management_endpoints.py | 4 +- .../test_key_management_endpoints.py | 164 ++++++++++++++++++ 2 files changed, 166 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 6e8a691ce93..24396fbd18e 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, ) @@ -3721,7 +3721,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/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..a16bc078cf3 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 @@ -8746,3 +8746,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 From 17e145a083de6e9d94da6e2c7d91a19d37f51264 Mon Sep 17 00:00:00 2001 From: Darien Kindlund Date: Sat, 11 Apr 2026 22:37:58 -0400 Subject: [PATCH 2/9] fix(proxy): use model_group for model_max_budget spend tracking cache key (#25549) 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 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/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, From e3d160f158ac33348699f4196e09d19401bbcc41 Mon Sep 17 00:00:00 2001 From: Utsab Dahal <250059@softwarica.edu.np> Date: Sun, 12 Apr 2026 08:24:19 +0545 Subject: [PATCH 3/9] fix(embedding): omit null encoding_format for openai requests (#25395) --- litellm/main.py | 3 -- .../test_openai_embeddings_encoding_format.py | 34 +++++++++++++++++++ 2 files changed, 34 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/llms/openai/embeddings/test_openai_embeddings_encoding_format.py diff --git a/litellm/main.py b/litellm/main.py index ddd37b47536..cf360855d11 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4913,9 +4913,6 @@ def embedding( # noqa: PLR0915 if encoding_format is not None: optional_params["encoding_format"] = encoding_format - else: - # Omiting causes openai sdk to add default value of "float" - optional_params["encoding_format"] = None api_version = None diff --git a/tests/test_litellm/llms/openai/embeddings/test_openai_embeddings_encoding_format.py b/tests/test_litellm/llms/openai/embeddings/test_openai_embeddings_encoding_format.py new file mode 100644 index 00000000000..617037865c0 --- /dev/null +++ b/tests/test_litellm/llms/openai/embeddings/test_openai_embeddings_encoding_format.py @@ -0,0 +1,34 @@ +from unittest.mock import patch + +import litellm + + +@patch("litellm.main.openai_chat_completions.embedding", return_value={"ok": True}) +def test_openai_embedding_does_not_send_encoding_format_when_unset(mock_embedding): + """Regression test: do not send encoding_format=null to OpenAI-compatible APIs.""" + litellm.embedding( + model="text-embedding-3-small", + input=["hello"], + api_base="https://example.com/v1", + api_key="test-key", + custom_llm_provider="openai", + ) + + optional_params = mock_embedding.call_args.kwargs["optional_params"] + assert "encoding_format" not in optional_params + + +@patch("litellm.main.openai_chat_completions.embedding", return_value={"ok": True}) +def test_openai_embedding_preserves_explicit_encoding_format(mock_embedding): + """Explicit encoding_format should still be forwarded.""" + litellm.embedding( + model="text-embedding-3-small", + input=["hello"], + api_base="https://example.com/v1", + api_key="test-key", + custom_llm_provider="openai", + encoding_format="float", + ) + + optional_params = mock_embedding.call_args.kwargs["optional_params"] + assert optional_params["encoding_format"] == "float" From e1bf1145919136ab1d6b6159f8ab735bab13dc0d Mon Sep 17 00:00:00 2001 From: Dmitriy Alergant <93501479+DmitriyAlergant@users.noreply.github.com> Date: Sat, 11 Apr 2026 22:45:23 -0400 Subject: [PATCH 4/9] fix(budget): align budget table reset times with standardized calendar schedule (#25440) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Budget table entries (team members, end-users) used duration_in_seconds() for a sliding-window reset, while keys/users/teams used calendar-aligned get_budget_reset_time(). This made "30d" and "1mo" mean different things depending on entity type. Now both paths use get_budget_reset_time() for consistent calendar-aligned resets (e.g. "30d" → 1st of next month). Fixes #25432 Co-authored-by: Claude Opus 4.6 (1M context) --- .../proxy/common_utils/reset_budget_job.py | 21 +-- .../budget_management_endpoints.py | 8 +- .../common_utils/test_reset_budget_job.py | 134 ++++++++++++++++++ 3 files changed, 142 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bcfaed24398..fe169f56a9d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -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/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/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 ): From ee06b9278a32617b8a66ed5cfe6abc13ea286f90 Mon Sep 17 00:00:00 2001 From: csoni-cweave Date: Sat, 11 Apr 2026 19:46:40 -0700 Subject: [PATCH 5/9] feat(model):add wandb model offerings to include kimi-k2.5 and minimax-m2.5 (#25409) --- litellm/constants.py | 3 ++ ...odel_prices_and_context_window_backup.json | 28 +++++++++++++++++++ model_prices_and_context_window.json | 28 +++++++++++++++++++ tests/test_litellm/test_cost_calculator.py | 17 +++++++++++ 4 files changed, 76 insertions(+) diff --git a/litellm/constants.py b/litellm/constants.py index 337cb1243fb..300cd8a391b 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 63ca003a26d..ae4a79c909d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -32396,6 +32396,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/model_prices_and_context_window.json b/model_prices_and_context_window.json index 90ff7d1103c..0e7e2557d1e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -32381,6 +32381,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/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="") From e6771feace5e33377cf1896206f082268331a19e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 14 Apr 2026 20:36:28 +0530 Subject: [PATCH 6/9] Revert "fix(embedding): omit null encoding_format for openai requests (#25395)" This reverts commit e3d160f158ac33348699f4196e09d19401bbcc41. --- litellm/main.py | 3 ++ .../test_openai_embeddings_encoding_format.py | 34 ------------------- 2 files changed, 3 insertions(+), 34 deletions(-) delete mode 100644 tests/test_litellm/llms/openai/embeddings/test_openai_embeddings_encoding_format.py diff --git a/litellm/main.py b/litellm/main.py index cf360855d11..ddd37b47536 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4913,6 +4913,9 @@ def embedding( # noqa: PLR0915 if encoding_format is not None: optional_params["encoding_format"] = encoding_format + else: + # Omiting causes openai sdk to add default value of "float" + optional_params["encoding_format"] = None api_version = None diff --git a/tests/test_litellm/llms/openai/embeddings/test_openai_embeddings_encoding_format.py b/tests/test_litellm/llms/openai/embeddings/test_openai_embeddings_encoding_format.py deleted file mode 100644 index 617037865c0..00000000000 --- a/tests/test_litellm/llms/openai/embeddings/test_openai_embeddings_encoding_format.py +++ /dev/null @@ -1,34 +0,0 @@ -from unittest.mock import patch - -import litellm - - -@patch("litellm.main.openai_chat_completions.embedding", return_value={"ok": True}) -def test_openai_embedding_does_not_send_encoding_format_when_unset(mock_embedding): - """Regression test: do not send encoding_format=null to OpenAI-compatible APIs.""" - litellm.embedding( - model="text-embedding-3-small", - input=["hello"], - api_base="https://example.com/v1", - api_key="test-key", - custom_llm_provider="openai", - ) - - optional_params = mock_embedding.call_args.kwargs["optional_params"] - assert "encoding_format" not in optional_params - - -@patch("litellm.main.openai_chat_completions.embedding", return_value={"ok": True}) -def test_openai_embedding_preserves_explicit_encoding_format(mock_embedding): - """Explicit encoding_format should still be forwarded.""" - litellm.embedding( - model="text-embedding-3-small", - input=["hello"], - api_base="https://example.com/v1", - api_key="test-key", - custom_llm_provider="openai", - encoding_format="float", - ) - - optional_params = mock_embedding.call_args.kwargs["optional_params"] - assert optional_params["encoding_format"] == "float" From f6e526c5bedd4b28d3aff387f364a17a7902a10b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 14 Apr 2026 20:46:21 +0530 Subject: [PATCH 7/9] Fix bulk update tests --- .../test_key_management_endpoints.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) 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 a16bc078cf3..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" From ef94f5fc4d98df48ec6678211972268503970619 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 14 Apr 2026 20:50:42 +0530 Subject: [PATCH 8/9] Fix budget reset test --- tests/test_budget_management.py | 41 ++++++++++++++++++++++----------- 1 file changed, 27 insertions(+), 14 deletions(-) diff --git a/tests/test_budget_management.py b/tests/test_budget_management.py index 5863cea9d3b..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,30 +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. - expected_reset_at_date = datetime.fromisoformat( - budget_setup["created_at"] - ) + 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 - actual_reset_at_date = datetime.fromisoformat(budget_setup["budget_reset_at"]) 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." ) From a0e61a9d495e2f0a14c4ea32bb72cefa8f57c9ba Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 14 Apr 2026 20:58:12 +0530 Subject: [PATCH 9/9] Fix code qa --- litellm/proxy/common_utils/reset_budget_job.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index fe169f56a9d..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