diff --git a/litellm/proxy/common_utils/resource_ownership.py b/litellm/proxy/common_utils/resource_ownership.py index 1bc769a8c05..9b55554fbc6 100644 --- a/litellm/proxy/common_utils/resource_ownership.py +++ b/litellm/proxy/common_utils/resource_ownership.py @@ -2,6 +2,8 @@ from typing import List, Optional from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +UNSCOPED_RESOURCE_OWNER_SCOPE = "__litellm_unscoped_proxy__" + def is_proxy_admin(user_api_key_dict: Optional[UserAPIKeyAuth]) -> bool: if user_api_key_dict is None: @@ -41,6 +43,10 @@ def get_resource_owner_scopes( _add(f"org:{user_api_key_dict.org_id}") if user_api_key_dict.api_key: _add(f"key:{user_api_key_dict.api_key}") + if user_api_key_dict.token: + _add(f"key:{user_api_key_dict.token}") + if not scopes: + _add(UNSCOPED_RESOURCE_OWNER_SCOPE) return scopes @@ -59,7 +65,9 @@ def get_primary_resource_owner_scope( return f"org:{user_api_key_dict.org_id}" if user_api_key_dict.api_key: return f"key:{user_api_key_dict.api_key}" - return None + if user_api_key_dict.token: + return f"key:{user_api_key_dict.token}" + return UNSCOPED_RESOURCE_OWNER_SCOPE def user_can_access_resource_owner( diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 72a7c6c8746..b67e1e28d8b 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -122,6 +122,11 @@ async def create_container( user_api_base=user_api_base, version=version, ) + return await record_container_owner( + response=response, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) except Exception as e: raise await processor._handle_llm_api_exception( e=e, @@ -129,11 +134,6 @@ async def create_container( proxy_logging_obj=proxy_logging_obj, version=version, ) - return await record_container_owner( - response=response, - user_api_key_dict=user_api_key_dict, - custom_llm_provider=custom_llm_provider, - ) @router.get( @@ -220,6 +220,11 @@ async def list_containers( user_api_base=user_api_base, version=version, ) + return await filter_container_list_response( + response=response, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) except Exception as e: raise await processor._handle_llm_api_exception( e=e, @@ -227,11 +232,6 @@ async def list_containers( proxy_logging_obj=proxy_logging_obj, version=version, ) - return await filter_container_list_response( - response=response, - user_api_key_dict=user_api_key_dict, - custom_llm_provider=custom_llm_provider, - ) @router.get( diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index b769b2c0953..edb1b65c715 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -94,10 +94,14 @@ async def record_container_owner( ) -> Any: container_id = _get_response_id(response) owner = get_primary_resource_owner_scope(user_api_key_dict) - prisma_client = await _get_prisma_client() if is_proxy_admin(user_api_key_dict) and (container_id is None or owner is None): return response - if container_id is None or owner is None: + if container_id is None: + verbose_proxy_logger.warning( + "Skipping container ownership tracking because provider response has no id" + ) + return response + if owner is None: raise HTTPException(status_code=500, detail="Unable to track container") original_container_id, resolved_provider = decode_container_id_for_ownership( @@ -112,16 +116,17 @@ async def record_container_owner( file_object["custom_llm_provider"] = resolved_provider file_object["provider_container_id"] = original_container_id - if prisma_client is None: - existing_owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) - if existing_owner is not None and not user_can_access_resource_owner( - existing_owner, user_api_key_dict - ): - raise HTTPException(status_code=403, detail="Forbidden") - _IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner - return response - try: + prisma_client = await _get_prisma_client() + if prisma_client is None: + existing_owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) + if existing_owner is not None and not user_can_access_resource_owner( + existing_owner, user_api_key_dict + ): + raise HTTPException(status_code=403, detail="Forbidden") + _IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner + return response + existing = await prisma_client.db.litellm_managedobjecttable.find_unique( where={"model_object_id": model_object_id} ) @@ -169,25 +174,30 @@ async def _get_container_owner( original_container_id: str, custom_llm_provider: str, ) -> Optional[str]: - prisma_client = await _get_prisma_client() - if prisma_client is None: - return _IN_MEMORY_CONTAINER_OWNERS.get( - _container_model_object_id( - original_container_id, - custom_llm_provider, - ) - ) - - row = await prisma_client.db.litellm_managedobjecttable.find_first( - where={ - "model_object_id": _container_model_object_id( - original_container_id, - custom_llm_provider, - ), - "file_purpose": CONTAINER_OBJECT_PURPOSE, - } + model_object_id = _container_model_object_id( + original_container_id, + custom_llm_provider, ) - return getattr(row, "created_by", None) if row is not None else None + try: + prisma_client = await _get_prisma_client() + if prisma_client is None: + return _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) + + row = await prisma_client.db.litellm_managedobjecttable.find_first( + where={ + "model_object_id": model_object_id, + "file_purpose": CONTAINER_OBJECT_PURPOSE, + } + ) + return getattr(row, "created_by", None) if row is not None else None + except Exception as e: + verbose_proxy_logger.warning( + "Failed to load container ownership for container_id=%s; " + "falling back to in-process tracking: %s", + model_object_id, + e, + ) + return _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) async def assert_user_can_access_container( @@ -257,30 +267,38 @@ async def _get_allowed_container_ids( user_api_key_dict: UserAPIKeyAuth, custom_llm_provider: str, ) -> Set[str]: - prisma_client = await _get_prisma_client() - if prisma_client is None: - owner_scopes = get_resource_owner_scopes(user_api_key_dict) - return { - model_object_id - for model_object_id, owner in _IN_MEMORY_CONTAINER_OWNERS.items() - if owner in owner_scopes - } - owner_scopes = get_resource_owner_scopes(user_api_key_dict) if not owner_scopes: return set() - rows = await prisma_client.db.litellm_managedobjecttable.find_many( - where={ - "file_purpose": CONTAINER_OBJECT_PURPOSE, - "created_by": {"in": owner_scopes}, - } - ) - return { - row.model_object_id - for row in rows - if getattr(row, "model_object_id", None) is not None + in_memory_allowed_ids = { + model_object_id + for model_object_id, owner in _IN_MEMORY_CONTAINER_OWNERS.items() + if owner in owner_scopes } + try: + prisma_client = await _get_prisma_client() + if prisma_client is None: + return in_memory_allowed_ids + + rows = await prisma_client.db.litellm_managedobjecttable.find_many( + where={ + "file_purpose": CONTAINER_OBJECT_PURPOSE, + "created_by": {"in": owner_scopes}, + } + ) + return { + row.model_object_id + for row in rows + if getattr(row, "model_object_id", None) is not None + } + except Exception as e: + verbose_proxy_logger.warning( + "Failed to load allowed container ids; falling back to in-process " + "tracking: %s", + e, + ) + return in_memory_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 5ec28e81155..e801300b294 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -82,6 +82,83 @@ async def test_should_record_team_owner_for_keys_without_user_id(monkeypatch): assert data["updated_by"] == "team:team-1" +@pytest.mark.asyncio +async def test_should_record_token_owner_for_keys_without_user_team_or_org(monkeypatch): + 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), + ) + auth = UserAPIKeyAuth(token="hashed-token") + + await ownership.record_container_owner( + response=_container("cntr_provider"), + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + data = table.create.await_args.kwargs["data"] + assert data["created_by"] == "key:hashed-token" + assert data["updated_by"] == "key:hashed-token" + + +@pytest.mark.asyncio +async def test_should_record_unscoped_owner_for_identityless_proxy_auth(monkeypatch): + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=None), + ) + auth = UserAPIKeyAuth() + + await ownership.record_container_owner( + response=_container("cntr_provider"), + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert ( + ownership._IN_MEMORY_CONTAINER_OWNERS["container:openai:cntr_provider"] + == "__litellm_unscoped_proxy__" + ) + original_id, provider = await ownership.assert_user_can_access_container( + container_id="cntr_provider", + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + assert original_id == "cntr_provider" + assert provider == "openai" + + +@pytest.mark.asyncio +async def test_should_skip_owner_record_when_provider_response_has_no_id(monkeypatch): + table = AsyncMock() + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + response = {"object": "container"} + + returned = await ownership.record_container_owner( + response=response, + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + custom_llm_provider="openai", + ) + + assert returned == response + table.find_unique.assert_not_awaited() + table.create.assert_not_awaited() + + @pytest.mark.asyncio async def test_should_fallback_to_memory_when_persistent_owner_record_fails( monkeypatch, @@ -178,6 +255,55 @@ async def test_should_deny_untracked_container_access_by_default(monkeypatch): assert exc.value.status_code == 403 +@pytest.mark.asyncio +async def test_should_fallback_to_memory_when_owner_lookup_fails(monkeypatch): + table = AsyncMock() + table.find_first.side_effect = Exception("db unavailable") + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + ownership._IN_MEMORY_CONTAINER_OWNERS["container:openai:cntr_owned"] = "user-1" + auth = UserAPIKeyAuth(user_id="user-1") + + original_id, provider = await ownership.assert_user_can_access_container( + container_id="cntr_owned", + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert original_id == "cntr_owned" + assert provider == "openai" + + +@pytest.mark.asyncio +async def test_should_fail_closed_when_owner_lookup_fails_without_memory(monkeypatch): + table = AsyncMock() + table.find_first.side_effect = Exception("db unavailable") + 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") + + with pytest.raises(HTTPException) as exc: + await ownership.assert_user_can_access_container( + container_id="cntr_owned", + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert exc.value.status_code == 403 + + @pytest.mark.asyncio async def test_should_allow_untracked_container_access_when_enabled(monkeypatch): monkeypatch.setattr( @@ -362,6 +488,38 @@ async def test_should_filter_container_list_with_in_memory_ownership(monkeypatch assert [item.id for item in filtered.data] == ["cntr_owned"] +@pytest.mark.asyncio +async def test_should_filter_container_list_with_memory_when_db_lookup_fails( + monkeypatch, +): + table = AsyncMock() + table.find_many.side_effect = Exception("db unavailable") + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + ownership._IN_MEMORY_CONTAINER_OWNERS["container:openai:cntr_owned"] = "user-1" + auth = UserAPIKeyAuth(user_id="user-1") + response = ContainerListResponse( + object="list", + data=[_container("cntr_owned"), _container("cntr_other")], + has_more=True, + ) + + filtered = await ownership.filter_container_list_response( + response=response, + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert [item.id for item in filtered.data] == ["cntr_owned"] + assert filtered.has_more is False + + @pytest.mark.asyncio async def test_should_forward_decoded_container_id_for_proxy_forwarding(monkeypatch): from litellm.proxy.container_endpoints import handler_factory diff --git a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py index fde1806f92c..f6946e519d1 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -44,6 +44,57 @@ async def test_should_store_team_owner_for_keys_without_user_id(monkeypatch): assert table.create.await_args.kwargs["data"]["updated_by"] == "team:team-1" +@pytest.mark.asyncio +async def test_should_store_token_owner_for_keys_without_user_team_or_org(monkeypatch): + table = AsyncMock() + table.create.side_effect = lambda data: _skill(data["skill_id"], data["created_by"]) + prisma_client = type( + "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} + )() + monkeypatch.setattr( + LiteLLMSkillsHandler, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + + auth = UserAPIKeyAuth(token="hashed-token") + + skill = await LiteLLMSkillsHandler.create_skill( + data=NewSkillRequest(display_title="skill"), + user_api_key_dict=auth, + ) + + assert skill.created_by == "key:hashed-token" + assert table.create.await_args.kwargs["data"]["updated_by"] == "key:hashed-token" + + +@pytest.mark.asyncio +async def test_should_store_unscoped_owner_for_identityless_proxy_auth(monkeypatch): + table = AsyncMock() + table.create.side_effect = lambda data: _skill(data["skill_id"], data["created_by"]) + prisma_client = type( + "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} + )() + monkeypatch.setattr( + LiteLLMSkillsHandler, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + + auth = UserAPIKeyAuth() + + skill = await LiteLLMSkillsHandler.create_skill( + data=NewSkillRequest(display_title="skill"), + user_api_key_dict=auth, + ) + + assert skill.created_by == "__litellm_unscoped_proxy__" + assert ( + table.create.await_args.kwargs["data"]["updated_by"] + == "__litellm_unscoped_proxy__" + ) + + @pytest.mark.asyncio async def test_should_filter_list_skills_to_authenticated_owner_scopes(monkeypatch): table = AsyncMock()