宜博 / lemonboss

/**
 * ClaudeStreamProtocol - Claude CLI stream-json 协议实现
 *
 * 职责：只关心 Claude stream-json 消息格式和业务语义。
 * 不关心传输方式和启动方式 — 那是 Transport 和 Launcher 的事。
 */

const { Protocol } = require('../Protocol');
const { StreamLogger } = require('./logger');
const { handleMessage, registerHandler, registerContentBlockHandler } = require('./handlers');

class ClaudeStreamProtocol extends Protocol {
  constructor(config, options = {}) {
    super(config, options);
    this.workingDir = null;
    this.outputChunks = [];
    this.transport = null;

    // 组件
    this.logger = new StreamLogger({
      logDir: options.logDir || '.autodev-logs',
      prefix: options.logPrefix || 'claude-stream'
    });

    this._customHandlers = new Map();
  }

  get name() {
    return 'claude-stream';
  }

  /**
   * 注册自定义消息处理器
   */
  registerMessageHandler(type, handler) {
    registerHandler(type, handler);
    return this;
  }

  /**
   * 注册自定义 content block 处理器
   */
  registerBlockHandler(type, handler) {
    registerContentBlockHandler(type, handler);
    return this;
  }

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

    const agentName = this.config.id || this.config.name || 'agent';
    this.logger.init(workingDir, agentName);

    // 监听 Transport 事件
    this.transport.on('message', (msg) => this._handleMessage(msg));
    this.transport.on('output', (line) => this.emit('output', line));
    this.transport.on('stderr', (text) => {
      this.logger.logInternal('stderr', { content: text });
      this._log(`stderr: ${text}`);
      this.emit('stderr', text);
    });
    this.transport.on('close', (code, signal) => {
      this.logger.logInternal('process_close', { code, signal });
      this.logger.close();
      this.isConnected = false;
      this.emit('disconnected', { code, signal });
    });

    this.isConnected = true;
    this._log('连接成功');
  }

  /**
   * 初始化（Claude 不需要显式初始化）
   */
  async initialize() {
    this._log('initialize (no-op for claude)');
    return { sessionId: this.sessionId };
  }

  /**
   * 创建会话（Claude 不支持多会话）
   */
  async newSession(options = {}) {
    this._log('newSession (no-op for claude)');
    return { sessionId: this.sessionId };
  }

  /**
   * 发送提示消息
   */
  async sendPrompt(text, files = []) {
    if (!this.isConnected) {
      throw new Error('Not connected');
    }

    this.outputChunks = [];

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

    const msg = {
      type: 'user',
      message: { role: 'user', content }
    };

    this._log('发送 user message');
    this.logger.logOutgoing(msg);
    this.transport.send(msg);

    return this._waitForResult();
  }

  /**
   * 等待结果
   */
  _waitForResult() {
    return new Promise((resolve, reject) => {
      const handler = (msg) => {
        if (msg.type === 'result') {
          this.removeListener('_message', handler);
          this.removeListener('disconnected', exitHandler);
          resolve({
            result: msg.result,
            is_error: msg.is_error,
            duration_ms: msg.duration_ms
          });
        }
      };

      const exitHandler = ({ code, signal } = {}) => {
        this.removeListener('_message', handler);
        reject(new Error(`Claude process exited unexpectedly (code: ${code}, signal: ${signal})`));
      };

      this.on('_message', handler);
      this.once('disconnected', exitHandler);
    });
  }

  /**
   * 响应权限请求
   */
  respondPermission(requestId, optionId) {
    this._log('respondPermission (not supported in bypass mode)');
  }

  /**
   * 断开连接 — 完整清理协议层资源
   *
   * 清理 Transport 层的事件监听器，防止已关闭的 transport
   * 继续触发回调导致 event loop 无法退出。
   */
  disconnect() {
    // 清理 Transport 层（解除流上的事件监听器）
    if (this.transport) {
      try { this.transport.close(); } catch { }
      this.transport = null;
    }

    // 清理 logger
    try { this.logger.close(); } catch { }

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

  /** @private */
  _handleMessage(msg) {
    this.logger.logIncoming(msg);
    this.emit('_message', msg);

    const ctx = this._createContext();
    handleMessage(msg, ctx);

    this.emit('acp_event', {
      type: msg.type,
      subtype: msg.subtype,
      data: msg,
      timestamp: Date.now()
    });
  }

  /** @private */
  _createContext() {
    return {
      protocol: this,
      emit: (event, data) => this.emit(event, data),
      log: (message) => this._log(message),
      appendOutput: (text) => this.outputChunks.push(text),
      getOutput: () => this.outputChunks.join('')
    };
  }

  /** @private */
  _log(message) {
    console.log(`[ClaudeStreamProtocol] ${message}`);
  }

  getLogStats() {
    return this.logger.getStats();
  }
}

module.exports = { ClaudeStreamProtocol };