mirror of
https://github.com/jxxghp/MoviePilot.git
synced 2026-05-10 07:12:12 +08:00
c762628217
Return only a 10KB preview to the agent so large command results do not overwhelm conversations while keeping the full output available for follow-up reads. Add pytest to the project dependencies to make the regression tests runnable in the project venv.
346 lines
12 KiB
Python
346 lines
12 KiB
Python
"""执行Shell命令工具"""
|
||
|
||
import asyncio
|
||
import os
|
||
import signal
|
||
import subprocess
|
||
from dataclasses import dataclass, field
|
||
from tempfile import NamedTemporaryFile
|
||
from typing import Optional, TextIO, Type
|
||
|
||
from pydantic import BaseModel, Field
|
||
|
||
from app.agent.tools.base import MoviePilotTool
|
||
from app.log import logger
|
||
|
||
|
||
DEFAULT_TIMEOUT_SECONDS = 60
|
||
MAX_TIMEOUT_SECONDS = 300
|
||
MAX_OUTPUT_PREVIEW_BYTES = 10 * 1024
|
||
READ_CHUNK_SIZE = 4096
|
||
KILL_GRACE_SECONDS = 3
|
||
COMMAND_CONCURRENCY_LIMIT = 2
|
||
|
||
_command_semaphore = asyncio.Semaphore(COMMAND_CONCURRENCY_LIMIT)
|
||
|
||
|
||
@dataclass
|
||
class _CommandOutput:
|
||
"""保存前 10KB 预览，并在超限时将完整输出写入临时文件。"""
|
||
|
||
preview_limit_bytes: int
|
||
preview_entries: list[tuple[str, str]] = field(default_factory=list)
|
||
captured_bytes: int = 0
|
||
preview_truncated: bool = False
|
||
temp_file_path: Optional[str] = None
|
||
temp_file_handle: Optional[TextIO] = None
|
||
last_written_stream: Optional[str] = None
|
||
|
||
@staticmethod
|
||
def _clip_text_to_bytes(text: str, byte_limit: int) -> str:
|
||
if byte_limit <= 0:
|
||
return ""
|
||
return text.encode("utf-8")[:byte_limit].decode("utf-8", errors="ignore")
|
||
|
||
def _write_chunk(self, stream_name: str, text: str) -> None:
|
||
if not self.temp_file_handle or not text:
|
||
return
|
||
|
||
if self.last_written_stream != stream_name:
|
||
if self.temp_file_handle.tell() > 0:
|
||
self.temp_file_handle.write("\n")
|
||
title = "标准输出" if stream_name == "stdout" else "错误输出"
|
||
self.temp_file_handle.write(f"[{title}]\n")
|
||
self.last_written_stream = stream_name
|
||
|
||
self.temp_file_handle.write(text)
|
||
|
||
def _ensure_temp_file(self) -> None:
|
||
if self.temp_file_handle:
|
||
return
|
||
|
||
temp_file = NamedTemporaryFile(
|
||
mode="w",
|
||
encoding="utf-8",
|
||
suffix=".log",
|
||
prefix="moviepilot-command-",
|
||
delete=False,
|
||
)
|
||
self.temp_file_path = temp_file.name
|
||
self.temp_file_handle = temp_file
|
||
for stream_name, chunk in self.preview_entries:
|
||
self._write_chunk(stream_name, chunk)
|
||
|
||
def close(self) -> None:
|
||
if not self.temp_file_handle:
|
||
return
|
||
self.temp_file_handle.flush()
|
||
self.temp_file_handle.close()
|
||
self.temp_file_handle = None
|
||
|
||
def append(self, stream_name: str, text: str) -> None:
|
||
if not text:
|
||
return
|
||
|
||
if self.temp_file_handle:
|
||
self._write_chunk(stream_name, text)
|
||
return
|
||
|
||
chunk_bytes = len(text.encode("utf-8"))
|
||
remaining = self.preview_limit_bytes - self.captured_bytes
|
||
if chunk_bytes <= remaining:
|
||
self.preview_entries.append((stream_name, text))
|
||
self.captured_bytes += chunk_bytes
|
||
return
|
||
|
||
self.preview_truncated = True
|
||
self._ensure_temp_file()
|
||
self._write_chunk(stream_name, text)
|
||
|
||
preview = self._clip_text_to_bytes(text, remaining)
|
||
if preview:
|
||
self.preview_entries.append((stream_name, preview))
|
||
self.captured_bytes += len(preview.encode("utf-8"))
|
||
|
||
@property
|
||
def stdout(self) -> str:
|
||
return "".join(
|
||
text for stream_name, text in self.preview_entries if stream_name == "stdout"
|
||
).strip()
|
||
|
||
@property
|
||
def stderr(self) -> str:
|
||
return "".join(
|
||
text for stream_name, text in self.preview_entries if stream_name == "stderr"
|
||
).strip()
|
||
|
||
|
||
class ExecuteCommandInput(BaseModel):
|
||
"""执行Shell命令工具的输入参数模型"""
|
||
|
||
explanation: str = Field(
|
||
..., description="Clear explanation of why this command is being executed"
|
||
)
|
||
command: str = Field(..., description="The shell command to execute")
|
||
timeout: Optional[int] = Field(
|
||
60, description="Max execution time in seconds (default: 60)"
|
||
)
|
||
|
||
|
||
class ExecuteCommandTool(MoviePilotTool):
|
||
name: str = "execute_command"
|
||
description: str = (
|
||
"Safely execute shell commands on the server. Useful for system "
|
||
"maintenance, checking status, or running custom scripts. Includes "
|
||
"timeout, concurrency, and output preview limits."
|
||
)
|
||
args_schema: Type[BaseModel] = ExecuteCommandInput
|
||
require_admin: bool = True
|
||
|
||
def get_tool_message(self, **kwargs) -> Optional[str]:
|
||
"""根据命令生成友好的提示消息"""
|
||
command = kwargs.get("command", "")
|
||
return f"执行系统命令: {command}"
|
||
|
||
@staticmethod
|
||
def _normalize_timeout(timeout: Optional[int]) -> tuple[int, Optional[str]]:
|
||
"""限制命令最长运行时间，避免 Agent 传入过大的 timeout。"""
|
||
try:
|
||
normalized = int(timeout or DEFAULT_TIMEOUT_SECONDS)
|
||
except (TypeError, ValueError):
|
||
normalized = DEFAULT_TIMEOUT_SECONDS
|
||
|
||
if normalized <= 0:
|
||
return DEFAULT_TIMEOUT_SECONDS, "timeout 参数无效，已使用默认 60 秒"
|
||
if normalized > MAX_TIMEOUT_SECONDS:
|
||
return (
|
||
MAX_TIMEOUT_SECONDS,
|
||
f"timeout 参数超过上限，已从 {normalized} 秒限制为 {MAX_TIMEOUT_SECONDS} 秒",
|
||
)
|
||
return normalized, None
|
||
|
||
@staticmethod
|
||
def _subprocess_kwargs() -> dict:
|
||
"""为子进程创建独立进程组，便于超时场景清理整棵子进程。"""
|
||
kwargs = {
|
||
"stdin": subprocess.DEVNULL,
|
||
"stdout": asyncio.subprocess.PIPE,
|
||
"stderr": asyncio.subprocess.PIPE,
|
||
}
|
||
if os.name == "posix":
|
||
kwargs["start_new_session"] = True
|
||
elif os.name == "nt":
|
||
kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
|
||
return kwargs
|
||
|
||
@staticmethod
|
||
async def _read_stream(
|
||
stream: asyncio.StreamReader,
|
||
stream_name: str,
|
||
output: _CommandOutput,
|
||
) -> None:
|
||
"""按块读取输出，始终只把前 10KB 保留在返回结果中。"""
|
||
while True:
|
||
chunk = await stream.read(READ_CHUNK_SIZE)
|
||
if not chunk:
|
||
break
|
||
|
||
output.append(stream_name, chunk.decode("utf-8", errors="replace"))
|
||
|
||
@staticmethod
|
||
def _terminate_process(process: asyncio.subprocess.Process, sig: int):
|
||
"""向进程组发送终止信号；不支持进程组的平台回退为单进程终止。"""
|
||
try:
|
||
if os.name == "posix":
|
||
os.killpg(process.pid, sig)
|
||
elif sig == getattr(signal, "SIGKILL", None):
|
||
process.kill()
|
||
else:
|
||
process.terminate()
|
||
except ProcessLookupError:
|
||
pass
|
||
|
||
@classmethod
|
||
async def _cleanup_process(
|
||
cls,
|
||
process: asyncio.subprocess.Process,
|
||
wait_task: asyncio.Task,
|
||
) -> None:
|
||
"""先温和终止，失败后强杀，避免超时 shell 遗留子进程。"""
|
||
if wait_task.done():
|
||
return
|
||
|
||
cls._terminate_process(process, signal.SIGTERM)
|
||
try:
|
||
await asyncio.wait_for(
|
||
asyncio.shield(wait_task), timeout=KILL_GRACE_SECONDS
|
||
)
|
||
return
|
||
except asyncio.TimeoutError:
|
||
pass
|
||
|
||
kill_signal = getattr(signal, "SIGKILL", signal.SIGTERM)
|
||
cls._terminate_process(process, kill_signal)
|
||
try:
|
||
await asyncio.wait_for(
|
||
asyncio.shield(wait_task), timeout=KILL_GRACE_SECONDS
|
||
)
|
||
except asyncio.TimeoutError:
|
||
logger.warning("命令进程强制清理超时: pid=%s", process.pid)
|
||
|
||
@staticmethod
|
||
async def _finish_reader_tasks(reader_tasks: list[asyncio.Task]) -> None:
|
||
"""等待输出读取任务退出，异常只记录不影响工具返回。"""
|
||
if not reader_tasks:
|
||
return
|
||
done, pending = await asyncio.wait(reader_tasks, timeout=1)
|
||
for task in pending:
|
||
task.cancel()
|
||
results = await asyncio.gather(*done, *pending, return_exceptions=True)
|
||
for result in results:
|
||
if isinstance(result, Exception) and not isinstance(
|
||
result, asyncio.CancelledError
|
||
):
|
||
logger.debug("命令输出读取任务异常: %s", result)
|
||
|
||
@staticmethod
|
||
def _format_result(
|
||
*,
|
||
exit_code: Optional[int],
|
||
output: _CommandOutput,
|
||
timeout: int,
|
||
timed_out: bool,
|
||
timeout_note: Optional[str],
|
||
) -> str:
|
||
if timed_out:
|
||
result = f"命令执行超时 (限制: {timeout}秒，已终止进程)"
|
||
else:
|
||
result = f"命令执行完成 (退出码: {exit_code})"
|
||
|
||
if timeout_note:
|
||
result += f"\n\n提示:\n{timeout_note}"
|
||
if output.temp_file_path:
|
||
file_note = (
|
||
"截至命令终止前的完整输出"
|
||
if timed_out
|
||
else "完整输出"
|
||
)
|
||
result += (
|
||
"\n\n提示:\n"
|
||
f"命令输出超过 10KB，仅返回前 {MAX_OUTPUT_PREVIEW_BYTES} 字节内容。\n"
|
||
f"{file_note}已写入临时文件: {output.temp_file_path}\n"
|
||
"如需完整内容，请继续读取该文件。"
|
||
)
|
||
if output.stdout:
|
||
result += f"\n\n标准输出:\n{output.stdout}"
|
||
if output.stderr:
|
||
result += f"\n\n错误输出:\n{output.stderr}"
|
||
if output.preview_truncated:
|
||
result += "\n\n...(仅展示前 10KB 内容)"
|
||
if not output.stdout and not output.stderr:
|
||
result += "\n\n(无输出内容)"
|
||
return result
|
||
|
||
async def run(self, command: str, timeout: Optional[int] = 60, **kwargs) -> str:
|
||
logger.info(
|
||
f"执行工具: {self.name}, 参数: command={command}, timeout={timeout}"
|
||
)
|
||
|
||
# 简单安全过滤
|
||
forbidden_keywords = [
|
||
"rm -rf /",
|
||
":(){ :|:& };:",
|
||
"dd if=/dev/zero",
|
||
"mkfs",
|
||
"reboot",
|
||
"shutdown",
|
||
]
|
||
for keyword in forbidden_keywords:
|
||
if keyword in command:
|
||
return f"错误：命令包含禁止使用的关键字 '{keyword}'"
|
||
|
||
normalized_timeout, timeout_note = self._normalize_timeout(timeout)
|
||
|
||
try:
|
||
async with _command_semaphore:
|
||
# 命令输出可能非常大，必须边读边落盘，不能使用 communicate() 一次性收集。
|
||
process = await asyncio.create_subprocess_shell(
|
||
command, **self._subprocess_kwargs()
|
||
)
|
||
output = _CommandOutput(preview_limit_bytes=MAX_OUTPUT_PREVIEW_BYTES)
|
||
wait_task = asyncio.create_task(process.wait())
|
||
reader_tasks = [
|
||
asyncio.create_task(
|
||
self._read_stream(process.stdout, "stdout", output)
|
||
),
|
||
asyncio.create_task(
|
||
self._read_stream(process.stderr, "stderr", output)
|
||
),
|
||
]
|
||
|
||
timed_out = False
|
||
try:
|
||
await asyncio.wait_for(
|
||
asyncio.shield(wait_task), timeout=normalized_timeout
|
||
)
|
||
except asyncio.TimeoutError:
|
||
timed_out = True
|
||
await self._cleanup_process(process, wait_task)
|
||
|
||
try:
|
||
await self._finish_reader_tasks(reader_tasks)
|
||
finally:
|
||
output.close()
|
||
|
||
return self._format_result(
|
||
exit_code=process.returncode,
|
||
output=output,
|
||
timeout=normalized_timeout,
|
||
timed_out=timed_out,
|
||
timeout_note=timeout_note,
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"执行命令失败: {e}", exc_info=True)
|
||
return f"执行命令时发生错误: {str(e)}"
|