宜博 / lemonboss

/**
 * AcpProtocol - Agent Communication Protocol 实现
 *
 * 职责：只关心 ACP JSON-RPC 消息格式和业务语义。
 * 不关心传输方式和启动方式 — 那是 Transport 和 Launcher 的事。
 */

const fs = require('fs').promises;
const path = require('path');
const { Protocol } = require('./Protocol');
const { JsonRpcHandler } = require('./JsonRpcHandler');

class AcpProtocol extends Protocol {
  constructor(config, options = {}) {
    super(config, options);
    this.transport = null;
    this.rpc = null;
    this.workingDir = null;
  }

  get name() {
    return 'acp';
  }

  /**
   * 绑定已连接的 transport
   * @param {Transport} transport - 已连接的 Transport 实例
   * @param {string} workingDir - 工作目录
   */
  async connect(transport, workingDir) {
    this.transport = transport;
    this.workingDir = workingDir;

    // 创建 JSON-RPC 处理器
    this.rpc = new JsonRpcHandler(this.transport, {
      defaultTimeout: this.options.defaultTimeout,
    });

    this._setupEventHandlers();
    this.isConnected = true;
  }

  /**
   * 初始化协议
   */
  async initialize() {
    console.log('[AcpProtocol] 发送 initialize 请求...');
    const result = await this.rpc.request('initialize', {
      protocolVersion: 1,
      clientCapabilities: {
        fs: {
          readTextFile: true,
          writeTextFile: true,
        },
      },
    }, this.options.initTimeout || 120000);
    console.log('[AcpProtocol] initialize 成功');
    return result;
  }

  /**
   * 创建新会话
   */
  async newSession(options = {}) {
    console.log('[AcpProtocol] 发送 session/new 请求...');
    const params = {
      cwd: options.cwd || this.workingDir || '.',
      mcpServers: options.mcpServers || [],
    };

    console.log('[AcpProtocol] session/new 参数:', JSON.stringify(params, null, 2));

    const result = await this.rpc.request('session/new', params);
    this.sessionId = result.sessionId;
    console.log('[AcpProtocol] session/new 成功, sessionId:', this.sessionId);
    return result;
  }

  /**
   * 发送消息
   */
  async sendPrompt(text, files = []) {
    if (!this.sessionId) {
      throw new Error('No active session');
    }

    const prompt = [{ type: 'text', text }];
    for (const file of files) {
      prompt.push({ type: 'file', path: file });
    }

    return this.rpc.request('session/prompt', {
      sessionId: this.sessionId,
      prompt,
    }, this.options.promptTimeout);
  }

  /**
   * 响应权限请求
   */
  respondPermission(requestId, optionId) {
    this.rpc.resumeTimeouts(req => req.method === 'session/prompt');

    const outcome = optionId.includes('reject') ? 'rejected' : 'selected';
    this.rpc.respond(requestId, {
      outcome: {
        outcome,
        optionId,
      },
    });
  }

  /**
   * 断开连接 — 完整清理协议层资源
   *
   * 清理顺序：
   *   1. RPC 层：reject 所有 pending requests（释放定时器和 Promise）
   *   2. Transport 层：解除底层流上的事件监听器
   *   3. 自身状态：重置连接标记
   *
   * 注意：不终止子进程，进程生命周期由 Launcher.stop() 管理。
   */
  disconnect() {
    // 1. 清理 RPC 层
    if (this.rpc) {
      this.rpc._rejectAll('Protocol disconnected');
      this.rpc.removeAllListeners();
      this.rpc = null;
    }

    // 2. 清理 Transport 层（解除流上的事件监听器）
    if (this.transport) {
      try { this.transport.close(); } catch { }
      this.transport = null;
    }

    this.isConnected = false;
    this.sessionId = null;
  }

  /** @private */
  _setupEventHandlers() {
    this.rpc.on('notification', (msg) => this._handleNotification(msg));
    this.rpc.on('request', (msg) => this._handleRequest(msg));
    this.rpc.on('error', (err) => this.emit('error', err));
    this.rpc.on('close', (code, signal) => {
      this.isConnected = false;
      this.emit('disconnected', { code, signal });
    });
    this.rpc.on('output', (line) => this.emit('output', line));
    this.rpc.on('stderr', (line) => {
      console.log('[AcpProtocol] stderr:', line);
      this.emit('stderr', line);
    });
  }

  /** @private */
  _handleNotification(msg) {
    const { method, params } = msg;

    if (method === 'session/update') {
      this.rpc.resetTimeouts(req => req.method === 'session/prompt');
      this.emit('session_update', params);
    }

    this.emit('acp_event', {
      type: 'notification',
      method,
      params,
      timestamp: Date.now(),
    });
  }

  /** @private */
  async _handleRequest(msg) {
    const { id, method, params } = msg;

    console.log(`[AcpProtocol] 收到请求: ${method}`, JSON.stringify(params, null, 2));

    this.emit('acp_event', {
      type: 'request',
      id,
      method,
      params,
      timestamp: Date.now(),
    });

    try {
      if (method === 'session/request_permission') {
        this.rpc.pauseTimeouts(req => req.method === 'session/prompt');

        if (this.options.yolo) {
          console.log('[AcpProtocol] YOLO mode: Auto-approving permission request');
          const { options } = params;
          const allowOption = options.find(o => o.kind === 'allow_always' || o.optionId === 'approve_for_session')
            || options.find(o => o.kind === 'allow_once' || o.optionId === 'approve');

          if (allowOption) {
            this.respondPermission(id, allowOption.optionId);
            return;
          }
        }

        this.emit('permission_request', { id, params });

      } else if (method === 'fs/read_text_file') {
        try {
          const filePath = path.isAbsolute(params.path)
            ? params.path
            : path.join(this.workingDir || process.cwd(), params.path);
          let content = await fs.readFile(filePath, 'utf8');

          // JSON-RPC 响应是单行 JSON，保守设置 48KB 避免超长内容导致解析异常
          const MAX_CONTENT_BYTES = 48 * 1024;
          if (Buffer.byteLength(content, 'utf8') > MAX_CONTENT_BYTES) {
            content = content.substring(0, MAX_CONTENT_BYTES) +
              `\n\n... [truncated: file too large (${Buffer.byteLength(content, 'utf8')} bytes, limit ${MAX_CONTENT_BYTES})]`;
          }

          this.rpc.respond(id, { content });
        } catch (e) {
          this.rpc.respond(id, null, { code: -32603, message: e.message });
        }

      } else if (method === 'fs/write_text_file') {
        try {
          const filePath = path.isAbsolute(params.path)
            ? params.path
            : path.join(this.workingDir || process.cwd(), params.path);
          await fs.mkdir(path.dirname(filePath), { recursive: true });
          await fs.writeFile(filePath, params.content, 'utf8');
          this.rpc.respond(id, {});
          this.emit('file_changed', { path: filePath });
        } catch (e) {
          this.rpc.respond(id, null, { code: -32603, message: e.message });
        }

      } else {
        this.rpc.respond(id, null, {
          code: -32601,
          message: `Method not supported: ${method}`,
        });
      }
    } catch (err) {
      this.rpc.respond(id, null, {
        code: -32603,
        message: `Internal Error: ${err.message}`,
      });
    }
  }
}

module.exports = { AcpProtocol };