ReMe/reme/memory/file_based/tools/shell.py

# -*- coding: utf-8 -*-
# flake8: noqa: E501
# pylint: disable=line-too-long
"""The shell command tool."""

import asyncio
import locale
import subprocess
import sys
from pathlib import Path

from agentscope.message import TextBlock
from agentscope.tool import ToolResponse


def _execute_subprocess_sync(
    cmd: str,
    cwd: str,
    timeout: int,
) -> tuple[int, str, str]:
    """Execute subprocess synchronously in a thread.

    This function runs in a separate thread to avoid Windows asyncio
    subprocess limitations.

    Args:
        cmd (`str`):
            The shell command to execute.
        cwd (`str`):
            The working directory for the command execution.
        timeout (`int`):
            The maximum time (in seconds) allowed for the command to run.

    Returns:
        `tuple[int, str, str]`:
            A tuple containing the return code, standard output, and
            standard error of the executed command. If timeout occurs, the
            return code will be -1 and stderr will contain timeout information.
    """
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            capture_output=True,
            text=True,
            cwd=cwd,
            timeout=timeout,
            encoding=locale.getpreferredencoding(False) or "utf-8",
            errors="replace",
            check=True,
        )
        return (
            result.returncode,
            result.stdout.strip("\n"),
            result.stderr.strip("\n"),
        )
    except subprocess.TimeoutExpired:
        return (
            -1,
            "",
            f"Command execution exceeded the timeout of {timeout} seconds.",
        )
    except Exception as e:
        return -1, "", str(e)


class Shell:
    """Shell command execution with a configurable working directory."""

    def __init__(self, working_dir: str | Path):
        """Initialize Shell with a working directory.

        Args:
            working_dir (`str | Path`):
                The working directory for command execution.
        """
        self.working_dir = Path(working_dir)

    # pylint: disable=too-many-branches, too-many-statements
    async def execute_shell_command(
        self,
        command: str,
        timeout: int = 60,
    ) -> ToolResponse:
        """Execute given command and return the return code, standard output and
        error within <returncode></returncode>, <stdout></stdout> and
        <stderr></stderr> tags.

        Args:
            command (`str`):
                The shell command to execute.
            timeout (`int`, defaults to `60`):
                The maximum time (in seconds) allowed for the command to run.
                Default is 60 seconds.

        Returns:
            `ToolResponse`:
                The tool response containing the return code, standard output, and
                standard error of the executed command. If timeout occurs, the
                return code will be -1 and stderr will contain timeout information.
        """

        cmd = (command or "").strip()

        # Set working directory
        working_dir = self.working_dir

        try:
            if sys.platform == "win32":
                # Windows: use thread pool to avoid asyncio subprocess limitations
                returncode, stdout_str, stderr_str = await asyncio.to_thread(
                    _execute_subprocess_sync,
                    cmd,
                    str(working_dir),
                    timeout,
                )
            else:
                proc = await asyncio.create_subprocess_shell(
                    cmd,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                    bufsize=0,
                    cwd=str(working_dir),
                )

                try:
                    # Apply timeout to communicate directly; wait()+communicate()
                    # can hang if descendants keep stdout/stderr pipes open.
                    stdout, stderr = await asyncio.wait_for(
                        proc.communicate(),
                        timeout=timeout,
                    )
                    encoding = locale.getpreferredencoding(False) or "utf-8"
                    stdout_str = stdout.decode(encoding, errors="replace").strip(
                        "\n",
                    )
                    stderr_str = stderr.decode(encoding, errors="replace").strip(
                        "\n",
                    )
                    returncode = proc.returncode

                except asyncio.TimeoutError:
                    # Handle timeout
                    stderr_suffix = (
                        f"⚠️ TimeoutError: The command execution exceeded "
                        f"the timeout of {timeout} seconds. "
                        f"Please consider increasing the timeout value if this command "
                        f"requires more time to complete."
                    )
                    returncode = -1
                    try:
                        proc.terminate()
                        # Wait a bit for graceful termination
                        try:
                            await asyncio.wait_for(proc.wait(), timeout=1)
                        except asyncio.TimeoutError:
                            # Force kill if graceful termination fails
                            proc.kill()
                            await proc.wait()

                        # Avoid hanging forever while draining pipes after timeout.
                        try:
                            stdout, stderr = await asyncio.wait_for(
                                proc.communicate(),
                                timeout=1,
                            )
                        except asyncio.TimeoutError:
                            stdout, stderr = b"", b""
                        encoding = locale.getpreferredencoding(False) or "utf-8"
                        stdout_str = stdout.decode(
                            encoding,
                            errors="replace",
                        ).strip(
                            "\n",
                        )
                        stderr_str = stderr.decode(
                            encoding,
                            errors="replace",
                        ).strip(
                            "\n",
                        )
                        if stderr_str:
                            stderr_str += f"\n{stderr_suffix}"
                        else:
                            stderr_str = stderr_suffix
                    except ProcessLookupError:
                        stdout_str = ""
                        stderr_str = stderr_suffix

            # Format the response in a human-friendly way
            if returncode == 0:
                # Success case: just show the output
                if stdout_str:
                    response_text = stdout_str
                else:
                    response_text = "Command executed successfully (no output)."
            else:
                # Error case: show detailed information
                response_parts = [f"Command failed with exit code {returncode}."]
                if stdout_str:
                    response_parts.append(f"\n[stdout]\n{stdout_str}")
                if stderr_str:
                    response_parts.append(f"\n[stderr]\n{stderr_str}")
                response_text = "".join(response_parts)

            return ToolResponse(
                content=[
                    TextBlock(
                        type="text",
                        text=response_text,
                    ),
                ],
            )

        except Exception as e:
            return ToolResponse(
                content=[
                    TextBlock(
                        type="text",
                        text=f"Error: Shell command execution failed due to \n{e}",
                    ),
                ],
            )