手搓生产级 AI Agent 系统（4）：A2A 协议——让两个 Agent 互相通信

┌─────────────────────────────────────────────┐
│           A2A 三种通信模式                    │
├─────────────────────────────────────────────┤
│                                              │
│  Mode 1: 广播（Broadcast）                    │
│  Manager → [Agent A, Agent B, Agent C]       │
│  用途：一对多分发任务，各自处理，汇总结果      │
│                                              │
│  Mode 2: 点对点（Peer-to-Peer）               │
│  Agent A ⇄ Agent B                          │
│  用途：两个 Agent 直接对话协商                │
│                                              │
│  Mode 3: 委托（Delegation）                   │
│  Agent A → Agent B（子任务）→ 返回结果给 A    │
│  用途：主 Agent 把子任务外包给专业 Agent      │
│                                              │
└─────────────────────────────────────────────┘

import json
import uuid
import time
from enum import Enum
from dataclasses import dataclass, field, asdict
from typing import Optional

class MessageType(Enum):
    TASK_REQUEST = "task_request"       # 请求执行任务
    TASK_RESULT = "task_result"         # 返回任务结果
    QUERY = "query"                     # 询问信息
    RESPONSE = "response"              # 回复询问
    ERROR = "error"                     # 错误通知
    HEARTBEAT = "heartbeat"            # 心跳检测

class MessagePriority(Enum):
    HIGH = 1
    NORMAL = 2
    LOW = 3

@dataclass
class A2AMessage:
    """A2A 协议消息体"""
    id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
    type: MessageType = MessageType.TASK_REQUEST
    sender: str = ""          # 发送方 Agent ID
    receiver: str = ""        # 接收方 Agent ID（空=广播）
    priority: MessagePriority = MessagePriority.NORMAL
    task: str = ""            # 任务描述
    context: dict = field(default_factory=dict)  # 上下文数据
    payload: str = ""         # 具体内容
    reply_to: str = ""        # 回复关联的消息 ID
    timestamp: float = field(default_factory=time.time)
    ttl: int = 60             # 消息有效期（秒），超时丢弃

    def to_json(self) -> str:
        return json.dumps({
            **asdict(self),
            "type": self.type.value,
            "priority": self.priority.value,
        }, ensure_ascii=False)

    @classmethod
    def from_json(cls, data: str) -> "A2AMessage":
        d = json.loads(data)
        d["type"] = MessageType(d["type"])
        d["priority"] = MessagePriority(d["priority"])
        return cls(**d)

    def is_expired(self) -> bool:
        return time.time() - self.timestamp > self.ttl

class A2ARouter:
    """A2A 消息路由器——Agent 之间的通信中枢"""

    def __init__(self):
        self.agents: dict[str, "A2AAgent"] = {}   # 注册的 Agent
        self.message_queue: list[A2AMessage] = []  # 消息队列
        self.message_log: list[dict] = []          # 消息日志（用于调试）

    def register(self, agent: "A2AAgent"):
        """注册 Agent 到路由"""
        self.agents[agent.id] = agent
        agent.router = self

    def send(self, message: A2AMessage) -> bool:
        """发送消息。如果指定 receiver，点对点；否则广播。"""
        self.message_log.append({
            "id": message.id, "type": message.type.value,
            "from": message.sender, "to": message.receiver or "ALL",
            "task": message.task[:80], "ts": message.timestamp,
        })

        if message.receiver:
            # 点对点
            target = self.agents.get(message.receiver)
            if target and not message.is_expired():
                target.receive(message)
                return True
            return False
        else:
            # 广播（排除发送者）
            for agent_id, agent in self.agents.items():
                if agent_id != message.sender and not message.is_expired():
                    agent.receive(message)
            return True

    def request_task(self, sender_id: str, task: str,
                     context: dict = None,
                     preferred_agent: str = None) -> str:
        """发送任务请求并等待结果。同步阻塞，实际生产应改为异步。"""
        msg = A2AMessage(
            type=MessageType.TASK_REQUEST,
            sender=sender_id,
            receiver=preferred_agent or "",
            task=task,
            context=context or {},
        )
        self.send(msg)

        # 等待接收方回复（简化实现：轮询）
        deadline = time.time() + msg.ttl
        while time.time()  list[dict]:
        return self.message_log[-20:]  # 最近 20 条

class A2AAgent(IntegratedAgent):
    """支持 A2A 通信的 Agent"""

    def __init__(self, agent_id: str, role: str, tools: dict,
                 llm_client, model: str = "deepseek-chat"):
        super().__init__(tools, llm_client, model)
        self.id = agent_id
        self.role = role          # "security_expert" / "report_writer" / "manager"
        self.router: Optional[A2ARouter] = None
        self._last_result: Optional[str] = None
        self._last_reply_to: str = ""

    def receive(self, message: A2AMessage):
        """接收消息并处理"""
        print(f"\n[{self.id}] 📨 收到来自 {message.sender} 的 {message.type.value}")

        if message.type == MessageType.HEARTBEAT:
            self._send(message.sender, MessageType.RESPONSE, payload="alive")
            return

        if message.type == MessageType.TASK_REQUEST:
            try:
                result = self.run(message.task)
                answer = result.get("answer", json.dumps(result))
            except Exception as e:
                answer = json.dumps({"error": str(e)})

            self._last_result = answer
            self._last_reply_to = message.id
            self._send(message.sender, MessageType.TASK_RESULT,
                       payload=answer, reply_to=message.id)

        elif message.type == MessageType.QUERY:
            # 回答询问
            response = self.run(f"回答以下问题，简洁准确：{message.payload}")
            self._send(message.sender, MessageType.RESPONSE,
                       payload=response.get("answer", ""),
                       reply_to=message.id)

    def _send(self, receiver: str, msg_type: MessageType,
              payload: str = "", reply_to: str = ""):
        """发送消息"""
        if self.router:
            msg = A2AMessage(
                type=msg_type, sender=self.id,
                receiver=receiver, payload=payload, reply_to=reply_to
            )
            self.router.send(msg)

    def delegate(self, target_agent: str, task: str,
                 context: dict = None) -> str:
        """委托任务给其他 Agent"""
        if self.router:
            return self.router.request_task(
                self.id, task, context, preferred_agent=target_agent
            )
        return json.dumps({"error": "no router connected"})

# ====== 初始化 ======
router = A2ARouter()

# 安全专家 Agent
security_agent = A2AAgent(
    agent_id="security_expert",
    role="security_expert",
    tools={
        "search": ToolDefinition(...),
        "read_file": ToolDefinition(...),
    },
    llm_client=client,
)

# 报告生成 Agent
report_agent = A2AAgent(
    agent_id="report_writer",
    role="report_writer",
    tools={
        "write_file": ToolDefinition(...),
        "format_markdown": ToolDefinition(...),
    },
    llm_client=client,
)

router.register(security_agent)
router.register(report_agent)

# ====== Manager Agent 协调 ======
async def security_review_pipeline(repo_description: str):
    """Manager 协调两个 Agent 完成安全审查流程"""

    # Step 1: 委托安全专家做分析
    print("🔍 Step 1: 安全分析")
    security_result = security_agent.run(
        f"分析以下仓库的安全问题：{repo_description}"
    )

    # Step 2: 委托报告生成器写报告
    print("📝 Step 2: 生成报告")
    report_task = f"""基于以下安全分析结果，生成一份 Markdown 格式的安全审查报告。
包含：摘要、漏洞清单（按严重程度排序）、修复建议、时间线。

分析结果：
{security_result['answer']}"""

    report = security_agent.delegate("report_writer", report_task)

    return {
        "security_analysis": security_result,
        "report": report,
        "message_log": router.get_log(),
    }

class A2AErrorHandler:
    """A2A 通信错误处理器"""

    @staticmethod
    def handle(message: A2AMessage, agent: A2AAgent) -> str:
        """根据错误类型决定处理策略"""
        error_str = message.payload

        # Agent 不可用 → 找替代
        if "unavailable" in error_str or "timeout" in error_str:
            if agent.router:
                available = [aid for aid in agent.router.agents
                            if aid != message.sender and aid != message.receiver]
                if available:
                    new_target = available[0]
                    return agent.delegate(new_target, message.task, message.context)
            return json.dumps({"error": "no_fallback_available"})

        # 能力不足 → 拆分任务
        if "cannot" in error_str or "超出能力" in error_str:
            return agent.router.request_task(
                message.sender,
                f"请将以下任务拆分为更小的子任务后逐一处理：{message.task}"
            ) if agent.router else json.dumps({"error": "no_router"})

        # 通用错误 → 重试一次
        return "RETRY"