feat: update local code executor to support powershell (#5884)

To support powershell on the local code executor.
Closes #5518

.github/workflows/checks.yml

@@ -181,6 +181,44 @@ jobs:
           name: coverage-autogen-ext-grpc
           path: ./python/coverage_autogen-ext-grpc.xml
 
+  test-autogen-ext-pwsh:
+    runs-on: windows-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+          version: "0.5.18"
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install Python deps
+        run: |
+          uv sync --locked --all-extras
+        shell: pwsh
+        working-directory: ./python
+
+      - name: Run tests for Windows
+        run: |
+          .venv/Scripts/activate.ps1
+          poe --directory ./packages/autogen-ext test-windows
+        shell: pwsh
+        working-directory: ./python
+
+      - name: Move coverage file
+        run: |
+          mv ./packages/autogen-ext/coverage.xml coverage_autogen_ext_windows.xml
+        working-directory: ./python
+
+      - name: Upload coverage artifact
+        uses: actions/upload-artifact@v4
+        with:
+          name: coverage-autogen-ext-windows
+          path: ./python/coverage_autogen_ext_windows.xml
+
   codecov:
     runs-on: ubuntu-latest
     needs: [test, test-grpc]

python/packages/autogen-ext/pyproject.toml

@@ -168,6 +168,7 @@ test.sequence = [
 ]
 test.default_item_type = "cmd"
 test-grpc = "pytest -n 1 --cov=src --cov-report=term-missing --cov-report=xml --grpc"
+test-windows = "pytest -n 1 --cov=src --cov-report=term-missing --cov-report=xml -m 'windows'"
 mypy = "mypy --config-file ../../pyproject.toml --exclude src/autogen_ext/runtimes/grpc/protos --exclude tests/protos src tests"
 
 [tool.mypy]

python/packages/autogen-ext/src/autogen_ext/code_executors/_common.py

@@ -158,6 +158,8 @@ def lang_to_cmd(lang: str) -> str:
         return lang
     if lang in ["shell"]:
         return "sh"
+    if lang in ["pwsh", "powershell", "ps1"]:
+        return "pwsh"
     else:
         raise ValueError(f"Unsupported language: {lang}")
 

python/packages/autogen-ext/src/autogen_ext/code_executors/local/__init__.py

@@ -302,26 +302,33 @@ $functions"""
     async def _execute_code_dont_check_setup(
         self, code_blocks: List[CodeBlock], cancellation_token: CancellationToken
     ) -> CommandLineCodeResult:
+        """
+        Execute the provided code blocks in the local command line without re-checking setup.
+        Returns a CommandLineCodeResult indicating success or failure.
+        """
         logs_all: str = ""
         file_names: List[Path] = []
         exitcode = 0
+
         for code_block in code_blocks:
             lang, code = code_block.language, code_block.code
             lang = lang.lower()
 
+            # Remove pip output where possible
             code = silence_pip(code, lang)
 
+            # Normalize python variants to "python"
             if lang in PYTHON_VARIANTS:
                 lang = "python"
 
+            # Abort if not supported
             if lang not in self.SUPPORTED_LANGUAGES:
-                # In case the language is not supported, we return an error message.
                 exitcode = 1
                 logs_all += "\n" + f"unknown language {lang}"
                 break
 
+            # Try extracting a filename (if present)
             try:
-                # Check if there is a filename comment
                 filename = get_file_name_from_content(code, self._work_dir)
             except ValueError:
                 return CommandLineCodeResult(
@@ -330,32 +337,57 @@ $functions"""
                     code_file=None,
                 )
 
+            # If no filename is found, create one
             if filename is None:
-                # create a file with an automatically generated name
                 code_hash = sha256(code.encode()).hexdigest()
-                filename = f"tmp_code_{code_hash}.{'py' if lang.startswith('python') else lang}"
+                if lang.startswith("python"):
+                    ext = "py"
+                elif lang in ["pwsh", "powershell", "ps1"]:
+                    ext = "ps1"
+                else:
+                    ext = lang
+
+                filename = f"tmp_code_{code_hash}.{ext}"
 
             written_file = (self._work_dir / filename).resolve()
             with written_file.open("w", encoding="utf-8") as f:
                 f.write(code)
             file_names.append(written_file)
 
+            # Build environment
             env = os.environ.copy()
-
             if self._virtual_env_context:
-                virtual_env_exe_abs_path = os.path.abspath(self._virtual_env_context.env_exe)
                 virtual_env_bin_abs_path = os.path.abspath(self._virtual_env_context.bin_path)
                 env["PATH"] = f"{virtual_env_bin_abs_path}{os.pathsep}{env['PATH']}"
 
-                program = virtual_env_exe_abs_path if lang.startswith("python") else lang_to_cmd(lang)
+            # Decide how to invoke the script
+            if lang == "python":
+                program = (
+                    os.path.abspath(self._virtual_env_context.env_exe) if self._virtual_env_context else sys.executable
+                )
+                extra_args = [str(written_file.absolute())]
             else:
-                program = sys.executable if lang.startswith("python") else lang_to_cmd(lang)
+                # Get the appropriate command for the language
+                program = lang_to_cmd(lang)
 
-            # Wrap in a task to make it cancellable
+                # Special handling for PowerShell
+                if program == "pwsh":
+                    extra_args = [
+                        "-NoProfile",
+                        "-ExecutionPolicy",
+                        "Bypass",
+                        "-File",
+                        str(written_file.absolute()),
+                    ]
+                else:
+                    # Shell commands (bash, sh, etc.)
+                    extra_args = [str(written_file.absolute())]
+
+            # Create a subprocess and run
             task = asyncio.create_task(
                 asyncio.create_subprocess_exec(
                     program,
-                    str(written_file.absolute()),
+                    *extra_args,
                     cwd=self._work_dir,
                     stdout=asyncio.subprocess.PIPE,
                     stderr=asyncio.subprocess.PIPE,
@@ -363,31 +395,27 @@ $functions"""
                 )
             )
             cancellation_token.link_future(task)
+
             try:
                 proc = await task
                 stdout, stderr = await asyncio.wait_for(proc.communicate(), self._timeout)
                 exitcode = proc.returncode or 0
-
             except asyncio.TimeoutError:
-                logs_all += "\n Timeout"
-                # Same exit code as the timeout command on linux.
+                logs_all += "\nTimeout"
                 exitcode = 124
                 break
             except asyncio.CancelledError:
-                logs_all += "\n Cancelled"
-                # TODO: which exit code? 125 is Operation Canceled
+                logs_all += "\nCancelled"
                 exitcode = 125
                 break
 
-            self._running_cmd_task = None
-
             logs_all += stderr.decode()
             logs_all += stdout.decode()
 
             if exitcode != 0:
                 break
 
-        code_file = str(file_names[0]) if len(file_names) > 0 else None
+        code_file = str(file_names[0]) if file_names else None
         return CommandLineCodeResult(exit_code=exitcode, output=logs_all, code_file=code_file)
 
     async def restart(self) -> None:

python/packages/autogen-ext/tests/code_executors/test_commandline_code_executor.py

@@ -4,6 +4,7 @@
 import asyncio
 import os
 import shutil
+import platform
 import sys
 import tempfile
 import venv
@@ -18,6 +19,11 @@ from autogen_core.code_executor import CodeBlock
 from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor
 
 
+HAS_POWERSHELL: bool = platform.system() == "Windows" and (
+    shutil.which("powershell") is not None or shutil.which("pwsh") is not None
+)
+
+
 @pytest_asyncio.fixture(scope="function")  # type: ignore
 async def executor_and_temp_dir(
     request: pytest.FixtureRequest,
@@ -203,3 +209,25 @@ def test_serialize_deserialize() -> None:
         executor_config = executor.dump_component()
         loaded_executor = LocalCommandLineCodeExecutor.load_component(executor_config)
         assert executor.work_dir == loaded_executor.work_dir
+
+
+@pytest.mark.asyncio
+@pytest.mark.windows
+@pytest.mark.skipif(
+    not HAS_POWERSHELL,
+    reason="No PowerShell interpreter (powershell or pwsh) found on this environment.",
+)
+@pytest.mark.parametrize("executor_and_temp_dir", ["local"], indirect=True)
+async def test_ps1_script(executor_and_temp_dir: ExecutorFixture) -> None:
+    """
+    Test execution of a simple PowerShell script.
+    This test is skipped if powershell/pwsh is not installed.
+    """
+    executor, _ = executor_and_temp_dir
+    cancellation_token = CancellationToken()
+    code = 'Write-Host "hello from powershell!"'
+    code_blocks = [CodeBlock(code=code, language="powershell")]
+    result = await executor.execute_code_blocks(code_blocks, cancellation_token)
+    assert result.exit_code == 0
+    assert "hello from powershell!" in result.output
+    assert result.code_file is not None

python/packages/autogen-ext/tests/conftest.py

@@ -0,0 +1,9 @@
+import pytest
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption("--windows", action="store_true", default=False, help="Run tests for Windows")
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    config.addinivalue_line("markers", "windows: mark test as requiring Windows")