diff --git a/litellm/caching/in_memory_cache.py b/litellm/caching/in_memory_cache.py index 7edf4672963..284e780cc6f 100644 --- a/litellm/caching/in_memory_cache.py +++ b/litellm/caching/in_memory_cache.py @@ -163,7 +163,12 @@ class InMemoryCache(BaseCache): return self.cache_dict[key] = value - if self.allow_ttl_override(key): # if ttl is not set, set it to default ttl + # refresh_ttl bypasses allow_ttl_override's "leave a still-live ttl + # alone" guard -- a caller only sets it for a counter whose ttl must + # keep extending on every write (e.g. a concurrency reservation's + # crash-safety-net ttl), never for one that must stay fixed to its + # original epoch window (e.g. a fixed-period rate-limit bucket). + if kwargs.get("refresh_ttl") or self.allow_ttl_override(key): # if ttl is not set, set it to default ttl if "ttl" in kwargs and kwargs["ttl"] is not None: self.ttl_dict[key] = time.time() + float(kwargs["ttl"]) heapq.heappush(self.expiration_heap, (self.ttl_dict[key], key)) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index f292caf8cc3..1a7cb3f12f7 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -447,7 +447,20 @@ def _build_limits_index(model_list: Sequence[Mapping[str, object]]) -> _LimitsIn sorted_by_model_name: Final = sorted(model_list, key=lambda deployment: deployment["model_name"]) by_model_name: Final[Mapping[str, tuple[_ConfiguredLimit, ...]]] = MappingProxyType( { - model_name: configured + model_name: ( + # `Router.should_include_deployment` lets a same-team caller + # reach a team-owned deployment by its own internal + # model_name, not only its team_public_model_name alias + # (litellm auto-generates a name unique per (team_id, uuid), + # so every deployment in this group shares one team_id when + # any does) -- stamping the identical team_scope here as the + # alias entry below gets keeps both paths resolving to the + # same bucket, so a caller can't split its usage across two + # independent counters just by alternating which name it calls. + tuple(replace(limit, team_scope=team_scope) for limit in configured) + if (team_scope := next((key[0] for dep in group if (key := _team_alias_key(dep))), None)) is not None + else configured + ) for model_name, deployment_group in groupby( sorted_by_model_name, key=lambda deployment: deployment["model_name"] ) @@ -963,7 +976,9 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] if current + increment > limit: return False, current new_value: Final = current + increment - await cache.async_set_cache(key=key, value=new_value, ttl=ttl, litellm_parent_otel_span=None) + await cache.async_set_cache( + key=key, value=new_value, ttl=ttl, refresh_ttl=refresh_ttl, litellm_parent_otel_span=None + ) return True, new_value async def _decrement_floor_zero(self, cache: InternalUsageCache, key: str, delta: float) -> None: diff --git a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py index 4a0f46e72aa..ac24039758e 100644 --- a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py @@ -3383,6 +3383,32 @@ async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_co await redis_cache.async_delete_cache(key=key) +@pytest.mark.asyncio +async def test_in_memory_concurrency_ttl_refreshes_on_every_admission(time_controller): + """ + Bugbot finding: the Redis path's refresh_ttl fix above was never mirrored + onto the in-memory fallback, which called async_set_cache unconditionally + -- InMemoryCache.allow_ttl_override leaves a still-live ttl untouched, so + a concurrency counter's expiry stayed fixed from its first admission even + with refresh_ttl=True, the same silent-past-the-cap failure mode the + Redis fix closed. + """ + limiter = _make_limiter(time_controller) + cache = limiter.internal_usage_cache + in_memory_cache = cache.dual_cache.in_memory_cache + key = f"tag_rl:test:in-memory-ttl-refresh:{uuid.uuid4().hex}" + + admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_first_admission = in_memory_cache.ttl_dict[key] + + admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True) + assert admitted + ttl_after_second_admission = in_memory_cache.ttl_dict[key] + + assert ttl_after_second_admission > ttl_after_first_admission + + # --------------------------------------------------------------------------- # team_public_model_name alias -- index lookup must not miss # --------------------------------------------------------------------------- @@ -3396,6 +3422,12 @@ def test_build_limits_index_is_also_keyed_by_team_public_model_name(): model_group_alias). The index must resolve either name to the same configured limits, or a team-aliased chain's limits are silently never checked. + + Security regression: Router.should_include_deployment also lets a + same-team (or team-unconstrained) caller reach this deployment by its + own internal model_name, not only the alias. Both paths must resolve to + the identical team_scope, or a caller could split its usage across two + independent buckets just by alternating which name it calls with. """ deployment = _deployment( "real-model-name", @@ -3409,10 +3441,7 @@ def test_build_limits_index_is_also_keyed_by_team_public_model_name(): by_alias = index.resolve("team-alias-name", team_id="team-1") assert by_name != () assert [c.entry for c in by_name] == [c.entry for c in by_alias] - # The alias resolution must carry the team_id into the bucket scope -- - # see test_build_limits_index_keeps_different_teams_same_alias_separate - # for why (two teams can publish the identical alias string). - assert by_name[0].team_scope is None + assert by_name[0].team_scope == "team-1" assert by_alias[0].team_scope == "team-1"