agent-Specialization/modules/multi_agent/state.py

"""多智能体会话状态机。

一个 MultiAgentState 绑定到一个多智能体对话的 conversation_id，维护：
- 已创建的子智能体实例（agent_id ↔ role_id ↔ display_name ↔ task_id ↔ status）
- 待插入到主对话的待发 user 消息队列（pending_master_messages）
- 主智能体工具调用 answer_sub_agent_question / answer_other_agent 写回答案的 futomap
- 子智能体调用 ask_master / ask_other_agent 时挂起的 futomap

关键约定（来自 .astrion/memory/multi_agent_mode_design.md）：
- 消息格式：`来自 {显示名} 的{类型}\\nid: {消息id}\\n\\n<{显示名}>\\n<{标签}>\\n{内容}\\n</{标签}>\\n</{显示名}>`
- 接收方决定插入方式：
  - 子智能体 ask 阻塞等待 → main 调 answer_* 返回到工具结果
  - 子智能体 idle 状态 → 主对话的 pending_master_messages 直接插入新轮 user 消息
  - 子智能体 running 中 → inline 插入到当前末尾（在下一轮 model 调用前合并 messages）
- 通信是「工具调用提问」+「回答返回到工具结果」；其他场景（输出/进度/完成/任务发布/消息/回答）
  才以 user 消息格式插入对话。
"""
from __future__ import annotations

import asyncio
import json
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from modules.sub_agent.task import SubAgentTask

# ---------- 消息类型常量 ----------
TYPE_TASK = "Task"               # 主→子 任务发布
TYPE_OUTPUT = "Output"           # 子→主 进度/完成输出（统一）
TYPE_ASK = "Ask"                 # 子→主 / 子→子 提问
TYPE_ANSWER = "Answer"           # 主→子 / 子→子 回答（不插入对话，仅做工具结果）
TYPE_MESSAGE = "Message"         # 任意方向 消息
# 内部枚举到此

QUESTION_PREFIX_ASK_MASTER = "ask_master"
QUESTION_PREFIX_ASK_OTHER = "ask_other"


def format_multi_agent_message(
    *,
    display_name: str,
    msg_type: str,
    content: str,
    msg_id: Optional[str] = None,
    target: Optional[str] = None,
    extra_attrs: Optional[Dict[str, str]] = None,
) -> str:
    """按统一格式构造 user 消息字符串。

    Args:
        display_name: 发出方显示名（如 UI Operator_1 / Team Leader）
        msg_type: 消息类型，对应上方 TYPE_* 常量
        content: 消息正文
        msg_id: 消息 id；不传则自动生成
        target: 接收方显示名（用于子→子 提问时标明对谁提问）
        extra_attrs: 额外标签属性（如 question_id="ask_xxx"）
    """
    if not msg_id:
        msg_id = f"msg_{uuid.uuid4().hex[:10]}"

    # 第一行：自然语言前缀（含 target 标识）
    if target:
        prefix = f"来自 {display_name} 向 {target} 的{msg_type_to_text(msg_type)}"
    else:
        prefix = f"来自 {display_name} 的{msg_type_to_text(msg_type)}"

    # 第二行：id
    id_line = f"id: {msg_id}"

    # 属性 attr 字符串
    attrs = ""
    if target:
        attrs += f' target="{target}"'
    if extra_attrs:
        for k, v in extra_attrs.items():
            attrs += f' {k}="{v}"'

    # XML 包裹
    tag = msg_type
    xml = (
        f"<{display_name}>\n"
        f"<{tag}{attrs}>\n"
        f"{content}\n"
        f"</{tag}>\n"
        f"</{display_name}>"
    )

    return f"{prefix}\n{id_line}\n\n{xml}"


def msg_type_to_text(msg_type: str) -> str:
    """把 TYPE_* 转为中文短语，用于 prompt 前缀。"""
    mapping = {
        TYPE_TASK: "任务发布",
        TYPE_OUTPUT: "任务进度输出",
        TYPE_ASK: "提问",
        TYPE_ANSWER: "回答",
        TYPE_MESSAGE: "消息",
    }
    return mapping.get(msg_type, msg_type)


def build_master_dispatch_text(task: str, msg_id: Optional[str] = None) -> str:
    """主智能体发布任务时插入到子智能体对话的 user 消息文本。"""
    return format_multi_agent_message(
        display_name="Team Leader",
        msg_type=TYPE_TASK,
        content=task,
        msg_id=msg_id,
    )


def build_sub_agent_output_text(display_name: str, content: str, msg_id: Optional[str] = None) -> str:
    """子智能体输出（进度或完成）插入到主对话的 user 消息文本。"""
    return format_multi_agent_message(
        display_name=display_name,
        msg_type=TYPE_OUTPUT,
        content=content,
        msg_id=msg_id,
    )


def build_sub_agent_ask_master_text(display_name: str, question: str, question_id: str) -> str:
    """子智能体向主智能体提问时插入到主对话的 user 消息文本。"""
    return format_multi_agent_message(
        display_name=display_name,
        msg_type=TYPE_ASK,
        content=question,
        msg_id=question_id,
    )


def build_sub_agent_ask_other_text(
    display_name: str,
    target_display: str,
    question: str,
    question_id: str,
) -> str:
    """子智能体向另一个子智能体提问时插入到目标子智能体对话的文本。"""
    return format_multi_agent_message(
        display_name=display_name,
        msg_type=TYPE_ASK,
        content=question,
        msg_id=question_id,
        target=target_display,
    )


def build_master_message_to_sub_agent(message: str, msg_id: Optional[str] = None) -> str:
    """主智能体 send_message_to_sub_agent 时插入子对话的 user 消息文本。"""
    return format_multi_agent_message(
        display_name="Team Leader",
        msg_type=TYPE_MESSAGE,
        content=message,
        msg_id=msg_id,
    )


def build_master_answer_to_sub_agent(
    display_name: str,
    target_display: str,
    answer: str,
    question_id: str,
) -> str:
    """主智能体回答插入到子对话（仅当子智能体 not waiting 或 idle 时走 user 消息路径）。"""
    return format_multi_agent_message(
        display_name=display_name,
        msg_type=TYPE_ANSWER,
        content=answer,
        msg_id=question_id,
        target=target_display,
        extra_attrs={"question_id": question_id},
    )


# ---------- 运行态状态机 ----------
@dataclass
class AgentInstance:
    """一个多智能体会话中已创建的子智能体实例。"""

    agent_id: int
    role_id: str
    display_name: str
    task_id: str
    status: str = "running"          # running / idle / terminated / failed / timeout
    summary: str = ""
    created_at: float = field(default_factory=lambda: datetime.now().timestamp())
    last_output: str = ""

    def to_dict(self) -> Dict[str, Any]:
        return {
            "agent_id": self.agent_id,
            "role_id": self.role_id,
            "display_name": self.display_name,
            "task_id": self.task_id,
            "status": self.status,
            "summary": self.summary,
            "created_at": self.created_at,
            "last_output": self.last_output,
        }


class MultiAgentState:
    """绑到一个 conversation_id 的多智能体运行态。

    线程安全：所有 pubic 方法均假设在 SubAgentManager 的事件循环线程中调用，
    或者由 chat task 主线程通过 manager 的 _run_coro 进入此循环。
    跨线程访问通过 manager._run_coro 桥接，避免直接调用。
    """

    def __init__(self, conversation_id: str):
        self.conversation_id = conversation_id
        # agent_id 映射；同一会话里 agent_id 唯一
        self.agents: Dict[int, AgentInstance] = {}
        # task_id -> agent_id（便于在 SubAgentTask 完成时回写）
        self.task_id_to_agent_id: Dict[str, int] = {}
        # 主智能体待插入消息队列（每条都是字符串，由 chat task 取走）
        self.pending_master_messages: List[str] = []
        # ask_master / ask_other_agent 的等待 future
        # key = question_id, value = asyncio.Future (结果为 answer str 或 Exception)
        self.pending_questions: Dict[str, asyncio.Future] = {}
        # 一个 agent 可能同时只阻塞在一个 ask 工具上（最简实现）
        # key = agent_id, value = question_id（表示当前 agent 正阻塞等待）
        self.agent_blocking_question: Dict[int, str] = {}
        # 角色实例计数：role_id -> 已分配的最大 agent_id（数字）
        # 用于创建新实例时自动递增编号，但允许调用方显式指定
        self.role_counters: Dict[str, int] = {}

    # ----- 创建/查询 -----
    def next_agent_id_for_role(self, role_id: str) -> int:
        """为指定角色分配下一个 agent_id 编号。"""
        n = self.role_counters.get(role_id, 0) + 1
        self.role_counters[role_id] = n
        return n

    def register_instance(self, instance: AgentInstance) -> None:
        if instance.agent_id in self.agents:
            raise ValueError(f"agent_id {instance.agent_id} 已存在")
        self.agents[instance.agent_id] = instance
        self.task_id_to_agent_id[instance.task_id] = instance.agent_id

    def get_instance(self, agent_id: int) -> Optional[AgentInstance]:
        return self.agents.get(agent_id)

    def get_instance_by_task_id(self, task_id: str) -> Optional[AgentInstance]:
        aid = self.task_id_to_agent_id.get(task_id)
        if aid is None:
            return None
        return self.agents.get(aid)

    def list_active(self) -> List[AgentInstance]:
        return [a for a in self.agents.values() if a.status == "running" or a.status == "idle"]

    def list_all(self) -> List[AgentInstance]:
        return list(self.agents.values())

    def mark_status(self, agent_id: int, status: str, last_output: str = "") -> None:
        a = self.agents.get(agent_id)
        if a:
            a.status = status
            if last_output:
                a.last_output = last_output

    # ----- 主对话注入 -----
    def push_master_message(self, message_text: str) -> None:
        """把一条 user 消息追加到主对话待插入队列。"""
        self.pending_master_messages.append(message_text)

    def drain_master_messages(self) -> List[str]:
        """取出（清空）所有待插入主对话的消息。"""
        msgs = self.pending_master_messages
        self.pending_master_messages = []
        return msgs

    def has_pending_master_messages(self) -> bool:
        return len(self.pending_master_messages) > 0

    # ----- 阻塞问答 -----
    async def wait_for_answer(self, question_id: str, agent_id: int, timeout: float = 600.0) -> str:
        """子智能体 ask_* 工具调用后阻塞等待答案。

        返回 answer 字符串；超时/取消抛 asyncio.TimeoutError 或 CancelledError。
        """
        if question_id in self.pending_questions:
            raise RuntimeError(f"question_id 已存在: {question_id}")
        fut: asyncio.Future = asyncio.get_event_loop().create_future()
        self.pending_questions[question_id] = fut
        self.agent_blocking_question[agent_id] = question_id
        try:
            return await asyncio.wait_for(fut, timeout=timeout)
        except asyncio.TimeoutError:
            raise
        finally:
            self.pending_questions.pop(question_id, None)
            if self.agent_blocking_question.get(agent_id) == question_id:
                self.agent_blocking_question.pop(agent_id, None)

    def provide_answer(self, question_id: str, answer: str) -> bool:
        """主/其他子智能体 answer_* 工具调用时回写答案。

        返回 True 表示找到等待中的 future；False 表示无等待方或已超时。
        """
        fut = self.pending_questions.get(question_id)
        if not fut or fut.done():
            return False
        try:
            fut.set_result(answer)
        except asyncio.InvalidStateError:
            return False
        return True

    def is_agent_blocking(self, agent_id: int) -> bool:
        return agent_id in self.agent_blocking_question

    def get_blocking_question_id(self, agent_id: int) -> Optional[str]:
        return self.agent_blocking_question.get(agent_id)

    # ----- 持久化（最简版） -----
    def to_snapshot(self) -> Dict[str, Any]:
        return {
            "conversation_id": self.conversation_id,
            "agents": [a.to_dict() for a in self.agents.values()],
            "role_counters": self.role_counters,
        }

    @classmethod
    def from_snapshot(cls, snapshot: Dict[str, Any]) -> "MultiAgentState":
        state = cls(conversation_id=snapshot.get("conversation_id", ""))
        state.role_counters = dict(snapshot.get("role_counters") or {})
        for a_data in snapshot.get("agents") or []:
            a = AgentInstance(**a_data)
            state.agents[a.agent_id] = a
            if a.task_id:
                state.task_id_to_agent_id[a.task_id] = a.agent_id
        return state