feat(timeout): add timeout control for command execution to enhance user control over long-running processes

docs(README): update documentation to include timeout feature and examples for command execution
test(timeout): add tests for command execution with timeout to ensure functionality works as expected

Author: tumf

README.md

@@ -8,6 +8,7 @@ A secure shell command execution server implementing the Model Context Protocol
 * **Standard Input Support**: Pass input to commands via stdin
 * **Comprehensive Output**: Returns stdout, stderr, exit status, and execution time
 * **Shell Operator Safety**: Validates commands after shell operators (; , &&, ||, |)
+* **Timeout Control**: Set maximum execution time for commands
 
 ## MCP client setting in your Claude.app
 
@@ -71,6 +72,19 @@ ALLOW_COMMANDS="ls,  cat  , echo"     # Multiple spaces
     "command": ["cat"],
     "stdin": "Hello, World!"
 }
+
+# Command with timeout
+{
+    "command": ["long-running-process"],
+    "timeout": 30  # Maximum execution time in seconds
+}
+
+# Command with working directory and timeout
+{
+    "command": ["grep", "-r", "pattern"],
+    "directory": "/path/to/search",
+    "timeout": 60
+}
 ```
 
 ### Response Format
@@ -133,10 +147,12 @@ pytest
 
 ### Request Arguments
 
-| Field    | Type       | Required | Description                                   |
-|----------|------------|----------|-----------------------------------------------|
-| command  | string[]   | Yes      | Command and its arguments as array elements   |
-| stdin    | string     | No       | Input to be passed to the command            |
+| Field     | Type       | Required | Description                                   |
+|-----------|------------|----------|-----------------------------------------------|
+| command   | string[]   | Yes      | Command and its arguments as array elements   |
+| stdin     | string     | No       | Input to be passed to the command            |
+| directory | string     | No       | Working directory for command execution       |
+| timeout   | integer    | No       | Maximum execution time in seconds             |
 
 ### Response Fields
 

mcp_shell_server/server.py

@@ -49,6 +49,11 @@ def get_tool_description(self) -> Tool:
                         "type": "string",
                         "description": "Working directory where the command will be executed",
                     },
+                    "timeout": {
+                        "type": "integer",
+                        "description": "Maximum execution time in seconds",
+                        "minimum": 0,
+                    },
                 },
                 "required": ["command"],
             },
@@ -59,11 +64,12 @@ async def run_tool(self, arguments: dict) -> Sequence[TextContent]:
         command = arguments.get("command", [])
         stdin = arguments.get("stdin")
         directory = arguments.get("directory")
+        timeout = arguments.get("timeout")
 
         if not command:
             raise ValueError("No command provided")
 
-        result = await self.executor.execute(command, stdin, directory)
+        result = await self.executor.execute(command, stdin, directory, timeout)
 
         # Raise error if command execution failed
         if result.get("error"):

mcp_shell_server/shell_executor.py

@@ -103,6 +103,7 @@ async def execute(
         command: List[str],
         stdin: Optional[str] = None,
         directory: Optional[str] = None,
+        timeout: Optional[int] = None,
     ) -> Dict[str, Any]:
         """
         Execute a shell command with optional stdin input and working directory.
@@ -111,6 +112,7 @@ async def execute(
             command (List[str]): Command and its arguments
             stdin (Optional[str]): Input to be passed to the command via stdin
             directory (Optional[str]): Working directory for command execution
+            timeout (Optional[int]): Timeout for command execution in seconds
 
         Returns:
             Dict[str, Any]: Execution result containing stdout, stderr, status code, and execution time.
@@ -146,17 +148,36 @@ async def execute(
                 cwd=directory,  # Set working directory if specified
             )
 
-            stdin_bytes = stdin.encode() if stdin else None
-            stdout, stderr = await process.communicate(input=stdin_bytes)
+            try:
+                stdin_bytes = stdin.encode() if stdin else None
+                stdout, stderr = await asyncio.wait_for(
+                    process.communicate(input=stdin_bytes), timeout=timeout
+                )
+
+                return {
+                    "error": None,
+                    "stdout": stdout.decode() if stdout else "",
+                    "stderr": stderr.decode() if stderr else "",
+                    "status": process.returncode,
+                    "execution_time": time.time() - start_time,
+                    "directory": directory,
+                }
+
+            except asyncio.TimeoutError:
+                try:
+                    process.kill()
+                    await process.wait()
+                except ProcessLookupError:
+                    pass
+
+                return {
+                    "error": f"Command timed out after {timeout} seconds",
+                    "status": -1,
+                    "stdout": "",
+                    "stderr": f"Command timed out after {timeout} seconds",
+                    "execution_time": time.time() - start_time,
+                }
 
-            return {
-                "error": None,  # Set error field to None for success case
-                "stdout": stdout.decode() if stdout else "",
-                "stderr": stderr.decode() if stderr else "",
-                "status": process.returncode,
-                "execution_time": time.time() - start_time,
-                "directory": directory,  # Include working directory in response
-            }
         except FileNotFoundError:
             return {
                 "error": f"Command not found: {cleaned_command[0]}",

tests/test_server.py

@@ -168,3 +168,20 @@ async def test_call_tool_with_nested_directory(temp_test_dir, monkeypatch):
     )
     assert isinstance(result[0], TextContent)
     assert result[0].text.strip() == nested_real_path
+
+
+@pytest.mark.asyncio
+async def test_call_tool_with_timeout(monkeypatch):
+    """Test command execution with timeout"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "sleep")
+    with pytest.raises(RuntimeError) as excinfo:
+        await call_tool("shell_execute", {"command": ["sleep", "2"], "timeout": 1})
+    assert "Command timed out after 1 seconds" in str(excinfo.value)
+
+
+@pytest.mark.asyncio
+async def test_call_tool_completes_within_timeout(monkeypatch):
+    """Test command that completes within timeout period"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "sleep")
+    result = await call_tool("shell_execute", {"command": ["sleep", "1"], "timeout": 2})
+    assert len(result) == 0  # sleep command produces no output

tests/test_shell_executor.py

@@ -163,3 +163,24 @@ async def test_execute_with_nested_directory(executor, temp_test_dir, monkeypatc
     assert result["error"] is None
     assert result["status"] == 0
     assert result["stdout"].strip() == nested_real_path
+
+
+@pytest.mark.asyncio
+async def test_command_timeout(executor, monkeypatch):
+    """Test command timeout functionality"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "sleep")
+    result = await executor.execute(["sleep", "2"], timeout=1)
+    assert result["error"] == "Command timed out after 1 seconds"
+    assert result["status"] == -1
+    assert result["stdout"] == ""
+    assert result["stderr"] == "Command timed out after 1 seconds"
+
+
+@pytest.mark.asyncio
+async def test_command_completes_within_timeout(executor, monkeypatch):
+    """Test command that completes within timeout period"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "sleep")
+    result = await executor.execute(["sleep", "1"], timeout=2)
+    assert result["error"] is None
+    assert result["status"] == 0
+    assert result["stdout"] == ""