Improve process management and timeout handling

- Enhance process termination reliability
- Improve timeout handling with graceful shutdown
- Add proper async/await support
- Fix process cleanup in tests
- Improve test coverage

Author: tumf

src/mcp_shell_server/process_manager.py

@@ -65,12 +65,13 @@ async def start_process_async(
         process = await self.create_process(
             " ".join(cmd), directory=None, timeout=timeout
         )
+        process.is_running = lambda self=process: self.returncode is None  # type: ignore
         return process
 
-    def start_process(
+    async def start_process(
         self, cmd: List[str], timeout: Optional[int] = None
     ) -> asyncio.subprocess.Process:
-        """Start a new process synchronously.
+        """Start a new process asynchronously.
 
         Args:
             cmd: Command to execute as list of strings
@@ -79,8 +80,8 @@ def start_process(
         Returns:
             Process object
         """
-        process = asyncio.get_event_loop().run_until_complete(
-            self.start_process_async(cmd, timeout)
+        process = await self.create_process(
+            " ".join(cmd), directory=None, timeout=timeout
         )
         process.is_running = lambda self=process: self.returncode is None  # type: ignore
         return process
@@ -186,22 +187,38 @@ async def execute_with_timeout(
         """
         stdin_bytes = stdin.encode() if stdin else None
 
-        async def communicate_with_timeout():
+        async def _kill_process():
+            if process.returncode is not None:
+                return
+
             try:
-                return await process.communicate(input=stdin_bytes)
-            except asyncio.TimeoutError:
-                process.kill()  # Kill process on timeout
-                raise
+                # Try graceful termination first
+                process.terminate()
+                for _ in range(5):  # Wait up to 0.5 seconds
+                    if process.returncode is not None:
+                        return
+                    await asyncio.sleep(0.1)
+
+                # Force kill if still running
+                if process.returncode is None:
+                    process.kill()
+                    await asyncio.wait_for(process.wait(), timeout=1.0)
             except Exception as e:
-                try:
-                    await process.wait()
-                except Exception:
-                    pass
-                raise e
+                logging.warning(f"Error killing process: {e}")
 
-        if timeout:
-            return await asyncio.wait_for(communicate_with_timeout(), timeout=timeout)
-        return await communicate_with_timeout()
+        try:
+            if timeout:
+                try:
+                    return await asyncio.wait_for(
+                        process.communicate(input=stdin_bytes), timeout=timeout
+                    )
+                except asyncio.TimeoutError:
+                    await _kill_process()
+                    raise
+            return await process.communicate(input=stdin_bytes)
+        except Exception as e:
+            await _kill_process()
+            raise e
 
     async def execute_pipeline(
         self,

tests/test_process_manager.py

@@ -1,5 +1,6 @@
 """Tests for the ProcessManager class."""
 
+import asyncio
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -13,6 +14,7 @@ def create_mock_process():
     process.returncode = 0
     process.communicate = AsyncMock(return_value=(b"output", b"error"))
     process.wait = AsyncMock(return_value=0)
+    process.terminate = MagicMock()
     process.kill = MagicMock()
     return process
 
@@ -65,16 +67,23 @@ async def test_execute_with_timeout_success(process_manager):
 async def test_execute_with_timeout_timeout(process_manager):
     """Test executing a process that times out."""
     mock_proc = create_mock_process()
-    exc = TimeoutError("Process timed out")
+    exc = asyncio.TimeoutError("Process timed out")
     mock_proc.communicate.side_effect = exc
+    mock_proc.returncode = None  # プロセスがまだ実行中の状態をシミュレート
+
+    # プロセスの終了状態をシミュレート
+    async def set_returncode():
+        mock_proc.returncode = -15  # SIGTERM
+
+    mock_proc.wait.side_effect = set_returncode
 
     with pytest.raises(TimeoutError):
         await process_manager.execute_with_timeout(
             mock_proc,
             timeout=1,
         )
 
-    mock_proc.kill.assert_called_once()
+    mock_proc.terminate.assert_called_once()
 
 
 @pytest.mark.asyncio

tests/test_process_manager_macos.py

@@ -1,110 +1,149 @@
 # tests/test_process_manager_macos.py
-import pytest
+import asyncio
 import platform
-import signal
-import time
-import os
 import subprocess
-from typing import Generator
+
+import pytest
 
 pytestmark = [
     pytest.mark.skipif(
-        platform.system() != "Darwin",
-        reason="These tests only run on macOS"
+        platform.system() != "Darwin", reason="These tests only run on macOS"
     ),
     pytest.mark.macos,
-    pytest.mark.slow
+    pytest.mark.slow,
 ]
 
+
 @pytest.fixture
 def process_manager():
     from mcp_shell_server.process_manager import ProcessManager
+
     pm = ProcessManager()
-    yield pm
-    pm.cleanup_all()
+    try:
+        yield pm
+    finally:
+        asyncio.run(pm.cleanup_all())
+
 
 def get_process_status(pid: int) -> str:
     """Get process status using ps command."""
     try:
         ps = subprocess.run(
-            ['ps', '-o', 'stat=', '-p', str(pid)],
-            capture_output=True,
-            text=True
+            ["ps", "-o", "stat=", "-p", str(pid)], capture_output=True, text=True
         )
         return ps.stdout.strip()
     except subprocess.CalledProcessError:
         return ""
 
-def test_zombie_process_cleanup(process_manager):
+
+@pytest.mark.asyncio
+async def test_zombie_process_cleanup(process_manager):
     """Test that background processes don't become zombies."""
     cmd = ["sh", "-c", "sleep 0.5 & wait"]
-    process = process_manager.start_process(cmd)
-    
+    process = await process_manager.start_process(cmd)
+
     # Wait for the background process to finish
-    time.sleep(1)
-    
+    await asyncio.sleep(1)
+
     # Get process status
     status = get_process_status(process.pid)
-    
+
     # Verify process is either gone or not zombie (Z state)
-    assert 'Z' not in status, f"Process {process.pid} is zombie (status: {status})"
+    assert "Z" not in status, f"Process {process.pid} is zombie (status: {status})"
 
-def test_process_timeout(process_manager):
+
+@pytest.mark.asyncio
+async def test_process_timeout(process_manager):
     """Test process timeout functionality."""
     # Start a process that should timeout
     cmd = ["sleep", "10"]
-    process = process_manager.start_process(cmd, timeout=1)
-    
-    # Wait slightly longer than the timeout
-    time.sleep(1.5)
-    
-    # Verify process was terminated
-    assert not process.is_running()
-    assert process.returncode is not None
-
-def test_multiple_process_cleanup(process_manager):
+    process = await process_manager.start_process(cmd)
+
+    try:
+        # Communicate with timeout
+        with pytest.raises(TimeoutError):
+            _, _ = await process_manager.execute_with_timeout(process, timeout=1)
+
+        # プロセスが終了するまで待つ
+        try:
+            await asyncio.wait_for(process.wait(), timeout=1.0)
+        except asyncio.TimeoutError:
+            process.kill()  # Force kill
+
+        # Wait for termination
+        await asyncio.wait_for(process.wait(), timeout=0.5)
+
+        # Verify process was terminated
+        assert process.returncode is not None
+        assert not process.is_running()
+    finally:
+        if process.returncode is None:
+            try:
+                process.kill()
+                await asyncio.wait_for(process.wait(), timeout=0.5)
+            except (ProcessLookupError, asyncio.TimeoutError):
+                pass
+
+
+@pytest.mark.asyncio
+async def test_multiple_process_cleanup(process_manager):
     """Test cleanup of multiple processes."""
     # Start multiple background processes
-    processes = [
-        process_manager.start_process(["sleep", "2"])
-        for _ in range(3)
-    ]
-    
+    # Start multiple processes in parallel
+    processes = await asyncio.gather(
+        *[process_manager.start_process(["sleep", "2"]) for _ in range(3)]
+    )
+
     # Give them a moment to start
-    time.sleep(0.1)
-    
-    # Verify they're all running
-    assert all(p.is_running() for p in processes)
-    
-    # Cleanup
-    process_manager.cleanup_all()
-    
-    # Give cleanup a moment to complete
-    time.sleep(0.1)
-    
-    # Verify all processes are gone
-    for p in processes:
-        status = get_process_status(p.pid)
-        assert status == "", f"Process {p.pid} still exists with status: {status}"
-
-def test_process_group_termination(process_manager):
+    await asyncio.sleep(0.1)
+
+    try:
+        # Verify they're all running
+        assert all(p.is_running() for p in processes)
+
+        # Cleanup
+        await process_manager.cleanup_all()
+
+        # Give cleanup a moment to complete
+        await asyncio.sleep(0.1)
+
+        # Verify all processes are gone
+        for p in processes:
+            status = get_process_status(p.pid)
+            assert status == "", f"Process {p.pid} still exists with status: {status}"
+    finally:
+        # Ensure cleanup in case of test failure
+        for p in processes:
+            if p.returncode is None:
+                try:
+                    p.kill()
+                except ProcessLookupError:
+                    pass
+
+
+@pytest.mark.asyncio
+async def test_process_group_termination(process_manager):
     """Test that entire process group is terminated."""
     # Create a process that spawns children
     cmd = ["sh", "-c", "sleep 10 & sleep 10 & sleep 10 & wait"]
-    process = process_manager.start_process(cmd)
-    
-    # Give processes time to start
-    time.sleep(0.5)
-    
-    # Kill the main process
-    process.kill()
-    
-    # Wait a moment for cleanup
-    time.sleep(0.5)
-    
-    # Check if any processes from the group remain
-    ps = subprocess.run(
-        ["pgrep", "-g", str(process.pid)],
-        capture_output=True
-    )
-    assert ps.returncode != 0, "Process group still exists"
+    process = await process_manager.start_process(cmd)
+
+    try:
+        # Give processes time to start
+        await asyncio.sleep(0.5)
+
+        # Kill the main process
+        process.kill()
+
+        # Wait a moment for cleanup
+        await asyncio.sleep(0.5)
+
+        # Check if any processes from the group remain
+        ps = subprocess.run(["pgrep", "-g", str(process.pid)], capture_output=True)
+        assert ps.returncode != 0, "Process group still exists"
+    finally:
+        if process.returncode is None:
+            try:
+                process.kill()
+            except ProcessLookupError:
+                pass

tests/test_server.py

@@ -44,7 +44,14 @@ def setup_mock_subprocess(monkeypatch):
     """Set up mock subprocess to avoid interactive shell warnings"""
 
     async def mock_create_subprocess_shell(
-        cmd, stdin=None, stdout=None, stderr=None, env=None, cwd=None
+        cmd,
+        stdin=None,
+        stdout=None,
+        stderr=None,
+        env=None,
+        cwd=None,
+        preexec_fn=None,
+        start_new_session=None,
     ):
         # Return appropriate output based on command
         if "echo" in cmd: