agent-Specialization/modules/sub_agent/manager.py

"""子智能体任务管理（主进程内协程模式）。

子智能体不再作为独立子进程启动，而是作为 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
        # 多智能体模式：为每个启用 multi_agent_mode 的会话维护一个 MultiAgentState
        # key = conversation_id, value = MultiAgentState
        self.multi_agent_states: Dict[str, Any] = {}

        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()
        # agent_id -> SubAgentTask 映射（供多智能体消息注入使用）
        self._sub_agent_instances: Dict[int, Any] = {}

        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,
        multi_agent_mode: bool = False,
        role_id: Optional[str] = None,
        display_name: Optional[str] = None,
    ) -> Dict:
        """创建子智能体任务并启动协程。
        
        参数 multi_agent_mode: True 时启用多智能体模式。
        参数 role_id: 多智能体模式下的角色标诶。
        参数 display_name: 多智能体模式下的显示名（如 UI Operator_1）。
        """
        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()

        # 多智能体模式：为该会话创建或复用 MultiAgentState
        multi_agent_state = None
        if multi_agent_mode:
            multi_agent_state = self.get_or_create_multi_agent_state(conversation_id)
            # 把实例注册到 state
            from modules.multi_agent.state import AgentInstance
            inst = AgentInstance(
                agent_id=agent_id,
                role_id=role_id or "",
                display_name=display_name or f"Agent_{agent_id}",
                task_id=task_id,
                status="running",
                summary=summary,
            )
            try:
                multi_agent_state.register_instance(inst)
            except ValueError:
                return {"success": False, "error": f"agent_id {agent_id} 已在该会话中使用"}

        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,
            multi_agent_mode=multi_agent_mode,
            multi_agent_state=multi_agent_state,
            display_name=display_name,
        )
        task_coro = sub_agent.run()
        asyncio_task = self._run_coro(task_coro)
        sub_agent._task = asyncio_task
        self._running_tasks[task_id] = asyncio_task
        # 缓存 sub_agent 实例供给多智能体模式 Poli注入使用
        self._sub_agent_instances[agent_id] = sub_agent

        def _on_done(fut):
            self._running_tasks.pop(task_id, None)
            self._sub_agent_instances.pop(agent_id, None)
            self.reconcile_task_states(conversation_id=conversation_id)
            # 多智能体模式：结束时把状态写回 MultiAgentState
            if multi_agent_mode and multi_agent_state:
                self._on_multi_agent_task_done(task_id, agent_id, multi_agent_state, sub_agent)

        asyncio_task.add_done_callback(_on_done)

        message = f"子智能体{agent_id} 已创建，任务ID: {task_id}"
        if multi_agent_mode and display_name:
            message = f"{display_name} 已创建，任务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,
            "display_name": display_name,
        }

    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:
            # 多智能体模式常见问答工具已在 SubAgentTask._execute_multi_agent_tool 中处理
            # 这里只处理实际通过主进程执行的工具
            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_or_create_multi_agent_state(self, conversation_id: str):
        """获取或为该会话创建 MultiAgentState。"""
        from modules.multi_agent.state import MultiAgentState
        state = self.multi_agent_states.get(conversation_id)
        if state:
            return state
        state = MultiAgentState(conversation_id=conversation_id)
        self.multi_agent_states[conversation_id] = state
        return state

    def get_multi_agent_state(self, conversation_id: str):
        """获取该会话的多智能体状态。"""
        return self.multi_agent_states.get(conversation_id)

    def drop_multi_agent_state(self, conversation_id: str) -> None:
        """删除会话状态（会话结束时调用）。"""
        self.multi_agent_states.pop(conversation_id, None)

    def inject_message_to_sub_agent(self, agent_id: int, message_text: str) -> bool:
        """同事件循环中向子智能体上下文插入 user 消息。
        
        适用于 ask_other_agent / send_message_to_sub_agent / answer_sub_agent_question_
        （非阻塞到工具结果的路径）。返回 True 表示成功注入。
        """
        # 查找该 agent_id 对应的 running SubAgentTask
        sub_agent = self._find_sub_agent_task_by_agent_id(agent_id)
        if not sub_agent:
            return False
        sub_agent.messages.append({"role": "user", "content": message_text})
        return True

    def _find_sub_agent_task_by_agent_id(self, agent_id: int) -> Optional[Any]:
        """通过遍历创建中的 task 查找活 SubAgentTask 实例。
        
        这是个 helper：在主实现中我们需要保留从 agent_id 到 SubAgentTask 的引用。
        理论上可以在 create_sub_agent 时把 sub_agent 存起来，这里使用 rs safer贪心法：
        遊历 _running_tasks 不为可行，因为 asyncio.Task 不抽不包含 SubAgentTask引用。
        我们改为 `SubAgentTask` 对象列表供查询。
        """
        # 优先查缓存：create_sub_agent 时的字段
        for inst in self._sub_agent_instances.values():
            if inst.agent_id == agent_id:
                return inst
        return None

    def _on_multi_agent_task_done(self, task_id: str, agent_id: int, state: Any, sub_agent: Any) -> None:
        """SubAgentTask 结束回调会调这个更新 MultiAgentState 实例状态。"""
        # 取出当前 status（由 _finalize_task 设置）
        final_task = self.tasks.get(task_id) or {}
        final_status = final_task.get("status") or "terminated"
        # 允许的自然退出：running -> idle（未走 _finalize） vs failed/timeout
        # SubAgentTask 在多智能体模式下没有 _finalize_task，我们手动赋值 idle
        # 如果状态被结 _finalize 为 failed/timeout，则保持该状态
        if final_status in TERMINAL_STATUSES:
            state.mark_status(agent_id, final_status, last_output=str(final_task.get("final_result") or ""))
        elif final_status == "terminated":
            state.mark_status(agent_id, "terminated")
        else:
            state.mark_status(agent_id, "idle")

    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()))