agent-Specialization/modules/terminal_ops/run.py

527 lines
21 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# modules/terminal_ops.py - 终端操作模块修复Python命令检测
import os
import sys
import asyncio
import subprocess
import shutil
import time
import signal
import re
from pathlib import Path
from typing import Dict, Optional, Tuple, TYPE_CHECKING
from types import SimpleNamespace
try:
from config import (
TERMINAL_COMMAND_TIMEOUT,
FORBIDDEN_COMMANDS,
OUTPUT_FORMATS,
MAX_RUN_COMMAND_CHARS,
TOOLBOX_TERMINAL_IDLE_SECONDS,
HOST_SANDBOX_NETWORK_PERMISSION,
)
except ImportError:
project_root = Path(__file__).resolve().parents[1]
if str(project_root) not in sys.path:
sys.path.insert(0, str(project_root))
from config import (
TERMINAL_COMMAND_TIMEOUT,
FORBIDDEN_COMMANDS,
OUTPUT_FORMATS,
MAX_RUN_COMMAND_CHARS,
TOOLBOX_TERMINAL_IDLE_SECONDS,
HOST_SANDBOX_NETWORK_PERMISSION,
)
from modules.toolbox_container import ToolboxContainer
from modules.host_sandbox_runner import (
HostSandboxError,
NETWORK_PERMISSION_RESTRICTED,
build_host_sandbox_plan,
build_host_sandbox_readonly_plan,
host_sandbox_enabled,
)
if TYPE_CHECKING:
from modules.user_container_manager import ContainerHandle
from modules.terminal_manager import TerminalManager
class RunMixin:
"""TerminalOperator run 能力 mixin。"""
async def check_environment(self) -> Dict:
"""检查Python环境"""
print(f"{OUTPUT_FORMATS['info']} 检查Python环境...")
env_info = {
"python_command": self.python_cmd,
"python_version": "",
"pip_version": "",
"installed_packages": [],
"working_directory": str(self.project_path)
}
# 获取Python版本
version_result = await self.run_command(
'python3 --version',
timeout=5
)
if version_result["success"]:
env_info["python_version"] = version_result["output"].strip()
# 获取pip版本
pip_result = await self.run_command(
'python3 -m pip --version',
timeout=5
)
if pip_result["success"]:
env_info["pip_version"] = pip_result["output"].strip()
# 获取已安装的包
packages_result = await self.run_command(
'python3 -m pip list --format=json',
timeout=10
)
if packages_result["success"]:
try:
import json
packages = json.loads(packages_result["output"])
env_info["installed_packages"] = [
f"{p['name']}=={p['version']}" for p in packages
]
except:
pass
return {
"success": True,
"environment": env_info
}
async def _run_command_subprocess(
self,
command: str,
work_path: Path,
timeout: int,
session_override: Optional["ContainerHandle"] = None,
sandbox_write_access: bool = True,
network_permission: Optional[str] = None,
) -> Dict:
start_ts = time.time()
try:
process = None
exec_cmd = None
use_shell = True
session = session_override or self.container_session
# 如果存在容器会话且模式为docker则在容器内执行
if session and getattr(session, "mode", None) == "docker":
container_name = getattr(session, "container_name", None)
mount_path = getattr(session, "mount_path", "/workspace") or "/workspace"
docker_bin = shutil.which("docker") or "docker"
try:
relative = work_path.relative_to(self.project_path).as_posix()
except ValueError:
relative = ""
container_workdir = mount_path.rstrip("/")
if relative:
container_workdir = f"{container_workdir}/{relative}"
exec_cmd = [
docker_bin,
"exec",
"-e",
"PATH=/opt/agent-venv/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin",
"-e",
"VIRTUAL_ENV=/opt/agent-venv",
"-w",
container_workdir,
container_name,
"/bin/bash",
"-lc",
command,
]
use_shell = False
# 统一环境,确保 Python 输出无缓冲
env = os.environ.copy()
env.setdefault("PYTHONUNBUFFERED", "1")
if self._python_env:
env.update(self._python_env)
if use_shell:
use_host_sandbox = self.host_execution_mode != "direct"
if use_host_sandbox and host_sandbox_enabled():
if sandbox_write_access:
plan = build_host_sandbox_plan(
command, work_path, env, network_permission=network_permission
)
else:
plan = build_host_sandbox_readonly_plan(
command, work_path, env, network_permission=network_permission
)
cmd_args, pass_fds, seccomp_fd = self._materialize_seccomp_fd(
plan.command,
plan.seccomp_bpf_path,
)
try:
process = await asyncio.create_subprocess_exec(
*cmd_args,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=plan.cwd,
env=plan.env,
start_new_session=True,
pass_fds=pass_fds,
)
finally:
if seccomp_fd is not None:
try:
os.close(seccomp_fd)
except OSError:
pass
elif use_host_sandbox:
return {
"success": False,
"status": "error",
"error": "宿主机命令执行被拒绝HOST_SANDBOX_ENABLED=0",
"output": "",
"return_code": -1,
"timeout": timeout,
"elapsed_ms": int((time.time() - start_ts) * 1000)
}
else:
process = await asyncio.create_subprocess_shell(
command,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
cwd=str(work_path),
shell=True,
env=env,
start_new_session=True,
)
else:
process = await asyncio.create_subprocess_exec(
*exec_cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env,
start_new_session=True,
)
stdout_buf: list[bytes] = []
stderr_buf: list[bytes] = []
async def _read_stream(stream, collector):
try:
async for chunk in stream:
if chunk:
collector.append(chunk)
except asyncio.CancelledError:
pass
except Exception:
pass
stdout_task = asyncio.create_task(_read_stream(process.stdout, stdout_buf))
stderr_task = asyncio.create_task(_read_stream(process.stderr, stderr_buf))
async def _finish_reader_tasks(force: bool = False) -> None:
"""收口 reader 任务,避免因后台子进程持续占用管道而卡死。"""
join_timeout = 0.8 if force else 3.0
try:
await asyncio.wait_for(
asyncio.gather(stdout_task, stderr_task, return_exceptions=True),
timeout=join_timeout
)
except asyncio.TimeoutError:
stdout_task.cancel()
stderr_task.cancel()
try:
await asyncio.gather(stdout_task, stderr_task, return_exceptions=True)
except Exception:
pass
except Exception:
pass
timed_out = False
try:
await asyncio.wait_for(process.wait(), timeout=timeout)
except asyncio.TimeoutError:
timed_out = True
try:
os.killpg(process.pid, signal.SIGINT)
except Exception:
pass
try:
await asyncio.wait_for(process.wait(), timeout=2)
except asyncio.TimeoutError:
try:
os.killpg(process.pid, signal.SIGKILL)
except Exception:
process.kill()
await process.wait()
except asyncio.CancelledError:
# 用户主动停止任务或会话断开,立即终止子进程
try:
os.killpg(process.pid, signal.SIGKILL)
except Exception:
try:
process.kill()
except Exception:
pass
# 取消读取任务,避免孤儿任务
stdout_task.cancel()
stderr_task.cancel()
try:
await asyncio.gather(stdout_task, stderr_task, return_exceptions=True)
except Exception:
pass
raise
# 确保读取协程结束(超时场景强制收口,避免永久等待 EOF
await _finish_reader_tasks(force=timed_out)
# 非超时场景下兜底再读一次,防止剩余缓冲未被读取
if not timed_out:
try:
remaining_out = await asyncio.wait_for(process.stdout.read(), timeout=0.2)
if remaining_out:
stdout_buf.append(remaining_out)
except Exception:
pass
try:
remaining_err = await asyncio.wait_for(process.stderr.read(), timeout=0.2)
if remaining_err:
stderr_buf.append(remaining_err)
except Exception:
pass
stdout = b"".join(stdout_buf)
stderr = b"".join(stderr_buf)
stdout_text = stdout.decode('utf-8', errors='replace') if stdout else ""
stderr_text = stderr.decode('utf-8', errors='replace') if stderr else ""
success = (process.returncode == 0) and not timed_out
status = "completed" if success else ("timeout" if timed_out else "error")
output_parts = []
if stdout_text:
output_parts.append(stdout_text)
if stderr_text:
output_parts.append(stderr_text)
combined_output = "\n".join(output_parts)
truncated = False
if MAX_RUN_COMMAND_CHARS and len(combined_output) > MAX_RUN_COMMAND_CHARS:
truncated = True
combined_output = combined_output[-MAX_RUN_COMMAND_CHARS:]
response = {
"success": success,
"status": status,
"command": command,
"output": combined_output,
"return_code": process.returncode,
"truncated": truncated,
"timeout": timeout,
"elapsed_ms": int((time.time() - start_ts) * 1000)
}
if not success and timed_out:
response["message"] = f"命令执行超时 ({timeout}秒)"
elif not success and process.returncode is not None:
response["message"] = f"命令执行失败 (返回码: {process.returncode})"
if stderr_text:
response["stderr"] = stderr_text
return response
except HostSandboxError as exc:
return {
"success": False,
"status": "error",
"error": f"宿主机沙箱不可用,拒绝执行: {str(exc)}",
"output": "",
"return_code": -1,
"timeout": timeout,
"elapsed_ms": int((time.time() - start_ts) * 1000)
}
except Exception as exc:
return {
"success": False,
"status": "error",
"error": f"执行失败: {str(exc)}",
"output": "",
"return_code": -1,
"timeout": timeout,
"elapsed_ms": int((time.time() - start_ts) * 1000)
}
async def install_package(self, package: str) -> Dict:
"""
安装Python包
Args:
package: 包名
Returns:
安装结果
"""
print(f"{OUTPUT_FORMATS['terminal']} 安装包: {package}")
# 使用通用 python3 占位,由 run_command 根据执行环境重写
command = f'python3 -m pip install {package}'
result = await self.run_command(command, timeout=120)
if result["success"]:
print(f"{OUTPUT_FORMATS['success']} 包安装成功: {package}")
else:
print(f"{OUTPUT_FORMATS['error']} 包安装失败: {package}")
return result
async def run_command(
self,
command: str,
working_dir: str = None,
timeout: int = None,
sandbox_write_access: bool = True,
network_permission: Optional[str] = None,
) -> Dict:
"""
执行终端命令
Args:
command: 要执行的命令
working_dir: 工作目录
timeout: 超时时间(秒)
Returns:
执行结果字典
"""
if timeout is None or timeout <= 0:
return {
"success": False,
"error": "timeout 参数必填且需大于0",
"status": "error",
"output": "timeout 参数缺失",
"return_code": -1
}
# 每次执行前重置工具容器(保持隔离),但下面改用一次性子进程执行,仍保留重置以兼容后续逻辑
self._reset_toolbox()
# 尝试复用活动终端的容器CLI 场景与 terminal_input 环境保持一致)
session_override = None
if not self.container_session:
session_override = self._resolve_active_container_session()
# 判断本次应在容器中执行:已绑定容器会话、复用终端容器或将使用工具箱容器
execution_in_container = self._will_use_container(session_override)
python_rewrite = self.container_python_cmd if execution_in_container else self.python_cmd
pip_rewrite = self._derive_pip_from_python(python_rewrite)
# 命令通过 ssh 在远程主机执行时,本地 Python/pip 路径对远程不适用,跳过重写
is_remote_ssh = command.lstrip().startswith("ssh ") or re.search(r"[;&|]\s*ssh\s", command) is not None
if not is_remote_ssh:
# 统一替换 python/python3为保障虚拟环境命中只替换独立单词
if re.search(r"\bpython3?\b", command):
command = re.sub(r"\bpython3?\b", python_rewrite, command)
# 同步替换 pip/pip3确保指向同一环境
if re.search(r"(?<![/.\w-])pip3?\b", command):
command = re.sub(r"(?<![/.\w-])pip3?\b", pip_rewrite, command)
# 验证命令
valid, error = self._validate_command(command)
if not valid:
return {
"success": False,
"error": error,
"output": "",
"return_code": -1
}
# 设置工作目录
try:
work_path = self._resolve_work_path(working_dir)
except ValueError:
return {
"success": False,
"error": "工作目录必须在项目文件夹内",
"output": "",
"return_code": -1
}
# 默认10s上限由 TERMINAL_COMMAND_TIMEOUT 控制
timeout = self._clamp_timeout(timeout, default=10, max_limit=TERMINAL_COMMAND_TIMEOUT)
print(f"{OUTPUT_FORMATS['terminal']} 执行命令: {command}")
print(f"{OUTPUT_FORMATS['info']} 工作目录: {work_path}")
start_ts = time.time()
# 优先在绑定的容器或活动终端的容器内执行,保证与实时终端环境一致
try:
if self.container_session or session_override:
result_payload = await self._run_command_subprocess(
command,
work_path,
timeout,
session_override=session_override,
sandbox_write_access=sandbox_write_access,
network_permission=network_permission,
)
else:
# 若未绑定用户容器,则使用工具箱容器(与终端相同镜像/预装包)
toolbox = self._get_toolbox()
try:
payload = await toolbox.run(command, work_path, timeout)
except asyncio.CancelledError:
# 任务被取消时强制关闭工具箱终端,避免后台命令继续运行
try:
toolbox.shutdown()
except Exception:
pass
raise
result_payload = self._format_toolbox_output(payload)
# 追加耗时信息以对齐接口
result_payload["elapsed_ms"] = int((time.time() - start_ts) * 1000)
result_payload["timeout"] = timeout
# 字符数检查(与主流程一致)
if result_payload.get("success") and "output" in result_payload:
char_count = len(result_payload["output"])
if char_count > MAX_RUN_COMMAND_CHARS:
return {
"success": False,
"error": f"结果内容过大,有{char_count}字符请使用限制字符数的获取内容方式根据程度选择10k以内的数",
"char_count": char_count,
"limit": MAX_RUN_COMMAND_CHARS,
"command": command
}
return result_payload
except asyncio.CancelledError:
return {
"success": False,
"message": "命令执行被用户取消",
"output": "",
"status": "cancelled",
"return_code": -1,
"timeout": timeout,
"elapsed_ms": int((time.time() - start_ts) * 1000)
}
# 改为一次性子进程执行,确保等待到超时或命令结束
result_payload = result_payload if result_payload is not None else await self._run_command_subprocess(
command, work_path, timeout, sandbox_write_access=sandbox_write_access, network_permission=network_permission
)
# 字符数检查
if result_payload.get("success") and "output" in result_payload:
char_count = len(result_payload["output"])
if char_count > MAX_RUN_COMMAND_CHARS:
return {
"success": False,
"error": f"结果内容过大,有{char_count}字符请使用限制字符数的获取内容方式根据程度选择10k以内的数",
"char_count": char_count,
"limit": MAX_RUN_COMMAND_CHARS,
"command": command
}
result_payload.setdefault("status", "completed" if result_payload.get("success") else "error")
result_payload["timeout"] = timeout
result_payload["elapsed_ms"] = int((time.time() - start_ts) * 1000)
return result_payload