From abc0b498228eb5a7f8c6da907c322417879de0d8 Mon Sep 17 00:00:00 2001 From: Yunnglin Date: Thu, 9 Jul 2026 11:53:48 +0800 Subject: [PATCH] Fix Docker sandbox timeout handling --- ms_enclave/sandbox/boxes/docker_sandbox.py | 100 ++++++++++++-- .../sandbox_tools/multi_code_executor.py | 4 +- .../tools/sandbox_tools/notebook_executor.py | 53 +++++++- .../tools/sandbox_tools/python_executor.py | 4 +- .../tools/sandbox_tools/shell_executor.py | 7 + tests/test_tool.py | 122 ++++++++++++++++++ 6 files changed, 272 insertions(+), 18 deletions(-) diff --git a/ms_enclave/sandbox/boxes/docker_sandbox.py b/ms_enclave/sandbox/boxes/docker_sandbox.py index 62cfc69..ad6ea38 100644 --- a/ms_enclave/sandbox/boxes/docker_sandbox.py +++ b/ms_enclave/sandbox/boxes/docker_sandbox.py @@ -18,6 +18,7 @@ logger = get_logger() _QUEUE_SENTINEL = object() +_TIMEOUT_EXIT_CODES = {124, 137} @register_sandbox(SandboxType.DOCKER) @@ -194,7 +195,31 @@ def _producer(): yield item async def _kill_exec_safe(self, exec_id: str) -> None: - """Best-effort exec_kill so the producer iterator unblocks.""" + """Best-effort termination of the exec process and its children.""" + pid = None + try: + inspect = await self._run_blocking(self.client.api.exec_inspect, exec_id) + pid = inspect.get('Pid') + except Exception as e: + logger.debug(f'exec_inspect failed before exec_kill: {e}') + + if pid and self.container: + kill_tree_cmd = ( + 'kill_tree() { ' + 'for child in $(cat /proc/$1/task/$1/children 2>/dev/null); do ' + 'kill_tree "$child" "$2"; ' + 'done; ' + 'kill "$2" "$1" 2>/dev/null || true; ' + '}; ' + f'kill_tree {pid} -TERM; ' + 'sleep 0.2; ' + f'kill_tree {pid} -KILL' + ) + try: + await self._run_blocking(self.container.exec_run, ['sh', '-c', kill_tree_cmd]) + except Exception as e: + logger.debug(f'exec process-tree kill failed: {e}') + try: await self._run_blocking(self.client.api.exec_kill, exec_id, 'SIGKILL') except Exception as e: @@ -221,6 +246,36 @@ def _run_buffered(self, command: Union[str, List[str]]) -> Tuple[int, str, str]: stderr = err_bytes.decode('utf-8', errors='replace') if err_bytes else '' return res.exit_code, stdout, stderr + @staticmethod + def _wrap_command_timeout(command: Union[str, List[str]], timeout: Optional[int]) -> Union[str, List[str]]: + """Run the command under GNU timeout inside the container.""" + if timeout is None: + return command + timeout_s = float(timeout) + if timeout_s <= 0: + return command + + prefix = ['timeout', '-s', 'TERM', '-k', '5s', f'{timeout_s}s'] + if isinstance(command, list): + return [*prefix, *command] + return [*prefix, 'sh', '-c', command] + + @staticmethod + def _outer_timeout(timeout: Optional[int]) -> Optional[float]: + if timeout is None: + return None + timeout_s = float(timeout) + return timeout_s + 10.0 if timeout_s > 0 else None + + @staticmethod + def _is_timeout_exit(exit_code: int, timeout: Optional[int], started_at: float) -> bool: + if exit_code not in _TIMEOUT_EXIT_CODES or timeout is None: + return False + timeout_s = float(timeout) + if timeout_s <= 0: + return False + return time.monotonic() - started_at + 0.05 >= timeout_s + async def execute_command( self, command: Union[str, List[str]], timeout: Optional[int] = None, stream: bool = True ) -> CommandResult: @@ -241,16 +296,28 @@ async def execute_command( if not self.container or not self.client: raise RuntimeError('Container is not running') + original_command = command + exec_command = self._wrap_command_timeout(command, timeout) + wait_timeout = self._outer_timeout(timeout) + if not stream: try: + started_at = time.monotonic() exit_code, stdout, stderr = await asyncio.wait_for( - self._run_blocking(self._run_buffered, command), timeout=timeout + self._run_blocking(self._run_buffered, exec_command), timeout=wait_timeout + ) + if self._is_timeout_exit(exit_code, timeout, started_at): + status = ExecutionStatus.TIMEOUT + exit_code = -1 + stderr = stderr or f'Command timed out after {timeout} seconds' + else: + status = ExecutionStatus.SUCCESS if exit_code == 0 else ExecutionStatus.ERROR + return CommandResult( + command=original_command, status=status, exit_code=exit_code, stdout=stdout, stderr=stderr ) - status = ExecutionStatus.SUCCESS if exit_code == 0 else ExecutionStatus.ERROR - return CommandResult(command=command, status=status, exit_code=exit_code, stdout=stdout, stderr=stderr) except asyncio.TimeoutError: return CommandResult( - command=command, + command=original_command, status=ExecutionStatus.TIMEOUT, exit_code=-1, stdout='', @@ -258,12 +325,12 @@ async def execute_command( ) except Exception as e: return CommandResult( - command=command, status=ExecutionStatus.ERROR, exit_code=-1, stdout='', stderr=str(e) + command=original_command, status=ExecutionStatus.ERROR, exit_code=-1, stdout='', stderr=str(e) ) # Streaming path: keep exec_id so we can exec_kill on cancel/timeout. exec_meta = await self._run_blocking( - self.client.api.exec_create, container=self.container.id, cmd=command, tty=False + self.client.api.exec_create, container=self.container.id, cmd=exec_command, tty=False ) exec_id = exec_meta['Id'] @@ -284,23 +351,30 @@ async def _consume() -> None: logger.error(f'[📦 {self.id}] {line}') try: - await asyncio.wait_for(_consume(), timeout=timeout) + started_at = time.monotonic() + await asyncio.wait_for(_consume(), timeout=wait_timeout) inspect = await self._run_blocking(self.client.api.exec_inspect, exec_id) exit_code = inspect.get('ExitCode') if exit_code is None: exit_code = -1 - status = ExecutionStatus.SUCCESS if exit_code == 0 else ExecutionStatus.ERROR + if self._is_timeout_exit(exit_code, timeout, started_at): + status = ExecutionStatus.TIMEOUT + exit_code = -1 + stderr = ''.join(stderr_parts) or f'Command timed out after {timeout} seconds' + else: + status = ExecutionStatus.SUCCESS if exit_code == 0 else ExecutionStatus.ERROR + stderr = ''.join(stderr_parts) return CommandResult( - command=command, + command=original_command, status=status, exit_code=exit_code, stdout=''.join(stdout_parts), - stderr=''.join(stderr_parts), + stderr=stderr, ) except asyncio.TimeoutError: await self._kill_exec_safe(exec_id) return CommandResult( - command=command, + command=original_command, status=ExecutionStatus.TIMEOUT, exit_code=-1, stdout=''.join(stdout_parts), @@ -312,7 +386,7 @@ async def _consume() -> None: except Exception as e: await self._kill_exec_safe(exec_id) return CommandResult( - command=command, + command=original_command, status=ExecutionStatus.ERROR, exit_code=-1, stdout=''.join(stdout_parts), 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 c8d3e53..74efbfb 100644 --- a/ms_enclave/sandbox/tools/sandbox_tools/multi_code_executor.py +++ b/ms_enclave/sandbox/tools/sandbox_tools/multi_code_executor.py @@ -193,7 +193,9 @@ async def execute( # Run phase run_res = await self._exec_in_dir(sandbox_context, workdir, run_cmd, timeout=run_timeout or 30) - status = ExecutionStatus.SUCCESS if run_res.exit_code == 0 else ExecutionStatus.ERROR + status = ExecutionStatus.TIMEOUT if run_res.status == ExecutionStatus.TIMEOUT else ( + ExecutionStatus.SUCCESS if run_res.exit_code == 0 else ExecutionStatus.ERROR + ) return ToolResult( tool_name=self.name, status=status, diff --git a/ms_enclave/sandbox/tools/sandbox_tools/notebook_executor.py b/ms_enclave/sandbox/tools/sandbox_tools/notebook_executor.py index b20b41b..257866c 100644 --- a/ms_enclave/sandbox/tools/sandbox_tools/notebook_executor.py +++ b/ms_enclave/sandbox/tools/sandbox_tools/notebook_executor.py @@ -49,7 +49,9 @@ async def execute(self, sandbox_context: 'Sandbox', code: str, timeout: Optional # Execute code using the sandbox's Jupyter kernel result = await self._execute_in_kernel(sandbox_context, code, timeout) - if result.exit_code == 0: + if result.status == ExecutionStatus.TIMEOUT: + status = ExecutionStatus.TIMEOUT + elif result.exit_code == 0: status = ExecutionStatus.SUCCESS else: status = ExecutionStatus.ERROR @@ -82,6 +84,7 @@ async def _execute_in_kernel(self, sandbox_context: 'Sandbox', code: str, timeou outputs = [] result = None error_occurred = False + timed_out = False error_msg = '' actual_timeout = timeout or 30 start_time = time.time() @@ -90,8 +93,11 @@ async def _execute_in_kernel(self, sandbox_context: 'Sandbox', code: str, timeou elapsed = time.time() - start_time if elapsed >= actual_timeout: error_occurred = True + timed_out = True error_msg = f'Execution timed out after {actual_timeout} seconds' logger.error(error_msg) + await self._interrupt_kernel(sandbox_context) + self._drain_execution_messages(sandbox_context, msg_id) break try: @@ -104,7 +110,7 @@ async def _execute_in_kernel(self, sandbox_context: 'Sandbox', code: str, timeou if parent_msg_id != msg_id: continue - msg_type = msg.get('msg_type', '') + msg_type = msg.get('msg_type') or msg.get('header', {}).get('msg_type', '') msg_content = msg.get('content', {}) if msg_type == 'stream': @@ -119,6 +125,8 @@ async def _execute_in_kernel(self, sandbox_context: 'Sandbox', code: str, timeou break except Exception as e: + if type(e).__name__ == 'WebSocketTimeoutException': + continue logger.error(f'Error receiving message: {e}') error_occurred = True error_msg = str(e) @@ -131,13 +139,52 @@ async def _execute_in_kernel(self, sandbox_context: 'Sandbox', code: str, timeou output_text += f'\n{error_msg}' return CommandResult( - status=ExecutionStatus.SUCCESS if not error_occurred else ExecutionStatus.ERROR, + status=ExecutionStatus.TIMEOUT if timed_out else + (ExecutionStatus.SUCCESS if not error_occurred else ExecutionStatus.ERROR), command=code, exit_code=1 if error_occurred else 0, stdout=output_text if not error_occurred else '', stderr=output_text if error_occurred else '' ) + async def _interrupt_kernel(self, sandbox_context: 'Sandbox') -> None: + if not getattr(sandbox_context, 'base_url', None) or not getattr(sandbox_context, 'kernel_id', None): + return + + def _post_interrupt(): + import requests + url = f'{sandbox_context.base_url}/api/kernels/{sandbox_context.kernel_id}/interrupt' + return requests.post(url, timeout=5) + + try: + import asyncio + response = await asyncio.to_thread(_post_interrupt) + if response.status_code not in (200, 204): + logger.warning(f'Kernel interrupt returned HTTP {response.status_code}: {response.text}') + except Exception as e: + logger.warning(f'Failed to interrupt notebook kernel after timeout: {e}') + + def _drain_execution_messages(self, sandbox_context: 'Sandbox', msg_id: str, timeout: float = 2.0) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + try: + remaining = max(0.1, min(0.5, deadline - time.time())) + sandbox_context.ws.settimeout(remaining) + msg = json.loads(sandbox_context.ws.recv()) + except Exception as e: + if type(e).__name__ == 'WebSocketTimeoutException': + continue + logger.debug(f'Error draining notebook timeout messages: {e}') + return + + if msg.get('parent_header', {}).get('msg_id') != msg_id: + continue + + msg_type = msg.get('msg_type') or msg.get('header', {}).get('msg_type', '') + msg_content = msg.get('content', {}) + if msg_type == 'status' and msg_content.get('execution_state') == 'idle': + return + def _send_execute_request(self, sandbox_context: 'Sandbox', code: str) -> str: """Send code execution request to kernel.""" # Generate a unique message ID diff --git a/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py b/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py index cdb4c8f..77f82f4 100644 --- a/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py +++ b/ms_enclave/sandbox/tools/sandbox_tools/python_executor.py @@ -64,7 +64,9 @@ async def execute(self, sandbox_context: 'DockerSandbox', code: str, timeout: Op command = f'python {script_path}' result = await sandbox_context.execute_command(command, timeout=timeout) - if result.exit_code == 0: + if result.status == ExecutionStatus.TIMEOUT: + status = ExecutionStatus.TIMEOUT + elif result.exit_code == 0: status = ExecutionStatus.SUCCESS else: status = ExecutionStatus.ERROR diff --git a/ms_enclave/sandbox/tools/sandbox_tools/shell_executor.py b/ms_enclave/sandbox/tools/sandbox_tools/shell_executor.py index a92d4f1..46dac38 100644 --- a/ms_enclave/sandbox/tools/sandbox_tools/shell_executor.py +++ b/ms_enclave/sandbox/tools/sandbox_tools/shell_executor.py @@ -66,6 +66,13 @@ async def execute( try: result = await sandbox_context.execute_command(command, timeout=timeout) + if result.status == ExecutionStatus.TIMEOUT: + return ToolResult( + tool_name=self.name, + status=ExecutionStatus.TIMEOUT, + output=result.stdout, + error=result.stderr if result.stderr else f'Command timed out after {timeout} seconds' + ) if result.exit_code == 0: return ToolResult( tool_name=self.name, diff --git a/tests/test_tool.py b/tests/test_tool.py index 4e51263..1fcb3ec 100644 --- a/tests/test_tool.py +++ b/tests/test_tool.py @@ -74,6 +74,109 @@ async def test_shell_executor(self): self.assertEqual(result.status, ExecutionStatus.SUCCESS) self.assertIn('Hello, Shell!', result.output) + async def test_execute_command_timeout_wrapper_keeps_original_command(self): + command = ['bash', '-c', 'echo wrapped'] + + result = await self.docker_sandbox.execute_command(command, timeout=5) + + self.assertEqual(result.status, ExecutionStatus.SUCCESS) + self.assertEqual(result.command, command) + self.assertIn('wrapped', result.stdout) + + async def test_execute_command_zero_timeout_disables_timeout(self): + result = await self.docker_sandbox.execute_command(['bash', '-c', 'echo no-timeout'], timeout=0) + + self.assertEqual(result.status, ExecutionStatus.SUCCESS) + self.assertIn('no-timeout', result.stdout) + + async def test_shell_executor_exit_124_is_not_timeout(self): + self.docker_sandbox.add_tool( + ToolFactory.create_tool('shell_executor') + ) + + result = await self.docker_sandbox.execute_tool( + 'shell_executor', {'command': ['bash', '-c', 'exit 124'], 'timeout': 5} + ) + + self.assertEqual(result.status, ExecutionStatus.ERROR) + + async def test_shell_executor_sigkill_is_not_timeout(self): + self.docker_sandbox.add_tool( + ToolFactory.create_tool('shell_executor') + ) + + result = await self.docker_sandbox.execute_tool( + 'shell_executor', {'command': ['bash', '-c', 'kill -9 $$'], 'timeout': 5} + ) + + self.assertEqual(result.status, ExecutionStatus.ERROR) + + async def test_shell_executor_timeout_terminates_child_process(self): + self.docker_sandbox.add_tool( + ToolFactory.create_tool('shell_executor') + ) + marker = '/tmp/ms-enclave-shell-timeout-marker' + + await self.docker_sandbox.execute_tool( + 'shell_executor', {'command': ['bash', '-c', f'rm -f {marker}'], 'timeout': 5} + ) + result = await self.docker_sandbox.execute_tool( + 'shell_executor', {'command': ['bash', '-c', f'sleep 2; touch {marker}'], 'timeout': 0.3} + ) + + self.assertEqual(result.status, ExecutionStatus.TIMEOUT) + await asyncio.sleep(2.5) + check = await self.docker_sandbox.execute_tool( + 'shell_executor', {'command': ['bash', '-c', f'test -e {marker} && echo PRESENT || echo ABSENT']} + ) + self.assertEqual(check.output.strip(), 'ABSENT') + + async def test_shell_executor_timeout_terminates_grandchild_process(self): + self.docker_sandbox.add_tool( + ToolFactory.create_tool('shell_executor') + ) + marker = '/tmp/ms-enclave-shell-grandchild-timeout-marker' + + await self.docker_sandbox.execute_tool( + 'shell_executor', {'command': ['bash', '-c', f'rm -f {marker}'], 'timeout': 5} + ) + result = await self.docker_sandbox.execute_tool( + 'shell_executor', + { + 'command': ['bash', '-c', f'bash -c "bash -c \\"sleep 2; touch {marker}\\" & wait"'], + 'timeout': 0.3 + } + ) + + self.assertEqual(result.status, ExecutionStatus.TIMEOUT) + await asyncio.sleep(2.5) + check = await self.docker_sandbox.execute_tool( + 'shell_executor', {'command': ['bash', '-c', f'test -e {marker} && echo PRESENT || echo ABSENT']} + ) + self.assertEqual(check.output.strip(), 'ABSENT') + + async def test_python_executor_timeout_terminates_child_process(self): + self.docker_sandbox.add_tool( + ToolFactory.create_tool('shell_executor') + ) + self.docker_sandbox.add_tool( + ToolFactory.create_tool('python_executor') + ) + marker = '/tmp/ms-enclave-python-timeout-marker' + + await self.docker_sandbox.execute_tool( + 'shell_executor', {'command': ['bash', '-c', f'rm -f {marker}'], 'timeout': 5} + ) + code = f"import os\nos.system('sleep 2; touch {marker}')\n" + result = await self.docker_sandbox.execute_tool('python_executor', {'code': code, 'timeout': 0.3}) + + self.assertEqual(result.status, ExecutionStatus.TIMEOUT) + await asyncio.sleep(2.5) + check = await self.docker_sandbox.execute_tool( + 'shell_executor', {'command': ['bash', '-c', f'test -e {marker} && echo PRESENT || echo ABSENT']} + ) + self.assertEqual(check.output.strip(), 'ABSENT') + async def test_file_operation(self): self.docker_sandbox.add_tool( ToolFactory.create_tool('file_operation') @@ -169,6 +272,25 @@ async def test_notebook_executor_state_persistence(self): self.assertEqual(result.output.strip(), result2.output.strip()) self.assertEqual(result2.status, ExecutionStatus.SUCCESS) + async def test_notebook_executor_timeout_interrupts_kernel(self): + self.docker_sandbox.add_tool( + ToolFactory.create_tool('notebook_executor') + ) + + await self.docker_sandbox.execute_tool( + 'notebook_executor', {'code': "globals().pop('timeout_marker', None)"} + ) + result = await self.docker_sandbox.execute_tool( + 'notebook_executor', {'code': "import time\ntime.sleep(5)\ntimeout_marker = 'PRESENT'", 'timeout': 1} + ) + + self.assertEqual(result.status, ExecutionStatus.TIMEOUT) + check = await self.docker_sandbox.execute_tool( + 'notebook_executor', {'code': "globals().get('timeout_marker', 'ABSENT')"} + ) + self.assertEqual(check.status, ExecutionStatus.SUCCESS) + self.assertIn('ABSENT', check.output) + if __name__ == '__main__': import asyncio unittest.main()