Merge e1b7f39fe8 into a5164fe210

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

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