mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-17 01:49:15 +08:00
chore: auto commit
This commit is contained in:
LIghtJUNction
@@ -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():
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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 <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:
|
||||
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,
|
||||
}
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user