mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-01 18:20:16 +08:00
Compare commits
10 Commits
v4.26.3
...
perf/shell
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
fddfbb946c | ||
|
|
451d5450ee | ||
|
|
33eeed1739 | ||
|
|
3f52c7aaa0 | ||
|
|
44e31bd49c | ||
|
|
ffa7508a8b | ||
|
|
988961a9a4 | ||
|
|
d7a54aed76 | ||
|
|
f1206d987b | ||
|
|
77040cabcd |
@@ -90,7 +90,7 @@ class LocalShellComponent(ShellComponent):
|
||||
command: str,
|
||||
cwd: str | None = None,
|
||||
env: dict[str, str] | None = None,
|
||||
timeout: int | None = 30,
|
||||
timeout: int | None = 300,
|
||||
shell: bool = True,
|
||||
background: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
@@ -123,7 +123,7 @@ class LocalShellComponent(ShellComponent):
|
||||
shell=shell,
|
||||
cwd=working_dir,
|
||||
env=run_env,
|
||||
timeout=timeout,
|
||||
timeout=timeout or 300,
|
||||
capture_output=True,
|
||||
)
|
||||
return {
|
||||
|
||||
18
astrbot/core/computer/booters/shell_background.py
Normal file
18
astrbot/core/computer/booters/shell_background.py
Normal file
@@ -0,0 +1,18 @@
|
||||
import shlex
|
||||
|
||||
_BACKGROUND_SPAWN_SCRIPT = (
|
||||
"import subprocess, sys; "
|
||||
"p = subprocess.Popen("
|
||||
"['bash', '-lc', sys.argv[1]], "
|
||||
"stdin=subprocess.DEVNULL, "
|
||||
"stdout=subprocess.DEVNULL, "
|
||||
"stderr=subprocess.DEVNULL, "
|
||||
"start_new_session=True, "
|
||||
"close_fds=True"
|
||||
"); "
|
||||
"print(p.pid)"
|
||||
)
|
||||
|
||||
|
||||
def build_detached_shell_command(command: str) -> str:
|
||||
return f"python3 -c {shlex.quote(_BACKGROUND_SPAWN_SCRIPT)} {shlex.quote(command)}"
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import shlex
|
||||
from typing import Any
|
||||
|
||||
from shipyard import FileSystemComponent as ShipyardFileSystemComponent
|
||||
@@ -9,9 +10,93 @@ from astrbot.api import logger
|
||||
|
||||
from ..olayer import FileSystemComponent, PythonComponent, ShellComponent
|
||||
from .base import ComputerBooter
|
||||
from .shell_background import build_detached_shell_command
|
||||
from .shipyard_search_file_util import search_files_via_shell
|
||||
|
||||
|
||||
def _maybe_model_dump(value: Any) -> dict[str, Any]:
|
||||
if isinstance(value, dict):
|
||||
return value
|
||||
if hasattr(value, "model_dump"):
|
||||
dumped = value.model_dump()
|
||||
if isinstance(dumped, dict):
|
||||
return dumped
|
||||
return {}
|
||||
|
||||
|
||||
class ShipyardShellWrapper:
|
||||
def __init__(self, _shipyard_shell: ShellComponent):
|
||||
self._shell = _shipyard_shell
|
||||
|
||||
async def exec(
|
||||
self,
|
||||
command: str,
|
||||
cwd: str | None = None,
|
||||
env: dict[str, str] | None = None,
|
||||
timeout: int | None = 300,
|
||||
shell: bool = True,
|
||||
background: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
if not shell:
|
||||
return {
|
||||
"stdout": "",
|
||||
"stderr": "error: only shell mode is supported in shipyard booter.",
|
||||
"exit_code": 2,
|
||||
"success": False,
|
||||
}
|
||||
|
||||
run_command = command
|
||||
if env:
|
||||
env_prefix = " ".join(
|
||||
f"{k}={shlex.quote(str(v))}" for k, v in sorted(env.items())
|
||||
)
|
||||
run_command = f"{env_prefix} {run_command}"
|
||||
|
||||
if background:
|
||||
run_command = build_detached_shell_command(run_command)
|
||||
|
||||
result = await self._shell.exec(
|
||||
run_command,
|
||||
timeout=timeout or 300,
|
||||
cwd=cwd,
|
||||
)
|
||||
payload = _maybe_model_dump(result)
|
||||
|
||||
stdout = payload.get("output", payload.get("stdout", "")) or ""
|
||||
stderr = payload.get("error", payload.get("stderr", "")) or ""
|
||||
exit_code = payload.get("exit_code")
|
||||
if background:
|
||||
pid: int | None = None
|
||||
try:
|
||||
pid = int(str(stdout).strip().splitlines()[-1])
|
||||
except Exception:
|
||||
pid = None
|
||||
return {
|
||||
"pid": pid,
|
||||
"stdout": (
|
||||
f"Command is running in the background. pid={pid}"
|
||||
if pid is not None
|
||||
else "Command was submitted in the background."
|
||||
),
|
||||
"stderr": stderr,
|
||||
"exit_code": exit_code,
|
||||
"success": bool(payload.get("success", not stderr)),
|
||||
"execution_id": payload.get("execution_id"),
|
||||
"execution_time_ms": payload.get("execution_time_ms"),
|
||||
"command": payload.get("command"),
|
||||
}
|
||||
|
||||
return {
|
||||
"stdout": stdout,
|
||||
"stderr": stderr,
|
||||
"exit_code": exit_code,
|
||||
"success": bool(payload.get("success", not stderr)),
|
||||
"execution_id": payload.get("execution_id"),
|
||||
"execution_time_ms": payload.get("execution_time_ms"),
|
||||
"command": payload.get("command"),
|
||||
}
|
||||
|
||||
|
||||
class ShipyardFileSystemWrapper:
|
||||
def __init__(
|
||||
self, _shipyard_fs: ShipyardFileSystemComponent, _shipyard_shell: ShellComponent
|
||||
@@ -107,7 +192,8 @@ class ShipyardBooter(ComputerBooter):
|
||||
)
|
||||
logger.info(f"Got sandbox ship: {ship.id} for session: {session_id}")
|
||||
self._ship = ship
|
||||
self._fs = ShipyardFileSystemWrapper(self._ship.fs, self._ship.shell)
|
||||
self._shell = ShipyardShellWrapper(self._ship.shell)
|
||||
self._fs = ShipyardFileSystemWrapper(self._ship.fs, self._shell)
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
logger.info("[Computer] Shipyard booter shutdown.")
|
||||
@@ -122,7 +208,7 @@ class ShipyardBooter(ComputerBooter):
|
||||
|
||||
@property
|
||||
def shell(self) -> ShellComponent:
|
||||
return self._ship.shell
|
||||
return self._shell
|
||||
|
||||
async def upload_file(self, path: str, file_name: str) -> dict:
|
||||
"""Upload file to sandbox"""
|
||||
|
||||
@@ -13,6 +13,7 @@ from ..olayer import (
|
||||
ShellComponent,
|
||||
)
|
||||
from .base import ComputerBooter
|
||||
from .shell_background import build_detached_shell_command
|
||||
from .shipyard_search_file_util import search_files_via_shell
|
||||
|
||||
try:
|
||||
@@ -96,7 +97,7 @@ class NeoShellComponent(ShellComponent):
|
||||
command: str,
|
||||
cwd: str | None = None,
|
||||
env: dict[str, str] | None = None,
|
||||
timeout: int | None = 30,
|
||||
timeout: int | None = 300,
|
||||
shell: bool = True,
|
||||
background: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
@@ -116,11 +117,11 @@ class NeoShellComponent(ShellComponent):
|
||||
run_command = f"{env_prefix} {run_command}"
|
||||
|
||||
if background:
|
||||
run_command = f"nohup sh -lc {shlex.quote(run_command)} >/tmp/astrbot_bg.log 2>&1 & echo $!"
|
||||
run_command = build_detached_shell_command(run_command)
|
||||
|
||||
result = await self._sandbox.shell.exec(
|
||||
run_command,
|
||||
timeout=timeout or 30,
|
||||
timeout=timeout or 300,
|
||||
cwd=cwd,
|
||||
)
|
||||
payload = _maybe_model_dump(result)
|
||||
@@ -136,7 +137,11 @@ class NeoShellComponent(ShellComponent):
|
||||
pid = None
|
||||
return {
|
||||
"pid": pid,
|
||||
"stdout": stdout,
|
||||
"stdout": (
|
||||
f"Command is running in the background. pid={pid}"
|
||||
if pid is not None
|
||||
else "Command was submitted in the background."
|
||||
),
|
||||
"stderr": stderr,
|
||||
"exit_code": exit_code,
|
||||
"success": bool(payload.get("success", not stderr)),
|
||||
|
||||
@@ -13,7 +13,7 @@ class ShellComponent(Protocol):
|
||||
command: str,
|
||||
cwd: str | None = None,
|
||||
env: dict[str, str] | None = None,
|
||||
timeout: int | None = 30,
|
||||
timeout: int | None = 300,
|
||||
shell: bool = True,
|
||||
background: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from astrbot.api import FunctionTool
|
||||
@@ -8,6 +11,7 @@ from astrbot.core.agent.run_context import ContextWrapper
|
||||
from astrbot.core.agent.tool import ToolExecResult
|
||||
from astrbot.core.astr_agent_context import AstrAgentContext
|
||||
from astrbot.core.computer.computer_client import get_booter
|
||||
from astrbot.core.utils.astrbot_path import get_astrbot_system_tmp_path
|
||||
|
||||
from ..registry import builtin_tool
|
||||
from .util import check_admin_permission, is_local_runtime, workspace_root
|
||||
@@ -17,6 +21,32 @@ _COMPUTER_RUNTIME_TOOL_CONFIG = {
|
||||
}
|
||||
|
||||
|
||||
def _quote_redirect_path(path: str, *, local_runtime: bool) -> str:
|
||||
if local_runtime and os.name == "nt":
|
||||
escaped_path = path.replace('"', '""')
|
||||
else:
|
||||
escaped_path = path.replace("\\", "\\\\").replace('"', '\\"')
|
||||
return f'"{escaped_path}"'
|
||||
|
||||
|
||||
def _build_background_output_path(*, local_runtime: bool) -> str:
|
||||
file_name = f"astrbot_shell_stdout_{uuid.uuid4().hex[:8]}.log"
|
||||
if local_runtime:
|
||||
output_dir = Path(get_astrbot_system_tmp_path()) / "shell"
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
return str((output_dir / file_name).resolve(strict=False))
|
||||
return f"/tmp/{file_name}"
|
||||
|
||||
|
||||
def _redirect_background_stdout_command(
|
||||
command: str,
|
||||
*,
|
||||
output_path: str,
|
||||
local_runtime: bool,
|
||||
) -> str:
|
||||
return f"({command}) > {_quote_redirect_path(output_path, local_runtime=local_runtime)} 2>&1"
|
||||
|
||||
|
||||
@builtin_tool(config=_COMPUTER_RUNTIME_TOOL_CONFIG)
|
||||
@dataclass
|
||||
class ExecuteShellTool(FunctionTool):
|
||||
@@ -32,12 +62,17 @@ class ExecuteShellTool(FunctionTool):
|
||||
},
|
||||
"background": {
|
||||
"type": "boolean",
|
||||
"description": "Whether to run the command in the background.",
|
||||
"description": "Run the command in the background. Use the file read tool to read the output later. For long running commands, using this option.",
|
||||
"default": False,
|
||||
},
|
||||
"timeout": {
|
||||
"type": "integer",
|
||||
"description": "Optional timeout in seconds for the command execution.",
|
||||
"default": 300,
|
||||
},
|
||||
"env": {
|
||||
"type": "object",
|
||||
"description": "Optional environment variables to set for the file creation process.",
|
||||
"description": "Optional environment variables to set.",
|
||||
"additionalProperties": {"type": "string"},
|
||||
"default": {},
|
||||
},
|
||||
@@ -51,6 +86,7 @@ class ExecuteShellTool(FunctionTool):
|
||||
context: ContextWrapper[AstrAgentContext],
|
||||
command: str,
|
||||
background: bool = False,
|
||||
timeout: int | None = 300,
|
||||
env: dict[str, Any] | None = None,
|
||||
) -> ToolExecResult:
|
||||
if permission_error := check_admin_permission(context, "Shell execution"):
|
||||
@@ -69,6 +105,18 @@ class ExecuteShellTool(FunctionTool):
|
||||
current_workspace_root.mkdir(parents=True, exist_ok=True)
|
||||
cwd = str(current_workspace_root)
|
||||
|
||||
stdout_file: str | None = None
|
||||
if background:
|
||||
local_runtime = is_local_runtime(context)
|
||||
stdout_file = _build_background_output_path(
|
||||
local_runtime=local_runtime,
|
||||
)
|
||||
command = _redirect_background_stdout_command(
|
||||
command,
|
||||
output_path=stdout_file,
|
||||
local_runtime=local_runtime,
|
||||
)
|
||||
|
||||
env = dict(env or {})
|
||||
effective_background = background and not _is_self_detached_command(command)
|
||||
result = await sb.shell.exec(
|
||||
@@ -76,7 +124,13 @@ class ExecuteShellTool(FunctionTool):
|
||||
cwd=cwd,
|
||||
background=effective_background,
|
||||
env=env,
|
||||
timeout=timeout or 300,
|
||||
)
|
||||
if stdout_file:
|
||||
result["stdout"] = (
|
||||
f"Command is running in the background. stdout/stderr is being "
|
||||
f"written to `{stdout_file}`. Use astrbot_file_read_tool to read it."
|
||||
)
|
||||
return json.dumps(result, ensure_ascii=False)
|
||||
except Exception as e:
|
||||
detail = str(e) or type(e).__name__
|
||||
|
||||
Reference in New Issue
Block a user