diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 6da5fc07e80..6bc44995804 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -151,6 +151,7 @@ jobs: tests/test_litellm/proxy/google_endpoints tests/test_litellm/proxy/openai_files_endpoint tests/test_litellm/proxy/batches_endpoints + tests/test_litellm/proxy/container_endpoints tests/test_litellm/proxy/fine_tuning_endpoints tests/test_litellm/proxy/vector_store_files_endpoints tests/test_litellm/proxy/video_endpoints diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 9e0f6fa741f..f3a2abc1225 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -15,10 +15,11 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin from litellm.proxy.container_endpoints.ownership import ( assert_user_can_access_container, - filter_container_list_response, get_container_forwarding_params, + list_owned_containers, record_container_owner, ) @@ -209,61 +210,54 @@ async def list_containers( version, ) - # Read query parameters query_params: Final = dict(request.query_params) - data: Final[dict[str, Any]] = { - "query_params": query_params, - "model": query_params.get("model"), - "after": after, - "limit": limit, - "order": order, - } - - # Extract custom_llm_provider using priority chain custom_llm_provider: Final = ( get_custom_llm_provider_from_request_headers(request=request) or get_custom_llm_provider_from_request_query(request=request) or "openai" ) + data: Final[dict[str, Any]] = { + "query_params": query_params, + "model": query_params.get("model"), + "order": order, + "custom_llm_provider": custom_llm_provider, + } - # Add custom_llm_provider to data - data["custom_llm_provider"] = custom_llm_provider + async def fetch_page(page_after: str | None, page_limit: int | None) -> object: + processor: Final = ProxyBaseLLMRequestProcessing(data={**data, "after": page_after, "limit": page_limit}) + try: + return await processor.base_process_llm_request( + request=request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="alist_containers", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) + except Exception as e: + raise await processor._handle_llm_api_exception( + e=e, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + version=version, + ) - # Process request using ProxyBaseLLMRequestProcessing - processor: Final = ProxyBaseLLMRequestProcessing(data=data) - try: - response: Final = await processor.base_process_llm_request( - request=request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - route_type="alist_containers", - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - model=None, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, - ) - except Exception as e: - raise await processor._handle_llm_api_exception( - e=e, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, - 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, + if is_proxy_admin(user_api_key_dict): + return await fetch_page(after, limit) + return await list_owned_containers( + fetch_page=fetch_page, + after=after, + limit=limit, user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index e3088771c82..14480232a4a 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -1,9 +1,10 @@ import json -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence from collections.abc import Set as AbstractSet -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypeAlias from fastapi import HTTPException +from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache @@ -46,6 +47,12 @@ _CONTAINER_STORED_ID_CACHE: Final = InMemoryCache(max_size_in_memory=10000, defa # different users with different scopes get disjoint cache entries. _ALLOWED_CONTAINER_IDS_CACHE: Final = InMemoryCache(max_size_in_memory=2048, default_ttl=60) +DEFAULT_CONTAINER_LIST_LIMIT: Final = 20 +OWNED_CONTAINER_LIST_PAGE_SIZE: Final = 100 +OWNED_CONTAINER_LIST_MAX_PAGES: Final = 5 + +FetchContainerListPage: TypeAlias = Callable[[str | None, int | None], Awaitable[object]] + def _allowed_container_ids_cache_key(owner_scopes: Sequence[str]) -> str: """JSON-encode the sorted scope list — using a separator like ``|`` @@ -337,27 +344,23 @@ def _get_container_list_data(response: object) -> Sequence[object] | None: return data if isinstance(data, list) else None -def _set_container_list_data(response: Any, data: list[object], removed_filtered_items: bool = False) -> object: +def _get_has_more(response: object) -> bool: 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 - response["has_more"] = False - if removed_filtered_items: - response["has_more"] = False - return response + return response.get("has_more") is True + return getattr(response, "has_more", None) is True - 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 - if removed_filtered_items and hasattr(response, "has_more"): - response.has_more = False + +def _with_container_list_page(response: object, data: Sequence[object], has_more: bool) -> object: + page: Final = { + "data": list(data), + "first_id": _get_response_id(data[0]) if data else None, + "last_id": _get_response_id(data[-1]) if data else None, + "has_more": has_more, + } + if isinstance(response, dict): + return {**response, **page} + if isinstance(response, BaseModel): + return response.model_copy(update=page) return response @@ -366,16 +369,16 @@ async def _get_allowed_container_ids( ) -> AbstractSet[str]: owner_scopes: Final = get_resource_owner_scopes(user_api_key_dict) if not owner_scopes: - return set() + return frozenset() cache_key: Final = _allowed_container_ids_cache_key(owner_scopes) cached: Final = _ALLOWED_CONTAINER_IDS_CACHE.get_cache(cache_key) if cached is not None: - return set(cached) + return frozenset(cached) prisma_client: Final = await _get_prisma_client() if prisma_client is None: - return set() + return frozenset() table: Final = ManagedObjectRepository(prisma_client).table rows: Final[Sequence[prisma_models.LiteLLM_ManagedObjectTable]] = await table.find_many( @@ -384,34 +387,69 @@ async def _get_allowed_container_ids( "created_by": {"in": owner_scopes}, } ) - allowed_ids: Final = {row.model_object_id for row in rows if getattr(row, "model_object_id", None) is not None} - # ``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)) + allowed_ids: Final = frozenset( + row.model_object_id for row in rows if getattr(row, "model_object_id", None) is not None + ) + _ALLOWED_CONTAINER_IDS_CACHE.set_cache(cache_key, tuple(allowed_ids)) return allowed_ids -async def filter_container_list_response( - response: object, +def _is_owned_container(item: object, allowed_container_ids: AbstractSet[str], custom_llm_provider: str) -> bool: + container_id: Final = _get_response_id(item) + if container_id is None: + return False + original_container_id, resolved_provider = decode_container_id_for_ownership(container_id, custom_llm_provider) + return _container_model_object_id(original_container_id, resolved_provider) in allowed_container_ids + + +async def _collect_owned_containers( + fetch_page: FetchContainerListPage, + after: str | None, + needed: int, + allowed_container_ids: AbstractSet[str], + custom_llm_provider: str, + pages_left: int, + collected: tuple[object, ...], +) -> tuple[object, tuple[object, ...]]: + page: Final = await fetch_page(after, OWNED_CONTAINER_LIST_PAGE_SIZE) + page_data: Final = _get_container_list_data(page) or () + owned: Final = collected + tuple( + item for item in page_data if _is_owned_container(item, allowed_container_ids, custom_llm_provider) + ) + upstream_last_id: Final = _get_response_id(page_data[-1]) if page_data else None + if len(owned) >= needed or upstream_last_id is None or pages_left <= 1 or not _get_has_more(page): + return page, owned + return await _collect_owned_containers( + fetch_page=fetch_page, + after=upstream_last_id, + needed=needed, + allowed_container_ids=allowed_container_ids, + custom_llm_provider=custom_llm_provider, + pages_left=pages_left - 1, + collected=owned, + ) + + +async def list_owned_containers( + fetch_page: FetchContainerListPage, + after: str | None, + limit: int | None, user_api_key_dict: UserAPIKeyAuth, custom_llm_provider: str, ) -> object: - if is_proxy_admin(user_api_key_dict): - return response - - data: Final = _get_container_list_data(response) - if data is None: - return response - allowed_container_ids: Final = await _get_allowed_container_ids(user_api_key_dict) - filtered: Final[list[object]] = [] - 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, removed_filtered_items=len(filtered) != len(data)) + page_limit: Final = limit if limit is not None else DEFAULT_CONTAINER_LIST_LIMIT + last_page, owned = await _collect_owned_containers( + fetch_page=fetch_page, + after=after, + needed=page_limit + 1, + allowed_container_ids=allowed_container_ids, + custom_llm_provider=custom_llm_provider, + pages_left=OWNED_CONTAINER_LIST_MAX_PAGES, + collected=(), + ) + return _with_container_list_page( + last_page, + owned[:page_limit], + has_more=len(owned) > page_limit or _get_has_more(last_page), + ) diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 176405bb9ca..38988d65c04 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -1,7 +1,7 @@ import json import sys from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException @@ -242,110 +242,156 @@ async def test_should_not_reassign_existing_container_to_different_owner(monkeyp table.update.assert_not_awaited() -@pytest.mark.asyncio -async def test_should_filter_container_list_to_owned_records(monkeypatch): +def _owned_containers_in_db(monkeypatch, *model_object_ids: str) -> AsyncMock: table = AsyncMock() - table.find_many.return_value = [ - SimpleNamespace(model_object_id="container:openai:cntr_owned"), - ] - prisma_client = SimpleNamespace( - db=SimpleNamespace(litellm_managedobjecttable=table) - ) + table.find_many.return_value = [SimpleNamespace(model_object_id=object_id) for object_id in model_object_ids] monkeypatch.setattr( ownership, "_get_prisma_client", - AsyncMock(return_value=prisma_client), + AsyncMock(return_value=SimpleNamespace(db=SimpleNamespace(litellm_managedobjecttable=table))), ) - auth = UserAPIKeyAuth(user_id="user-1") - response = ContainerListResponse( + return table + + +def _upstream(pages_by_after): + calls = [] + + async def fetch_page(after, limit): + calls.append((after, limit)) + return pages_by_after[after] + + return fetch_page, calls + + +def _page(*container_ids: str, has_more: bool) -> ContainerListResponse: + return ContainerListResponse( object="list", - data=[_container("cntr_owned"), _container("cntr_other")], - has_more=True, + data=[_container(container_id) for container_id in container_ids], + has_more=has_more, ) - filtered = await ownership.filter_container_list_response( - response=response, - user_api_key_dict=auth, + +async def _list_owned(fetch_page, after=None, limit=None): + return await ownership.list_owned_containers( + fetch_page=fetch_page, + after=after, + limit=limit, + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), 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" - assert filtered.has_more is False + +@pytest.mark.asyncio +async def test_should_page_upstream_until_owned_containers_fill_the_limit(monkeypatch): + table = _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned") + fetch_page, calls = _upstream( + { + None: _page("cntr_other_1", "cntr_other_2", has_more=True), + "cntr_other_2": _page("cntr_owned", has_more=False), + } + ) + + listed = await _list_owned(fetch_page, limit=1) + + assert [item.id for item in listed.data] == ["cntr_owned"] + assert listed.first_id == "cntr_owned" + assert listed.last_id == "cntr_owned" + assert listed.has_more is False + assert calls == [(None, 100), ("cntr_other_2", 100)] 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_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, - ) +async def test_should_trim_owned_containers_to_the_limit_without_mutating_the_upstream_page(monkeypatch): + _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned_1", "container:openai:cntr_owned_2") + upstream_page = _page("cntr_owned_1", "cntr_other", "cntr_owned_2", has_more=False) + fetch_page, calls = _upstream({None: upstream_page}) - filtered = await ownership.filter_container_list_response( - response=response, - user_api_key_dict=auth, - custom_llm_provider="openai", - ) + listed = await _list_owned(fetch_page, limit=1) - assert filtered.data == [] - assert filtered.first_id is None - assert filtered.last_id is None - assert filtered.has_more is False + assert [item.id for item in listed.data] == ["cntr_owned_1"] + assert listed.first_id == "cntr_owned_1" + assert listed.last_id == "cntr_owned_1" + assert listed.has_more is True + assert calls == [(None, 100)] + assert [item.id for item in upstream_page.data] == ["cntr_owned_1", "cntr_other", "cntr_owned_2"] + assert upstream_page.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) +async def test_should_start_paging_from_the_requested_cursor(monkeypatch): + _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned_2") + fetch_page, calls = _upstream({"cntr_owned_1": _page("cntr_other", "cntr_owned_2", has_more=False)}) + + listed = await _list_owned(fetch_page, after="cntr_owned_1", limit=1) + + assert [item.id for item in listed.data] == ["cntr_owned_2"] + assert listed.has_more is False + assert calls == [("cntr_owned_1", 100)] + + +@pytest.mark.asyncio +async def test_should_default_to_twenty_owned_containers_per_page(monkeypatch): + owned_ids = tuple(f"cntr_owned_{index}" for index in range(21)) + _owned_containers_in_db(monkeypatch, *(f"container:openai:{container_id}" for container_id in owned_ids)) + fetch_page, _ = _upstream({None: _page(*owned_ids, has_more=False)}) + + listed = await _list_owned(fetch_page) + + assert [item.id for item in listed.data] == list(owned_ids[:20]) + assert listed.last_id == "cntr_owned_19" + assert listed.has_more is True + + +@pytest.mark.asyncio +async def test_should_stop_after_five_upstream_pages_and_keep_has_more(monkeypatch): + _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned") + fetch_page, calls = _upstream( + { + None: _page("cntr_other_0", has_more=True), + **{f"cntr_other_{index}": _page(f"cntr_other_{index + 1}", has_more=True) for index in range(6)}, + } ) - monkeypatch.setattr( - ownership, - "_get_prisma_client", - AsyncMock(return_value=prisma_client), - ) - auth = UserAPIKeyAuth(user_id="user-1") - response = { + + listed = await _list_owned(fetch_page, limit=1) + + assert listed.data == [] + assert listed.first_id is None + assert listed.last_id is None + assert listed.has_more is True + assert len(calls) == 5 + + +@pytest.mark.asyncio +async def test_should_stop_when_upstream_has_no_more_pages(monkeypatch): + _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned") + fetch_page, calls = _upstream({None: _page("cntr_other", has_more=False)}) + + listed = await _list_owned(fetch_page, limit=1) + + assert listed.data == [] + assert listed.has_more is False + assert calls == [(None, 100)] + + +@pytest.mark.asyncio +async def test_should_build_dict_pages_without_mutating_the_upstream_page(monkeypatch): + _owned_containers_in_db(monkeypatch, "container:openai:cntr_owned") + upstream_page = {"object": "list", "data": [{"id": "cntr_other"}, {"id": "cntr_owned"}], "has_more": False} + fetch_page, _ = _upstream({None: upstream_page}) + + listed = await _list_owned(fetch_page, limit=1) + + assert listed == { "object": "list", - "data": [{"id": "cntr_other"}], - "has_more": True, + "data": [{"id": "cntr_owned"}], + "first_id": "cntr_owned", + "last_id": "cntr_owned", + "has_more": False, } - - 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 + assert [item["id"] for item in upstream_page["data"]] == ["cntr_other", "cntr_owned"] @pytest.mark.asyncio @@ -647,7 +693,7 @@ async def test_should_return_response_when_owner_recording_raises_unexpected( @pytest.mark.asyncio -async def test_should_filter_container_list_inside_list_endpoint(monkeypatch): +async def test_should_list_owned_containers_inside_list_endpoint(monkeypatch): from litellm.proxy.container_endpoints import endpoints proxy_server_stub = SimpleNamespace( @@ -665,42 +711,37 @@ async def test_should_filter_container_list_inside_list_endpoint(monkeypatch): ) 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, + upstream_page = _page("cntr_provider", has_more=False) + processor_cls = MagicMock( + side_effect=lambda data: SimpleNamespace(base_process_llm_request=AsyncMock(return_value=upstream_page)) ) + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", processor_cls) + list_owned = AsyncMock(return_value=upstream_page) + monkeypatch.setattr(endpoints, "list_owned_containers", list_owned) result = await endpoints.list_containers( request=SimpleNamespace(query_params={}, headers={}), fastapi_response=SimpleNamespace(), user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + after="cntr_prev", + limit=2, + order="desc", ) - assert result == response - filter_response.assert_awaited_once_with( - response=response, - user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), - custom_llm_provider="openai", - ) + assert result == upstream_page + kwargs = list_owned.await_args.kwargs + assert kwargs["after"] == "cntr_prev" + assert kwargs["limit"] == 2 + assert kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_id="user-1") + assert kwargs["custom_llm_provider"] == "openai" + processor_cls.assert_not_called() + + assert await kwargs["fetch_page"]("cntr_page_cursor", 100) == upstream_page + forwarded = processor_cls.call_args.kwargs["data"] + assert forwarded["after"] == "cntr_page_cursor" + assert forwarded["limit"] == 100 + assert forwarded["order"] == "desc" + assert forwarded["custom_llm_provider"] == "openai" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/container_endpoints/test_endpoints.py b/tests/test_litellm/proxy/container_endpoints/test_endpoints.py index 3604da65258..1beff4c82ba 100644 --- a/tests/test_litellm/proxy/container_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/container_endpoints/test_endpoints.py @@ -1,14 +1,15 @@ import sys from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock +import pytest from fastapi import FastAPI from fastapi.testclient import TestClient -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.container_endpoints import endpoints -from litellm.types.containers.main import ContainerListResponse +from litellm.proxy.container_endpoints import endpoints, ownership +from litellm.types.containers.main import ContainerListResponse, ContainerObject PROXY_SERVER_STUB = SimpleNamespace( general_settings={}, @@ -24,41 +25,110 @@ PROXY_SERVER_STUB = SimpleNamespace( user_temperature=None, version="test", ) +ADMIN = UserAPIKeyAuth(user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN) +NON_ADMIN = UserAPIKeyAuth(user_id="user-1") -def _client() -> TestClient: +@pytest.fixture(autouse=True) +def clear_allowed_container_ids_cache(): + ownership._ALLOWED_CONTAINER_IDS_CACHE.cache_dict.clear() + ownership._ALLOWED_CONTAINER_IDS_CACHE.ttl_dict.clear() + yield + ownership._ALLOWED_CONTAINER_IDS_CACHE.cache_dict.clear() + ownership._ALLOWED_CONTAINER_IDS_CACHE.ttl_dict.clear() + + +def _client(auth: UserAPIKeyAuth) -> TestClient: app = FastAPI() app.include_router(endpoints.router) - app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="user-1") + app.dependency_overrides[user_api_key_auth] = lambda: auth return TestClient(app) -def test_list_containers_forwards_typed_pagination_params(monkeypatch): +def _container(container_id: str) -> ContainerObject: + return ContainerObject(id=container_id, object="container", created_at=1, status="active") + + +def _page(*container_ids: str, has_more: bool) -> ContainerListResponse: + return ContainerListResponse( + object="list", + data=[_container(container_id) for container_id in container_ids], + has_more=has_more, + ) + + +def _upstream_pages(monkeypatch, pages_by_after) -> MagicMock: + processor_cls = MagicMock( + side_effect=lambda data: SimpleNamespace( + base_process_llm_request=AsyncMock(return_value=pages_by_after[data["after"]]) + ) + ) + monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", processor_cls) + return processor_cls + + +def _forwarded_pages(processor_cls: MagicMock): + return [(call.kwargs["data"]["after"], call.kwargs["data"]["limit"]) for call in processor_cls.call_args_list] + + +def test_list_containers_forwards_typed_pagination_params_for_admins(monkeypatch): monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", PROXY_SERVER_STUB) - upstream = ContainerListResponse(object="list", data=[], has_more=True) - captured = {} + processor_cls = _upstream_pages(monkeypatch, {"cntr_prev": _page("cntr_next", has_more=True)}) - class FakeProcessor: - def __init__(self, data): - captured["data"] = data - - async def base_process_llm_request(self, **kwargs): - return upstream - - async def _handle_llm_api_exception(self, **kwargs): - raise kwargs["e"] - - monkeypatch.setattr(endpoints, "ProxyBaseLLMRequestProcessing", FakeProcessor) - monkeypatch.setattr(endpoints, "filter_container_list_response", AsyncMock(return_value=upstream)) - - response = _client().get( + response = _client(ADMIN).get( "/v1/containers", params={"limit": "1", "order": "desc", "after": "cntr_prev"}, headers={"Authorization": "Bearer sk-test"}, ) assert response.status_code == 200 + assert [item["id"] for item in response.json()["data"]] == ["cntr_next"] assert response.json()["has_more"] is True - assert captured["data"]["limit"] == 1 - assert captured["data"]["order"] == "desc" - assert captured["data"]["after"] == "cntr_prev" + assert _forwarded_pages(processor_cls) == [("cntr_prev", 1)] + assert processor_cls.call_args.kwargs["data"]["order"] == "desc" + + +def test_list_containers_rejects_a_non_integer_limit(monkeypatch): + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", PROXY_SERVER_STUB) + processor_cls = _upstream_pages(monkeypatch, {}) + + response = _client(ADMIN).get( + "/v1/containers", + params={"limit": "abc"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 422 + processor_cls.assert_not_called() + + +def test_list_containers_pages_upstream_until_non_admin_keys_see_their_containers(monkeypatch): + monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", PROXY_SERVER_STUB) + table = AsyncMock() + table.find_many.return_value = [SimpleNamespace(model_object_id="container:openai:cntr_owned")] + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=SimpleNamespace(db=SimpleNamespace(litellm_managedobjecttable=table))), + ) + processor_cls = _upstream_pages( + monkeypatch, + { + None: _page("cntr_other", has_more=True), + "cntr_other": _page("cntr_owned", has_more=False), + }, + ) + + response = _client(NON_ADMIN).get( + "/v1/containers", + params={"limit": "1"}, + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + body = response.json() + assert [item["id"] for item in body["data"]] == ["cntr_owned"] + assert body["first_id"] == "cntr_owned" + assert body["last_id"] == "cntr_owned" + assert body["has_more"] is False + assert _forwarded_pages(processor_cls) == [(None, 100), ("cntr_other", 100)] diff --git a/tests/test_litellm/proxy/container_endpoints/test_handler_factory.py b/tests/test_litellm/proxy/container_endpoints/test_handler_factory.py index 01a03f6ef03..a471f915071 100644 --- a/tests/test_litellm/proxy/container_endpoints/test_handler_factory.py +++ b/tests/test_litellm/proxy/container_endpoints/test_handler_factory.py @@ -1,6 +1,6 @@ import sys from types import SimpleNamespace -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, MagicMock from fastapi import FastAPI from fastapi.testclient import TestClient @@ -39,20 +39,11 @@ def test_list_container_files_forwards_declared_query_params(monkeypatch): "assert_user_can_access_container", AsyncMock(return_value=("cntr_123", "openai")), ) - captured = {} - - class FakeProcessor: - def __init__(self, data): - captured["data"] = data - - async def base_process_llm_request(self, **kwargs): - captured["route_type"] = kwargs["route_type"] - return {"object": "list", "data": [], "has_more": True} - - async def _handle_llm_api_exception(self, **kwargs): - raise kwargs["e"] - - monkeypatch.setattr(handler_factory, "ProxyBaseLLMRequestProcessing", FakeProcessor) + processor_cls = MagicMock() + processor_cls.return_value.base_process_llm_request = AsyncMock( + return_value={"object": "list", "data": [], "has_more": True} + ) + monkeypatch.setattr(handler_factory, "ProxyBaseLLMRequestProcessing", processor_cls) response = _client().get( "/v1/containers/cntr_123/files", @@ -61,9 +52,11 @@ def test_list_container_files_forwards_declared_query_params(monkeypatch): ) assert response.status_code == 200 - assert captured["route_type"] == "alist_container_files" - assert captured["data"]["container_id"] == "cntr_123" - assert captured["data"]["limit"] == "1" - assert captured["data"]["order"] == "desc" - assert captured["data"]["after"] == "cfile_prev" - assert "unknown" not in captured["data"] + assert response.json()["has_more"] is True + assert processor_cls.return_value.base_process_llm_request.await_args.kwargs["route_type"] == "alist_container_files" + forwarded = processor_cls.call_args.kwargs["data"] + assert forwarded["container_id"] == "cntr_123" + assert forwarded["limit"] == "1" + assert forwarded["order"] == "desc" + assert forwarded["after"] == "cfile_prev" + assert "unknown" not in forwarded