fix(cache_warming): derive a record's partition from the caller's full identity

Anything stored under a partition is also attributed to that caller, so the
partition has to cover every dimension the attribution does. It did not. The
attribution carried seven identity fields while the partition returned the first
of four, and the gap between those two lists was a cross-tenant collision channel:
principals differing only in an omitted dimension shared one record, and whoever
wrote last owned the payload and the attribution while state accumulated under it
stayed.

End users are the reachable case, since they share a virtual key by design and a
key hash cannot tell them apart. One end user could seed a session's touched-model
set, the next end user's turn would overwrite the payload and attribution while
inheriting those models, and the replays then sent that caller's prompt to a model
they had never used and billed it to their end-user budget. Projects were the same
bug, unreported.

The partition is now composed from every dimension rather than the strongest one,
and both it and the attribution read the same list, taken from the type that
declares those fields rather than spelled out in either place, so they cannot
drift apart again. That also removes the need for a rule about state surviving an
attribution change: a different attribution is a different partition, so there is
no shared state to clear.

Values are hashed rather than concatenated. End-user and project identifiers are
customer-chosen strings that would otherwise sit in Redis key names and in every
log line quoting one, and hashing identifiers before they become cache keys is
what DeploymentAffinityCheck already does.

Operational note: this changes partition strings, so existing cache-warming records
and complexity-router session-affinity pins are not read after deploy. Both are
caches with TTLs; pins re-derive on the next turn and records re-capture.
This commit is contained in:
Tin Chi Lo 2026-07-30 17:39:47 -07:00
parent 18dc38f9f6
commit 4ff2ecd4c8
3 changed files with 59 additions and 37 deletions

View file

@ -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:

View file

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

View file

@ -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"