513 lines
20 KiB
Python
513 lines
20 KiB
Python
"""子智能体任务管理（主进程内协程模式）。
|
||
|
||
子智能体不再作为独立子进程启动，而是作为 SubAgentManager 所在事件循环中的
|
||
asyncio.Task 运行。所有实际工具调用都通过主 WebTerminal 执行，因此自然复用
|
||
主进程的宿主机沙箱 / Docker 容器链路。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import threading
|
||
import time
|
||
from pathlib import Path, PurePosixPath
|
||
from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
||
|
||
from config import (
|
||
OUTPUT_FORMATS,
|
||
SUB_AGENT_DEFAULT_TIMEOUT,
|
||
SUB_AGENT_MAX_ACTIVE,
|
||
SUB_AGENT_MODELS_CONFIG_FILE,
|
||
SUB_AGENT_STATE_FILE,
|
||
SUB_AGENT_STATUS_POLL_INTERVAL,
|
||
SUB_AGENT_TASKS_BASE_DIR,
|
||
)
|
||
from utils.logger import setup_logger
|
||
from modules.sub_agent.task import SubAgentTask
|
||
from modules.sub_agent.prompts import build_user_message, build_system_prompt
|
||
from modules.sub_agent.tools import handle_search_workspace, handle_read_mediafile
|
||
from modules.sub_agent.state import SubAgentStateMixin
|
||
from modules.sub_agent.stats import SubAgentStatsMixin
|
||
from modules.sub_agent.creation import SubAgentCreationMixin
|
||
|
||
if TYPE_CHECKING:
|
||
from core.web_terminal import WebTerminal
|
||
from modules.user_container_manager import ContainerHandle
|
||
|
||
logger = setup_logger(__name__)
|
||
TERMINAL_STATUSES = {"completed", "failed", "timeout"}
|
||
|
||
|
||
class SubAgentManager(SubAgentStateMixin, SubAgentStatsMixin, SubAgentCreationMixin):
|
||
"""负责主智能体与子智能体的任务调度（协程模式）。"""
|
||
|
||
def __init__(
|
||
self,
|
||
project_path: str,
|
||
data_dir: str,
|
||
container_session: Optional["ContainerHandle"] = None,
|
||
):
|
||
self.project_path = Path(project_path).resolve()
|
||
self.data_dir = Path(data_dir).resolve()
|
||
self.base_dir = Path(SUB_AGENT_TASKS_BASE_DIR).resolve()
|
||
self.state_file = Path(SUB_AGENT_STATE_FILE).resolve()
|
||
self.models_config_file = SUB_AGENT_MODELS_CONFIG_FILE
|
||
self.container_session: Optional["ContainerHandle"] = container_session
|
||
self.host_execution_mode: str = "sandbox"
|
||
self.terminal: Optional["WebTerminal"] = None
|
||
|
||
self.base_dir.mkdir(parents=True, exist_ok=True)
|
||
self.state_file.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
self.tasks: Dict[str, Dict[str, Any]] = {}
|
||
self.conversation_agents: Dict[str, List[int]] = {}
|
||
self._running_tasks: Dict[str, asyncio.Task] = {}
|
||
self._event_loop: Optional[asyncio.AbstractEventLoop] = None
|
||
self._loop_thread: Optional[threading.Thread] = None
|
||
self._state_lock = threading.Lock()
|
||
|
||
self._load_state()
|
||
try:
|
||
self.reconcile_task_states()
|
||
except Exception:
|
||
pass
|
||
|
||
# ------------------------------------------------------------------
|
||
# 生命周期与事件循环
|
||
# ------------------------------------------------------------------
|
||
def _ensure_event_loop(self) -> asyncio.AbstractEventLoop:
|
||
"""确保有一个独立的后台事件循环供子智能体使用。"""
|
||
if self._event_loop is not None and not self._event_loop.is_closed():
|
||
return self._event_loop
|
||
|
||
loop = asyncio.new_event_loop()
|
||
self._event_loop = loop
|
||
|
||
def run_loop():
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
loop.run_forever()
|
||
finally:
|
||
try:
|
||
loop.close()
|
||
except Exception:
|
||
pass
|
||
|
||
thread = threading.Thread(target=run_loop, name="sub-agent-loop", daemon=True)
|
||
thread.start()
|
||
self._loop_thread = thread
|
||
return loop
|
||
|
||
async def _create_task(self, coro):
|
||
"""在事件循环内部把协程包装为 Task。"""
|
||
return asyncio.create_task(coro)
|
||
|
||
def _run_coro(self, coro):
|
||
"""在后台事件循环中调度一个协程并返回 asyncio.Task。"""
|
||
loop = self._ensure_event_loop()
|
||
# 先提交创建 Task 的协程，阻塞等待拿到 Task 句柄
|
||
future = asyncio.run_coroutine_threadsafe(self._create_task(coro), loop)
|
||
return future.result(timeout=10)
|
||
|
||
def set_terminal(self, terminal: "WebTerminal") -> None:
|
||
"""注入主终端引用，用于工具执行代理。"""
|
||
self.terminal = terminal
|
||
|
||
def set_container_session(self, session: Optional["ContainerHandle"]):
|
||
"""更新容器会话信息。"""
|
||
self.container_session = session
|
||
|
||
def set_host_execution_mode(self, mode: str) -> None:
|
||
normalized = str(mode or "").strip().lower()
|
||
self.host_execution_mode = "direct" if normalized == "direct" else "sandbox"
|
||
|
||
# ------------------------------------------------------------------
|
||
# 公共方法
|
||
# ------------------------------------------------------------------
|
||
def create_sub_agent(
|
||
self,
|
||
*,
|
||
agent_id: int,
|
||
summary: str,
|
||
task: str,
|
||
deliverables_dir: str,
|
||
timeout_seconds: Optional[int] = None,
|
||
conversation_id: Optional[str] = None,
|
||
run_in_background: bool = False,
|
||
model_key: Optional[str] = None,
|
||
thinking_mode: Optional[str] = None,
|
||
) -> Dict:
|
||
"""创建子智能体任务并启动协程。"""
|
||
validation_error = self._validate_create_params(agent_id, summary, task, deliverables_dir)
|
||
if validation_error:
|
||
return {"success": False, "error": validation_error}
|
||
|
||
if not thinking_mode:
|
||
return {"success": False, "error": "缺少 thinking_mode 参数，必须指定 fast 或 thinking"}
|
||
if thinking_mode not in {"fast", "thinking"}:
|
||
return {"success": False, "error": "thinking_mode 仅支持 fast 或 thinking"}
|
||
|
||
if not conversation_id:
|
||
return {"success": False, "error": "缺少对话ID，无法创建子智能体"}
|
||
|
||
if not self._ensure_agent_slot_available(conversation_id, agent_id):
|
||
return {
|
||
"success": False,
|
||
"error": f"该对话已使用过编号 {agent_id}，请更换新的子智能体代号。"
|
||
}
|
||
|
||
if self._active_task_count(conversation_id) >= SUB_AGENT_MAX_ACTIVE:
|
||
return {
|
||
"success": False,
|
||
"error": f"该对话已存在 {SUB_AGENT_MAX_ACTIVE} 个运行中的子智能体，请稍后再试。",
|
||
}
|
||
|
||
task_id = self._generate_task_id(agent_id)
|
||
task_root = self.base_dir / task_id
|
||
task_root.mkdir(parents=True, exist_ok=True)
|
||
|
||
try:
|
||
deliverables_path = self._resolve_deliverables_dir(deliverables_dir)
|
||
except ValueError as exc:
|
||
return {"success": False, "error": str(exc)}
|
||
|
||
task_file = task_root / "task.txt"
|
||
system_prompt_file = task_root / "system_prompt.txt"
|
||
output_file = task_root / "output.json"
|
||
stats_file = task_root / "stats.json"
|
||
progress_file = task_root / "progress.jsonl"
|
||
conversation_file = task_root / "conversation.json"
|
||
|
||
prompt_workspace = self._get_runtime_path(self.project_path)
|
||
deliverables_display = self._get_runtime_path(deliverables_path)
|
||
user_message = build_user_message(agent_id, summary, task, deliverables_display, timeout_seconds or SUB_AGENT_DEFAULT_TIMEOUT)
|
||
task_file.write_text(user_message, encoding="utf-8")
|
||
|
||
system_prompt = build_system_prompt(prompt_workspace)
|
||
system_prompt_file.write_text(system_prompt, encoding="utf-8")
|
||
|
||
timeout_seconds = timeout_seconds or SUB_AGENT_DEFAULT_TIMEOUT
|
||
|
||
task_record = {
|
||
"task_id": task_id,
|
||
"agent_id": agent_id,
|
||
"summary": summary,
|
||
"task": task,
|
||
"status": "running",
|
||
"deliverables_dir": str(deliverables_path),
|
||
"timeout_seconds": timeout_seconds,
|
||
"thinking_mode": thinking_mode,
|
||
"created_at": time.time(),
|
||
"updated_at": time.time(),
|
||
"conversation_id": conversation_id,
|
||
"run_in_background": run_in_background,
|
||
"task_root": str(task_root),
|
||
"output_file": str(output_file),
|
||
"stats_file": str(stats_file),
|
||
"progress_file": str(progress_file),
|
||
"conversation_file": str(conversation_file),
|
||
"execution_mode": "in_process",
|
||
"container_name": None,
|
||
}
|
||
self.tasks[task_id] = task_record
|
||
self._mark_agent_id_used(conversation_id, agent_id)
|
||
self._save_state()
|
||
|
||
sub_agent = SubAgentTask(
|
||
manager=self,
|
||
task_record=task_record,
|
||
task_message=user_message,
|
||
system_prompt=system_prompt,
|
||
model_key=model_key,
|
||
thinking_mode=thinking_mode,
|
||
)
|
||
task_coro = sub_agent.run()
|
||
asyncio_task = self._run_coro(task_coro)
|
||
sub_agent._task = asyncio_task
|
||
self._running_tasks[task_id] = asyncio_task
|
||
|
||
def _on_done(fut):
|
||
self._running_tasks.pop(task_id, None)
|
||
self.reconcile_task_states(conversation_id=conversation_id)
|
||
|
||
asyncio_task.add_done_callback(_on_done)
|
||
|
||
message = f"子智能体{agent_id} 已创建，任务ID: {task_id}"
|
||
print(f"{OUTPUT_FORMATS['info']} {message}")
|
||
|
||
return {
|
||
"success": True,
|
||
"task_id": task_id,
|
||
"agent_id": agent_id,
|
||
"status": "running",
|
||
"message": message,
|
||
"deliverables_dir": str(deliverables_path),
|
||
"run_in_background": run_in_background,
|
||
}
|
||
|
||
def wait_for_completion(
|
||
self,
|
||
*,
|
||
task_id: Optional[str] = None,
|
||
agent_id: Optional[int] = None,
|
||
timeout_seconds: Optional[int] = None,
|
||
) -> Dict:
|
||
"""阻塞等待子智能体完成或超时。"""
|
||
task = self._select_task(task_id, agent_id)
|
||
if not task:
|
||
return {"success": False, "error": "未找到对应的子智能体任务"}
|
||
|
||
if task.get("status") in TERMINAL_STATUSES or task.get("status") == "terminated":
|
||
if task.get("final_result"):
|
||
return task["final_result"]
|
||
return {"success": False, "status": task.get("status"), "message": "子智能体已结束。"}
|
||
|
||
real_task_id = task["task_id"]
|
||
deadline = time.time() + (timeout_seconds or task.get("timeout_seconds") or SUB_AGENT_DEFAULT_TIMEOUT)
|
||
|
||
while time.time() < deadline:
|
||
self.reconcile_task_states()
|
||
# 关键：其他线程（如前端轮询 /api/sub_agents）可能调用 _load_state()
|
||
# 并替换 self.tasks 字典，导致旧 task 引用失效。每次循环重新获取引用。
|
||
task = self.tasks.get(real_task_id)
|
||
if not task:
|
||
return {"success": False, "error": "未找到对应的子智能体任务"}
|
||
running_task = self._running_tasks.get(real_task_id)
|
||
status = task.get("status")
|
||
|
||
# 已到达终态：返回最终结果（持续 reconcile 直到 final_result 就绪）
|
||
if status in TERMINAL_STATUSES or status == "terminated":
|
||
if task.get("final_result"):
|
||
return task["final_result"]
|
||
# 终态但 final_result 尚未写入，短暂等待后重试
|
||
time.sleep(SUB_AGENT_STATUS_POLL_INTERVAL)
|
||
self.reconcile_task_states()
|
||
task = self.tasks.get(real_task_id) or task
|
||
if task.get("final_result"):
|
||
return task["final_result"]
|
||
return {"success": False, "status": status, "message": "子智能体已结束，但未获取到结果。"}
|
||
|
||
# asyncio Task 已结束但状态可能还没同步：等待 final_result 就绪
|
||
if running_task and running_task.done():
|
||
self.reconcile_task_states()
|
||
task = self.tasks.get(real_task_id) or task
|
||
if task.get("final_result"):
|
||
return task["final_result"]
|
||
# 结果尚未落盘，继续轮询，避免把「已创建」误判为失败
|
||
time.sleep(SUB_AGENT_STATUS_POLL_INTERVAL)
|
||
continue
|
||
|
||
time.sleep(SUB_AGENT_STATUS_POLL_INTERVAL)
|
||
|
||
return self._handle_timeout(task)
|
||
|
||
def terminate_sub_agent(
|
||
self,
|
||
*,
|
||
task_id: Optional[str] = None,
|
||
agent_id: Optional[int] = None,
|
||
) -> Dict:
|
||
"""强制关闭指定子智能体。"""
|
||
task = self._select_task(task_id, agent_id)
|
||
if not task:
|
||
return {"success": False, "error": "未找到对应的子智能体任务"}
|
||
|
||
task_id = task["task_id"]
|
||
running_task = self._running_tasks.pop(task_id, None)
|
||
if running_task and not running_task.done():
|
||
running_task.cancel()
|
||
deadline = time.time() + 5
|
||
while not running_task.done() and time.time() < deadline:
|
||
time.sleep(0.05)
|
||
|
||
self._mark_task_terminated(
|
||
task,
|
||
message="子智能体已被强制关闭。",
|
||
system_message=f"🛑 子智能体{task.get('agent_id')} 已被手动关闭。",
|
||
notified=True,
|
||
)
|
||
self._save_state()
|
||
|
||
return {
|
||
"success": True,
|
||
"task_id": task_id,
|
||
"message": "子智能体已被强制关闭。",
|
||
"system_message": f"🛑 子智能体{task.get('agent_id')} 已被手动关闭。",
|
||
}
|
||
|
||
def get_sub_agent_status(
|
||
self,
|
||
*,
|
||
agent_ids: Optional[List[int]] = None,
|
||
) -> Dict:
|
||
"""获取指定子智能体的详细状态。
|
||
|
||
对于已结束（completed/failed/timeout/terminated）的子智能体，同样返回其
|
||
最终状态，而不是返回「不存在」。
|
||
"""
|
||
if not agent_ids:
|
||
return {"success": False, "error": "必须指定至少一个agent_id"}
|
||
|
||
def _find_task_by_agent_id(aid: int):
|
||
# 先查运行中/待运行的任务
|
||
task = self._select_task(None, aid)
|
||
if task:
|
||
return task
|
||
# 再查已结束的任务（按创建时间取最新一条）
|
||
candidates = [
|
||
t for t in self.tasks.values()
|
||
if t.get("agent_id") == aid
|
||
]
|
||
if not candidates:
|
||
return None
|
||
candidates.sort(key=lambda item: item.get("created_at", 0), reverse=True)
|
||
return candidates[0]
|
||
|
||
results = []
|
||
for agent_id in agent_ids:
|
||
task = _find_task_by_agent_id(agent_id)
|
||
if not task:
|
||
results.append({
|
||
"agent_id": agent_id,
|
||
"found": False,
|
||
"error": "子智能体不存在",
|
||
})
|
||
continue
|
||
|
||
status = task.get("status")
|
||
if status not in TERMINAL_STATUSES.union({"terminated"}):
|
||
self._check_task_status(task)
|
||
status = task.get("status")
|
||
|
||
stats = {}
|
||
stats_file = Path(task.get("stats_file", ""))
|
||
if stats_file.exists():
|
||
try:
|
||
stats = json.loads(stats_file.read_text(encoding="utf-8"))
|
||
except Exception:
|
||
pass
|
||
stats_summary = self._build_stats_summary(stats)
|
||
|
||
results.append({
|
||
"agent_id": agent_id,
|
||
"found": True,
|
||
"task_id": task["task_id"],
|
||
"status": status,
|
||
"summary": task.get("summary"),
|
||
"created_at": task.get("created_at"),
|
||
"updated_at": task.get("updated_at"),
|
||
"deliverables_dir": task.get("deliverables_dir"),
|
||
"stats": stats,
|
||
"stats_summary": stats_summary,
|
||
"final_result": task.get("final_result"),
|
||
})
|
||
|
||
return {"success": True, "results": results}
|
||
|
||
def poll_updates(self) -> List[Dict]:
|
||
"""检查运行中的子智能体任务，返回新完成的结果。"""
|
||
updates: List[Dict] = []
|
||
self.reconcile_task_states()
|
||
pending_tasks = [
|
||
task for task in self.tasks.values()
|
||
if task.get("status") not in TERMINAL_STATUSES.union({"terminated"})
|
||
]
|
||
if not pending_tasks:
|
||
return updates
|
||
|
||
state_changed = False
|
||
for task in pending_tasks:
|
||
result = self._check_task_status(task)
|
||
if result["status"] in TERMINAL_STATUSES:
|
||
updates.append(result)
|
||
state_changed = True
|
||
|
||
if state_changed:
|
||
self._save_state()
|
||
return updates
|
||
|
||
def lookup_task(self, *, task_id: Optional[str] = None, agent_id: Optional[int] = None) -> Optional[Dict]:
|
||
"""只读查询任务信息。"""
|
||
task = self._select_task(task_id, agent_id)
|
||
if not task:
|
||
return None
|
||
return {
|
||
"task_id": task.get("task_id"),
|
||
"agent_id": task.get("agent_id"),
|
||
"status": task.get("status"),
|
||
"timeout_seconds": task.get("timeout_seconds"),
|
||
"conversation_id": task.get("conversation_id"),
|
||
}
|
||
|
||
def get_overview(self, conversation_id: Optional[str] = None) -> List[Dict[str, Any]]:
|
||
"""返回子智能体任务概览，用于前端展示。"""
|
||
self.reconcile_task_states(conversation_id=conversation_id)
|
||
overview: List[Dict[str, Any]] = []
|
||
for task_id, task in self.tasks.items():
|
||
if conversation_id and task.get("conversation_id") != conversation_id:
|
||
continue
|
||
|
||
snapshot = {
|
||
"task_id": task_id,
|
||
"agent_id": task.get("agent_id"),
|
||
"summary": task.get("summary"),
|
||
"status": task.get("status"),
|
||
"created_at": task.get("created_at"),
|
||
"updated_at": task.get("updated_at"),
|
||
"target_dir": task.get("target_project_dir"),
|
||
"last_tool": task.get("last_tool"),
|
||
"deliverables_dir": task.get("deliverables_dir"),
|
||
"copied_path": task.get("copied_path"),
|
||
"conversation_id": task.get("conversation_id"),
|
||
"sub_conversation_id": task.get("sub_conversation_id"),
|
||
}
|
||
|
||
if snapshot["status"] in TERMINAL_STATUSES or snapshot["status"] == "terminated":
|
||
final_result = task.get("final_result") or {}
|
||
snapshot["final_message"] = final_result.get("system_message") or final_result.get("message")
|
||
snapshot["success"] = final_result.get("success")
|
||
|
||
overview.append(snapshot)
|
||
|
||
overview.sort(key=lambda item: item.get("created_at") or 0, reverse=True)
|
||
return overview
|
||
|
||
# ------------------------------------------------------------------
|
||
# 工具执行代理
|
||
# ------------------------------------------------------------------
|
||
async def execute_tool_for_sub_agent(self, tool_name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
|
||
"""代表子智能体在主进程中执行工具。"""
|
||
if not self.terminal:
|
||
return {"success": False, "error": "子智能体管理器未绑定终端，无法执行工具"}
|
||
|
||
try:
|
||
if tool_name == "search_workspace":
|
||
return await handle_search_workspace(self.project_path, self.terminal, arguments)
|
||
if tool_name == "read_mediafile":
|
||
return await handle_read_mediafile(self.project_path, arguments)
|
||
|
||
# 其余工具直接走主进程 handle_tool_call，自然经过沙箱/容器/权限链路
|
||
result_text = await self.terminal.handle_tool_call(tool_name, arguments)
|
||
try:
|
||
return json.loads(result_text)
|
||
except Exception:
|
||
return {"success": True, "output": result_text}
|
||
except Exception as exc:
|
||
logger.exception(f"[SubAgent] 工具执行异常: {tool_name}")
|
||
return {"success": False, "error": f"工具执行异常: {exc}"}
|
||
|
||
def _get_runtime_path(self, host_path: Path) -> str:
|
||
"""将宿主机路径映射为容器内路径（仅用于提示展示）。"""
|
||
if not self.container_session or getattr(self.container_session, "mode", None) != "docker":
|
||
return str(host_path)
|
||
mount_path = (getattr(self.container_session, "mount_path", None) or "/workspace").rstrip("/") or "/workspace"
|
||
try:
|
||
relative = host_path.resolve().relative_to(self.project_path)
|
||
except Exception:
|
||
return mount_path
|
||
if str(relative) in {"", "."}:
|
||
return mount_path
|
||
return str(PurePosixPath(mount_path) / PurePosixPath(relative.as_posix()))
|
||
|