feat(tags): reset a tag's spend through /tag/update (#50323)

* feat(tags): reset a tag's spend through /tag/update

Co-authored-by: Arnold Galovics <arnold_galovics@docktape.com>

* test(tags): assert spend update response contract

* fix(tags): seed spend counter after evicting tag cache

* refactor(spend): expose shared counter reset helper

* fix(spend): sanitize counter key in spend counter warning log

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: Arnold Galovics <arnold_galovics@docktape.com>

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,

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,

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

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

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()

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

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;
         };