宜博 / lemonboss

/**
 * ═══════════════════════════════════════════════════════════════
 *  AgentBus — Agent 中央消息总线
 * ═══════════════════════════════════════════════════════════════
 *
 *  职责:
 *    1. Agent 之间的消息路由和分发
 *    2. 消息持久化 (SQLite)
 *    3. 消息类型管理和校验
 *
 *  通讯模式:
 *    上级→下级: task.assign    分派任务
 *    下级→上级: task.report    汇报进度/结果
 *    同级→同级: task.handover  Pipeline 节点流转
 *    同级→同级: task.feedback  经验回流
 *    CXO↔CXO:  cxo.review     制衡审核
 *    广播:      status.update  状态变更通知
 *
 *  用法:
 *    const bus = new AgentBus(db);
 *
 *    // Agent 订阅自己的消息
 *    bus.on('cxo-ceo-001', (message) => { ... });
 *
 *    // Agent 发送消息
 *    bus.send({
 *      from: 'cxo-ceo-001',
 *      to: 'role-dramatist-001',
 *      method: 'task.assign',
 *      payload: { task: '创作一个都市爱情短剧剧本' }
 *    });
 */

import { EventEmitter } from 'node:events';
import { randomUUID } from 'node:crypto';

// ─────────────────────────────────────────────
//  消息类型定义
// ─────────────────────────────────────────────

const MESSAGE_METHODS = {
  // 上下级通讯
  'task.assign':    { direction: 'downward',  desc: '分派任务' },
  'task.report':    { direction: 'upward',    desc: '汇报结果/进度' },

  // 同级通讯 (Pipeline)
  'task.handover':  { direction: 'lateral',   desc: 'Pipeline 节点流转' },
  'task.feedback':  { direction: 'lateral',   desc: '经验/数据反馈回流' },

  // CXO 制衡
  'cxo.review':     { direction: 'lateral',   desc: '方案审核请求' },
  'cxo.approve':    { direction: 'lateral',   desc: '方案审核通过' },
  'cxo.reject':     { direction: 'lateral',   desc: '方案审核驳回' },

  // 状态通知
  'status.update':  { direction: 'broadcast', desc: '状态变更通知' },
  'status.heartbeat': { direction: 'upward',  desc: '心跳状态上报' },

  // 数据共享
  'data.share':     { direction: 'any',       desc: '共享数据/资源' },
  'data.request':   { direction: 'any',       desc: '请求数据/资源' },
};

export class AgentBus extends EventEmitter {

  /**
   * @param {Object} [db] - 可选的 SQLite 数据库实例 (better-sqlite3)
   */
  constructor(db = null) {
    super();
    this.db = db;
    this.agents = new Map(); // agentId → { handler, type, parentId }

    if (this.db) {
      this._initDB();
    }
  }

  // ─────────────────────────────────────────
  //  Agent 注册/注销
  // ─────────────────────────────────────────

  /**
   * Agent 接入消息总线
   * @param {string} agentId
   * @param {Object} meta - { type, parentId }
   * @param {Function} handler - 收到消息时的回调
   */
  register(agentId, meta, handler) {
    this.agents.set(agentId, { ...meta, handler });
    this.emit('agent:registered', agentId);
  }

  /**
   * Agent 断开消息总线
   */
  unregister(agentId) {
    this.agents.delete(agentId);
    this.emit('agent:unregistered', agentId);
  }

  // ─────────────────────────────────────────
  //  消息发送
  // ─────────────────────────────────────────

  /**
   * 发送消息
   *
   * @param {Object} message
   * @param {string} message.from     - 发送者 Agent ID
   * @param {string} message.to       - 接收者 Agent ID (广播时为 '*' 或 'children:xxx')
   * @param {string} message.method   - 消息类型 (如 'task.assign')
   * @param {Object} message.payload  - 消息内容
   * @param {string} [message.replyTo]  - 回复的消息 ID (用于关联对话)
   *
   * @returns {Object} 发送的消息 (含生成的 id 和 timestamp)
   */
  send(message) {
    const msg = {
      id: randomUUID(),
      timestamp: new Date().toISOString(),
      from: message.from,
      to: message.to,
      method: message.method,
      payload: message.payload || {},
      replyTo: message.replyTo || null,
      status: 'delivered',
    };

    // 校验消息类型
    if (!MESSAGE_METHODS[msg.method]) {
      throw new Error(`未知的消息类型: ${msg.method}。可用: ${Object.keys(MESSAGE_METHODS).join(', ')}`);
    }

    // 持久化
    if (this.db) {
      this._persistMessage(msg);
    }

    // 路由分发
    this._route(msg);

    // 全局事件 (用于监控和日志)
    this.emit('message', msg);

    return msg;
  }

  /**
   * 分派任务 (上级→下级 的便捷方法)
   */
  assignTask(fromId, toId, task) {
    return this.send({
      from: fromId,
      to: toId,
      method: 'task.assign',
      payload: task,
    });
  }

  /**
   * 汇报结果 (下级→上级 的便捷方法)
   */
  reportResult(fromId, toId, result, replyTo = null) {
    return this.send({
      from: fromId,
      to: toId,
      method: 'task.report',
      payload: result,
      replyTo,
    });
  }

  /**
   * 流转到下一个节点 (同级 Pipeline 的便捷方法)
   */
  handover(fromId, toId, deliverable, pipelineId = null) {
    return this.send({
      from: fromId,
      to: toId,
      method: 'task.handover',
      payload: { ...deliverable, pipeline_id: pipelineId },
    });
  }

  /**
   * 共享数据
   */
  shareData(fromId, toId, data) {
    return this.send({
      from: fromId,
      to: toId,
      method: 'data.share',
      payload: data,
    });
  }

  // ─────────────────────────────────────────
  //  消息查询
  // ─────────────────────────────────────────

  /**
   * 获取某个 Agent 的消息历史
   */
  getMessages(agentId, options = {}) {
    if (!this.db) return [];

    const { method, direction, limit = 50 } = options;
    let sql = `SELECT * FROM messages WHERE (from_id = ? OR to_id = ?)`;
    const params = [agentId, agentId];

    if (method) {
      sql += ` AND method = ?`;
      params.push(method);
    }

    sql += ` ORDER BY timestamp DESC LIMIT ?`;
    params.push(limit);

    return this.db.prepare(sql).all(...params);
  }

  /**
   * 获取 Agent 之间的对话链 (通过 replyTo 关联)
   */
  getThread(messageId) {
    if (!this.db) return [];

    const thread = [];
    let currentId = messageId;

    while (currentId) {
      const msg = this.db.prepare('SELECT * FROM messages WHERE id = ?').get(currentId);
      if (!msg) break;
      thread.unshift(msg);
      currentId = msg.reply_to;
    }

    // 也查找后续回复
    const replies = this.db.prepare(
      'SELECT * FROM messages WHERE reply_to = ? ORDER BY timestamp'
    ).all(messageId);
    thread.push(...replies);

    return thread;
  }

  // ─────────────────────────────────────────
  //  内部方法
  // ─────────────────────────────────────────

  /**
   * 消息路由
   */
  _route(msg) {
    // 广播: to = '*' 或 to = 'children:xxx'
    if (msg.to === '*') {
      for (const [id, agent] of this.agents) {
        if (id !== msg.from) {
          agent.handler(msg);
        }
      }
      return;
    }

    // 下级广播: to = 'children:cxo-ceo-001' → 发给所有 parentId 匹配的 agent
    if (msg.to.startsWith('children:')) {
      const parentId = msg.to.replace('children:', '');
      for (const [id, agent] of this.agents) {
        if (agent.parentId === parentId) {
          agent.handler(msg);
        }
      }
      return;
    }

    // 点对点
    const target = this.agents.get(msg.to);
    if (target) {
      target.handler(msg);
    } else {
      msg.status = 'undeliverable';
      this.emit('undeliverable', msg);
    }
  }

  /**
   * 初始化消息表
   */
  _initDB() {
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS messages (
        id          TEXT PRIMARY KEY,
        timestamp   TEXT NOT NULL,
        from_id     TEXT NOT NULL,
        to_id       TEXT NOT NULL,
        method      TEXT NOT NULL,
        payload     TEXT,
        reply_to    TEXT,
        status      TEXT DEFAULT 'delivered',
        created_at  DATETIME DEFAULT CURRENT_TIMESTAMP
      );

      CREATE INDEX IF NOT EXISTS idx_messages_from ON messages(from_id);
      CREATE INDEX IF NOT EXISTS idx_messages_to ON messages(to_id);
      CREATE INDEX IF NOT EXISTS idx_messages_method ON messages(method);
      CREATE INDEX IF NOT EXISTS idx_messages_reply ON messages(reply_to);
    `);
  }

  /**
   * 消息持久化
   */
  _persistMessage(msg) {
    this.db.prepare(`
      INSERT INTO messages (id, timestamp, from_id, to_id, method, payload, reply_to, status)
      VALUES (?, ?, ?, ?, ?, ?, ?, ?)
    `).run(
      msg.id,
      msg.timestamp,
      msg.from,
      msg.to,
      msg.method,
      JSON.stringify(msg.payload),
      msg.replyTo,
      msg.status,
    );
  }
}

export default AgentBus;