diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 8fdca178388..164c8608922 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -229,6 +229,27 @@ def iter_request_metadata_dicts(request_kwargs: Mapping[str, object]) -> tuple[M ) +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 + 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" + + def get_request_metadata_field(request_kwargs: Mapping[str, object], field: str) -> str | None: """First stringified value of ``field`` across the top-level metadata slots.""" for metadata in iter_request_metadata_dicts(request_kwargs): diff --git a/litellm/router_strategy/complexity_router/cache_warming/capture.py b/litellm/router_strategy/complexity_router/cache_warming/capture.py index 877cd109b94..9c7bb3f8f5f 100644 --- a/litellm/router_strategy/complexity_router/cache_warming/capture.py +++ b/litellm/router_strategy/complexity_router/cache_warming/capture.py @@ -18,6 +18,7 @@ from litellm.router_strategy.complexity_router.cache_warming.types import ( compress_payload, ) from litellm.litellm_core_utils.core_helpers import ( + get_caller_scope, get_request_metadata_field, iter_request_metadata_dicts, ) @@ -210,7 +211,7 @@ async def capture_session( store = strategy.get_cache_warming_store() if store is None: return - caller_scope = get_request_metadata_field(request_kwargs, "user_api_key_hash") or "unscoped" + caller_scope = get_caller_scope(request_kwargs) await store.upsert_session( caller_scope=caller_scope, session_id=session_id, diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index d129b2a50f7..157f994af2c 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -27,7 +27,7 @@ from pydantic import BaseModel from litellm._logging import verbose_router_logger from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import get_request_metadata_field +from litellm.litellm_core_utils.core_helpers import get_caller_scope, get_request_metadata_field from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.types.utils import ( ModelResponse, @@ -664,7 +664,7 @@ class ComplexityRouter(CustomLogger): return None from litellm.router_strategy.complexity_router.cache_warming.types import is_cache_fresh - caller_scope = get_request_metadata_field(request_kwargs, "user_api_key_hash") or "unscoped" + caller_scope = get_caller_scope(request_kwargs) record_key = store.record_key(self.model_name, caller_scope, session_id) record = await store.get_record(record_key) if record is None: @@ -1108,7 +1108,7 @@ class ComplexityRouter(CustomLogger): # same client-supplied session_id can't poison each other's routing pin. Falls # back to "unscoped" only when there's no authenticated caller to scope by # (e.g. direct Router usage without the proxy layer). - caller_scope = get_request_metadata_field(request_kwargs, "user_api_key_hash") or "unscoped" + caller_scope = get_caller_scope(request_kwargs) return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" async def async_pre_routing_hook( 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 d189b060a76..4e3929290ae 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 @@ -99,7 +99,7 @@ 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", "hash-1", "sess-1") + key = CacheWarmingStore.record_key("smart-router", "key:hash-1", "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) @@ -177,3 +177,17 @@ async def test_warming_does_not_accept_a_forged_billing_identity(): record = _stored_records(redis)[0] assert record["attribution"]["user_api_key"] == "hash-1" 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 + + 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) + 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"