diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 9d4dd90299..93587262f5 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 67c1d4b0f5..bb6f8a8db6 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 1a62abc9f3..e3d8da8b0d 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.proxy.common_utils.resource_ownership import ( 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 760e522199..aed84b4fbe 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 fastapi import HTTPException 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"