fix(proxy): declare the file usage counters as rate limit calls

Main's redis producer gate requires every module that writes a shared cache to name its key family, and the file usage counters wrote theirs without one.
This commit is contained in:
mateo-berri 2026-10-03 12:48:34 -07:00
parent 0e61b32516
commit c74eb7d1c9
2 changed files with 21 additions and 0 deletions

View file

@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, NoReturn, TypeAlias
from pydantic import Field, TypeAdapter, ValidationError
from typing_extensions import assert_never
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
EMPTY_MAPPING,
@ -148,6 +149,7 @@ async def _increment_all_or_none(
return exceeded
@with_service_target("rate_limits")
async def consume_file_usage(
cache: "InternalUsageCache",
limits: tuple[ScopedFileUsageLimit, ...],

View file

@ -2,6 +2,7 @@ from typing import Final
import pytest
from litellm._internal_context import current_service_target
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
@ -141,6 +142,24 @@ async def test_a_rejection_does_not_use_up_the_scope_that_rejected_it():
assert await consume_file_usage(cache, raised_limit, DAY, "", MIDDAY) is None
class _ServiceTargetSeen(Exception):
pass
class _TargetReportingCache(InternalUsageCache):
async def async_increment_cache(self, key, value, litellm_parent_otel_span, local_only=False, **kwargs):
raise _ServiceTargetSeen(current_service_target())
async def test_counter_writes_are_declared_as_rate_limit_calls_so_their_redis_spans_are_named():
cache: Final = _TargetReportingCache(dual_cache=DualCache())
with pytest.raises(_ServiceTargetSeen) as seen:
await consume_file_usage(cache, (_scoped("key", "hashed", 1),), DAY, "", MIDDAY)
assert seen.value.args == ("rate_limits",)
async def test_download_counters_are_per_file():
cache: Final = _cache()
limits: Final = (_scoped("key", "hashed", 1, setting="max_file_downloads_per_minute"),)