宜博 / lemonboss

/**
 * JsonRpcHandler - JSON-RPC 2.0 协议处理器
 *
 * 处理请求/响应关联、超时管理等 JSON-RPC 协议相关逻辑。
 * ACP 和 MCP 协议都基于 JSON-RPC，因此这是共享的基础层。
 */

const { EventEmitter } = require('events');

const JSONRPC_VERSION = '2.0';
const DEBUG = !!process.env.DEBUG_ACP;

/**
 * JSON-RPC 会话处理器
 *
 * Events:
 * - 'request': 收到服务端请求（需要响应）
 * - 'notification': 收到服务端通知（无需响应）
 * - 'error': 发生错误
 * - 'close': 连接关闭
 * - 'output': 收到非 JSON 输出
 * - 'stderr': 收到标准错误输出
 */
class JsonRpcHandler extends EventEmitter {
  constructor(transport, options = {}) {
    super();
    this.transport = transport;
    this.options = {
      defaultTimeout: 60000,
      ...options,
    };
    this.requestId = 1;
    this.pendingRequests = new Map();

    // 监听传输层消息
    this.transport.on('message', (msg) => this._handleMessage(msg));

    // 转发传输层事件
    this.transport.on('error', (err) => this.emit('error', err));
    this.transport.on('close', (code, signal) => {
      this._rejectAll('Transport closed');
      this.emit('close', code, signal);
    });
    this.transport.on('output', (line) => this.emit('output', line));
    this.transport.on('stderr', (line) => this.emit('stderr', line));
  }

  /**
   * 发送请求并等待响应
   * @param {string} method - 方法名
   * @param {object} params - 参数
   * @param {number} timeoutMs - 超时时间
   * @returns {Promise<any>}
   */
  request(method, params = {}, timeoutMs = this.options.defaultTimeout) {
    return new Promise((resolve, reject) => {
      const id = this.requestId++;

      const requestObj = {
        resolve,
        reject,
        method,
        startTime: Date.now(),
        duration: timeoutMs,
        isPaused: false,
        timeout: timeoutMs > 0 ? setTimeout(() => this._handleTimeout(id), timeoutMs) : null,
      };

      this.pendingRequests.set(id, requestObj);

      const message = {
        jsonrpc: JSONRPC_VERSION,
        id,
        method,
        params,
      };

      if (DEBUG) console.log('[JsonRpcHandler] >>> 发送请求:', JSON.stringify(message));

      try {
        this.transport.send(message);
      } catch (e) {
        this.pendingRequests.delete(id);
        if (requestObj.timeout) clearTimeout(requestObj.timeout);
        reject(e);
      }
    });
  }

  /**
   * 发送通知（无需响应）
   * @param {string} method - 方法名
   * @param {object} params - 参数
   */
  notify(method, params = {}) {
    try {
      this.transport.send({
        jsonrpc: JSONRPC_VERSION,
        method,
        params,
      });
    } catch (e) {
      this.emit('error', e);
    }
  }

  /**
   * 发送响应
   * @param {number|string} id - 请求 ID
   * @param {any} result - 结果
   * @param {object|null} error - 错误信息
   */
  respond(id, result, error = null) {
    try {
      const msg = {
        jsonrpc: JSONRPC_VERSION,
        id,
      };
      if (error) {
        msg.error = error;
      } else {
        msg.result = result;
      }
      this.transport.send(msg);
    } catch (e) {
      this.emit('error', e);
    }
  }

  /**
   * 暂停指定请求的超时计时器
   * 用于用户交互时防止超时
   * @param {function} filterFn - 筛选函数
   */
  pauseTimeouts(filterFn) {
    for (const [id, req] of this.pendingRequests.entries()) {
      if (filterFn(req) && !req.isPaused) {
        clearTimeout(req.timeout);
        req.isPaused = true;
        req.timeout = null;
      }
    }
  }

  /**
   * 恢复指定请求的超时计时器
   * @param {function} filterFn - 筛选函数
   */
  resumeTimeouts(filterFn) {
    const now = Date.now();
    for (const [id, req] of this.pendingRequests.entries()) {
      if (filterFn(req) && req.isPaused) {
        req.isPaused = false;
        // duration <= 0 表示永不超时
        if (req.duration > 0) {
          const elapsed = now - req.startTime;
          const remaining = Math.max(0, req.duration - elapsed);
          req.timeout = setTimeout(() => this._handleTimeout(id), remaining);
        }
      }
    }
  }

  /**
   * 重置指定请求的超时计时器
   * 用于收到流式更新时保持连接
   * @param {function} filterFn - 筛选函数
   */
  resetTimeouts(filterFn) {
    for (const [id, req] of this.pendingRequests.entries()) {
      if (filterFn(req)) {
        if (req.timeout) clearTimeout(req.timeout);
        req.startTime = Date.now();
        req.isPaused = false;
        // duration <= 0 表示永不超时
        req.timeout = req.duration > 0
          ? setTimeout(() => this._handleTimeout(id), req.duration)
          : null;
      }
    }
  }

  /**
   * 处理接收到的消息
   * @private
   */
  _handleMessage(msg) {
    if (DEBUG) console.log('[JsonRpcHandler] <<< 收到消息:', JSON.stringify(msg).substring(0, 200));

    if (msg.method) {
      // 来自服务端的请求或通知
      if (msg.id !== undefined) {
        this.emit('request', msg);
      } else {
        this.emit('notification', msg);
      }
    } else if (msg.id !== undefined) {
      // 我们发送请求的响应
      const req = this.pendingRequests.get(msg.id);
      if (req) {
        clearTimeout(req.timeout);
        this.pendingRequests.delete(msg.id);

        if (msg.error) {
          req.reject(new Error(msg.error.message || 'Unknown RPC error'));
        } else {
          req.resolve(msg.result);
        }
      }
    }
  }

  /**
   * 处理超时
   * @private
   */
  _handleTimeout(id) {
    const req = this.pendingRequests.get(id);
    if (req) {
      this.pendingRequests.delete(id);
      req.reject(new Error(`Request ${req.method} timed out after ${req.duration}ms`));
    }
  }

  /**
   * 拒绝所有未完成的请求
   * @private
   */
  _rejectAll(reason) {
    for (const req of this.pendingRequests.values()) {
      clearTimeout(req.timeout);
      req.reject(new Error(reason));
    }
    this.pendingRequests.clear();
  }
}

module.exports = { JsonRpcHandler };