Tool for executing Python code in an isolated Docker sandbox. Features: - Complete Docker container isolation - Resource limits (memory, CPU) - File operations - Data directory mounting for accessing input/output files
| 81 | |
| 82 | |
| 83 | class CodeExecutionTool(ToolBase): |
| 84 | """ |
| 85 | Tool for executing Python code in an isolated Docker sandbox. |
| 86 | |
| 87 | Features: |
| 88 | - Complete Docker container isolation |
| 89 | - Resource limits (memory, CPU) |
| 90 | - File operations |
| 91 | - Data directory mounting for accessing input/output files |
| 92 | """ |
| 93 | |
| 94 | def __init__(self, config): |
| 95 | logger.info('Installing ms-enclave package...') |
| 96 | try: |
| 97 | install_package( |
| 98 | package_name='ms-enclave', import_name='ms_enclave') |
| 99 | except Exception as e: |
| 100 | raise e |
| 101 | |
| 102 | super().__init__(config) |
| 103 | self.manager: Optional['SandboxManager'] = None |
| 104 | self.sandbox_id: Optional[str] = None |
| 105 | self._initialized = False |
| 106 | self._original_port: Optional[int] = None |
| 107 | self.sandbox_type: Optional['SandboxType'] = None |
| 108 | |
| 109 | # Extract sandbox configuration |
| 110 | self.sandbox_config = self._build_sandbox_config(config) |
| 111 | |
| 112 | self.exclude_func(getattr(config.tools, 'code_executor', None)) |
| 113 | |
| 114 | logger.info('CodeExecutionTool initialized (ms-enclave based)') |
| 115 | |
| 116 | def _build_sandbox_config( |
| 117 | self, |
| 118 | config) -> Union['DockerNotebookConfig', 'DockerSandboxConfig']: |
| 119 | """Build sandbox configuration from agent config""" |
| 120 | from ms_enclave.sandbox.model import DockerNotebookConfig, DockerSandboxConfig, SandboxType |
| 121 | |
| 122 | # Get sandbox-specific config or use defaults |
| 123 | if isinstance(config, DictConfig) and hasattr( |
| 124 | config, 'tools') and hasattr(config.tools, 'code_executor'): |
| 125 | sandbox_cfg = getattr(config.tools.code_executor, 'sandbox', {}) |
| 126 | else: |
| 127 | sandbox_cfg = getattr(config, 'sandbox', {}) or getattr( |
| 128 | config, 'tools', {}).get('sandbox', {}) |
| 129 | |
| 130 | # Get output directory for data mounting |
| 131 | output_dir = Path(getattr(config, 'output_dir', DEFAULT_OUTPUT_DIR)) |
| 132 | output_dir.mkdir(parents=True, exist_ok=True) |
| 133 | |
| 134 | # Build volumes configuration |
| 135 | volumes = {str(output_dir.absolute()): {'bind': '/data', 'mode': 'rw'}} |
| 136 | |
| 137 | # Build environment variables |
| 138 | env_vars = { |
| 139 | 'DATA_DIR': '/data', |
| 140 | 'PYTHONUNBUFFERED': '1', |