fix: revert shell execution logic back to master implementation

Author: LIghtJUNction

astrbot/core/computer/tools/shell.py

@@ -1,45 +1,17 @@
-"""
-ExecuteShellTool - subprocess-based shell execution with per-session state.
-
-Replaces previous plumbum-based implementation with a subprocess-based,
-per-session state manager that tracks current working directory and
-per-session environment variables.
-
-Behavior:
-- Each session has its own `cwd` and `env` stored in-memory.
-- `cd` commands are interpreted and update the session `cwd`.
-  Supports constructs like `cd /path && ls` or `cd rel/path; echo hi`.
-- Foreground commands run to completion with a configurable timeout.
-- Background commands spawn a subprocess and return immediately with the pid.
-- Environment variables passed in `env` are merged with the session env.
-- Returns JSON string describing result to match existing tool contract.
-"""
-
-from __future__ import annotations
-
-import asyncio
 import json
-import os
-import shlex
-import subprocess
 from dataclasses import dataclass, field
-from typing import Any
 
+from astrbot.api import FunctionTool
 from astrbot.core.agent.run_context import ContextWrapper
-from astrbot.core.agent.tool import FunctionTool, ToolExecResult
+from astrbot.core.agent.tool import ToolExecResult
 from astrbot.core.astr_agent_context import AstrAgentContext
 
+from ..computer_client import get_booter, get_local_booter
 from .permissions import check_admin_permission
 
 
 @dataclass
 class ExecuteShellTool(FunctionTool):
-    """
-    Stateful shell execution tool using subprocess.
-
-    Each agent session keeps its own working directory and environment mapping.
-    """
-
     name: str = "astrbot_execute_shell"
     description: str = "Execute a command in the shell."
     parameters: dict = field(
@@ -48,7 +20,7 @@ class ExecuteShellTool(FunctionTool):
             "properties": {
                 "command": {
                     "type": "string",
-                    "description": "The shell command to execute in the current runtime shell (for example, cmd.exe on Windows). Equivalent to running 'cd {working_dir} && {your_command}'.",
+                    "description": "The shell command to execute in the current runtime shell (for example, cmd.exe on Windows). Equal to 'cd {working_dir} && {your_command}'.",
                 },
                 "background": {
                     "type": "boolean",
@@ -57,7 +29,7 @@ class ExecuteShellTool(FunctionTool):
                 },
                 "env": {
                     "type": "object",
-                    "description": "Optional environment variables to set for the command (merged with session env).",
+                    "description": "Optional environment variables to set for the file creation process.",
                     "additionalProperties": {"type": "string"},
                     "default": {},
                 },
@@ -67,226 +39,26 @@ class ExecuteShellTool(FunctionTool):
     )
 
     is_local: bool = False
-    is_stateful: bool = True
-    # session_id -> {"cwd": str, "env": dict}
-    _sessions: dict[str, dict[str, Any]] = field(
-        default_factory=dict, init=False, repr=False
-    )
-
-    def _get_session_state(self, session_id: str) -> dict[str, Any]:
-        """
-        Initialize or return the per-session state.
-        State contains:
-          - cwd: current working directory for session
-          - env: environment variables dict for session
-        """
-        if session_id not in self._sessions:
-            # start from current process cwd and a copy of os.environ
-            self._sessions[session_id] = {
-                "cwd": os.getcwd(),
-                "env": dict(os.environ),
-            }
-        return self._sessions[session_id]
-
-    def _get_framework_session_state(
-        self, context: ContextWrapper[AstrAgentContext], session_id: str
-    ) -> dict[str, Any]:
-        """
-        Get session state via the framework's ToolSessionManager if available.
-        Falls back to the tool's own _sessions if session_manager is not set.
-        """
-        session_mgr = getattr(context, "session_manager", None)
-        if session_mgr is None:
-            return self._get_session_state(session_id)
-        return session_mgr.get_state(session_id, self.name)
 
     async def call(
-        self, context: ContextWrapper[AstrAgentContext], **kwargs: Any
+        self,
+        context: ContextWrapper[AstrAgentContext],
+        command: str,
+        background: bool = False,
+        env: dict = {},
     ) -> ToolExecResult:
-        """
-        Execute a shell command for the session.
-
-        Parameters are accepted via kwargs for compatibility with FunctionTool.call:
-        - command (str): the shell command to execute
-        - background (bool): whether to run in background
-        - env (dict): environment variables to merge for this execution
-        """
-        # Use the context directly - already typed as ContextWrapper[AstrAgentContext]
-        astr_ctx = context
-
-        # Permission check (use the cast wrapper)
-        if permission_error := check_admin_permission(astr_ctx, "Shell execution"):
+        if permission_error := check_admin_permission(context, "Shell execution"):
             return permission_error
 
-        # Extract parameters with defaults for backward compatibility
-        command: str = kwargs.get("command", "")
-        background: bool = bool(kwargs.get("background", False))
-        env: dict | None = kwargs.get("env")
-
-        # Resolve session id and session state (use the cast wrapper)
-        session_id = astr_ctx.context.event.unified_msg_origin
-        # Use framework ToolSessionManager if available, otherwise fall back
-        if astr_ctx.session_manager is not None:
-            state = self._get_framework_session_state(astr_ctx, session_id)
+        if self.is_local:
+            sb = get_local_booter()
         else:
-            state = self._get_session_state(session_id)
-        session_cwd = state.get("cwd", os.getcwd())
-        session_env = state.get("env", dict(os.environ)).copy()
-
-        # Merge provided env into execution env (do not mutate saved session env)
-        if env:
-            exec_env = session_env.copy()
-            exec_env.update({k: str(v) for k, v in env.items()})
-        else:
-            exec_env = session_env
-
-        # Determine timeout from config (fall back to 30)
-        config = astr_ctx.context.context.get_config(umo=session_id)
-        provider_settings: dict = {}
-        if isinstance(config, dict):
-            provider_settings = config.get("provider_settings") or {}
-        try:
-            timeout = int(provider_settings.get("tool_call_timeout", 30))
-        except (ValueError, TypeError):
-            timeout = 30
-
-        # Single atomic try block for overall execution to satisfy anti-nested-try rule.
-        try:
-            # Quick handling for explicit `cd` constructs that should change session cwd.
-            # We support leading cd followed by && or ;: e.g. "cd dir && ls", "cd dir; ls"
-            cmd_str = command.strip()
-
-            # Helper to split by shell '&&' or ';' while preserving remainder.
-            remainder_cmd = ""
-            cd_handled = False
-            # Handle forms like: cd <path> && rest OR cd <path>; rest
-            for sep in ("&&", ";"):
-                if sep in cmd_str:
-                    left, right = cmd_str.split(sep, 1)
-                    left_strip = left.strip()
-                    if left_strip.startswith("cd"):
-                        remainder_cmd = right.strip()
-                        cd_part = left_strip
-                        cd_handled = True
-                        break
-            else:
-                # No separator case, but single 'cd' command or just 'cd /path'
-                if cmd_str.startswith("cd"):
-                    cd_part = cmd_str
-                    remainder_cmd = ""
-                    cd_handled = True
-
-            if cd_handled:
-                # parse cd argument
-                parts = shlex.split(cd_part)
-                # cd with no args -> home
-                if len(parts) == 1:
-                    target = await asyncio.to_thread(os.path.expanduser, "~")
-                else:
-                    target_raw = parts[1]
-                    # expand ~ and variables
-                    target_raw = await asyncio.to_thread(os.path.expanduser, target_raw)
-                    if await asyncio.to_thread(os.path.isabs, target_raw):
-                        target = target_raw
-                    else:
-                        target = await asyncio.to_thread(
-                            os.path.normpath, os.path.join(session_cwd, target_raw)
-                        )
-
-                target_exists = await asyncio.to_thread(os.path.exists, target)
-                target_isdir = await asyncio.to_thread(os.path.isdir, target)
-                if not target_exists or not target_isdir:
-                    result = {
-                        "success": False,
-                        "exit_code": -1,
-                        "stdout": "",
-                        "stderr": f"cd: no such directory: {target}",
-                        "cwd": session_cwd,
-                    }
-                    return json.dumps(result)
-
-                # Update session cwd permanently
-                state["cwd"] = target
-                session_cwd = target
-
-                # If there is no remaining command, just return success and new cwd
-                if not remainder_cmd:
-                    result = {
-                        "success": True,
-                        "exit_code": 0,
-                        "stdout": "",
-                        "stderr": "",
-                        "cwd": session_cwd,
-                    }
-                    return json.dumps(result)
-
-                # Otherwise we'll execute the remainder using the updated cwd
-                # Use the remainder command as the command to run below
-                command_to_run = remainder_cmd
-            else:
-                command_to_run = cmd_str
-
-            # Background execution: spawn process and return pid immediately.
-            if background:
-                # Start background process; do not wait. Use shell to support pipes/redirects.
-                popen = await asyncio.to_thread(
-                    subprocess.Popen,
-                    ["/bin/sh", "-c", command_to_run],
-                    cwd=session_cwd,
-                    env=exec_env,
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.PIPE,
-                )
-                result = {
-                    "success": True,
-                    "background": True,
-                    "pid": popen.pid,
-                    "cwd": session_cwd,
-                }
-                return json.dumps(result)
-
-            # Foreground execution: run to completion, capture output.
-            completed = await asyncio.to_thread(
-                subprocess.run,
-                ["/bin/sh", "-c", command_to_run],
-                cwd=session_cwd,
-                env=exec_env,
-                timeout=timeout,
-                capture_output=True,
-                text=True,
+            sb = await get_booter(
+                context.context.context,
+                context.context.event.unified_msg_origin,
             )
-
-            exit_code = completed.returncode
-            stdout = completed.stdout if completed.stdout is not None else ""
-            stderr = completed.stderr if completed.stderr is not None else ""
-
-            result = {
-                "success": exit_code == 0,
-                "exit_code": exit_code,
-                "stdout": stdout,
-                "stderr": stderr,
-                "cwd": session_cwd,
-            }
+        try:
+            result = await sb.shell.exec(command, background=background, env=env)
             return json.dumps(result)
-
-        except subprocess.TimeoutExpired as e:
-            return json.dumps(
-                {
-                    "success": False,
-                    "exit_code": -1,
-                    "stdout": e.stdout or "",
-                    "stderr": f"Command timed out after {timeout} seconds",
-                    "cwd": session_cwd,
-                }
-            )
         except Exception as e:
-            # Do not silently swallow errors; return an explicit failure payload.
-            return json.dumps(
-                {
-                    "success": False,
-                    "exit_code": -1,
-                    "stdout": "",
-                    "stderr": f"Error executing command: {e!s}",
-                    "cwd": session_cwd,
-                }
-            )
+            return f"Error executing command: {str(e)}"