diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 137366c955..41dc1f34e0 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -22,6 +22,13 @@ CONTAINER_OBJECT_PURPOSE = "container" _NEGATIVE_OWNER_SENTINEL = "__litellm_container_no_owner__" _CONTAINER_OWNER_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 +# tuple — different keys for the same user share the same allow-set, but +# different users with different scopes get disjoint cache entries. +_ALLOWED_CONTAINER_IDS_CACHE = InMemoryCache(max_size_in_memory=2048, default_ttl=60) + def _container_model_object_id( original_container_id: str, custom_llm_provider: str @@ -86,19 +93,20 @@ async def record_container_owner( custom_llm_provider: str, ) -> Any: container_id = _get_response_id(response) - owner = get_primary_resource_owner_scope(user_api_key_dict) - if is_proxy_admin(user_api_key_dict) and (container_id is None or owner is None): - return response if container_id is None: verbose_proxy_logger.warning( "Skipping container ownership tracking because provider response has no id" ) return response + owner = get_primary_resource_owner_scope(user_api_key_dict) if owner is None: - # Identity-less callers (no user_id / team_id / org_id / api_key / - # token) can't be uniquely stamped on the row. Stamping a + # Admins with identity (the common path: master-key auth populates + # ``user_id`` + ``api_key``) flow through the normal record path + # below so admin-created containers are still tracked. Truly + # identity-less admins (no user_id / team_id / org_id / api_key / + # token) can't be uniquely stamped on the row — stamping a # placeholder would collapse every such caller into a shared - # owner — the cross-tenant primitive we explicitly avoid. + # owner, the cross-tenant primitive we explicitly avoid. raise HTTPException( status_code=403, detail="Unable to record container ownership: caller has no identity scope.", @@ -151,6 +159,14 @@ async def record_container_owner( ) _CONTAINER_OWNER_CACHE.set_cache(model_object_id, owner) + # 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 + # tuples self-correct on the 60s TTL. + caller_scope_key = "|".join(sorted(get_resource_owner_scopes(user_api_key_dict))) + if caller_scope_key: + _ALLOWED_CONTAINER_IDS_CACHE.cache_dict.pop(caller_scope_key, None) + _ALLOWED_CONTAINER_IDS_CACHE.ttl_dict.pop(caller_scope_key, None) return response @@ -249,6 +265,11 @@ async def _get_allowed_container_ids( if not owner_scopes: return set() + cache_key = "|".join(sorted(owner_scopes)) + cached = _ALLOWED_CONTAINER_IDS_CACHE.get_cache(cache_key) + if cached is not None: + return set(cached) + prisma_client = await _get_prisma_client() if prisma_client is None: return set() @@ -259,11 +280,15 @@ async def _get_allowed_container_ids( "created_by": {"in": owner_scopes}, } ) - return { + allowed_ids = { row.model_object_id for row in rows if getattr(row, "model_object_id", None) is not None } + # ``InMemoryCache`` json-encodes values; sets aren't JSON-serializable, + # so store as a list and rehydrate above. + _ALLOWED_CONTAINER_IDS_CACHE.set_cache(cache_key, sorted(allowed_ids)) + return allowed_ids async def filter_container_list_response( diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index b046fa1536..7a6232d7ce 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -13,11 +13,19 @@ from litellm.types.containers.main import ContainerListResponse, ContainerObject @pytest.fixture(autouse=True) def clear_container_owner_cache(): - ownership._CONTAINER_OWNER_CACHE.cache_dict.clear() - ownership._CONTAINER_OWNER_CACHE.ttl_dict.clear() + for cache in ( + ownership._CONTAINER_OWNER_CACHE, + ownership._ALLOWED_CONTAINER_IDS_CACHE, + ): + cache.cache_dict.clear() + cache.ttl_dict.clear() yield - ownership._CONTAINER_OWNER_CACHE.cache_dict.clear() - ownership._CONTAINER_OWNER_CACHE.ttl_dict.clear() + for cache in ( + ownership._CONTAINER_OWNER_CACHE, + ownership._ALLOWED_CONTAINER_IDS_CACHE, + ): + cache.cache_dict.clear() + cache.ttl_dict.clear() def _container(container_id: str) -> ContainerObject: @@ -801,3 +809,112 @@ async def test_get_container_owner_caches_negative_lookups(monkeypatch): assert await ownership._get_container_owner("cntr_x", "openai") is None assert await ownership._get_container_owner("cntr_x", "openai") is None assert table.find_first.await_count == 1 + + +@pytest.mark.asyncio +async def test_allowed_container_ids_uses_cache_after_first_db_hit(monkeypatch): + """``GET /v1/containers`` filtering must not issue a fresh ``find_many`` + on every list call within the cache TTL window.""" + table = AsyncMock() + table.find_many.return_value = [ + SimpleNamespace(model_object_id="container:openai:cntr_a"), + ] + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + auth = UserAPIKeyAuth(user_id="user-1") + + first = await ownership._get_allowed_container_ids(auth) + second = await ownership._get_allowed_container_ids(auth) + third = await ownership._get_allowed_container_ids(auth) + + assert first == {"container:openai:cntr_a"} + assert second == first + assert third == first + # Single DB call across three list filterings — the cache absorbs the rest. + assert table.find_many.await_count == 1 + + +@pytest.mark.asyncio +async def test_record_container_owner_invalidates_caller_list_cache(monkeypatch): + """A just-created container must show up on the caller's next ``GET + /v1/containers`` — recording the owner has to drop the caller's + list-cache entry, otherwise the new container is invisible for up + to the cache TTL.""" + table = AsyncMock() + table.find_unique.return_value = None + table.find_many.return_value = [ + SimpleNamespace(model_object_id="container:openai:cntr_old"), + ] + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + auth = UserAPIKeyAuth(user_id="user-1") + + # Prime the list cache. + await ownership._get_allowed_container_ids(auth) + assert table.find_many.await_count == 1 + + # Recording a new owner invalidates the caller's list-cache entry. + table.find_many.return_value = [ + SimpleNamespace(model_object_id="container:openai:cntr_old"), + SimpleNamespace(model_object_id="container:openai:cntr_new"), + ] + await ownership.record_container_owner( + response=_container("cntr_new"), + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + # Next list call refreshes from DB and picks up the new container. + refreshed = await ownership._get_allowed_container_ids(auth) + assert "container:openai:cntr_new" in refreshed + assert table.find_many.await_count == 2 + + +@pytest.mark.asyncio +async def test_admin_with_identity_records_container_ownership(monkeypatch): + """The admin early-return only short-circuits when there's literally no + container ID to stamp. An admin with identity (the master-key path + populates ``user_id`` + ``api_key``) creates an owned row like any + other caller, so admin-created containers aren't permanently + untracked.""" + table = AsyncMock() + table.find_unique.return_value = None + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + admin_auth = UserAPIKeyAuth( + user_id="proxy-admin", + user_role=ownership.is_proxy_admin.__module__.split(".")[0] + and "proxy_admin", # placeholder; the create flow doesn't actually gate on the role + ) + # Use the real role enum value. + from litellm.proxy._types import LitellmUserRoles + + admin_auth.user_role = LitellmUserRoles.PROXY_ADMIN.value + + await ownership.record_container_owner( + response=_container("cntr_admin"), + user_api_key_dict=admin_auth, + custom_llm_provider="openai", + ) + + table.create.assert_awaited_once() + created_data = table.create.await_args.kwargs["data"] + assert created_data["created_by"] == "proxy-admin"