Fix Docker sandbox timeout handling

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,29 +296,41 @@ 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='',
                     stderr=f'Command timed out after {timeout} seconds',
                 )
             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),

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,

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

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

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,