501 lines
19 KiB
Python
501 lines
19 KiB
Python
"""后台 run_command 任务管理。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import re
|
||
import shutil
|
||
import signal
|
||
import subprocess
|
||
import threading
|
||
import time
|
||
import uuid
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
|
||
from config import MAX_RUN_COMMAND_CHARS
|
||
|
||
|
||
TERMINAL_STATUSES = {"completed", "failed", "timeout", "cancelled"}
|
||
|
||
|
||
class BackgroundCommandManager:
|
||
"""管理 run_command 的后台执行、轮询与等待。"""
|
||
|
||
def __init__(self, project_path: str):
|
||
self.project_path = Path(project_path).resolve()
|
||
self._records: Dict[str, Dict[str, Any]] = {}
|
||
self._lock = threading.RLock()
|
||
self._cv = threading.Condition(self._lock)
|
||
|
||
def create_background_command(
|
||
self,
|
||
*,
|
||
terminal_ops,
|
||
command: str,
|
||
timeout: Optional[int],
|
||
conversation_id: Optional[str],
|
||
wait_seconds: float = 5.0,
|
||
) -> Dict[str, Any]:
|
||
"""启动后台命令；先等待一小段时间返回已有输出。"""
|
||
if timeout is None or float(timeout) <= 0:
|
||
return {
|
||
"success": False,
|
||
"error": "timeout 参数必填且需大于0",
|
||
"status": "error",
|
||
"output": "timeout 参数缺失",
|
||
"return_code": -1,
|
||
}
|
||
|
||
timeout_value = min(int(timeout), 3600)
|
||
if timeout_value <= 0:
|
||
timeout_value = 1
|
||
|
||
session_override = None
|
||
if not getattr(terminal_ops, "container_session", None):
|
||
session_override = terminal_ops._resolve_active_container_session()
|
||
|
||
execution_in_container = terminal_ops._will_use_container(session_override)
|
||
python_rewrite = terminal_ops.container_python_cmd if execution_in_container else terminal_ops.python_cmd
|
||
pip_rewrite = terminal_ops._derive_pip_from_python(python_rewrite)
|
||
final_command = command
|
||
if re.search(r"\bpython3?\b", final_command):
|
||
final_command = re.sub(r"\bpython3?\b", python_rewrite, final_command)
|
||
if re.search(r"(?<![/.\w-])pip3?\b", final_command):
|
||
final_command = re.sub(r"(?<![/.\w-])pip3?\b", pip_rewrite, final_command)
|
||
|
||
valid, error = terminal_ops._validate_command(final_command)
|
||
if not valid:
|
||
return {
|
||
"success": False,
|
||
"error": error,
|
||
"status": "error",
|
||
"output": "",
|
||
"return_code": -1,
|
||
}
|
||
|
||
try:
|
||
work_path = terminal_ops._resolve_work_path(None)
|
||
except ValueError:
|
||
return {
|
||
"success": False,
|
||
"error": "工作目录必须在项目文件夹内",
|
||
"status": "error",
|
||
"output": "",
|
||
"return_code": -1,
|
||
}
|
||
|
||
command_id = f"cmd_{int(time.time())}_{uuid.uuid4().hex[:8]}"
|
||
now = time.time()
|
||
with self._lock:
|
||
self._records[command_id] = {
|
||
"command_id": command_id,
|
||
"conversation_id": conversation_id,
|
||
"status": "running",
|
||
"command": final_command,
|
||
"timeout": timeout_value,
|
||
"created_at": now,
|
||
"updated_at": now,
|
||
"finished_at": None,
|
||
"stdout_chunks": [],
|
||
"stderr_chunks": [],
|
||
"truncated": False,
|
||
"result": None,
|
||
"notified": False,
|
||
"claimed_by_sleep": False,
|
||
"pid": None,
|
||
}
|
||
|
||
thread = threading.Thread(
|
||
target=self._run_command_thread,
|
||
kwargs={
|
||
"command_id": command_id,
|
||
"command": final_command,
|
||
"work_path": work_path,
|
||
"timeout": timeout_value,
|
||
"session": session_override or getattr(terminal_ops, "container_session", None),
|
||
"python_env": getattr(terminal_ops, "_python_env", None) or {},
|
||
},
|
||
name=f"bg-run-command-{command_id}",
|
||
daemon=True,
|
||
)
|
||
thread.start()
|
||
|
||
wait_seconds = max(0.1, min(float(wait_seconds or 5.0), 5.0))
|
||
deadline = time.time() + wait_seconds
|
||
with self._cv:
|
||
while time.time() < deadline:
|
||
rec = self._records.get(command_id)
|
||
if not rec:
|
||
break
|
||
if rec.get("status") in TERMINAL_STATUSES:
|
||
break
|
||
remaining = deadline - time.time()
|
||
if remaining <= 0:
|
||
break
|
||
self._cv.wait(timeout=remaining)
|
||
|
||
with self._lock:
|
||
rec = self._records.get(command_id)
|
||
if not rec:
|
||
return {
|
||
"success": False,
|
||
"status": "error",
|
||
"error": "后台任务记录丢失",
|
||
"output": "",
|
||
"return_code": -1,
|
||
}
|
||
|
||
if rec.get("status") in TERMINAL_STATUSES and isinstance(rec.get("result"), dict):
|
||
rec["claimed_by_sleep"] = True
|
||
rec["notified"] = True
|
||
rec["updated_at"] = time.time()
|
||
result = dict(rec["result"])
|
||
result["command_id"] = command_id
|
||
result["run_in_background"] = True
|
||
result["background_task_created"] = False
|
||
result["message"] = "命令在5秒内完成，未创建后台任务"
|
||
return result
|
||
|
||
output = self._build_current_output(rec)
|
||
return {
|
||
"success": True,
|
||
"status": "running_background",
|
||
"command_id": command_id,
|
||
"command": final_command,
|
||
"message": "后台命令已创建；以下为当前已捕获输出。",
|
||
"output": output,
|
||
"return_code": None,
|
||
"timeout": timeout_value,
|
||
"elapsed_ms": int((time.time() - now) * 1000),
|
||
"run_in_background": True,
|
||
"background_task_created": True,
|
||
}
|
||
|
||
def _run_command_thread(
|
||
self,
|
||
*,
|
||
command_id: str,
|
||
command: str,
|
||
work_path: Path,
|
||
timeout: int,
|
||
session,
|
||
python_env: Dict[str, str],
|
||
) -> None:
|
||
start_ts = time.time()
|
||
process: Optional[subprocess.Popen] = None
|
||
stdout_buf: List[str] = []
|
||
stderr_buf: List[str] = []
|
||
status = "failed"
|
||
return_code = -1
|
||
message: Optional[str] = None
|
||
|
||
try:
|
||
exec_cmd: Optional[List[str]] = None
|
||
use_shell = True
|
||
env = os.environ.copy()
|
||
env.setdefault("PYTHONUNBUFFERED", "1")
|
||
if python_env:
|
||
env.update(python_env)
|
||
|
||
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"
|
||
if not container_name:
|
||
raise RuntimeError("容器模式下缺少 container_name")
|
||
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
|
||
|
||
if use_shell:
|
||
process = subprocess.Popen(
|
||
command,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
cwd=str(work_path),
|
||
shell=True,
|
||
env=env,
|
||
start_new_session=True,
|
||
text=True,
|
||
errors="replace",
|
||
bufsize=1,
|
||
)
|
||
else:
|
||
process = subprocess.Popen(
|
||
exec_cmd,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
env=env,
|
||
start_new_session=True,
|
||
text=True,
|
||
errors="replace",
|
||
bufsize=1,
|
||
)
|
||
|
||
with self._lock:
|
||
rec = self._records.get(command_id)
|
||
if rec is not None:
|
||
rec["pid"] = process.pid
|
||
rec["updated_at"] = time.time()
|
||
|
||
def _reader(stream, collector, rec_key: str):
|
||
try:
|
||
while True:
|
||
line = stream.readline()
|
||
if line == "":
|
||
break
|
||
collector.append(line)
|
||
with self._lock:
|
||
rec = self._records.get(command_id)
|
||
if rec is not None:
|
||
rec[rec_key].append(line)
|
||
rec["updated_at"] = time.time()
|
||
except Exception:
|
||
return
|
||
|
||
t_out = threading.Thread(target=_reader, args=(process.stdout, stdout_buf, "stdout_chunks"), daemon=True)
|
||
t_err = threading.Thread(target=_reader, args=(process.stderr, stderr_buf, "stderr_chunks"), daemon=True)
|
||
t_out.start()
|
||
t_err.start()
|
||
|
||
try:
|
||
process.wait(timeout=timeout)
|
||
return_code = process.returncode if process.returncode is not None else -1
|
||
if return_code == 0:
|
||
status = "completed"
|
||
else:
|
||
status = "failed"
|
||
message = f"命令执行失败 (返回码: {return_code})"
|
||
except subprocess.TimeoutExpired:
|
||
status = "timeout"
|
||
message = f"命令执行超时 ({timeout}秒)"
|
||
try:
|
||
os.killpg(process.pid, signal.SIGINT)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
process.wait(timeout=2)
|
||
except subprocess.TimeoutExpired:
|
||
try:
|
||
os.killpg(process.pid, signal.SIGKILL)
|
||
except Exception:
|
||
try:
|
||
process.kill()
|
||
except Exception:
|
||
pass
|
||
process.wait(timeout=2)
|
||
return_code = process.returncode if process.returncode is not None else -1
|
||
|
||
t_out.join(timeout=1)
|
||
t_err.join(timeout=1)
|
||
|
||
except Exception as exc:
|
||
status = "failed"
|
||
message = f"执行失败: {exc}"
|
||
finally:
|
||
combined_output = "".join(stdout_buf + stderr_buf)
|
||
truncated = False
|
||
if MAX_RUN_COMMAND_CHARS and len(combined_output) > MAX_RUN_COMMAND_CHARS:
|
||
combined_output = combined_output[-MAX_RUN_COMMAND_CHARS:]
|
||
truncated = True
|
||
|
||
success = status == "completed"
|
||
result = {
|
||
"success": success,
|
||
"status": status,
|
||
"command": command,
|
||
"output": combined_output,
|
||
"return_code": return_code,
|
||
"truncated": truncated,
|
||
"timeout": timeout,
|
||
"elapsed_ms": int((time.time() - start_ts) * 1000),
|
||
"command_id": command_id,
|
||
"run_in_background": True,
|
||
}
|
||
if message:
|
||
result["message"] = message
|
||
|
||
with self._cv:
|
||
rec = self._records.get(command_id)
|
||
if rec is not None:
|
||
rec["status"] = status
|
||
rec["result"] = result
|
||
rec["truncated"] = truncated
|
||
rec["updated_at"] = time.time()
|
||
rec["finished_at"] = time.time()
|
||
self._cv.notify_all()
|
||
|
||
def _build_current_output(self, rec: Dict[str, Any]) -> str:
|
||
output = "".join((rec.get("stdout_chunks") or []) + (rec.get("stderr_chunks") or []))
|
||
if MAX_RUN_COMMAND_CHARS and len(output) > MAX_RUN_COMMAND_CHARS:
|
||
return output[-MAX_RUN_COMMAND_CHARS:]
|
||
return output
|
||
|
||
def wait_for_completion(self, command_id: str, timeout_seconds: Optional[float] = None, claim: bool = False) -> Dict[str, Any]:
|
||
"""阻塞等待后台命令完成。"""
|
||
with self._cv:
|
||
rec = self._records.get(command_id)
|
||
if not rec:
|
||
return {"success": False, "error": f"未找到后台命令: {command_id}", "status": "error"}
|
||
|
||
status = rec.get("status")
|
||
if status in TERMINAL_STATUSES and isinstance(rec.get("result"), dict):
|
||
if claim:
|
||
rec["claimed_by_sleep"] = True
|
||
return dict(rec["result"])
|
||
|
||
wait_limit = timeout_seconds
|
||
if wait_limit is None:
|
||
created = float(rec.get("created_at") or time.time())
|
||
timeout_val = float(rec.get("timeout") or 0)
|
||
if timeout_val > 0:
|
||
wait_limit = max(1.0, (created + timeout_val + 5.0) - time.time())
|
||
else:
|
||
wait_limit = 3605.0
|
||
|
||
deadline = time.time() + max(0.1, float(wait_limit))
|
||
while time.time() < deadline:
|
||
rec = self._records.get(command_id)
|
||
if not rec:
|
||
return {"success": False, "error": f"未找到后台命令: {command_id}", "status": "error"}
|
||
status = rec.get("status")
|
||
if status in TERMINAL_STATUSES and isinstance(rec.get("result"), dict):
|
||
if claim:
|
||
rec["claimed_by_sleep"] = True
|
||
return dict(rec["result"])
|
||
self._cv.wait(timeout=min(0.5, deadline - time.time()))
|
||
|
||
rec = self._records.get(command_id)
|
||
if not rec:
|
||
return {"success": False, "error": f"未找到后台命令: {command_id}", "status": "error"}
|
||
return {
|
||
"success": False,
|
||
"status": "timeout",
|
||
"command_id": command_id,
|
||
"message": "等待后台命令完成超时",
|
||
"output": self._build_current_output(rec),
|
||
"return_code": rec.get("result", {}).get("return_code") if isinstance(rec.get("result"), dict) else None,
|
||
}
|
||
|
||
def poll_updates(self, conversation_id: Optional[str] = None) -> List[Dict[str, Any]]:
|
||
"""获取未通知且未被 sleep 领取的已完成任务。"""
|
||
updates: List[Dict[str, Any]] = []
|
||
with self._lock:
|
||
for rec in self._records.values():
|
||
if conversation_id and rec.get("conversation_id") != conversation_id:
|
||
continue
|
||
if rec.get("status") not in TERMINAL_STATUSES:
|
||
continue
|
||
if rec.get("notified") or rec.get("claimed_by_sleep"):
|
||
continue
|
||
payload = rec.get("result")
|
||
if isinstance(payload, dict):
|
||
updates.append(dict(payload))
|
||
|
||
updates.sort(key=lambda item: self._records.get(item.get("command_id"), {}).get("updated_at", 0))
|
||
return updates
|
||
|
||
def mark_notified(self, command_id: str):
|
||
with self._lock:
|
||
rec = self._records.get(command_id)
|
||
if rec:
|
||
rec["notified"] = True
|
||
rec["updated_at"] = time.time()
|
||
|
||
def mark_claimed(self, command_id: str):
|
||
with self._lock:
|
||
rec = self._records.get(command_id)
|
||
if rec:
|
||
rec["claimed_by_sleep"] = True
|
||
rec["updated_at"] = time.time()
|
||
|
||
def get_record(self, command_id: str) -> Optional[Dict[str, Any]]:
|
||
with self._lock:
|
||
rec = self._records.get(command_id)
|
||
if not rec:
|
||
return None
|
||
return dict(rec)
|
||
|
||
def get_record_with_output(self, command_id: str) -> Optional[Dict[str, Any]]:
|
||
"""获取单条后台命令记录，并附带当前可读输出。"""
|
||
with self._lock:
|
||
rec = self._records.get(command_id)
|
||
if not rec:
|
||
return None
|
||
payload = dict(rec)
|
||
payload["output"] = self._build_current_output(rec)
|
||
return payload
|
||
|
||
def list_records(
|
||
self,
|
||
*,
|
||
conversation_id: Optional[str] = None,
|
||
limit: int = 200,
|
||
) -> List[Dict[str, Any]]:
|
||
"""列出后台命令记录（按创建时间倒序）。"""
|
||
with self._lock:
|
||
items: List[Dict[str, Any]] = []
|
||
for rec in self._records.values():
|
||
if conversation_id and rec.get("conversation_id") != conversation_id:
|
||
continue
|
||
item = dict(rec)
|
||
item["output"] = self._build_current_output(rec)
|
||
items.append(item)
|
||
items.sort(key=lambda x: float(x.get("created_at") or 0), reverse=True)
|
||
max_limit = max(1, min(int(limit or 200), 1000))
|
||
return items[:max_limit]
|
||
|
||
def has_pending_for_conversation(self, conversation_id: Optional[str]) -> bool:
|
||
if not conversation_id:
|
||
return False
|
||
with self._lock:
|
||
for rec in self._records.values():
|
||
if rec.get("conversation_id") != conversation_id:
|
||
continue
|
||
if rec.get("status") == "running":
|
||
return True
|
||
if rec.get("status") in TERMINAL_STATUSES and (not rec.get("notified")) and (not rec.get("claimed_by_sleep")):
|
||
return True
|
||
return False
|
||
|
||
def list_waiting_items(self, conversation_id: Optional[str]) -> List[Dict[str, Any]]:
|
||
items: List[Dict[str, Any]] = []
|
||
if not conversation_id:
|
||
return items
|
||
with self._lock:
|
||
for rec in self._records.values():
|
||
if rec.get("conversation_id") != conversation_id:
|
||
continue
|
||
if rec.get("status") == "running" or (
|
||
rec.get("status") in TERMINAL_STATUSES
|
||
and not rec.get("notified")
|
||
and not rec.get("claimed_by_sleep")
|
||
):
|
||
items.append({
|
||
"command_id": rec.get("command_id"),
|
||
"command": rec.get("command"),
|
||
"status": rec.get("status"),
|
||
})
|
||
items.sort(key=lambda x: x.get("command_id") or "")
|
||
return items
|