mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
425c8d37fc
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.
1150 lines
37 KiB
Python
1150 lines
37 KiB
Python
import json
|
|
import sys
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
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
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clear_container_owner_cache():
|
|
for cache in (
|
|
ownership._CONTAINER_OWNER_CACHE,
|
|
ownership._ALLOWED_CONTAINER_IDS_CACHE,
|
|
):
|
|
cache.cache_dict.clear()
|
|
cache.ttl_dict.clear()
|
|
yield
|
|
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:
|
|
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_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"]
|
|
file_obj = json.loads(data["file_object"])
|
|
assert file_obj["custom_llm_provider"] == "openai"
|
|
assert file_obj["provider_container_id"] == "cntr_provider"
|
|
|
|
|
|
@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_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_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",
|
|
AsyncMock(return_value=None),
|
|
)
|
|
auth = UserAPIKeyAuth()
|
|
|
|
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)
|
|
|
|
|
|
@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_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_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()
|
|
|
|
|
|
def _owned_containers_in_db(monkeypatch, *model_object_ids: str) -> AsyncMock:
|
|
table = AsyncMock()
|
|
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=SimpleNamespace(db=SimpleNamespace(litellm_managedobjecttable=table))),
|
|
)
|
|
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(container_id) for container_id in container_ids],
|
|
has_more=has_more,
|
|
)
|
|
|
|
|
|
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",
|
|
)
|
|
|
|
|
|
@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_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})
|
|
|
|
listed = await _list_owned(fetch_page, limit=1)
|
|
|
|
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_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)},
|
|
}
|
|
)
|
|
|
|
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_owned"}],
|
|
"first_id": "cntr_owned",
|
|
"last_id": "cntr_owned",
|
|
"has_more": False,
|
|
}
|
|
assert [item["id"] for item in upstream_page["data"]] == ["cntr_other", "cntr_owned"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_validate_owner_and_forward_decoded_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,
|
|
)
|
|
access_check = AsyncMock(return_value=("cntr_provider", "azure"))
|
|
monkeypatch.setattr(
|
|
handler_factory,
|
|
"assert_user_can_access_container",
|
|
access_check,
|
|
)
|
|
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,
|
|
)
|
|
|
|
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"
|
|
assert result["file"] == "file-data"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_forward_decoded_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"] == "cntr_provider"
|
|
assert result["custom_llm_provider"] == "azure"
|
|
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_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_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_list_owned_containers_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)
|
|
|
|
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 == 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
|
|
async def test_should_forward_decoded_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"] == "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_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=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"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_record_containers_from_responses_output_for_service_account(
|
|
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")
|
|
encoded_container_id = (
|
|
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmR"
|
|
"lZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl"
|
|
)
|
|
responses_payload = {
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"content": [
|
|
{
|
|
"type": "output_text",
|
|
"annotations": [
|
|
{
|
|
"type": "container_file_citation",
|
|
"container_id": encoded_container_id,
|
|
"file_id": "cfile_abc",
|
|
}
|
|
],
|
|
}
|
|
],
|
|
}
|
|
],
|
|
"_hidden_params": {"custom_llm_provider": "azure"},
|
|
}
|
|
|
|
await ownership.record_container_owners_from_responses_response(
|
|
response=responses_payload,
|
|
user_api_key_dict=auth,
|
|
)
|
|
|
|
table.create.assert_awaited_once()
|
|
created_data = table.create.await_args.kwargs["data"]
|
|
assert created_data["created_by"] == "team:team-1"
|
|
assert created_data["unified_object_id"] == encoded_container_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_service_account_can_access_container_after_responses_tracking(
|
|
monkeypatch,
|
|
):
|
|
encoded_container_id = (
|
|
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmR"
|
|
"lZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl"
|
|
)
|
|
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_owners_from_responses_response(
|
|
response={
|
|
"output": [
|
|
{
|
|
"type": "code_interpreter_call",
|
|
"container_id": encoded_container_id,
|
|
}
|
|
],
|
|
"_hidden_params": {"custom_llm_provider": "azure"},
|
|
},
|
|
user_api_key_dict=auth,
|
|
)
|
|
|
|
original_id, provider = await ownership.assert_user_can_access_container(
|
|
container_id=encoded_container_id,
|
|
user_api_key_dict=auth,
|
|
custom_llm_provider="azure",
|
|
)
|
|
assert original_id == "cntr_native"
|
|
assert provider == "azure"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_record_container_ownership_after_streaming_responses_finish(
|
|
monkeypatch,
|
|
):
|
|
"""Streaming /v1/responses calls return through the
|
|
``select_data_generator`` branch and never reach the non-streaming
|
|
container-ownership tail. The wrapper must read
|
|
``completed_response`` off the upstream iterator once iteration
|
|
finishes and write the row, otherwise code-interpreter containers
|
|
created during the stream stay unregistered and follow-up file API
|
|
calls 403.
|
|
"""
|
|
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
|
|
|
encoded_container_id = (
|
|
"cntr_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOmF6dXJlO21vZGVsX2lkOmR"
|
|
"lZi0xMjM7Y29udGFpbmVyX2lkOmNudHJfbmF0aXZl"
|
|
)
|
|
response_body = SimpleNamespace(
|
|
output=[
|
|
SimpleNamespace(
|
|
type="code_interpreter_call",
|
|
container_id=encoded_container_id,
|
|
code_interpreter_call=None,
|
|
)
|
|
]
|
|
)
|
|
stream_response = SimpleNamespace(
|
|
completed_response=SimpleNamespace(response=response_body),
|
|
_hidden_params={"custom_llm_provider": "azure"},
|
|
)
|
|
|
|
async def fake_sse_generator():
|
|
yield "data: chunk-1\n\n"
|
|
yield "data: chunk-2\n\n"
|
|
|
|
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")
|
|
|
|
wrapped = (
|
|
ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership(
|
|
original_stream_response=stream_response,
|
|
wrapped_generator=fake_sse_generator(),
|
|
user_api_key_dict=auth,
|
|
)
|
|
)
|
|
|
|
chunks = [chunk async for chunk in wrapped]
|
|
assert chunks == ["data: chunk-1\n\n", "data: chunk-2\n\n"]
|
|
|
|
table.create.assert_awaited_once()
|
|
created_data = table.create.await_args.kwargs["data"]
|
|
assert created_data["created_by"] == "team:team-1"
|
|
assert created_data["unified_object_id"] == encoded_container_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_streaming_ownership_wrap_no_op_when_stream_did_not_complete(
|
|
monkeypatch,
|
|
):
|
|
"""If the stream errored before ``response.completed``,
|
|
``completed_response`` is ``None`` — we must skip the ownership
|
|
write rather than crash the response generator."""
|
|
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
|
|
|
stream_response = SimpleNamespace(completed_response=None)
|
|
|
|
async def fake_sse_generator():
|
|
yield "data: chunk-1\n\n"
|
|
|
|
record = AsyncMock()
|
|
monkeypatch.setattr(
|
|
ownership,
|
|
"record_container_owners_from_responses_response",
|
|
record,
|
|
)
|
|
|
|
wrapped = (
|
|
ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership(
|
|
original_stream_response=stream_response,
|
|
wrapped_generator=fake_sse_generator(),
|
|
user_api_key_dict=UserAPIKeyAuth(user_id="user-1"),
|
|
)
|
|
)
|
|
chunks = [chunk async for chunk in wrapped]
|
|
|
|
assert chunks == ["data: chunk-1\n\n"]
|
|
record.assert_not_awaited()
|