宜博 / lemonboss

/**
 * Orchestrator - 多 Agent 编排器
 *
 * 主循环编排逻辑：Master ↔ Sub-Agent 交互、格式自动纠正、结果累积、完成判定。
 * 提示词模板由 prompts.js 管理，结果状态由 result-tracker.js 管理。
 *
 * 关键概念区分：
 * - loopCount（循环次数）：主循环的安全上限，防止死循环
 * - repairCount（修复次数）：validator FAIL 后触发的修复次数，受 maxRepairs 限制
 * - 正常的首次执行和通过验证不消耗修复额度
 */

const fs = require('fs');
const { EventEmitter } = require('events');
const { Parser } = require('./parser');
const { Scheduler } = require('./scheduler');
const { Prompts } = require('./prompts');
const { ResultTracker } = require('./result-tracker');
const { getGlobalLogManager, waitForGlobalLogManager } = require('./logging');
const { FailoverError } = require('../engine');

/** 默认最大修复次数 */
const MAX_REPAIRS = 5;
/** 主循环安全上限（防止死循环） */
const SAFETY_MAX_LOOPS = 15;
/** 连续全部失败的熔断阈值 */
const CONSECUTIVE_FAILURE_LIMIT = 3;
/** Agent 输出格式不对时的最大重试次数 */
const MAX_FORMAT_RETRIES = 2;
/** 递归编排最大深度 */
const MAX_ORCHESTRATOR_DEPTH = 2;

class Orchestrator extends EventEmitter {
  constructor(config, workDir, options = {}) {
    super();
    this.agents = config.agents;
    this.config = config;
    this.registry = options.registry || null;
    this.runRegistry = options.runRegistry || null;
    this.workDir = workDir;
    this.manager = null;
    this.depth = options.depth || 0;

    this.sessionId = options.sessionId || `orch-${Date.now()}`;
    this.logManager = options.logManager || getGlobalLogManager();
    this.logSession = options.logSession;
  }

  // ── 初始化 ──

  getLogSession() {
    if (this.logSession) return this.logSession;

    if (this.logManager.initialized) {
      this.logSession = this.logManager.getOrCreateSession(this.sessionId, {
        workDir: this.workDir,
        config: this.config,
      });
      return this.logSession;
    }

    return null;
  }

  async initManager(overrides = {}) {
    if (!this.manager || Object.keys(overrides).length > 0) {
      await waitForGlobalLogManager();

      const AgentManager = require('./agent-manager');
      this.manager = new AgentManager(this.config, this.workDir, {
        sessionId: overrides.sessionId || this.sessionId,
        logSession: overrides.logSession || this.getLogSession(),
        logManager: overrides.logManager || this.logManager,
        registry: overrides.registry || this.registry,
        runRegistry: overrides.runRegistry || this.runRegistry,
      });
    }
  }

  async shutdown() {
    if (this.manager && typeof this.manager.closeAllSessions === 'function') {
      await this.manager.closeAllSessions();
    }
  }

  // ── 解析 ──

  parseCalls(output) {
    return Parser.parseCalls(output);
  }

  parseResult(output) {
    return Parser.parseResult(output);
  }

  // ── Agent 执行 ──

  /**
   * 调用单个 Agent，带格式自动纠正重试
   * - type=orchestrator → 走子 Orchestrator 路径
   * - 其他类型 → 最多调用 MAX_FORMAT_RETRIES + 1 次（首次 + 重试次数）
   */
  async callAgentWithRetry(call, loopCount, idx) {
    // orchestrator 类型走子编排路径
    const agentConfig = this.agents[call.agent];
    if (agentConfig?.type === 'orchestrator') {
      return this.handleOrchestratorCall(call, loopCount);
    }

    // 通知：Sub Agent 开始执行
    this.emit('agent_start', { agent: call.agent, loop: loopCount });

    let task = call.task;
    let lastResult = null;

    for (let retry = 0; retry <= MAX_FORMAT_RETRIES; retry++) {
      const { output } = await this.manager.callAgent(call.agent, task, {
        metadata: { loop: loopCount, index: idx, formatRetry: retry },
        onData: (chunk) => {
          this.emit('agent_data', { agent: call.agent, chunk, loop: loopCount });
        },
      });
      const parsed = this.parseResult(output);

      // 格式解析成功（状态不是 ERROR，或虽然 ERROR 但不是因为格式问题）
      if (parsed.status !== 'ERROR' || parsed.message !== 'No result marker found') {
        this.emit('agent_complete', { agent: call.agent, status: parsed.status, loop: loopCount });
        return { agent: call.agent, status: parsed.status, message: parsed.message, output };
      }

      // 记录最后一次结果，用于重试耗尽时返回
      lastResult = { agent: call.agent, status: parsed.status, message: parsed.message, output };

      // 格式不对且还有重试机会，构造纠正提示
      if (retry < MAX_FORMAT_RETRIES) {
        const agentType = this.agents[call.agent]?.type;
        console.warn(`[FORMAT_RETRY] ${call.agent} 返回格式不符合预期，自动重试 (${retry + 1}/${MAX_FORMAT_RETRIES})`);

        task = Prompts.formatCorrection({
          agentType,
          retryCount: retry + 1,
          maxRetries: MAX_FORMAT_RETRIES,
          originalTask: call.task,
        });
      }
    }

    // 重试耗尽
    this.emit('agent_complete', { agent: call.agent, status: 'ERROR', loop: loopCount });
    return lastResult;
  }

  // ── 递归编排 ──

  /**
   * 处理 type=orchestrator 的 Agent 调用
   * 创建子 Orchestrator 实例，执行完整编排循环
   *
   * @param {{agent: string, task: string}} call - 调用指令
   * @param {number} loopCount - 父编排的当前轮次
   * @returns {Promise<{agent: string, status: string, message: string}>}
   */
  async handleOrchestratorCall(call, loopCount) {
    const agentConfig = this.agents[call.agent];
    const groupName = agentConfig.agentGroup;

    if (!groupName || !this.config.agentGroups?.[groupName]) {
      return {
        agent: call.agent,
        status: 'ERROR',
        message: `Orchestrator agent "${call.agent}" references unknown agentGroup "${groupName}"`,
      };
    }

    // 递归深度限制
    if (this.depth >= MAX_ORCHESTRATOR_DEPTH) {
      return {
        agent: call.agent,
        status: 'ERROR',
        message: `Max orchestrator depth (${MAX_ORCHESTRATOR_DEPTH}) exceeded`,
      };
    }

    const group = this.config.agentGroups[groupName];
    const subConfig = this.buildSubConfig(group);
    const childOrch = new Orchestrator(subConfig, this.workDir, {
      registry: this.registry,
      runRegistry: this.runRegistry,
      logManager: this.logManager,
      depth: this.depth + 1,
      sessionId: `${this.sessionId}:${call.agent}`,
    });

    // 事件冒泡：子编排的 progress/completed 透传给父，标记 source
    childOrch.on('progress', (data) => {
      this.emit('progress', { ...data, source: call.agent });
    });

    try {
      await childOrch.initManager();
      const maxRepairs = group.maxRepairs || MAX_REPAIRS;

      // Phase 6: 执行前快照工作目录
      const filesBefore = new Set(this._listWorkDirFiles());

      console.log(`[SUB-ORCH] ${call.agent} (group=${groupName}, depth=${this.depth + 1}, maxRepairs=${maxRepairs})`);

      const childResult = await childOrch.execute(call.task, {
        maxIterations: maxRepairs,
      });

      // Phase 6: 执行后 diff → artifacts
      const filesAfter = this._listWorkDirFiles();
      const artifacts = filesAfter.filter(f => !filesBefore.has(f));

      const status = childResult.completed ? 'DONE' : 'ERROR';
      const message = childResult.completed
        ? `子编排完成 (${childResult.iterations} 轮, ${childResult.repairs} 次修复)`
        : `子编排失败: ${childResult.output}`;

      return { agent: call.agent, status, message, output: childResult.output, artifacts };
    } catch (err) {
      return {
        agent: call.agent,
        status: 'ERROR',
        message: `子编排异常: ${err.message}`,
      };
    } finally {
      await childOrch.shutdown();
    }
  }

  /**
   * 为子编排构建配置：继承父配置，但只包含 group 中的 agents + master
   */
  buildSubConfig(group) {
    // 优先从全量 agents (allAgents) 取，fallback 到当前层 agents
    const allAgents = this.config.allAgents || this.config.agents || {};
    const subAgents = {};

    // master 总是包含
    if (allAgents['master']) {
      subAgents['master'] = { ...allAgents['master'] };
    }

    // group 里的 agents
    for (const name of (group.agents || [])) {
      if (allAgents[name]) {
        subAgents[name] = { ...allAgents[name] };
      }
    }

    return {
      ...this.config,
      agents: subAgents,
      allAgents,  // 继承全量 agents（支持多层递归）
      agentGroups: this.config.agentGroups,
    };
  }

  /**
   * 按依赖分组执行多个 Agent 调用
   * 使用 Promise.allSettled 实现错误降级：单个 Agent 失败不会终结整组
   */
  async executeWithDependencies(calls, loopCount) {
    const groups = Scheduler.getExecutionGroups(this.agents, calls);
    const allResults = [];

    for (const group of groups) {
      const groupCalls = calls.filter(c => group.includes(c.agent));
      const settled = await Promise.allSettled(
        groupCalls.map((call, idx) => this.callAgentWithRetry(call, loopCount, idx))
      );

      for (let i = 0; i < settled.length; i++) {
        if (settled[i].status === 'fulfilled') {
          allResults.push(settled[i].value);
        } else {
          // Agent 调用抛出异常 → 降级为 ERROR 结果
          const err = settled[i].reason;
          const failover = FailoverError.from(err);
          console.error(`[AGENT_ERROR] ${groupCalls[i].agent} 执行失败 (retryable=${failover.retryable}): ${err.message}`);
          allResults.push({
            agent: groupCalls[i].agent,
            status: 'ERROR',
            message: err.message,
            output: '',
            retryable: failover.retryable,
          });
        }
      }
    }

    return allResults;
  }

  // ── 检测 ──

  /**
   * 检测 Master 是否越权执行（直接在工作目录创建了文件）
   * @param {Set<string>} [baseline] - 任务开始时的文件快照，用于增量检测
   */
  detectMasterOverreach(baseline) {
    try {
      const entries = fs.readdirSync(this.workDir);
      if (baseline) {
        // 增量检测：只报告新增的非隐藏文件
        return entries.filter(e => !baseline.has(e) && !e.startsWith('.'));
      }
      return entries.filter(e => !e.startsWith('.'));
    } catch {
      return [];
    }
  }

  // ── 日志辅助 ──

  /**
   * 递归列出 workDir 内所有文件（相对路径）
   * 排除隐藏文件、node_modules、__pycache__
   */
  _listWorkDirFiles(dir, prefix = '') {
    const base = dir || this.workDir;
    const results = [];
    try {
      const entries = fs.readdirSync(base, { withFileTypes: true });
      for (const entry of entries) {
        if (entry.name.startsWith('.') || entry.name === 'node_modules' || entry.name === '__pycache__') continue;
        const relPath = prefix ? `${prefix}/${entry.name}` : entry.name;
        if (entry.isDirectory()) {
          results.push(...this._listWorkDirFiles(path.join(base, entry.name), relPath));
        } else {
          results.push(relPath);
        }
      }
    } catch {
      // 目录不存在或权限不足
    }
    return results;
  }

  async logEvent(session, event, data, level) {
    if (session) {
      await session.logEvent(null, event, data, level);
    }
  }

  // ── 主循环 ──

  buildInitialPrompt(requirement) {
    const agentList = Object.entries(this.agents)
      .filter(([name]) => name !== 'master')
      .map(([name, cfg]) => {
        const typeLabel = cfg.type === 'orchestrator'
          ? `子编排, agentGroup=${cfg.agentGroup}`
          : cfg.type;
        return `- ${name} (${typeLabel}): ${cfg.description || ''}`;
      })
      .join('\n');

    return Prompts.initial(requirement, agentList);
  }

  /**
   * 执行编排任务
   * @param {string} requirement - 用户需求
   * @param {Object} options
   * @param {number} options.maxIterations - 最大修复次数（validator FAIL 后的重试），默认 5
   */
  async execute(requirement, options = {}) {
    await this.initManager();
    const maxRepairs = options.maxIterations || MAX_REPAIRS;
    const logSession = this.getLogSession();
    const tracker = new ResultTracker(this.agents);
    const resumeState = options.resumeState || null;
    const onCheckpoint = options.onCheckpoint;

    // 记录工作目录基线快照（用于越权检测增量对比）
    let baselineFiles;
    try {
      baselineFiles = new Set(fs.readdirSync(this.workDir));
    } catch {
      baselineFiles = new Set();
    }

    let loopCount = resumeState?.loopCount || 0;
    let repairCount = resumeState?.repairCount || 0;
    let consecutiveAllFail = 0;
    if (Array.isArray(resumeState?.results) && resumeState.results.length > 0) {
      tracker.merge(resumeState.results);
    }

    await this.logEvent(logSession, 'orchestrator_start', {
      requirement,
      maxRepairs,
      availableAgents: Object.keys(this.agents),
    });

    let prompt = resumeState?.prompt || this.buildInitialPrompt(requirement);

    if (resumeState) {
      await this.logEvent(logSession, 'orchestrator_resumed', {
        loop: loopCount,
        repairCount,
        recoveredResults: tracker.getAll().length,
      });
    }

    try {
      while (loopCount < SAFETY_MAX_LOOPS) {
        loopCount++;

        await this.logEvent(logSession, 'loop_start', {
          loop: loopCount,
          repairCount,
        });

        // 1. 调用 Master Agent（含流式输出 → Reflection 可见）
        const { output: masterOutput } = await this.manager.callAgent('master', prompt, {
          metadata: { loop: loopCount, repairCount },
          onData: (chunk) => {
            this.emit('agent_data', { agent: 'master', chunk, loop: loopCount });
          },
        });
        const calls = this.parseCalls(masterOutput);

        // 通知：Master 决定调用哪些 agent
        if (calls.length > 0) {
          this.emit('master_calls', {
            loop: loopCount,
            agents: calls.map(c => c.agent),
          });
        }

        await this.logEvent(logSession, 'master_completed', {
          loop: loopCount,
          calls: calls.map(c => ({ agent: c.agent, task: c.task.substring(0, 200) })),
        });

        // 2. Master 没有输出 CALL 指令
        if (calls.length === 0) {
          prompt = this.handleNoCalls(masterOutput, tracker, logSession, loopCount, baselineFiles);
          if (prompt === null) {
            return this.buildResult(true, loopCount, repairCount, masterOutput, tracker, logSession);
          }
          if (onCheckpoint) {
            await onCheckpoint({
              version: 1,
              requirement,
              loopCount,
              repairCount,
              prompt,
              results: tracker.getAll(),
            });
          }
          continue;
        }

        // 3. 执行子 Agent
        const results = await this.executeWithDependencies(calls, loopCount);
        tracker.merge(results, loopCount);

        // 3.5 熔断检测：所有 Agent 都失败 → 连续失败计数器+1
        const allFailed = results.length > 0 && results.every(r => r.status === 'ERROR');
        if (allFailed) {
          consecutiveAllFail++;
          console.warn(`[CIRCUIT_BREAKER] 所有 Sub Agent 失败 (连续 ${consecutiveAllFail}/${CONSECUTIVE_FAILURE_LIMIT} 次)`);

          if (consecutiveAllFail >= CONSECUTIVE_FAILURE_LIMIT) {
            await this.logEvent(logSession, 'orchestrator_failed', {
              loops: loopCount,
              repairs: repairCount,
              reason: 'circuit_breaker_all_agents_failed',
              consecutiveFailures: consecutiveAllFail,
            }, 'error');

            return {
              completed: false,
              iterations: loopCount,
              repairs: repairCount,
              output: `熔断：所有 Sub Agent 连续 ${consecutiveAllFail} 轮全部失败，终止编排`,
              results: tracker.getAll(),
            };
          }
        } else {
          consecutiveAllFail = 0; // 有任意成功则重置
        }

        // 4. 统计修复次数：有 validator FAIL → repairCount++
        const hasValidatorFail = results.some(r =>
          this.agents[r.agent]?.type === 'validator' && r.status === 'FAIL'
        );
        if (hasValidatorFail) {
          repairCount++;
          console.log(`[REPAIR] 修复次数: ${repairCount}/${maxRepairs}`);
        }

        await this.logEvent(logSession, 'loop_completed', {
          loop: loopCount,
          repairCount,
          consecutiveAllFail,
          results: results.map(r => ({
            agent: r.agent,
            status: r.status,
            message: r.message?.substring(0, 200),
          })),
        });

        // 5. 不再提前退出 — 始终将结果反馈给 Master 做 Reflection
        // Master 看到结果后会做 Reflection，如果满意则输出 RESULT: DONE（无 CALL），
        // 下一轮 handleNoCalls 检测到 RESULT 标记后自然退出。

        // 6. 检查修复次数是否耗尽
        if (repairCount >= maxRepairs) {
          await this.logEvent(logSession, 'orchestrator_failed', {
            loops: loopCount,
            repairs: repairCount,
            reason: 'max_repairs_reached',
          }, 'warn');

          return {
            completed: false,
            iterations: loopCount,
            repairs: repairCount,
            output: `修复次数耗尽 (${repairCount}/${maxRepairs})`,
            results: tracker.getAll(),
          };
        }

        // 7. 进度通知
        this.emit('progress', {
          loop: loopCount,
          repairCount,
          results: results.map(r => ({ agent: r.agent, status: r.status })),
          passedValidators: tracker.passedValidators,
          unpassedValidators: tracker.unpassedValidators,
        });

        // 7. 构建下一轮反馈
        prompt = Prompts.followup({
          currentResults: results,
          passedValidators: tracker.passedValidators,
          unpassedValidators: tracker.unpassedValidators,
        });
        if (onCheckpoint) {
          await onCheckpoint({
            version: 1,
            requirement,
            loopCount,
            repairCount,
            prompt,
            results: tracker.getAll(),
          });
        }
      }

      // 安全上限耗尽
      await this.logEvent(logSession, 'orchestrator_failed', {
        loops: loopCount,
        repairs: repairCount,
        reason: 'safety_loop_limit_reached',
      }, 'warn');

      return {
        completed: false,
        iterations: loopCount,
        repairs: repairCount,
        output: 'Safety loop limit reached',
        results: tracker.getAll(),
      };

    } catch (err) {
      await this.logEvent(logSession, 'orchestrator_error', {
        error: err.message,
        stack: err.stack,
      }, 'error');

      // 异常时也保存 checkpoint，便于恢复
      if (onCheckpoint) {
        try {
          await onCheckpoint({
            version: 1,
            requirement,
            loopCount,
            repairCount,
            prompt,
            results: tracker.getAll(),
            error: err.message,
          });
        } catch (cpErr) {
          console.error('[CHECKPOINT] 异常时保存 checkpoint 失败:', cpErr.message);
        }
      }

      throw err;
    }
  }

  /**
   * 构建最终返回结果
   */
  async buildResult(completed, loopCount, repairCount, output, tracker, logSession) {
    const reason = completed ? 'all_validators_passed' : 'incomplete';
    await this.logEvent(logSession, completed ? 'orchestrator_completed' : 'orchestrator_failed', {
      loops: loopCount,
      repairs: repairCount,
      reason,
    });

    const result = {
      completed,
      iterations: loopCount,
      repairs: repairCount,
      output,
      results: tracker.getAll(),
    };

    if (completed) {
      this.emit('completed', result);
    }

    return result;
  }

  /**
   * 处理 Master 没有输出 CALL 指令的情况
   * @returns {string|null} 下一轮 prompt，null 表示任务完成
   */
  handleNoCalls(masterOutput, tracker, logSession, loopCount, baselineFiles) {
    // 越权检测（使用基线增量对比）
    const filesInWorkDir = this.detectMasterOverreach(baselineFiles);
    if (filesInWorkDir.length > 0 && tracker.isEmpty()) {
      console.warn(`[WARNING] Master 可能越权执行了任务，检测到文件: ${filesInWorkDir.join(', ')}`);
      this.logEvent(logSession, 'master_overreach_detected', { loop: loopCount, filesFound: filesInWorkDir }, 'warn');

      // 构造虚拟 executor 结果
      const firstExecutor = tracker.allExecutors[0];
      if (firstExecutor) {
        tracker.set(firstExecutor, {
          agent: firstExecutor,
          status: 'DONE',
          message: `Master 直接执行产生的文件: ${filesInWorkDir.join(', ')}`,
          output: masterOutput,
        });
      }

      return Prompts.forceValidation(filesInWorkDir, tracker.allValidators);
    }

    // 检查是否已全部完成
    if (tracker.isComplete()) {
      return null; // 信号：任务完成
    }

    // 未完成，强制继续
    this.logEvent(logSession, 'loop_incomplete', {
      loop: loopCount,
      unpassedValidators: tracker.unpassedValidators,
    });

    return Prompts.forceResume(tracker.unpassedValidators);
  }
}

module.exports = { Orchestrator };