diff --git a/litellm/router_strategy/complexity_router/cache_warming/capture.py b/litellm/router_strategy/complexity_router/cache_warming/capture.py index 94840e3a135..66ce0d5493c 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/capture.py +++ b/litellm/router_strategy/complexity_router/cache_warming/capture.py @@ -101,18 +101,6 @@ def _capture_allowed(kwargs: Mapping[str, object]) -> bool: return _should_store_prompts_and_responses_in_spend_logs() -def _caller_tags(metadata_dicts: Sequence[Mapping[str, object]]) -> tuple[str, ...]: - """The caller's own request tags, kept so a replay presents the same body the gates read them from. - Tag budgets, tag budget reservation and the limiter's tag descriptors all resolve tags through - get_tags_from_request_body, so a replay that drops them is not refused by those ceilings, it is simply - invisible to them.""" - for metadata in metadata_dicts: - tags = metadata.get("tags") - if isinstance(tags, list): - return tuple(tag for tag in tags if isinstance(tag, str)) - return () - - def _is_replay(metadata_dicts: Sequence[Mapping[str, object]]) -> bool: for metadata in metadata_dicts: if metadata.get(CACHE_WARMING_REPLAY_MARKER_KEY): @@ -254,7 +242,6 @@ async def capture_session( await store.upsert_session( caller_scope=caller_scope, session_id=session_id, - tags=_caller_tags(metadata_dicts), payload_compressed=blob, payload_sha256=sha, token_estimate=token_estimate, diff --git a/litellm/router_strategy/complexity_router/cache_warming/refresher.py b/litellm/router_strategy/complexity_router/cache_warming/refresher.py index 8832a247b69..cc0c107df72 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/refresher.py +++ b/litellm/router_strategy/complexity_router/cache_warming/refresher.py @@ -152,7 +152,6 @@ def _replay_body( data[get_metadata_variable_name_from_kwargs(data)] = { # mutable-ok: request metadata, never retained CACHE_WARMING_REPLAY_MARKER_KEY: True, **({"session_id": record.session_id} if record.session_id is not None else {}), - **({"tags": list(record.tags)} if record.tags else {}), "spend_logs_metadata": {CACHE_WARMING_REPLAY_TAG: "true"}, # mutable-ok: request metadata, never retained } return data @@ -290,6 +289,32 @@ def _tenancy_no_longer_holds( return "reassigned" if captured != current else None +def _apply_tenant_limits( + principal: "UserAPIKeyAuth", user: "LiteLLM_UserTable | None", end_user: "LiteLLM_EndUserTable | None" +) -> None: + """The v3 limiter builds its tenant descriptors from limit fields on the principal, so a field the + principal lacks is a ceiling that silently does not apply rather than one that refuses. + + Most of them need nothing here: LiteLLM_VerificationTokenView joins the team and organization onto the + key row, so get_key_object already returns team, team-member and organization limits and those descriptors + are built on a replay unchanged. User and end-user limits are the two the request path adds during auth + from objects it loads, so warming adds them from the objects it loads for the same gates, through the same + owner where there is one.""" + from litellm.proxy.auth.user_api_key_auth import ( + _apply_budget_limits_to_end_user_params, # pyright: ignore[reportPrivateUsage] # the owner of this mapping; no public form exists + ) + + if user is not None: + principal.user_tpm_limit = user.tpm_limit + principal.user_rpm_limit = user.rpm_limit + budget = getattr(end_user, "litellm_budget_table", None) + if budget is not None: + end_user_params: dict[str, object] = {} # mutable-ok: the owner's out-parameter shape + _apply_budget_limits_to_end_user_params(end_user_params, budget, principal.end_user_id) + for field, value in end_user_params.items(): + setattr(principal, field, value) + + def _excluded_from_warming(key_state: "UserAPIKeyAuth", now: "datetime", proxy_logging_obj: "ProxyLogging") -> bool: """Mirrors the canonical auth checks (user_api_key_auth.py:1717 and :2891) because common_checks owns that policy for real traffic but needs a FastAPI Request. datetime.fromisoformat only accepts a @@ -871,6 +896,7 @@ class CacheWarmingRefresher: skip_budget_checks = _should_skip_budget_checks( request_data=data, route=route, request=None, llm_router=llm_router ) + _apply_tenant_limits(principal, user, end_user) await _authorize_replay( principal=principal, data=data, diff --git a/litellm/router_strategy/complexity_router/cache_warming/store.py b/litellm/router_strategy/complexity_router/cache_warming/store.py index 9a836a3e0b4..17efa5dd78b 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/store.py +++ b/litellm/router_strategy/complexity_router/cache_warming/store.py @@ -189,7 +189,6 @@ class CacheWarmingStore: payload_sha256: str, token_estimate: int, served_model: str, - tags: tuple[str, ...], attribution: CacheWarmingAttribution, ttl_seconds: int, max_sessions: int, @@ -207,7 +206,6 @@ class CacheWarmingStore: last_activity=now, served_model=served_model, session_id=session_id, - tags=tags, attribution=attribution, auto_router_model_name=self.auto_router_model_name, ) diff --git a/litellm/router_strategy/complexity_router/cache_warming/types.py b/litellm/router_strategy/complexity_router/cache_warming/types.py index 6dbd9bf8ddf..62966976f8b 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/types.py +++ b/litellm/router_strategy/complexity_router/cache_warming/types.py @@ -80,7 +80,6 @@ class CacheWarmingRecord(BaseModel): last_activity: float served_model: str session_id: str | None = None - tags: tuple[str, ...] = () attribution: CacheWarmingAttribution auto_router_model_name: str diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py index af9473cc067..40299b1f711 100644 --- a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_refresher.py @@ -521,18 +521,42 @@ async def test_a_session_is_dropped_once_its_key_leaves_the_tenant_it_was_captur @pytest.mark.asyncio -async def test_a_replay_presents_the_callers_own_tags_to_the_gates_that_read_them(): - """Tag budgets, tag budget reservation and the limiter's tag descriptors all resolve tags through one - owner, get_tags_from_request_body, which reads them off the request body. Warming already enters through - all three, so a replay that drops the caller's tags is not refused by those ceilings, it is invisible to - them, and its spend lands outside the tag it belongs to. The marker stays out of this channel because - tags feed deployment selection.""" +async def test_a_replay_carries_the_key_and_team_tags_its_owner_re_derives_and_no_client_tags(): + """Tags have two owners and warming is neither. Key and team tags are re-derived on every replay by + add_key_team_project_metadata, so capturing them would snapshot a value that is recomputed anyway; and + re-presenting captured tags as request metadata turns proxy-derived tags into client-supplied ones, which + _reject_clientside_metadata_tags_check refuses outright when the operator forbids them. So a replay + carries exactly what the owner puts there and nothing warming remembered.""" from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body llm_router, redis = warming_rig(redis=FakeRedisCache()) - seed_session(redis, tags=("cost-center-7",), touched=_VISITED_BOTH_TIERS) - await tick(llm_router) + seed_session(redis, user_api_key="k", touched=_VISITED_BOTH_TIERS) + keys = FakeKeyDirectory({"k": key_state(token="k", metadata={"tags": ["cost-center-7"]})}) + await tick(llm_router, active=refresher(keys=keys)) assert llm_router.completion_calls, "expected a replay" for call in llm_router.completion_calls: assert get_tags_from_request_body(call) == ["cost-center-7"] assert CACHE_WARMING_REPLAY_TAG not in get_tags_from_request_body(call) + + +@pytest.mark.asyncio +async def test_a_users_own_rate_limit_binds_on_a_replay(): + """The v3 limiter builds tenant descriptors from limit fields on the principal, so one the principal lacks + is a ceiling that silently does not apply. Team, team-member and organization limits ride the key row + through LiteLLM_VerificationTokenView, but user limits are added during auth from the user object, so + warming adds them from the object it loads for the same gates.""" + from litellm.proxy._types import LiteLLM_UserTable + + limiter, counters = real_limiter() + key_cache = DualCache() + llm_router, redis = warming_rig(redis=FakeRedisCache()) + await key_cache.async_set_cache( + key="u", value=LiteLLM_UserTable(user_id="u", max_budget=None, spend=0.0, rpm_limit=1) + ) + await counters.async_set_cache(key="{user:u}:window", value=str(int(time.time()))) + await counters.async_set_cache(key="{user:u}:requests", value=99) + seed_session(redis, user_api_key="k", user_id="u", touched=_VISITED_BOTH_TIERS) + keys = FakeKeyDirectory({"k": key_state(token="k", user_id="u")}) + with registered_callbacks(limiter): + await tick(llm_router, active=refresher(keys=keys, limiter=limiter), user_api_key_cache=key_cache) + assert llm_router.completion_calls == [], "a user at their RPM limit must not be warmed" diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py index 35ddbddff60..d5f130bdd05 100644 --- a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_store.py @@ -109,7 +109,6 @@ def _record_json(**overrides: object) -> str: token_estimate=2048, last_activity=1000.0, served_model="sonnet", - tags=(), attribution=CacheWarmingAttribution(user_api_key="hashed"), auto_router_model_name="smart-router", ).model_dump() @@ -128,7 +127,6 @@ async def _upsert(store: CacheWarmingStore, session_id: str = "s1", max_sessions payload_sha256="sha2", token_estimate=4096, served_model="sonnet", - tags=(), attribution=CacheWarmingAttribution(), ttl_seconds=1800, max_sessions=max_sessions, diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py index 15de55cbd6a..9d4050073e7 100644 --- a/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/warming_rig.py @@ -221,7 +221,6 @@ def seed_session( org_id: str | None = None, project_id: str | None = None, touched: tuple[str, ...] | None = None, - tags: tuple[str, ...] = (), ) -> str: payload = CacheWarmingPayload( model=served_model, @@ -240,7 +239,6 @@ def seed_session( last_activity=last_activity if last_activity is not None else time.time(), served_model=served_model, session_id=session_id, - tags=tags, attribution=CacheWarmingAttribution( user_api_key=user_api_key, user_api_key_team_id=team_id,