多Agent通信架构实战：从NATS消息总线到五大编排模式的生产落地——Agent通信协议篇

工具数量    1个Agent的工具选择准确率    5个Agent（每人10个工具）
  10               94%                      98%
  30               78%                      95%
  50               62%                      93%
 100               41%                      90%

场景                            单Agent方案             多Agent方案
─────────────────────────────────────────────────────────────────
代码生成 + 简单分类             Opus (昂贵)             Worker用Haiku + Orchestrator用Sonnet
                                                           成本降低 60%
长文档分析 + 摘要               Sonnet全程              Reader用Sonnet + Summarizer用Mini
                                                           成本降低 40%
安全审计 + 文档生成             GPT-4全程               Auditor用Opus + Writer用4o-mini
                                                           成本降低 55%

# Pipeline with per-stage validation gates
from typing import TypedDict

class PipelineState(TypedDict):
    input: str
    stage_outputs: list
    validation_passed: list

def stage_with_gate(state: PipelineState, stage_name: str, agent, validator=None):
    """每个 stage 执行后必须过验证门，防止脏数据流向下游"""
    output = agent.invoke(state["input"])

    if validator:
        is_valid = validator.invoke(f"Check if this output is valid: {output}")
        if not is_valid:
            raise ValueError(f"Stage {stage_name} validation failed")

    state["stage_outputs"].append(output)
    state["validation_passed"].append(stage_name)
    return state

MAX_DEBATE_ROUNDS = 3  # 硬上限，防止成本爆炸

def debate_with_guard(agent_a, agent_b, judge, question: str, max_rounds=MAX_DEBATE_ROUNDS):
    a_answer = agent_a.invoke(question)
    b_answer = agent_b.invoke(question)

    for round_num in range(max_rounds):
        # 互相批判
        a_critique = agent_a.invoke(f"Critique this answer: {b_answer}")
        b_critique = agent_b.invoke(f"Critique this answer: {a_answer}")

        # 裁决
        verdict = judge.invoke(f"""
        Question: {question}
        Agent A: {a_answer}
        Agent B: {b_answer}
        A's critique of B: {a_critique}
        B's critique of A: {b_critique}

        Decide: which answer is better, or do we need another round?
        """)

        if "ACCEPT" in verdict or round_num == max_rounds - 1:
            return verdict

    return verdict  # 强制终止

# 用 NATS CLI 测试延迟
$ nats bench --pub 1 --sub 1 --size 128 test.subject
# 典型结果: ~50µs P50, ~300µs P99

# 对比 Kafka（需 Java 客户端）：
# 即使调优到极致，P99 也在 5-10ms 级别

# 创建 JetStream stream，为 agent task 提供持久化
$ nats stream add AGENT_TASKS \
  --subjects "agents.tasks.*" \
  --storage file \
  --retention work \
  --max-msgs 100000 \
  --max-age 24h \
  --ack \
  --replicas 3

# 查看 stream 状态
$ nats stream info AGENT_TASKS

# 启动 NATS server（单二进制，无需 JVM/Erlang）
$ nats-server -js -m 8222 &

# Agent 注册为 micro service
$ nats micro add my-agent \
  --endpoint process \
  --subject "agents.prompt.process" \
  --description "Code review agent"

# 其他 Agent 发现可用的 agents
$ nats micro list
# 输出:
# my-agent (1 endpoint)
#   process - agents.prompt.process

┌──────────────────────────────────────┐
│  Agent ←→ Agent (跨组织/跨云)       │  → A2A 协议
├──────────────────────────────────────┤
│  Agent ←→ Agent (同一集群内)        │  → NATS Agent Protocol
├──────────────────────────────────────┤
│  Agent ←→ 工具/数据库/API           │  → MCP 协议
└──────────────────────────────────────┘

# Step 1: 安装 NATS Server（单二进制，2 秒搞定）
$ curl -sf https://binaries.nats.dev/nats-io/nats-server/v2@main \
  | PREFIX=/usr/local bash

# Step 2: 启动 NATS Server（带 JetStream）
$ nats-server -js -m 8222 &
# 输出: Server is ready, JetStream is enabled

# Step 3: 安装 Python 依赖
$ pip install nats-py openai  # nats-py 是 NATS 官方 Python 客户端

# Step 4: 验证 NATS 连通性
$ nats server check connection
# 输出: OK: connected to nats://localhost:4222

┌──────────────────────────────────────────────┐
│                  NATS Server                  │
│                                              │
│  subjects:                                    │
│  agents.tasks.{agent_id}      (任务分发)      │
│  agents.results.{agent_id}    (结果汇聚)      │
│  agents.hb.*                  (心跳)          │
│  agents.events.>              (事件总线)      │
└──────┬──────────┬──────────┬─────────────────┘
       │          │          │
  ┌────┴────┐ ┌───┴────┐ ┌──┴─────┐
  │Research │ │ Code   │ │ Review │
  │ Worker  │ │ Worker │ │ Worker │
  └─────────┘ └────────┘ └────────┘
       │          │          │
       └──────────┼──────────┘
                  │
          ┌───────┴───────┐
          │ Orchestrator  │
          │ (任务拆解+聚合) │
          └───────────────┘

#!/usr/bin/env python3
"""
多 Agent 协作系统 — NATS 消息总线 + Supervisor 模式
运行方式: python multi_agent_system.py
前置条件: nats-server -js -m 8222 & (已在 Step 2 启动)
"""

import asyncio
import json
import time
from dataclasses import dataclass, field
from typing import Optional
from nats.aio.client import Client as NATS
from nats.aio.errors import ErrTimeout


# ── 消息格式 ────────────────────────────────────
@dataclass
class AgentTask:
    task_id: str
    agent_type: str      # "research" | "code" | "review"
    prompt: str
    context: dict = field(default_factory=dict)

@dataclass
class AgentResult:
    task_id: str
    agent_type: str
    output: str
    latency_ms: float
    success: bool


# ── Worker Agent 基类 ────────────────────────────
class WorkerAgent:
    """每个 Worker 订阅自己的任务 subject，处理完后发布结果"""

    def __init__(self, agent_id: str, agent_type: str):
        self.agent_id = agent_id
        self.agent_type = agent_type
        self.task_subject = f"agents.tasks.{agent_id}"
        self.result_subject = f"agents.results.{agent_id}"
        self.hb_subject = f"agents.hb.{agent_type}.{agent_id}"
        self.nc = NATS()

    async def start(self):
        await self.nc.connect("nats://localhost:4222")

        # 订阅任务队列（JetStream push consumer）
        await self.nc.subscribe(
            self.task_subject,
            cb=self.handle_task,
            queue="workers"  # 同一类型的 worker 组成 queue group
        )

        # 发送心跳
        asyncio.create_task(self._heartbeat())
        print(f"[{self.agent_type}] Worker {self.agent_id} ready on {self.task_subject}")

    async def _heartbeat(self):
        """每 5 秒发心跳，离线超过 15 秒被判定为 dead"""
        while True:
            await self.nc.publish(
                self.hb_subject,
                json.dumps({
                    "agent_id": self.agent_id,
                    "agent_type": self.agent_type,
                    "ts": time.time()
                }).encode()
            )
            await asyncio.sleep(5)

    async def handle_task(self, msg):
        """处理任务——实际生产中这里调 LLM"""
        task = json.loads(msg.data.decode())
        t_start = time.time()

        try:
            # ⚡ 实际调用 LLM 的地方
            # response = await llm_client.chat(task["prompt"])
            # 这里用模拟延迟代替
            await asyncio.sleep(0.5)  # 模拟 LLM 调用

            result = AgentResult(
                task_id=task["task_id"],
                agent_type=self.agent_type,
                output=f"[{self.agent_type}] Processed: {task['prompt'][:50]}...",
                latency_ms=(time.time() - t_start) * 1000,
                success=True
            )
        except Exception as e:
            result = AgentResult(
                task_id=task["task_id"],
                agent_type=self.agent_type,
                output=str(e),
                latency_ms=(time.time() - t_start) * 1000,
                success=False
            )

        # 发布结果
        await self.nc.publish(
            self.result_subject,
            json.dumps({
                "task_id": result.task_id,
                "agent_type": result.agent_type,
                "output": result.output,
                "latency_ms": result.latency_ms,
                "success": result.success
            }).encode()
        )
        await self.nc.publish(msg.reply, b"ACK")  # 回复确认

    async def stop(self):
        await self.nc.close()


# ── Orchestrator ─────────────────────────────────
class Orchestrator:
    """接收用户请求 → 拆解任务 → 分发给 Workers → 聚合结果"""

    def __init__(self):
        self.nc = NATS()
        self.results: dict = {}  # task_id → AgentResult
        self.pending: set = set()
        self.done_event = asyncio.Event()

    async def start(self):
        await self.nc.connect("nats://localhost:4222")

    async def decompose_and_dispatch(self, user_request: str) -> list:
        """
        实际生产中用 LLM 做任务拆解，这里用规则模拟
        LLM 拆解 prompt:
        "Decompose the following task into subtasks for research, code, and review agents:
         {user_request}"
        """
        # 模拟 LLM 拆解（生产中用真实的 LLM 调用）
        subtasks = [
            AgentTask(task_id=f"task-{int(time.time()*1000)}-1",
                      agent_type="research",
                      prompt=f"Research: {user_request}"),
            AgentTask(task_id=f"task-{int(time.time()*1000)}-2",
                      agent_type="code",
                      prompt=f"Implement: {user_request}"),
            AgentTask(task_id=f"task-{int(time.time()*1000)}-3",
                      agent_type="review",
                      prompt=f"Review: {user_request}"),
        ]

        # 并行发布任务给所有 Worker
        for task in subtasks:
            subject = f"agents.tasks.{task.agent_type}_worker_1"
            await self.nc.publish(subject, json.dumps({
                "task_id": task.task_id,
                "agent_type": task.agent_type,
                "prompt": task.prompt
            }).encode())
            self.pending.add(task.task_id)
            print(f"[Orchestrator] Dispatched {task.task_id} → {subject}")

        return subtasks

    async def collect_results(self, timeout: float = 30.0):
        """收集所有 Worker 结果，支持超时"""
        async def on_result(msg):
            data = json.loads(msg.data.decode())
            task_id = data["task_id"]
            self.results[task_id] = data
            self.pending.discard(task_id)
            print(f"[Orchestrator] Received result for {task_id} "
                  f"({len(self.pending)} remaining)")

            if not self.pending:
                self.done_event.set()  # 全部完成

        # 订阅所有 result subjects
        for agent_type in ["research", "code", "review"]:
            subject = f"agents.results.{agent_type}_worker_1"
            await self.nc.subscribe(subject, cb=on_result)

        # 等待全部完成或超时
        try:
            await asyncio.wait_for(self.done_event.wait(), timeout=timeout)
        except asyncio.TimeoutError:
            print(f"[Orchestrator] Timeout! {len(self.pending)} tasks still pending: {self.pending}")

        return self.results

    async def aggregate(self, results: dict, user_request: str) -> str:
        """
        聚合 Worker 结果。实际生产中用 LLM 做 synthesis：
        "Synthesize these agent outputs into one coherent answer: {results}"
        """
        success_count = sum(1 for r in results.values() if r["success"])
        total_latency = sum(r["latency_ms"] for r in results.values())

        summary = f"""
        ═══════════════════════════════════════
        Multi-Agent Execution Report
        ═══════════════════════════════════════
        Request: {user_request}
        Tasks: {len(results)} dispatched, {success_count} succeeded

        Individual Results:
        {"".join(f"  [{r['agent_type']}] {'✅' if r['success'] else '❌'} "
                 f"{r['latency_ms']:.0f}ms → {r['output']}
        " for r in results.values())}

        Total Wall Time: {total_latency:.0f}ms
        ═══════════════════════════════════════
        """
        return summary

    async def stop(self):
        await self.nc.close()


# ── 监控：统一心跳收集 ──────────────────────────
class AgentMonitor:
    """订阅 agents.hb.> 通配符，实时监控所有 Agent 存活状态"""

    def __init__(self):
        self.nc = NATS()
        self.agents: dict = {}  # agent_id → last heartbeat

    async def start(self):
        await self.nc.connect("nats://localhost:4222")
        await self.nc.subscribe("agents.hb.>", cb=self._on_heartbeat)
        asyncio.create_task(self._check_liveness())

    async def _on_heartbeat(self, msg):
        data = json.loads(msg.data.decode())
        agent_id = data["agent_id"]
        self.agents[agent_id] = {"last_beat": time.time(), **data}

    async def _check_liveness(self):
        """每 10 秒检查一次——超过 15 秒无心跳 = 离线"""
        while True:
            await asyncio.sleep(10)
            now = time.time()
            for agent_id, info in list(self.agents.items()):
                if now - info["last_beat"] > 15:
                    print(f"⚠️  [Monitor] Agent {agent_id} is DEAD "
                          f"(last heartbeat {now - info['last_beat']:.0f}s ago)")
                    del self.agents[agent_id]


# ── 主函数 ───────────────────────────────────────
async def main():
    # 启动 Workers
    research_worker = WorkerAgent("research_worker_1", "research")
    code_worker = WorkerAgent("code_worker_1", "code")
    review_worker = WorkerAgent("review_worker_1", "review")

    # 启动监控
    monitor = AgentMonitor()

    # 启动 Orchestrator
    orchestrator = Orchestrator()

    # 并发启动所有组件
    await asyncio.gather(
        research_worker.start(),
        code_worker.start(),
        review_worker.start(),
        monitor.start(),
        orchestrator.start(),
    )

    # 用户请求
    user_request = "Build a REST API for a todo app with authentication"

    # Orchestrator 拆解 + 分发
    await orchestrator.decompose_and_dispatch(user_request)

    # 收集结果（30 秒超时）
    results = await orchestrator.collect_results(timeout=30.0)

    # 聚合并输出
    report = await orchestrator.aggregate(results, user_request)
    print(report)

    # 清理
    await research_worker.stop()
    await code_worker.stop()
    await review_worker.stop()
    await monitor.nc.close()
    await orchestrator.stop()

if __name__ == "__main__":
    asyncio.run(main())

# 终端 1: 启动 NATS
$ nats-server -js -m 8222 &

# 终端 2: 启动监控（可选）
$ nats sub "agents.hb.>" 
# 你会看到每 5 秒一次的心跳：
# [research_worker_1]: {"agent_id": "...", "ts": 1751700000}

# 终端 3: 运行多 Agent 系统
$ python multi_agent_system.py

# 预期输出:
# [research] Worker research_worker_1 ready on agents.tasks.research_worker_1
# [code] Worker code_worker_1 ready on agents.tasks.code_worker_1
# [review] Worker review_worker_1 ready on agents.tasks.review_worker_1
# [Orchestrator] Dispatched task-1751700000-1 → agents.tasks.research_worker_1
# [Orchestrator] Dispatched task-1751700000-2 → agents.tasks.code_worker_1
# [Orchestrator] Dispatched task-1751700000-3 → agents.tasks.review_worker_1
# [Orchestrator] Received result for task-1751700000-1 (2 remaining)
# [Orchestrator] Received result for task-1751700000-2 (1 remaining)
# [Orchestrator] Received result for task-1751700000-3 (0 remaining)
# ═══════════════════════════════════════
# Multi-Agent Execution Report
# ═══════════════════════════════════════
# ...

# 用 NATS bench 测试消息吞吐
$ nats bench --pub 10 --sub 10 --size 256 --msgs 10000 agents.tasks.test

# 典型输出（单节点）：
# Pub:   1,234,567 msgs/sec
# Sub:   1,234,567 msgs/sec
# P50:   0.05ms
# P99:   0.30ms
# Max:   15.2ms

# 10MB 大消息测试（模拟带附件的 Agent 间通信）
$ nats bench --pub 1 --sub 1 --size 10485760 --msgs 100 agents.tasks.large
# P50 ~50ms（大消息需序列化/反序列化开销）

# ❌ 危险代码
while not converged:
    result = agent_a.invoke(...)  # 可能无限循环

# ✅ 生产安全
MAX_ITERATIONS = 5
COST_BUDGET_USD = 2.00  # 单次请求最大成本

for i in range(MAX_ITERATIONS):
    if cost_tracker.total > COST_BUDGET_USD:
        raise BudgetExceededError(f"Cost {cost_tracker.total} exceeds budget {COST_BUDGET_USD}")
    result = agent.invoke(...)

# 创建工作队列 stream
$ nats stream add AGENT_WORK_QUEUE \
  --subjects "agents.tasks.>" \
  --retention work \
  --storage file \
  --max-msgs-per-subject 1000 \
  --max-age 1h \
  --ack \
  --discard old

# 创建 KV store 用于共享状态
$ nats kv add agent-state \
  --history 5 \
  --replicas 3

# 查看持久化状态
$ nats kv get agent-state orchestrator.context

# 创建 DLQ
$ nats stream add AGENT_DLQ \
  --subjects "agents.dlq.>" \
  --storage file \
  --max-msgs 10000 \
  --retention limits

# 消费失败超过 3 次的消息路由到 DLQ
# 在 Python 代码中配置 JetStream consumer:
# consumer_config = {
#     "max_deliver": 3,
#     "ack_wait": 30,  # 30 秒未 ack 算超时
#     "backoff": [1, 2, 4, 8]  # 重试间隔：1s, 2s, 4s, 8s
# }

# 启用 NATS 分布式追踪（需要 NATS 2.11+）
$ nats server request '$SYS.REQ.SERVER.PING.TRACING'
# 在每个消息中注入 trace_id，贯穿整个 Agent 调用链

import uuid

trace_id = str(uuid.uuid4())

# 所有 Agent 间消息携带 trace_id
headers = {
    "X-Trace-ID": trace_id,
    "X-Parent-Span": parent_span_id,
}
await nc.publish(subject, payload, headers=headers)

# 在生产 NATS config 中设置 ACL
# /etc/nats/nats-server.conf
authorization {
    AGENT_WORKERS = {
        subscribe = ["agents.tasks.research_worker_1",
                      "agents.tasks.code_worker_1",
                      "agents.tasks.review_worker_1"]
        publish = ["agents.results.*", "agents.hb.*"]
    }

    ORCHESTRATOR = {
        publish = ["agents.tasks.>"]
        subscribe = ["agents.results.>", "agents.hb.>"]
    }

    MONITOR = {
        subscribe = ["agents.hb.>"]
        publish = []  # 只读，不能发消息
    }
}

症状：Orchestrator 下发 3 个任务，只收到 2 个 result，一直等第 3 个
根因：No ACK → JetStream 自动重投递（如果没有 JetStream → 直接丢失）
排查：
  $ nats consumer info AGENT_TASKS research_consumer
  # 查看 pending 消息数和 redelivery 次数
修复：
  1. 确认 JetStream consumer 的 ack_wait 和 max_deliver 配置正确
  2. 检查 Worker 日志——是否在处理中卡死（LLM API 超时）

症状：Agent C 收到 Agent A 的结果，但 Agent B 的还没到
根因：NATS 不保证跨 subject 的顺序
修复：
  1. 在消息中添加 sequence_id
  2. 用 JetStream stream 的 ordered consumer
  3. 或者在 Orchestrator 层面做 barrier sync（所有结果到齐再聚合）

症状：50 个 Agent、5 秒心跳 = 每秒 10 条消息，监控 channel 爆炸
根因：所有 Agent 用同一个 heartbeat interval
修复：
  1. 随机 jitter：interval = random.uniform(4, 6)
  2. 或者用 NATS 的 connection status 代替应用层心跳：
     $ nats event '>'
     # 订阅连接事件，nats-server 自动通知 connect/disconnect