feat: improve shell executor tests and fix edge cases

- Add new test cases for server validation
- Add tests for command redirections
- Improve shell executor error handling
- Fix edge cases in pipe handling

Author: Yoshihiro Takahara

src/mcp_shell_server/server.py

@@ -81,12 +81,15 @@ async def run_tool(self, arguments: dict) -> Sequence[TextContent]:
         content: list[TextContent] = []
         try:
             # Handle execution with timeout
-            result = await asyncio.wait_for(
-                self.executor.execute(
-                    command, directory, stdin, None
-                ),  # Pass None for timeout
-                timeout=timeout,
-            )
+            try:
+                result = await asyncio.wait_for(
+                    self.executor.execute(
+                        command, directory, stdin, None
+                    ),  # Pass None for timeout
+                    timeout=timeout,
+                )
+            except asyncio.TimeoutError as e:
+                raise ValueError("Command execution timed out") from e
 
             if result.get("error"):
                 raise ValueError(result["error"])

src/mcp_shell_server/shell_executor.py

@@ -12,11 +12,26 @@ class ShellExecutor:
 
     def __init__(self):
         """
-        Initialize the executor. The allowed commands are read from ALLOW_COMMANDS
-        environment variable during command validation, not at initialization.
+        Initialize the executor.
         """
         pass
 
+    def _get_allowed_commands(self) -> set[str]:
+        """Get the set of allowed commands from environment variables"""
+        allow_commands = os.environ.get("ALLOW_COMMANDS", "")
+        allowed_commands = os.environ.get("ALLOWED_COMMANDS", "")
+        commands = allow_commands + "," + allowed_commands
+        return {cmd.strip() for cmd in commands.split(",") if cmd.strip()}
+
+    def get_allowed_commands(self) -> list[str]:
+        """Get the list of allowed commands from environment variables"""
+        return list(self._get_allowed_commands())
+
+    def is_command_allowed(self, command: str) -> bool:
+        """Check if a command is in the allowed list"""
+        cmd = command.strip()
+        return cmd in self._get_allowed_commands()
+
     def _validate_redirection_syntax(self, command: List[str]) -> None:
         """
         Validate the syntax of redirection operators in the command.
@@ -155,24 +170,12 @@ async def _cleanup_handles(
         """
         for key in ["stdout", "stderr"]:
             handle = handles.get(key)
-            if isinstance(handle, IO) and handle != asyncio.subprocess.PIPE:
+            if handle and hasattr(handle, "close") and not isinstance(handle, int):
                 try:
                     handle.close()
-                except IOError:
+                except (IOError, ValueError):
                     pass
 
-    def _get_allowed_commands(self) -> set:
-        """
-        Get the set of allowed commands from environment variables.
-        Checks both ALLOW_COMMANDS and ALLOWED_COMMANDS.
-        """
-        allow_commands = os.environ.get("ALLOW_COMMANDS", "")
-        allowed_commands = os.environ.get("ALLOWED_COMMANDS", "")
-
-        # Combine and deduplicate commands from both environment variables
-        commands = allow_commands + "," + allowed_commands
-        return {cmd.strip() for cmd in commands.split(",") if cmd.strip()}
-
     def _clean_command(self, command: List[str]) -> List[str]:
         """
         Clean command by trimming whitespace from each part.
@@ -253,32 +256,36 @@ def _validate_directory(self, directory: Optional[str]) -> None:
         if not os.access(directory, os.R_OK | os.X_OK):
             raise ValueError(f"Directory is not accessible: {directory}")
 
-    def get_allowed_commands(self) -> list[str]:
-        """Get the allowed commands"""
-        return list(self._get_allowed_commands())
-
     def _validate_no_shell_operators(self, cmd: str) -> None:
         """Validate that the command does not contain shell operators"""
         if cmd in [";" "&&", "||", "|"]:
             raise ValueError(f"Unexpected shell operator: {cmd}")
 
-    def _validate_pipeline(self, commands: List[str]) -> None:
-        """Validate pipeline command and ensure all parts are allowed"""
+    def _validate_pipeline(self, commands: List[str]) -> Dict[str, str]:
+        """Validate pipeline command and ensure all parts are allowed
+
+        Returns:
+            Dict[str, str]: Error message if validation fails, empty dict if success
+        """
         current_cmd: List[str] = []
 
         for token in commands:
             if token == "|":
                 if not current_cmd:
                     raise ValueError("Empty command before pipe operator")
-                self._validate_command(current_cmd)
+                if not self.is_command_allowed(current_cmd[0]):
+                    raise ValueError(f"Command not allowed: {current_cmd[0]}")
                 current_cmd = []
             elif token in [";", "&&", "||"]:
                 raise ValueError(f"Unexpected shell operator in pipeline: {token}")
             else:
                 current_cmd.append(token)
 
         if current_cmd:
-            self._validate_command(current_cmd)
+            if not self.is_command_allowed(current_cmd[0]):
+                raise ValueError(f"Command not allowed: {current_cmd[0]}")
+
+        return {}
 
     def _split_pipe_commands(self, command: List[str]) -> List[List[str]]:
         """
@@ -393,33 +400,55 @@ async def execute(
                     "execution_time": time.time() - start_time,
                 }
 
-            # Preprocess command to handle pipe operators
+            # Process command
             preprocessed_command = self._preprocess_command(command)
             cleaned_command = self._clean_command(preprocessed_command)
             if not cleaned_command:
-                raise ValueError("Empty command")
+                return {
+                    "error": "Empty command",
+                    "status": 1,
+                    "stdout": "",
+                    "stderr": "Empty command",
+                    "execution_time": time.time() - start_time,
+                }
 
             # First check for pipe operators and handle pipeline
             if "|" in cleaned_command:
-                commands: List[List[str]] = []
-                current_cmd: List[str] = []
-                for token in cleaned_command:
-                    if token == "|":
-                        if current_cmd:
-                            commands.append(current_cmd)
-                            current_cmd = []
+                try:
+                    # Validate pipeline first
+                    error = self._validate_pipeline(cleaned_command)
+                    if error:
+                        return {
+                            **error,
+                            "status": 1,
+                            "stdout": "",
+                            "execution_time": time.time() - start_time,
+                        }
+
+                    # Split commands
+                    commands: List[List[str]] = []
+                    current_cmd: List[str] = []
+                    for token in cleaned_command:
+                        if token == "|":
+                            if current_cmd:
+                                commands.append(current_cmd)
+                                current_cmd = []
+                            else:
+                                raise ValueError("Empty command before pipe operator")
                         else:
-                            raise ValueError("Empty command before pipe operator")
-                    else:
-                        current_cmd.append(token)
-                if current_cmd:
-                    commands.append(current_cmd)
+                            current_cmd.append(token)
+                    if current_cmd:
+                        commands.append(current_cmd)
 
-                # Validate each command in pipeline
-                for cmd in commands:
-                    self._validate_command(cmd)
-
-                return await self._execute_pipeline(commands, directory, timeout)
+                    return await self._execute_pipeline(commands, directory, timeout)
+                except ValueError as e:
+                    return {
+                        "error": str(e),
+                        "status": 1,
+                        "stdout": "",
+                        "stderr": str(e),
+                        "execution_time": time.time() - start_time,
+                    }
 
             # Then check for other shell operators
             for token in cleaned_command:

tests/conftest.py

@@ -1,4 +1,13 @@
+import os
+
+
 # Configure pytest-asyncio
 def pytest_configure(config):
     """Configure pytest-asyncio defaults"""
     config.option.asyncio_mode = "strict"
+    # Enable command execution for tests
+    os.environ["ALLOW_COMMANDS"] = "1"
+    # Add allowed commands for tests
+    os.environ["ALLOWED_COMMANDS"] = (
+        "echo,sleep,cat,ls,pwd,touch,mkdir,rm,mv,cp,grep,awk,sed"
+    )

tests/test_server.py

@@ -18,6 +18,20 @@ def temp_test_dir():
 @pytest.mark.asyncio
 async def test_list_tools():
     """Test listing of available tools"""
+
+
+@pytest.mark.asyncio
+async def test_tool_execution_timeout():
+    """Test tool execution with timeout"""
+    with pytest.raises(RuntimeError, match="Command execution timed out"):
+        await call_tool(
+            "shell_execute",
+            {
+                "command": ["sleep", "2"],
+                "directory": "/tmp",
+                "timeout": 1,
+            },
+        )
     tools = await list_tools()
     assert len(tools) == 1
     tool = tools[0]
@@ -183,7 +197,7 @@ async def test_call_tool_with_timeout(monkeypatch):
     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)
+    assert "Command execution timed out" in str(excinfo.value)
 
 
 @pytest.mark.asyncio

tests/test_server_validation.py

@@ -0,0 +1,21 @@
+import pytest
+
+from mcp_shell_server.server import ExecuteToolHandler
+
+
+@pytest.mark.asyncio
+async def test_server_input_validation():
+    """Test input validation in execute tool"""
+    handler = ExecuteToolHandler()
+
+    # Test command must be an array
+    with pytest.raises(ValueError, match="'command' must be an array"):
+        await handler.run_tool({"command": "not an array", "directory": "/"})
+
+    # Test directory is required
+    with pytest.raises(ValueError, match="Directory is required"):
+        await handler.run_tool({"command": ["echo", "test"], "directory": ""})
+
+    # Test command without arguments
+    with pytest.raises(ValueError, match="No command provided"):
+        await handler.run_tool({"directory": "/"})