diff --git a/astrbot/_internal/protocols/mcp/client.py b/astrbot/_internal/protocols/mcp/client.py index 757a61429..04780b9be 100644 --- a/astrbot/_internal/protocols/mcp/client.py +++ b/astrbot/_internal/protocols/mcp/client.py @@ -270,7 +270,7 @@ class McpClient(BaseAstrbotMcpClient): mcp.ClientSession( *streams, read_timeout_seconds=read_timeout, - logging_callback=logging_callback, # type: ignore + logging_callback=cast(Any, logging_callback), ), ) else: @@ -322,12 +322,12 @@ class McpClient(BaseAstrbotMcpClient): stdio_transport = await self.exit_stack.enter_async_context( mcp.stdio_client( server_params, - errlog=LogPipe( + errlog=cast(Any, LogPipe( level=logging.INFO, logger=logger, identifier=f"MCPServer-{name}", callback=callback, - ), # type: ignore + )), ), ) self.process_pid = self._extract_stdio_process_pid(stdio_transport) @@ -421,7 +421,7 @@ class McpClient(BaseAstrbotMcpClient): retry=retry_if_exception_type(anyio.ClosedResourceError), stop=stop_after_attempt(2), wait=wait_exponential(multiplier=1, min=1, max=3), - before_sleep=before_sleep_log(logger, logging.WARNING), # type: ignore[arg-type] + before_sleep=cast(Any, before_sleep_log(logger, logging.WARNING)), reraise=True, ) async def _call_with_retry(): diff --git a/astrbot/cli/commands/cmd_init.py b/astrbot/cli/commands/cmd_init.py index 3c93c0f8a..ca4e8267a 100644 --- a/astrbot/cli/commands/cmd_init.py +++ b/astrbot/cli/commands/cmd_init.py @@ -3,6 +3,7 @@ import json import os import re from pathlib import Path +from typing import Any, cast import click from filelock import FileLock, Timeout @@ -112,7 +113,7 @@ async def initialize_astrbot( effective_admin_username = ( admin_username.strip() if admin_username - else str(DEFAULT_CONFIG["dashboard"]["username"]) # type: ignore[index] + else str(cast(dict[str, Any], DEFAULT_CONFIG)["dashboard"]["username"]) ) if admin_username: config = ensure_config_file() diff --git a/astrbot/core/computer/tools/shell.py b/astrbot/core/computer/tools/shell.py index 7b4932288..1124ce869 100644 --- a/astrbot/core/computer/tools/shell.py +++ b/astrbot/core/computer/tools/shell.py @@ -1,11 +1,28 @@ +""" +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 json +import os +import shlex +import subprocess from dataclasses import dataclass, field -from typing import Any - -from plumbum import local -from plumbum.commands.processes import ProcessExecutionError +from typing import Any, cast from astrbot.api import FunctionTool from astrbot.core.agent.run_context import ContextWrapper @@ -17,10 +34,10 @@ from .permissions import check_admin_permission @dataclass class ExecuteShellTool(FunctionTool): - """Stateful shell execution tool based on plumbum. + """ + Stateful shell execution tool using subprocess. - Each session maintains its own shell instance, ensuring state isolation - across different sessions (such as working directory, environment variables, etc.). + Each agent session keeps its own working directory and environment mapping. """ name: str = "astrbot_execute_shell" @@ -31,7 +48,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). Equal to 'cd {working_dir} && {your_command}'.", + "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}'.", }, "background": { "type": "boolean", @@ -40,7 +57,7 @@ class ExecuteShellTool(FunctionTool): }, "env": { "type": "object", - "description": "Optional environment variables to set for the file creation process.", + "description": "Optional environment variables to set for the command (merged with session env).", "additionalProperties": {"type": "string"}, "default": {}, }, @@ -50,111 +67,204 @@ class ExecuteShellTool(FunctionTool): ) is_local: bool = False - _session_shells: dict[str, Any] = field( + # session_id -> {"cwd": str, "env": dict} + _sessions: dict[str, dict[str, Any]] = field( default_factory=dict, init=False, repr=False ) - def _get_session_shell(self, session_id: str) -> Any: - """Get or create a shell instance for a specific session. - - Args: - session_id: The unique session identifier. - - Returns: - plumbum shell instance for this session. + def _get_session_state(self, session_id: str) -> dict[str, Any]: """ - import os + 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] - if session_id not in self._session_shells: - self._session_shells[session_id] = local.cwd(os.getcwd()) - return self._session_shells[session_id] - - async def call( - self, - context: ContextWrapper[AstrAgentContext], - command: str, - background: bool = False, - env: dict = {}, + async def call( # type: ignore[override] + self, context: ContextWrapper[TContext], **kwargs: Any ) -> ToolExecResult: - if permission_error := check_admin_permission(context, "Shell execution"): + """ + 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 + """ + # Cast the generic ContextWrapper to the concrete AstrAgentContext wrapper so + # subsequent permission checks and attribute access use the expected type. + astr_ctx = cast(ContextWrapper[AstrAgentContext], context) + + # Permission check (use the cast wrapper) + if permission_error := check_admin_permission(astr_ctx, "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 + state = self._get_session_state(session_id) + session_cwd = state["cwd"] + session_env = state["env"].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) — use the cast wrapper's context + config = astr_ctx.context.context.get_config(umo=session_id) try: - # Get session ID for per-session shell isolation - session_id = context.context.event.unified_msg_origin - shell = self._get_session_shell(session_id) + timeout = int( + config.get("provider_settings", {}).get("tool_call_timeout", 30) + ) + except (ValueError, TypeError): + timeout = 30 - config = context.context.context.get_config(umo=session_id) - try: - timeout = int( - config.get("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() - # Merge environment variables - if env: - full_env = shell.env.copy() - full_env.update(env) + # Helper to split by shell '&&' or ';' while preserving remainder. + remainder_cmd = "" + cd_handled = False + # Handle forms like: cd && rest OR cd ; 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: - full_env = None + # 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 background: - # Background execution: use & to run in background, no waiting - cmd_line = f"{command} &" - if full_env: - _, stdout, stderr = shell.run( - ["/bin/sh", "-c", cmd_line], - env=full_env, - timeout=timeout, - ) + if cd_handled: + # parse cd argument + parts = shlex.split(cd_part) + # cd with no args -> home + if len(parts) == 1: + target = os.path.expanduser("~") else: - _, stdout, stderr = shell.run( - ["/bin/sh", "-c", cmd_line], - timeout=timeout, + target_raw = parts[1] + # expand ~ and variables + target_raw = os.path.expanduser(target_raw) + target = ( + target_raw + if os.path.isabs(target_raw) + else os.path.normpath(os.path.join(session_cwd, target_raw)) ) - result = { - "success": True, - "background": True, - "stdout": stdout, - "stderr": stderr, - } - else: - # Foreground execution: wait for command completion - try: - if full_env: - exit_code, stdout, stderr = shell.run( - ["/bin/sh", "-c", command], - env=full_env, - timeout=timeout, - ) - else: - exit_code, stdout, stderr = shell.run( - ["/bin/sh", "-c", command], - timeout=timeout, - ) - result = { - "success": exit_code == 0, - "exit_code": exit_code, - "stdout": stdout, - "stderr": stderr, - } - except ProcessExecutionError as e: - result = { - "success": False, - "exit_code": e.exit_code, - "stdout": e.stdout, - "stderr": e.stderr, - } - except Exception as e: + + if not os.path.exists(target) or not os.path.isdir(target): result = { "success": False, "exit_code": -1, "stdout": "", - "stderr": str(e), + "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 = 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 = subprocess.run( + ["/bin/sh", "-c", command_to_run], + cwd=session_cwd, + env=exec_env, + timeout=timeout, + capture_output=True, + text=True, + ) + + 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, + } return json.dumps(result) + except subprocess.TimeoutExpired as e: + return json.dumps( + { + "success": False, + "exit_code": -1, + "stdout": getattr(e, "output", "") or "", + "stderr": f"Command timed out after {timeout} seconds", + "cwd": session_cwd, + } + ) except Exception as e: - return f"Error executing command: {e!s}" + # 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, + } + ) diff --git a/examples/abp_demo.py b/examples/abp_demo.py index 73f9b4c1f..e6407db02 100644 --- a/examples/abp_demo.py +++ b/examples/abp_demo.py @@ -12,11 +12,11 @@ Run with: uv run python examples/abp_demo.py from __future__ import annotations -import anyio -import asyncio from dataclasses import dataclass from typing import Any +import anyio + # Mock Star class that implements the expected interface @dataclass