Merge branch 'develop'

Author: tumf

src/mcp_shell_server/shell_executor.py

@@ -1,5 +1,6 @@
 import asyncio
 import os
+import pwd
 import shlex
 import time
 from typing import IO, Any, Dict, List, Optional, Tuple, Union
@@ -14,7 +15,7 @@ def __init__(self):
         """
         Initialize the executor.
         """
-        pass
+        pass  # pragma: no cover
 
     def _get_allowed_commands(self) -> set[str]:
         """Get the set of allowed commands from environment variables"""
@@ -378,6 +379,13 @@ def _preprocess_command(self, command: List[str]) -> List[str]:
                 preprocessed_command.append(token)
         return preprocessed_command
 
+    def _get_default_shell(self) -> str:
+        """Get the login shell of the current user"""
+        try:
+            return pwd.getpwuid(os.getuid()).pw_shell
+        except (ImportError, KeyError):
+            return os.environ.get("SHELL", "/bin/sh")
+
     async def execute(
         self,
         command: List[str],
@@ -517,14 +525,17 @@ async def execute(
                 except IOError as e:
                     raise ValueError(f"Failed to open output file: {e}") from e
 
-            # Execute the command
+            # Execute the command with interactive shell
+            shell = self._get_default_shell()
             shell_cmd = self._create_shell_command(cmd)
+            shell_cmd = f"{shell} -i -c {shlex.quote(shell_cmd)}"
+
             process = await asyncio.create_subprocess_shell(
                 shell_cmd,
                 stdin=asyncio.subprocess.PIPE if stdin else None,
                 stdout=stdout_handle,
                 stderr=asyncio.subprocess.PIPE,
-                env={"PATH": os.environ.get("PATH", "")},
+                env=os.environ,  # Use all environment variables
                 cwd=directory,
             )
 
@@ -642,12 +653,17 @@ async def _execute_pipeline(
             for i, cmd in enumerate(parsed_commands):
                 shell_cmd = self._create_shell_command(cmd)
 
+                # Get default shell for the first command and set interactive mode
+                if i == 0:
+                    shell = self._get_default_shell()
+                    shell_cmd = f"{shell} -i -c {shlex.quote(shell_cmd)}"
+
                 process = await asyncio.create_subprocess_shell(
                     shell_cmd,
                     stdin=asyncio.subprocess.PIPE if prev_stdout is not None else None,
                     stdout=asyncio.subprocess.PIPE,
                     stderr=asyncio.subprocess.PIPE,
-                    env={"PATH": os.environ.get("PATH", "")},
+                    env=os.environ,  # Use all environment variables
                     cwd=directory,
                 )
 

tests/test_server.py

@@ -322,3 +322,26 @@ def stdio_server_impl():
         await main()
 
     assert str(exc.value) == "Test error"
+
+
+@pytest.mark.asyncio
+async def test_shell_startup(monkeypatch, temp_test_dir):
+    """Test shell startup and environment"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "ps")
+    result = await call_tool(
+        "shell_execute",
+        {"command": ["ps", "-p", "$$", "-o", "command="], "directory": temp_test_dir},
+    )
+    assert len(result) == 1
+    assert result[0].type == "text"
+
+
+@pytest.mark.asyncio
+async def test_environment_variables(monkeypatch, temp_test_dir):
+    """Test to check environment variables during test execution"""
+    monkeypatch.setenv("ALLOW_COMMANDS", "env")
+    result = await call_tool(
+        "shell_execute",
+        {"command": ["env"], "directory": temp_test_dir},
+    )
+    assert len(result) == 1

tests/test_shell_executor.py

@@ -600,3 +600,45 @@ def test_preprocess_command_pipeline(executor):
         "grep",
         "pattern",
     ]
+
+
+@pytest.mark.asyncio
+async def test_command_cleanup_on_error(executor, temp_test_dir, monkeypatch):
+    """Test cleanup of processes when error occurs"""
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOW_COMMANDS", "sleep")
+
+    async def execute_with_keyboard_interrupt():
+        # Simulate keyboard interrupt during execution
+        result = await executor.execute(["sleep", "5"], temp_test_dir, timeout=1)
+        return result
+
+    result = await execute_with_keyboard_interrupt()
+    assert result["error"] == "Command timed out after 1 seconds"
+    assert result["status"] == -1
+    assert "execution_time" in result
+
+
+@pytest.mark.asyncio
+async def test_output_redirection_with_append(executor, temp_test_dir, monkeypatch):
+    """Test output redirection with append mode"""
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOW_COMMANDS", "echo,cat")
+    output_file = os.path.join(temp_test_dir, "test.txt")
+
+    # Write initial content
+    await executor.execute(["echo", "hello", ">", output_file], directory=temp_test_dir)
+
+    # Append content
+    result = await executor.execute(
+        ["echo", "world", ">>", output_file], directory=temp_test_dir
+    )
+    assert result["error"] is None
+    assert result["status"] == 0
+
+    # Verify contents
+    result = await executor.execute(["cat", output_file], directory=temp_test_dir)
+    lines = result["stdout"].strip().split("\n")
+    assert len(lines) == 2
+    assert lines[0] == "hello"
+    assert lines[1] == "world"

tests/test_shell_executor_pipeline.py

@@ -1,13 +1,34 @@
+"""Test pipeline execution and cleanup scenarios."""
+
+import os
+import tempfile
+
 import pytest
 
 from mcp_shell_server.shell_executor import ShellExecutor
 
 
+def clear_env(monkeypatch):
+    monkeypatch.delenv("ALLOW_COMMANDS", raising=False)
+    monkeypatch.delenv("ALLOWED_COMMANDS", raising=False)
+
+
+@pytest.fixture
+def executor():
+    return ShellExecutor()
+
+
+@pytest.fixture
+def temp_test_dir():
+    """Create a temporary directory for testing"""
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        # Return the real path to handle macOS /private/tmp symlink
+        yield os.path.realpath(tmpdirname)
+
+
 @pytest.mark.asyncio
-async def test_pipeline_split():
+async def test_pipeline_split(executor):
     """Test pipeline command splitting functionality"""
-    executor = ShellExecutor()
-
     # Test basic pipe command
     commands = executor._split_pipe_commands(["echo", "hello", "|", "grep", "h"])
     assert len(commands) == 2
@@ -35,19 +56,43 @@ async def test_pipeline_split():
 
 
 @pytest.mark.asyncio
-async def test_pipeline_execution_success():
+async def test_pipeline_execution_success(executor, temp_test_dir, monkeypatch):
     """Test successful pipeline execution with proper return value"""
-    executor = ShellExecutor()
-    import os
-
-    os.environ["ALLOWED_COMMANDS"] = "echo,grep"
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOWED_COMMANDS", "echo,grep")
 
     result = await executor.execute(
-        ["echo", "hello world", "|", "grep", "world"], directory="/tmp", timeout=5
+        ["echo", "hello world", "|", "grep", "world"],
+        directory=temp_test_dir,
+        timeout=5,
     )
 
     assert result["error"] is None
     assert result["status"] == 0
     assert "world" in result["stdout"]
     assert "execution_time" in result
-    assert result["directory"] == "/tmp"
+
+
+@pytest.mark.asyncio
+async def test_pipeline_cleanup_and_timeouts(executor, temp_test_dir, monkeypatch):
+    """Test cleanup of processes in pipelines and timeout handling"""
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOW_COMMANDS", "cat,tr,head,sleep")
+
+    # Test pipeline with early termination
+    test_file = os.path.join(temp_test_dir, "test.txt")
+    with open(test_file, "w") as f:
+        f.write("test\n" * 1000)
+
+    result = await executor.execute(
+        ["cat", test_file, "|", "tr", "[:lower:]", "[:upper:]", "|", "head", "-n", "1"],
+        temp_test_dir,
+        timeout=2,
+    )
+    assert result["status"] == 0
+    assert result["stdout"].strip() == "TEST"
+
+    # Test timeout handling in pipeline
+    result = await executor.execute(["sleep", "5"], temp_test_dir, timeout=1)
+    assert result["status"] == -1
+    assert "timed out" in result["error"].lower()  # タイムアウトエラーの確認