Merge branch 'feat/support-custom-env' into develop

Author: tumf

src/mcp_shell_server/shell_executor.py

@@ -392,6 +392,7 @@ async def execute(
         directory: str,
         stdin: Optional[str] = None,
         timeout: Optional[int] = None,
+        envs: Optional[Dict[str, str]] = None,
     ) -> Dict[str, Any]:
         start_time = time.time()
         process = None  # Initialize process variable
@@ -449,7 +450,9 @@ async def execute(
                     if current_cmd:
                         commands.append(current_cmd)
 
-                    return await self._execute_pipeline(commands, directory, timeout)
+                    return await self._execute_pipeline(
+                        commands, directory, timeout, envs
+                    )
                 except ValueError as e:
                     return {
                         "error": str(e),
@@ -535,7 +538,7 @@ async def execute(
                 stdin=asyncio.subprocess.PIPE if stdin else None,
                 stdout=stdout_handle,
                 stderr=asyncio.subprocess.PIPE,
-                env=os.environ,  # Use all environment variables
+                env={**os.environ, **(envs or {})},  # Merge environment variables
                 cwd=directory,
             )
 
@@ -613,6 +616,7 @@ async def _execute_pipeline(
         commands: List[List[str]],
         directory: Optional[str] = None,
         timeout: Optional[int] = None,
+        envs: Optional[Dict[str, str]] = None,
     ) -> Dict[str, Any]:
         start_time = time.time()
         processes: List[asyncio.subprocess.Process] = []
@@ -663,7 +667,7 @@ async def _execute_pipeline(
                     stdin=asyncio.subprocess.PIPE if prev_stdout is not None else None,
                     stdout=asyncio.subprocess.PIPE,
                     stderr=asyncio.subprocess.PIPE,
-                    env=os.environ,  # Use all environment variables
+                    env={**os.environ, **(envs or {})},  # Merge environment variables
                     cwd=directory,
                 )
 

tests/test_shell_executor.py

@@ -621,4 +621,53 @@ async def test_output_redirection_with_append(executor, temp_test_dir, monkeypat
     lines = result["stdout"].strip().split("\n")
     assert len(lines) == 2
     assert lines[0] == "hello"
-    assert lines[1] == "world"
+
+
+@pytest.mark.asyncio
+async def test_execute_with_custom_env(executor, temp_test_dir, monkeypatch):
+    """Test command execution with custom environment variables"""
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOW_COMMANDS", "env,printenv")
+
+    custom_env = {"TEST_VAR1": "test_value1", "TEST_VAR2": "test_value2"}
+
+    # Test env command
+    result = await executor.execute(["env"], directory=temp_test_dir, envs=custom_env)
+    assert "TEST_VAR1=test_value1" in result["stdout"]
+    assert "TEST_VAR2=test_value2" in result["stdout"]
+
+    # Test specific variable
+    result = await executor.execute(
+        ["printenv", "TEST_VAR1"], directory=temp_test_dir, envs=custom_env
+    )
+    assert result["stdout"].strip() == "test_value1"
+
+
+@pytest.mark.asyncio
+async def test_execute_env_override(executor, temp_test_dir, monkeypatch):
+    """Test that custom environment variables override system variables"""
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOW_COMMANDS", "env")
+    monkeypatch.setenv("TEST_VAR", "original_value")
+
+    # Override system environment variable
+    result = await executor.execute(
+        ["env"], directory=temp_test_dir, envs={"TEST_VAR": "new_value"}
+    )
+
+    assert "TEST_VAR=new_value" in result["stdout"]
+    assert "TEST_VAR=original_value" not in result["stdout"]
+
+
+@pytest.mark.asyncio
+async def test_execute_with_empty_env(executor, temp_test_dir, monkeypatch):
+    """Test command execution with empty environment variables"""
+    clear_env(monkeypatch)
+    monkeypatch.setenv("ALLOW_COMMANDS", "env")
+
+    result = await executor.execute(["env"], directory=temp_test_dir, envs={})
+
+    # Command should still work with system environment
+    assert result["error"] is None
+    assert result["status"] == 0
+    assert len(result["stdout"]) > 0