Repository: vndee/llm-sandbox

Issue:
New timeout handling breaks async
**Describe the bug**
A recent update improves timeout handling but also breaks async execution
https://github.com/vndee/llm-sandbox/blob/90087b10c6b3e81488eee184200d20256b1dbf5b/docs/integrations.md?plain=1#L630-L641

Now the code returns `signal only works in main thread of the main interpreter`.

Hints from issue comments:
@github-actions: A recent update to improve timeout handling in the codebase has inadvertently broken asynchronous execution, as it now triggers the error: "signal only works in main thread of the main interpreter." This issue likely stems from using signals, which are not compatible with threads other than the main thread in Python. The user links to a specific section in the documentation for context and highlights the unintended consequences of the changes, suggesting the need for further investigation or adjustments to restore async compatibility.@vndee: Fixed with https://github.com/vndee/llm-sandbox/pull/69. Please help me check if it resolve your issue.@Nifury: Yes it has been fixed. Thank you!

Solution patch:
diff --git a/docs/integrations.md b/docs/integrations.md
index 4922d8d..68b8382 100644
--- a/docs/integrations.md
+++ b/docs/integrations.md
@@ -568,98 +568,7 @@ if response.choices[0].message.function_call:
 ### Generic Integration Pattern
 
 ```python
-from typing import Protocol, Optional, Dict, Any
-import asyncio
-
-class CodeExecutor(Protocol):
-    """Protocol for code execution integration"""
-
-    def execute(
-        self,
-        code: str,
-        language: str = "python",
-        libraries: Optional[List[str]] = None,
-        **kwargs
-    ) -> Dict[str, Any]:
-        ...
-
-class SandboxCodeExecutor:
-    """Sandbox implementation of CodeExecutor"""
-
-    def __init__(self, default_security_policy=None):
-        self.default_security_policy = default_security_policy
-
-    def execute(
-        self,
-        code: str,
-        language: str = "python",
-        libraries: Optional[List[str]] = None,
-        security_policy=None,
-        **kwargs
-    ) -> Dict[str, Any]:
-        """
-        Execute code in sandbox.
-
-        Returns:
-            Dictionary with stdout, stderr, exit_code, and plots
-        """
-        policy = security_policy or self.default_security_policy
-
-        try:
-            with SandboxSession(
-                lang=language,
-                security_policy=policy,
-                **kwargs
-            ) as session:
-                result = session.run(code, libraries)
-
-                return {
-                    "success": result.exit_code == 0,
-                    "stdout": result.stdout,
-                    "stderr": result.stderr,
-                    "exit_code": result.exit_code,
-                    "plots": getattr(result, "plots", [])
-                }
-        except Exception as e:
-            return {
-                "success": False,
-                "error": str(e),
-                "exit_code": -1
-            }
-
-    async def execute_async(
-        self,
-        code: str,
-        **kwargs
-    ) -> Dict[str, Any]:
-        """Async wrapper for execution"""
-        loop = asyncio.get_event_loop()
-        return await loop.run_in_executor(
-            None,
-            lambda: self.execute(code, **kwargs)
-        )
-
-# Use in any framework
-executor = SandboxCodeExecutor(
-    default_security_policy=get_security_policy("production")
-)
-
-# Sync execution
-result = executor.execute(
-    "print('Hello, World!')",
-    language="python"
-)
-
-# Async execution
-async def main():
-    result = await executor.execute_async(
-        "print('Hello, Async!')",
-        language="python"
-    )
-    print(result["stdout"])
-
-# Run async
-asyncio.run(main())
+--8<-- "examples/async_example.py"
 ```
 
 ### Middleware Pattern
diff --git a/examples/async_example.py b/examples/async_example.py
new file mode 100644
index 0000000..487a6bc
--- /dev/null
+++ b/examples/async_example.py
@@ -0,0 +1,78 @@
+"""Async execution example for LLM Sandbox."""
+
+import asyncio
+from typing import Any, Protocol
+
+from llm_sandbox import SandboxSession
+
+
+class CodeExecutor(Protocol):
+    """Protocol for code execution integration."""
+
+    def execute(
+        self, code: str, language: str = "python", libraries: list[str] | None = None, **kwargs: Any
+    ) -> dict[str, Any]:
+        """Execute code and return results."""
+        ...
+
+
+class SandboxCodeExecutor:
+    """Sandbox implementation of CodeExecutor."""
+
+    def __init__(self, default_security_policy: Any = None) -> None:
+        """Initialize the executor with optional security policy."""
+        self.default_security_policy = default_security_policy
+
+    def execute(
+        self,
+        code: str,
+        language: str = "python",
+        libraries: list[str] | None = None,
+        security_policy: Any = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Execute code in sandbox.
+
+        Returns:
+            Dictionary with stdout, stderr, exit_code, and plots
+
+        """
+        policy = security_policy or self.default_security_policy
+
+        try:
+            with SandboxSession(lang=language, security_policy=policy, **kwargs) as session:
+                result = session.run(code, libraries)
+
+                return {
+                    "success": result.exit_code == 0,
+                    "stdout": result.stdout,
+                    "stderr": result.stderr,
+                    "exit_code": result.exit_code,
+                    "plots": getattr(result, "plots", []),
+                }
+        except Exception as e:  # noqa: BLE001
+            return {"success": False, "error": str(e), "exit_code": -1}
+
+    async def execute_async(self, code: str, **kwargs: Any) -> dict[str, Any]:
+        """Async wrapper for execution."""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(None, lambda: self.execute(code, **kwargs))
+
+
+async def main() -> None:
+    """Demonstrate both sync and async execution."""
+    # Use in any framework
+    executor = SandboxCodeExecutor(default_security_policy=None)
+
+    # Sync execution
+    result = executor.execute("print('Hello, World!')", language="python")
+    print(f"Sync result: {result['stdout']}")  # noqa: T201
+
+    # Async execution
+    result = await executor.execute_async("print('Hello, Async!')", language="python")
+    print(f"Async result: {result['stdout']}")  # noqa: T201
+
+
+# Run async
+if __name__ == "__main__":
+    asyncio.run(main())
diff --git a/llm_sandbox/core/mixins.py b/llm_sandbox/core/mixins.py
index c13cc90..d019fbc 100644
--- a/llm_sandbox/core/mixins.py
+++ b/llm_sandbox/core/mixins.py
@@ -1,17 +1,21 @@
 """Mixins for common functionality."""
 
 import io
-import signal
 import tarfile
 import threading
-import types
 from abc import abstractmethod
+from contextlib import suppress
 from pathlib import Path
-from typing import Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol
+
+if TYPE_CHECKING:
+    from types import TracebackType
 
 from llm_sandbox.data import ConsoleOutput
 from llm_sandbox.exceptions import CommandEmptyError, NotOpenSessionError, SandboxTimeoutError
 
+CLEANUP_THREAD_TIMEOUT = 0.1
+
 
 class ContainerAPI(Protocol):
     """Protocol for container operations."""
@@ -44,14 +48,32 @@ class ContainerAPI(Protocol):
 class TimeoutMixin:
     """Mixin for timeout functionality."""
 
-    def _execute_with_timeout(self, func: Any, timeout: float | None = None, *args: Any, **kwargs: Any) -> Any:
+    logger: Any
+
+    def _execute_with_timeout(
+        self, func: Any, *args: Any, timeout: float | None = None, force_kill_on_timeout: bool = True, **kwargs: Any
+    ) -> Any:
         """Execute a function with timeout monitoring.
 
-        Uses signal-based timeout on Unix systems and threading-based timeout on Windows.
+        Uses threading-based timeout that works in all contexts with proper cleanup:
+        - Main thread
+        - Worker threads
+        - Async contexts (asyncio.run_in_executor)
+        - Any other execution context
+        - Includes thread cleanup to prevent resource leaks
+
+        **Thread Limitation**: This implementation cannot forcefully terminate running threads.
+        If the operation involves long-running C extensions or I/O-bound code, the daemon
+        thread may continue executing in the background even after timeout. True cancellation
+        is achieved at the container level via `_handle_timeout()` which kills/stops the
+        entire container process.
 
         Args:
             func: The function to execute
             timeout: Timeout in seconds
+            force_kill_on_timeout: If True, calls `_handle_timeout()` directly on timeout
+                for immediate container-level cancellation. Use with caution as this
+                terminates the entire container.
             *args: Arguments to pass to the function
             **kwargs: Keyword arguments to pass to the function
 
@@ -59,52 +81,61 @@ class TimeoutMixin:
             The result from the function execution
 
         Raises:
-            SandboxTimeoutError: If the function execution times out
+            SandboxTimeoutError: If the function execution times out. The calling code
+                should handle this by invoking container-level cleanup via `_handle_timeout()`.
 
         """
         if timeout is None:
             return func(*args, **kwargs)
 
-        # For Unix systems with SIGALRM, use signal-based timeout
-        if hasattr(signal, "SIGALRM"):
-
-            def timeout_handler(signum: int, frame: types.FrameType | None) -> None:  # noqa: ARG001
-                msg = f"Operation timed out after {timeout} seconds"
-                raise SandboxTimeoutError(msg)
-
-            old_handler = signal.signal(signal.SIGALRM, timeout_handler)
-            signal.alarm(int(timeout))
-
-            try:
-                return func(*args, **kwargs)
-            finally:
-                signal.alarm(0)
-                signal.signal(signal.SIGALRM, old_handler)
-
-        # For Windows or systems without SIGALRM, use threading
         result: list[Any] = [None]
-        exception: list[Exception | None] = [None]
+        exception_info: list[tuple[type[BaseException], BaseException, TracebackType | Any] | None] = [None]
+        completed = threading.Event()
 
         def target() -> None:
             try:
                 result[0] = func(*args, **kwargs)
-            except Exception as e:  # noqa: BLE001
-                exception[0] = e
+            except BaseException as e:  # noqa: BLE001 # NOSONAR
+                exception_info[0] = (type(e), e, e.__traceback__)
+            finally:
+                completed.set()
 
-        thread = threading.Thread(target=target)
-        thread.daemon = True
+        thread = threading.Thread(target=target, daemon=True)
         thread.start()
-        thread.join(timeout)
-
-        if thread.is_alive():
-            # Note: We can't actually kill the thread, but we can timeout
-            msg = f"Operation timed out after {timeout} seconds"
-            raise SandboxTimeoutError(msg)
 
-        if exception[0]:
-            raise exception[0]
+        try:
+            # Wait for completion or timeout
+            if not completed.wait(timeout):
+                msg = f"Operation timed out after {timeout} seconds"
 
-        return result[0]
+                # Optional: Force container-level kill for true cancellation
+                handler = getattr(self, "_handle_timeout", None)
+                if force_kill_on_timeout and callable(handler):
+
+                    def cleanup_async() -> None:
+                        try:
+                            handler()  # pyright: ignore[reportAttributeAccess]
+                        except Exception:  # noqa: BLE001
+                            self.logger.warning("Failed to cleanup container after timeout")
+
+                    # Run cleanup in a separate daemon thread to avoid blocking
+                    cleanup_thread = threading.Thread(target=cleanup_async, daemon=True)
+                    cleanup_thread.start()
+
+                raise SandboxTimeoutError(msg, timeout_duration=timeout)
+
+            if exception_info[0]:
+                _, exc_value, exc_traceback = exception_info[0]
+                raise exc_value.with_traceback(exc_traceback)
+
+            return result[0]
+        finally:
+            # Best-effort thread cleanup to prevent resource leaks
+            # This will reclaim finished threads promptly
+            with suppress(Exception):
+                thread.join(
+                    timeout=0.0 if completed.is_set() else CLEANUP_THREAD_TIMEOUT
+                )  # Non-blocking join for cleanup
 
 
 class FileOperationsMixin:
diff --git a/llm_sandbox/core/session_base.py b/llm_sandbox/core/session_base.py
index 1a16324..8511a6d 100644
--- a/llm_sandbox/core/session_base.py
+++ b/llm_sandbox/core/session_base.py
@@ -382,7 +382,7 @@ class BaseSession(
                 )
 
         try:
-            result = self._execute_with_timeout(_run_code, actual_timeout)
+            result = self._execute_with_timeout(_run_code, timeout=actual_timeout)
             return cast("ConsoleOutput", result)
         except SandboxTimeoutError:
             self._handle_timeout()


Test patch:
diff --git a/tests/test_mixins.py b/tests/test_mixins.py
index dd33a56..ab0f737 100644
--- a/tests/test_mixins.py
+++ b/tests/test_mixins.py
@@ -1,9 +1,8 @@
-# ruff: noqa: SLF001, PLR2004, ARG002, PT011
+# ruff: noqa: SLF001, PLR2004, ARG002, PT011, FBT003
 
 """Test cases for the new architecture mixins."""
 
 import io
-import signal
 import tarfile
 import time
 from pathlib import Path
@@ -39,7 +38,7 @@ class TestTimeoutMixin:
         def test_func(x: int, y: int) -> int:
             return x + y
 
-        result = mixin._execute_with_timeout(test_func, None, 2, 3)
+        result = mixin._execute_with_timeout(test_func, 2, 3, timeout=None, force_kill_on_timeout=False)
         assert result == 5
 
     def test_execute_with_timeout_success(self) -> None:
@@ -49,50 +48,113 @@ class TestTimeoutMixin:
         def fast_func() -> str:
             return "completed"
 
-        result = mixin._execute_with_timeout(fast_func, 2.0)
+        result = mixin._execute_with_timeout(fast_func, timeout=2.0)
         assert result == "completed"
 
-    @pytest.mark.skipif(not hasattr(signal, "SIGALRM"), reason="SIGALRM not available on this platform")
-    def test_execute_with_timeout_signal_based(self) -> None:
-        """Test timeout using signal-based approach (Unix)."""
+    def test_execute_with_timeout_exception_propagation(self) -> None:
+        """Test that exceptions are properly propagated."""
+        mixin = TimeoutMixin()
+
+        def error_func() -> str:
+            raise ValueError
+
+        with pytest.raises(ValueError):
+            mixin._execute_with_timeout(error_func, timeout=2.0)
+
+    def test_execute_with_timeout_actual_timeout(self) -> None:
+        """Test actual timeout scenario with SandboxTimeoutError."""
         mixin = TimeoutMixin()
 
         def slow_func() -> str:
-            time.sleep(2.0)
-            return "should not reach here"
+            time.sleep(0.5)  # Simulated sleep
+            return "should not complete"
 
-        with pytest.raises(SandboxTimeoutError):
-            mixin._execute_with_timeout(slow_func, 1)
+        with pytest.raises(SandboxTimeoutError, match="Operation timed out after 0.1 seconds"):
+            mixin._execute_with_timeout(slow_func, timeout=0.1)
 
-    @patch("llm_sandbox.core.mixins.hasattr")
-    @patch("llm_sandbox.core.mixins.threading.Thread")
-    def test_execute_with_timeout_threading_based(self, mock_thread: MagicMock, mock_hasattr: MagicMock) -> None:
-        """Test timeout using threading-based approach."""
-        # Simulate Windows environment (no SIGALRM)
-        mock_hasattr.return_value = False
+    def test_execute_with_timeout_with_handler_success(self) -> None:
+        """Test timeout with force_kill_on_timeout=True and successful handler call."""
 
-        # Mock thread to simulate timeout
-        mock_thread_instance = Mock()
-        mock_thread_instance.is_alive.return_value = True
-        mock_thread.return_value = mock_thread_instance
+        class TimeoutMixinWithHandler(TimeoutMixin):
+            def __init__(self) -> None:
+                self.logger = Mock()
+                self.handler_called = False
 
-        mixin = TimeoutMixin()
+            def _handle_timeout(self) -> None:
+                self.handler_called = True
 
-        def test_func() -> str:
-            return "test"
+        mixin = TimeoutMixinWithHandler()
 
-        with pytest.raises(SandboxTimeoutError):
-            mixin._execute_with_timeout(test_func, 0.1)
+        def slow_func() -> str:
+            time.sleep(0.5)  # Simulated sleep
+            return "should not complete"
 
-    def test_execute_with_timeout_exception_propagation(self) -> None:
-        """Test that exceptions are properly propagated."""
+        with pytest.raises(SandboxTimeoutError, match="Operation timed out after 0.1 seconds"):
+            mixin._execute_with_timeout(slow_func, timeout=0.1, force_kill_on_timeout=True)
+
+        # Verify the handler was called
+        assert mixin.handler_called
+
+    def test_execute_with_timeout_with_handler_exception(self) -> None:
+        """Test timeout with force_kill_on_timeout=True when handler raises exception."""
+
+        class TimeoutMixinWithFailingHandler(TimeoutMixin):
+            def __init__(self) -> None:
+                self.logger = Mock()
+
+            def _handle_timeout(self) -> None:
+                msg = "Handler failed"
+                raise RuntimeError(msg)
+
+        mixin = TimeoutMixinWithFailingHandler()
+
+        def slow_func() -> str:
+            time.sleep(0.5)  # Simulated sleep
+            return "should not complete"
+
+        with pytest.raises(SandboxTimeoutError, match="Operation timed out after 0.1 seconds"):
+            mixin._execute_with_timeout(slow_func, timeout=0.1, force_kill_on_timeout=True)
+
+        # Verify the warning was logged when handler failed
+        mixin.logger.warning.assert_called_once_with("Failed to cleanup container after timeout")
+
+    def test_execute_with_timeout_no_handler(self) -> None:
+        """Test timeout with force_kill_on_timeout=True but no _handle_timeout method."""
         mixin = TimeoutMixin()
+        mixin.logger = Mock()
 
-        def error_func() -> str:
-            raise ValueError
+        def slow_func() -> str:
+            time.sleep(0.5)  # Simulated sleep
+            return "should not complete"
 
-        with pytest.raises(ValueError):
-            mixin._execute_with_timeout(error_func, 2.0)
+        with pytest.raises(SandboxTimeoutError, match="Operation timed out after 0.1 seconds"):
+            mixin._execute_with_timeout(slow_func, timeout=0.1, force_kill_on_timeout=True)
+
+        # Handler should not be called since it doesn't exist
+        mixin.logger.warning.assert_not_called()
+
+    def test_execute_with_timeout_force_kill_disabled(self) -> None:
+        """Test timeout with force_kill_on_timeout=False."""
+
+        class TimeoutMixinWithHandler(TimeoutMixin):
+            def __init__(self) -> None:
+                self.logger = Mock()
+                self.handler_called = False
+
+            def _handle_timeout(self) -> None:
+                self.handler_called = True
+
+        mixin = TimeoutMixinWithHandler()
+
+        def slow_func() -> str:
+            time.sleep(0.5)  # Simulated sleep
+            return "should not complete"
+
+        with pytest.raises(SandboxTimeoutError, match="Operation timed out after 0.1 seconds"):
+            mixin._execute_with_timeout(slow_func, timeout=0.1, force_kill_on_timeout=False)
+
+        # Handler should not be called when force_kill_on_timeout=False
+        assert not mixin.handler_called
 
 
 class TestFileOperationsMixin:
diff --git a/tests/test_session_base.py b/tests/test_session_base.py
index c7102e8..b460595 100644
--- a/tests/test_session_base.py
+++ b/tests/test_session_base.py
@@ -543,9 +543,10 @@ class TestBaseSessionCodeExecution:
             self.session.run("print('hello')", timeout=5.0)
 
             mock_timeout.assert_called_once()
-            args, _ = mock_timeout.call_args
-            assert len(args) >= 2  # function and timeout
-            assert abs(args[1] - 5.0) < 0.0001  # timeout value
+            args, kwargs = mock_timeout.call_args
+            assert len(args) >= 1  # function
+            assert "timeout" in kwargs
+            assert abs(kwargs["timeout"] - 5.0) < 0.0001  # timeout value
 
     @patch.object(MockBaseSession, "_handle_timeout")
     def test_run_timeout_exception(self, mock_handle_timeout: Mock) -> None: