宜博 / lemonboss

/**
 * TaskQueue — Sub-Agent 任务调度器
 *
 * 管理后台 Sub-Agent 任务的生命周期：
 * - dispatch(): 创建独立 Claude CLI session 执行任务
 * - 15s 快速超时策略：秒级返回的同步回；超时的转后台
 * - 心跳检查 + 完成回调
 * - 最大并发控制
 */

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

const MAX_CONCURRENT = 3;           // 最大并发 SubAgent
const HEARTBEAT_INTERVAL = 60_000;  // 60s 心跳检查
const TASK_TIMEOUT = 10 * 60_000;   // 10 分钟任务超时
const FAST_TIMEOUT = 15_000;        // 15s 快速超时（同步→异步切换点）

/**
 * Task 状态: queued → running → completed | failed | timeout
 */
class TaskQueue extends EventEmitter {
  constructor(chatService) {
    super();
    this.chatService = chatService;
    this.tasks = new Map();
    this._heartbeatTimer = null;
  }

  /**
   * 启动心跳
   */
  start() {
    if (this._heartbeatTimer) return;
    this._heartbeatTimer = setInterval(() => this._heartbeat(), HEARTBEAT_INTERVAL);
    this._heartbeatTimer.unref();
    console.log('[TaskQueue] 🚀 已启动 (心跳 60s, 最大并发 ' + MAX_CONCURRENT + ')');
  }

  /**
   * 停止心跳
   */
  stop() {
    if (this._heartbeatTimer) {
      clearInterval(this._heartbeatTimer);
      this._heartbeatTimer = null;
    }
  }

  /**
   * 发送消息，带 15s 快速超时策略
   *
   * - 15s 内完成 → 返回 { mode: 'sync', result, sessionId }
   * - 15s 超时   → 后台继续执行，返回 { mode: 'async', taskId, sessionId }
   *
   * @param {string} message 用户消息
   * @param {object} opts { sessionId?, provider?, parentSessionId? }
   * @returns {Promise<object>}
   */
  async sendWithTimeout(message, opts = {}) {
    const { sessionId: reqSessionId, provider = 'claude' } = opts;

    // 获取或创建主会话
    let session;
    let sessionId = reqSessionId;

    if (sessionId) {
      session = this.chatService.getSession(sessionId);
    }
    if (!session) {
      sessionId = sessionId || `bot-${Date.now().toString(36)}`;
      session = this.chatService.createSession(sessionId);
    }

    // 如果未连接，先连接
    if (session.status !== 'connected') {
      await session.connect(process.cwd(), { provider });
    }

    // 竞赛: sendMessage vs 15s 超时
    let resolved = false;

    // 保存原始 promise，用于后台完成检测
    const rawSendPromise = session.sendMessage(message);

    const resultPromise = rawSendPromise.then(result => {
      if (resolved) return null; // 超时后忽略同步路径
      resolved = true;
      return { mode: 'sync', result, sessionId };
    });

    const timeoutPromise = new Promise(resolve => {
      setTimeout(() => {
        if (resolved) return resolve(null);
        resolved = true;
        resolve('TIMEOUT');
      }, FAST_TIMEOUT);
    });

    const winner = await Promise.race([resultPromise, timeoutPromise]);

    if (winner === 'TIMEOUT') {
      // 转为后台任务 — sendMessage 仍在运行中
      const taskId = `task-${Date.now().toString(36)}-${Math.random().toString(36).substr(2, 4)}`;

      const task = {
        taskId,
        sessionId,
        message: message.substring(0, 200),
        status: 'running',
        result: null,
        error: null,
        createdAt: new Date().toISOString(),
        updatedAt: new Date().toISOString(),
        completedAt: null,
      };

      this.tasks.set(taskId, task);

      // 用原始 promise 等待完成（不是 resultPromise，那个会返回 null）
      rawSendPromise.then(result => {
        task.status = 'completed';
        task.result = result?.result || '';
        task.completedAt = new Date().toISOString();
        task.updatedAt = task.completedAt;
        this.emit('taskComplete', task);
        console.log(`[TaskQueue] ✅ Task ${taskId} 完成 (${task.result.length} chars)`);
      }).catch(err => {
        task.status = 'failed';
        task.error = err.message;
        task.updatedAt = new Date().toISOString();
        this.emit('taskFailed', task);
        console.error(`[TaskQueue] ❌ Task ${taskId} 失败:`, err.message);
      });

      console.log(`[TaskQueue] 📋 Task ${taskId} 已创建 (15s 超时转后台)`);

      return {
        mode: 'async',
        taskId,
        sessionId,
        message: `⏳ 任务已启动 (ID: ${taskId})，正在后台处理...`,
      };
    }

    // 15s 内完成 — 同步返回
    return winner;
  }

  /**
   * 派发独立 SubAgent 任务（创建新 session）
   */
  async dispatch(message, opts = {}) {
    const { provider = 'claude', parentSessionId } = opts;

    // 检查并发
    const runningCount = [...this.tasks.values()].filter(t => t.status === 'running').length;
    if (runningCount >= MAX_CONCURRENT) {
      return {
        mode: 'rejected',
        message: `⚠️ 当前有 ${runningCount} 个任务在运行中，请稍后再试`,
      };
    }

    const taskId = `task-${Date.now().toString(36)}-${Math.random().toString(36).substr(2, 4)}`;
    const sessionId = `sub-${taskId}`;

    const task = {
      taskId,
      sessionId,
      parentSessionId: parentSessionId || null,
      message: message.substring(0, 200),
      status: 'queued',
      result: null,
      error: null,
      createdAt: new Date().toISOString(),
      updatedAt: new Date().toISOString(),
      completedAt: null,
    };

    this.tasks.set(taskId, task);
    console.log(`[TaskQueue] 📋 Dispatched task ${taskId}`);

    // 后台异步执行
    this._executeTask(task, message, provider).catch(err => {
      console.error(`[TaskQueue] ❌ Task ${taskId} execution error:`, err.message);
    });

    return {
      mode: 'async',
      taskId,
      sessionId,
      message: `⏳ SubAgent 任务已启动 (ID: ${taskId})`,
    };
  }

  /**
   * 后台执行任务
   */
  async _executeTask(task, message, provider) {
    try {
      task.status = 'running';
      task.updatedAt = new Date().toISOString();

      // 创建独立 session
      const session = this.chatService.createSession(task.sessionId);
      await session.connect(process.cwd(), { provider });

      // 发送消息
      const result = await session.sendMessage(message);

      task.status = 'completed';
      task.result = result.result || '';
      task.completedAt = new Date().toISOString();
      task.updatedAt = task.completedAt;

      this.emit('taskComplete', task);
      console.log(`[TaskQueue] ✅ Task ${task.taskId} completed (${task.result.length} chars)`);

    } catch (err) {
      task.status = 'failed';
      task.error = err.message;
      task.updatedAt = new Date().toISOString();

      this.emit('taskFailed', task);
      console.error(`[TaskQueue] ❌ Task ${task.taskId} failed:`, err.message);
    }
  }

  /**
   * 查询任务
   */
  getTask(taskId) {
    return this.tasks.get(taskId) || null;
  }

  /**
   * 列出任务
   */
  listTasks(opts = {}) {
    let tasks = [...this.tasks.values()];

    if (opts.sessionId) {
      tasks = tasks.filter(t => t.sessionId === opts.sessionId || t.parentSessionId === opts.sessionId);
    }
    if (opts.status) {
      tasks = tasks.filter(t => t.status === opts.status);
    }

    // 按创建时间倒序
    tasks.sort((a, b) => new Date(b.createdAt) - new Date(a.createdAt));

    return tasks.slice(0, opts.limit || 50);
  }

  /**
   * 心跳检查
   */
  _heartbeat() {
    const now = Date.now();

    for (const [taskId, task] of this.tasks) {
      if (task.status !== 'running') continue;

      const elapsed = now - new Date(task.createdAt).getTime();
      if (elapsed > TASK_TIMEOUT) {
        console.log(`[TaskQueue] ⏰ Task ${taskId} 超时 (${Math.round(elapsed / 1000)}s)`);
        task.status = 'timeout';
        task.error = `任务超时 (${Math.round(elapsed / 60000)} 分钟)`;
        task.updatedAt = new Date().toISOString();
        this.emit('taskTimeout', task);
      }
    }

    // 清理 1 小时前的已完成任务
    const ONE_HOUR = 60 * 60 * 1000;
    for (const [taskId, task] of this.tasks) {
      if (['completed', 'failed', 'timeout'].includes(task.status)) {
        if (now - new Date(task.updatedAt).getTime() > ONE_HOUR) {
          this.tasks.delete(taskId);
        }
      }
    }
  }
}

module.exports = { TaskQueue, FAST_TIMEOUT, MAX_CONCURRENT };