diff --git a/litellm/proxy/common_utils/resource_ownership.py b/litellm/proxy/common_utils/resource_ownership.py index 936c4e18bde..ff5517cb853 100644 --- a/litellm/proxy/common_utils/resource_ownership.py +++ b/litellm/proxy/common_utils/resource_ownership.py @@ -90,3 +90,21 @@ def user_can_access_resource_owner( if owner is None: return False return owner in get_resource_owner_scopes(user_api_key_dict) + + +def user_can_access_resource_owner_or_team( + owner: Optional[str], + resource_team_id: Optional[str], + user_api_key_dict: Optional[UserAPIKeyAuth], +) -> bool: + """Owner-scope access, widened to members of a team the resource is + shared with. ``resource_team_id`` is only set when the creator opted the + resource into team visibility, so a private resource (``resource_team_id`` + None) falls back to owner-only access and the team branch never widens it. + This is what lets a team's service account (whose only scope is + ``team:``) read a container a teammate created.""" + if user_can_access_resource_owner(owner, user_api_key_dict): + return True + if user_api_key_dict is None or not resource_team_id: + return False + return f"team:{resource_team_id}" in get_resource_owner_scopes(user_api_key_dict) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 8118d53b9f6..b861f909786 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -10,7 +10,7 @@ from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, get_resource_owner_scopes, is_proxy_admin, - user_can_access_resource_owner, + user_can_access_resource_owner_or_team, ) from litellm.repositories.table_repositories import ManagedObjectRepository from litellm.responses.utils import ResponsesAPIRequestUtils @@ -31,6 +31,13 @@ _CONTAINER_OWNER_CACHE = InMemoryCache(max_size_in_memory=10000, default_ttl=60) _NEGATIVE_STORED_ID_SENTINEL = "__litellm_container_no_stored_id__" _CONTAINER_STORED_ID_CACHE = InMemoryCache(max_size_in_memory=10000, default_ttl=60) +# Caches the team a container is shared with (its ``team_id``), populated from +# the same row fetch as the owner cache. A negative sentinel marks a private +# container (no team), so the access check can distinguish "private" from a +# cache miss. +_NEGATIVE_TEAM_SENTINEL = "__litellm_container_no_team__" +_CONTAINER_TEAM_CACHE = InMemoryCache(max_size_in_memory=10000, default_ttl=60) + # Per-caller-scope cache for ``GET /v1/containers`` list filtering. Without # this, every list call issues a fresh ``find_many`` against # ``litellm_managedobjecttable``. The cache key is the sorted owner-scope @@ -163,13 +170,27 @@ async def record_container_owners_from_responses_response( # batch — other containers in the same response should still # get recorded so their follow-up file API calls don't 403. verbose_proxy_logger.exception( - "Failed to record container ownership from responses output " - "for container_id=%s: %s", + "Failed to record container ownership from responses output for container_id=%s: %s", container_id, e, ) +def _resolve_container_team_id(user_api_key_dict: UserAPIKeyAuth) -> Optional[str]: + """Return the team to share a newly-created container with, or ``None`` to + keep it private to its creator. A container is team-visible only when the + creator belongs to a team whose metadata opts in + (``container_visibility == "team"``); this preserves the private-by-default + behaviour and makes team sharing an explicit team-level switch.""" + team_id = getattr(user_api_key_dict, "team_id", None) + if not team_id: + return None + team_metadata = getattr(user_api_key_dict, "team_metadata", None) or {} + if team_metadata.get("container_visibility") == "team": + return team_id + return None + + async def record_container_owner( response: Any, user_api_key_dict: UserAPIKeyAuth, @@ -214,15 +235,22 @@ async def record_container_owner( ) return response + container_team_id = _resolve_container_team_id(user_api_key_dict) table = ManagedObjectRepository(prisma_client).table existing = await table.find_unique(where={"model_object_id": model_object_id}) if existing is not None: if getattr(existing, "file_purpose", None) != CONTAINER_OBJECT_PURPOSE: raise HTTPException(status_code=500, detail="Unable to track container") - if not user_can_access_resource_owner( - getattr(existing, "created_by", None), user_api_key_dict + if not user_can_access_resource_owner_or_team( + getattr(existing, "created_by", None), + getattr(existing, "team_id", None), + user_api_key_dict, ): raise HTTPException(status_code=403, detail="Forbidden") + # Visibility is fixed at create time; a re-record (e.g. a colliding + # container id) refreshes the encoded id/object but must not silently + # flip an existing container's visibility, so ``team_id`` is left as-is. + effective_team_id = getattr(existing, "team_id", None) await table.update( where={"model_object_id": model_object_id}, data={ @@ -232,6 +260,7 @@ async def record_container_owner( }, ) else: + effective_team_id = container_team_id await table.create( data={ "unified_object_id": container_id, @@ -239,12 +268,16 @@ async def record_container_owner( "file_object": file_object_json, "file_purpose": CONTAINER_OBJECT_PURPOSE, "created_by": owner, + "team_id": container_team_id, "updated_by": owner, } ) _CONTAINER_OWNER_CACHE.set_cache(model_object_id, owner) _CONTAINER_STORED_ID_CACHE.set_cache(model_object_id, container_id) + _CONTAINER_TEAM_CACHE.set_cache( + model_object_id, effective_team_id or _NEGATIVE_TEAM_SENTINEL + ) # Drop the caller's own list-cache entry so the just-created container # shows up on their next ``GET /v1/containers``. Other callers with # disjoint scope tuples have their own entries; intersecting-scope @@ -293,6 +326,11 @@ async def _get_container_owner( else _NEGATIVE_STORED_ID_SENTINEL ), ) + team_id = getattr(row, "team_id", None) if row is not None else None + _CONTAINER_TEAM_CACHE.set_cache( + model_object_id, + team_id if isinstance(team_id, str) and team_id else _NEGATIVE_TEAM_SENTINEL, + ) return owner @@ -338,6 +376,40 @@ async def _get_stored_container_id( return stored_id if isinstance(stored_id, str) and stored_id else None +async def _get_container_team_id( + original_container_id: str, custom_llm_provider: str +) -> Optional[str]: + """Return the team a container is shared with, or ``None`` if it is + private to its creator. ``_get_container_owner`` populates this cache from + the same row, so calling it right after an owner lookup is a cache hit.""" + model_object_id = _container_model_object_id( + original_container_id, custom_llm_provider + ) + + cached = _CONTAINER_TEAM_CACHE.get_cache(model_object_id) + if cached == _NEGATIVE_TEAM_SENTINEL: + return None + if isinstance(cached, str) and cached: + return cached + + prisma_client = await _get_prisma_client() + if prisma_client is None: + return None + + row = await ManagedObjectRepository(prisma_client).table.find_first( + where={ + "model_object_id": model_object_id, + "file_purpose": CONTAINER_OBJECT_PURPOSE, + } + ) + team_id = getattr(row, "team_id", None) if row is not None else None + _CONTAINER_TEAM_CACHE.set_cache( + model_object_id, + team_id if isinstance(team_id, str) and team_id else _NEGATIVE_TEAM_SENTINEL, + ) + return team_id if isinstance(team_id, str) and team_id else None + + async def assert_user_can_access_container( container_id: str, user_api_key_dict: UserAPIKeyAuth, @@ -354,7 +426,12 @@ async def assert_user_can_access_container( # that pre-date this enforcement need an admin to either re-create # via the now-tracked flow or assign ``created_by`` on the row. owner = await _get_container_owner(original_container_id, resolved_provider) - if not user_can_access_resource_owner(owner, user_api_key_dict): + resource_team_id = await _get_container_team_id( + original_container_id, resolved_provider + ) + if not user_can_access_resource_owner_or_team( + owner, resource_team_id, user_api_key_dict + ): raise HTTPException(status_code=403, detail="Forbidden") return original_container_id, resolved_provider @@ -412,11 +489,20 @@ async def _get_allowed_container_ids( if prisma_client is None: return set() + # A caller sees containers they own (created_by in their scopes) plus any + # team-visible container shared with their team. ``owner_scopes`` already + # varies by team, so the cache key stays correct. + where_filter: Dict[str, Any] = {"file_purpose": CONTAINER_OBJECT_PURPOSE} + caller_team_id = getattr(user_api_key_dict, "team_id", None) + if caller_team_id: + where_filter["OR"] = [ + {"created_by": {"in": owner_scopes}}, + {"team_id": caller_team_id}, + ] + else: + where_filter["created_by"] = {"in": owner_scopes} rows = await ManagedObjectRepository(prisma_client).table.find_many( - where={ - "file_purpose": CONTAINER_OBJECT_PURPOSE, - "created_by": {"in": owner_scopes}, - } + where=where_filter ) allowed_ids = { row.model_object_id diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 176405bb9ca..35a865bdc68 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -14,17 +14,17 @@ from litellm.types.containers.main import ContainerListResponse, ContainerObject @pytest.fixture(autouse=True) def clear_container_owner_cache(): - for cache in ( + caches = ( ownership._CONTAINER_OWNER_CACHE, + ownership._CONTAINER_STORED_ID_CACHE, + ownership._CONTAINER_TEAM_CACHE, ownership._ALLOWED_CONTAINER_IDS_CACHE, - ): + ) + for cache in caches: cache.cache_dict.clear() cache.ttl_dict.clear() yield - for cache in ( - ownership._CONTAINER_OWNER_CACHE, - ownership._ALLOWED_CONTAINER_IDS_CACHE, - ): + for cache in caches: cache.cache_dict.clear() cache.ttl_dict.clear() @@ -932,10 +932,7 @@ async def test_should_record_containers_from_responses_output_for_service_accoun AsyncMock(return_value=prisma_client), ) auth = UserAPIKeyAuth(team_id="team-1") - encoded_container_id = ( - "cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmR" - "lZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl" - ) + encoded_container_id = "cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmRlZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl" responses_payload = { "output": [ { @@ -972,10 +969,7 @@ async def test_should_record_containers_from_responses_output_for_service_accoun async def test_service_account_can_access_container_after_responses_tracking( monkeypatch, ): - encoded_container_id = ( - "cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmR" - "lZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl" - ) + encoded_container_id = "cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmRlZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl" table = AsyncMock() table.find_unique.return_value = None prisma_client = SimpleNamespace( @@ -1024,10 +1018,7 @@ async def test_should_record_container_ownership_after_streaming_responses_finis """ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing - encoded_container_id = ( - "cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmR" - "lZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl" - ) + encoded_container_id = "cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmRlZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl" response_body = SimpleNamespace( output=[ SimpleNamespace( @@ -1107,3 +1098,127 @@ async def test_streaming_ownership_wrap_no_op_when_stream_did_not_complete( assert chunks == ["data: chunk-1\n\n"] record.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Team-visibility opt-in (cross-principal access). +# +# Regression coverage for the Disney case: a container created by a user key +# (owner = user_id) was unreachable by the team's service account (scope only +# ``team:``). The fix lets a creator's team opt into team visibility, which +# stamps the container's ``team_id`` so any same-team principal can read it, +# while keeping private-to-creator the default and preserving cross-team +# isolation. +# --------------------------------------------------------------------------- + + +def _row(created_by, team_id, container_id="cntr_x"): + return SimpleNamespace( + created_by=created_by, + team_id=team_id, + unified_object_id=container_id, + file_object="{}", + file_purpose="container", + ) + + +def _prisma(monkeypatch, *, find_unique=None, find_first=None): + table = AsyncMock() + table.find_unique.return_value = find_unique + table.find_first.return_value = find_first + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, "_get_prisma_client", AsyncMock(return_value=prisma_client) + ) + return table + + +@pytest.mark.asyncio +async def test_user_created_container_is_private_by_default(monkeypatch): + """A user whose team has NOT opted in creates a private container: + team_id is left unset so no teammate (or service account) inherits access.""" + table = _prisma(monkeypatch, find_unique=None) + creator = UserAPIKeyAuth(user_id="alice", team_id="team-1", team_metadata={}) + + await ownership.record_container_owner( + response=_container("cntr_priv"), + user_api_key_dict=creator, + custom_llm_provider="openai", + ) + + data = table.create.call_args.kwargs["data"] + assert data["created_by"] == "alice" + assert data["team_id"] is None + + +@pytest.mark.asyncio +async def test_team_opt_in_stamps_team_id_but_keeps_creator(monkeypatch): + """With the team opted in, the container is stamped with the team while + still recording the individual creator in created_by.""" + table = _prisma(monkeypatch, find_unique=None) + creator = UserAPIKeyAuth( + user_id="alice", + team_id="team-1", + team_metadata={"container_visibility": "team"}, + ) + + await ownership.record_container_owner( + response=_container("cntr_shared"), + user_api_key_dict=creator, + custom_llm_provider="openai", + ) + + data = table.create.call_args.kwargs["data"] + assert data["created_by"] == "alice" + assert data["team_id"] == "team-1" + + +@pytest.mark.asyncio +async def test_same_team_service_account_can_access_team_visible_container(monkeypatch): + """The Disney case, fixed: a user-created, team-visible container is + reachable by the team's service account (user_id=None, scope team:team-1).""" + _prisma(monkeypatch, find_first=_row("alice", "team-1", "cntr_shared")) + service_account = UserAPIKeyAuth(team_id="team-1") + + original_id, provider = await ownership.assert_user_can_access_container( + container_id="cntr_shared", + user_api_key_dict=service_account, + custom_llm_provider="openai", + ) + + assert original_id == "cntr_shared" + assert provider == "openai" + + +@pytest.mark.asyncio +async def test_service_account_denied_on_private_user_container(monkeypatch): + """Regression guard: when the team has NOT opted in (team_id None), a + same-team service account is still forbidden from a user's container.""" + _prisma(monkeypatch, find_first=_row("alice", None, "cntr_priv")) + service_account = UserAPIKeyAuth(team_id="team-1") + + with pytest.raises(HTTPException) as exc: + await ownership.assert_user_can_access_container( + container_id="cntr_priv", + user_api_key_dict=service_account, + custom_llm_provider="openai", + ) + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_other_team_cannot_access_team_visible_container(monkeypatch): + """Cross-team isolation is preserved: a service account from a different + team cannot read a team-visible container.""" + _prisma(monkeypatch, find_first=_row("alice", "team-1", "cntr_shared")) + other_team = UserAPIKeyAuth(team_id="team-2") + + with pytest.raises(HTTPException) as exc: + await ownership.assert_user_can_access_container( + container_id="cntr_shared", + user_api_key_dict=other_team, + custom_llm_provider="openai", + ) + assert exc.value.status_code == 403