diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 164c8608922..fe560b359ea 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -1,7 +1,10 @@ # What is this? ## Helper utilities import copy +import hashlib +import json from collections.abc import Mapping +from functools import lru_cache from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union import httpx @@ -229,25 +232,43 @@ def iter_request_metadata_dicts(request_kwargs: Mapping[str, object]) -> tuple[M ) +@lru_cache(maxsize=1) +def proxy_identity_fields() -> "tuple[str, ...]": + """Every tenant dimension the proxy stamps, read off the type that defines them rather than listed here, + so a dimension litellm adds is picked up instead of silently omitted.""" + from litellm.types.utils import StandardLoggingUserAPIKeyMetadata + + return ("user_api_key_hash",) + tuple( + sorted( + field + for field in StandardLoggingUserAPIKeyMetadata.__annotations__ + if field.startswith("user_api_key_") and field.endswith("_id") + ) + ) + + def get_caller_scope(request_kwargs: Mapping[str, object]) -> str: - """A cache/affinity partition key for the authenticated caller, scoped by the strongest identity the - request actually carries. Virtual keys carry a key hash; JWT and other keyless proxy principals do not, - so they scope by the tenancy the proxy stamped instead, mirroring how the v3 rate limiter builds an - api_key descriptor only when api_key is present and separate descriptors per user, team and organization - (parallel_request_limiter_v3.py:1988). Collapsing keyless callers onto one shared literal would merge - distinct tenants that reuse a session id into a single partition. "unscoped" is returned only when the - request carries no proxy identity at all, which is direct SDK use where there is one tenant by + """A cache/affinity partition for the authenticated caller, derived from the caller's COMPLETE identity. + + Anything stored under this partition is also attributed to the caller, so the partition has to cover + every dimension the attribution does. A partition built from a subset is a collision channel: two + principals that differ only in an omitted dimension share one entry, and whoever writes last owns the + payload and the attribution while any accumulated state stays. End users are the reachable case, since + they share a virtual key by design and the key hash alone cannot tell them apart, but the rule is what + matters rather than that instance. Composing every dimension also means a tenancy change (a key removed + from a team, say) lands on a new partition instead of inheriting the old one's state. + + The dimensions come from the type that declares them, so this cannot drift from what the proxy stamps. + Values are hashed rather than concatenated: end-user and project identifiers are customer-chosen strings + that would otherwise appear in Redis key names and in every log line quoting one, and hashing identifiers + before they become cache keys is what DeploymentAffinityCheck already does. "unscoped" is returned only + when the request carries no proxy identity at all, which is direct SDK use with one tenant by construction.""" - for prefix, field in ( - ("key", "user_api_key_hash"), - ("user", "user_api_key_user_id"), - ("team", "user_api_key_team_id"), - ("org", "user_api_key_org_id"), - ): - value = get_request_metadata_field(request_kwargs, field) - if value is not None: - return f"{prefix}:{value}" - return "unscoped" + identity = tuple((field, get_request_metadata_field(request_kwargs, field)) for field in proxy_identity_fields()) + if all(value is None for _, value in identity): + return "unscoped" + digest = hashlib.sha256(json.dumps(identity, sort_keys=True, default=str).encode("utf-8")).hexdigest() + return f"id:{digest[:32]}" def get_request_metadata_field(request_kwargs: Mapping[str, object], field: str) -> str | None: diff --git a/litellm/router_strategy/complexity_router/cache_warming/capture.py b/litellm/router_strategy/complexity_router/cache_warming/capture.py index 1ee35de3758..57e28cb4c5e 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/capture.py +++ b/litellm/router_strategy/complexity_router/cache_warming/capture.py @@ -10,6 +10,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_caller_scope, get_request_metadata_field, iter_request_metadata_dicts, + proxy_identity_fields, ) from litellm.router_strategy.complexity_router.cache_warming.eligibility import ( min_prompt_cache_tokens_for_warm_set, @@ -27,15 +28,11 @@ if TYPE_CHECKING: from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter _MAX_UNCOMPRESSED_RATIO = 8 -_ATTRIBUTION_KEYS = ( - "user_api_key", - "user_api_key_hash", - "user_api_key_user_id", - "user_api_key_team_id", - "user_api_key_org_id", - "user_api_key_project_id", - "user_api_key_end_user_id", -) +# The partition a record is stored under and the identity it is attributed to must cover the same +# dimensions, or two principals differing only in an omitted one share a record. Both are derived from +# proxy_identity_fields() for that reason; user_api_key is the key-state lookup handle rather than a +# tenant dimension, so it rides along here and not in the partition. +_ATTRIBUTION_KEYS = ("user_api_key", *proxy_identity_fields()) @lru_cache(maxsize=64) diff --git a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py index 672348e792d..922ea6a4def 100644 --- a/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py +++ b/tests/test_litellm/router_strategy/complexity_router/cache_warming/test_capture.py @@ -94,7 +94,9 @@ async def test_second_turn_overwrites_payload_and_preserves_other_model_warmth() redis = FakeRedisCache() router = _complexity_router(redis) await router._capture_session(_kwargs(), MESSAGES, "claude-sonnet-4-5") - key = CacheWarmingStore.record_key("smart-router", "key:hash-1", "sess-1") + from litellm.litellm_core_utils.core_helpers import get_caller_scope + + key = CacheWarmingStore.record_key("smart-router", get_caller_scope(_kwargs()), "sess-1") first = json.loads(redis.hashes[SESSIONS_KEY][key]) other_stamp = WarmthStamp(at=123.0, warmed=True).model_dump() redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")] = json.dumps(other_stamp) @@ -166,15 +168,17 @@ async def test_warming_does_not_accept_a_forged_billing_identity(): assert record["attribution"]["user_api_key_hash"] == "hash-1" -def test_keyless_tenants_reusing_a_session_id_do_not_share_a_record(): - """JWT and other keyless proxy principals carry no key hash. Collapsing them onto one literal would merge - distinct tenants that reuse a session id into a single record, so the last writer's payload and - attribution would win and later replays could spend under the wrong team.""" - from litellm.litellm_core_utils.core_helpers import get_caller_scope +def test_two_tenants_sharing_a_session_id_never_share_a_partition(): + """Whatever is stored under a partition is also attributed to that caller, so the partition has to cover + every dimension the attribution does. A partition built from a subset is a collision channel: principals + differing only in an omitted dimension share one record, and the last writer owns the payload and the + attribution while any state accumulated under it stays. Every dimension is driven independently here, + including the two the old strongest-identity partition dropped, end user and project.""" + from litellm.litellm_core_utils.core_helpers import get_caller_scope, proxy_identity_fields - tenant_a = {"litellm_metadata": {"session_id": "shared", "user_api_key_team_id": "team-a"}} - tenant_b = {"litellm_metadata": {"session_id": "shared", "user_api_key_team_id": "team-b"}} - assert get_caller_scope(tenant_a) != get_caller_scope(tenant_b) + base = {field: "same" for field in proxy_identity_fields()} + scopes = {get_caller_scope({"metadata": {**base, field: "other"}}) for field in proxy_identity_fields()} + assert len(scopes) == len(proxy_identity_fields()), "each identity dimension must move the partition" + assert get_caller_scope({"metadata": base}) not in scopes assert get_caller_scope({"metadata": {}}) == "unscoped" - keyed = {"metadata": {"user_api_key_hash": "hash-1", "user_api_key_team_id": "team-a"}} - assert get_caller_scope(keyed) == "key:hash-1" + assert "same" not in get_caller_scope({"metadata": base}), "identifiers must not appear in the key"