scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

fix(containers): page upstream until a non-admin container list fills its limit

Browse source 
Forwarding limit to OpenAI made the ownership filter cut the page down after
the fact, so a key that owned an older container got an empty first page and
its cursor never moved. Non-admin lists now walk upstream pages of 100 until
they have enough owned containers (or five pages), trim to the requested
limit, and report first_id, last_id and has_more off what the caller keeps.

Also assigns tests/test_litellm/proxy/container_endpoints to a CI shard.
This commit is contained in:
mateo-berri 2026-09-02 21:26:34 -07:00
parent 872e115295
commit 425c8d37fc
6 changed files with 389 additions and 252 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit.yml vendored
View file

@ -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

90
litellm/proxy/container_endpoints/endpoints.py
View file

@ -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,
)

134
litellm/proxy/container_endpoints/ownership.py
View file

@ -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),
)

259
tests/test_litellm/containers/test_container_proxy_ownership.py
View file

@ -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

122
tests/test_litellm/proxy/container_endpoints/test_endpoints.py
View file

@ -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)]

35
tests/test_litellm/proxy/container_endpoints/test_handler_factory.py
View file

@ -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