From 6aac4552f930fd9f7fc94f3bd22bc0a8de9f9fc8 Mon Sep 17 00:00:00 2001 From: user <70670632+stuxf@users.noreply.github.com> Date: Thu, 30 Apr 2026 19:13:20 -0700 Subject: [PATCH] fix(proxy): forward decoded container ids --- .../proxy/container_endpoints/endpoints.py | 27 ++++++---- .../container_endpoints/handler_factory.py | 46 +++++++++++------ .../proxy/container_endpoints/ownership.py | 35 +++++++++++-- .../test_container_proxy_ownership.py | 51 +++++++++++++++---- 4 files changed, 121 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/container_endpoints/endpoints.py b/litellm/proxy/container_endpoints/endpoints.py index 93587262f5..72a7c6c874 100644 --- a/litellm/proxy/container_endpoints/endpoints.py +++ b/litellm/proxy/container_endpoints/endpoints.py @@ -17,6 +17,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.container_endpoints.ownership import ( assert_user_can_access_container, filter_container_list_response, + get_container_forwarding_params, record_container_owner, ) @@ -295,15 +296,18 @@ async def retrieve_container( ) # Add custom_llm_provider to data - container_access = await assert_user_can_access_container( + original_container_id, custom_llm_provider = 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 + data.update( + get_container_forwarding_params( + container_id, + original_container_id, + custom_llm_provider, + ) + ) # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) @@ -397,15 +401,18 @@ async def delete_container( ) # Add custom_llm_provider to data - container_access = await assert_user_can_access_container( + original_container_id, custom_llm_provider = 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 + data.update( + get_container_forwarding_params( + container_id, + original_container_id, + custom_llm_provider, + ) + ) # Process request using ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=data) diff --git a/litellm/proxy/container_endpoints/handler_factory.py b/litellm/proxy/container_endpoints/handler_factory.py index bb6f8a8db6..ee79ebd96e 100644 --- a/litellm/proxy/container_endpoints/handler_factory.py +++ b/litellm/proxy/container_endpoints/handler_factory.py @@ -21,6 +21,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ) from litellm.proxy.container_endpoints.ownership import ( assert_user_can_access_container, + get_container_forwarding_params, ) @@ -183,6 +184,13 @@ async def _process_binary_request( user_api_key_dict=user_api_key_dict, custom_llm_provider=custom_llm_provider, ) + forwarding_params = get_container_forwarding_params( + container_id, + original_container_id, + custom_llm_provider, + ) + if "model_id" in forwarding_params: + litellm_params["model_id"] = forwarding_params["model_id"] # Get the provider config container_provider_config = _get_container_provider_config(custom_llm_provider) @@ -283,18 +291,19 @@ async def _process_multipart_upload_request( or "openai" ) - container_access = await assert_user_can_access_container( + original_container_id, custom_llm_provider = 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 + data.update( + get_container_forwarding_params( + container_id, + original_container_id, + custom_llm_provider, + ) + ) processor = ProxyBaseLLMRequestProcessing(data=data) try: @@ -361,15 +370,22 @@ async def _process_request( # Validate container_id ownership if present in path_params. if "container_id" in path_params: - 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, + original_container_id, custom_llm_provider = ( + 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 + data.update( + get_container_forwarding_params( + path_params["container_id"], + original_container_id, + custom_llm_provider, + ) + ) + else: + data["custom_llm_provider"] = custom_llm_provider processor = ProxyBaseLLMRequestProcessing(data=data) try: diff --git a/litellm/proxy/container_endpoints/ownership.py b/litellm/proxy/container_endpoints/ownership.py index 2260d67b62..b769b2c095 100644 --- a/litellm/proxy/container_endpoints/ownership.py +++ b/litellm/proxy/container_endpoints/ownership.py @@ -45,6 +45,22 @@ def decode_container_id_for_ownership( return original_container_id, custom_llm_provider +def get_container_forwarding_params( + container_id: str, + original_container_id: str, + custom_llm_provider: str, +) -> Dict[str, str]: + params = { + "container_id": original_container_id, + "custom_llm_provider": custom_llm_provider, + } + decoded = ResponsesAPIRequestUtils._decode_container_id(container_id) + model_id = decoded.get("model_id") + if isinstance(model_id, str) and model_id: + params["model_id"] = model_id + return params + + def _get_response_id(response: Any) -> Optional[str]: if response is None: return None @@ -139,11 +155,12 @@ async def record_container_owner( raise except Exception as e: verbose_proxy_logger.warning( - "Failed to record container ownership for container_id=%s: %s", + "Failed to persist container ownership for container_id=%s; " + "falling back to in-process tracking: %s", model_object_id, e, ) - raise HTTPException(status_code=500, detail="Unable to track container") + _IN_MEMORY_CONTAINER_OWNERS[model_object_id] = owner return response @@ -210,7 +227,9 @@ def _get_container_list_data(response: Any) -> Optional[List[Any]]: return data if isinstance(data, list) else None -def _set_container_list_data(response: Any, data: List[Any]) -> Any: +def _set_container_list_data( + response: Any, data: List[Any], removed_filtered_items: bool = False +) -> Any: if isinstance(response, dict): response["data"] = data if data: @@ -220,6 +239,8 @@ def _set_container_list_data(response: Any, data: List[Any]) -> Any: response["first_id"] = None response["last_id"] = None response["has_more"] = False + if removed_filtered_items: + response["has_more"] = False return response response.data = data @@ -227,6 +248,8 @@ def _set_container_list_data(response: Any, data: List[Any]) -> Any: 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 return response @@ -291,4 +314,8 @@ async def filter_container_list_response( ): filtered.append(item) - return _set_container_list_data(response, filtered) + return _set_container_list_data( + response, + filtered, + removed_filtered_items=len(filtered) != len(data), + ) diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 40397a1886..5ec28e8115 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -82,6 +82,34 @@ 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_fallback_to_memory_when_persistent_owner_record_fails( + monkeypatch, +): + table = AsyncMock() + table.find_unique.side_effect = Exception("db unavailable") + 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") + + await ownership.record_container_owner( + response=_container("cntr_provider"), + user_api_key_dict=auth, + custom_llm_provider="openai", + ) + + assert ( + ownership._IN_MEMORY_CONTAINER_OWNERS["container:openai:cntr_provider"] + == "user-1" + ) + + @pytest.mark.asyncio async def test_should_track_container_owner_in_memory_without_prisma(monkeypatch): monkeypatch.setattr( @@ -216,7 +244,7 @@ async def test_should_filter_container_list_to_owned_records(monkeypatch): response = ContainerListResponse( object="list", data=[_container("cntr_owned"), _container("cntr_other")], - has_more=False, + has_more=True, ) filtered = await ownership.filter_container_list_response( @@ -228,6 +256,7 @@ async def test_should_filter_container_list_to_owned_records(monkeypatch): 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 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"] @@ -334,7 +363,7 @@ async def test_should_filter_container_list_with_in_memory_ownership(monkeypatch @pytest.mark.asyncio -async def test_should_preserve_managed_container_id_for_proxy_forwarding(monkeypatch): +async def test_should_forward_decoded_container_id_for_proxy_forwarding(monkeypatch): from litellm.proxy.container_endpoints import handler_factory proxy_server_stub = SimpleNamespace( @@ -388,12 +417,13 @@ async def test_should_preserve_managed_container_id_for_proxy_forwarding(monkeyp path_params={"container_id": encoded_id}, ) - assert result["container_id"] == encoded_id + 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_preserve_managed_container_id_for_multipart_upload(monkeypatch): +async def test_should_forward_decoded_container_id_for_multipart_upload(monkeypatch): from litellm.proxy.common_utils import http_parsing_utils from litellm.proxy.container_endpoints import handler_factory @@ -458,13 +488,14 @@ async def test_should_preserve_managed_container_id_for_multipart_upload(monkeyp container_id=encoded_id, ) - assert result["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_preserve_managed_container_id_for_proxy_retrieve(monkeypatch): +async def test_should_forward_decoded_container_id_for_proxy_retrieve(monkeypatch): from litellm.proxy.container_endpoints import endpoints proxy_server_stub = SimpleNamespace( @@ -513,12 +544,13 @@ async def test_should_preserve_managed_container_id_for_proxy_retrieve(monkeypat user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), ) - assert result["container_id"] == encoded_id + 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_preserve_managed_container_id_for_proxy_delete(monkeypatch): +async def test_should_forward_decoded_container_id_for_proxy_delete(monkeypatch): from litellm.proxy.container_endpoints import endpoints proxy_server_stub = SimpleNamespace( @@ -567,5 +599,6 @@ async def test_should_preserve_managed_container_id_for_proxy_delete(monkeypatch user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), ) - assert result["container_id"] == encoded_id + assert result["container_id"] == "cntr_provider" assert result["custom_llm_provider"] == "azure" + assert result["model_id"] == "router-gpt"