diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 11733ce4cee..d98da24529c 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -63,6 +63,7 @@ class Cache: port: Optional[str] = None, password: Optional[str] = None, namespace: Optional[str] = None, + add_team_id_to_cache_key: bool = False, ttl: Optional[float] = None, default_in_memory_ttl: Optional[float] = None, default_in_redis_ttl: Optional[float] = None, @@ -257,6 +258,7 @@ class Cache: self.supported_call_types = supported_call_types # default to ["completion", "acompletion", "embedding", "aembedding"] self.type = type self.namespace = namespace + self.add_team_id_to_cache_key = add_team_id_to_cache_key self.redis_flush_size = redis_flush_size self.ttl = ttl self.mode: CacheMode = mode or CacheMode.default_on @@ -309,6 +311,8 @@ class Cache: param_value = kwargs[param] cache_key += f"{str(param)}: {str(param_value)}" + cache_key += self._get_team_scope_for_cache_key(**kwargs) + verbose_logger.debug("\nCreated cache key: %s", cache_key) hashed_cache_key = Cache._get_hashed_cache_key(cache_key) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) @@ -420,6 +424,44 @@ class Cache: verbose_logger.debug("Hashed cache key (SHA-256): %s", hash_hex) return hash_hex + def _get_team_scope_for_cache_key(self, **kwargs) -> str: + """Optionally scope the cache key by the requesting tenant. + + On a multi-tenant proxy the cache key is otherwise derived only from the + request parameters, so two different teams (tenants) sending the same + request share cache entries - one team can be served another's cached + response. When ``add_team_id_to_cache_key`` is enabled (via + ``cache_params``), the authenticated team id is folded into the cache key + so entries are not reused across teams; same-team requests still share the + cache. A request whose key has no team falls back to the (hashed) virtual + key, so it is still isolated rather than silently sharing the global entry. + Opt-in - the default preserves the existing behavior. + + Security: the scope is read only from ``litellm_params["metadata"]``, which + the proxy populates from the authenticated key in + ``litellm.proxy.litellm_pre_call_utils`` - it strips any client-supplied + ``user_api_key_*`` fields from the request-body metadata before writing the + authenticated values, so the team/key used here cannot be forged by the + caller. The authenticated ``user_api_key_auth`` object is preferred. Direct + SDK calls (no proxy metadata) fall through to "" and keep today's behavior. + """ + if not self.add_team_id_to_cache_key: + return "" + litellm_params = kwargs.get("litellm_params") or {} + metadata = litellm_params.get("metadata") or {} + user_api_key_auth = metadata.get("user_api_key_auth") + team_id = getattr(user_api_key_auth, "team_id", None) or metadata.get( + "user_api_key_team_id" + ) + if team_id: + return f"user_api_key_team_id: {team_id}" + api_key = getattr(user_api_key_auth, "api_key", None) or metadata.get( + "user_api_key" + ) + if api_key: + return f"user_api_key: {api_key}" + return "" + def _add_namespace_to_cache_key(self, hash_hex: str, **kwargs) -> str: """ If a redis namespace is provided, add it to the cache key diff --git a/tests/test_litellm/caching/test_team_scoped_cache_key.py b/tests/test_litellm/caching/test_team_scoped_cache_key.py new file mode 100644 index 00000000000..b439def0360 --- /dev/null +++ b/tests/test_litellm/caching/test_team_scoped_cache_key.py @@ -0,0 +1,93 @@ +"""Tests for the opt-in team-scoped cache key (Cache.add_team_id_to_cache_key). + +On a multi-tenant proxy the response-cache key is otherwise derived only from the +request params, so two teams sending the same request share cache entries - one +team can be served another's cached response. With add_team_id_to_cache_key=True +the authenticated team id (or, with no team, the hashed virtual key) is folded +into the cache key so entries are not reused across tenants; same-team requests +still share the cache. The flag defaults to False, preserving existing behavior. + +The scope is read only from the proxy-trusted litellm_params["metadata"] (which +the proxy populates from the authenticated key and strips of any client-supplied +user_api_key_* fields), never from the caller-supplied top-level metadata, so a +client cannot forge the team/key used for scoping. +""" + +from litellm.caching.caching import Cache + + +def _key_for_team(cache: Cache, team_id: str) -> str: + return cache.get_cache_key( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hello"}], + litellm_params={"metadata": {"user_api_key_team_id": team_id}}, + ) + + +def test_team_scoped_cache_key_isolates_teams(): + cache = Cache(add_team_id_to_cache_key=True) + assert _key_for_team(cache, "team-a") != _key_for_team( + cache, "team-b" + ) # different teams -> different keys + assert _key_for_team(cache, "team-a") == _key_for_team( + cache, "team-a" + ) # same team -> same key + + +def test_cache_key_shared_across_teams_by_default(): + cache = Cache() # flag defaults to False -> existing behavior preserved + assert _key_for_team(cache, "team-a") == _key_for_team( + cache, "team-b" + ) # team ignored -> shared key + + +def test_team_scoped_cache_key_falls_back_to_api_key_when_no_team(): + cache = Cache(add_team_id_to_cache_key=True) + + def key_for(api_key: str) -> str: + return cache.get_cache_key( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hello"}], + litellm_params={"metadata": {"user_api_key": api_key}}, + ) + + # no team -> fall back to the hashed virtual key, so callers stay isolated + assert key_for("hashed-key-1") != key_for("hashed-key-2") + + +def test_team_scope_ignores_caller_supplied_top_level_metadata(): + # Security: the scope must come only from the proxy-trusted litellm_params + # metadata, never from caller-supplied top-level metadata. A client must not + # be able to forge a team by putting user_api_key_team_id in the request body. + cache = Cache(add_team_id_to_cache_key=True) + base = cache.get_cache_key( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hello"}], + ) + forged = cache.get_cache_key( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hello"}], + metadata={"user_api_key_team_id": "victim-team"}, + ) + assert base == forged # caller-supplied team is ignored -> no cross-tenant forge + + +def test_team_scope_reads_authenticated_user_api_key_auth_object(): + # The proxy attaches the authenticated UserAPIKeyAuth as + # litellm_params["metadata"]["user_api_key_auth"]; its team_id is the + # un-forgeable source and takes precedence over the flat field. + cache = Cache(add_team_id_to_cache_key=True) + + class _Auth: + def __init__(self, team_id: str) -> None: + self.team_id = team_id + self.api_key = "hashed-key" + + def key_for(team_id: str) -> str: + return cache.get_cache_key( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hello"}], + litellm_params={"metadata": {"user_api_key_auth": _Auth(team_id)}}, + ) + + assert key_for("team-a") != key_for("team-b")