diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c4be2187292..5d2dfdbae58 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -196,6 +196,22 @@ def _collect_response_file_search_vector_store_ids(data: Dict[str, Any]) -> set[ return vector_store_ids +def _collect_response_code_interpreter_container_ids(data: Dict[str, Any]) -> set[str]: + container_ids: set[str] = set() + tools = data.get("tools") + if not isinstance(tools, list): + return container_ids + + for tool in tools: + if not isinstance(tool, dict) or tool.get("type") != "code_interpreter": + continue + container_id = tool.get("container") + if isinstance(container_id, str) and container_id: + container_ids.add(container_id) + + return container_ids + + async def _authorize_response_file_search_vector_stores( data: Dict[str, Any], user_api_key_dict: UserAPIKeyAuth, @@ -215,6 +231,30 @@ async def _authorize_response_file_search_vector_stores( ) +async def _authorize_response_code_interpreter_containers( + data: Dict[str, Any], + user_api_key_dict: UserAPIKeyAuth, +) -> None: + container_ids = _collect_response_code_interpreter_container_ids(data) + if not container_ids: + return + + from litellm.proxy.container_endpoints.ownership import ( + assert_user_can_access_container, + ) + + custom_llm_provider = data.get("custom_llm_provider") + if not isinstance(custom_llm_provider, str) or not custom_llm_provider: + custom_llm_provider = "openai" + + for container_id in sorted(container_ids): + await assert_user_can_access_container( + container_id=container_id, + user_api_key_dict=user_api_key_dict, + custom_llm_provider=custom_llm_provider, + ) + + async def _parse_event_data_for_error(event_line: Union[str, bytes]) -> Optional[int]: """Parses an event line and returns an error code if present, else None.""" event_line = event_line.decode("utf-8") if isinstance(event_line, bytes) else event_line @@ -1062,6 +1102,10 @@ class ProxyBaseLLMRequestProcessing: data=self.data, user_api_key_dict=user_api_key_dict, ) + await _authorize_response_code_interpreter_containers( + data=self.data, + user_api_key_dict=user_api_key_dict, + ) # Calculate request queue time after add_litellm_data_to_request # which sets arrival_time in proxy_server_request diff --git a/tests/test_litellm/containers/test_container_proxy_ownership.py b/tests/test_litellm/containers/test_container_proxy_ownership.py index 176405bb9ca..17168660f1b 100644 --- a/tests/test_litellm/containers/test_container_proxy_ownership.py +++ b/tests/test_litellm/containers/test_container_proxy_ownership.py @@ -214,6 +214,87 @@ async def test_should_deny_untracked_container_access_by_default(monkeypatch): assert exc.value.status_code == 403 +@pytest.mark.asyncio +async def test_response_code_interpreter_container_requires_owner_access(monkeypatch): + from litellm.proxy.common_request_processing import ( + _authorize_response_code_interpreter_containers, + ) + + table = AsyncMock() + table.find_first.return_value = SimpleNamespace(created_by="user-2") + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + encoded_container_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="azure-deployment-id", + container_id="cntr_native_123", + ) + + with pytest.raises(HTTPException) as exc: + await _authorize_response_code_interpreter_containers( + data={ + "tools": [ + { + "type": "code_interpreter", + "container": encoded_container_id, + } + ] + }, + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + ) + + assert exc.value.status_code == 403 + table.find_first.assert_awaited_once_with( + where={ + "model_object_id": "container:azure:cntr_native_123", + "file_purpose": ownership.CONTAINER_OBJECT_PURPOSE, + } + ) + + +@pytest.mark.asyncio +async def test_response_code_interpreter_container_allows_owner(monkeypatch): + from litellm.proxy.common_request_processing import ( + _authorize_response_code_interpreter_containers, + ) + + table = AsyncMock() + table.find_first.return_value = SimpleNamespace(created_by="user-1") + prisma_client = SimpleNamespace( + db=SimpleNamespace(litellm_managedobjecttable=table) + ) + monkeypatch.setattr( + ownership, + "_get_prisma_client", + AsyncMock(return_value=prisma_client), + ) + encoded_container_id = ResponsesAPIRequestUtils._build_container_id( + custom_llm_provider="azure", + model_id="azure-deployment-id", + container_id="cntr_native_123", + ) + + await _authorize_response_code_interpreter_containers( + data={ + "tools": [ + { + "type": "code_interpreter", + "container": encoded_container_id, + } + ] + }, + user_api_key_dict=UserAPIKeyAuth(user_id="user-1"), + ) + + table.find_first.assert_awaited_once() + + @pytest.mark.asyncio async def test_should_not_reassign_existing_container_to_different_owner(monkeypatch): table = AsyncMock()