From ad9aa43e86390f4e95bb133f86aad11e24c689b2 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 18:23:48 -0700 Subject: [PATCH 01/27] chore(proxy): scope skills and container resources --- litellm/llms/litellm_proxy/skills/handler.py | 60 +++- .../litellm_proxy/skills/transformation.py | 65 ++++- .../proxy/common_utils/resource_ownership.py | 75 +++++ .../proxy/container_endpoints/endpoints.py | 29 +- .../container_endpoints/handler_factory.py | 47 ++-- .../proxy/container_endpoints/ownership.py | 256 ++++++++++++++++++ litellm/proxy/hooks/litellm_skills/main.py | 14 +- litellm/skills/main.py | 29 +- .../test_container_proxy_ownership.py | 159 +++++++++++ .../litellm_proxy/test_skills_ownership.py | 118 ++++++++ 10 files changed, 797 insertions(+), 55 deletions(-) create mode 100644 litellm/proxy/common_utils/resource_ownership.py create mode 100644 litellm/proxy/container_endpoints/ownership.py create mode 100644 tests/test_litellm/containers/test_container_proxy_ownership.py create mode 100644 tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 8e5070c2724e..ddd5ab78c9d8 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -9,7 +9,13 @@ from typing import Any, Dict, List, Optional from litellm._logging import verbose_logger -from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest +from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth +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, +) def _prisma_skill_to_litellm(prisma_skill) -> LiteLLM_SkillsTable: @@ -58,6 +64,7 @@ async def _get_prisma_client(): async def create_skill( data: NewSkillRequest, user_id: Optional[str] = None, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> LiteLLM_SkillsTable: """ Create a new skill in the LiteLLM database. @@ -72,6 +79,7 @@ async def create_skill( prisma_client = await LiteLLMSkillsHandler._get_prisma_client() skill_id = f"litellm_skill_{uuid.uuid4()}" + owner = get_primary_resource_owner_scope(user_api_key_dict) or user_id skill_data: Dict[str, Any] = { "skill_id": skill_id, @@ -79,8 +87,8 @@ async def create_skill( "description": data.description, "instructions": data.instructions, "source": "custom", - "created_by": user_id, - "updated_by": user_id, + "created_by": owner, + "updated_by": owner, } # Handle metadata @@ -111,6 +119,7 @@ async def create_skill( async def list_skills( limit: int = 20, offset: int = 0, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> List[LiteLLM_SkillsTable]: """ List skills from the LiteLLM database. @@ -128,16 +137,28 @@ async def list_skills( f"LiteLLMSkillsHandler: Listing skills with limit={limit}, offset={offset}" ) + find_many_kwargs: Dict[str, Any] = { + "take": limit, + "skip": offset, + "order": {"created_at": "desc"}, + } + if user_api_key_dict is not None and not is_proxy_admin(user_api_key_dict): + owner_scopes = get_resource_owner_scopes(user_api_key_dict) + if not owner_scopes: + return [] + find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} + skills = await prisma_client.db.litellm_skillstable.find_many( - take=limit, - skip=offset, - order={"created_at": "desc"}, + **find_many_kwargs ) return [_prisma_skill_to_litellm(s) for s in skills] @staticmethod - async def get_skill(skill_id: str) -> LiteLLM_SkillsTable: + async def get_skill( + skill_id: str, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, + ) -> LiteLLM_SkillsTable: """ Get a skill by ID from the LiteLLM database. @@ -161,10 +182,18 @@ async def get_skill(skill_id: str) -> LiteLLM_SkillsTable: if skill is None: raise ValueError(f"Skill not found: {skill_id}") + if not user_can_access_resource_owner( + getattr(skill, "created_by", None), user_api_key_dict + ): + raise ValueError(f"Skill not found: {skill_id}") + return _prisma_skill_to_litellm(skill) @staticmethod - async def delete_skill(skill_id: str) -> Dict[str, str]: + async def delete_skill( + skill_id: str, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, + ) -> Dict[str, str]: """ Delete a skill by ID from the LiteLLM database. @@ -189,13 +218,21 @@ async def delete_skill(skill_id: str) -> Dict[str, str]: if skill is None: raise ValueError(f"Skill not found: {skill_id}") + if not user_can_access_resource_owner( + getattr(skill, "created_by", None), user_api_key_dict + ): + raise ValueError(f"Skill not found: {skill_id}") + # Delete the skill await prisma_client.db.litellm_skillstable.delete(where={"skill_id": skill_id}) return {"id": skill_id, "type": "skill_deleted"} @staticmethod - async def fetch_skill_from_db(skill_id: str) -> Optional[LiteLLM_SkillsTable]: + async def fetch_skill_from_db( + skill_id: str, + user_api_key_dict: Optional[UserAPIKeyAuth] = None, + ) -> Optional[LiteLLM_SkillsTable]: """ Fetch a skill from the database (used by skills injection hook). @@ -209,7 +246,10 @@ async def fetch_skill_from_db(skill_id: str) -> Optional[LiteLLM_SkillsTable]: LiteLLM_SkillsTable or None if not found """ try: - return await LiteLLMSkillsHandler.get_skill(skill_id) + return await LiteLLMSkillsHandler.get_skill( + skill_id, + user_api_key_dict=user_api_key_dict, + ) except ValueError: return None except Exception as e: diff --git a/litellm/llms/litellm_proxy/skills/transformation.py b/litellm/llms/litellm_proxy/skills/transformation.py index 4622bda4e804..199f13191fe3 100644 --- a/litellm/llms/litellm_proxy/skills/transformation.py +++ b/litellm/llms/litellm_proxy/skills/transformation.py @@ -18,6 +18,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth class LiteLLMSkillsTransformationHandler: @@ -44,6 +45,7 @@ def create_skill_handler( file_type: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, _is_async: bool = False, logging_obj: Optional["LiteLLMLoggingObj"] = None, litellm_call_id: Optional[str] = None, @@ -99,6 +101,7 @@ def create_skill_handler( file_type=file_type, metadata=metadata, user_id=user_id, + user_api_key_dict=user_api_key_dict, ) import asyncio @@ -113,6 +116,7 @@ def create_skill_handler( file_type=file_type, metadata=metadata, user_id=user_id, + user_api_key_dict=user_api_key_dict, ) ) @@ -126,6 +130,7 @@ async def _async_create_skill( file_type: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, ) -> Skill: """Async implementation of create_skill.""" # Lazy import to avoid SDK dependency on proxy @@ -145,6 +150,7 @@ async def _async_create_skill( db_skill = await LiteLLMSkillsHandler.create_skill( data=skill_request, user_id=user_id, + user_api_key_dict=user_api_key_dict, ) return self._db_skill_to_response(db_skill) @@ -156,6 +162,7 @@ def list_skills_handler( _is_async: bool = False, logging_obj: Optional["LiteLLMLoggingObj"] = None, litellm_call_id: Optional[str] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, **kwargs, ) -> Union[ListSkillsResponse, Coroutine[Any, Any, ListSkillsResponse]]: """ @@ -182,18 +189,27 @@ def list_skills_handler( ) if _is_async: - return self._async_list_skills(limit=limit, offset=offset) + return self._async_list_skills( + limit=limit, + offset=offset, + user_api_key_dict=user_api_key_dict, + ) import asyncio return asyncio.get_event_loop().run_until_complete( - self._async_list_skills(limit=limit, offset=offset) + self._async_list_skills( + limit=limit, + offset=offset, + user_api_key_dict=user_api_key_dict, + ) ) async def _async_list_skills( self, limit: int = 20, offset: int = 0, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, ) -> ListSkillsResponse: """Async implementation of list_skills.""" # Lazy import to avoid SDK dependency on proxy @@ -202,6 +218,7 @@ async def _async_list_skills( db_skills = await LiteLLMSkillsHandler.list_skills( limit=limit, offset=offset, + user_api_key_dict=user_api_key_dict, ) skills = [self._db_skill_to_response(s) for s in db_skills] @@ -217,6 +234,7 @@ def get_skill_handler( _is_async: bool = False, logging_obj: Optional["LiteLLMLoggingObj"] = None, litellm_call_id: Optional[str] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, **kwargs, ) -> Union[Skill, Coroutine[Any, Any, Skill]]: """ @@ -242,20 +260,33 @@ def get_skill_handler( ) if _is_async: - return self._async_get_skill(skill_id=skill_id) + return self._async_get_skill( + skill_id=skill_id, + user_api_key_dict=user_api_key_dict, + ) import asyncio return asyncio.get_event_loop().run_until_complete( - self._async_get_skill(skill_id=skill_id) + self._async_get_skill( + skill_id=skill_id, + user_api_key_dict=user_api_key_dict, + ) ) - async def _async_get_skill(self, skill_id: str) -> Skill: + async def _async_get_skill( + self, + skill_id: str, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + ) -> Skill: """Async implementation of get_skill.""" # Lazy import to avoid SDK dependency on proxy from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler - db_skill = await LiteLLMSkillsHandler.get_skill(skill_id=skill_id) + db_skill = await LiteLLMSkillsHandler.get_skill( + skill_id=skill_id, + user_api_key_dict=user_api_key_dict, + ) return self._db_skill_to_response(db_skill) def delete_skill_handler( @@ -264,6 +295,7 @@ def delete_skill_handler( _is_async: bool = False, logging_obj: Optional["LiteLLMLoggingObj"] = None, litellm_call_id: Optional[str] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, **kwargs, ) -> Union[DeleteSkillResponse, Coroutine[Any, Any, DeleteSkillResponse]]: """ @@ -289,20 +321,33 @@ def delete_skill_handler( ) if _is_async: - return self._async_delete_skill(skill_id=skill_id) + return self._async_delete_skill( + skill_id=skill_id, + user_api_key_dict=user_api_key_dict, + ) import asyncio return asyncio.get_event_loop().run_until_complete( - self._async_delete_skill(skill_id=skill_id) + self._async_delete_skill( + skill_id=skill_id, + user_api_key_dict=user_api_key_dict, + ) ) - async def _async_delete_skill(self, skill_id: str) -> DeleteSkillResponse: + async def _async_delete_skill( + self, + skill_id: str, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + ) -> DeleteSkillResponse: """Async implementation of delete_skill.""" # Lazy import to avoid SDK dependency on proxy from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler - result = await LiteLLMSkillsHandler.delete_skill(skill_id=skill_id) + result = await LiteLLMSkillsHandler.delete_skill( + skill_id=skill_id, + user_api_key_dict=user_api_key_dict, + ) return DeleteSkillResponse( id=result["id"], type=result.get("type", "skill_deleted"), diff --git a/litellm/proxy/common_utils/resource_ownership.py b/litellm/proxy/common_utils/resource_ownership.py new file mode 100644 index 000000000000..1bc769a8c053 --- /dev/null +++ b/litellm/proxy/common_utils/resource_ownership.py @@ -0,0 +1,75 @@ +from typing import List, Optional + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +def is_proxy_admin(user_api_key_dict: Optional[UserAPIKeyAuth]) -> bool: + if user_api_key_dict is None: + return False + + return ( + user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value + ) + + +def get_resource_owner_scopes( + user_api_key_dict: Optional[UserAPIKeyAuth], +) -> List[str]: + """ + Return ownership scopes that may access a user-created proxy resource. + + Raw user_id is included for rows created before scope prefixes existed. + Prefixes avoid collisions when falling back to team/org/key ownership for + keys that do not have a user_id. + """ + if user_api_key_dict is None: + return [] + + scopes: List[str] = [] + + def _add(scope: Optional[str]) -> None: + if scope and scope not in scopes: + scopes.append(scope) + + if user_api_key_dict.user_id: + _add(user_api_key_dict.user_id) + _add(f"user:{user_api_key_dict.user_id}") + if user_api_key_dict.team_id: + _add(f"team:{user_api_key_dict.team_id}") + if user_api_key_dict.org_id: + _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}") + + return scopes + + +def get_primary_resource_owner_scope( + user_api_key_dict: Optional[UserAPIKeyAuth], +) -> Optional[str]: + if user_api_key_dict is None: + return None + + if user_api_key_dict.user_id: + return user_api_key_dict.user_id + if user_api_key_dict.team_id: + return f"team:{user_api_key_dict.team_id}" + if user_api_key_dict.org_id: + 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 + + +def user_can_access_resource_owner( + owner: Optional[str], + user_api_key_dict: Optional[UserAPIKeyAuth], +) -> bool: + if user_api_key_dict is None: + return True + if is_proxy_admin(user_api_key_dict): + return True + if owner is None: + return False + return owner in get_resource_owner_scopes(user_api_key_dict) diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 45870d73d4b8..9d4dd9029925 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -14,6 +14,11 @@ get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) +from litellm.proxy.container_endpoints.ownership import ( + assert_user_can_access_container, + filter_container_list_response, + record_container_owner, +) router = APIRouter() @@ -98,7 +103,7 @@ async def create_container( # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -123,6 +128,11 @@ 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( @@ -191,7 +201,7 @@ async def list_containers( # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) try: - return await processor.base_process_llm_request( + response = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -216,6 +226,11 @@ 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( @@ -280,6 +295,11 @@ async def retrieve_container( ) # Add custom_llm_provider to data + _, custom_llm_provider = await assert_user_can_access_container( + container_id=container_id, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) data["custom_llm_provider"] = custom_llm_provider # Process request using ProxyBaseLLMRequestProcessing @@ -374,6 +394,11 @@ async def delete_container( ) # Add custom_llm_provider to data + _, custom_llm_provider = await assert_user_can_access_container( + container_id=container_id, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) data["custom_llm_provider"] = custom_llm_provider # Process request using ProxyBaseLLMRequestProcessing diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index fae7f939aed1..67c1d4b0f5ab 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -19,7 +19,9 @@ get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) -from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.proxy.container_endpoints.ownership import ( + assert_user_can_access_container, +) def _load_endpoints_config() -> Dict: @@ -176,14 +178,11 @@ async def _process_binary_request( # Build litellm_params - credentials are resolved by provider config from env litellm_params = GenericLiteLLMParams() - # Decode container ID and extract provider info - decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) - original_container_id = decoded.get("response_id", container_id) - - # If container ID has encoded provider info and user didn't explicitly set provider, use it - decoded_provider = decoded.get("custom_llm_provider") - if decoded_provider and custom_llm_provider == "openai": - custom_llm_provider = decoded_provider + original_container_id, custom_llm_provider = await assert_user_can_access_container( + container_id=container_id, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) # Get the provider config container_provider_config = _get_container_provider_config(custom_llm_provider) @@ -284,16 +283,13 @@ async def _process_multipart_upload_request( or "openai" ) - # Decode container ID and extract provider info - decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) - original_container_id = decoded.get("response_id", container_id) - - # If container ID has encoded provider info and user didn't explicitly set provider, use it - decoded_provider = decoded.get("custom_llm_provider") - if decoded_provider and custom_llm_provider == "openai": - custom_llm_provider = decoded_provider + _, custom_llm_provider = await assert_user_can_access_container( + container_id=container_id, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) - data["container_id"] = original_container_id # Use decoded original ID + data["container_id"] = container_id data["custom_llm_provider"] = custom_llm_provider processor = ProxyBaseLLMRequestProcessing(data=data) @@ -361,18 +357,11 @@ async def _process_request( # Decode container_id if present in path_params if "container_id" in path_params: - decoded = ResponsesAPIRequestUtils._decode_container_id( - path_params["container_id"] + _, custom_llm_provider = await assert_user_can_access_container( + container_id=path_params["container_id"], + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, ) - original_container_id = decoded.get("response_id", path_params["container_id"]) - - # If container ID has encoded provider info and user didn't explicitly set provider, use it - decoded_provider = decoded.get("custom_llm_provider") - if decoded_provider and custom_llm_provider == "openai": - custom_llm_provider = decoded_provider - - # Update path_params with decoded original ID - data["container_id"] = original_container_id data["custom_llm_provider"] = custom_llm_provider diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py new file mode 100644 index 000000000000..1a62abc9f36a --- /dev/null +++ b/litellm/proxy/container_endpoints/ownership.py @@ -0,0 +1,256 @@ +from typing import Any, Dict, List, Optional, Set, Tuple + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UserAPIKeyAuth +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, +) +from litellm.responses.utils import ResponsesAPIRequestUtils + +CONTAINER_OBJECT_PURPOSE = "container" + + +def _container_model_object_id( + original_container_id: str, + custom_llm_provider: str, +) -> str: + return f"{CONTAINER_OBJECT_PURPOSE}:{custom_llm_provider}:{original_container_id}" + + +def decode_container_id_for_ownership( + container_id: str, + custom_llm_provider: str, +) -> Tuple[str, str]: + decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) + original_container_id = decoded.get("response_id", container_id) + decoded_provider = decoded.get("custom_llm_provider") + if decoded_provider and custom_llm_provider == "openai": + custom_llm_provider = decoded_provider + return original_container_id, custom_llm_provider + + +def _get_response_id(response: Any) -> Optional[str]: + if response is None: + return None + if isinstance(response, dict): + value = response.get("id") + else: + value = getattr(response, "id", None) + return value if isinstance(value, str) else None + + +def _dump_response(response: Any) -> Dict[str, Any]: + if isinstance(response, dict): + return response + if hasattr(response, "model_dump"): + return response.model_dump() + if hasattr(response, "dict"): + return response.dict() + return {"id": _get_response_id(response)} + + +async def _get_prisma_client(): + from litellm.proxy.proxy_server import prisma_client + + return prisma_client + + +async def record_container_owner( + response: Any, + user_api_key_dict: UserAPIKeyAuth, + custom_llm_provider: str, +) -> 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 or prisma_client is None + ): + return response + if container_id is None or owner is None or prisma_client is None: + raise HTTPException(status_code=500, detail="Unable to track container") + + original_container_id, resolved_provider = decode_container_id_for_ownership( + container_id, + custom_llm_provider, + ) + model_object_id = _container_model_object_id( + original_container_id, + resolved_provider, + ) + file_object = _dump_response(response) + file_object["custom_llm_provider"] = resolved_provider + file_object["provider_container_id"] = original_container_id + + try: + existing = await prisma_client.db.litellm_managedobjecttable.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 + ): + raise HTTPException(status_code=403, detail="Forbidden") + await prisma_client.db.litellm_managedobjecttable.update( + where={"model_object_id": model_object_id}, + data={ + "unified_object_id": container_id, + "file_object": file_object, + "updated_by": owner, + }, + ) + else: + await prisma_client.db.litellm_managedobjecttable.create( + data={ + "unified_object_id": container_id, + "model_object_id": model_object_id, + "file_object": file_object, + "file_purpose": CONTAINER_OBJECT_PURPOSE, + "created_by": owner, + "updated_by": owner, + } + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.warning( + "Failed to record container ownership for container_id=%s: %s", + model_object_id, + e, + ) + raise HTTPException(status_code=500, detail="Unable to track container") + + return response + + +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 None + + 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, + } + ) + return getattr(row, "created_by", None) if row is not None else None + + +async def assert_user_can_access_container( + container_id: str, + user_api_key_dict: UserAPIKeyAuth, + custom_llm_provider: str, +) -> Tuple[str, str]: + original_container_id, resolved_provider = decode_container_id_for_ownership( + container_id, + custom_llm_provider, + ) + + if is_proxy_admin(user_api_key_dict): + return original_container_id, resolved_provider + + owner = await _get_container_owner(original_container_id, resolved_provider) + if not user_can_access_resource_owner(owner, user_api_key_dict): + raise HTTPException(status_code=403, detail="Forbidden") + + return original_container_id, resolved_provider + + +def _get_container_list_data(response: Any) -> Optional[List[Any]]: + if response is None: + return None + if isinstance(response, dict): + data = response.get("data") + else: + data = getattr(response, "data", None) + return data if isinstance(data, list) else None + + +def _set_container_list_data(response: Any, data: List[Any]) -> Any: + if isinstance(response, dict): + response["data"] = data + if data: + response["first_id"] = _get_response_id(data[0]) + response["last_id"] = _get_response_id(data[-1]) + else: + response["first_id"] = None + response["last_id"] = None + return response + + response.data = data + response.first_id = _get_response_id(data[0]) if data else None + response.last_id = _get_response_id(data[-1]) if data else None + return response + + +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: + return set() + + 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 + } + + +async def filter_container_list_response( + response: Any, + user_api_key_dict: UserAPIKeyAuth, + custom_llm_provider: str, +) -> Any: + if is_proxy_admin(user_api_key_dict): + return response + + data = _get_container_list_data(response) + if data is None: + return response + + allowed_container_ids = await _get_allowed_container_ids( + user_api_key_dict, + custom_llm_provider, + ) + filtered: List[Any] = [] + for item in data: + container_id = _get_response_id(item) + if container_id is None: + continue + original_container_id, resolved_provider = decode_container_id_for_ownership( + container_id, + custom_llm_provider, + ) + if ( + _container_model_object_id(original_container_id, resolved_provider) + in allowed_container_ids + ): + filtered.append(item) + + return _set_container_list_data(response, filtered) diff --git a/litellm/proxy/hooks/litellm_skills/main.py b/litellm/proxy/hooks/litellm_skills/main.py index 7c6bfbd6b2b5..21e8bbbd3085 100644 --- a/litellm/proxy/hooks/litellm_skills/main.py +++ b/litellm/proxy/hooks/litellm_skills/main.py @@ -110,7 +110,10 @@ async def async_pre_call_hook( skill_id = skill.get("skill_id", "") if skill_id.startswith("litellm_"): # Fetch from LiteLLM DB - db_skill = await self._fetch_skill_from_db(skill_id) + db_skill = await self._fetch_skill_from_db( + skill_id, + user_api_key_dict=user_api_key_dict, + ) if db_skill: litellm_skills.append(db_skill) else: @@ -276,7 +279,9 @@ def _process_non_anthropic_model( return data async def _fetch_skill_from_db( - self, skill_id: str + self, + skill_id: str, + user_api_key_dict: UserAPIKeyAuth, ) -> Optional[LiteLLM_SkillsTable]: """ Fetch a skill from the LiteLLM database. @@ -290,7 +295,10 @@ async def _fetch_skill_from_db( try: from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler - return await LiteLLMSkillsHandler.fetch_skill_from_db(skill_id) + return await LiteLLMSkillsHandler.fetch_skill_from_db( + skill_id, + user_api_key_dict=user_api_key_dict, + ) except Exception as e: verbose_proxy_logger.warning( f"SkillsInjectionHook: Error fetching skill {skill_id}: {e}" diff --git a/litellm/skills/main.py b/litellm/skills/main.py index 3ff0f52c6417..cee811d84fd5 100644 --- a/litellm/skills/main.py +++ b/litellm/skills/main.py @@ -34,6 +34,29 @@ _litellm_skills_handler = None +def _get_user_api_key_auth_from_kwargs(kwargs: Dict[str, Any]) -> Optional[Any]: + for metadata_key in ("metadata", "litellm_metadata"): + metadata = kwargs.get(metadata_key) + if isinstance(metadata, dict) and metadata.get("user_api_key_auth") is not None: + return metadata["user_api_key_auth"] + return None + + +def _get_skill_request_metadata( + kwargs: Dict[str, Any], + extra_body: Optional[Dict[str, Any]], +) -> Optional[Dict[str, Any]]: + if extra_body and isinstance(extra_body.get("metadata"), dict): + return extra_body["metadata"] + + metadata = kwargs.get("metadata") + if isinstance(metadata, dict) and isinstance( + metadata.get("requester_metadata"), dict + ): + return metadata["requester_metadata"] + return None + + def _get_litellm_skills_handler(): """Lazy initialization of LiteLLM skills handler to avoid import overhead.""" global _litellm_skills_handler @@ -165,8 +188,9 @@ def create_skill( return _get_litellm_skills_handler().create_skill_handler( display_title=display_title, files=files, - metadata=extra_body.get("metadata") if extra_body else None, + metadata=_get_skill_request_metadata(kwargs, extra_body), user_id=kwargs.get("user_id"), + user_api_key_dict=_get_user_api_key_auth_from_kwargs(kwargs), _is_async=_is_async, logging_obj=litellm_logging_obj, litellm_call_id=litellm_call_id, @@ -348,6 +372,7 @@ def list_skills( return _get_litellm_skills_handler().list_skills_handler( limit=limit or 20, offset=0, + user_api_key_dict=_get_user_api_key_auth_from_kwargs(kwargs), _is_async=_is_async, logging_obj=litellm_logging_obj, litellm_call_id=litellm_call_id, @@ -523,6 +548,7 @@ def get_skill( if custom_llm_provider == LlmProviders.LITELLM_PROXY.value: return _get_litellm_skills_handler().get_skill_handler( skill_id=skill_id, + user_api_key_dict=_get_user_api_key_auth_from_kwargs(kwargs), _is_async=_is_async, logging_obj=litellm_logging_obj, litellm_call_id=litellm_call_id, @@ -690,6 +716,7 @@ def delete_skill( if custom_llm_provider == LlmProviders.LITELLM_PROXY.value: return _get_litellm_skills_handler().delete_skill_handler( skill_id=skill_id, + user_api_key_dict=_get_user_api_key_auth_from_kwargs(kwargs), _is_async=_is_async, logging_obj=litellm_logging_obj, litellm_call_id=litellm_call_id, diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py new file mode 100644 index 000000000000..760e522199d9 --- /dev/null +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -0,0 +1,159 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.container_endpoints import ownership +from litellm.types.containers.main import ContainerListResponse, ContainerObject + + +def _container(container_id: str) -> ContainerObject: + return ContainerObject( + id=container_id, + object="container", + created_at=1, + status="active", + ) + + +@pytest.mark.asyncio +async def test_should_record_container_owner_with_original_provider_id(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(user_id="user-1") + + response = _container("cntr_provider") + + await ownership.record_container_owner( + response=response, + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + table.create.assert_awaited_once() + data = table.create.await_args.kwargs["data"] + assert data["model_object_id"] == "container:openai:cntr_provider" + assert data["file_purpose"] == ownership.CONTAINER_OBJECT_PURPOSE + assert data["created_by"] == "user-1" + + +@pytest.mark.asyncio +async def test_should_record_team_owner_for_keys_without_user_id(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(team_id="team-1") + + 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"] == "team:team-1" + assert data["updated_by"] == "team:team-1" + + +@pytest.mark.asyncio +async def test_should_deny_container_access_for_different_owner(monkeypatch): + table = AsyncMock() + table.find_first.return_value = SimpleNamespace(created_by="user-2") + 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_provider", + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert exc.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_should_not_reassign_existing_container_to_different_owner(monkeypatch): + table = AsyncMock() + table.find_unique.return_value = SimpleNamespace( + file_purpose=ownership.CONTAINER_OBJECT_PURPOSE, + created_by="user-2", + ) + 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.record_container_owner( + response=_container("cntr_existing"), + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert exc.value.status_code == 403 + table.update.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_should_filter_container_list_to_owned_records(monkeypatch): + table = AsyncMock() + table.find_many.return_value = [ + SimpleNamespace(model_object_id="container:openai:cntr_owned"), + ] + 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") + response = ContainerListResponse( + object="list", + data=[_container("cntr_owned"), _container("cntr_other")], + has_more=False, + ) + + 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.first_id == "cntr_owned" + assert filtered.last_id == "cntr_owned" + where = table.find_many.await_args.kwargs["where"] + assert where["file_purpose"] == ownership.CONTAINER_OBJECT_PURPOSE + assert where["created_by"]["in"] == ["user-1", "user:user-1"] diff --git a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py new file mode 100644 index 000000000000..d5151b47a6f1 --- /dev/null +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -0,0 +1,118 @@ +from unittest.mock import AsyncMock + +import pytest + +from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler +from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth + + +def _skill(skill_id: str, created_by: str | None) -> LiteLLM_SkillsTable: + return LiteLLM_SkillsTable( + skill_id=skill_id, + display_title="skill", + created_by=created_by, + ) + + +@pytest.mark.asyncio +async def test_should_store_team_owner_for_keys_without_user_id(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(team_id="team-1") + + skill = await LiteLLMSkillsHandler.create_skill( + data=NewSkillRequest(display_title="skill"), + user_api_key_dict=auth, + ) + + assert skill.created_by == "team:team-1" + assert table.create.await_args.kwargs["data"]["updated_by"] == "team:team-1" + + +@pytest.mark.asyncio +async def test_should_filter_list_skills_to_authenticated_owner_scopes(monkeypatch): + table = AsyncMock() + table.find_many.return_value = [_skill("litellm_skill_owner", "user-1")] + prisma_client = type( + "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} + )() + monkeypatch.setattr( + LiteLLMSkillsHandler, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + + auth = UserAPIKeyAuth(user_id="user-1", team_id="team-1") + + skills = await LiteLLMSkillsHandler.list_skills(user_api_key_dict=auth) + + assert [skill.skill_id for skill in skills] == ["litellm_skill_owner"] + table.find_many.assert_awaited_once() + where = table.find_many.await_args.kwargs["where"] + assert where["created_by"]["in"] == [ + "user-1", + "user:user-1", + "team:team-1", + ] + + +@pytest.mark.asyncio +async def test_should_hide_skill_from_different_owner(monkeypatch): + table = AsyncMock() + table.find_unique.return_value = _skill("litellm_skill_other", "user-2") + prisma_client = type( + "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} + )() + monkeypatch.setattr( + LiteLLMSkillsHandler, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + + auth = UserAPIKeyAuth(user_id="user-1") + + with pytest.raises(ValueError, match="Skill not found"): + await LiteLLMSkillsHandler.get_skill( + "litellm_skill_other", + user_api_key_dict=auth, + ) + + +@pytest.mark.asyncio +async def test_should_scope_skill_injection_fetch_to_authenticated_user(monkeypatch): + from litellm.proxy.hooks.litellm_skills.main import SkillsInjectionHook + + fetch = AsyncMock(return_value=None) + monkeypatch.setattr(LiteLLMSkillsHandler, "fetch_skill_from_db", fetch) + + auth = UserAPIKeyAuth(user_id="user-1") + hook = SkillsInjectionHook() + data = { + "container": { + "skills": [ + {"skill_id": "litellm_skill_other"}, + ] + } + } + + response = await hook.async_pre_call_hook( + user_api_key_dict=auth, + cache=AsyncMock(), + data=data, + call_type="completion", + ) + + assert response == data + fetch.assert_awaited_once_with( + "litellm_skill_other", + user_api_key_dict=auth, + ) From 95c8406b587c4ea0125dc177050ba3460b3e3d43 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 18:34:12 -0700 Subject: [PATCH 02/27] chore(proxy): stabilize openapi snapshot ids --- litellm/proxy/_lazy_openapi_snapshot.json | 34 +++++++++++------------ litellm/proxy/proxy_server.py | 9 ++++++ 2 files changed, 26 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 8331f748c6ea..b8e9eb6c261e 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3572,7 +3572,7 @@ "/anthropic/{endpoint}": { "delete": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -3616,7 +3616,7 @@ }, "get": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -3660,7 +3660,7 @@ }, "patch": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -3704,7 +3704,7 @@ }, "post": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -3748,7 +3748,7 @@ }, "put": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -13260,7 +13260,7 @@ "/langfuse/{endpoint}": { "delete": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -13299,7 +13299,7 @@ }, "get": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -13338,7 +13338,7 @@ }, "patch": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -13377,7 +13377,7 @@ }, "post": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -13416,7 +13416,7 @@ }, "put": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -26883,7 +26883,7 @@ "/toolset/{toolset_name}/mcp": { "delete": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -26922,7 +26922,7 @@ }, "get": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -26961,7 +26961,7 @@ }, "head": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -27000,7 +27000,7 @@ }, "options": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -27039,7 +27039,7 @@ }, "patch": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -27078,7 +27078,7 @@ }, "post": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -27117,7 +27117,7 @@ }, "put": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6cba6a3e96b7..b6cec92a1d11 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6,6 +6,7 @@ import io import os import random +import re import secrets import shutil import subprocess @@ -950,6 +951,13 @@ async def _run_pw_migration(): await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] +def _generate_stable_operation_id(route: Any) -> str: + operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}") + if route.methods: + operation_id = f"{operation_id}_{sorted(route.methods)[0].lower()}" + return operation_id + + app = FastAPI( docs_url=_get_docs_url(), redoc_url=_get_redoc_url(), @@ -959,6 +967,7 @@ async def _run_pw_migration(): version=version, root_path=server_root_path, lifespan=proxy_startup_event, # type: ignore[reportGeneralTypeIssues] + generate_unique_id_function=_generate_stable_operation_id, ) vertex_live_passthrough_vertex_base = VertexBase() From d2a322074faf7619c05642ba2ec89489795af4ce Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 18:40:26 -0700 Subject: [PATCH 03/27] fix(proxy): preserve container ownership compatibility --- .../proxy/container_endpoints/endpoints.py | 10 +- .../container_endpoints/handler_factory.py | 12 +- .../proxy/container_endpoints/ownership.py | 47 +++- .../test_container_proxy_ownership.py | 234 ++++++++++++++++++ 4 files changed, 292 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 9d4dd9029925..93587262f588 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -295,11 +295,14 @@ async def retrieve_container( ) # Add custom_llm_provider to data - _, custom_llm_provider = await assert_user_can_access_container( + container_access = await assert_user_can_access_container( container_id=container_id, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) + custom_llm_provider = container_access[1] + # Keep the managed id in request data so downstream container utilities can + # preserve encoded routing metadata while decoding before the provider call. data["custom_llm_provider"] = custom_llm_provider # Process request using ProxyBaseLLMRequestProcessing @@ -394,11 +397,14 @@ async def delete_container( ) # Add custom_llm_provider to data - _, custom_llm_provider = await assert_user_can_access_container( + container_access = await assert_user_can_access_container( container_id=container_id, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) + custom_llm_provider = container_access[1] + # Keep the managed id in request data so downstream container utilities can + # preserve encoded routing metadata while decoding before the provider call. data["custom_llm_provider"] = custom_llm_provider # Process request using ProxyBaseLLMRequestProcessing diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index 67c1d4b0f5ab..bb6f8a8db6f8 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -283,12 +283,16 @@ async def _process_multipart_upload_request( or "openai" ) - _, custom_llm_provider = await assert_user_can_access_container( + container_access = await assert_user_can_access_container( container_id=container_id, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) + custom_llm_provider = container_access[1] + # Keep the managed container id in the forwarded request. The container API + # layer decodes it before the upstream provider call and uses embedded + # routing metadata to preserve model/deployment affinity. data["container_id"] = container_id data["custom_llm_provider"] = custom_llm_provider @@ -355,13 +359,15 @@ async def _process_request( or "openai" ) - # Decode container_id if present in path_params + # Validate container_id ownership if present in path_params. if "container_id" in path_params: - _, custom_llm_provider = await assert_user_can_access_container( + container_access = await assert_user_can_access_container( container_id=path_params["container_id"], user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) + custom_llm_provider = container_access[1] + # Preserve the managed id for downstream container decoding/routing. data["custom_llm_provider"] = custom_llm_provider diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 1a62abc9f36a..e3d8da8b0def 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,3 +1,4 @@ +import os from typing import Any, Dict, List, Optional, Set, Tuple from fastapi import HTTPException @@ -13,6 +14,16 @@ from litellm.responses.utils import ResponsesAPIRequestUtils CONTAINER_OBJECT_PURPOSE = "container" +ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV = "LITELLM_ALLOW_UNTRACKED_CONTAINER_ACCESS" +_IN_MEMORY_CONTAINER_OWNERS: Dict[str, str] = {} + + +def _allow_untracked_container_access() -> bool: + return os.getenv(ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, "").lower() in { + "1", + "true", + "yes", + } def _container_model_object_id( @@ -68,11 +79,9 @@ async def record_container_owner( 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 or prisma_client is None - ): + 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 or prisma_client is None: + if container_id is None or owner is None: raise HTTPException(status_code=500, detail="Unable to track container") original_container_id, resolved_provider = decode_container_id_for_ownership( @@ -87,6 +96,15 @@ 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: existing = await prisma_client.db.litellm_managedobjecttable.find_unique( where={"model_object_id": model_object_id} @@ -136,7 +154,12 @@ async def _get_container_owner( ) -> Optional[str]: prisma_client = await _get_prisma_client() if prisma_client is None: - return 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={ @@ -164,6 +187,13 @@ async def assert_user_can_access_container( return original_container_id, resolved_provider owner = await _get_container_owner(original_container_id, resolved_provider) + if owner is None and _allow_untracked_container_access(): + verbose_proxy_logger.warning( + "Allowing untracked container access because %s is enabled", + ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, + ) + return original_container_id, resolved_provider + if not user_can_access_resource_owner(owner, user_api_key_dict): raise HTTPException(status_code=403, detail="Forbidden") @@ -203,7 +233,12 @@ async def _get_allowed_container_ids( ) -> Set[str]: prisma_client = await _get_prisma_client() if prisma_client is None: - return set() + 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: diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 760e522199d9..aed84b4fbe6e 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -1,3 +1,4 @@ +import sys from types import SimpleNamespace from unittest.mock import AsyncMock @@ -6,9 +7,18 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.container_endpoints import ownership +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.containers.main import ContainerListResponse, ContainerObject +@pytest.fixture(autouse=True) +def clear_in_memory_container_owners(monkeypatch): + ownership._IN_MEMORY_CONTAINER_OWNERS.clear() + monkeypatch.delenv(ownership.ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, raising=False) + yield + ownership._IN_MEMORY_CONTAINER_OWNERS.clear() + + def _container(container_id: str) -> ContainerObject: return ContainerObject( id=container_id, @@ -72,6 +82,31 @@ 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_track_container_owner_in_memory_without_prisma(monkeypatch): + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=None), + ) + auth = UserAPIKeyAuth(user_id="user-1") + + await ownership.record_container_owner( + response=_container("cntr_provider"), + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + 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_deny_container_access_for_different_owner(monkeypatch): table = AsyncMock() @@ -96,6 +131,45 @@ async def test_should_deny_container_access_for_different_owner(monkeypatch): assert exc.value.status_code == 403 +@pytest.mark.asyncio +async def test_should_deny_untracked_container_access_by_default(monkeypatch): + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=None), + ) + auth = UserAPIKeyAuth(user_id="user-1") + + with pytest.raises(HTTPException) as exc: + await ownership.assert_user_can_access_container( + container_id="cntr_untracked", + 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( + ownership, + "_get_prisma_client", + AsyncMock(return_value=None), + ) + monkeypatch.setenv(ownership.ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, "true") + auth = UserAPIKeyAuth(user_id="user-1") + + original_id, provider = await ownership.assert_user_can_access_container( + container_id="cntr_untracked", + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert original_id == "cntr_untracked" + assert provider == "openai" + + @pytest.mark.asyncio async def test_should_not_reassign_existing_container_to_different_owner(monkeypatch): table = AsyncMock() @@ -157,3 +231,163 @@ async def test_should_filter_container_list_to_owned_records(monkeypatch): where = table.find_many.await_args.kwargs["where"] assert where["file_purpose"] == ownership.CONTAINER_OBJECT_PURPOSE assert where["created_by"]["in"] == ["user-1", "user:user-1"] + + +@pytest.mark.asyncio +async def test_should_filter_container_list_with_in_memory_ownership(monkeypatch): + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=None), + ) + auth = UserAPIKeyAuth(user_id="user-1") + + await ownership.record_container_owner( + response=_container("cntr_owned"), + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + response = ContainerListResponse( + object="list", + data=[_container("cntr_owned"), _container("cntr_other")], + has_more=False, + ) + + 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"] + + +@pytest.mark.asyncio +async def test_should_preserve_managed_container_id_for_proxy_forwarding(monkeypatch): + from litellm.proxy.container_endpoints import handler_factory + + proxy_server_stub = SimpleNamespace( + general_settings={}, + llm_router=None, + proxy_config=None, + proxy_logging_obj=None, + select_data_generator=None, + user_api_base=None, + user_max_tokens=None, + user_model=None, + user_request_timeout=None, + user_temperature=None, + version="test", + ) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub) + + captured = {} + + class FakeProcessor: + def __init__(self, data): + captured["data"] = data + + async def base_process_llm_request(self, **kwargs): + return captured["data"] + + async def _handle_llm_api_exception(self, **kwargs): + raise kwargs["e"] + + monkeypatch.setattr( + handler_factory, + "ProxyBaseLLMRequestProcessing", + FakeProcessor, + ) + monkeypatch.setattr( + handler_factory, + "assert_user_can_access_container", + AsyncMock(return_value=("cntr_provider", "azure")), + ) + encoded_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="router-gpt", + container_id="cntr_provider", + ) + + result = await handler_factory._process_request( + request=SimpleNamespace(query_params={}, headers={}), + fastapi_response=SimpleNamespace(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + route_type="alist_container_files", + path_params={"container_id": encoded_id}, + ) + + assert result["container_id"] == encoded_id + assert result["custom_llm_provider"] == "azure" + + +@pytest.mark.asyncio +async def test_should_preserve_managed_container_id_for_multipart_upload(monkeypatch): + from litellm.proxy.common_utils import http_parsing_utils + from litellm.proxy.container_endpoints import handler_factory + + proxy_server_stub = SimpleNamespace( + general_settings={}, + llm_router=None, + proxy_config=None, + proxy_logging_obj=None, + select_data_generator=None, + user_api_base=None, + user_max_tokens=None, + user_model=None, + user_request_timeout=None, + user_temperature=None, + version="test", + ) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub) + + captured = {} + + class FakeProcessor: + def __init__(self, data): + captured["data"] = data + + async def base_process_llm_request(self, **kwargs): + return captured["data"] + + async def _handle_llm_api_exception(self, **kwargs): + raise kwargs["e"] + + monkeypatch.setattr( + handler_factory, + "ProxyBaseLLMRequestProcessing", + FakeProcessor, + ) + monkeypatch.setattr( + handler_factory, + "assert_user_can_access_container", + AsyncMock(return_value=("cntr_provider", "azure")), + ) + monkeypatch.setattr( + http_parsing_utils, + "get_form_data", + AsyncMock(return_value={}), + ) + monkeypatch.setattr( + http_parsing_utils, + "convert_upload_files_to_file_data", + AsyncMock(return_value={"file": ["file-data"]}), + ) + encoded_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="router-gpt", + container_id="cntr_provider", + ) + + result = await handler_factory._process_multipart_upload_request( + request=SimpleNamespace(query_params={}, headers={}), + fastapi_response=SimpleNamespace(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + route_type="aupload_container_file", + container_id=encoded_id, + ) + + assert result["container_id"] == encoded_id + assert result["custom_llm_provider"] == "azure" + assert result["file"] == "file-data" From 0745b1872eb08e07091b82bcaa9a25f4ded0d1ce Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 18:46:14 -0700 Subject: [PATCH 04/27] test(proxy): cover container endpoint forwarding --- .../test_container_proxy_ownership.py | 108 ++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index aed84b4fbe6e..d0b755645da1 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -391,3 +391,111 @@ async def _handle_llm_api_exception(self, **kwargs): assert result["container_id"] == encoded_id assert result["custom_llm_provider"] == "azure" assert result["file"] == "file-data" + + +@pytest.mark.asyncio +async def test_should_preserve_managed_container_id_for_proxy_retrieve(monkeypatch): + from litellm.proxy.container_endpoints import endpoints + + proxy_server_stub = SimpleNamespace( + general_settings={}, + llm_router=None, + proxy_config=None, + proxy_logging_obj=None, + select_data_generator=None, + user_api_base=None, + user_max_tokens=None, + user_model=None, + user_request_timeout=None, + user_temperature=None, + version="test", + ) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub) + + captured = {} + + class FakeProcessor: + def __init__(self, data): + captured["data"] = data + + async def base_process_llm_request(self, **kwargs): + return captured["data"] + + async def _handle_llm_api_exception(self, **kwargs): + raise kwargs["e"] + + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", FakeProcessor) + monkeypatch.setattr( + endpoints, + "assert_user_can_access_container", + AsyncMock(return_value=("cntr_provider", "azure")), + ) + encoded_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="router-gpt", + container_id="cntr_provider", + ) + + result = await endpoints.retrieve_container( + request=SimpleNamespace(query_params={}, headers={}), + container_id=encoded_id, + fastapi_response=SimpleNamespace(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + ) + + assert result["container_id"] == encoded_id + assert result["custom_llm_provider"] == "azure" + + +@pytest.mark.asyncio +async def test_should_preserve_managed_container_id_for_proxy_delete(monkeypatch): + from litellm.proxy.container_endpoints import endpoints + + proxy_server_stub = SimpleNamespace( + general_settings={}, + llm_router=None, + proxy_config=None, + proxy_logging_obj=None, + select_data_generator=None, + user_api_base=None, + user_max_tokens=None, + user_model=None, + user_request_timeout=None, + user_temperature=None, + version="test", + ) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub) + + captured = {} + + class FakeProcessor: + def __init__(self, data): + captured["data"] = data + + async def base_process_llm_request(self, **kwargs): + return captured["data"] + + async def _handle_llm_api_exception(self, **kwargs): + raise kwargs["e"] + + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", FakeProcessor) + monkeypatch.setattr( + endpoints, + "assert_user_can_access_container", + AsyncMock(return_value=("cntr_provider", "azure")), + ) + encoded_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="router-gpt", + container_id="cntr_provider", + ) + + result = await endpoints.delete_container( + request=SimpleNamespace(query_params={}, headers={}), + container_id=encoded_id, + fastapi_response=SimpleNamespace(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + ) + + assert result["container_id"] == encoded_id + assert result["custom_llm_provider"] == "azure" From e4741fbb2b985df2155e5d97061cc5c33a33d66e Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 18:53:44 -0700 Subject: [PATCH 05/27] fix(proxy): add legacy skill ownership opt-out --- litellm/llms/litellm_proxy/skills/handler.py | 41 +++++++++- .../litellm_proxy/test_skills_ownership.py | 78 +++++++++++++++++++ 2 files changed, 116 insertions(+), 3 deletions(-) diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index ddd5ab78c9d8..7093874acd93 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -5,6 +5,7 @@ Used by the transformation layer and skills injection hook. """ +import os import uuid from typing import Any, Dict, List, Optional @@ -17,6 +18,32 @@ user_can_access_resource_owner, ) +ALLOW_UNOWNED_SKILL_ACCESS_ENV = "LITELLM_ALLOW_UNOWNED_SKILL_ACCESS" + + +def _allow_unowned_skill_access() -> bool: + return os.getenv(ALLOW_UNOWNED_SKILL_ACCESS_ENV, "").lower() in { + "1", + "true", + "yes", + } + + +def _user_can_access_skill_owner( + owner: Optional[str], + user_api_key_dict: Optional[UserAPIKeyAuth], +) -> bool: + if owner is None and user_api_key_dict is not None: + if is_proxy_admin(user_api_key_dict): + return True + if _allow_unowned_skill_access(): + verbose_logger.warning( + "Allowing unowned skill access because %s is enabled", + ALLOW_UNOWNED_SKILL_ACCESS_ENV, + ) + return True + return user_can_access_resource_owner(owner, user_api_key_dict) + def _prisma_skill_to_litellm(prisma_skill) -> LiteLLM_SkillsTable: """ @@ -146,7 +173,15 @@ async def list_skills( owner_scopes = get_resource_owner_scopes(user_api_key_dict) if not owner_scopes: return [] - find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} + if _allow_unowned_skill_access(): + find_many_kwargs["where"] = { + "OR": [ + {"created_by": {"in": owner_scopes}}, + {"created_by": None}, + ] + } + else: + find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} skills = await prisma_client.db.litellm_skillstable.find_many( **find_many_kwargs @@ -182,7 +217,7 @@ async def get_skill( if skill is None: raise ValueError(f"Skill not found: {skill_id}") - if not user_can_access_resource_owner( + if not _user_can_access_skill_owner( getattr(skill, "created_by", None), user_api_key_dict ): raise ValueError(f"Skill not found: {skill_id}") @@ -218,7 +253,7 @@ async def delete_skill( if skill is None: raise ValueError(f"Skill not found: {skill_id}") - if not user_can_access_resource_owner( + if not _user_can_access_skill_owner( getattr(skill, "created_by", None), user_api_key_dict ): raise ValueError(f"Skill not found: {skill_id}") 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 d5151b47a6f1..fde1806f92cb 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -3,9 +3,15 @@ import pytest from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler +from litellm.llms.litellm_proxy.skills import handler as skills_handler from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth +@pytest.fixture(autouse=True) +def clear_skill_ownership_env(monkeypatch): + monkeypatch.delenv(skills_handler.ALLOW_UNOWNED_SKILL_ACCESS_ENV, raising=False) + + def _skill(skill_id: str, created_by: str | None) -> LiteLLM_SkillsTable: return LiteLLM_SkillsTable( skill_id=skill_id, @@ -87,6 +93,78 @@ async def test_should_hide_skill_from_different_owner(monkeypatch): ) +@pytest.mark.asyncio +async def test_should_hide_unowned_skill_by_default(monkeypatch): + table = AsyncMock() + table.find_unique.return_value = _skill("litellm_skill_unowned", None) + prisma_client = type( + "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} + )() + monkeypatch.setattr( + LiteLLMSkillsHandler, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + + auth = UserAPIKeyAuth(user_id="user-1") + + with pytest.raises(ValueError, match="Skill not found"): + await LiteLLMSkillsHandler.get_skill( + "litellm_skill_unowned", + user_api_key_dict=auth, + ) + + +@pytest.mark.asyncio +async def test_should_allow_unowned_skill_when_enabled(monkeypatch): + table = AsyncMock() + table.find_unique.return_value = _skill("litellm_skill_unowned", None) + prisma_client = type( + "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} + )() + monkeypatch.setattr( + LiteLLMSkillsHandler, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + monkeypatch.setenv(skills_handler.ALLOW_UNOWNED_SKILL_ACCESS_ENV, "true") + + auth = UserAPIKeyAuth(user_id="user-1") + + skill = await LiteLLMSkillsHandler.get_skill( + "litellm_skill_unowned", + user_api_key_dict=auth, + ) + + assert skill.skill_id == "litellm_skill_unowned" + + +@pytest.mark.asyncio +async def test_should_include_unowned_skills_in_list_when_enabled(monkeypatch): + table = AsyncMock() + table.find_many.return_value = [_skill("litellm_skill_unowned", None)] + prisma_client = type( + "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} + )() + monkeypatch.setattr( + LiteLLMSkillsHandler, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + monkeypatch.setenv(skills_handler.ALLOW_UNOWNED_SKILL_ACCESS_ENV, "true") + + auth = UserAPIKeyAuth(user_id="user-1") + + skills = await LiteLLMSkillsHandler.list_skills(user_api_key_dict=auth) + + assert [skill.skill_id for skill in skills] == ["litellm_skill_unowned"] + where = table.find_many.await_args.kwargs["where"] + assert where["OR"] == [ + {"created_by": {"in": ["user-1", "user:user-1"]}}, + {"created_by": None}, + ] + + @pytest.mark.asyncio async def test_should_scope_skill_injection_fetch_to_authenticated_user(monkeypatch): from litellm.proxy.hooks.litellm_skills.main import SkillsInjectionHook From 9376b30bca95f146adfb2af0e89760aa0496d735 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 19:01:37 -0700 Subject: [PATCH 06/27] fix(proxy): reset filtered container pagination --- .../proxy/container_endpoints/ownership.py | 3 + .../test_container_proxy_ownership.py | 70 +++++++++++++++++++ 2 files changed, 73 insertions(+) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index e3d8da8b0def..2260d67b6296 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -219,11 +219,14 @@ def _set_container_list_data(response: Any, data: List[Any]) -> Any: else: response["first_id"] = None response["last_id"] = None + response["has_more"] = False return response response.data = data response.first_id = _get_response_id(data[0]) if data else None response.last_id = _get_response_id(data[-1]) if data else None + if not data and hasattr(response, "has_more"): + response.has_more = False return response diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index d0b755645da1..40397a188606 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -233,6 +233,76 @@ async def test_should_filter_container_list_to_owned_records(monkeypatch): assert where["created_by"]["in"] == ["user-1", "user:user-1"] +@pytest.mark.asyncio +async def test_should_clear_has_more_when_filtered_container_list_is_empty( + monkeypatch, +): + table = AsyncMock() + table.find_many.return_value = [ + SimpleNamespace(model_object_id="container:openai:cntr_owned"), + ] + 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") + response = ContainerListResponse( + object="list", + data=[_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 filtered.data == [] + assert filtered.first_id is None + assert filtered.last_id is None + assert filtered.has_more is False + + +@pytest.mark.asyncio +async def test_should_clear_dict_has_more_when_filtered_container_list_is_empty( + monkeypatch, +): + table = AsyncMock() + table.find_many.return_value = [ + SimpleNamespace(model_object_id="container:openai:cntr_owned"), + ] + 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") + response = { + "object": "list", + "data": [{"id": "cntr_other"}], + "has_more": True, + } + + filtered = await ownership.filter_container_list_response( + response=response, + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert filtered["data"] == [] + assert filtered["first_id"] is None + assert filtered["last_id"] is None + assert filtered["has_more"] is False + + @pytest.mark.asyncio async def test_should_filter_container_list_with_in_memory_ownership(monkeypatch): monkeypatch.setattr( From 6aac4552f930fd9f7fc94f3bd22bc0a8de9f9fc8 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 19:13:20 -0700 Subject: [PATCH 07/27] fix(proxy): forward decoded container ids --- .../proxy/container_endpoints/endpoints.py | 27 ++++++---- .../container_endpoints/handler_factory.py | 46 +++++++++++------ .../proxy/container_endpoints/ownership.py | 35 +++++++++++-- .../test_container_proxy_ownership.py | 51 +++++++++++++++---- 4 files changed, 121 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 93587262f588..72a7c6c8746c 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -17,6 +17,7 @@ from litellm.proxy.container_endpoints.ownership import ( assert_user_can_access_container, filter_container_list_response, + get_container_forwarding_params, record_container_owner, ) @@ -295,15 +296,18 @@ async def retrieve_container( ) # Add custom_llm_provider to data - container_access = await assert_user_can_access_container( + original_container_id, custom_llm_provider = await assert_user_can_access_container( container_id=container_id, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) - custom_llm_provider = container_access[1] - # Keep the managed id in request data so downstream container utilities can - # preserve encoded routing metadata while decoding before the provider call. - data["custom_llm_provider"] = custom_llm_provider + data.update( + get_container_forwarding_params( + container_id, + original_container_id, + custom_llm_provider, + ) + ) # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) @@ -397,15 +401,18 @@ async def delete_container( ) # Add custom_llm_provider to data - container_access = await assert_user_can_access_container( + original_container_id, custom_llm_provider = await assert_user_can_access_container( container_id=container_id, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) - custom_llm_provider = container_access[1] - # Keep the managed id in request data so downstream container utilities can - # preserve encoded routing metadata while decoding before the provider call. - data["custom_llm_provider"] = custom_llm_provider + data.update( + get_container_forwarding_params( + container_id, + original_container_id, + custom_llm_provider, + ) + ) # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index bb6f8a8db6f8..ee79ebd96edc 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -21,6 +21,7 @@ ) from litellm.proxy.container_endpoints.ownership import ( assert_user_can_access_container, + get_container_forwarding_params, ) @@ -183,6 +184,13 @@ async def _process_binary_request( user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) + forwarding_params = get_container_forwarding_params( + container_id, + original_container_id, + custom_llm_provider, + ) + if "model_id" in forwarding_params: + litellm_params["model_id"] = forwarding_params["model_id"] # Get the provider config container_provider_config = _get_container_provider_config(custom_llm_provider) @@ -283,18 +291,19 @@ async def _process_multipart_upload_request( or "openai" ) - container_access = await assert_user_can_access_container( + original_container_id, custom_llm_provider = await assert_user_can_access_container( container_id=container_id, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) - custom_llm_provider = container_access[1] - # Keep the managed container id in the forwarded request. The container API - # layer decodes it before the upstream provider call and uses embedded - # routing metadata to preserve model/deployment affinity. - data["container_id"] = container_id - data["custom_llm_provider"] = custom_llm_provider + data.update( + get_container_forwarding_params( + container_id, + original_container_id, + custom_llm_provider, + ) + ) processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -361,15 +370,22 @@ async def _process_request( # Validate container_id ownership if present in path_params. if "container_id" in path_params: - container_access = await assert_user_can_access_container( - container_id=path_params["container_id"], - user_api_key_dict=user_api_key_dict, - custom_llm_provider=custom_llm_provider, + original_container_id, custom_llm_provider = ( + await assert_user_can_access_container( + container_id=path_params["container_id"], + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) ) - custom_llm_provider = container_access[1] - # Preserve the managed id for downstream container decoding/routing. - - data["custom_llm_provider"] = custom_llm_provider + data.update( + get_container_forwarding_params( + path_params["container_id"], + original_container_id, + custom_llm_provider, + ) + ) + else: + data["custom_llm_provider"] = custom_llm_provider processor = ProxyBaseLLMRequestProcessing(data=data) try: diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 2260d67b6296..b769b2c09538 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -45,6 +45,22 @@ def decode_container_id_for_ownership( return original_container_id, custom_llm_provider +def get_container_forwarding_params( + container_id: str, + original_container_id: str, + custom_llm_provider: str, +) -> Dict[str, str]: + params = { + "container_id": original_container_id, + "custom_llm_provider": custom_llm_provider, + } + decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) + model_id = decoded.get("model_id") + if isinstance(model_id, str) and model_id: + params["model_id"] = model_id + return params + + def _get_response_id(response: Any) -> Optional[str]: if response is None: return None @@ -139,11 +155,12 @@ async def record_container_owner( raise except Exception as e: verbose_proxy_logger.warning( - "Failed to record container ownership for container_id=%s: %s", + "Failed to persist container ownership for container_id=%s; " + "falling back to in-process tracking: %s", model_object_id, e, ) - raise HTTPException(status_code=500, detail="Unable to track container") + _IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner return response @@ -210,7 +227,9 @@ def _get_container_list_data(response: Any) -> Optional[List[Any]]: return data if isinstance(data, list) else None -def _set_container_list_data(response: Any, data: List[Any]) -> Any: +def _set_container_list_data( + response: Any, data: List[Any], removed_filtered_items: bool = False +) -> Any: if isinstance(response, dict): response["data"] = data if data: @@ -220,6 +239,8 @@ def _set_container_list_data(response: Any, data: List[Any]) -> Any: response["first_id"] = None response["last_id"] = None response["has_more"] = False + if removed_filtered_items: + response["has_more"] = False return response response.data = data @@ -227,6 +248,8 @@ def _set_container_list_data(response: Any, data: List[Any]) -> Any: response.last_id = _get_response_id(data[-1]) if data else None if not data and hasattr(response, "has_more"): response.has_more = False + if removed_filtered_items and hasattr(response, "has_more"): + response.has_more = False return response @@ -291,4 +314,8 @@ async def filter_container_list_response( ): filtered.append(item) - return _set_container_list_data(response, filtered) + return _set_container_list_data( + response, + filtered, + removed_filtered_items=len(filtered) != len(data), + ) diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 40397a188606..5ec28e81155a 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -82,6 +82,34 @@ 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_fallback_to_memory_when_persistent_owner_record_fails( + monkeypatch, +): + table = AsyncMock() + table.find_unique.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") + + 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"] + == "user-1" + ) + + @pytest.mark.asyncio async def test_should_track_container_owner_in_memory_without_prisma(monkeypatch): monkeypatch.setattr( @@ -216,7 +244,7 @@ async def test_should_filter_container_list_to_owned_records(monkeypatch): response = ContainerListResponse( object="list", data=[_container("cntr_owned"), _container("cntr_other")], - has_more=False, + has_more=True, ) filtered = await ownership.filter_container_list_response( @@ -228,6 +256,7 @@ async def test_should_filter_container_list_to_owned_records(monkeypatch): assert [item.id for item in filtered.data] == ["cntr_owned"] assert filtered.first_id == "cntr_owned" assert filtered.last_id == "cntr_owned" + assert filtered.has_more is False where = table.find_many.await_args.kwargs["where"] assert where["file_purpose"] == ownership.CONTAINER_OBJECT_PURPOSE assert where["created_by"]["in"] == ["user-1", "user:user-1"] @@ -334,7 +363,7 @@ async def test_should_filter_container_list_with_in_memory_ownership(monkeypatch @pytest.mark.asyncio -async def test_should_preserve_managed_container_id_for_proxy_forwarding(monkeypatch): +async def test_should_forward_decoded_container_id_for_proxy_forwarding(monkeypatch): from litellm.proxy.container_endpoints import handler_factory proxy_server_stub = SimpleNamespace( @@ -388,12 +417,13 @@ async def _handle_llm_api_exception(self, **kwargs): path_params={"container_id": encoded_id}, ) - assert result["container_id"] == encoded_id + assert result["container_id"] == "cntr_provider" assert result["custom_llm_provider"] == "azure" + assert result["model_id"] == "router-gpt" @pytest.mark.asyncio -async def test_should_preserve_managed_container_id_for_multipart_upload(monkeypatch): +async def test_should_forward_decoded_container_id_for_multipart_upload(monkeypatch): from litellm.proxy.common_utils import http_parsing_utils from litellm.proxy.container_endpoints import handler_factory @@ -458,13 +488,14 @@ async def _handle_llm_api_exception(self, **kwargs): container_id=encoded_id, ) - assert result["container_id"] == encoded_id + assert result["container_id"] == "cntr_provider" assert result["custom_llm_provider"] == "azure" + assert result["model_id"] == "router-gpt" assert result["file"] == "file-data" @pytest.mark.asyncio -async def test_should_preserve_managed_container_id_for_proxy_retrieve(monkeypatch): +async def test_should_forward_decoded_container_id_for_proxy_retrieve(monkeypatch): from litellm.proxy.container_endpoints import endpoints proxy_server_stub = SimpleNamespace( @@ -513,12 +544,13 @@ async def _handle_llm_api_exception(self, **kwargs): user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), ) - assert result["container_id"] == encoded_id + assert result["container_id"] == "cntr_provider" assert result["custom_llm_provider"] == "azure" + assert result["model_id"] == "router-gpt" @pytest.mark.asyncio -async def test_should_preserve_managed_container_id_for_proxy_delete(monkeypatch): +async def test_should_forward_decoded_container_id_for_proxy_delete(monkeypatch): from litellm.proxy.container_endpoints import endpoints proxy_server_stub = SimpleNamespace( @@ -567,5 +599,6 @@ async def _handle_llm_api_exception(self, **kwargs): user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), ) - assert result["container_id"] == encoded_id + assert result["container_id"] == "cntr_provider" assert result["custom_llm_provider"] == "azure" + assert result["model_id"] == "router-gpt" From f3cff9338e2c8f0dc7f95beee320f08b36a85b19 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 19:29:08 -0700 Subject: [PATCH 08/27] fix(proxy): harden resource ownership fallbacks --- .../proxy/common_utils/resource_ownership.py | 10 +- .../proxy/container_endpoints/endpoints.py | 20 +-- .../proxy/container_endpoints/ownership.py | 114 +++++++------ .../test_container_proxy_ownership.py | 158 ++++++++++++++++++ .../litellm_proxy/test_skills_ownership.py | 51 ++++++ 5 files changed, 294 insertions(+), 59 deletions(-) diff --git a/litellm/proxy/common_utils/resource_ownership.py b/litellm/proxy/common_utils/resource_ownership.py index 1bc769a8c053..9b55554fbc69 100644 --- a/litellm/proxy/common_utils/resource_ownership.py +++ b/litellm/proxy/common_utils/resource_ownership.py @@ -2,6 +2,8 @@ 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 _add(scope: Optional[str]) -> None: _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 72a7c6c8746c..b67e1e28d8b7 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 b769b2c09538..edb1b65c715e 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 5ec28e81155a..e801300b2944 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 fde1806f92cb..f6946e519d10 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() From 7b1e3f278be77d2f6cfca12b2d9190d2b280c8c2 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 19:34:20 -0700 Subject: [PATCH 09/27] test(proxy): cover container endpoint post processing --- .../test_container_proxy_ownership.py | 111 ++++++++++++++++++ 1 file changed, 111 insertions(+) diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index e801300b2944..576e25bbbadf 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -707,6 +707,117 @@ async def _handle_llm_api_exception(self, **kwargs): assert result["model_id"] == "router-gpt" +@pytest.mark.asyncio +async def test_should_record_container_owner_inside_create_endpoint(monkeypatch): + from litellm.proxy.container_endpoints import endpoints + + proxy_server_stub = SimpleNamespace( + general_settings={}, + llm_router=None, + proxy_config=None, + proxy_logging_obj=None, + select_data_generator=None, + user_api_base=None, + user_max_tokens=None, + user_model=None, + user_request_timeout=None, + user_temperature=None, + version="test", + ) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub) + + response = _container("cntr_provider") + + class FakeProcessor: + def __init__(self, data): + pass + + async def base_process_llm_request(self, **kwargs): + return response + + async def _handle_llm_api_exception(self, **kwargs): + raise kwargs["e"] + + record_owner = AsyncMock(return_value=response) + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", FakeProcessor) + monkeypatch.setattr(endpoints, "record_container_owner", record_owner) + + result = await endpoints.create_container( + request=SimpleNamespace( + query_params={}, + headers={}, + json=AsyncMock(return_value={}), + body=AsyncMock(return_value=b"{}"), + ), + fastapi_response=SimpleNamespace(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + ) + + assert result == response + record_owner.assert_awaited_once_with( + response=response, + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + custom_llm_provider="openai", + ) + + +@pytest.mark.asyncio +async def test_should_filter_container_list_inside_list_endpoint(monkeypatch): + from litellm.proxy.container_endpoints import endpoints + + proxy_server_stub = SimpleNamespace( + general_settings={}, + llm_router=None, + proxy_config=None, + proxy_logging_obj=None, + select_data_generator=None, + user_api_base=None, + user_max_tokens=None, + user_model=None, + user_request_timeout=None, + user_temperature=None, + version="test", + ) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub) + + response = ContainerListResponse( + object="list", + data=[_container("cntr_provider")], + has_more=False, + ) + + class FakeProcessor: + def __init__(self, data): + pass + + async def base_process_llm_request(self, **kwargs): + return response + + async def _handle_llm_api_exception(self, **kwargs): + raise kwargs["e"] + + filter_response = AsyncMock(return_value=response) + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", FakeProcessor) + monkeypatch.setattr( + endpoints, + "filter_container_list_response", + filter_response, + ) + + result = await endpoints.list_containers( + request=SimpleNamespace(query_params={}, headers={}), + fastapi_response=SimpleNamespace(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + ) + + assert result == response + filter_response.assert_awaited_once_with( + response=response, + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + custom_llm_provider="openai", + ) + + @pytest.mark.asyncio async def test_should_forward_decoded_container_id_for_proxy_delete(monkeypatch): from litellm.proxy.container_endpoints import endpoints From 20eb9c96caba8d6460cf54bedcc1885b58c3831b Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 19:43:20 -0700 Subject: [PATCH 10/27] fix(proxy): avoid mutating container responses --- .../proxy/container_endpoints/ownership.py | 2 +- .../test_container_proxy_ownership.py | 29 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index edb1b65c715e..05634d506b75 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -73,7 +73,7 @@ def _get_response_id(response: Any) -> Optional[str]: def _dump_response(response: Any) -> Dict[str, Any]: if isinstance(response, dict): - return response + return dict(response) if hasattr(response, "model_dump"): return response.model_dump() if hasattr(response, "dict"): diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 576e25bbbadf..8ebfd3ab6f72 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -57,6 +57,35 @@ async def test_should_record_container_owner_with_original_provider_id(monkeypat assert data["created_by"] == "user-1" +@pytest.mark.asyncio +async def test_should_not_mutate_dict_container_response_when_recording_owner( + 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(user_id="user-1") + response = {"id": "cntr_provider", "object": "container"} + + returned = await ownership.record_container_owner( + response=response, + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert returned == {"id": "cntr_provider", "object": "container"} + data = table.create.await_args.kwargs["data"] + assert data["file_object"]["custom_llm_provider"] == "openai" + assert data["file_object"]["provider_container_id"] == "cntr_provider" + + @pytest.mark.asyncio async def test_should_record_team_owner_for_keys_without_user_id(monkeypatch): table = AsyncMock() From 3a566c3938f6623cecd7190c04f09bacdd0feca9 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 19:55:00 -0700 Subject: [PATCH 11/27] fix(proxy): stabilize ownership fallback and openapi ids --- litellm/proxy/_lazy_openapi_snapshot.json | 90 +++++++++---------- litellm/proxy/_lazy_openapi_snapshot.py | 4 +- .../proxy/container_endpoints/ownership.py | 6 +- litellm/proxy/proxy_server.py | 73 +++++++++++++++ .../test_container_proxy_ownership.py | 57 ++++++++++++ .../proxy/test_swagger_chat_completions.py | 74 ++++++++++++++- 6 files changed, 255 insertions(+), 49 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index b8e9eb6c261e..e8b2d701691d 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3616,7 +3616,7 @@ }, "get": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", + "operationId": "anthropic_proxy_route_anthropic__endpoint__get", "parameters": [ { "in": "path", @@ -3660,7 +3660,7 @@ }, "patch": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", + "operationId": "anthropic_proxy_route_anthropic__endpoint__patch", "parameters": [ { "in": "path", @@ -3704,7 +3704,7 @@ }, "post": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", + "operationId": "anthropic_proxy_route_anthropic__endpoint__post", "parameters": [ { "in": "path", @@ -3748,7 +3748,7 @@ }, "put": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", + "operationId": "anthropic_proxy_route_anthropic__endpoint__put", "parameters": [ { "in": "path", @@ -13299,7 +13299,7 @@ }, "get": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", + "operationId": "langfuse_proxy_route_langfuse__endpoint__get", "parameters": [ { "in": "path", @@ -13338,7 +13338,7 @@ }, "patch": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", + "operationId": "langfuse_proxy_route_langfuse__endpoint__patch", "parameters": [ { "in": "path", @@ -13377,7 +13377,7 @@ }, "post": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", + "operationId": "langfuse_proxy_route_langfuse__endpoint__post", "parameters": [ { "in": "path", @@ -13416,7 +13416,7 @@ }, "put": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", + "operationId": "langfuse_proxy_route_langfuse__endpoint__put", "parameters": [ { "in": "path", @@ -14008,7 +14008,7 @@ "/mcp-rest/test/connection": { "post": { "description": "Test if we can connect to the provided MCP server before adding it", - "operationId": "test_connection_mcp_rest_test_connection_post", + "operationId": "test_connection_mcp_rest_test_connection_post_2", "requestBody": { "content": { "application/json": { @@ -14053,7 +14053,7 @@ "/mcp-rest/test/tools/list": { "post": { "description": "Preview tools available from MCP server before adding it", - "operationId": "test_tools_list_mcp_rest_test_tools_list_post", + "operationId": "test_tools_list_mcp_rest_test_tools_list_post_2", "requestBody": { "content": { "application/json": { @@ -14098,7 +14098,7 @@ "/mcp-rest/tools/call": { "post": { "description": "REST API to call a specific MCP tool with the provided arguments", - "operationId": "call_tool_rest_api_mcp_rest_tools_call_post", + "operationId": "call_tool_rest_api_mcp_rest_tools_call_post_2", "responses": { "200": { "content": { @@ -14123,7 +14123,7 @@ "/mcp-rest/tools/list": { "get": { "description": "List all available tools with information about the server they belong to.\n\nExample response:\n{\n \"tools\": [\n {\n \"name\": \"create_zap\",\n \"description\": \"Create a new zap\",\n \"inputSchema\": \"tool_input_schema\",\n \"mcp_info\": {\n \"server_name\": \"zapier\",\n \"logo_url\": \"https://www.zapier.com/logo.png\",\n }\n }\n ],\n \"error\": null,\n \"message\": \"Successfully retrieved tools\"\n}", - "operationId": "list_tool_rest_api_mcp_rest_tools_list_get", + "operationId": "list_tool_rest_api_mcp_rest_tools_list_get_2", "parameters": [ { "description": "The server id to list tools for", @@ -21896,7 +21896,7 @@ "/policies/usage/overview": { "get": { "description": "Return policy performance overview for the dashboard.", - "operationId": "policies_usage_overview_policies_usage_overview_get", + "operationId": "policies_usage_overview_policies_usage_overview_get_2", "parameters": [ { "description": "YYYY-MM-DD", @@ -22521,7 +22521,7 @@ "/policies/attachments/estimate-impact": { "post": { "description": "Estimate how many keys and teams would be affected by a policy attachment.\n\nUse this before creating an attachment to preview the blast radius.\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/policies/attachments/estimate-impact\" \\\n -H \"Authorization: Bearer \" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"policy_name\": \"hipaa-compliance\",\n \"tags\": [\"healthcare\", \"health-*\"]\n }'\n```", - "operationId": "estimate_attachment_impact_policies_attachments_estimate_impact_post", + "operationId": "estimate_attachment_impact_policies_attachments_estimate_impact_post_2", "requestBody": { "content": { "application/json": { @@ -22568,7 +22568,7 @@ "/policies/resolve": { "post": { "description": "Resolve which policies and guardrails apply for a given context.\n\nUse this endpoint to debug \"what guardrails would apply to a request\nwith this team/key/model/tags combination?\"\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/policies/resolve\" \\\n -H \"Authorization: Bearer \" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"tags\": [\"healthcare\"],\n \"model\": \"gpt-4\"\n }'\n```", - "operationId": "resolve_policies_for_context_policies_resolve_post", + "operationId": "resolve_policies_for_context_policies_resolve_post_2", "parameters": [ { "description": "Force a DB sync before resolving. Default uses in-memory cache.", @@ -26922,7 +26922,7 @@ }, "get": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_get", "parameters": [ { "in": "path", @@ -26961,7 +26961,7 @@ }, "head": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_head", "parameters": [ { "in": "path", @@ -27000,7 +27000,7 @@ }, "options": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_options", "parameters": [ { "in": "path", @@ -27039,7 +27039,7 @@ }, "patch": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_patch", "parameters": [ { "in": "path", @@ -27078,7 +27078,7 @@ }, "post": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_post", "parameters": [ { "in": "path", @@ -27117,7 +27117,7 @@ }, "put": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", "parameters": [ { "in": "path", @@ -28329,7 +28329,7 @@ "/v1/vector_stores": { "get": { "description": "List vector stores.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/list", - "operationId": "vector_store_list_v1_vector_stores_get", + "operationId": "vector_store_list_v1_vector_stores_get_2", "parameters": [ { "in": "query", @@ -28430,7 +28430,7 @@ }, "post": { "description": "Create a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/create\n\nSupports target_model_names parameter for creating vector stores across multiple models:\n```json\n{\n \"name\": \"my-vector-store\",\n \"target_model_names\": \"gpt-4,gemini-2.0\"\n}\n```", - "operationId": "vector_store_create_v1_vector_stores_post", + "operationId": "vector_store_create_v1_vector_stores_post_2", "responses": { "200": { "content": { @@ -28455,7 +28455,7 @@ "/v1/vector_stores/{vector_store_id}": { "delete": { "description": "Delete a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/delete", - "operationId": "vector_store_delete_v1_vector_stores__vector_store_id__delete", + "operationId": "vector_store_delete_v1_vector_stores__vector_store_id__delete_2", "parameters": [ { "in": "path", @@ -28499,7 +28499,7 @@ }, "get": { "description": "Retrieve a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/retrieve", - "operationId": "vector_store_retrieve_v1_vector_stores__vector_store_id__get", + "operationId": "vector_store_retrieve_v1_vector_stores__vector_store_id__get_2", "parameters": [ { "in": "path", @@ -28543,7 +28543,7 @@ }, "post": { "description": "Update a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/modify", - "operationId": "vector_store_update_v1_vector_stores__vector_store_id__post", + "operationId": "vector_store_update_v1_vector_stores__vector_store_id__post_2", "parameters": [ { "in": "path", @@ -28588,7 +28588,7 @@ }, "/v1/vector_stores/{vector_store_id}/files": { "get": { - "operationId": "vector_store_file_list_v1_vector_stores__vector_store_id__files_get", + "operationId": "vector_store_file_list_v1_vector_stores__vector_store_id__files_get_2", "parameters": [ { "in": "path", @@ -28631,7 +28631,7 @@ ] }, "post": { - "operationId": "vector_store_file_create_v1_vector_stores__vector_store_id__files_post", + "operationId": "vector_store_file_create_v1_vector_stores__vector_store_id__files_post_2", "parameters": [ { "in": "path", @@ -28676,7 +28676,7 @@ }, "/v1/vector_stores/{vector_store_id}/files/{file_id}": { "delete": { - "operationId": "vector_store_file_delete_v1_vector_stores__vector_store_id__files__file_id__delete", + "operationId": "vector_store_file_delete_v1_vector_stores__vector_store_id__files__file_id__delete_2", "parameters": [ { "in": "path", @@ -28728,7 +28728,7 @@ ] }, "get": { - "operationId": "vector_store_file_retrieve_v1_vector_stores__vector_store_id__files__file_id__get", + "operationId": "vector_store_file_retrieve_v1_vector_stores__vector_store_id__files__file_id__get_2", "parameters": [ { "in": "path", @@ -28780,7 +28780,7 @@ ] }, "post": { - "operationId": "vector_store_file_update_v1_vector_stores__vector_store_id__files__file_id__post", + "operationId": "vector_store_file_update_v1_vector_stores__vector_store_id__files__file_id__post_2", "parameters": [ { "in": "path", @@ -28834,7 +28834,7 @@ }, "/v1/vector_stores/{vector_store_id}/files/{file_id}/content": { "get": { - "operationId": "vector_store_file_content_v1_vector_stores__vector_store_id__files__file_id__content_get", + "operationId": "vector_store_file_content_v1_vector_stores__vector_store_id__files__file_id__content_get_2", "parameters": [ { "in": "path", @@ -28889,7 +28889,7 @@ "/v1/vector_stores/{vector_store_id}/search": { "post": { "description": "Search a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/search", - "operationId": "vector_store_search_v1_vector_stores__vector_store_id__search_post", + "operationId": "vector_store_search_v1_vector_stores__vector_store_id__search_post_2", "parameters": [ { "in": "path", @@ -28935,7 +28935,7 @@ "/vector_stores": { "get": { "description": "List vector stores.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/list", - "operationId": "vector_store_list_vector_stores_get", + "operationId": "vector_store_list_vector_stores_get_2", "parameters": [ { "in": "query", @@ -29036,7 +29036,7 @@ }, "post": { "description": "Create a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/create\n\nSupports target_model_names parameter for creating vector stores across multiple models:\n```json\n{\n \"name\": \"my-vector-store\",\n \"target_model_names\": \"gpt-4,gemini-2.0\"\n}\n```", - "operationId": "vector_store_create_vector_stores_post", + "operationId": "vector_store_create_vector_stores_post_2", "responses": { "200": { "content": { @@ -29061,7 +29061,7 @@ "/vector_stores/{vector_store_id}": { "delete": { "description": "Delete a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/delete", - "operationId": "vector_store_delete_vector_stores__vector_store_id__delete", + "operationId": "vector_store_delete_vector_stores__vector_store_id__delete_2", "parameters": [ { "in": "path", @@ -29105,7 +29105,7 @@ }, "get": { "description": "Retrieve a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/retrieve", - "operationId": "vector_store_retrieve_vector_stores__vector_store_id__get", + "operationId": "vector_store_retrieve_vector_stores__vector_store_id__get_2", "parameters": [ { "in": "path", @@ -29149,7 +29149,7 @@ }, "post": { "description": "Update a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/modify", - "operationId": "vector_store_update_vector_stores__vector_store_id__post", + "operationId": "vector_store_update_vector_stores__vector_store_id__post_2", "parameters": [ { "in": "path", @@ -29194,7 +29194,7 @@ }, "/vector_stores/{vector_store_id}/files": { "get": { - "operationId": "vector_store_file_list_vector_stores__vector_store_id__files_get", + "operationId": "vector_store_file_list_vector_stores__vector_store_id__files_get_2", "parameters": [ { "in": "path", @@ -29237,7 +29237,7 @@ ] }, "post": { - "operationId": "vector_store_file_create_vector_stores__vector_store_id__files_post", + "operationId": "vector_store_file_create_vector_stores__vector_store_id__files_post_2", "parameters": [ { "in": "path", @@ -29282,7 +29282,7 @@ }, "/vector_stores/{vector_store_id}/files/{file_id}": { "delete": { - "operationId": "vector_store_file_delete_vector_stores__vector_store_id__files__file_id__delete", + "operationId": "vector_store_file_delete_vector_stores__vector_store_id__files__file_id__delete_2", "parameters": [ { "in": "path", @@ -29334,7 +29334,7 @@ ] }, "get": { - "operationId": "vector_store_file_retrieve_vector_stores__vector_store_id__files__file_id__get", + "operationId": "vector_store_file_retrieve_vector_stores__vector_store_id__files__file_id__get_2", "parameters": [ { "in": "path", @@ -29386,7 +29386,7 @@ ] }, "post": { - "operationId": "vector_store_file_update_vector_stores__vector_store_id__files__file_id__post", + "operationId": "vector_store_file_update_vector_stores__vector_store_id__files__file_id__post_2", "parameters": [ { "in": "path", @@ -29440,7 +29440,7 @@ }, "/vector_stores/{vector_store_id}/files/{file_id}/content": { "get": { - "operationId": "vector_store_file_content_vector_stores__vector_store_id__files__file_id__content_get", + "operationId": "vector_store_file_content_vector_stores__vector_store_id__files__file_id__content_get_2", "parameters": [ { "in": "path", @@ -29495,7 +29495,7 @@ "/vector_stores/{vector_store_id}/search": { "post": { "description": "Search a vector store.\n\nAPI Reference:\nhttps://platform.openai.com/docs/api-reference/vector-stores/search", - "operationId": "vector_store_search_vector_stores__vector_store_id__search_post", + "operationId": "vector_store_search_vector_stores__vector_store_id__search_post_2", "parameters": [ { "in": "path", diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index 315f6a9742a0..2aa9dcca742b 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -31,7 +31,7 @@ def generate_snapshot() -> Dict[str, Dict]: from fastapi.openapi.utils import get_openapi from litellm.proxy._lazy_features import LAZY_FEATURES - from litellm.proxy.proxy_server import app + from litellm.proxy.proxy_server import app, ensure_unique_openapi_operation_ids for feat in LAZY_FEATURES: if feat.module_path in sys.modules: @@ -43,6 +43,7 @@ def generate_snapshot() -> Dict[str, Dict]: sys.stderr.write(f"warning: skip {feat.name}: {exc}\n") fragments: Dict[str, Dict] = {} + used_operation_ids = set() for feat in LAZY_FEATURES: feat_routes = [ r @@ -57,6 +58,7 @@ def generate_snapshot() -> Dict[str, Dict]: for op in path_ops.values(): if isinstance(op, dict): op["tags"] = [feat.name] + full = ensure_unique_openapi_operation_ids(full, used_operation_ids) fragments[feat.name] = { "paths": full.get("paths", {}), "components": {"schemas": full.get("components", {}).get("schemas", {})}, diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 05634d506b75..05fe7e648467 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -189,7 +189,9 @@ async def _get_container_owner( "file_purpose": CONTAINER_OBJECT_PURPOSE, } ) - return getattr(row, "created_by", None) if row is not None else None + if row is not None: + return getattr(row, "created_by", None) + return _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) except Exception as e: verbose_proxy_logger.warning( "Failed to load container ownership for container_id=%s; " @@ -287,7 +289,7 @@ async def _get_allowed_container_ids( "created_by": {"in": owner_scopes}, } ) - return { + return in_memory_allowed_ids | { row.model_object_id for row in rows if getattr(row, "model_object_id", None) is not None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b6cec92a1d11..ac7f86b20285 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -958,6 +958,77 @@ def _generate_stable_operation_id(route: Any) -> str: return operation_id +_OPENAPI_HTTP_METHODS = { + "delete", + "get", + "head", + "options", + "patch", + "post", + "put", + "trace", +} + + +def _strip_operation_id_method_suffix(operation_id: str) -> str: + base, separator, suffix = operation_id.rpartition("_") + if separator and suffix in _OPENAPI_HTTP_METHODS: + return base + return operation_id + + +def ensure_unique_openapi_operation_ids( + openapi_schema: Dict[str, Any], + reserved_operation_ids: Optional[Set[str]] = None, +) -> Dict[str, Any]: + operation_entries = [] + operation_id_counts: Dict[str, int] = {} + for path_item in openapi_schema.get("paths", {}).values(): + if not isinstance(path_item, dict): + continue + for method, operation in path_item.items(): + if method not in _OPENAPI_HTTP_METHODS or not isinstance(operation, dict): + continue + operation_id = operation.get("operationId") + if not isinstance(operation_id, str): + continue + operation_entries.append((method, operation, operation_id)) + operation_id_counts[operation_id] = ( + operation_id_counts.get(operation_id, 0) + 1 + ) + + used_operation_ids = set(reserved_operation_ids or set()) + seen_operation_ids: Set[str] = set() + for method, operation, operation_id in operation_entries: + should_rewrite = ( + operation_id_counts[operation_id] > 1 + or operation_id in used_operation_ids + or operation_id in seen_operation_ids + ) + if not should_rewrite: + seen_operation_ids.add(operation_id) + used_operation_ids.add(operation_id) + continue + + base_operation_id = _strip_operation_id_method_suffix(operation_id) + new_operation_id = f"{base_operation_id}_{method}" + suffix = 2 + while ( + new_operation_id in used_operation_ids + or new_operation_id in seen_operation_ids + ): + new_operation_id = f"{base_operation_id}_{method}_{suffix}" + suffix += 1 + operation["operationId"] = new_operation_id + seen_operation_ids.add(new_operation_id) + used_operation_ids.add(new_operation_id) + + if reserved_operation_ids is not None: + reserved_operation_ids.update(used_operation_ids) + + return openapi_schema + + app = FastAPI( docs_url=_get_docs_url(), redoc_url=_get_redoc_url(), @@ -1047,6 +1118,7 @@ def get_openapi_schema(): from litellm.proxy._lazy_features import inject_lazy_stubs openapi_schema = inject_lazy_stubs(openapi_schema) + openapi_schema = ensure_unique_openapi_operation_ids(openapi_schema) # Fix Swagger UI execute path error when server_root_path is set if server_root_path: @@ -1078,6 +1150,7 @@ def custom_openapi(): from litellm.proxy._lazy_features import inject_lazy_stubs openapi_schema = inject_lazy_stubs(openapi_schema) + openapi_schema = ensure_unique_openapi_operation_ids(openapi_schema) # Fix Swagger UI execute path error when server_root_path is set if server_root_path: diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 8ebfd3ab6f72..fbce5d698597 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -309,6 +309,31 @@ async def test_should_fallback_to_memory_when_owner_lookup_fails(monkeypatch): assert provider == "openai" +@pytest.mark.asyncio +async def test_should_use_memory_owner_when_db_recovers_without_row(monkeypatch): + table = AsyncMock() + table.find_first.return_value = None + 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() @@ -549,6 +574,38 @@ async def test_should_filter_container_list_with_memory_when_db_lookup_fails( assert filtered.has_more is False +@pytest.mark.asyncio +async def test_should_include_memory_container_list_when_db_recovers_without_row( + monkeypatch, +): + table = AsyncMock() + table.find_many.return_value = [] + 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/proxy/test_swagger_chat_completions.py b/tests/test_litellm/proxy/test_swagger_chat_completions.py index 4723034114ac..2a773bd3168b 100644 --- a/tests/test_litellm/proxy/test_swagger_chat_completions.py +++ b/tests/test_litellm/proxy/test_swagger_chat_completions.py @@ -5,7 +5,7 @@ for the /chat/completions endpoint, showing all expected fields in the Swagger documentation. """ -from unittest.mock import Mock, patch +from unittest.mock import patch import pytest from fastapi.testclient import TestClient @@ -414,3 +414,75 @@ def test_openapi_schema_servers_url_with_root_path(self): assert ( schema["servers"][0]["url"] == expected_url ), f"Expected servers URL '{expected_url}' in custom_openapi, got '{schema['servers'][0]['url']}'" + + def test_should_make_duplicate_operation_ids_unique_by_method(self): + from litellm.proxy.proxy_server import ensure_unique_openapi_operation_ids + + schema = { + "paths": { + "/anthropic/{endpoint}": { + "delete": { + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete" + }, + "get": { + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete" + }, + "post": { + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete" + }, + }, + "/models": { + "get": { + "operationId": "list_models_models_get", + } + }, + } + } + + result = ensure_unique_openapi_operation_ids(schema) + + assert ( + result["paths"]["/anthropic/{endpoint}"]["delete"]["operationId"] + == "anthropic_proxy_route_anthropic__endpoint__delete" + ) + assert ( + result["paths"]["/anthropic/{endpoint}"]["get"]["operationId"] + == "anthropic_proxy_route_anthropic__endpoint__get" + ) + assert ( + result["paths"]["/anthropic/{endpoint}"]["post"]["operationId"] + == "anthropic_proxy_route_anthropic__endpoint__post" + ) + operation_ids = [ + operation["operationId"] + for path_item in result["paths"].values() + for operation in path_item.values() + ] + assert len(operation_ids) == len(set(operation_ids)) + + def test_should_reserve_operation_ids_across_lazy_fragments(self): + from litellm.proxy.proxy_server import ensure_unique_openapi_operation_ids + + used_operation_ids = set() + first_schema = { + "paths": { + "/tools/list": { + "get": {"operationId": "list_tool_rest_api_mcp_rest_tools_list_get"} + } + } + } + second_schema = { + "paths": { + "/v1/tools/list": { + "get": {"operationId": "list_tool_rest_api_mcp_rest_tools_list_get"} + } + } + } + + ensure_unique_openapi_operation_ids(first_schema, used_operation_ids) + result = ensure_unique_openapi_operation_ids(second_schema, used_operation_ids) + + assert ( + result["paths"]["/v1/tools/list"]["get"]["operationId"] + == "list_tool_rest_api_mcp_rest_tools_list_get_2" + ) From a02aacf1b5012e2e27e5d3a427b5bf654f5ed3b9 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 20:01:43 -0700 Subject: [PATCH 12/27] fix(proxy): type openapi snapshot operation ids --- litellm/proxy/_lazy_openapi_snapshot.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index 2aa9dcca742b..51cbd6eb989e 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -10,7 +10,7 @@ import json import sys from pathlib import Path -from typing import Dict, Optional +from typing import Dict, Optional, Set SNAPSHOT_FILE = Path(__file__).parent / "_lazy_openapi_snapshot.json" @@ -43,7 +43,7 @@ def generate_snapshot() -> Dict[str, Dict]: sys.stderr.write(f"warning: skip {feat.name}: {exc}\n") fragments: Dict[str, Dict] = {} - used_operation_ids = set() + used_operation_ids: Set[str] = set() for feat in LAZY_FEATURES: feat_routes = [ r From 2ecc79b9e9b41f04230048892164b32f4d62cd43 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 20:06:44 -0700 Subject: [PATCH 13/27] test(proxy): cover skill ownership propagation --- .../litellm_proxy/test_skills_ownership.py | 194 +++++++++++++++++- 1 file changed, 191 insertions(+), 3 deletions(-) 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 f6946e519d10..45ed172a8451 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -1,10 +1,20 @@ -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, Mock import pytest -from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler from litellm.llms.litellm_proxy.skills import handler as skills_handler -from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth +from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler +from litellm.llms.litellm_proxy.skills.transformation import ( + LiteLLMSkillsTransformationHandler, +) +from litellm.proxy._types import ( + LiteLLM_SkillsTable, + LitellmUserRoles, + NewSkillRequest, + UserAPIKeyAuth, +) +from litellm.proxy.common_utils import resource_ownership +from litellm.skills import main as skills_main @pytest.fixture(autouse=True) @@ -20,6 +30,184 @@ def _skill(skill_id: str, created_by: str | None) -> LiteLLM_SkillsTable: ) +def test_should_extract_skill_auth_from_supported_metadata_fields(): + auth = UserAPIKeyAuth(user_id="user-1") + + assert ( + skills_main._get_user_api_key_auth_from_kwargs( + {"metadata": {"user_api_key_auth": auth}} + ) + is auth + ) + assert ( + skills_main._get_user_api_key_auth_from_kwargs( + {"metadata": {}, "litellm_metadata": {"user_api_key_auth": auth}} + ) + is auth + ) + assert skills_main._get_user_api_key_auth_from_kwargs({"metadata": "bad"}) is None + + +def test_should_extract_skill_request_metadata_with_extra_body_precedence(): + body_metadata = {"purpose": "body"} + requester_metadata = {"purpose": "requester"} + + assert ( + skills_main._get_skill_request_metadata( + {"metadata": {"requester_metadata": requester_metadata}}, + {"metadata": body_metadata}, + ) + == body_metadata + ) + assert ( + skills_main._get_skill_request_metadata( + {"metadata": {"requester_metadata": requester_metadata}}, + None, + ) + == requester_metadata + ) + assert ( + skills_main._get_skill_request_metadata( + {"metadata": {}}, + {"metadata": "bad"}, + ) + is None + ) + + +def test_should_forward_skill_auth_through_sdk_entrypoints(monkeypatch): + auth = UserAPIKeyAuth(user_id="user-1") + handler = Mock() + handler.create_skill_handler.return_value = "created" + handler.list_skills_handler.return_value = "listed" + handler.get_skill_handler.return_value = "got" + handler.delete_skill_handler.return_value = "deleted" + monkeypatch.setattr(skills_main, "_get_litellm_skills_handler", lambda: handler) + + assert ( + skills_main.create_skill( + display_title="skill", + extra_body={"metadata": {"source": "request"}}, + custom_llm_provider="litellm_proxy", + metadata={"user_api_key_auth": auth}, + user_id="user-1", + ) + == "created" + ) + assert ( + skills_main.list_skills( + custom_llm_provider="litellm_proxy", + metadata={"user_api_key_auth": auth}, + ) + == "listed" + ) + assert ( + skills_main.get_skill( + "litellm_skill_1", + custom_llm_provider="litellm_proxy", + metadata={"user_api_key_auth": auth}, + ) + == "got" + ) + assert ( + skills_main.delete_skill( + "litellm_skill_1", + custom_llm_provider="litellm_proxy", + metadata={"user_api_key_auth": auth}, + ) + == "deleted" + ) + + assert handler.create_skill_handler.call_args.kwargs["metadata"] == { + "source": "request" + } + assert handler.create_skill_handler.call_args.kwargs["user_api_key_dict"] is auth + assert handler.list_skills_handler.call_args.kwargs["user_api_key_dict"] is auth + assert handler.get_skill_handler.call_args.kwargs["user_api_key_dict"] is auth + assert handler.delete_skill_handler.call_args.kwargs["user_api_key_dict"] is auth + + +def test_should_build_resource_owner_scopes_for_auth_context(): + auth = UserAPIKeyAuth( + user_id="user-1", + team_id="team-1", + org_id="org-1", + api_key="api-key-hash", + token="token-hash", + ) + + assert resource_ownership.get_resource_owner_scopes(auth) == [ + "user-1", + "user:user-1", + "team:team-1", + "org:org-1", + "key:api-key-hash", + ] + assert resource_ownership.get_primary_resource_owner_scope(auth) == "user-1" + assert resource_ownership.user_can_access_resource_owner("team:team-1", auth) + assert resource_ownership.get_resource_owner_scopes( + UserAPIKeyAuth(token="token-hash") + ) == ["key:token-hash"] + assert resource_ownership.get_resource_owner_scopes(UserAPIKeyAuth()) == [ + resource_ownership.UNSCOPED_RESOURCE_OWNER_SCOPE + ] + + +def test_should_allow_admin_and_anonymous_resource_owner_paths(): + admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value) + + assert resource_ownership.is_proxy_admin(admin) + assert resource_ownership.user_can_access_resource_owner(None, admin) + assert resource_ownership.user_can_access_resource_owner(None, None) + assert not resource_ownership.user_can_access_resource_owner( + None, UserAPIKeyAuth(user_id="user-1") + ) + + +@pytest.mark.asyncio +async def test_should_forward_skill_auth_through_transformation_handler(monkeypatch): + handler = LiteLLMSkillsTransformationHandler() + auth = UserAPIKeyAuth(user_id="user-1") + + create_skill = AsyncMock(return_value=_skill("litellm_skill_created", "user-1")) + list_skills = AsyncMock(return_value=[_skill("litellm_skill_listed", "user-1")]) + get_skill = AsyncMock(return_value=_skill("litellm_skill_got", "user-1")) + delete_skill = AsyncMock(return_value={"id": "litellm_skill_deleted"}) + monkeypatch.setattr(LiteLLMSkillsHandler, "create_skill", create_skill) + monkeypatch.setattr(LiteLLMSkillsHandler, "list_skills", list_skills) + monkeypatch.setattr(LiteLLMSkillsHandler, "get_skill", get_skill) + monkeypatch.setattr(LiteLLMSkillsHandler, "delete_skill", delete_skill) + + created = await handler._async_create_skill( + display_title="skill", + metadata={"source": "request"}, + user_id="user-1", + user_api_key_dict=auth, + ) + listed = await handler._async_list_skills( + limit=10, + offset=2, + user_api_key_dict=auth, + ) + got = await handler._async_get_skill( + "litellm_skill_got", + user_api_key_dict=auth, + ) + deleted = await handler._async_delete_skill( + "litellm_skill_deleted", + user_api_key_dict=auth, + ) + + assert created.id == "litellm_skill_created" + assert [skill.id for skill in listed.data] == ["litellm_skill_listed"] + assert got.id == "litellm_skill_got" + assert deleted.id == "litellm_skill_deleted" + assert create_skill.await_args.kwargs["user_api_key_dict"] is auth + assert list_skills.await_args.kwargs["user_api_key_dict"] is auth + assert get_skill.await_args.kwargs["user_api_key_dict"] is auth + assert delete_skill.await_args.kwargs["user_api_key_dict"] is auth + + @pytest.mark.asyncio async def test_should_store_team_owner_for_keys_without_user_id(monkeypatch): table = AsyncMock() From f18ee0319d17fccbcbc3f2a6efc0d65b81583c4d Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 20:25:40 -0700 Subject: [PATCH 14/27] fix(proxy): isolate ownership persistence paths --- litellm/llms/litellm_proxy/skills/handler.py | 21 +++-- litellm/llms/litellm_proxy/skills/store.py | 22 ++++++ .../proxy/container_endpoints/endpoints.py | 10 +-- .../proxy/container_endpoints/ownership.py | 64 +++++++++------- .../container_endpoints/ownership_store.py | 55 ++++++++++++++ .../test_container_proxy_ownership.py | 76 +++++++++++++++++++ 6 files changed, 204 insertions(+), 44 deletions(-) create mode 100644 litellm/llms/litellm_proxy/skills/store.py create mode 100644 litellm/proxy/container_endpoints/ownership_store.py diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 7093874acd93..6eb6f40fc64e 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -10,6 +10,7 @@ from typing import Any, Dict, List, Optional from litellm._logging import verbose_logger +from litellm.llms.litellm_proxy.skills.store import LiteLLMSkillsStore from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, @@ -104,6 +105,7 @@ async def create_skill( LiteLLM_SkillsTable record """ prisma_client = await LiteLLMSkillsHandler._get_prisma_client() + store = LiteLLMSkillsStore(prisma_client) skill_id = f"litellm_skill_{uuid.uuid4()}" owner = get_primary_resource_owner_scope(user_api_key_dict) or user_id @@ -138,7 +140,7 @@ async def create_skill( f"LiteLLMSkillsHandler: Creating skill {skill_id} with title={data.display_title}" ) - new_skill = await prisma_client.db.litellm_skillstable.create(data=skill_data) + new_skill = await store.create_skill(skill_data) return _prisma_skill_to_litellm(new_skill) @@ -159,6 +161,7 @@ async def list_skills( List of LiteLLM_SkillsTable records """ prisma_client = await LiteLLMSkillsHandler._get_prisma_client() + store = LiteLLMSkillsStore(prisma_client) verbose_logger.debug( f"LiteLLMSkillsHandler: Listing skills with limit={limit}, offset={offset}" @@ -183,9 +186,7 @@ async def list_skills( else: find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} - skills = await prisma_client.db.litellm_skillstable.find_many( - **find_many_kwargs - ) + skills = await store.list_skills(find_many_kwargs) return [_prisma_skill_to_litellm(s) for s in skills] @@ -207,12 +208,11 @@ async def get_skill( ValueError: If skill not found """ prisma_client = await LiteLLMSkillsHandler._get_prisma_client() + store = LiteLLMSkillsStore(prisma_client) verbose_logger.debug(f"LiteLLMSkillsHandler: Getting skill {skill_id}") - skill = await prisma_client.db.litellm_skillstable.find_unique( - where={"skill_id": skill_id} - ) + skill = await store.find_skill(skill_id) if skill is None: raise ValueError(f"Skill not found: {skill_id}") @@ -242,13 +242,12 @@ async def delete_skill( ValueError: If skill not found """ prisma_client = await LiteLLMSkillsHandler._get_prisma_client() + store = LiteLLMSkillsStore(prisma_client) verbose_logger.debug(f"LiteLLMSkillsHandler: Deleting skill {skill_id}") # Check if skill exists - skill = await prisma_client.db.litellm_skillstable.find_unique( - where={"skill_id": skill_id} - ) + skill = await store.find_skill(skill_id) if skill is None: raise ValueError(f"Skill not found: {skill_id}") @@ -259,7 +258,7 @@ async def delete_skill( raise ValueError(f"Skill not found: {skill_id}") # Delete the skill - await prisma_client.db.litellm_skillstable.delete(where={"skill_id": skill_id}) + await store.delete_skill(skill_id) return {"id": skill_id, "type": "skill_deleted"} diff --git a/litellm/llms/litellm_proxy/skills/store.py b/litellm/llms/litellm_proxy/skills/store.py new file mode 100644 index 000000000000..6e69a7a96251 --- /dev/null +++ b/litellm/llms/litellm_proxy/skills/store.py @@ -0,0 +1,22 @@ +from typing import Any, Dict, List, Optional + + +class LiteLLMSkillsStore: + def __init__(self, prisma_client: Any): + self.prisma_client = prisma_client + + @property + def _table(self) -> Any: + return self.prisma_client.db.litellm_skillstable + + async def create_skill(self, data: Dict[str, Any]) -> Any: + return await self._table.create(data=data) + + async def list_skills(self, find_many_kwargs: Dict[str, Any]) -> List[Any]: + return await self._table.find_many(**find_many_kwargs) + + async def find_skill(self, skill_id: str) -> Optional[Any]: + return await self._table.find_unique(where={"skill_id": skill_id}) + + async def delete_skill(self, skill_id: str) -> None: + await self._table.delete(where={"skill_id": skill_id}) diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index b67e1e28d8b7..089f22a23a57 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -122,11 +122,6 @@ 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, @@ -134,6 +129,11 @@ 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( diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 05fe7e648467..86188c02414e 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,4 +1,5 @@ import os +from collections import OrderedDict from typing import Any, Dict, List, Optional, Set, Tuple from fastapi import HTTPException @@ -11,11 +12,15 @@ is_proxy_admin, user_can_access_resource_owner, ) +from litellm.proxy.container_endpoints.ownership_store import ( + CONTAINER_OBJECT_PURPOSE, + ContainerOwnershipStore, +) from litellm.responses.utils import ResponsesAPIRequestUtils -CONTAINER_OBJECT_PURPOSE = "container" ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV = "LITELLM_ALLOW_UNTRACKED_CONTAINER_ACCESS" -_IN_MEMORY_CONTAINER_OWNERS: Dict[str, str] = {} +MAX_IN_MEMORY_CONTAINER_OWNERS = 10000 +_IN_MEMORY_CONTAINER_OWNERS: "OrderedDict[str, str]" = OrderedDict() def _allow_untracked_container_access() -> bool: @@ -26,6 +31,15 @@ def _allow_untracked_container_access() -> bool: } +def _remember_container_owner(model_object_id: str, owner: str) -> None: + existing_owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) + if existing_owner is not None: + _IN_MEMORY_CONTAINER_OWNERS.move_to_end(model_object_id) + _IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner + while len(_IN_MEMORY_CONTAINER_OWNERS) > MAX_IN_MEMORY_CONTAINER_OWNERS: + _IN_MEMORY_CONTAINER_OWNERS.popitem(last=False) + + def _container_model_object_id( original_container_id: str, custom_llm_provider: str, @@ -124,12 +138,11 @@ async def record_container_owner( existing_owner, user_api_key_dict ): raise HTTPException(status_code=403, detail="Forbidden") - _IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner + _remember_container_owner(model_object_id, owner) return response - existing = await prisma_client.db.litellm_managedobjecttable.find_unique( - where={"model_object_id": model_object_id} - ) + store = ContainerOwnershipStore(prisma_client) + existing = await store.find_by_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") @@ -137,8 +150,8 @@ async def record_container_owner( getattr(existing, "created_by", None), user_api_key_dict ): raise HTTPException(status_code=403, detail="Forbidden") - await prisma_client.db.litellm_managedobjecttable.update( - where={"model_object_id": model_object_id}, + await store.update_owner_record( + model_object_id=model_object_id, data={ "unified_object_id": container_id, "file_object": file_object, @@ -146,7 +159,7 @@ async def record_container_owner( }, ) else: - await prisma_client.db.litellm_managedobjecttable.create( + await store.create_owner_record( data={ "unified_object_id": container_id, "model_object_id": model_object_id, @@ -165,7 +178,12 @@ async def record_container_owner( model_object_id, e, ) - _IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner + 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") + _remember_container_owner(model_object_id, owner) return response @@ -183,14 +201,9 @@ async def _get_container_owner( 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, - } - ) - if row is not None: - return getattr(row, "created_by", None) + owner = await ContainerOwnershipStore(prisma_client).get_owner(model_object_id) + if owner is not None: + return owner return _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) except Exception as e: verbose_proxy_logger.warning( @@ -283,17 +296,12 @@ async def _get_allowed_container_ids( 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}, - } + db_allowed_ids = await ContainerOwnershipStore( + prisma_client + ).list_model_object_ids_for_owners( + owner_scopes=owner_scopes, ) - return in_memory_allowed_ids | { - row.model_object_id - for row in rows - if getattr(row, "model_object_id", None) is not None - } + return in_memory_allowed_ids | db_allowed_ids except Exception as e: verbose_proxy_logger.warning( "Failed to load allowed container ids; falling back to in-process " diff --git a/litellm/proxy/container_endpoints/ownership_store.py b/litellm/proxy/container_endpoints/ownership_store.py new file mode 100644 index 000000000000..b3c406f618d0 --- /dev/null +++ b/litellm/proxy/container_endpoints/ownership_store.py @@ -0,0 +1,55 @@ +from typing import Any, Dict, List, Optional, Set + +CONTAINER_OBJECT_PURPOSE = "container" + + +class ContainerOwnershipStore: + def __init__(self, prisma_client: Any): + self.prisma_client = prisma_client + + @property + def _table(self) -> Any: + return self.prisma_client.db.litellm_managedobjecttable + + async def find_by_model_object_id(self, model_object_id: str) -> Optional[Any]: + return await self._table.find_unique(where={"model_object_id": model_object_id}) + + async def create_owner_record(self, data: Dict[str, Any]) -> None: + await self._table.create(data=data) + + async def update_owner_record( + self, + model_object_id: str, + data: Dict[str, Any], + ) -> None: + await self._table.update( + where={"model_object_id": model_object_id}, + data=data, + ) + + async def get_owner(self, model_object_id: str) -> Optional[str]: + row = await self._table.find_first( + where={ + "model_object_id": model_object_id, + "file_purpose": CONTAINER_OBJECT_PURPOSE, + } + ) + if row is None: + return None + return getattr(row, "created_by", None) + + async def list_model_object_ids_for_owners( + self, + owner_scopes: List[str], + ) -> Set[str]: + rows = await self._table.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 + } diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index fbce5d698597..19682649bb91 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -241,6 +241,29 @@ async def test_should_track_container_owner_in_memory_without_prisma(monkeypatch assert provider == "openai" +@pytest.mark.asyncio +async def test_should_bound_in_memory_container_owner_tracking(monkeypatch): + monkeypatch.setattr(ownership, "MAX_IN_MEMORY_CONTAINER_OWNERS", 2) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=None), + ) + auth = UserAPIKeyAuth(user_id="user-1") + + for container_id in ("cntr_1", "cntr_2", "cntr_3"): + await ownership.record_container_owner( + response=_container(container_id), + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert list(ownership._IN_MEMORY_CONTAINER_OWNERS.keys()) == [ + "container:openai:cntr_2", + "container:openai:cntr_3", + ] + + @pytest.mark.asyncio async def test_should_deny_container_access_for_different_owner(monkeypatch): table = AsyncMock() @@ -847,6 +870,59 @@ async def _handle_llm_api_exception(self, **kwargs): ) +@pytest.mark.asyncio +async def test_should_not_route_owner_record_errors_through_llm_error_handler( + monkeypatch, +): + from litellm.proxy.container_endpoints import endpoints + + proxy_server_stub = SimpleNamespace( + general_settings={}, + llm_router=None, + proxy_config=None, + proxy_logging_obj=None, + select_data_generator=None, + user_api_base=None, + user_max_tokens=None, + user_model=None, + user_request_timeout=None, + user_temperature=None, + version="test", + ) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub) + + class FakeProcessor: + def __init__(self, data): + pass + + async def base_process_llm_request(self, **kwargs): + return _container("cntr_provider") + + async def _handle_llm_api_exception(self, **kwargs): + raise AssertionError("ownership errors should not use LLM error handler") + + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", FakeProcessor) + monkeypatch.setattr( + endpoints, + "record_container_owner", + AsyncMock(side_effect=HTTPException(status_code=403, detail="Forbidden")), + ) + + with pytest.raises(HTTPException) as exc: + await endpoints.create_container( + request=SimpleNamespace( + query_params={}, + headers={}, + json=AsyncMock(return_value={}), + body=AsyncMock(return_value=b"{}"), + ), + fastapi_response=SimpleNamespace(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + ) + + assert exc.value.status_code == 403 + + @pytest.mark.asyncio async def test_should_filter_container_list_inside_list_endpoint(monkeypatch): from litellm.proxy.container_endpoints import endpoints From e9fb89b90c35b28b2ebbb87d18f739b194f12b8b Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 20:44:14 -0700 Subject: [PATCH 15/27] fix(proxy): avoid misleading multi-method operation ids --- litellm/proxy/proxy_server.py | 5 ++-- .../proxy/test_swagger_chat_completions.py | 25 +++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ac7f86b20285..ed8365f95ce2 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -953,8 +953,9 @@ async def _run_pw_migration(): def _generate_stable_operation_id(route: Any) -> str: operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}") - if route.methods: - operation_id = f"{operation_id}_{sorted(route.methods)[0].lower()}" + route_methods = sorted(route.methods or []) + if len(route_methods) == 1: + operation_id = f"{operation_id}_{route_methods[0].lower()}" return operation_id diff --git a/tests/test_litellm/proxy/test_swagger_chat_completions.py b/tests/test_litellm/proxy/test_swagger_chat_completions.py index 2a773bd3168b..36ead2f41c79 100644 --- a/tests/test_litellm/proxy/test_swagger_chat_completions.py +++ b/tests/test_litellm/proxy/test_swagger_chat_completions.py @@ -460,6 +460,31 @@ def test_should_make_duplicate_operation_ids_unique_by_method(self): ] assert len(operation_ids) == len(set(operation_ids)) + def test_should_not_add_method_suffix_to_multi_method_route_base_id(self): + from types import SimpleNamespace + + from litellm.proxy.proxy_server import _generate_stable_operation_id + + multi_method_route = SimpleNamespace( + name="anthropic_proxy_route", + path_format="/anthropic/{endpoint}", + methods={"DELETE", "GET", "POST"}, + ) + single_method_route = SimpleNamespace( + name="list_models", + path_format="/models", + methods={"GET"}, + ) + + assert ( + _generate_stable_operation_id(multi_method_route) + == "anthropic_proxy_route_anthropic__endpoint_" + ) + assert ( + _generate_stable_operation_id(single_method_route) + == "list_models_models_get" + ) + def test_should_reserve_operation_ids_across_lazy_fragments(self): from litellm.proxy.proxy_server import ensure_unique_openapi_operation_ids From d812a5356ea00258d9d1f1c01ebd74570e2cc8d2 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 11:44:33 -0700 Subject: [PATCH 16/27] fix managed container routing after staging merge --- .../container_endpoints/handler_factory.py | 45 +++++-------------- .../test_azure_container_transformation.py | 27 ++++++++++- .../test_container_proxy_ownership.py | 30 ++++++++----- 3 files changed, 57 insertions(+), 45 deletions(-) diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index 5a642e3fc9dd..f4887f812cbd 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -19,10 +19,7 @@ get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) -from litellm.proxy.container_endpoints.ownership import ( - assert_user_can_access_container, - get_container_forwarding_params, -) +from litellm.proxy.container_endpoints.ownership import assert_user_can_access_container def _load_endpoints_config() -> Dict: @@ -188,18 +185,15 @@ async def _process_binary_request( or "openai" ) - original_container_id, custom_llm_provider = await assert_user_can_access_container( + await assert_user_can_access_container( container_id=container_id, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) data: Dict[str, Any] = { + "container_id": container_id, "file_id": file_id, - **get_container_forwarding_params( - container_id, - original_container_id, - custom_llm_provider, - ), + "custom_llm_provider": custom_llm_provider, } processor = ProxyBaseLLMRequestProcessing(data=data) @@ -308,19 +302,14 @@ async def _process_multipart_upload_request( or "openai" ) - original_container_id, custom_llm_provider = await assert_user_can_access_container( + await assert_user_can_access_container( container_id=container_id, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) - data.update( - get_container_forwarding_params( - container_id, - original_container_id, - custom_llm_provider, - ) - ) + data["container_id"] = container_id + data["custom_llm_provider"] = custom_llm_provider processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -387,22 +376,12 @@ async def _process_request( # Validate container_id ownership if present in path_params. if "container_id" in path_params: - original_container_id, custom_llm_provider = ( - await assert_user_can_access_container( - container_id=path_params["container_id"], - user_api_key_dict=user_api_key_dict, - custom_llm_provider=custom_llm_provider, - ) - ) - data.update( - get_container_forwarding_params( - path_params["container_id"], - original_container_id, - custom_llm_provider, - ) + await assert_user_can_access_container( + container_id=path_params["container_id"], + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, ) - else: - data["custom_llm_provider"] = custom_llm_provider + data["custom_llm_provider"] = custom_llm_provider processor = ProxyBaseLLMRequestProcessing(data=data) try: diff --git a/tests/test_litellm/containers/test_azure_container_transformation.py b/tests/test_litellm/containers/test_azure_container_transformation.py index 45fa23bcb6ec..7623c5b0a178 100644 --- a/tests/test_litellm/containers/test_azure_container_transformation.py +++ b/tests/test_litellm/containers/test_azure_container_transformation.py @@ -1,6 +1,6 @@ import os import sys -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock from urllib.parse import parse_qs, urlparse import httpx @@ -13,7 +13,6 @@ from litellm.llms.base_llm.containers.transformation import BaseContainerConfig from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.containers.main import ( - ContainerFileListResponse, ContainerListResponse, ContainerObject, DeleteContainerResult, @@ -556,6 +555,12 @@ async def _mock_base_process_llm_request( "base_process_llm_request", _mock_base_process_llm_request, ) + access_check = AsyncMock(return_value=("cntr_123", "azure")) + monkeypatch.setattr( + handler_factory, + "assert_user_can_access_container", + access_check, + ) request = Request( { @@ -576,6 +581,8 @@ async def _mock_base_process_llm_request( path_params={"container_id": encoded_id}, ) + access_check.assert_awaited_once() + assert access_check.await_args.kwargs["container_id"] == encoded_id assert captured["route_type"] == "alist_container_files" assert captured["data"]["container_id"] == encoded_id assert captured["data"]["custom_llm_provider"] == "openai" @@ -620,6 +627,12 @@ async def _mock_base_process_llm_request( "base_process_llm_request", _mock_base_process_llm_request, ) + access_check = AsyncMock(return_value=("cntr_123", "azure")) + monkeypatch.setattr( + handler_factory, + "assert_user_can_access_container", + access_check, + ) request = Request( { @@ -640,6 +653,8 @@ async def _mock_base_process_llm_request( user_api_key_dict=MagicMock(), ) + access_check.assert_awaited_once() + assert access_check.await_args.kwargs["container_id"] == encoded_id assert captured["route_type"] == "aretrieve_container_file_content" assert captured["data"]["container_id"] == encoded_id assert captured["data"]["file_id"] == "cfile_abc" @@ -700,6 +715,12 @@ async def _mock_base_process_llm_request( "base_process_llm_request", _mock_base_process_llm_request, ) + access_check = AsyncMock(return_value=("cntr_123", "azure")) + monkeypatch.setattr( + handler_factory, + "assert_user_can_access_container", + access_check, + ) request = Request( { @@ -719,6 +740,8 @@ async def _mock_base_process_llm_request( container_id=encoded_id, ) + access_check.assert_awaited_once() + assert access_check.await_args.kwargs["container_id"] == encoded_id assert captured["route_type"] == "aupload_container_file" assert captured["data"]["container_id"] == encoded_id assert captured["data"]["custom_llm_provider"] == "openai" diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 19682649bb91..b696f8fc48f5 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -630,7 +630,9 @@ async def test_should_include_memory_container_list_when_db_recovers_without_row @pytest.mark.asyncio -async def test_should_forward_decoded_container_id_for_proxy_forwarding(monkeypatch): +async def test_should_validate_owner_and_preserve_managed_id_for_proxy_forwarding( + monkeypatch, +): from litellm.proxy.container_endpoints import handler_factory proxy_server_stub = SimpleNamespace( @@ -665,10 +667,11 @@ async def _handle_llm_api_exception(self, **kwargs): "ProxyBaseLLMRequestProcessing", FakeProcessor, ) + access_check = AsyncMock(return_value=("cntr_provider", "azure")) monkeypatch.setattr( handler_factory, "assert_user_can_access_container", - AsyncMock(return_value=("cntr_provider", "azure")), + access_check, ) encoded_id = ResponsesAPIRequestUtils._build_container_id( custom_llm_provider="azure", @@ -684,13 +687,17 @@ async def _handle_llm_api_exception(self, **kwargs): path_params={"container_id": encoded_id}, ) - assert result["container_id"] == "cntr_provider" - assert result["custom_llm_provider"] == "azure" - assert result["model_id"] == "router-gpt" + access_check.assert_awaited_once() + assert access_check.await_args.kwargs["container_id"] == encoded_id + assert result["container_id"] == encoded_id + assert result["custom_llm_provider"] == "openai" + assert "model_id" not in result @pytest.mark.asyncio -async def test_should_forward_decoded_container_id_for_multipart_upload(monkeypatch): +async def test_should_validate_owner_and_preserve_managed_id_for_multipart_upload( + monkeypatch, +): from litellm.proxy.common_utils import http_parsing_utils from litellm.proxy.container_endpoints import handler_factory @@ -726,10 +733,11 @@ async def _handle_llm_api_exception(self, **kwargs): "ProxyBaseLLMRequestProcessing", FakeProcessor, ) + access_check = AsyncMock(return_value=("cntr_provider", "azure")) monkeypatch.setattr( handler_factory, "assert_user_can_access_container", - AsyncMock(return_value=("cntr_provider", "azure")), + access_check, ) monkeypatch.setattr( http_parsing_utils, @@ -755,9 +763,11 @@ async def _handle_llm_api_exception(self, **kwargs): container_id=encoded_id, ) - assert result["container_id"] == "cntr_provider" - assert result["custom_llm_provider"] == "azure" - assert result["model_id"] == "router-gpt" + access_check.assert_awaited_once() + assert access_check.await_args.kwargs["container_id"] == encoded_id + assert result["container_id"] == encoded_id + assert result["custom_llm_provider"] == "openai" + assert "model_id" not in result assert result["file"] == "file-data" From 6c37e9d8d571a4683ad9e068e6b9c16047515495 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Sat, 2 May 2026 02:23:45 +0000 Subject: [PATCH 17/27] fix(proxy): handle ownership-recording failures after upstream create MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit If record_container_owner raises after the upstream container is created, the user previously got a 500 with no usable container — they were billed for an unreachable resource. Move ownership recording into the create path's exception handling and split the two failure modes: - HTTPException from the recorder (auth conflicts) propagates verbatim so the client sees the real status code, not a generic LLM error. - Unexpected exceptions are logged and swallowed; the response is returned to the caller so they aren't billed for a container they can't address. The DB row stays untracked until an operator reconciles. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../proxy/container_endpoints/endpoints.py | 33 +++++++++-- .../test_container_proxy_ownership.py | 59 +++++++++++++++++++ 2 files changed, 86 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 089f22a23a57..7e097802632d 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -2,9 +2,10 @@ from typing import Any, Dict -from fastapi import APIRouter, Depends, Request, Response +from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import ORJSONResponse +from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -129,11 +130,31 @@ 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, - ) + + # Ownership recording sits between the upstream create and the response + # to the caller. The recorder swallows DB errors internally (falling + # back to in-memory tracking) and only surfaces HTTPException on auth + # conflicts; we let those propagate verbatim so the client sees the + # real status code rather than a generic LLM error wrapper. + try: + return await record_container_owner( + response=response, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) + except HTTPException: + raise + except Exception as e: + # Unexpected (non-HTTPException) failure after upstream create. + # The container exists upstream but is now untracked — log loudly + # so an operator can reconcile, and return the response so the + # caller does not get charged for a resource they cannot use. + verbose_proxy_logger.exception( + "Container ownership recording failed after upstream create; " + "returning response with untracked ownership: %s", + e, + ) + return response @router.get( diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index b696f8fc48f5..13216ad7c59f 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -933,6 +933,65 @@ async def _handle_llm_api_exception(self, **kwargs): assert exc.value.status_code == 403 +@pytest.mark.asyncio +async def test_should_return_response_when_owner_recording_raises_unexpected( + monkeypatch, +): + """If record_container_owner raises a non-HTTPException after upstream create, + the upstream container exists but is untracked. The caller still gets the + response (not a 500) so they aren't billed for an unusable resource — an + operator reconciles via logs. + """ + from litellm.proxy.container_endpoints import endpoints + + proxy_server_stub = SimpleNamespace( + general_settings={}, + llm_router=None, + proxy_config=None, + proxy_logging_obj=None, + select_data_generator=None, + user_api_base=None, + user_max_tokens=None, + user_model=None, + user_request_timeout=None, + user_temperature=None, + version="test", + ) + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub) + + created = _container("cntr_provider") + + class FakeProcessor: + def __init__(self, data): + pass + + async def base_process_llm_request(self, **kwargs): + return created + + async def _handle_llm_api_exception(self, **kwargs): + raise AssertionError("upstream-create errors only") + + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", FakeProcessor) + monkeypatch.setattr( + endpoints, + "record_container_owner", + AsyncMock(side_effect=RuntimeError("transient db blip")), + ) + + response = await endpoints.create_container( + request=SimpleNamespace( + query_params={}, + headers={}, + json=AsyncMock(return_value={}), + body=AsyncMock(return_value=b"{}"), + ), + fastapi_response=SimpleNamespace(), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + ) + + assert response is created + + @pytest.mark.asyncio async def test_should_filter_container_list_inside_list_endpoint(monkeypatch): from litellm.proxy.container_endpoints import endpoints From 6194028f798cd49cd9eca3eafa8371fd661365e6 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Sat, 2 May 2026 03:26:58 +0000 Subject: [PATCH 18/27] perf(proxy): cache container/skill ownership reads on the hot path Container ownership and skill rows are looked up on every retrieve / delete / list / file-content / chat-completion-with-skill call. The new stores wrapped raw Prisma queries with no cache, putting one DB round-trip on each request. Add an in-process TTL'd cache mirroring the _byok_cred_cache pattern in mcp_server/server.py: per-key (value, monotonic_timestamp), 60s TTL, 10000-entry cap with full-clear on overflow, invalidated by every write. Negative results (`None`) are cached too so untracked-resource checks also skip the DB. Tests cover: cache-after-first-hit, negative caching, write invalidation, no-caching-on-DB-error, TTL expiry, capacity eviction. 56 tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- litellm/llms/litellm_proxy/skills/handler.py | 62 +++++++- .../proxy/container_endpoints/ownership.py | 52 ++++++- .../test_container_proxy_ownership.py | 133 ++++++++++++++++++ .../litellm_proxy/test_skills_ownership.py | 105 ++++++++++++++ 4 files changed, 342 insertions(+), 10 deletions(-) diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 6eb6f40fc64e..7a6dbb175fb7 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -6,8 +6,9 @@ """ import os +import time import uuid -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple from litellm._logging import verbose_logger from litellm.llms.litellm_proxy.skills.store import LiteLLMSkillsStore @@ -21,6 +22,39 @@ ALLOW_UNOWNED_SKILL_ACCESS_ENV = "LITELLM_ALLOW_UNOWNED_SKILL_ACCESS" +# Skills are looked up on every chat completion that has skills enabled +# (`LiteLLMSkillsHandler.fetch_skill_from_db` in the injection hook). Cache +# the Prisma skill row for a short window so the hot path doesn't issue a DB +# round-trip per request. Same shape as `_byok_cred_cache` and the container +# ownership cache: (value, monotonic_timestamp). `None` is cached as a true +# negative ("skill does not exist") so repeated misses also avoid the DB. +_SKILL_CACHE: Dict[str, Tuple[Optional[Any], float]] = {} +_SKILL_CACHE_TTL = 60 # seconds +_SKILL_CACHE_MAX_SIZE = 10000 + + +def _read_skill_cache(skill_id: str) -> Tuple[bool, Optional[Any]]: + """Return (hit, value). hit=False means caller must consult the DB.""" + entry = _SKILL_CACHE.get(skill_id) + if entry is None: + return False, None + value, timestamp = entry + if time.monotonic() - timestamp > _SKILL_CACHE_TTL: + _SKILL_CACHE.pop(skill_id, None) + return False, None + return True, value + + +def _write_skill_cache(skill_id: str, skill: Optional[Any]) -> None: + if len(_SKILL_CACHE) >= _SKILL_CACHE_MAX_SIZE: + _SKILL_CACHE.clear() + _SKILL_CACHE[skill_id] = (skill, time.monotonic()) + + +def _invalidate_skill_cache(skill_id: str) -> None: + """Drop the cache entry after a write so the next read sees the new row.""" + _SKILL_CACHE.pop(skill_id, None) + def _allow_unowned_skill_access() -> bool: return os.getenv(ALLOW_UNOWNED_SKILL_ACCESS_ENV, "").lower() in { @@ -190,6 +224,24 @@ async def list_skills( return [_prisma_skill_to_litellm(s) for s in skills] + @staticmethod + async def _load_skill(skill_id: str) -> Optional[Any]: + """Cache-first read of the Prisma skill row. + + Caching here keeps `fetch_skill_from_db` (called per chat completion in + the skills injection hook) off the DB. Owner-scope filtering happens + on the cached row, so the cache is per-skill and not per-caller. + """ + cached_hit, cached_skill = _read_skill_cache(skill_id) + if cached_hit: + return cached_skill + + prisma_client = await LiteLLMSkillsHandler._get_prisma_client() + store = LiteLLMSkillsStore(prisma_client) + skill = await store.find_skill(skill_id) + _write_skill_cache(skill_id, skill) + return skill + @staticmethod async def get_skill( skill_id: str, @@ -207,12 +259,9 @@ async def get_skill( Raises: ValueError: If skill not found """ - prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - store = LiteLLMSkillsStore(prisma_client) - verbose_logger.debug(f"LiteLLMSkillsHandler: Getting skill {skill_id}") - skill = await store.find_skill(skill_id) + skill = await LiteLLMSkillsHandler._load_skill(skill_id) if skill is None: raise ValueError(f"Skill not found: {skill_id}") @@ -247,7 +296,7 @@ async def delete_skill( verbose_logger.debug(f"LiteLLMSkillsHandler: Deleting skill {skill_id}") # Check if skill exists - skill = await store.find_skill(skill_id) + skill = await LiteLLMSkillsHandler._load_skill(skill_id) if skill is None: raise ValueError(f"Skill not found: {skill_id}") @@ -259,6 +308,7 @@ async def delete_skill( # Delete the skill await store.delete_skill(skill_id) + _invalidate_skill_cache(skill_id) return {"id": skill_id, "type": "skill_deleted"} diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 86188c02414e..e5577c4052e3 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,4 +1,5 @@ import os +import time from collections import OrderedDict from typing import Any, Dict, List, Optional, Set, Tuple @@ -22,6 +23,38 @@ MAX_IN_MEMORY_CONTAINER_OWNERS = 10000 _IN_MEMORY_CONTAINER_OWNERS: "OrderedDict[str, str]" = OrderedDict() +# Short-lived cache keeps every container access check from hitting the DB +# (`_get_container_owner` is invoked on retrieve / delete / list / file-content +# paths). Mirrors the `_byok_cred_cache` pattern in mcp_server/server.py: +# (value, monotonic_timestamp) tuples, TTL'd, capped, invalidated by writes. +# A `None` value caches "untracked" so repeated negative lookups also avoid DB. +_CONTAINER_OWNER_CACHE: Dict[str, Tuple[Optional[str], float]] = {} +_CONTAINER_OWNER_CACHE_TTL = 60 # seconds +_CONTAINER_OWNER_CACHE_MAX_SIZE = 10000 + + +def _read_container_owner_cache(model_object_id: str) -> Tuple[bool, Optional[str]]: + """Return (hit, value). hit=False means caller must consult the DB.""" + entry = _CONTAINER_OWNER_CACHE.get(model_object_id) + if entry is None: + return False, None + value, timestamp = entry + if time.monotonic() - timestamp > _CONTAINER_OWNER_CACHE_TTL: + _CONTAINER_OWNER_CACHE.pop(model_object_id, None) + return False, None + return True, value + + +def _write_container_owner_cache(model_object_id: str, owner: Optional[str]) -> None: + if len(_CONTAINER_OWNER_CACHE) >= _CONTAINER_OWNER_CACHE_MAX_SIZE: + _CONTAINER_OWNER_CACHE.clear() + _CONTAINER_OWNER_CACHE[model_object_id] = (owner, time.monotonic()) + + +def _invalidate_container_owner_cache(model_object_id: str) -> None: + """Drop a cache entry after a write so the next read sees the new owner.""" + _CONTAINER_OWNER_CACHE.pop(model_object_id, None) + def _allow_untracked_container_access() -> bool: return os.getenv(ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, "").lower() in { @@ -139,6 +172,7 @@ async def record_container_owner( ): raise HTTPException(status_code=403, detail="Forbidden") _remember_container_owner(model_object_id, owner) + _invalidate_container_owner_cache(model_object_id) return response store = ContainerOwnershipStore(prisma_client) @@ -185,6 +219,7 @@ async def record_container_owner( raise HTTPException(status_code=403, detail="Forbidden") _remember_container_owner(model_object_id, owner) + _invalidate_container_owner_cache(model_object_id) return response @@ -196,15 +231,23 @@ async def _get_container_owner( original_container_id, custom_llm_provider, ) + + cached_hit, cached_value = _read_container_owner_cache(model_object_id) + if cached_hit: + return cached_value + try: prisma_client = await _get_prisma_client() if prisma_client is None: - return _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) + owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) + _write_container_owner_cache(model_object_id, owner) + return owner owner = await ContainerOwnershipStore(prisma_client).get_owner(model_object_id) - if owner is not None: - return owner - return _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) + if owner is None: + owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) + _write_container_owner_cache(model_object_id, owner) + return owner except Exception as e: verbose_proxy_logger.warning( "Failed to load container ownership for container_id=%s; " @@ -212,6 +255,7 @@ async def _get_container_owner( model_object_id, e, ) + # Don't cache transient DB errors — let the next request retry. return _IN_MEMORY_CONTAINER_OWNERS.get(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 13216ad7c59f..63e5e16d0d60 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -14,9 +14,11 @@ @pytest.fixture(autouse=True) def clear_in_memory_container_owners(monkeypatch): ownership._IN_MEMORY_CONTAINER_OWNERS.clear() + ownership._CONTAINER_OWNER_CACHE.clear() monkeypatch.delenv(ownership.ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, raising=False) yield ownership._IN_MEMORY_CONTAINER_OWNERS.clear() + ownership._CONTAINER_OWNER_CACHE.clear() def _container(container_id: str) -> ContainerObject: @@ -1102,3 +1104,134 @@ async def _handle_llm_api_exception(self, **kwargs): assert result["container_id"] == "cntr_provider" assert result["custom_llm_provider"] == "azure" assert result["model_id"] == "router-gpt" + + +# ── Cache layer ──────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_get_container_owner_uses_cache_after_first_db_hit(monkeypatch): + """Repeated access checks within the TTL window must not hit the DB. + + Greptile's P1 was that ownership reads issued a Prisma query on every + request. The cache here mirrors `_byok_cred_cache`: TTL'd, capped, and + invalidated on writes. + """ + table = AsyncMock() + fake_row = SimpleNamespace( + created_by="user-1", file_purpose=ownership.CONTAINER_OBJECT_PURPOSE + ) + table.find_first.return_value = fake_row + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + + owner_first = await ownership._get_container_owner("cntr_x", "openai") + owner_second = await ownership._get_container_owner("cntr_x", "openai") + owner_third = await ownership._get_container_owner("cntr_x", "openai") + + assert owner_first == "user-1" + assert owner_second == "user-1" + assert owner_third == "user-1" + # Single DB call across three reads — the cache absorbs the rest. + assert table.find_first.await_count == 1 + + +@pytest.mark.asyncio +async def test_get_container_owner_caches_negative_lookups(monkeypatch): + """`None` (untracked) must also be cached so repeated misses don't query.""" + table = AsyncMock() + table.find_first.return_value = None + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + + 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_record_container_owner_invalidates_cache(monkeypatch): + """A recorded owner must drop the cached value so the next read re-fetches. + + Otherwise a stale `None` from a prior negative lookup would survive the + create and the new owner would be invisible until the TTL elapses. + """ + # Seed the cache with a stale negative result. + ownership._write_container_owner_cache("container:openai:cntr_new", None) + cached_hit, cached_value = ownership._read_container_owner_cache( + "container:openai:cntr_new" + ) + assert cached_hit and cached_value is None + + 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), + ) + + await ownership.record_container_owner( + response=_container("cntr_new"), + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + custom_llm_provider="openai", + ) + + # Invalidation drops the entry — next read goes to the DB. + cached_hit, _ = ownership._read_container_owner_cache("container:openai:cntr_new") + assert not cached_hit + + +@pytest.mark.asyncio +async def test_get_container_owner_does_not_cache_on_db_error(monkeypatch): + """DB errors must skip caching so transient failures don't pin a `None`.""" + 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), + ) + + result = await ownership._get_container_owner("cntr_x", "openai") + assert result is None + cached_hit, _ = ownership._read_container_owner_cache("container:openai:cntr_x") + assert not cached_hit + + +def test_container_owner_cache_expires_after_ttl(monkeypatch): + """Entries past the TTL count as misses so writes elsewhere are eventually + visible to this process.""" + monkeypatch.setattr(ownership, "_CONTAINER_OWNER_CACHE_TTL", 0.0) + ownership._write_container_owner_cache("k", "user-1") + cached_hit, _ = ownership._read_container_owner_cache("k") + # TTL of 0 means anything in the cache is already stale. + assert not cached_hit + + +def test_container_owner_cache_evicts_when_at_capacity(monkeypatch): + """The cache must not grow unbounded; reaching capacity clears all entries.""" + monkeypatch.setattr(ownership, "_CONTAINER_OWNER_CACHE_MAX_SIZE", 2) + ownership._write_container_owner_cache("a", "user-a") + ownership._write_container_owner_cache("b", "user-b") + ownership._write_container_owner_cache("c", "user-c") + # Reaching the cap clears everything — the new write is the only survivor. + assert ownership._CONTAINER_OWNER_CACHE.keys() == {"c"} 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 45ed172a8451..f98a04205b40 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -20,6 +20,9 @@ @pytest.fixture(autouse=True) def clear_skill_ownership_env(monkeypatch): monkeypatch.delenv(skills_handler.ALLOW_UNOWNED_SKILL_ACCESS_ENV, raising=False) + skills_handler._SKILL_CACHE.clear() + yield + skills_handler._SKILL_CACHE.clear() def _skill(skill_id: str, created_by: str | None) -> LiteLLM_SkillsTable: @@ -433,3 +436,105 @@ async def test_should_scope_skill_injection_fetch_to_authenticated_user(monkeypa "litellm_skill_other", user_api_key_dict=auth, ) + + +# ── Cache layer ──────────────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_load_skill_uses_cache_after_first_db_hit(monkeypatch): + """`fetch_skill_from_db` is hit per-chat-completion; the cache absorbs + repeats so we don't issue a Prisma query on every request.""" + fake_skill = Mock(created_by="user-1", skill_id="litellm_skill_a") + store_factory = AsyncMock() + store = Mock() + store.find_skill = AsyncMock(return_value=fake_skill) + monkeypatch.setattr( + skills_handler.LiteLLMSkillsHandler, + "_get_prisma_client", + AsyncMock(return_value=store_factory), + ) + monkeypatch.setattr( + skills_handler, + "LiteLLMSkillsStore", + Mock(return_value=store), + ) + + first = await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") + second = await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") + third = await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") + + assert first is fake_skill + assert second is fake_skill + assert third is fake_skill + assert store.find_skill.await_count == 1 + + +@pytest.mark.asyncio +async def test_load_skill_caches_negative_lookups(monkeypatch): + """Missing skills must cache as `None` so repeated lookups skip the DB.""" + store_factory = AsyncMock() + store = Mock() + store.find_skill = AsyncMock(return_value=None) + monkeypatch.setattr( + skills_handler.LiteLLMSkillsHandler, + "_get_prisma_client", + AsyncMock(return_value=store_factory), + ) + monkeypatch.setattr( + skills_handler, + "LiteLLMSkillsStore", + Mock(return_value=store), + ) + + assert await skills_handler.LiteLLMSkillsHandler._load_skill("missing") is None + assert await skills_handler.LiteLLMSkillsHandler._load_skill("missing") is None + assert store.find_skill.await_count == 1 + + +@pytest.mark.asyncio +async def test_delete_skill_invalidates_cache(monkeypatch): + """After delete, the next read must consult the DB rather than the cached + pre-delete row.""" + fake_skill = Mock(created_by="user-1", skill_id="litellm_skill_a") + store = Mock() + store.find_skill = AsyncMock(return_value=fake_skill) + store.delete_skill = AsyncMock() + monkeypatch.setattr( + skills_handler.LiteLLMSkillsHandler, + "_get_prisma_client", + AsyncMock(return_value=Mock()), + ) + monkeypatch.setattr( + skills_handler, + "LiteLLMSkillsStore", + Mock(return_value=store), + ) + + # Prime the cache via the read path. + await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") + cached_hit, _ = skills_handler._read_skill_cache("litellm_skill_a") + assert cached_hit + + auth = UserAPIKeyAuth(user_id="user-1") + await skills_handler.LiteLLMSkillsHandler.delete_skill( + "litellm_skill_a", user_api_key_dict=auth + ) + + cached_hit_after, _ = skills_handler._read_skill_cache("litellm_skill_a") + assert not cached_hit_after + + +def test_skill_cache_expires_after_ttl(monkeypatch): + monkeypatch.setattr(skills_handler, "_SKILL_CACHE_TTL", 0.0) + skills_handler._write_skill_cache("k", Mock()) + cached_hit, _ = skills_handler._read_skill_cache("k") + assert not cached_hit + + +def test_skill_cache_evicts_when_at_capacity(monkeypatch): + monkeypatch.setattr(skills_handler, "_SKILL_CACHE_MAX_SIZE", 2) + skills_handler._write_skill_cache("a", Mock()) + skills_handler._write_skill_cache("b", Mock()) + skills_handler._write_skill_cache("c", Mock()) + assert skills_handler._SKILL_CACHE.keys() == {"c"} From c01f20972354c625257ab797a6d66d4a37f49adb Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Fri, 1 May 2026 21:50:57 -0700 Subject: [PATCH 19/27] fix(proxy): forward decoded container ids after ownership checks --- .../container_endpoints/handler_factory.py | 50 +++++++++++++------ .../test_azure_container_transformation.py | 18 ++++--- .../test_container_proxy_ownership.py | 16 +++--- 3 files changed, 54 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index f4887f812cbd..4284cdd5d4a5 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -19,7 +19,10 @@ get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) -from litellm.proxy.container_endpoints.ownership import assert_user_can_access_container +from litellm.proxy.container_endpoints.ownership import ( + assert_user_can_access_container, + get_container_forwarding_params, +) def _load_endpoints_config() -> Dict: @@ -162,8 +165,9 @@ async def _process_binary_request( """ Process binary content requests through the standard proxy/router pipeline. - The router owns managed container ID decoding and deployment selection. This - handler only adapts the byte response to FastAPI. + Validate ownership before forwarding the provider-native container id through + the standard proxy/router pipeline. This handler only adapts the byte + response to FastAPI. """ from litellm.proxy.proxy_server import ( general_settings, @@ -185,15 +189,18 @@ async def _process_binary_request( or "openai" ) - await assert_user_can_access_container( + original_container_id, resolved_provider = await assert_user_can_access_container( container_id=container_id, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) data: Dict[str, Any] = { - "container_id": container_id, "file_id": file_id, - "custom_llm_provider": custom_llm_provider, + **get_container_forwarding_params( + container_id=container_id, + original_container_id=original_container_id, + custom_llm_provider=resolved_provider, + ), } processor = ProxyBaseLLMRequestProcessing(data=data) @@ -302,14 +309,19 @@ async def _process_multipart_upload_request( or "openai" ) - await assert_user_can_access_container( + original_container_id, resolved_provider = await assert_user_can_access_container( container_id=container_id, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) - data["container_id"] = container_id - data["custom_llm_provider"] = custom_llm_provider + data.update( + get_container_forwarding_params( + container_id=container_id, + original_container_id=original_container_id, + custom_llm_provider=resolved_provider, + ) + ) processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -376,12 +388,22 @@ async def _process_request( # Validate container_id ownership if present in path_params. if "container_id" in path_params: - await assert_user_can_access_container( - container_id=path_params["container_id"], - user_api_key_dict=user_api_key_dict, - custom_llm_provider=custom_llm_provider, + original_container_id, resolved_provider = ( + await assert_user_can_access_container( + container_id=path_params["container_id"], + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) + ) + data.update( + get_container_forwarding_params( + container_id=path_params["container_id"], + original_container_id=original_container_id, + custom_llm_provider=resolved_provider, + ) ) - data["custom_llm_provider"] = custom_llm_provider + else: + data["custom_llm_provider"] = custom_llm_provider processor = ProxyBaseLLMRequestProcessing(data=data) try: diff --git a/tests/test_litellm/containers/test_azure_container_transformation.py b/tests/test_litellm/containers/test_azure_container_transformation.py index 9a3a8bc20a94..70181f6f03df 100644 --- a/tests/test_litellm/containers/test_azure_container_transformation.py +++ b/tests/test_litellm/containers/test_azure_container_transformation.py @@ -543,7 +543,7 @@ def test_regression_proxy_resolves_azure_text_same_as_azure(self): assert isinstance(c1, AzureContainerConfig) @pytest.mark.asyncio - async def test_proxy_process_request_preserves_managed_container_id( + async def test_proxy_process_request_forwards_decoded_container_id( self, monkeypatch ): from starlette.requests import Request @@ -607,9 +607,9 @@ async def _mock_base_process_llm_request( access_check.assert_awaited_once() assert access_check.await_args.kwargs["container_id"] == encoded_id assert captured["route_type"] == "alist_container_files" - assert captured["data"]["container_id"] == encoded_id - assert captured["data"]["custom_llm_provider"] == "openai" - assert "model_id" not in captured["data"] + assert captured["data"]["container_id"] == "cntr_123" + assert captured["data"]["custom_llm_provider"] == "azure" + assert captured["data"]["model_id"] == "model_abc123" assert "api_base" not in captured["data"] @pytest.mark.asyncio @@ -679,9 +679,10 @@ async def _mock_base_process_llm_request( access_check.assert_awaited_once() assert access_check.await_args.kwargs["container_id"] == encoded_id assert captured["route_type"] == "aretrieve_container_file_content" - assert captured["data"]["container_id"] == encoded_id + assert captured["data"]["container_id"] == "cntr_123" assert captured["data"]["file_id"] == "cfile_abc" - assert captured["data"]["custom_llm_provider"] == "openai" + assert captured["data"]["custom_llm_provider"] == "azure" + assert captured["data"]["model_id"] == "model_abc123" assert response.status_code == 200 assert response.body == b"csv-bytes" assert response.headers["x-litellm-call-id"] == "call-123" @@ -766,5 +767,6 @@ async def _mock_base_process_llm_request( access_check.assert_awaited_once() assert access_check.await_args.kwargs["container_id"] == encoded_id assert captured["route_type"] == "aupload_container_file" - assert captured["data"]["container_id"] == encoded_id - assert captured["data"]["custom_llm_provider"] == "openai" + assert captured["data"]["container_id"] == "cntr_123" + assert captured["data"]["custom_llm_provider"] == "azure" + assert captured["data"]["model_id"] == "model_abc123" diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 63e5e16d0d60..2a760edf7db5 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -632,7 +632,7 @@ async def test_should_include_memory_container_list_when_db_recovers_without_row @pytest.mark.asyncio -async def test_should_validate_owner_and_preserve_managed_id_for_proxy_forwarding( +async def test_should_validate_owner_and_forward_decoded_id_for_proxy_forwarding( monkeypatch, ): from litellm.proxy.container_endpoints import handler_factory @@ -691,13 +691,13 @@ async def _handle_llm_api_exception(self, **kwargs): access_check.assert_awaited_once() assert access_check.await_args.kwargs["container_id"] == encoded_id - assert result["container_id"] == encoded_id - assert result["custom_llm_provider"] == "openai" - assert "model_id" not in result + assert result["container_id"] == "cntr_provider" + assert result["custom_llm_provider"] == "azure" + assert result["model_id"] == "router-gpt" @pytest.mark.asyncio -async def test_should_validate_owner_and_preserve_managed_id_for_multipart_upload( +async def test_should_validate_owner_and_forward_decoded_id_for_multipart_upload( monkeypatch, ): from litellm.proxy.common_utils import http_parsing_utils @@ -767,9 +767,9 @@ async def _handle_llm_api_exception(self, **kwargs): access_check.assert_awaited_once() assert access_check.await_args.kwargs["container_id"] == encoded_id - assert result["container_id"] == encoded_id - assert result["custom_llm_provider"] == "openai" - assert "model_id" not in result + assert result["container_id"] == "cntr_provider" + assert result["custom_llm_provider"] == "azure" + assert result["model_id"] == "router-gpt" assert result["file"] == "file-data" From 4fa577810ba1c6b3258e4242bd662e72b8ea7e8c Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Mon, 4 May 2026 22:43:18 +0000 Subject: [PATCH 20/27] fix(container): keep ownership-filter exceptions out of the LLM-error path filter_container_list_response runs after the upstream call has already succeeded; treating an ownership-lookup failure as an LLM-API error fires post_call_failure_hook for a successful upstream call and returns a misleading provider-shaped error to the client. Run the filter outside the try/except so genuine LLM errors stay scoped to the upstream call. --- litellm/proxy/container_endpoints/endpoints.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 7e097802632d..9650604bf81a 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -241,11 +241,6 @@ 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, @@ -254,6 +249,16 @@ async def list_containers( version=version, ) + # Ownership filtering runs OUTSIDE the LLM-exception scope: a DB error + # in the ownership lookup is not an LLM-API error and shouldn't be + # translated to a provider-shaped failure (which would also fire the + # post_call_failure_hook for what is in fact a successful upstream call). + return await filter_container_list_response( + response=response, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) + @router.get( "/v1/containers/{container_id}", From ec9b84d38c7053be561d387f03b33207482e4eb9 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Mon, 4 May 2026 22:52:54 +0000 Subject: [PATCH 21/27] chore(container,skills): LRU eviction for owner caches; widen file_purpose Literal Two cleanups from the /simplify pass: * ``_CONTAINER_OWNER_CACHE`` and ``_SKILL_CACHE`` now LRU-evict via ``OrderedDict.popitem(last=False)`` instead of full ``clear()`` at capacity. Full clears converted a steady-state cached workload into a periodic full-DB-load oscillation as the cache repopulated from zero and cleared again. Reads now ``move_to_end`` so the just-touched entry survives the next eviction. Mirrors the pre-existing LRU pattern in ``_remember_container_owner``. * ``LiteLLM_ManagedObjectTable.file_purpose`` Literal now includes ``"container"`` so Pydantic validation accepts rows written by the ownership store. --- litellm/llms/litellm_proxy/skills/handler.py | 12 ++++++++--- litellm/proxy/_types.py | 2 +- .../proxy/container_endpoints/ownership.py | 18 ++++++++++++----- .../test_container_proxy_ownership.py | 20 ++++++++++++++++--- .../litellm_proxy/test_skills_ownership.py | 3 ++- 5 files changed, 42 insertions(+), 13 deletions(-) diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 7a6dbb175fb7..e5caf3cffc7e 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -8,6 +8,7 @@ import os import time import uuid +from collections import OrderedDict from typing import Any, Dict, List, Optional, Tuple from litellm._logging import verbose_logger @@ -28,7 +29,7 @@ # round-trip per request. Same shape as `_byok_cred_cache` and the container # ownership cache: (value, monotonic_timestamp). `None` is cached as a true # negative ("skill does not exist") so repeated misses also avoid the DB. -_SKILL_CACHE: Dict[str, Tuple[Optional[Any], float]] = {} +_SKILL_CACHE: "OrderedDict[str, Tuple[Optional[Any], float]]" = OrderedDict() _SKILL_CACHE_TTL = 60 # seconds _SKILL_CACHE_MAX_SIZE = 10000 @@ -42,13 +43,18 @@ def _read_skill_cache(skill_id: str) -> Tuple[bool, Optional[Any]]: if time.monotonic() - timestamp > _SKILL_CACHE_TTL: _SKILL_CACHE.pop(skill_id, None) return False, None + _SKILL_CACHE.move_to_end(skill_id) return True, value def _write_skill_cache(skill_id: str, skill: Optional[Any]) -> None: - if len(_SKILL_CACHE) >= _SKILL_CACHE_MAX_SIZE: - _SKILL_CACHE.clear() + # LRU eviction (popitem(last=False)) instead of full ``clear()`` — + # see container ownership cache for rationale. + if skill_id in _SKILL_CACHE: + _SKILL_CACHE.move_to_end(skill_id) _SKILL_CACHE[skill_id] = (skill, time.monotonic()) + while len(_SKILL_CACHE) > _SKILL_CACHE_MAX_SIZE: + _SKILL_CACHE.popitem(last=False) def _invalidate_skill_cache(skill_id: str) -> None: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3cca23f07ab0..df0a2bcee5f1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -4609,7 +4609,7 @@ class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): unified_object_id: str model_object_id: str - file_purpose: Literal["batch", "fine-tune", "response"] + file_purpose: Literal["batch", "fine-tune", "response", "container"] file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse] diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index e5577c4052e3..4905c6537db0 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -26,9 +26,10 @@ # Short-lived cache keeps every container access check from hitting the DB # (`_get_container_owner` is invoked on retrieve / delete / list / file-content # paths). Mirrors the `_byok_cred_cache` pattern in mcp_server/server.py: -# (value, monotonic_timestamp) tuples, TTL'd, capped, invalidated by writes. -# A `None` value caches "untracked" so repeated negative lookups also avoid DB. -_CONTAINER_OWNER_CACHE: Dict[str, Tuple[Optional[str], float]] = {} +# (value, monotonic_timestamp) tuples, TTL'd, LRU-evicted at capacity, +# invalidated by writes. A ``None`` value caches "untracked" so repeated +# negative lookups also avoid DB. +_CONTAINER_OWNER_CACHE: "OrderedDict[str, Tuple[Optional[str], float]]" = OrderedDict() _CONTAINER_OWNER_CACHE_TTL = 60 # seconds _CONTAINER_OWNER_CACHE_MAX_SIZE = 10000 @@ -42,13 +43,20 @@ def _read_container_owner_cache(model_object_id: str) -> Tuple[bool, Optional[st if time.monotonic() - timestamp > _CONTAINER_OWNER_CACHE_TTL: _CONTAINER_OWNER_CACHE.pop(model_object_id, None) return False, None + _CONTAINER_OWNER_CACHE.move_to_end(model_object_id) return True, value def _write_container_owner_cache(model_object_id: str, owner: Optional[str]) -> None: - if len(_CONTAINER_OWNER_CACHE) >= _CONTAINER_OWNER_CACHE_MAX_SIZE: - _CONTAINER_OWNER_CACHE.clear() + # LRU eviction (popitem(last=False)) instead of full ``clear()`` — a + # full clear at capacity converts a steady-state cached workload into + # a periodic full-DB-load oscillation as the cache repopulates from + # zero and clears again. + if model_object_id in _CONTAINER_OWNER_CACHE: + _CONTAINER_OWNER_CACHE.move_to_end(model_object_id) _CONTAINER_OWNER_CACHE[model_object_id] = (owner, time.monotonic()) + while len(_CONTAINER_OWNER_CACHE) > _CONTAINER_OWNER_CACHE_MAX_SIZE: + _CONTAINER_OWNER_CACHE.popitem(last=False) def _invalidate_container_owner_cache(model_object_id: str) -> None: diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 2a760edf7db5..2d175ceedb63 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -1228,10 +1228,24 @@ def test_container_owner_cache_expires_after_ttl(monkeypatch): def test_container_owner_cache_evicts_when_at_capacity(monkeypatch): - """The cache must not grow unbounded; reaching capacity clears all entries.""" + """The cache must not grow unbounded; reaching capacity LRU-evicts the + oldest entry, not the entire cache.""" monkeypatch.setattr(ownership, "_CONTAINER_OWNER_CACHE_MAX_SIZE", 2) ownership._write_container_owner_cache("a", "user-a") ownership._write_container_owner_cache("b", "user-b") ownership._write_container_owner_cache("c", "user-c") - # Reaching the cap clears everything — the new write is the only survivor. - assert ownership._CONTAINER_OWNER_CACHE.keys() == {"c"} + # ``a`` was the oldest and is dropped; ``b`` and ``c`` survive. + assert list(ownership._CONTAINER_OWNER_CACHE.keys()) == ["b", "c"] + + +def test_container_owner_cache_read_marks_as_recently_used(monkeypatch): + """Reading an entry should reset its position so a subsequent eviction + drops a less-recently-used entry instead of the just-touched one.""" + monkeypatch.setattr(ownership, "_CONTAINER_OWNER_CACHE_MAX_SIZE", 2) + ownership._write_container_owner_cache("a", "user-a") + ownership._write_container_owner_cache("b", "user-b") + # Touch ``a`` so it becomes the most-recently-used. + ownership._read_container_owner_cache("a") + ownership._write_container_owner_cache("c", "user-c") + # ``b`` is the LRU at this point; ``a`` and ``c`` survive. + assert list(ownership._CONTAINER_OWNER_CACHE.keys()) == ["a", "c"] 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 f98a04205b40..d6c0601a57cd 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -537,4 +537,5 @@ def test_skill_cache_evicts_when_at_capacity(monkeypatch): skills_handler._write_skill_cache("a", Mock()) skills_handler._write_skill_cache("b", Mock()) skills_handler._write_skill_cache("c", Mock()) - assert skills_handler._SKILL_CACHE.keys() == {"c"} + # ``a`` was the oldest and is LRU-evicted; ``b`` and ``c`` survive. + assert list(skills_handler._SKILL_CACHE.keys()) == ["b", "c"] From de682c810e406f05be019e68b941550b91fff9b0 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Mon, 4 May 2026 23:22:19 +0000 Subject: [PATCH 22/27] chore(container,skills): drop legacy-access opt-out env vars MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit LITELLM_ALLOW_UNTRACKED_CONTAINER_ACCESS and LITELLM_ALLOW_UNOWNED_SKILL_ACCESS were operator-toggleable opt-outs for the cross-tenant access primitive this PR closes — flipping either on re-enabled exactly the VERIA-20 read path. Default-secure with no escape hatch matches sibling fixes (vector-store cred isolation, semantic cache key isolation, user_config strip): all rejected the opt-out-of-security pattern. Untracked containers and unowned skills (rows that pre-date this enforcement) are admin-only. Non-admin owners need to either re-create via the now-tracked flow or have an admin assign ``created_by`` on the existing row. Update tests to assert the strict-only behaviour. --- litellm/llms/litellm_proxy/skills/handler.py | 34 +++-------------- .../proxy/container_endpoints/ownership.py | 21 ++--------- .../test_container_proxy_ownership.py | 26 +++++++------ .../litellm_proxy/test_skills_ownership.py | 37 +++++++++---------- 4 files changed, 40 insertions(+), 78 deletions(-) diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index e5caf3cffc7e..086dccd40322 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -5,7 +5,6 @@ Used by the transformation layer and skills injection hook. """ -import os import time import uuid from collections import OrderedDict @@ -21,8 +20,6 @@ user_can_access_resource_owner, ) -ALLOW_UNOWNED_SKILL_ACCESS_ENV = "LITELLM_ALLOW_UNOWNED_SKILL_ACCESS" - # Skills are looked up on every chat completion that has skills enabled # (`LiteLLMSkillsHandler.fetch_skill_from_db` in the injection hook). Cache # the Prisma skill row for a short window so the hot path doesn't issue a DB @@ -62,27 +59,14 @@ def _invalidate_skill_cache(skill_id: str) -> None: _SKILL_CACHE.pop(skill_id, None) -def _allow_unowned_skill_access() -> bool: - return os.getenv(ALLOW_UNOWNED_SKILL_ACCESS_ENV, "").lower() in { - "1", - "true", - "yes", - } - - def _user_can_access_skill_owner( owner: Optional[str], user_api_key_dict: Optional[UserAPIKeyAuth], ) -> bool: - if owner is None and user_api_key_dict is not None: - if is_proxy_admin(user_api_key_dict): - return True - if _allow_unowned_skill_access(): - verbose_logger.warning( - "Allowing unowned skill access because %s is enabled", - ALLOW_UNOWNED_SKILL_ACCESS_ENV, - ) - return True + # Pre-isolation skills with no ``created_by`` are admin-only — same + # rule as untracked containers. Owners need to either re-create via + # the now-tracked flow or have an admin assign ``created_by`` on the + # row. return user_can_access_resource_owner(owner, user_api_key_dict) @@ -216,15 +200,7 @@ async def list_skills( owner_scopes = get_resource_owner_scopes(user_api_key_dict) if not owner_scopes: return [] - if _allow_unowned_skill_access(): - find_many_kwargs["where"] = { - "OR": [ - {"created_by": {"in": owner_scopes}}, - {"created_by": None}, - ] - } - else: - find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} + find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} skills = await store.list_skills(find_many_kwargs) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 4905c6537db0..6bfea87ebf64 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,4 +1,3 @@ -import os import time from collections import OrderedDict from typing import Any, Dict, List, Optional, Set, Tuple @@ -19,7 +18,6 @@ ) from litellm.responses.utils import ResponsesAPIRequestUtils -ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV = "LITELLM_ALLOW_UNTRACKED_CONTAINER_ACCESS" MAX_IN_MEMORY_CONTAINER_OWNERS = 10000 _IN_MEMORY_CONTAINER_OWNERS: "OrderedDict[str, str]" = OrderedDict() @@ -64,14 +62,6 @@ def _invalidate_container_owner_cache(model_object_id: str) -> None: _CONTAINER_OWNER_CACHE.pop(model_object_id, None) -def _allow_untracked_container_access() -> bool: - return os.getenv(ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, "").lower() in { - "1", - "true", - "yes", - } - - def _remember_container_owner(model_object_id: str, owner: str) -> None: existing_owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) if existing_owner is not None: @@ -280,14 +270,11 @@ async def assert_user_can_access_container( if is_proxy_admin(user_api_key_dict): return original_container_id, resolved_provider + # Untracked containers (no ownership row) are admin-only. Pre-isolation + # rows that pre-date this enforcement need the admin to either re-create + # via the now-tracked flow or explicitly assign ``created_by`` on the + # ``litellm_managedobjecttable`` row. owner = await _get_container_owner(original_container_id, resolved_provider) - if owner is None and _allow_untracked_container_access(): - verbose_proxy_logger.warning( - "Allowing untracked container access because %s is enabled", - ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, - ) - return original_container_id, resolved_provider - if not user_can_access_resource_owner(owner, user_api_key_dict): raise HTTPException(status_code=403, detail="Forbidden") diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 2d175ceedb63..2e82b5b6b2d3 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -12,10 +12,9 @@ @pytest.fixture(autouse=True) -def clear_in_memory_container_owners(monkeypatch): +def clear_in_memory_container_owners(): ownership._IN_MEMORY_CONTAINER_OWNERS.clear() ownership._CONTAINER_OWNER_CACHE.clear() - monkeypatch.delenv(ownership.ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, raising=False) yield ownership._IN_MEMORY_CONTAINER_OWNERS.clear() ownership._CONTAINER_OWNER_CACHE.clear() @@ -384,23 +383,26 @@ async def test_should_fail_closed_when_owner_lookup_fails_without_memory(monkeyp @pytest.mark.asyncio -async def test_should_allow_untracked_container_access_when_enabled(monkeypatch): +async def test_untracked_container_is_admin_only(monkeypatch): + """Pre-isolation containers with no ownership row are admin-only. + Non-admin callers see them as 403, with no opt-out flag re-opening + the cross-tenant access primitive.""" + from fastapi import HTTPException + monkeypatch.setattr( ownership, "_get_prisma_client", AsyncMock(return_value=None), ) - monkeypatch.setenv(ownership.ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, "true") auth = UserAPIKeyAuth(user_id="user-1") - original_id, provider = await ownership.assert_user_can_access_container( - container_id="cntr_untracked", - user_api_key_dict=auth, - custom_llm_provider="openai", - ) - - assert original_id == "cntr_untracked" - assert provider == "openai" + with pytest.raises(HTTPException) as exc: + await ownership.assert_user_can_access_container( + container_id="cntr_untracked", + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + assert exc.value.status_code == 403 @pytest.mark.asyncio 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 d6c0601a57cd..cef5f8a554b8 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -18,8 +18,7 @@ @pytest.fixture(autouse=True) -def clear_skill_ownership_env(monkeypatch): - monkeypatch.delenv(skills_handler.ALLOW_UNOWNED_SKILL_ACCESS_ENV, raising=False) +def clear_skill_cache(): skills_handler._SKILL_CACHE.clear() yield skills_handler._SKILL_CACHE.clear() @@ -358,7 +357,10 @@ async def test_should_hide_unowned_skill_by_default(monkeypatch): @pytest.mark.asyncio -async def test_should_allow_unowned_skill_when_enabled(monkeypatch): +async def test_unowned_skill_is_admin_only(monkeypatch): + """Pre-isolation skills with no ``created_by`` are admin-only — non-admin + callers see the same "not found" they'd see for a missing row, with no + opt-out env var that re-opens the cross-tenant access primitive.""" table = AsyncMock() table.find_unique.return_value = _skill("litellm_skill_unowned", None) prisma_client = type( @@ -369,22 +371,22 @@ async def test_should_allow_unowned_skill_when_enabled(monkeypatch): "_get_prisma_client", AsyncMock(return_value=prisma_client), ) - monkeypatch.setenv(skills_handler.ALLOW_UNOWNED_SKILL_ACCESS_ENV, "true") auth = UserAPIKeyAuth(user_id="user-1") - skill = await LiteLLMSkillsHandler.get_skill( - "litellm_skill_unowned", - user_api_key_dict=auth, - ) - - assert skill.skill_id == "litellm_skill_unowned" + with pytest.raises(ValueError, match="Skill not found"): + await LiteLLMSkillsHandler.get_skill( + "litellm_skill_unowned", + user_api_key_dict=auth, + ) @pytest.mark.asyncio -async def test_should_include_unowned_skills_in_list_when_enabled(monkeypatch): +async def test_list_skills_excludes_unowned_for_non_admin(monkeypatch): + """Non-admin list queries scope to ``created_by IN owner_scopes``; rows + with ``created_by IS NULL`` are excluded — admin-only.""" table = AsyncMock() - table.find_many.return_value = [_skill("litellm_skill_unowned", None)] + table.find_many.return_value = [] prisma_client = type( "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} )() @@ -393,18 +395,13 @@ async def test_should_include_unowned_skills_in_list_when_enabled(monkeypatch): "_get_prisma_client", AsyncMock(return_value=prisma_client), ) - monkeypatch.setenv(skills_handler.ALLOW_UNOWNED_SKILL_ACCESS_ENV, "true") auth = UserAPIKeyAuth(user_id="user-1") + await LiteLLMSkillsHandler.list_skills(user_api_key_dict=auth) - skills = await LiteLLMSkillsHandler.list_skills(user_api_key_dict=auth) - - assert [skill.skill_id for skill in skills] == ["litellm_skill_unowned"] where = table.find_many.await_args.kwargs["where"] - assert where["OR"] == [ - {"created_by": {"in": ["user-1", "user:user-1"]}}, - {"created_by": None}, - ] + # No OR fallback to ``created_by IS NULL`` — strict scope only. + assert where == {"created_by": {"in": ["user-1", "user:user-1"]}} @pytest.mark.asyncio From 758b4883269e8830dc9809789b36604f8afefab7 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Mon, 4 May 2026 23:40:22 +0000 Subject: [PATCH 23/27] fix(ownership): reject identity-less callers instead of sharing a sentinel scope MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit UNSCOPED_RESOURCE_OWNER_SCOPE collapsed every caller without an identity field (no user_id / team_id / org_id / api_key / token) into a single shared owner — a cross-tenant access primitive: any two such callers could see and delete each other's containers and skills. Drop the sentinel. ``get_primary_resource_owner_scope`` returns ``None`` and ``get_resource_owner_scopes`` returns ``[]`` for identity-less callers. ``record_container_owner`` and ``LiteLLMSkillsHandler.create_skill`` now reject creates from identity-less callers with a 403 instead of stamping the placeholder. Read paths already deny ``owner is None`` correctly so legacy rows (if any) are admin-only. --- litellm/llms/litellm_proxy/skills/handler.py | 11 ++++++ .../proxy/common_utils/resource_ownership.py | 23 +++++++++---- .../proxy/container_endpoints/ownership.py | 10 +++++- .../test_container_proxy_ownership.py | 33 ++++++++---------- .../litellm_proxy/test_skills_ownership.py | 34 +++++++++++-------- 5 files changed, 70 insertions(+), 41 deletions(-) diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 086dccd40322..1129af64e14e 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -133,6 +133,17 @@ async def create_skill( skill_id = f"litellm_skill_{uuid.uuid4()}" owner = get_primary_resource_owner_scope(user_api_key_dict) or user_id + if owner is None: + # Caller has no identity scope (no user_id / team_id / org_id / + # api_key / token). Stamping a placeholder would let any two + # identity-less callers see each other's skills via the shared + # owner — the cross-tenant primitive we avoid. + from fastapi import HTTPException + + raise HTTPException( + status_code=403, + detail="Unable to record skill ownership: caller has no identity scope.", + ) skill_data: Dict[str, Any] = { "skill_id": skill_id, diff --git a/litellm/proxy/common_utils/resource_ownership.py b/litellm/proxy/common_utils/resource_ownership.py index 9b55554fbc69..936c4e18bded 100644 --- a/litellm/proxy/common_utils/resource_ownership.py +++ b/litellm/proxy/common_utils/resource_ownership.py @@ -2,8 +2,6 @@ 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: @@ -22,8 +20,14 @@ def get_resource_owner_scopes( Return ownership scopes that may access a user-created proxy resource. Raw user_id is included for rows created before scope prefixes existed. - Prefixes avoid collisions when falling back to team/org/key ownership for - keys that do not have a user_id. + Prefixes avoid collisions when falling back to team/org/key ownership + for keys that do not have a user_id. + + Identity-less callers (no user_id, team_id, org_id, api_key, or token) + return ``[]`` — they share no scope with any other caller, so access + checks against an existing owner always fail and creates that depend + on a primary scope must reject up front. Returning a shared sentinel + here would let any two identity-less callers see each other's data. """ if user_api_key_dict is None: return [] @@ -45,8 +49,6 @@ def _add(scope: Optional[str]) -> None: _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 @@ -54,6 +56,13 @@ def _add(scope: Optional[str]) -> None: def get_primary_resource_owner_scope( user_api_key_dict: Optional[UserAPIKeyAuth], ) -> Optional[str]: + """Return the canonical owner scope to stamp on newly-created rows. + + ``None`` for identity-less callers — callers that depend on a primary + scope to record ownership must surface that as a hard error rather + than fall back to a shared sentinel (which would collapse every + identity-less caller into the same logical owner). + """ if user_api_key_dict is None: return None @@ -67,7 +76,7 @@ def get_primary_resource_owner_scope( return f"key:{user_api_key_dict.api_key}" if user_api_key_dict.token: return f"key:{user_api_key_dict.token}" - return UNSCOPED_RESOURCE_OWNER_SCOPE + return None def user_can_access_resource_owner( diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 6bfea87ebf64..b5b340fe2be9 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -147,7 +147,15 @@ async def record_container_owner( ) return response if owner is None: - raise HTTPException(status_code=500, detail="Unable to track container") + # Caller has no identity (no user_id / team_id / org_id / api_key / + # token) we can stamp on the row. Recording with a placeholder + # would collapse every such caller into a single shared owner — + # the cross-tenant data-access primitive we explicitly avoid. + # Reject with 403 rather than fall back to a sentinel. + raise HTTPException( + status_code=403, + detail="Unable to record container ownership: caller has no identity scope.", + ) original_container_id, resolved_provider = decode_container_id_for_ownership( container_id, diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 2e82b5b6b2d3..850c5a96e537 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -138,7 +138,12 @@ async def test_should_record_token_owner_for_keys_without_user_team_or_org(monke @pytest.mark.asyncio -async def test_should_record_unscoped_owner_for_identityless_proxy_auth(monkeypatch): +async def test_should_reject_record_for_identityless_proxy_auth(monkeypatch): + """Identity-less callers (no user_id / team_id / org_id / api_key / + token) cannot record ownership — stamping a shared sentinel would let + any two such callers see each other's containers.""" + from fastapi import HTTPException + monkeypatch.setattr( ownership, "_get_prisma_client", @@ -146,23 +151,15 @@ async def test_should_record_unscoped_owner_for_identityless_proxy_auth(monkeypa ) 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" + with pytest.raises(HTTPException) as exc: + await ownership.record_container_owner( + response=_container("cntr_provider"), + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + assert exc.value.status_code == 403 + assert "identity scope" in str(exc.value.detail) + assert ownership._IN_MEMORY_CONTAINER_OWNERS == {} @pytest.mark.asyncio 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 cef5f8a554b8..8b189be49219 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -150,9 +150,11 @@ def test_should_build_resource_owner_scopes_for_auth_context(): assert resource_ownership.get_resource_owner_scopes( UserAPIKeyAuth(token="token-hash") ) == ["key:token-hash"] - assert resource_ownership.get_resource_owner_scopes(UserAPIKeyAuth()) == [ - resource_ownership.UNSCOPED_RESOURCE_OWNER_SCOPE - ] + # Identity-less callers get an empty scope set — sharing a sentinel + # would collapse every identity-less caller into the same logical + # owner, which is a cross-tenant data-access primitive. + assert resource_ownership.get_resource_owner_scopes(UserAPIKeyAuth()) == [] + assert resource_ownership.get_primary_resource_owner_scope(UserAPIKeyAuth()) is None def test_should_allow_admin_and_anonymous_resource_owner_paths(): @@ -259,9 +261,13 @@ async def test_should_store_token_owner_for_keys_without_user_team_or_org(monkey @pytest.mark.asyncio -async def test_should_store_unscoped_owner_for_identityless_proxy_auth(monkeypatch): +async def test_should_reject_skill_create_for_identityless_proxy_auth(monkeypatch): + """Identity-less callers cannot create skills — stamping a shared + sentinel as ``created_by`` would let any two such callers see each + other's skills via the resulting shared owner scope.""" + from fastapi import HTTPException + 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})()} )() @@ -273,16 +279,14 @@ async def test_should_store_unscoped_owner_for_identityless_proxy_auth(monkeypat 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__" - ) + with pytest.raises(HTTPException) as exc: + await LiteLLMSkillsHandler.create_skill( + data=NewSkillRequest(display_title="skill"), + user_api_key_dict=auth, + ) + assert exc.value.status_code == 403 + assert "identity scope" in str(exc.value.detail) + table.create.assert_not_awaited() @pytest.mark.asyncio From 12fe945e7bbac2b9e7db6dbcefb73e9a4c5f0a13 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Mon, 4 May 2026 23:54:33 +0000 Subject: [PATCH 24/27] fix: keep skills handler FastAPI-free; fold gcs deny list into the body bouncer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two cleanups: * ``LiteLLMSkillsHandler.create_skill`` raised ``HTTPException`` for identity-less callers, importing FastAPI from a ``litellm/llms/`` module — that violates the project rule that FastAPI lives only under ``proxy/``. Switch to ``ValueError`` (the same shape the rest of the handler uses for not-found/forbidden) and update the test. * The proxy-auth body bouncer derived its observability ban list from ``_supported_callback_params`` only, missing ``_request_blocked_callback_params`` (where ``gcs_bucket_name`` and ``gcs_path_service_account`` live). Two recently-merged sibling PRs (#27019 added the deny list, #27081 added the test asserting these are rejected at the request body root) crossed without folding them together. Union the GCS deny list into the bouncer's derivation so the single source of truth covers both code paths. --- litellm/llms/litellm_proxy/skills/handler.py | 11 +++++------ litellm/proxy/auth/auth_utils.py | 13 +++++++++++-- .../llms/litellm_proxy/test_skills_ownership.py | 6 +----- 3 files changed, 17 insertions(+), 13 deletions(-) diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 1129af64e14e..48c02660a061 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -137,12 +137,11 @@ async def create_skill( # Caller has no identity scope (no user_id / team_id / org_id / # api_key / token). Stamping a placeholder would let any two # identity-less callers see each other's skills via the shared - # owner — the cross-tenant primitive we avoid. - from fastapi import HTTPException - - raise HTTPException( - status_code=403, - detail="Unable to record skill ownership: caller has no identity scope.", + # owner — the cross-tenant primitive we avoid. ValueError keeps + # this module FastAPI-free per the project layering rule + # (litellm_proxy provider integrations live outside proxy/). + raise ValueError( + "Unable to record skill ownership: caller has no identity scope." ) skill_data: Dict[str, Any] = { diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 9a6fc95f1451..917d3d992283 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -224,14 +224,23 @@ def _build_banned_observability_params() -> FrozenSet[str]: the extras the canonical allowlist hasn't caught up to yet. New integrations added to the canonical allowlist are banned by default, which is the safe failure mode. + + ``_request_blocked_callback_params`` (e.g. ``gcs_bucket_name``, + ``gcs_path_service_account``) is the GCS-logging-specific deny list + that lives alongside the allowlist; fold it in here so a single + declaration of "this field must not be caller-supplied" covers both + the request-body bouncer and the dynamic callback initializer. """ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( + _request_blocked_callback_params, _supported_callback_params, ) return ( - frozenset(_supported_callback_params) - _SAFE_CLIENT_CALLBACK_PARAMS - ) | _EXTRA_BANNED_OBSERVABILITY_PARAMS + (frozenset(_supported_callback_params) - _SAFE_CLIENT_CALLBACK_PARAMS) + | _EXTRA_BANNED_OBSERVABILITY_PARAMS + | frozenset(_request_blocked_callback_params) + ) _BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = ( 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 8b189be49219..c6dd6c4e33ca 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -265,8 +265,6 @@ async def test_should_reject_skill_create_for_identityless_proxy_auth(monkeypatc """Identity-less callers cannot create skills — stamping a shared sentinel as ``created_by`` would let any two such callers see each other's skills via the resulting shared owner scope.""" - from fastapi import HTTPException - table = AsyncMock() prisma_client = type( "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} @@ -279,13 +277,11 @@ async def test_should_reject_skill_create_for_identityless_proxy_auth(monkeypatc auth = UserAPIKeyAuth() - with pytest.raises(HTTPException) as exc: + with pytest.raises(ValueError, match="identity scope"): await LiteLLMSkillsHandler.create_skill( data=NewSkillRequest(display_title="skill"), user_api_key_dict=auth, ) - assert exc.value.status_code == 403 - assert "identity scope" in str(exc.value.detail) table.create.assert_not_awaited() From 6ce84effe1a49bb58d9f3087e95c26a141ae6102 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Tue, 5 May 2026 00:23:32 +0000 Subject: [PATCH 25/27] =?UTF-8?q?chore:=20simplify=20ownership=20tracking?= =?UTF-8?q?=20=E2=80=94=20drop=20thin=20stores,=20in-memory=20fallback,=20?= =?UTF-8?q?hand-rolled=20cache?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Substantial reduction (~765 LOC) without changing the security boundary: * Drop ContainerOwnershipStore and LiteLLMSkillsStore — both were one-method-per-Prisma-call wrappers. Inline the calls instead, matching the established pattern in vector_store_endpoints, agent_endpoints, and mcp_server/db.py. * Drop the prisma_client is None in-memory fallback. Production deploys always have Prisma; running ownership-critical paths on a process-local dict is a security footgun in the dev-mode case it was meant to support, and complicates every code path with a branch. Fail-secure: skip recording if Prisma is unavailable, and treat reads as "not found" (admin-only). * Drop the hand-rolled module-level cache. Replace with the existing litellm.caching.in_memory_cache.InMemoryCache, which already has TTL + max-size + eviction tested in its own module. Sentinel string for negative caching since InMemoryCache can't disambiguate "miss" from "cached as None". * Tests: drop coverage for removed code paths (in-memory fallback, hand-rolled cache internals). Keep tests for actual behavior (cache hit-rate, negative caching, owner check, list filtering, identity-less reject, admin bypass). --- litellm/llms/litellm_proxy/skills/handler.py | 209 ++------ litellm/llms/litellm_proxy/skills/store.py | 22 - .../proxy/container_endpoints/ownership.py | 286 ++++------- .../container_endpoints/ownership_store.py | 55 --- .../test_container_proxy_ownership.py | 457 +----------------- .../litellm_proxy/test_skills_ownership.py | 102 ++-- 6 files changed, 183 insertions(+), 948 deletions(-) delete mode 100644 litellm/llms/litellm_proxy/skills/store.py delete mode 100644 litellm/proxy/container_endpoints/ownership_store.py diff --git a/litellm/llms/litellm_proxy/skills/handler.py b/litellm/llms/litellm_proxy/skills/handler.py index 48c02660a061..37aabd8b477d 100644 --- a/litellm/llms/litellm_proxy/skills/handler.py +++ b/litellm/llms/litellm_proxy/skills/handler.py @@ -5,13 +5,11 @@ Used by the transformation layer and skills injection hook. """ -import time import uuid -from collections import OrderedDict -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Optional from litellm._logging import verbose_logger -from litellm.llms.litellm_proxy.skills.store import LiteLLMSkillsStore +from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy._types import LiteLLM_SkillsTable, NewSkillRequest, UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, @@ -21,88 +19,37 @@ ) # Skills are looked up on every chat completion that has skills enabled -# (`LiteLLMSkillsHandler.fetch_skill_from_db` in the injection hook). Cache -# the Prisma skill row for a short window so the hot path doesn't issue a DB -# round-trip per request. Same shape as `_byok_cred_cache` and the container -# ownership cache: (value, monotonic_timestamp). `None` is cached as a true -# negative ("skill does not exist") so repeated misses also avoid the DB. -_SKILL_CACHE: "OrderedDict[str, Tuple[Optional[Any], float]]" = OrderedDict() -_SKILL_CACHE_TTL = 60 # seconds -_SKILL_CACHE_MAX_SIZE = 10000 - - -def _read_skill_cache(skill_id: str) -> Tuple[bool, Optional[Any]]: - """Return (hit, value). hit=False means caller must consult the DB.""" - entry = _SKILL_CACHE.get(skill_id) - if entry is None: - return False, None - value, timestamp = entry - if time.monotonic() - timestamp > _SKILL_CACHE_TTL: - _SKILL_CACHE.pop(skill_id, None) - return False, None - _SKILL_CACHE.move_to_end(skill_id) - return True, value - - -def _write_skill_cache(skill_id: str, skill: Optional[Any]) -> None: - # LRU eviction (popitem(last=False)) instead of full ``clear()`` — - # see container ownership cache for rationale. - if skill_id in _SKILL_CACHE: - _SKILL_CACHE.move_to_end(skill_id) - _SKILL_CACHE[skill_id] = (skill, time.monotonic()) - while len(_SKILL_CACHE) > _SKILL_CACHE_MAX_SIZE: - _SKILL_CACHE.popitem(last=False) - - -def _invalidate_skill_cache(skill_id: str) -> None: - """Drop the cache entry after a write so the next read sees the new row.""" - _SKILL_CACHE.pop(skill_id, None) - - -def _user_can_access_skill_owner( - owner: Optional[str], - user_api_key_dict: Optional[UserAPIKeyAuth], -) -> bool: - # Pre-isolation skills with no ``created_by`` are admin-only — same - # rule as untracked containers. Owners need to either re-create via - # the now-tracked flow or have an admin assign ``created_by`` on the - # row. - return user_can_access_resource_owner(owner, user_api_key_dict) +# (`SkillsInjectionHook` calls ``fetch_skill_from_db``). 60s LRU/TTL cache +# absorbs the hot read before it reaches Prisma. ``_NEGATIVE_SKILL_SENTINEL`` +# lets us cache a true "skill does not exist" so repeated misses also +# avoid the DB — ``InMemoryCache`` returns ``None`` indistinguishably for +# "miss" and "cached as None". +_NEGATIVE_SKILL_SENTINEL = "__litellm_skill_not_found__" +_SKILL_CACHE = InMemoryCache(max_size_in_memory=10000, default_ttl=60) def _prisma_skill_to_litellm(prisma_skill) -> LiteLLM_SkillsTable: - """ - Convert a Prisma skill record to LiteLLM_SkillsTable. + """Convert a Prisma skill record to LiteLLM_SkillsTable. - Handles Base64 decoding of file_content field. + Handles Base64 decoding of file_content field — model_dump() converts + Base64 fields to base64-encoded strings. """ import base64 data = prisma_skill.model_dump() - # Decode Base64 file_content back to bytes - # model_dump() converts Base64 field to base64-encoded string if data.get("file_content") is not None: if isinstance(data["file_content"], str): data["file_content"] = base64.b64decode(data["file_content"]) - elif isinstance(data["file_content"], bytes): - # Already bytes, no conversion needed - pass return LiteLLM_SkillsTable(**data) class LiteLLMSkillsHandler: - """ - Handler for LiteLLM database-backed skills operations. - - This class provides static methods for CRUD operations on skills - stored in the LiteLLM proxy database (LiteLLM_SkillsTable). - """ + """CRUD for skills stored in ``litellm_skillstable``.""" @staticmethod async def _get_prisma_client(): - """Get the prisma client from proxy server.""" from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -118,28 +65,16 @@ async def create_skill( user_id: Optional[str] = None, user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> LiteLLM_SkillsTable: - """ - Create a new skill in the LiteLLM database. - - Args: - data: NewSkillRequest with skill details - user_id: Optional user ID for tracking - - Returns: - LiteLLM_SkillsTable record - """ prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - store = LiteLLMSkillsStore(prisma_client) skill_id = f"litellm_skill_{uuid.uuid4()}" owner = get_primary_resource_owner_scope(user_api_key_dict) or user_id if owner is None: - # Caller has no identity scope (no user_id / team_id / org_id / - # api_key / token). Stamping a placeholder would let any two - # identity-less callers see each other's skills via the shared - # owner — the cross-tenant primitive we avoid. ValueError keeps - # this module FastAPI-free per the project layering rule - # (litellm_proxy provider integrations live outside proxy/). + # Identity-less callers (no user_id / team_id / org_id / + # api_key / token) can't be uniquely stamped on the row. + # Stamping a placeholder would let any two such callers see + # each other's skills via the shared owner. ValueError keeps + # this module FastAPI-free per the project layering rule. raise ValueError( "Unable to record skill ownership: caller has no identity scope." ) @@ -154,13 +89,11 @@ async def create_skill( "updated_by": owner, } - # Handle metadata if data.metadata is not None: from litellm.litellm_core_utils.safe_json_dumps import safe_dumps skill_data["metadata"] = safe_dumps(data.metadata) - # Handle file content - wrap bytes in Base64 for Prisma if data.file_content is not None: from prisma.fields import Base64 @@ -174,8 +107,7 @@ async def create_skill( f"LiteLLMSkillsHandler: Creating skill {skill_id} with title={data.display_title}" ) - new_skill = await store.create_skill(skill_data) - + new_skill = await prisma_client.db.litellm_skillstable.create(data=skill_data) return _prisma_skill_to_litellm(new_skill) @staticmethod @@ -184,18 +116,7 @@ async def list_skills( offset: int = 0, user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> List[LiteLLM_SkillsTable]: - """ - List skills from the LiteLLM database. - - Args: - limit: Maximum number of skills to return - offset: Number of skills to skip - - Returns: - List of LiteLLM_SkillsTable records - """ prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - store = LiteLLMSkillsStore(prisma_client) verbose_logger.debug( f"LiteLLMSkillsHandler: Listing skills with limit={limit}, offset={offset}" @@ -212,26 +133,29 @@ async def list_skills( return [] find_many_kwargs["where"] = {"created_by": {"in": owner_scopes}} - skills = await store.list_skills(find_many_kwargs) - + skills = await prisma_client.db.litellm_skillstable.find_many( + **find_many_kwargs + ) return [_prisma_skill_to_litellm(s) for s in skills] @staticmethod async def _load_skill(skill_id: str) -> Optional[Any]: - """Cache-first read of the Prisma skill row. - - Caching here keeps `fetch_skill_from_db` (called per chat completion in - the skills injection hook) off the DB. Owner-scope filtering happens - on the cached row, so the cache is per-skill and not per-caller. + """Cache-first read of the Prisma skill row. Owner-scope filtering + happens on the cached row, so the cache is per-skill not per-caller. """ - cached_hit, cached_skill = _read_skill_cache(skill_id) - if cached_hit: - return cached_skill + cached = _SKILL_CACHE.get_cache(skill_id) + if cached == _NEGATIVE_SKILL_SENTINEL: + return None + if cached is not None: + return cached prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - store = LiteLLMSkillsStore(prisma_client) - skill = await store.find_skill(skill_id) - _write_skill_cache(skill_id, skill) + skill = await prisma_client.db.litellm_skillstable.find_unique( + where={"skill_id": skill_id} + ) + _SKILL_CACHE.set_cache( + skill_id, skill if skill is not None else _NEGATIVE_SKILL_SENTINEL + ) return skill @staticmethod @@ -239,26 +163,12 @@ async def get_skill( skill_id: str, user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> LiteLLM_SkillsTable: - """ - Get a skill by ID from the LiteLLM database. - - Args: - skill_id: The skill ID to retrieve - - Returns: - LiteLLM_SkillsTable record - - Raises: - ValueError: If skill not found - """ verbose_logger.debug(f"LiteLLMSkillsHandler: Getting skill {skill_id}") skill = await LiteLLMSkillsHandler._load_skill(skill_id) - - if skill is None: - raise ValueError(f"Skill not found: {skill_id}") - - if not _user_can_access_skill_owner( + # Same "not found" message for both "missing" and "cross-tenant" + # so callers can't enumerate skill IDs they don't own. + if skill is None or not user_can_access_resource_owner( getattr(skill, "created_by", None), user_api_key_dict ): raise ValueError(f"Skill not found: {skill_id}") @@ -270,37 +180,17 @@ async def delete_skill( skill_id: str, user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> Dict[str, str]: - """ - Delete a skill by ID from the LiteLLM database. - - Args: - skill_id: The skill ID to delete - - Returns: - Dict with id and type of deleted skill - - Raises: - ValueError: If skill not found - """ prisma_client = await LiteLLMSkillsHandler._get_prisma_client() - store = LiteLLMSkillsStore(prisma_client) - verbose_logger.debug(f"LiteLLMSkillsHandler: Deleting skill {skill_id}") - # Check if skill exists skill = await LiteLLMSkillsHandler._load_skill(skill_id) - - if skill is None: - raise ValueError(f"Skill not found: {skill_id}") - - if not _user_can_access_skill_owner( + if skill is None or not user_can_access_resource_owner( getattr(skill, "created_by", None), user_api_key_dict ): raise ValueError(f"Skill not found: {skill_id}") - # Delete the skill - await store.delete_skill(skill_id) - _invalidate_skill_cache(skill_id) + await prisma_client.db.litellm_skillstable.delete(where={"skill_id": skill_id}) + _SKILL_CACHE.set_cache(skill_id, _NEGATIVE_SKILL_SENTINEL) return {"id": skill_id, "type": "skill_deleted"} @@ -309,22 +199,11 @@ async def fetch_skill_from_db( skill_id: str, user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> Optional[LiteLLM_SkillsTable]: - """ - Fetch a skill from the database (used by skills injection hook). - - This is a convenience method that returns None instead of raising - an exception if the skill is not found. - - Args: - skill_id: The skill ID to fetch - - Returns: - LiteLLM_SkillsTable or None if not found - """ + """Skills-injection-hook helper: returns None instead of raising on + not-found / not-authorized so the hook can silently skip.""" try: return await LiteLLMSkillsHandler.get_skill( - skill_id, - user_api_key_dict=user_api_key_dict, + skill_id, user_api_key_dict=user_api_key_dict ) except ValueError: return None diff --git a/litellm/llms/litellm_proxy/skills/store.py b/litellm/llms/litellm_proxy/skills/store.py deleted file mode 100644 index 6e69a7a96251..000000000000 --- a/litellm/llms/litellm_proxy/skills/store.py +++ /dev/null @@ -1,22 +0,0 @@ -from typing import Any, Dict, List, Optional - - -class LiteLLMSkillsStore: - def __init__(self, prisma_client: Any): - self.prisma_client = prisma_client - - @property - def _table(self) -> Any: - return self.prisma_client.db.litellm_skillstable - - async def create_skill(self, data: Dict[str, Any]) -> Any: - return await self._table.create(data=data) - - async def list_skills(self, find_many_kwargs: Dict[str, Any]) -> List[Any]: - return await self._table.find_many(**find_many_kwargs) - - async def find_skill(self, skill_id: str) -> Optional[Any]: - return await self._table.find_unique(where={"skill_id": skill_id}) - - async def delete_skill(self, skill_id: str) -> None: - await self._table.delete(where={"skill_id": skill_id}) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index b5b340fe2be9..137366c955a5 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,10 +1,9 @@ -import time -from collections import OrderedDict from typing import Any, Dict, List, Optional, Set, Tuple from fastapi import HTTPException from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.resource_ownership import ( get_primary_resource_owner_scope, @@ -12,75 +11,26 @@ is_proxy_admin, user_can_access_resource_owner, ) -from litellm.proxy.container_endpoints.ownership_store import ( - CONTAINER_OBJECT_PURPOSE, - ContainerOwnershipStore, -) from litellm.responses.utils import ResponsesAPIRequestUtils -MAX_IN_MEMORY_CONTAINER_OWNERS = 10000 -_IN_MEMORY_CONTAINER_OWNERS: "OrderedDict[str, str]" = OrderedDict() - -# Short-lived cache keeps every container access check from hitting the DB -# (`_get_container_owner` is invoked on retrieve / delete / list / file-content -# paths). Mirrors the `_byok_cred_cache` pattern in mcp_server/server.py: -# (value, monotonic_timestamp) tuples, TTL'd, LRU-evicted at capacity, -# invalidated by writes. A ``None`` value caches "untracked" so repeated -# negative lookups also avoid DB. -_CONTAINER_OWNER_CACHE: "OrderedDict[str, Tuple[Optional[str], float]]" = OrderedDict() -_CONTAINER_OWNER_CACHE_TTL = 60 # seconds -_CONTAINER_OWNER_CACHE_MAX_SIZE = 10000 - - -def _read_container_owner_cache(model_object_id: str) -> Tuple[bool, Optional[str]]: - """Return (hit, value). hit=False means caller must consult the DB.""" - entry = _CONTAINER_OWNER_CACHE.get(model_object_id) - if entry is None: - return False, None - value, timestamp = entry - if time.monotonic() - timestamp > _CONTAINER_OWNER_CACHE_TTL: - _CONTAINER_OWNER_CACHE.pop(model_object_id, None) - return False, None - _CONTAINER_OWNER_CACHE.move_to_end(model_object_id) - return True, value - - -def _write_container_owner_cache(model_object_id: str, owner: Optional[str]) -> None: - # LRU eviction (popitem(last=False)) instead of full ``clear()`` — a - # full clear at capacity converts a steady-state cached workload into - # a periodic full-DB-load oscillation as the cache repopulates from - # zero and clears again. - if model_object_id in _CONTAINER_OWNER_CACHE: - _CONTAINER_OWNER_CACHE.move_to_end(model_object_id) - _CONTAINER_OWNER_CACHE[model_object_id] = (owner, time.monotonic()) - while len(_CONTAINER_OWNER_CACHE) > _CONTAINER_OWNER_CACHE_MAX_SIZE: - _CONTAINER_OWNER_CACHE.popitem(last=False) - - -def _invalidate_container_owner_cache(model_object_id: str) -> None: - """Drop a cache entry after a write so the next read sees the new owner.""" - _CONTAINER_OWNER_CACHE.pop(model_object_id, None) - - -def _remember_container_owner(model_object_id: str, owner: str) -> None: - existing_owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) - if existing_owner is not None: - _IN_MEMORY_CONTAINER_OWNERS.move_to_end(model_object_id) - _IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner - while len(_IN_MEMORY_CONTAINER_OWNERS) > MAX_IN_MEMORY_CONTAINER_OWNERS: - _IN_MEMORY_CONTAINER_OWNERS.popitem(last=False) +CONTAINER_OBJECT_PURPOSE = "container" + +# 60s LRU/TTL cache absorbs every container access check before it reaches +# Prisma. ``_NEGATIVE_OWNER_SENTINEL`` lets us cache a true "untracked" +# answer so repeated misses also avoid the DB — ``InMemoryCache`` returns +# ``None`` indistinguishably for "miss" and "cached as None". +_NEGATIVE_OWNER_SENTINEL = "__litellm_container_no_owner__" +_CONTAINER_OWNER_CACHE = InMemoryCache(max_size_in_memory=10000, default_ttl=60) def _container_model_object_id( - original_container_id: str, - custom_llm_provider: str, + original_container_id: str, custom_llm_provider: str ) -> str: return f"{CONTAINER_OBJECT_PURPOSE}:{custom_llm_provider}:{original_container_id}" def decode_container_id_for_ownership( - container_id: str, - custom_llm_provider: str, + container_id: str, custom_llm_provider: str ) -> Tuple[str, str]: decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) original_container_id = decoded.get("response_id", container_id) @@ -91,9 +41,7 @@ def decode_container_id_for_ownership( def get_container_forwarding_params( - container_id: str, - original_container_id: str, - custom_llm_provider: str, + container_id: str, original_container_id: str, custom_llm_provider: str ) -> Dict[str, str]: params = { "container_id": original_container_id, @@ -147,122 +95,93 @@ async def record_container_owner( ) return response if owner is None: - # Caller has no identity (no user_id / team_id / org_id / api_key / - # token) we can stamp on the row. Recording with a placeholder - # would collapse every such caller into a single shared owner — - # the cross-tenant data-access primitive we explicitly avoid. - # Reject with 403 rather than fall back to a sentinel. + # Identity-less callers (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. raise HTTPException( status_code=403, detail="Unable to record container ownership: caller has no identity scope.", ) original_container_id, resolved_provider = decode_container_id_for_ownership( - container_id, - custom_llm_provider, + container_id, custom_llm_provider ) model_object_id = _container_model_object_id( - original_container_id, - resolved_provider, + original_container_id, resolved_provider ) file_object = _dump_response(response) file_object["custom_llm_provider"] = resolved_provider file_object["provider_container_id"] = original_container_id - 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") - _remember_container_owner(model_object_id, owner) - _invalidate_container_owner_cache(model_object_id) - return response - - store = ContainerOwnershipStore(prisma_client) - existing = await store.find_by_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 - ): - raise HTTPException(status_code=403, detail="Forbidden") - await store.update_owner_record( - model_object_id=model_object_id, - data={ - "unified_object_id": container_id, - "file_object": file_object, - "updated_by": owner, - }, - ) - else: - await store.create_owner_record( - data={ - "unified_object_id": container_id, - "model_object_id": model_object_id, - "file_object": file_object, - "file_purpose": CONTAINER_OBJECT_PURPOSE, - "created_by": owner, - "updated_by": owner, - } - ) - except HTTPException: - raise - except Exception as e: + prisma_client = await _get_prisma_client() + if prisma_client is None: verbose_proxy_logger.warning( - "Failed to persist container ownership for container_id=%s; " - "falling back to in-process tracking: %s", - model_object_id, - e, + "Skipping container ownership tracking because 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 + return response + + table = prisma_client.db.litellm_managedobjecttable + 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 ): raise HTTPException(status_code=403, detail="Forbidden") - _remember_container_owner(model_object_id, owner) + await table.update( + where={"model_object_id": model_object_id}, + data={ + "unified_object_id": container_id, + "file_object": file_object, + "updated_by": owner, + }, + ) + else: + await table.create( + data={ + "unified_object_id": container_id, + "model_object_id": model_object_id, + "file_object": file_object, + "file_purpose": CONTAINER_OBJECT_PURPOSE, + "created_by": owner, + "updated_by": owner, + } + ) - _invalidate_container_owner_cache(model_object_id) + _CONTAINER_OWNER_CACHE.set_cache(model_object_id, owner) return response async def _get_container_owner( - original_container_id: str, - custom_llm_provider: str, + original_container_id: str, custom_llm_provider: str ) -> Optional[str]: model_object_id = _container_model_object_id( - original_container_id, - custom_llm_provider, + original_container_id, custom_llm_provider ) - cached_hit, cached_value = _read_container_owner_cache(model_object_id) - if cached_hit: - return cached_value - - try: - prisma_client = await _get_prisma_client() - if prisma_client is None: - owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) - _write_container_owner_cache(model_object_id, owner) - return owner - - owner = await ContainerOwnershipStore(prisma_client).get_owner(model_object_id) - if owner is None: - owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) - _write_container_owner_cache(model_object_id, owner) - return owner - 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, - ) - # Don't cache transient DB errors — let the next request retry. - return _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id) + cached = _CONTAINER_OWNER_CACHE.get_cache(model_object_id) + if cached == _NEGATIVE_OWNER_SENTINEL: + return None + if cached is not None: + return cached + + prisma_client = await _get_prisma_client() + if prisma_client is None: + return None + + row = await prisma_client.db.litellm_managedobjecttable.find_first( + where={ + "model_object_id": model_object_id, + "file_purpose": CONTAINER_OBJECT_PURPOSE, + } + ) + owner = getattr(row, "created_by", None) if row is not None else None + _CONTAINER_OWNER_CACHE.set_cache( + model_object_id, owner if owner is not None else _NEGATIVE_OWNER_SENTINEL + ) + return owner async def assert_user_can_access_container( @@ -271,17 +190,15 @@ async def assert_user_can_access_container( custom_llm_provider: str, ) -> Tuple[str, str]: original_container_id, resolved_provider = decode_container_id_for_ownership( - container_id, - custom_llm_provider, + container_id, custom_llm_provider ) if is_proxy_admin(user_api_key_dict): return original_container_id, resolved_provider - # Untracked containers (no ownership row) are admin-only. Pre-isolation - # rows that pre-date this enforcement need the admin to either re-create - # via the now-tracked flow or explicitly assign ``created_by`` on the - # ``litellm_managedobjecttable`` row. + # Untracked rows (no ownership) are admin-only. Pre-isolation rows + # 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): raise HTTPException(status_code=403, detail="Forbidden") @@ -327,35 +244,26 @@ def _set_container_list_data( async def _get_allowed_container_ids( user_api_key_dict: UserAPIKeyAuth, - custom_llm_provider: str, ) -> Set[str]: owner_scopes = get_resource_owner_scopes(user_api_key_dict) if not owner_scopes: return set() - in_memory_allowed_ids = { - model_object_id - for model_object_id, owner in _IN_MEMORY_CONTAINER_OWNERS.items() - if owner in owner_scopes + prisma_client = await _get_prisma_client() + if prisma_client is None: + 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 } - try: - prisma_client = await _get_prisma_client() - if prisma_client is None: - return in_memory_allowed_ids - - db_allowed_ids = await ContainerOwnershipStore( - prisma_client - ).list_model_object_ids_for_owners( - owner_scopes=owner_scopes, - ) - return in_memory_allowed_ids | db_allowed_ids - 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( @@ -370,18 +278,14 @@ async def filter_container_list_response( if data is None: return response - allowed_container_ids = await _get_allowed_container_ids( - user_api_key_dict, - custom_llm_provider, - ) + allowed_container_ids = await _get_allowed_container_ids(user_api_key_dict) filtered: List[Any] = [] for item in data: container_id = _get_response_id(item) if container_id is None: continue original_container_id, resolved_provider = decode_container_id_for_ownership( - container_id, - custom_llm_provider, + container_id, custom_llm_provider ) if ( _container_model_object_id(original_container_id, resolved_provider) @@ -390,7 +294,5 @@ async def filter_container_list_response( filtered.append(item) return _set_container_list_data( - response, - filtered, - removed_filtered_items=len(filtered) != len(data), + response, filtered, removed_filtered_items=len(filtered) != len(data) ) diff --git a/litellm/proxy/container_endpoints/ownership_store.py b/litellm/proxy/container_endpoints/ownership_store.py deleted file mode 100644 index b3c406f618d0..000000000000 --- a/litellm/proxy/container_endpoints/ownership_store.py +++ /dev/null @@ -1,55 +0,0 @@ -from typing import Any, Dict, List, Optional, Set - -CONTAINER_OBJECT_PURPOSE = "container" - - -class ContainerOwnershipStore: - def __init__(self, prisma_client: Any): - self.prisma_client = prisma_client - - @property - def _table(self) -> Any: - return self.prisma_client.db.litellm_managedobjecttable - - async def find_by_model_object_id(self, model_object_id: str) -> Optional[Any]: - return await self._table.find_unique(where={"model_object_id": model_object_id}) - - async def create_owner_record(self, data: Dict[str, Any]) -> None: - await self._table.create(data=data) - - async def update_owner_record( - self, - model_object_id: str, - data: Dict[str, Any], - ) -> None: - await self._table.update( - where={"model_object_id": model_object_id}, - data=data, - ) - - async def get_owner(self, model_object_id: str) -> Optional[str]: - row = await self._table.find_first( - where={ - "model_object_id": model_object_id, - "file_purpose": CONTAINER_OBJECT_PURPOSE, - } - ) - if row is None: - return None - return getattr(row, "created_by", None) - - async def list_model_object_ids_for_owners( - self, - owner_scopes: List[str], - ) -> Set[str]: - rows = await self._table.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 - } diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 850c5a96e537..b046fa1536ba 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -12,12 +12,12 @@ @pytest.fixture(autouse=True) -def clear_in_memory_container_owners(): - ownership._IN_MEMORY_CONTAINER_OWNERS.clear() - ownership._CONTAINER_OWNER_CACHE.clear() +def clear_container_owner_cache(): + ownership._CONTAINER_OWNER_CACHE.cache_dict.clear() + ownership._CONTAINER_OWNER_CACHE.ttl_dict.clear() yield - ownership._IN_MEMORY_CONTAINER_OWNERS.clear() - ownership._CONTAINER_OWNER_CACHE.clear() + ownership._CONTAINER_OWNER_CACHE.cache_dict.clear() + ownership._CONTAINER_OWNER_CACHE.ttl_dict.clear() def _container(container_id: str) -> ContainerObject: @@ -159,7 +159,6 @@ async def test_should_reject_record_for_identityless_proxy_auth(monkeypatch): ) assert exc.value.status_code == 403 assert "identity scope" in str(exc.value.detail) - assert ownership._IN_MEMORY_CONTAINER_OWNERS == {} @pytest.mark.asyncio @@ -186,106 +185,6 @@ async def test_should_skip_owner_record_when_provider_response_has_no_id(monkeyp table.create.assert_not_awaited() -@pytest.mark.asyncio -async def test_should_fallback_to_memory_when_persistent_owner_record_fails( - monkeypatch, -): - table = AsyncMock() - table.find_unique.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") - - 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"] - == "user-1" - ) - - -@pytest.mark.asyncio -async def test_should_track_container_owner_in_memory_without_prisma(monkeypatch): - monkeypatch.setattr( - ownership, - "_get_prisma_client", - AsyncMock(return_value=None), - ) - auth = UserAPIKeyAuth(user_id="user-1") - - await ownership.record_container_owner( - response=_container("cntr_provider"), - user_api_key_dict=auth, - custom_llm_provider="openai", - ) - - 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_bound_in_memory_container_owner_tracking(monkeypatch): - monkeypatch.setattr(ownership, "MAX_IN_MEMORY_CONTAINER_OWNERS", 2) - monkeypatch.setattr( - ownership, - "_get_prisma_client", - AsyncMock(return_value=None), - ) - auth = UserAPIKeyAuth(user_id="user-1") - - for container_id in ("cntr_1", "cntr_2", "cntr_3"): - await ownership.record_container_owner( - response=_container(container_id), - user_api_key_dict=auth, - custom_llm_provider="openai", - ) - - assert list(ownership._IN_MEMORY_CONTAINER_OWNERS.keys()) == [ - "container:openai:cntr_2", - "container:openai:cntr_3", - ] - - -@pytest.mark.asyncio -async def test_should_deny_container_access_for_different_owner(monkeypatch): - table = AsyncMock() - table.find_first.return_value = SimpleNamespace(created_by="user-2") - 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_provider", - user_api_key_dict=auth, - custom_llm_provider="openai", - ) - - assert exc.value.status_code == 403 - - @pytest.mark.asyncio async def test_should_deny_untracked_container_access_by_default(monkeypatch): monkeypatch.setattr( @@ -305,103 +204,6 @@ 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_use_memory_owner_when_db_recovers_without_row(monkeypatch): - table = AsyncMock() - table.find_first.return_value = None - 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_untracked_container_is_admin_only(monkeypatch): - """Pre-isolation containers with no ownership row are admin-only. - Non-admin callers see them as 403, with no opt-out flag re-opening - the cross-tenant access primitive.""" - from fastapi import HTTPException - - monkeypatch.setattr( - ownership, - "_get_prisma_client", - AsyncMock(return_value=None), - ) - auth = UserAPIKeyAuth(user_id="user-1") - - with pytest.raises(HTTPException) as exc: - await ownership.assert_user_can_access_container( - container_id="cntr_untracked", - user_api_key_dict=auth, - custom_llm_provider="openai", - ) - assert exc.value.status_code == 403 - - @pytest.mark.asyncio async def test_should_not_reassign_existing_container_to_different_owner(monkeypatch): table = AsyncMock() @@ -536,165 +338,6 @@ async def test_should_clear_dict_has_more_when_filtered_container_list_is_empty( assert filtered["has_more"] is False -@pytest.mark.asyncio -async def test_should_filter_container_list_with_in_memory_ownership(monkeypatch): - monkeypatch.setattr( - ownership, - "_get_prisma_client", - AsyncMock(return_value=None), - ) - auth = UserAPIKeyAuth(user_id="user-1") - - await ownership.record_container_owner( - response=_container("cntr_owned"), - user_api_key_dict=auth, - custom_llm_provider="openai", - ) - - response = ContainerListResponse( - object="list", - data=[_container("cntr_owned"), _container("cntr_other")], - has_more=False, - ) - - 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"] - - -@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_include_memory_container_list_when_db_recovers_without_row( - monkeypatch, -): - table = AsyncMock() - table.find_many.return_value = [] - 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_validate_owner_and_forward_decoded_id_for_proxy_forwarding( - monkeypatch, -): - from litellm.proxy.container_endpoints import handler_factory - - proxy_server_stub = SimpleNamespace( - general_settings={}, - llm_router=None, - proxy_config=None, - proxy_logging_obj=None, - select_data_generator=None, - user_api_base=None, - user_max_tokens=None, - user_model=None, - user_request_timeout=None, - user_temperature=None, - version="test", - ) - monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub) - - captured = {} - - class FakeProcessor: - def __init__(self, data): - captured["data"] = data - - async def base_process_llm_request(self, **kwargs): - return captured["data"] - - async def _handle_llm_api_exception(self, **kwargs): - raise kwargs["e"] - - monkeypatch.setattr( - handler_factory, - "ProxyBaseLLMRequestProcessing", - FakeProcessor, - ) - access_check = AsyncMock(return_value=("cntr_provider", "azure")) - monkeypatch.setattr( - handler_factory, - "assert_user_can_access_container", - access_check, - ) - encoded_id = ResponsesAPIRequestUtils._build_container_id( - custom_llm_provider="azure", - model_id="router-gpt", - container_id="cntr_provider", - ) - - result = await handler_factory._process_request( - request=SimpleNamespace(query_params={}, headers={}), - fastapi_response=SimpleNamespace(), - user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), - route_type="alist_container_files", - path_params={"container_id": encoded_id}, - ) - - access_check.assert_awaited_once() - assert access_check.await_args.kwargs["container_id"] == encoded_id - assert result["container_id"] == "cntr_provider" - assert result["custom_llm_provider"] == "azure" - assert result["model_id"] == "router-gpt" - - @pytest.mark.asyncio async def test_should_validate_owner_and_forward_decoded_id_for_multipart_upload( monkeypatch, @@ -1158,93 +801,3 @@ 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_record_container_owner_invalidates_cache(monkeypatch): - """A recorded owner must drop the cached value so the next read re-fetches. - - Otherwise a stale `None` from a prior negative lookup would survive the - create and the new owner would be invisible until the TTL elapses. - """ - # Seed the cache with a stale negative result. - ownership._write_container_owner_cache("container:openai:cntr_new", None) - cached_hit, cached_value = ownership._read_container_owner_cache( - "container:openai:cntr_new" - ) - assert cached_hit and cached_value is None - - 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), - ) - - await ownership.record_container_owner( - response=_container("cntr_new"), - user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), - custom_llm_provider="openai", - ) - - # Invalidation drops the entry — next read goes to the DB. - cached_hit, _ = ownership._read_container_owner_cache("container:openai:cntr_new") - assert not cached_hit - - -@pytest.mark.asyncio -async def test_get_container_owner_does_not_cache_on_db_error(monkeypatch): - """DB errors must skip caching so transient failures don't pin a `None`.""" - 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), - ) - - result = await ownership._get_container_owner("cntr_x", "openai") - assert result is None - cached_hit, _ = ownership._read_container_owner_cache("container:openai:cntr_x") - assert not cached_hit - - -def test_container_owner_cache_expires_after_ttl(monkeypatch): - """Entries past the TTL count as misses so writes elsewhere are eventually - visible to this process.""" - monkeypatch.setattr(ownership, "_CONTAINER_OWNER_CACHE_TTL", 0.0) - ownership._write_container_owner_cache("k", "user-1") - cached_hit, _ = ownership._read_container_owner_cache("k") - # TTL of 0 means anything in the cache is already stale. - assert not cached_hit - - -def test_container_owner_cache_evicts_when_at_capacity(monkeypatch): - """The cache must not grow unbounded; reaching capacity LRU-evicts the - oldest entry, not the entire cache.""" - monkeypatch.setattr(ownership, "_CONTAINER_OWNER_CACHE_MAX_SIZE", 2) - ownership._write_container_owner_cache("a", "user-a") - ownership._write_container_owner_cache("b", "user-b") - ownership._write_container_owner_cache("c", "user-c") - # ``a`` was the oldest and is dropped; ``b`` and ``c`` survive. - assert list(ownership._CONTAINER_OWNER_CACHE.keys()) == ["b", "c"] - - -def test_container_owner_cache_read_marks_as_recently_used(monkeypatch): - """Reading an entry should reset its position so a subsequent eviction - drops a less-recently-used entry instead of the just-touched one.""" - monkeypatch.setattr(ownership, "_CONTAINER_OWNER_CACHE_MAX_SIZE", 2) - ownership._write_container_owner_cache("a", "user-a") - ownership._write_container_owner_cache("b", "user-b") - # Touch ``a`` so it becomes the most-recently-used. - ownership._read_container_owner_cache("a") - ownership._write_container_owner_cache("c", "user-c") - # ``b`` is the LRU at this point; ``a`` and ``c`` survive. - assert list(ownership._CONTAINER_OWNER_CACHE.keys()) == ["a", "c"] 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 c6dd6c4e33ca..3ffba9723bd0 100644 --- a/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py +++ b/tests/test_litellm/llms/litellm_proxy/test_skills_ownership.py @@ -19,9 +19,11 @@ @pytest.fixture(autouse=True) def clear_skill_cache(): - skills_handler._SKILL_CACHE.clear() + skills_handler._SKILL_CACHE.cache_dict.clear() + skills_handler._SKILL_CACHE.ttl_dict.clear() yield - skills_handler._SKILL_CACHE.clear() + skills_handler._SKILL_CACHE.cache_dict.clear() + skills_handler._SKILL_CACHE.ttl_dict.clear() def _skill(skill_id: str, created_by: str | None) -> LiteLLM_SkillsTable: @@ -440,99 +442,75 @@ async def test_should_scope_skill_injection_fetch_to_authenticated_user(monkeypa @pytest.mark.asyncio async def test_load_skill_uses_cache_after_first_db_hit(monkeypatch): - """`fetch_skill_from_db` is hit per-chat-completion; the cache absorbs + """``fetch_skill_from_db`` runs per chat-completion; the cache absorbs repeats so we don't issue a Prisma query on every request.""" fake_skill = Mock(created_by="user-1", skill_id="litellm_skill_a") - store_factory = AsyncMock() - store = Mock() - store.find_skill = AsyncMock(return_value=fake_skill) + table = AsyncMock() + table.find_unique = AsyncMock(return_value=fake_skill) + prisma_client = type( + "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} + )() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", - AsyncMock(return_value=store_factory), - ) - monkeypatch.setattr( - skills_handler, - "LiteLLMSkillsStore", - Mock(return_value=store), + AsyncMock(return_value=prisma_client), ) - first = await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") - second = await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") - third = await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") - - assert first is fake_skill - assert second is fake_skill - assert third is fake_skill - assert store.find_skill.await_count == 1 + for _ in range(3): + assert ( + await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") + is fake_skill + ) + assert table.find_unique.await_count == 1 @pytest.mark.asyncio async def test_load_skill_caches_negative_lookups(monkeypatch): - """Missing skills must cache as `None` so repeated lookups skip the DB.""" - store_factory = AsyncMock() - store = Mock() - store.find_skill = AsyncMock(return_value=None) + """Missing skills cache as the negative sentinel so repeated misses skip + the DB and the caller still sees ``None``.""" + table = AsyncMock() + table.find_unique = AsyncMock(return_value=None) + prisma_client = type( + "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} + )() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", - AsyncMock(return_value=store_factory), - ) - monkeypatch.setattr( - skills_handler, - "LiteLLMSkillsStore", - Mock(return_value=store), + AsyncMock(return_value=prisma_client), ) assert await skills_handler.LiteLLMSkillsHandler._load_skill("missing") is None assert await skills_handler.LiteLLMSkillsHandler._load_skill("missing") is None - assert store.find_skill.await_count == 1 + assert table.find_unique.await_count == 1 @pytest.mark.asyncio async def test_delete_skill_invalidates_cache(monkeypatch): - """After delete, the next read must consult the DB rather than the cached - pre-delete row.""" + """After delete, the next read should not see the pre-delete cached row.""" fake_skill = Mock(created_by="user-1", skill_id="litellm_skill_a") - store = Mock() - store.find_skill = AsyncMock(return_value=fake_skill) - store.delete_skill = AsyncMock() + table = AsyncMock() + table.find_unique = AsyncMock(return_value=fake_skill) + table.delete = AsyncMock() + prisma_client = type( + "Prisma", (), {"db": type("DB", (), {"litellm_skillstable": table})()} + )() monkeypatch.setattr( skills_handler.LiteLLMSkillsHandler, "_get_prisma_client", - AsyncMock(return_value=Mock()), - ) - monkeypatch.setattr( - skills_handler, - "LiteLLMSkillsStore", - Mock(return_value=store), + AsyncMock(return_value=prisma_client), ) # Prime the cache via the read path. await skills_handler.LiteLLMSkillsHandler._load_skill("litellm_skill_a") - cached_hit, _ = skills_handler._read_skill_cache("litellm_skill_a") - assert cached_hit + assert skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") is fake_skill auth = UserAPIKeyAuth(user_id="user-1") await skills_handler.LiteLLMSkillsHandler.delete_skill( "litellm_skill_a", user_api_key_dict=auth ) - cached_hit_after, _ = skills_handler._read_skill_cache("litellm_skill_a") - assert not cached_hit_after - - -def test_skill_cache_expires_after_ttl(monkeypatch): - monkeypatch.setattr(skills_handler, "_SKILL_CACHE_TTL", 0.0) - skills_handler._write_skill_cache("k", Mock()) - cached_hit, _ = skills_handler._read_skill_cache("k") - assert not cached_hit - - -def test_skill_cache_evicts_when_at_capacity(monkeypatch): - monkeypatch.setattr(skills_handler, "_SKILL_CACHE_MAX_SIZE", 2) - skills_handler._write_skill_cache("a", Mock()) - skills_handler._write_skill_cache("b", Mock()) - skills_handler._write_skill_cache("c", Mock()) - # ``a`` was the oldest and is LRU-evicted; ``b`` and ``c`` survive. - assert list(skills_handler._SKILL_CACHE.keys()) == ["b", "c"] + # Post-delete, the cache holds the negative sentinel — not the stale row. + assert ( + skills_handler._SKILL_CACHE.get_cache("litellm_skill_a") + == skills_handler._NEGATIVE_SKILL_SENTINEL + ) From 2adfa96db273e71e22bd4d772b3613a20a6b1890 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Tue, 5 May 2026 00:39:53 +0000 Subject: [PATCH 26/27] fix(container): cache list-allow-set, track admin-created containers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address Greptile P2 follow-ups from the prior round: * Cache ``_get_allowed_container_ids`` (60s LRU/TTL keyed by sorted owner-scope tuple) so ``GET /v1/containers`` doesn't issue a fresh ``find_many`` against ``litellm_managedobjecttable`` on every list call. Invalidate the caller's own cache entry when they record a new owner so the just-created container shows up on their next list. * Tighten the admin early-return in ``record_container_owner`` to skip ONLY when there's literally no container ID to stamp. An admin with identity (the master-key path populates ``user_id`` + ``api_key``) flows through the normal record path so admin-created containers are tracked like any other caller's. The truly-identity-less admin case still falls through to the 403 below — correct fail-secure default. Skill-cache invalidation gap (also flagged by Greptile) is moot: there is no skill update endpoint exposed; ownership-affecting mutations are only delete (already invalidates) and create (new ID, no cache entry to update). --- .../proxy/container_endpoints/ownership.py | 39 +++++- .../test_container_proxy_ownership.py | 125 +++++++++++++++++- 2 files changed, 153 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 137366c955a5..41dc1f34e011 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -22,6 +22,13 @@ _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 b046fa1536ba..7a6232d7ce70 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 @@ @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" From 4699b3dc8106e442d7a852e58357567f2866a9f1 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Tue, 5 May 2026 00:43:47 +0000 Subject: [PATCH 27/27] chore(container): use delete_cache, json-encode scope key, clean test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit /simplify follow-ups: * Replace the two-``pop`` reach into ``cache_dict``/``ttl_dict`` with the existing public ``InMemoryCache.delete_cache(key)`` — the same idiom used elsewhere in the proxy. Bonus: ``delete_cache`` calls ``_remove_key`` which also handles ``expiration_heap`` consistency the direct pops were silently leaking. * JSON-encode the sorted scope list for the cache key instead of ``"|".join``. ``user_id`` / ``team_id`` / ``org_id`` / ``api_key`` are free-form strings and could contain a literal ``|`` — JSON quoting escapes any in-string separator unambiguously. * Extract ``_allowed_container_ids_cache_key()`` so the read and invalidation sites compute the key the same way. * Fix a placeholder-then-overwrite test construction: the ``__module__.split(".")[0] and "proxy_admin"`` line evaluated to a literal string that was immediately overwritten with the real enum value. Hoist the import and construct directly. --- .../proxy/container_endpoints/ownership.py | 27 +++++++++++++------ .../test_container_proxy_ownership.py | 9 ++----- 2 files changed, 21 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 41dc1f34e011..568eca523ae7 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,3 +1,4 @@ +import json from typing import Any, Dict, List, Optional, Set, Tuple from fastapi import HTTPException @@ -30,6 +31,14 @@ _ALLOWED_CONTAINER_IDS_CACHE = InMemoryCache(max_size_in_memory=2048, default_ttl=60) +def _allowed_container_ids_cache_key(owner_scopes: List[str]) -> str: + """JSON-encode the sorted scope list — using a separator like ``|`` + would collide for any tenant whose user_id / team_id / org_id / + api_key happens to contain the separator. JSON quoting escapes + every separator that matters.""" + return json.dumps(sorted(owner_scopes)) + + def _container_model_object_id( original_container_id: str, custom_llm_provider: str ) -> str: @@ -163,10 +172,11 @@ async def record_container_owner( # 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) + caller_scopes = get_resource_owner_scopes(user_api_key_dict) + if caller_scopes: + _ALLOWED_CONTAINER_IDS_CACHE.delete_cache( + _allowed_container_ids_cache_key(caller_scopes) + ) return response @@ -265,7 +275,7 @@ async def _get_allowed_container_ids( if not owner_scopes: return set() - cache_key = "|".join(sorted(owner_scopes)) + cache_key = _allowed_container_ids_cache_key(owner_scopes) cached = _ALLOWED_CONTAINER_IDS_CACHE.get_cache(cache_key) if cached is not None: return set(cached) @@ -285,9 +295,10 @@ async def _get_allowed_container_ids( 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)) + # ``InMemoryCache.get_cache`` attempts ``json.loads`` on the stored + # value; passing a set would round-trip through that path + # unnecessarily. Store as a list and rehydrate above. + _ALLOWED_CONTAINER_IDS_CACHE.set_cache(cache_key, list(allowed_ids)) return allowed_ids diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 7a6232d7ce70..c295805bdb34 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -5,7 +5,7 @@ import pytest from fastapi import HTTPException -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.container_endpoints import ownership from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.containers.main import ContainerListResponse, ContainerObject @@ -901,13 +901,8 @@ async def test_admin_with_identity_records_container_ownership(monkeypatch): ) 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 + user_role=LitellmUserRoles.PROXY_ADMIN.value, ) - # 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"),