宜博 / lemonboss

/**
 * SessionPool — 可复用会话池管理
 *
 * 职责：
 *   - 会话的创建 / 获取 / 重置 / 关闭
 *   - 生命周期事件通知
 *   - 会话统计
 *
 * 不负责：
 *   - Protocol 选择（由 protocol-factory 决定）
 *   - Prompt 执行（由 stream-executor 决定）
 *   - Agent 配置查询（由 AgentManager 决定）
 */

const { EventEmitter } = require('events');
const { randomUUID } = require('crypto');
const { createProtocol } = require('./protocol-factory');
const { IdleTimer } = require('../engine/idle-timer');

/**
 * 会话空闲上限。与 routes/agent-chat.js 的 `sessions` 池共用同一个旋钮和同一个含义
 * —— 两边存的都是**常驻**的 claude CLI 子进程，回收判据也一样。
 *
 * 在此之前这个池子只有三条回收路径：换新会话 / reset() / killActive()。
 * 「正常跑完」和「不可重试的报错」都不在其中，而 getKey 里带了 taskId
 * —— 每个图节点一个新 key，于是一节点一进程、永不释放，且不报错。
 * 实测 23 小时漏出 7 个图执行进程（另有 38 个来自 agent-chat 侧的另一个洞）。
 *
 * 每次读 env 而不是模块加载时读一次：单测要能改这个值，不然只能等 30 分钟。
 */
function idleMs() {
  return Number(process.env.AGENT_SESSION_IDLE_MS) || 30 * 60 * 1000;
}

/**
 * 机器级会话上限（27-P3C §15.5 第 2 条）。
 *
 * 池子里每个 session 就是一个**常驻** claude CLI 子进程（stream-json 模式），实测单进程 RSS
 * 约 200-300MB。在此之前 `this.sessions.size` **没有任何人过问** —— 一次心跳泄漏就攒到过
 * 82 个进程 / 17.5 GB。批 4 又刚把公司级并行放开（单公司上限 20），机器上同时能起多少
 * 完全没有闸门。
 *
 * **调用时读，不在模块加载时快照**（`company-paths.js` 顶上那条教训：快照一次就意味着
 * 「改了配置要重启才生效」，而且不报错）。
 */
const DEFAULT_MAX_SESSIONS = 12;
function maxSessions() {
  const raw = Number(process.env.AGENT_MAX_SESSIONS);
  return Number.isInteger(raw) && raw > 0 ? raw : DEFAULT_MAX_SESSIONS;
}

/** 池子满了、且**一个空闲的都没有**（全都在途）时抛这个 —— 调用方该重试或排队，不是崩 */
class SessionPoolFullError extends Error {
  constructor(size, max) {
    super(`会话池已满（${size}/${max}）且无空闲会话可回收：机器级并发上限，请稍后重试`);
    this.name = 'SessionPoolFullError';
    this.size = size;
    this.max = max;
  }
}

class SessionPool extends EventEmitter {
  constructor() {
    super();
    this.sessions = new Map();
    this.activeKey = null;
  }

  /**
   * 生成会话唯一 Key
   *
   * key 里必须带 (companyId, taskId)（28 §6 批次 3-12）：Agent 本体合并成一份之后
   * `agentName` 不再唯一标识一次运行 —— 两家公司调同一个本体，key 相同就会**复用同一个
   * CLI 会话**，A 公司的上下文直接漏进 B 公司的对话，而且两边都成功返回、不报错。
   * 不能靠 workDir 顺带带上：body 模式的 cwd 确实含这两段，legacy / scoped skill 模式的
   * 不含（cwd 分别是公司私有目录和本体目录），指望「路径里恰好有」就是又一份手抄的约定。
   * 缺省值与 agent-locator 的保留值同名，含义一致：没有公司 / 没有任务。
   */
  getKey(agentName, agentConfig, workDir, scope) {
    const runtime = `${agentConfig.companyId || '_global'}/${agentConfig.taskId || '_no-task'}`;
    const base = `acp:${agentName}:${agentConfig.client}:${agentConfig.deploy}:${runtime}:${workDir}`;
    return scope ? `${base}:${scope}` : base;
  }

  /**
   * 创建新会话（Launcher + Transport + Protocol 全链路）
   *
   * @param {string} agentName
   * @param {object} agentConfig - { type, client, deploy, ... }
   * @param {object} launcherConfig - 来自 runtime-resolver
   * @param {string} protocolName - acp / claude-stream / codex-jsonl
   * @param {string} workDir
   * @param {string} [scope] - 会话细分维度（如 conversationId）；必须与 getOrCreate 查询时用的一致
   * @returns {Promise<object>} session object
   */
  async create(agentName, agentConfig, launcherConfig, protocolName, workDir, scope) {
    // scope 必须透传：不透传的话 getOrCreate 用「带 scope 的 key」查、create 用
    // 「不带 scope 的 key」存，两个 key 永远对不上 —— 查不到就再建一个，会话池里
    // 堆的全是查不到的死会话，且每次调用都新起一个 CLI 进程，同样不报错。
    const sessionKey = this.getKey(agentName, agentConfig, workDir, scope);

    this._emit(agentName, protocolName, 'initializing', null, sessionKey);

    // 通过工厂创建 Protocol（无 if/else）
    const protocol = createProtocol(
      protocolName,
      { id: agentName, name: agentName },
      { yolo: true }
    );

    let launcher = null;

    if (protocolName === 'claude-cli') {
      // ClaudeCliProtocol 自行管理子进程（每次 sendPrompt 启动 claude -p）
      // 不需要 launcher/transport
      await protocol.connect(null, workDir);
      await protocol.initialize();
      await protocol.newSession({ cwd: workDir });
    } else {
      // 其他协议：组合 Launcher + Transport + JSON-RPC
      const { createConnection } = require('@mas/acp-core');
      // ⚠️ 起子进程**之前**先腾位置：起完再判等于该拦的那一个已经起来了
      this._evictForRoom();

      const { transport, launcher: lnch } = await createConnection(
        agentConfig, launcherConfig, workDir
      );
      launcher = lnch;

      await protocol.connect(transport, workDir);
      await protocol.initialize();
      await protocol.newSession({ cwd: workDir });
    }

    const sessionId = protocol.sessionId || `${protocolName}-${agentName}-${randomUUID()}`;
    if (!protocol.sessionId) {
      protocol.sessionId = sessionId;
    }

    this._emit(agentName, protocolName, 'ready', sessionId, sessionKey);

    const session = {
      key: sessionKey,
      agentName,
      protocolName,
      protocol,
      launcher,
      sessionId,
      callQueue: Promise.resolve(),
      idleTimer: null, // 由 armIdle 上弦；新建会话在首次调用期间故意不计时
    };

    this.sessions.set(sessionKey, session);
    return session;
  }

  /**
   * 获取已有会话或创建新会话
   */
  async getOrCreate(agentName, agentConfig, launcherConfig, protocolName, workDir, scope) {
    const key = this.getKey(agentName, agentConfig, workDir, scope);
    const existing = this.sessions.get(key);

    if (existing?.protocol?.isConnected) {
      existing.idleTimer?.cancel(); // 请求在途期间不计空闲
      existing.idleSince = null;    // 在途 = 不可被 _evictForRoom 挑走
      this._emit(agentName, protocolName, 'reused', existing.sessionId, key);
      return existing;
    }

    // 清理失效会话
    if (existing?.protocol) {
      try { existing.protocol.disconnect(); } catch { }
      if (existing.launcher) {
        try { await existing.launcher.stop(); } catch { }
      }
      this.sessions.delete(key);
    }

    return this.create(agentName, agentConfig, launcherConfig, protocolName, workDir, scope);
  }

  /**
   * 调用结束后给会话上弦；调用开始时由 getOrCreate 停表。
   *
   * 计时器只在「没有在途请求」的间隙走，所以再长的单次生成也不会被 TTL 打断
   * —— 与 agent-chat 侧同一套写法，比在流式回调里 touch() 少一处要维护的调用点。
   *
   * ponytail: 这是**兜底**不是精算 —— 一次性 taskId（图节点跑完就不会再用）也要
   * 白占一个 TTL 周期的内存，天花板是「并发会话数 × 单进程 RSS（实测约 300MB）× TTL」。
   * 要更紧的话，调用方在明知 taskId 一次性时直接 reset(key) 立刻释放；这里不替它判断，
   * 因为 devloop 的 retry-with-reflection 复用同一个 taskId，一刀切关掉会砍掉复用能力。
   */
  armIdle(sessionKey) {
    const session = this.sessions.get(sessionKey);
    if (!session) return;
    session.idleTimer?.cancel();
    // 空闲态用显式时间戳记，不去戳 IdleTimer 的私有 _timer —— `_evictForRoom` 靠它挑最久没用的那个
    session.idleSince = Date.now();
    const ms = idleMs();
    session.idleTimer = new IdleTimer(
      ms,
      () => {
        console.log(`[SessionPool] 🧹 会话空闲超过 ${Math.round(ms / 60000)} 分钟，回收 ${sessionKey}`);
        this.reset(sessionKey).catch(() => {});
      },
      { label: `pool-session:${sessionKey}` }
    ).start();
  }

  /**
   * 起新会话之前腾位置：池子满了就回收**最久没用**的空闲会话。
   *
   * 为什么是「回收空闲」而不是「直接拒绝」：session 是可复用的池化资源，杀掉一个空闲的
   * 只是让它下次冷启动一次（几秒），而拒绝会让一个本来能跑的任务因为资源原因失败 ——
   * 那是把可恢复的代价换成不可恢复的。
   *
   * 只有**一个空闲的都没有**（全部在途）时才抛 `SessionPoolFullError`：那时机器是真满了，
   * 再起就是超发。**刻意抛而不是静默放行** —— 静默放行正是这个仓库栽过的那一类
   * 「写不进去/超出去但不报错」（27-P3C §九.4quater）。
   *
   * 注意 `reset()` 是 async（要 await launcher.stop()），这里**故意不 await**：
   * 腾位置只需要把它从 Map 里摘掉（reset 第一件事就是），子进程的收尾可以在后台走完。
   * await 的话，一次腾位要串行等一次进程退出，而调用方还在等着起自己的会话。
   */
  _evictForRoom() {
    const max = maxSessions();
    while (this.sessions.size >= max) {
      let victim = null;
      for (const [key, sess] of this.sessions) {
        if (!sess.idleSince) continue;                       // 在途的不许动
        if (!victim || sess.idleSince < victim.sess.idleSince) victim = { key, sess };
      }
      if (!victim) throw new SessionPoolFullError(this.sessions.size, max);
      console.log(`[SessionPool] 🧹 池满（${this.sessions.size}/${max}），回收最久空闲的 ${victim.key}`);
      this.reset(victim.key).catch(() => {});
      if (this.sessions.has(victim.key)) {                   // reset 没能摘掉 = 防死循环
        this.sessions.delete(victim.key);
      }
    }
  }

  /**
   * 重置（断开 + 清理）单个会话
   */
  async reset(sessionKey) {
    const session = this.sessions.get(sessionKey);
    session?.idleTimer?.cancel();
    if (session?.protocol) {
      this._emit(
        session.agentName, session.protocolName, 'closed',
        session.sessionId, sessionKey, 'reset'
      );
      try { session.protocol.disconnect(); } catch { }
      if (session.launcher) {
        try { await session.launcher.stop(); } catch { }
      }
    }
    this.sessions.delete(sessionKey);
    if (this.activeKey === sessionKey) {
      this.activeKey = null;
    }
  }

  /**
   * 关闭所有会话（并行）
   */
  async closeAll() {
    const keys = [...this.sessions.keys()];
    if (keys.length === 0) return;

    console.log(`[AgentManager] 关闭 ${keys.length} 个会话...`);
    await Promise.allSettled(keys.map(key => this.reset(key)));
    this.sessions.clear();
    this.activeKey = null;
    console.log(`[AgentManager] 所有会话已关闭`);
  }

  /**
   * 会话池统计
   */
  getStats() {
    const list = [];
    for (const [key, session] of this.sessions) {
      list.push({
        key,
        agentName: session.agentName,
        protocol: session.protocolName,
        sessionId: session.sessionId,
        connected: !!session.protocol?.isConnected,
        active: this.activeKey === key,
      });
    }
    return {
      total: list.length,
      connected: list.filter(s => s.connected).length,
      active: this.activeKey,
      sessions: list,
    };
  }

  /**
   * 终止当前活跃会话
   * @returns {object|null} 被终止的会话信息，或 null
   */
  killActive() {
    if (!this.activeKey) return null;
    const session = this.sessions.get(this.activeKey);
    if (!session?.protocol) return null;
    session.idleTimer?.cancel();

    this._emit(
      session.agentName, session.protocolName, 'closed',
      session.sessionId, session.key, 'killed'
    );
    session.protocol.disconnect();
    this.sessions.delete(this.activeKey);
    this.activeKey = null;
    return session;
  }

  /** @private */
  _emit(agentName, protocol, status, sessionId, sessionKey, reason) {
    const payload = { agentName, protocol, status, sessionId, sessionKey };
    if (reason) payload.reason = reason;
    this.emit('session_event', payload);
  }
}

module.exports = { SessionPool, SessionPoolFullError, DEFAULT_MAX_SESSIONS, maxSessions };