diff --git a/ms_enclave/sandbox/boxes/base.py b/ms_enclave/sandbox/boxes/base.py index de6db2c..2318b8f 100644 --- a/ms_enclave/sandbox/boxes/base.py +++ b/ms_enclave/sandbox/boxes/base.py @@ -1,7 +1,10 @@ """Base sandbox interface and factory.""" import abc +import asyncio +import os from datetime import datetime +from pathlib import Path from typing import Any, Dict, List, Optional, Type, Union import shortuuid as uuid @@ -20,6 +23,7 @@ VolcengineSandboxConfig, ) from ..tools import Tool, ToolFactory +from ..utils.archive import tar_directory, tar_file logger = get_logger() @@ -153,6 +157,27 @@ async def execute_command( """ raise NotImplementedError('execute_command must be implemented by subclasses') + async def put_archive(self, target_dir: str, data: bytes) -> bool: + """Copy a tar archive into a sandbox directory.""" + raise NotImplementedError(f'{type(self).__name__} does not support put_archive') + + async def put_dir(self, source_dir: str | Path, target_dir: str) -> bool: + """Copy a host directory into a sandbox directory.""" + source = Path(source_dir).expanduser() + if not source.is_dir(): + raise FileNotFoundError(f'put_dir source is not a directory: {source}') + return await self.put_archive(target_dir, await asyncio.to_thread(tar_directory, source)) + + async def put_file(self, file_path: str, content: str, encoding: str = 'utf-8') -> bool: + """Write a text file into the sandbox.""" + if not os.path.isabs(file_path): + raise ValueError(f'file_path must be an absolute path: {file_path}') + target_dir = os.path.dirname(file_path) or '/' + file_name = os.path.basename(file_path) + if not file_name: + raise ValueError(f'file_path must include a file name: {file_path}') + return await self.put_archive(target_dir, tar_file(file_name, content.encode(encoding))) + @abc.abstractmethod async def get_execution_context(self) -> Any: """Get the execution context for tools (e.g., container, process, etc.).""" diff --git a/ms_enclave/sandbox/boxes/docker_sandbox.py b/ms_enclave/sandbox/boxes/docker_sandbox.py index a271475..62cfc69 100644 --- a/ms_enclave/sandbox/boxes/docker_sandbox.py +++ b/ms_enclave/sandbox/boxes/docker_sandbox.py @@ -150,6 +150,19 @@ async def get_execution_context(self) -> Any: """Return the container for tool execution.""" return self.container + async def put_archive(self, target_dir: str, data: bytes) -> bool: + """Extract a tar archive into the Docker container.""" + if not self.container: + raise RuntimeError('Container is not running') + if not target_dir: + raise ValueError('target_dir must not be empty') + + mkdir = await self._run_blocking(self.container.exec_run, ['mkdir', '-p', '--', target_dir]) + if mkdir.exit_code is None or mkdir.exit_code != 0: + raise RuntimeError(f'Failed to create target directory {target_dir}: {mkdir.output!r}') + + return bool(await self._run_blocking(self.container.put_archive, target_dir, data)) + async def _aiter_exec_output(self, exec_id: str) -> AsyncIterator[Tuple[Optional[bytes], Optional[bytes]]]: """Bridge ``exec_start(stream=True)`` chunks from a worker thread to async. diff --git a/ms_enclave/sandbox/manager/base.py b/ms_enclave/sandbox/manager/base.py index 39fe3a2..c555885 100644 --- a/ms_enclave/sandbox/manager/base.py +++ b/ms_enclave/sandbox/manager/base.py @@ -3,6 +3,7 @@ import asyncio from abc import ABC, abstractmethod from collections import deque +from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, List, Optional, Union from ..model import ( @@ -148,6 +149,18 @@ async def get_sandbox_tools(self, sandbox_id: str) -> Dict[str, Any]: """ pass + async def put_archive(self, sandbox_id: str, target_dir: str, data: bytes) -> bool: + """Copy a tar archive into a sandbox directory. + + Managers that cannot access the sandbox filesystem should leave this + unsupported rather than emulating transfer through shell commands. + """ + raise NotImplementedError(f'{type(self).__name__} does not support put_archive') + + async def put_dir(self, sandbox_id: str, source_dir: str | Path, target_dir: str) -> bool: + """Copy a host directory into a sandbox directory.""" + raise NotImplementedError(f'{type(self).__name__} does not support put_dir') + @abstractmethod async def get_stats(self) -> Dict[str, Any]: """Get manager statistics. diff --git a/ms_enclave/sandbox/manager/http_manager.py b/ms_enclave/sandbox/manager/http_manager.py index 472378f..3c62945 100644 --- a/ms_enclave/sandbox/manager/http_manager.py +++ b/ms_enclave/sandbox/manager/http_manager.py @@ -1,5 +1,7 @@ """HTTP-based sandbox manager for remote sandbox services.""" +import asyncio +from pathlib import Path from typing import Any, Dict, List, Optional, Union import aiohttp @@ -16,6 +18,7 @@ ToolExecutionRequest, ToolResult, ) +from ..utils.archive import tar_directory from .base import SandboxManager, register_manager logger = get_logger() @@ -322,6 +325,38 @@ async def get_sandbox_tools(self, sandbox_id: str) -> Dict[str, Any]: logger.error(f'HTTP client error getting sandbox tools: {e}') raise RuntimeError(f'Failed to get sandbox tools: {e}') + async def put_archive(self, sandbox_id: str, target_dir: str, data: bytes) -> bool: + """Copy a tar archive into a sandbox directory via HTTP API.""" + if not self._session: + raise RuntimeError('Manager not started') + try: + async with self._session.post( + f'{self.base_url}/sandbox/{sandbox_id}/archive', + params={'target_dir': target_dir}, + data=data, + headers={'Content-Type': 'application/x-tar'}, + ) as response: + if response.status == 200: + payload = await response.json() + return bool(payload.get('ok', True)) + error_data = await response.json() + detail = error_data.get('detail', 'Unknown error') + if response.status in (400, 404): + raise ValueError(detail) + if response.status == 501: + raise NotImplementedError(detail) + raise RuntimeError(f'HTTP {response.status}: {detail}') + except aiohttp.ClientError as e: + logger.error(f'HTTP client error uploading archive: {e}') + raise RuntimeError(f'Failed to upload archive: {e}') + + async def put_dir(self, sandbox_id: str, source_dir: str | Path, target_dir: str) -> bool: + """Copy a host directory into a sandbox directory via HTTP API.""" + source = Path(source_dir).expanduser() + if not source.is_dir(): + raise FileNotFoundError(f'put_dir source is not a directory: {source}') + return await self.put_archive(sandbox_id, target_dir, await asyncio.to_thread(tar_directory, source)) + async def cleanup_all_sandboxes(self) -> None: """Clean up all sandboxes created by this manager via HTTP API.""" if not self._session: diff --git a/ms_enclave/sandbox/manager/local_manager.py b/ms_enclave/sandbox/manager/local_manager.py index b7057da..7138570 100644 --- a/ms_enclave/sandbox/manager/local_manager.py +++ b/ms_enclave/sandbox/manager/local_manager.py @@ -3,6 +3,7 @@ import asyncio from collections import Counter from datetime import datetime, timedelta +from pathlib import Path from typing import Any, Dict, List, Optional, Union from ms_enclave.sandbox.model.constants import DEFAULT_POOL_EXECUTION_TIMEOUT @@ -130,6 +131,17 @@ async def get_sandbox(self, sandbox_id: str) -> Optional[Sandbox]: """ return self._sandboxes.get(sandbox_id) + def _get_running_sandbox(self, sandbox_id: str) -> Sandbox: + """Return a running sandbox or raise a consistent error.""" + if not self._running: + raise RuntimeError('Sandbox manager not started') + sandbox = self._sandboxes.get(sandbox_id) + if not sandbox: + raise ValueError(f'Sandbox {sandbox_id} not found') + if sandbox.status not in (SandboxStatus.RUNNING, SandboxStatus.IDLE, SandboxStatus.BUSY): + raise ValueError(f'Sandbox {sandbox_id} is not active (status: {sandbox.status})') + return sandbox + async def get_sandbox_info(self, sandbox_id: str) -> Optional[SandboxInfo]: """Get sandbox information. @@ -219,13 +231,7 @@ async def execute_tool(self, sandbox_id: str, tool_name: str, parameters: Dict[s Raises: ValueError: If sandbox or tool not found """ - sandbox = self._sandboxes.get(sandbox_id) - if not sandbox: - raise ValueError(f'Sandbox {sandbox_id} not found') - - if sandbox.status != SandboxStatus.RUNNING: - raise ValueError(f'Sandbox {sandbox_id} is not running (status: {sandbox.status})') - + sandbox = self._get_running_sandbox(sandbox_id) result = await sandbox.execute_tool(tool_name, parameters) return result @@ -247,6 +253,16 @@ async def get_sandbox_tools(self, sandbox_id: str) -> Dict[str, Any]: return sandbox.get_available_tools() + async def put_archive(self, sandbox_id: str, target_dir: str, data: bytes) -> bool: + """Copy a tar archive into a local sandbox directory.""" + sandbox = self._get_running_sandbox(sandbox_id) + return await sandbox.put_archive(target_dir, data) + + async def put_dir(self, sandbox_id: str, source_dir: str | Path, target_dir: str) -> bool: + """Copy a host directory into a local sandbox directory.""" + sandbox = self._get_running_sandbox(sandbox_id) + return await sandbox.put_dir(source_dir, target_dir) + async def cleanup_all_sandboxes(self) -> None: """Clean up all sandboxes.""" sandbox_ids = list(self._sandboxes.keys()) diff --git a/ms_enclave/sandbox/server/server.py b/ms_enclave/sandbox/server/server.py index 03be4d1..a1c3dc3 100644 --- a/ms_enclave/sandbox/server/server.py +++ b/ms_enclave/sandbox/server/server.py @@ -164,6 +164,23 @@ async def get_sandbox_tools(sandbox_id: str): except ValueError as e: raise HTTPException(status_code=404, detail=str(e)) + @self.app.post('/sandbox/{sandbox_id}/archive') + async def put_archive(sandbox_id: str, target_dir: str, request: Request): + """Copy a tar archive into a sandbox directory.""" + data = await request.body() + if not data: + raise HTTPException(status_code=400, detail='Archive body must not be empty') + try: + ok = await self.manager.put_archive(sandbox_id, target_dir, data) + return {'ok': ok} + except ValueError as e: + status_code = 404 if 'not found' in str(e).lower() else 400 + raise HTTPException(status_code=status_code, detail=str(e)) + except NotImplementedError as e: + raise HTTPException(status_code=501, detail=str(e)) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + # System info @self.app.get('/stats') async def get_stats(): diff --git a/ms_enclave/sandbox/tools/sandbox_tools/file_operation.py b/ms_enclave/sandbox/tools/sandbox_tools/file_operation.py index c3104f0..fcfbf44 100644 --- a/ms_enclave/sandbox/tools/sandbox_tools/file_operation.py +++ b/ms_enclave/sandbox/tools/sandbox_tools/file_operation.py @@ -1,8 +1,6 @@ """File operation tool for reading and writing files.""" -import io import os -import tarfile import uuid from typing import TYPE_CHECKING, Literal, Optional @@ -173,8 +171,8 @@ async def _write_file(self, sandbox_context: 'Sandbox', file_path: str, content: # Generate unique temporary file name to avoid conflicts temp_file = f'/tmp/file_op_{uuid.uuid4().hex}' - # Step 1: Write content to temporary location using tar archive - await self._write_file_to_container(sandbox_context, temp_file, content, encoding) + # Step 1: Write content to temporary location using sandbox file transfer + await sandbox_context.put_file(temp_file, content, encoding) # Step 2: Copy from temp to target location using cp command # This handles permission issues better than direct tar extraction @@ -296,36 +294,3 @@ async def _check_exists(self, sandbox_context: 'Sandbox', file_path: str) -> Too return ToolResult( tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Exists check failed: {str(e)}' ) - - async def _write_file_to_container( - self, sandbox_context: 'Sandbox', file_path: str, content: str, encoding: str - ) -> None: - """Write content to a file in the container using tar archive method. - - This is a low-level method that creates a tar archive containing the file - and extracts it to the container. Used internally by _write_file. - - Args: - sandbox_context: Sandbox instance with container access - file_path: Target file path in container - content: File content as string - encoding: Text encoding for content conversion - - Raises: - Exception: If tar creation or container extraction fails - """ - # Create a tar archive in memory to transfer file content - tar_stream = io.BytesIO() - tar = tarfile.TarFile(fileobj=tar_stream, mode='w') - - # Encode content using specified encoding and create tar entry - file_data = content.encode(encoding) - tarinfo = tarfile.TarInfo(name=os.path.basename(file_path)) - tarinfo.size = len(file_data) - tar.addfile(tarinfo, io.BytesIO(file_data)) - tar.close() - - # Extract tar archive to container filesystem - # Note: This writes to the directory containing the target file - tar_stream.seek(0) - sandbox_context.container.put_archive(os.path.dirname(file_path) or '/', tar_stream.getvalue()) diff --git a/ms_enclave/sandbox/tools/sandbox_tools/multi_code_executor.py b/ms_enclave/sandbox/tools/sandbox_tools/multi_code_executor.py index 301cf27..c8d3e53 100644 --- a/ms_enclave/sandbox/tools/sandbox_tools/multi_code_executor.py +++ b/ms_enclave/sandbox/tools/sandbox_tools/multi_code_executor.py @@ -137,7 +137,7 @@ async def execute( if files: for fname, content in files.items(): safe_name = os.path.basename(fname) - await self._write_file_to_container(sandbox_context, os.path.join(workdir, safe_name), content) + await sandbox_context.put_file(os.path.join(workdir, safe_name), content) # Pre-build setup (e.g., initialize a C# project) prebuild_error = await self._prebuild_setup(sandbox_context, lang, workdir, compile_timeout) @@ -154,7 +154,7 @@ async def execute( error=f'Unsupported language: {language}' ) main_path = os.path.join(workdir, main_filename) - await self._write_file_to_container(sandbox_context, main_path, code) + await sandbox_context.put_file(main_path, code) # Detect Scala classname early to provide a clear error scala_classname: Optional[str] = None @@ -218,26 +218,6 @@ async def _ensure_dir(self, sandbox_context: 'DockerSandbox', dir_path: str) -> if res.exit_code != 0: raise RuntimeError(f'Failed to create workdir: {dir_path} ({res.stderr})') - async def _write_file_to_container(self, sandbox_context: 'DockerSandbox', file_path: str, content: str) -> None: - """Write content to a file in the container using a tar archive; creates parent dir if missing.""" - import io - import tarfile - - dir_name = os.path.dirname(file_path) - base_name = os.path.basename(file_path) - - # Ensure parent directory exists - await self._ensure_dir(sandbox_context, dir_name) - - with io.BytesIO() as tar_stream: - with tarfile.TarFile(fileobj=tar_stream, mode='w') as tar: - data = content.encode('utf-8') - tarinfo = tarfile.TarInfo(name=base_name) - tarinfo.size = len(data) - tar.addfile(tarinfo, io.BytesIO(data)) - tar_stream.seek(0) - sandbox_context.container.put_archive(dir_name, tar_stream.getvalue()) - async def _python_env_prefix(self, sandbox_context: 'DockerSandbox') -> str: """Best-effort Python PATH adjustment similar to get_python_rt_env('sandbox-runtime').""" # Prefer miniconda at /root/miniconda3 when present; avoid noisy activation errors. @@ -279,7 +259,7 @@ async def _get_cpp_rt_flags(self, sandbox_context: 'DockerSandbox', workdir: str optional_flags = ['-lcrypto', '-lssl', '-lpthread'] # Write a tiny C++ file test_file = os.path.join(workdir, 'mce_rt_probe.cpp') - await self._write_file_to_container(sandbox_context, test_file, 'int main(){return 0;}') + await sandbox_context.put_file(test_file, 'int main(){return 0;}') detected: List[str] = [] for flag in optional_flags: diff --git a/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py b/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py index 3cf0c41..cdb4c8f 100644 --- a/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py +++ b/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py @@ -1,6 +1,5 @@ """Python code execution tool.""" -import os import uuid from typing import TYPE_CHECKING, Optional @@ -59,7 +58,7 @@ async def execute(self, sandbox_context: 'DockerSandbox', code: str, timeout: Op try: # Write script to container to avoid long code errors - await self._write_file_to_container(sandbox_context, script_path, code) + await sandbox_context.put_file(script_path, code) # Execute using python command = f'python {script_path}' @@ -80,20 +79,3 @@ async def execute(self, sandbox_context: 'DockerSandbox', code: str, timeout: Op return ToolResult( tool_name=self.name, status=ExecutionStatus.ERROR, output='', error=f'Execution failed: {str(e)}' ) - - async def _write_file_to_container(self, sandbox_context: 'DockerSandbox', file_path: str, content: str) -> None: - """Write content to a file in the container.""" - import io - import tarfile - - # Create a tar archive in memory using context managers - with io.BytesIO() as tar_stream: - with tarfile.TarFile(fileobj=tar_stream, mode='w') as tar: - file_data = content.encode('utf-8') - tarinfo = tarfile.TarInfo(name=os.path.basename(file_path)) - tarinfo.size = len(file_data) - tar.addfile(tarinfo, io.BytesIO(file_data)) - - # Reset stream position and put archive into container - tar_stream.seek(0) - sandbox_context.container.put_archive(os.path.dirname(file_path), tar_stream.getvalue()) diff --git a/ms_enclave/sandbox/utils/__init__.py b/ms_enclave/sandbox/utils/__init__.py new file mode 100644 index 0000000..a4af994 --- /dev/null +++ b/ms_enclave/sandbox/utils/__init__.py @@ -0,0 +1 @@ +"""Shared sandbox utility helpers.""" diff --git a/ms_enclave/sandbox/utils/archive.py b/ms_enclave/sandbox/utils/archive.py new file mode 100644 index 0000000..f837ead --- /dev/null +++ b/ms_enclave/sandbox/utils/archive.py @@ -0,0 +1,31 @@ +"""Archive helpers for sandbox file transfer.""" + +import io +import tarfile +import time +from pathlib import Path + + +def tar_directory(source: Path) -> bytes: + """Create a tar archive containing the contents of ``source``.""" + stream = io.BytesIO() + with tarfile.open(fileobj=stream, mode='w') as tar: + for path in sorted(source.rglob('*')): + if path.is_symlink(): + continue + arcname = path.relative_to(source).as_posix() + tar.add(path, arcname=arcname, recursive=False) + stream.seek(0) + return stream.getvalue() + + +def tar_file(name: str, data: bytes) -> bytes: + """Create a tar archive containing a single file.""" + stream = io.BytesIO() + with tarfile.open(fileobj=stream, mode='w') as tar: + tarinfo = tarfile.TarInfo(name=name) + tarinfo.size = len(data) + tarinfo.mtime = int(time.time()) + tar.addfile(tarinfo, io.BytesIO(data)) + stream.seek(0) + return stream.getvalue() diff --git a/tests/test_manager.py b/tests/test_manager.py index 33f5f21..613300e 100644 --- a/tests/test_manager.py +++ b/tests/test_manager.py @@ -1,7 +1,11 @@ """Unit tests for the sandbox manager functionality.""" import asyncio +import socket +import tempfile +import threading import unittest +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch from ms_enclave.sandbox.manager import HttpSandboxManager, LocalSandboxManager, SandboxManagerFactory @@ -12,6 +16,36 @@ SandboxStatus, SandboxType, ) +from ms_enclave.sandbox.server.server import create_server + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(('127.0.0.1', 0)) + return int(sock.getsockname()[1]) + + +async def _wait_for_port(host: str, port: int) -> None: + for _ in range(100): + try: + reader, writer = await asyncio.open_connection(host, port) + writer.close() + await writer.wait_closed() + return + except OSError: + await asyncio.sleep(0.1) + raise TimeoutError(f'Timed out waiting for {host}:{port}') + + +def _skip_if_docker_unavailable() -> None: + try: + import docker + + client = docker.from_env() + client.ping() + client.close() + except Exception as exc: + raise unittest.SkipTest(f'Docker daemon is not available: {exc}') class TestSandboxManagerFactory(unittest.TestCase): @@ -185,3 +219,94 @@ async def test_execute_tool_nonexistent_sandbox(self): 'python_executor', {'code': 'print("test")'} ) + + async def test_execute_tool_accepts_active_pool_statuses(self): + """Test direct operations can target pooled sandboxes.""" + for status in (SandboxStatus.IDLE, SandboxStatus.BUSY): + sandbox = MagicMock() + sandbox.status = status + sandbox.execute_tool = AsyncMock(return_value=f'{status.value}-result') + sandbox.stop = AsyncMock() + self.manager._sandboxes[status.value] = sandbox + + result = await self.manager.execute_tool(status.value, 'python_executor', {}) + + self.assertEqual(result, f'{status.value}-result') + sandbox.execute_tool.assert_awaited_once_with('python_executor', {}) + + +class TestSandboxFileTransfer(unittest.IsolatedAsyncioTestCase): + """Test manager file transfer helpers through real local and HTTP APIs.""" + + async def test_local_put_dir_copies_directory_into_docker_sandbox(self): + _skip_if_docker_unavailable() + manager = LocalSandboxManager() + await manager.start() + sandbox_id = None + try: + with tempfile.TemporaryDirectory() as tmp: + source = Path(tmp) / 'skills' + skill = source / 'demo' + skill.mkdir(parents=True) + (skill / 'SKILL.md').write_text('skill from local', encoding='utf-8') + sandbox_id = await manager.create_sandbox( + SandboxType.DOCKER, + DockerSandboxConfig(image='python:3.11-slim', tools_config={'shell_executor': {}}), + ) + ok = await manager.put_dir(sandbox_id, source, '/tmp/eval-skills') + result = await manager.execute_tool( + sandbox_id, + 'shell_executor', + {'command': ['cat', '/tmp/eval-skills/demo/SKILL.md'], 'timeout': 30}, + ) + finally: + if sandbox_id is not None: + await manager.delete_sandbox(sandbox_id) + await manager.stop() + + self.assertTrue(ok) + self.assertTrue(result.success) + self.assertIn('skill from local', result.output) + + async def test_http_put_dir_copies_directory_into_docker_sandbox(self): + _skip_if_docker_unavailable() + import uvicorn + + port = _free_port() + server = create_server(SandboxManagerConfig(cleanup_interval=60)) + uvicorn_server = uvicorn.Server( + uvicorn.Config(server.app, host='127.0.0.1', port=port, log_level='warning', ws='none') + ) + thread = threading.Thread(target=uvicorn_server.run, daemon=True) + thread.start() + await _wait_for_port('127.0.0.1', port) + + manager = HttpSandboxManager(SandboxManagerConfig(base_url=f'http://127.0.0.1:{port}', timeout=60.0)) + await manager.start() + sandbox_id = None + try: + with tempfile.TemporaryDirectory() as tmp: + source = Path(tmp) / 'skills' + skill = source / 'demo' + skill.mkdir(parents=True) + (skill / 'SKILL.md').write_text('skill from http', encoding='utf-8') + sandbox_id = await manager.create_sandbox( + SandboxType.DOCKER, + DockerSandboxConfig(image='python:3.11-slim', tools_config={'shell_executor': {}}), + ) + ok = await manager.put_dir(sandbox_id, source, '/tmp/eval-skills') + result = await manager.execute_tool( + sandbox_id, + 'shell_executor', + {'command': ['cat', '/tmp/eval-skills/demo/SKILL.md'], 'timeout': 30}, + ) + finally: + if sandbox_id is not None: + await manager.delete_sandbox(sandbox_id) + await manager.stop() + uvicorn_server.should_exit = True + thread.join(timeout=10) + + self.assertTrue(ok) + self.assertTrue(result.success) + self.assertIn('skill from http', result.output)