ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(proxy): preserve container ownership compatibility

Browse Source
This commit is contained in:
user 2026-04-30 18:40:26 -07:00
parent 95c8406b58
commit d2a322074f
4 changed files with 292 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
litellm/proxy/container_endpoints/endpoints.py
View File

@ -295,11 +295,14 @@ async def retrieve_container(
)
# Add custom_llm_provider to data
_, custom_llm_provider = await assert_user_can_access_container(
container_access = await assert_user_can_access_container(
container_id=container_id,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
custom_llm_provider = container_access[1]
# Keep the managed id in request data so downstream container utilities can
# preserve encoded routing metadata while decoding before the provider call.
data["custom_llm_provider"] = custom_llm_provider
# Process request using ProxyBaseLLMRequestProcessing
@ -394,11 +397,14 @@ async def delete_container(
)
# Add custom_llm_provider to data
_, custom_llm_provider = await assert_user_can_access_container(
container_access = await assert_user_can_access_container(
container_id=container_id,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
custom_llm_provider = container_access[1]
# Keep the managed id in request data so downstream container utilities can
# preserve encoded routing metadata while decoding before the provider call.
data["custom_llm_provider"] = custom_llm_provider
# Process request using ProxyBaseLLMRequestProcessing

12
litellm/proxy/container_endpoints/handler_factory.py
View File

@ -283,12 +283,16 @@ async def _process_multipart_upload_request(
or "openai"
)
_, custom_llm_provider = await assert_user_can_access_container(
container_access = await assert_user_can_access_container(
container_id=container_id,
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
custom_llm_provider = container_access[1]
# Keep the managed container id in the forwarded request. The container API
# layer decodes it before the upstream provider call and uses embedded
# routing metadata to preserve model/deployment affinity.
data["container_id"] = container_id
data["custom_llm_provider"] = custom_llm_provider
@ -355,13 +359,15 @@ async def _process_request(
or "openai"
)
# Decode container_id if present in path_params
# Validate container_id ownership if present in path_params.
if "container_id" in path_params:
_, custom_llm_provider = await assert_user_can_access_container(
container_access = await assert_user_can_access_container(
container_id=path_params["container_id"],
user_api_key_dict=user_api_key_dict,
custom_llm_provider=custom_llm_provider,
)
custom_llm_provider = container_access[1]
# Preserve the managed id for downstream container decoding/routing.
data["custom_llm_provider"] = custom_llm_provider

47
litellm/proxy/container_endpoints/ownership.py
View File

@ -1,3 +1,4 @@
import os
from typing import Any, Dict, List, Optional, Set, Tuple
from fastapi import HTTPException
@ -13,6 +14,16 @@ from litellm.proxy.common_utils.resource_ownership import (
from litellm.responses.utils import ResponsesAPIRequestUtils
CONTAINER_OBJECT_PURPOSE = "container"
ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV = "LITELLM_ALLOW_UNTRACKED_CONTAINER_ACCESS"
_IN_MEMORY_CONTAINER_OWNERS: Dict[str, str] = {}
def _allow_untracked_container_access() -> bool:
return os.getenv(ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, "").lower() in {
"1",
"true",
"yes",
}
def _container_model_object_id(
@ -68,11 +79,9 @@ async def record_container_owner(
container_id = _get_response_id(response)
owner = get_primary_resource_owner_scope(user_api_key_dict)
prisma_client = await _get_prisma_client()
if is_proxy_admin(user_api_key_dict) and (
container_id is None or owner is None or prisma_client is None
):
if is_proxy_admin(user_api_key_dict) and (container_id is None or owner is None):
return response
if container_id is None or owner is None or prisma_client is None:
if container_id is None or owner is None:
raise HTTPException(status_code=500, detail="Unable to track container")
original_container_id, resolved_provider = decode_container_id_for_ownership(
@ -87,6 +96,15 @@ async def record_container_owner(
file_object["custom_llm_provider"] = resolved_provider
file_object["provider_container_id"] = original_container_id
if prisma_client is None:
existing_owner = _IN_MEMORY_CONTAINER_OWNERS.get(model_object_id)
if existing_owner is not None and not user_can_access_resource_owner(
existing_owner, user_api_key_dict
):
raise HTTPException(status_code=403, detail="Forbidden")
_IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner
return response
try:
existing = await prisma_client.db.litellm_managedobjecttable.find_unique(
where={"model_object_id": model_object_id}
@ -136,7 +154,12 @@ async def _get_container_owner(
) -> Optional[str]:
prisma_client = await _get_prisma_client()
if prisma_client is None:
return None
return _IN_MEMORY_CONTAINER_OWNERS.get(
_container_model_object_id(
original_container_id,
custom_llm_provider,
)
)
row = await prisma_client.db.litellm_managedobjecttable.find_first(
where={
@ -164,6 +187,13 @@ async def assert_user_can_access_container(
return original_container_id, resolved_provider
owner = await _get_container_owner(original_container_id, resolved_provider)
if owner is None and _allow_untracked_container_access():
verbose_proxy_logger.warning(
"Allowing untracked container access because %s is enabled",
ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV,
)
return original_container_id, resolved_provider
if not user_can_access_resource_owner(owner, user_api_key_dict):
raise HTTPException(status_code=403, detail="Forbidden")
@ -203,7 +233,12 @@ async def _get_allowed_container_ids(
) -> Set[str]:
prisma_client = await _get_prisma_client()
if prisma_client is None:
return set()
owner_scopes = get_resource_owner_scopes(user_api_key_dict)
return {
model_object_id
for model_object_id, owner in _IN_MEMORY_CONTAINER_OWNERS.items()
if owner in owner_scopes
}
owner_scopes = get_resource_owner_scopes(user_api_key_dict)
if not owner_scopes:

234
tests/test_litellm/containers/test_container_proxy_ownership.py
View File

@ -1,3 +1,4 @@
import sys
from types import SimpleNamespace
from unittest.mock import AsyncMock
@ -6,9 +7,18 @@ from fastapi import HTTPException
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.container_endpoints import ownership
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.containers.main import ContainerListResponse, ContainerObject
@pytest.fixture(autouse=True)
def clear_in_memory_container_owners(monkeypatch):
ownership._IN_MEMORY_CONTAINER_OWNERS.clear()
monkeypatch.delenv(ownership.ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, raising=False)
yield
ownership._IN_MEMORY_CONTAINER_OWNERS.clear()
def _container(container_id: str) -> ContainerObject:
return ContainerObject(
id=container_id,
@ -72,6 +82,31 @@ async def test_should_record_team_owner_for_keys_without_user_id(monkeypatch):
assert data["updated_by"] == "team:team-1"
@pytest.mark.asyncio
async def test_should_track_container_owner_in_memory_without_prisma(monkeypatch):
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=None),
)
auth = UserAPIKeyAuth(user_id="user-1")
await ownership.record_container_owner(
response=_container("cntr_provider"),
user_api_key_dict=auth,
custom_llm_provider="openai",
)
original_id, provider = await ownership.assert_user_can_access_container(
container_id="cntr_provider",
user_api_key_dict=auth,
custom_llm_provider="openai",
)
assert original_id == "cntr_provider"
assert provider == "openai"
@pytest.mark.asyncio
async def test_should_deny_container_access_for_different_owner(monkeypatch):
table = AsyncMock()
@ -96,6 +131,45 @@ async def test_should_deny_container_access_for_different_owner(monkeypatch):
assert exc.value.status_code == 403
@pytest.mark.asyncio
async def test_should_deny_untracked_container_access_by_default(monkeypatch):
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=None),
)
auth = UserAPIKeyAuth(user_id="user-1")
with pytest.raises(HTTPException) as exc:
await ownership.assert_user_can_access_container(
container_id="cntr_untracked",
user_api_key_dict=auth,
custom_llm_provider="openai",
)
assert exc.value.status_code == 403
@pytest.mark.asyncio
async def test_should_allow_untracked_container_access_when_enabled(monkeypatch):
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=None),
)
monkeypatch.setenv(ownership.ALLOW_UNTRACKED_CONTAINER_ACCESS_ENV, "true")
auth = UserAPIKeyAuth(user_id="user-1")
original_id, provider = await ownership.assert_user_can_access_container(
container_id="cntr_untracked",
user_api_key_dict=auth,
custom_llm_provider="openai",
)
assert original_id == "cntr_untracked"
assert provider == "openai"
@pytest.mark.asyncio
async def test_should_not_reassign_existing_container_to_different_owner(monkeypatch):
table = AsyncMock()
@ -157,3 +231,163 @@ async def test_should_filter_container_list_to_owned_records(monkeypatch):
where = table.find_many.await_args.kwargs["where"]
assert where["file_purpose"] == ownership.CONTAINER_OBJECT_PURPOSE
assert where["created_by"]["in"] == ["user-1", "user:user-1"]
@pytest.mark.asyncio
async def test_should_filter_container_list_with_in_memory_ownership(monkeypatch):
monkeypatch.setattr(
ownership,
"_get_prisma_client",
AsyncMock(return_value=None),
)
auth = UserAPIKeyAuth(user_id="user-1")
await ownership.record_container_owner(
response=_container("cntr_owned"),
user_api_key_dict=auth,
custom_llm_provider="openai",
)
response = ContainerListResponse(
object="list",
data=[_container("cntr_owned"), _container("cntr_other")],
has_more=False,
)
filtered = await ownership.filter_container_list_response(
response=response,
user_api_key_dict=auth,
custom_llm_provider="openai",
)
assert [item.id for item in filtered.data] == ["cntr_owned"]
@pytest.mark.asyncio
async def test_should_preserve_managed_container_id_for_proxy_forwarding(monkeypatch):
from litellm.proxy.container_endpoints import handler_factory
proxy_server_stub = SimpleNamespace(
general_settings={},
llm_router=None,
proxy_config=None,
proxy_logging_obj=None,
select_data_generator=None,
user_api_base=None,
user_max_tokens=None,
user_model=None,
user_request_timeout=None,
user_temperature=None,
version="test",
)
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub)
captured = {}
class FakeProcessor:
def __init__(self, data):
captured["data"] = data
async def base_process_llm_request(self, **kwargs):
return captured["data"]
async def _handle_llm_api_exception(self, **kwargs):
raise kwargs["e"]
monkeypatch.setattr(
handler_factory,
"ProxyBaseLLMRequestProcessing",
FakeProcessor,
)
monkeypatch.setattr(
handler_factory,
"assert_user_can_access_container",
AsyncMock(return_value=("cntr_provider", "azure")),
)
encoded_id = ResponsesAPIRequestUtils._build_container_id(
custom_llm_provider="azure",
model_id="router-gpt",
container_id="cntr_provider",
)
result = await handler_factory._process_request(
request=SimpleNamespace(query_params={}, headers={}),
fastapi_response=SimpleNamespace(),
user_api_key_dict=UserAPIKeyAuth(user_id="user-1"),
route_type="alist_container_files",
path_params={"container_id": encoded_id},
)
assert result["container_id"] == encoded_id
assert result["custom_llm_provider"] == "azure"
@pytest.mark.asyncio
async def test_should_preserve_managed_container_id_for_multipart_upload(monkeypatch):
from litellm.proxy.common_utils import http_parsing_utils
from litellm.proxy.container_endpoints import handler_factory
proxy_server_stub = SimpleNamespace(
general_settings={},
llm_router=None,
proxy_config=None,
proxy_logging_obj=None,
select_data_generator=None,
user_api_base=None,
user_max_tokens=None,
user_model=None,
user_request_timeout=None,
user_temperature=None,
version="test",
)
monkeypatch.setitem(sys.modules, "litellm.proxy.proxy_server", proxy_server_stub)
captured = {}
class FakeProcessor:
def __init__(self, data):
captured["data"] = data
async def base_process_llm_request(self, **kwargs):
return captured["data"]
async def _handle_llm_api_exception(self, **kwargs):
raise kwargs["e"]
monkeypatch.setattr(
handler_factory,
"ProxyBaseLLMRequestProcessing",
FakeProcessor,
)
monkeypatch.setattr(
handler_factory,
"assert_user_can_access_container",
AsyncMock(return_value=("cntr_provider", "azure")),
)
monkeypatch.setattr(
http_parsing_utils,
"get_form_data",
AsyncMock(return_value={}),
)
monkeypatch.setattr(
http_parsing_utils,
"convert_upload_files_to_file_data",
AsyncMock(return_value={"file": ["file-data"]}),
)
encoded_id = ResponsesAPIRequestUtils._build_container_id(
custom_llm_provider="azure",
model_id="router-gpt",
container_id="cntr_provider",
)
result = await handler_factory._process_multipart_upload_request(
request=SimpleNamespace(query_params={}, headers={}),
fastapi_response=SimpleNamespace(),
user_api_key_dict=UserAPIKeyAuth(user_id="user-1"),
route_type="aupload_container_file",
container_id=encoded_id,
)
assert result["container_id"] == encoded_id
assert result["custom_llm_provider"] == "azure"
assert result["file"] == "file-data"