diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d56250a5504..f990f043dea 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -65,7 +65,6 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( evict_and_broadcast, - publish_auth_cache_invalidation, ) from litellm.proxy.common_utils.callback_config_validation import logging_metadata_config_error from litellm.proxy.common_utils.callback_utils import ( @@ -171,7 +170,6 @@ from litellm.proxy.utils import ( require_enterprise_license, ) from litellm.redis.stores.spend_counters import ( - SPEND_COUNTERS_TARGET, key_counter, key_window_counter, ) @@ -6809,38 +6807,6 @@ def _validate_reset_spend_value(reset_to: object, key_in_db: LiteLLM_Verificatio return reset_to -@with_service_target(SPEND_COUNTERS_TARGET) -async def _set_spend_counter_with_floor_and_broadcast(counter_key: str, value: float) -> None: - """ - Set a Redis-backed spend counter to `value`, mirror it into the short-lived - spend_db_floor marker `_authoritative_floor_spend` reads, and broadcast both - to every worker (LIT-3803 pattern: setting, not deleting, means a worker's - own self-delivered broadcast still carries the reset value forward). - - Without the floor marker, `_authoritative_floor_spend` can re-derive a - stale, pre-reset value from a marker another worker cached moments earlier - and raise the just-reset counter right back up via `_repair_stale_spend_counter`. - Without the broadcast, a worker that already cached the pre-reset key object - or floor marker keeps enforcing against it until its own TTL expires. - """ - from litellm.proxy.proxy_server import SPEND_DB_FLOOR_CACHE_TTL_SECONDS, spend_counter_cache, spend_counter_store - - try: - await spend_counter_store.set(counter_key, value, ttl=60) - except Exception as redis_err: - verbose_proxy_logger.warning( - "Failed to update spend counter %s in Redis: %s. Budget checks may use stale value until counter expires.", - counter_key, - redis_err, - ) - - floor_key: Final = f"spend_db_floor:{counter_key}" - spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) - - await publish_auth_cache_invalidation(cache_key=counter_key, new_value=value, ttl=60) - await publish_auth_cache_invalidation(cache_key=floor_key, new_value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) - - def _budget_limit_windows(budget_limits: Sequence[object] | str | None) -> tuple[Mapping[str, object], ...]: """Coerce a key's stored `budget_limits` into a tuple of plain window dicts. @@ -6925,6 +6891,8 @@ async def _reset_key_budget_windows( pre-reset boundary, re-sum the unchanged historical spend, and put the counter right back where it was before the write ever landed. """ + from litellm.proxy.proxy_server import set_spend_counter_with_floor_and_broadcast + windows: Final = _budget_limit_windows(budget_limits) if not windows: return @@ -6943,7 +6911,7 @@ async def _reset_key_budget_windows( duration = window.get("budget_duration") if isinstance(duration, str) and duration: counter_key = key_window_counter(hashed_api_key, duration) - await _set_spend_counter_with_floor_and_broadcast(counter_key=counter_key, value=0.0) + await set_spend_counter_with_floor_and_broadcast(counter_key=counter_key, value=0.0) @router.post( @@ -6966,6 +6934,7 @@ async def reset_key_spend_fn( hash_token, prisma_client, proxy_logging_obj, + set_spend_counter_with_floor_and_broadcast, user_api_key_cache, ) @@ -7017,7 +6986,7 @@ async def reset_key_spend_fn( # amount for every enforcement check immediately instead of the stale # pre-reset value. _counter_key: Final = key_counter(hashed_api_key) - await _set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to) + await set_spend_counter_with_floor_and_broadcast(counter_key=_counter_key, value=reset_to) await _reset_key_budget_windows( prisma_client=prisma_client, hashed_api_key=hashed_api_key, diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index 3ec92165a84..0b7a7ca43cd 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -30,6 +30,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, get_daily_activity, ) +from litellm.proxy.management_endpoints.common_utils import validate_finite_spend from litellm.proxy.management_helpers.utils import handle_budget_for_entity from litellm.redis.stores.spend_counters import ( tag_counter, @@ -415,6 +416,7 @@ async def update_tag( Parameters: - name: str - The name of the tag to update + - spend: Optional[float] - Updated spend for the tag - description: Optional[str] - Updated description - models: List[str] - Updated list of allowed LLM models - budget_id: Optional[str] - The id for a budget to associate with the tag @@ -428,11 +430,13 @@ async def update_tag( - model_max_budget: Optional[dict] - Max budget for a specific model - budget_duration: Optional[str] - Frequency of resetting tag budget """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, set_spend_counter_with_floor_and_broadcast if prisma_client is None: raise HTTPException(status_code=500, detail="Database not connected") + validate_finite_spend(tag.spend) + try: # Check if tag exists existing_tag: Final = await _table(TagRepository(prisma_client)).find_unique(where={"tag_name": tag.name}) @@ -459,12 +463,10 @@ async def update_tag( "description": tag.description, "models": tag.models or [], "model_info": json.dumps(model_info), + **({"spend": tag.spend} if tag.spend is not None else {}), + **({"budget_id": budget_id} if budget_id != existing_tag.budget_id else {}), } - # Add budget_id if it changed - if budget_id != existing_tag.budget_id: - update_data["budget_id"] = budget_id - # Update tag in database updated_tag_record: Final = await _table(TagRepository(prisma_client)).update( where={"tag_name": tag.name}, @@ -473,6 +475,9 @@ async def update_tag( await _evict_tag_cache_keys((tag_cache_key(tag.name),)) + if tag.spend is not None: + await set_spend_counter_with_floor_and_broadcast(counter_key=f"spend:tag:{tag.name}", value=tag.spend) + # Build response tag_config: Final = TagConfig( name=updated_tag_record.tag_name, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7f48d77ff09..781455c6815 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -517,6 +517,7 @@ from litellm.proxy.common_utils.anthropic_error_payload import ( ) from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, + publish_auth_cache_invalidation, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy from litellm.proxy.common_utils.codex_model_catalog import codex_model_list_body @@ -727,6 +728,7 @@ from litellm.proxy.litellm_pre_call_utils import ( SCHEDULER_PRIORITY_HEADER, LiteLLMProxyRequestSetup, add_litellm_data_to_request, + sanitize_for_log, ) from litellm.proxy.logging_endpoints.callback_logs_endpoints import ( rust_control_plane_router, @@ -4425,6 +4427,36 @@ async def invalidate_spend_counter(counter_key: str) -> None: _invalidate_spend_counter: Final = invalidate_spend_counter +@with_service_target(SPEND_COUNTERS_TARGET) +async def set_spend_counter_with_floor_and_broadcast(counter_key: str, value: float) -> None: + """ + Set a Redis-backed spend counter to `value`, mirror it into the short-lived + spend_db_floor marker `_authoritative_floor_spend` reads, and broadcast both + to every worker (LIT-3803 pattern: setting, not deleting, means a worker's + own self-delivered broadcast still carries the reset value forward). + + Without the floor marker, `_authoritative_floor_spend` can re-derive a + stale, pre-reset value from a marker another worker cached moments earlier + and raise the just-reset counter right back up via `_repair_stale_spend_counter`. + Without the broadcast, a worker that already cached the pre-reset key object + or floor marker keeps enforcing against it until its own TTL expires. + """ + try: + await spend_counter_store.set(counter_key, value, ttl=60) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to update spend counter %s in Redis: %s. Budget checks may use stale value until counter expires.", + sanitize_for_log(counter_key), + redis_err, + ) + + floor_key: Final = f"spend_db_floor:{counter_key}" + spend_counter_cache.in_memory_cache.set_cache(key=floor_key, value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) + + await publish_auth_cache_invalidation(cache_key=counter_key, new_value=value, ttl=60) + await publish_auth_cache_invalidation(cache_key=floor_key, new_value=value, ttl=SPEND_DB_FLOOR_CACHE_TTL_SECONDS) + + async def _apply_spend_counter_increments(pending: Sequence[PendingSpendIncrement]) -> None: if _defer_spend_counter_increments(pending): return diff --git a/litellm/types/tag_management.py b/litellm/types/tag_management.py index e59b02af71a..2958c07ce95 100644 --- a/litellm/types/tag_management.py +++ b/litellm/types/tag_management.py @@ -27,6 +27,7 @@ class TagNewRequest(TagBase): class TagUpdateRequest(TagBase): + spend: float | None = None budget_id: str | None = None # Budget fields - if provided, will update the budget max_budget: float | None = None diff --git a/tests/integration/spend/test_tag_spend_reset.py b/tests/integration/spend/test_tag_spend_reset.py new file mode 100644 index 00000000000..2104539f5d7 --- /dev/null +++ b/tests/integration/spend/test_tag_spend_reset.py @@ -0,0 +1,75 @@ +import json +import uuid +from typing import Final + +import httpx + +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.database import read_rows +from tests.integration._support.provider import PROVIDER_URL, SharedProvider +from tests.integration._support.wire import Reply + + +def test_admin_can_reset_tag_spend_and_resume_budget_blocked_requests( + gateway: Gateway, provider: SharedProvider +) -> None: + tag: Final = f"tag-reset-{uuid.uuid4().hex}" + completion: Final = Reply( + body=json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + with gateway.scenario() as scenario: + model: Final = scenario.model( + api_base=f"{PROVIDER_URL}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002 + ) + gateway.post("/tag/new", {"name": tag, "max_budget": 0.0001, "budget_duration": "30d"}) + + def delete_tag() -> None: + gateway.post("/tag/delete", {"name": tag}) + + scenario.cleanups.callback(delete_tag) + + def tagged_request() -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"tag reset {uuid.uuid4().hex}"}], + "metadata": {"tags": [tag]}, + }, + ) + + provider.expect(completion) + first: Final = tagged_request() + assert first.status_code == 200, first.text + assert len(provider.received()) == 1 + eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_TagTable" WHERE tag_name=%s', (tag,)), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0.0001, + seconds=70, + ) + blocked: Final = tagged_request() + assert blocked.status_code == 422, blocked.text + assert blocked.json()["error"]["type"] == "budget_exceeded", blocked.text + assert f"Tag={tag}" in blocked.text, blocked.text + assert provider.received() == () + reset: Final = gateway.request("POST", "/tag/update", {"name": tag, "spend": 0}) + assert reset.status_code == 200, reset.text + provider.expect(completion) + try: + resumed: Final = tagged_request() + assert resumed.status_code == 200, resumed.text + assert len(provider.received()) == 1 + finally: + provider.replies.clear() diff --git a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py index 531bf98b60e..813875bce7c 100644 --- a/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_tag_management_endpoints.py @@ -16,6 +16,7 @@ from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.proxy_server import app from litellm.proxy.utils import PrismaClient from litellm.types.tag_management import TagNewRequest +from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters client = TestClient(app) @@ -230,6 +231,138 @@ async def test_update_tag(): app.dependency_overrides.clear() +@pytest.mark.asyncio +@pytest.mark.parametrize("spend_fields", [{"spend": 0}, {}, {"spend": None}]) +async def test_update_tag_changes_spend_and_counter_only_when_provided( + spend_fields: Mapping[str, float | None], +) -> None: + from litellm.proxy.management_endpoints.tag_management_endpoints import update_tag + from litellm.types.tag_management import TagUpdateRequest + + updated_tag: Final = SimpleNamespace( + tag_name="reset-tag", + description=None, + models=[], + budget_id=None, + created_at=datetime(2024, 1, 1), + updated_at=datetime(2024, 1, 1), + created_by="admin", + ) + update_row: Final = AsyncMock(return_value=updated_tag) + spend: Final = FakeSpendCounters(redis=None) + mock_prisma: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_tagtable=SimpleNamespace( + find_unique=AsyncMock(return_value=updated_tag), update=update_row + ), + litellm_proxymodeltable=SimpleNamespace(find_many=AsyncMock(return_value=[])), + ) + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + _tag_cache_doubles(), + patch("litellm.proxy.proxy_server.invalidate_spend_counter", new_callable=AsyncMock) as invalidate, + patch("litellm.proxy.proxy_server.spend_counter_cache") as counter_cache, + patch("litellm.proxy.proxy_server.spend_counter_store", spend.store()), + ): + counter_cache.redis_cache = None + response: Final = await update_tag( + tag=TagUpdateRequest.model_validate({"name": "reset-tag", **spend_fields}), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert response["message"] == "Tag reset-tag updated successfully" + update_row.assert_awaited_once_with( + where={"tag_name": "reset-tag"}, + data={ + "description": None, + "models": [], + "model_info": "{}", + **({"spend": spend_fields["spend"]} if spend_fields.get("spend") is not None else {}), + }, + ) + if spend_fields.get("spend") is not None: + assert spend.memory["spend:tag:reset-tag"] == 0 + else: + assert "spend:tag:reset-tag" not in spend.memory + counter_cache.in_memory_cache.set_cache.assert_not_called() + invalidate.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("spend", [float("nan"), float("inf"), float("-inf")]) +async def test_update_tag_rejects_non_finite_spend_before_writing(spend: float) -> None: + from litellm.proxy.management_endpoints.tag_management_endpoints import update_tag + from litellm.types.tag_management import TagUpdateRequest + + update_row: Final = AsyncMock() + mock_prisma: Final = SimpleNamespace( + db=SimpleNamespace(litellm_tagtable=SimpleNamespace(update=update_row)) + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch("litellm.proxy.proxy_server.invalidate_spend_counter", new_callable=AsyncMock) as invalidate, + pytest.raises(HTTPException) as error, + ): + await update_tag( + tag=TagUpdateRequest(name="reset-tag", spend=spend), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert error.value.status_code == 400 + assert "spend must be a finite number" in str(error.value.detail) + update_row.assert_not_awaited() + invalidate.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("redis_state", ["available", "absent", "error"]) +async def test_update_tag_overwrites_stale_counter_repopulated_during_tag_cache_eviction(redis_state: str) -> None: + from litellm.proxy.management_endpoints.tag_management_endpoints import update_tag + from litellm.types.tag_management import TagUpdateRequest + + counter_key: Final = "spend:tag:overlap-reset-tag" + spend: Final = FakeSpendCounters(redis=None if redis_state == "absent" else {}, fail=redis_state == "error") + updated_tag: Final = SimpleNamespace( + tag_name="overlap-reset-tag", + description=None, + models=[], + budget_id=None, + created_at=datetime(2024, 1, 1), + updated_at=datetime(2024, 1, 1), + created_by="admin", + ) + mock_prisma: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_tagtable=SimpleNamespace( + find_unique=AsyncMock(return_value=updated_tag), update=AsyncMock(return_value=updated_tag) + ), + litellm_proxymodeltable=SimpleNamespace(find_many=AsyncMock(return_value=[])), + ) + ) + + async def overlapping_completion(cache_keys: Sequence[str]) -> None: + spend.set_local(counter_key, 1.5, ttl_seconds=60) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch( + "litellm.proxy.proxy_server.spend_counter_store", spend.store(), + ), + patch( + "litellm.proxy.management_endpoints.tag_management_endpoints._evict_tag_cache_keys", + overlapping_completion, + ), + ): + await update_tag( + tag=TagUpdateRequest(name="overlap-reset-tag", spend=0), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert spend.memory[counter_key] == 0 + if redis_state == "available": + assert spend.redis is not None + assert spend.redis[counter_key] == 0 + + @pytest.mark.asyncio async def test_new_tag_persists_a_budget(): from datetime import datetime diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2272b4bea43..c351e3b1b90 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -17159,6 +17159,7 @@ export interface paths { * * Parameters: * - name: str - The name of the tag to update + * - spend: Optional[float] - Updated spend for the tag * - description: Optional[str] - Updated description * - models: List[str] - Updated list of allowed LLM models * - budget_id: Optional[str] - The id for a budget to associate with the tag @@ -47982,6 +47983,8 @@ export interface components { rpm_limit?: number | null; /** Soft Budget */ soft_budget?: number | null; + /** Spend */ + spend?: number | null; /** Tpm Limit */ tpm_limit?: number | null; };