宜博 / lemonboss

/**
 * Unit Tests — 账户欠费门 + 自动恢复（27-P4 T5）
 *
 * 覆盖：
 *   - setBillingGate 未注入时对现有行为零影响（默认「未欠费」）
 *   - 欠费时 advanceOne / advanceBatch 都会 suspendRunTree，reason='billing-arrears'，
 *     写进 graph_runs.block_reason，ok:false
 *   - startRunSync「启动」时清空 block_reason
 *   - resumeBillingClearedRuns：欠费清零后自动恢复；仍欠费的留在 blocked 且不重复推进
 *   - advanceBatch 的排序 bug 回归：候选先查、budget/billing 后查——一个已经收敛
 *     （零 ready 节点）只是恰好预算也耗尽的 run，必须正确收敛到 completed，
 *     不能被误判成 budget-exhausted
 *
 * 隔离手法与 graph-engine-advance-batch.test.js 一致。
 */
import { describe, it, expect, beforeAll, afterAll, beforeEach } from 'vitest';
import fs from 'fs';
import os from 'os';
import path from 'path';

describe('账户欠费门 + 自动恢复（27-P4 T5）', () => {
  let wdb;
  let engine;
  let runsDir;
  const TEST_DB_PATH = '/tmp/test-graph-engine-billing-gate.sqlite';

  beforeAll(() => {
    for (const f of [TEST_DB_PATH, `${TEST_DB_PATH}-shm`, `${TEST_DB_PATH}-wal`]) {
      try { fs.unlinkSync(f); } catch {}
    }
    process.env.GRAPH_DB_PATH = TEST_DB_PATH;
    runsDir = fs.mkdtempSync(path.join(os.tmpdir(), 'graph-billing-gate-runs-'));
    process.env.GRAPH_RUNS_DIR = runsDir;
    for (const m of [
      '../../services/graph-db',
      '../../services/graph-service',
      '../../services/org-graph-service',
      '../../services/work-graph-db',
      '../../services/graph-run-sse',
      '../../services/graph-engine',
    ]) {
      delete require.cache[require.resolve(m)];
    }
    wdb = require('../../services/work-graph-db');
    engine = require('../../services/graph-engine');
  });

  afterAll(() => {
    for (const f of [TEST_DB_PATH, `${TEST_DB_PATH}-shm`, `${TEST_DB_PATH}-wal`]) {
      try { fs.unlinkSync(f); } catch {}
    }
    delete process.env.GRAPH_DB_PATH;
    try { fs.rmSync(runsDir, { recursive: true, force: true }); } catch {}
    delete process.env.GRAPH_RUNS_DIR;
  });

  beforeEach(() => {
    const db = wdb.getDb();
    db.pragma('foreign_keys = OFF');
    db.exec(`
      DELETE FROM graph_run_loop_records;
      DELETE FROM graph_run_nodes;
      DELETE FROM graph_runs;
      DELETE FROM work_graph_edges;
      DELETE FROM work_graph_nodes;
      DELETE FROM work_graphs;
      DELETE FROM system_agents;
    `);
    db.pragma('foreign_keys = ON');
    engine.setAgentExecutor(null);
    engine.setBillingGate(null); // 每个用例结束都要回默认桩子，不然会串到下一个用例
  });

  let seq = 0;
  const nextCompany = () => `p4t5-co-${++seq}`;

  function ensureAgents(ids) {
    const db = wdb.getDb();
    const ts = new Date().toISOString();
    const ins = db.prepare(`
      INSERT OR IGNORE INTO system_agents
        (id, name, scope, owner_company_id, legacy_path, status, created_at, updated_at)
      VALUES (?, ?, 'public', NULL, ?, 'active', ?, ?)
    `);
    for (const id of new Set(ids)) ins.run(id, id, `system/agents/${id}`, ts, ts);
  }

  function buildSoloRun({ companyId = nextCompany(), maxConcurrency = null } = {}) {
    ensureAgents(['solo']);
    const g = wdb.createWorkGraph({ companyId, nodes: ['solo'], edges: [] });
    expect(g.errors || []).toEqual([]);
    const run = wdb.createRun({
      graphKind: 'work', graphId: g.id, companyId, executionMode: 'manual', maxConcurrency,
    });
    engine.markRunReady(run.id);
    return { runId: run.id, companyId };
  }

  it('未注入 billingGateFn 时零影响：checkBilling 恒 ok，节点正常跑完', async () => {
    const { runId } = buildSoloRun();
    engine.setAgentExecutor(async () => ({ status: 'success' }));
    engine.startRun(runId);
    const res = await engine.advanceOne(runId);
    expect(res.ok).toBe(true);
    expect(wdb.getRun(runId).run_status).toBe('completed');
  });

  it('advanceOne：欠费时挂起整棵树，reason=billing-arrears，block_reason 落库', async () => {
    const { runId } = buildSoloRun();
    engine.setAgentExecutor(async () => ({ status: 'success' })); // 不该被调到
    engine.setBillingGate(() => true); // 恒欠费
    engine.startRun(runId);

    const res = await engine.advanceOne(runId);
    expect(res.ok).toBe(false);
    expect(res.reason).toBe('billing-arrears');
    expect(res.billing).toEqual({ ok: false, inArrears: true });

    const run = wdb.getRun(runId);
    expect(run.run_status).toBe('blocked');
    expect(run.block_reason).toBe('billing-arrears');
    // 欠费门跟 checkBudget 位置一致——都在「拾到 ready 节点」之后、真正 claim（ready→running
    // 的条件更新）之前查，所以挂起时节点还停在 ready，没有被 claim（更没有被结算成任何终态）。
    expect(wdb.getRunState(runId).nodes[0].status).toBe('ready');
  });

  it('advanceBatch（maxConcurrency>1）：欠费时同样挂起，reason=billing-arrears', async () => {
    const companyId = nextCompany();
    ensureAgents(['entry', 'a', 'b']);
    const g = wdb.createWorkGraph({
      companyId,
      nodes: ['entry', 'a', 'b'],
      edges: [{ from: 'entry', to: 'a' }, { from: 'entry', to: 'b' }], // 单入口，fan-out 出 a/b
    });
    expect(g.errors || []).toEqual([]);
    const run = wdb.createRun({ graphKind: 'work', graphId: g.id, companyId, executionMode: 'manual' });
    engine.markRunReady(run.id);
    engine.setAgentExecutor(async () => ({ status: 'success' }));
    engine.startRun(run.id);
    await engine.advanceOne(run.id); // entry → success，点亮 a/b（这一步不欠费）

    engine.setBillingGate(() => ({ inArrears: true })); // 现在才欠费
    const res = await engine.advanceBatch(run.id, { maxConcurrency: 3 });
    expect(res.ok).toBe(false);
    expect(res.reason).toBe('billing-arrears');
    expect(wdb.getRun(run.id).run_status).toBe('blocked');
    expect(wdb.getRun(run.id).block_reason).toBe('billing-arrears');
  });

  it('startRunSync「启动」清空 block_reason：欠费挂起 → 手动点启动 → block_reason 变 null', async () => {
    const { runId } = buildSoloRun();
    engine.setAgentExecutor(async () => ({ status: 'success' }));
    engine.setBillingGate(() => true);
    engine.startRun(runId);
    await engine.advanceOne(runId);
    expect(wdb.getRun(runId).block_reason).toBe('billing-arrears');

    engine.setBillingGate(null); // 欠费清零
    engine.startRun(runId); // 人工点「启动」
    expect(wdb.getRun(runId).block_reason).toBeNull();
    expect(wdb.getRun(runId).run_status).toBe('running');
  });

  it('resumeBillingClearedRuns：欠费清零的 run 被自动恢复并跑完，仍欠费的原地不动', async () => {
    const { runId: cleared } = buildSoloRun({ companyId: 'p4t5-cleared' });
    const { runId: stillArrears } = buildSoloRun({ companyId: 'p4t5-still' });

    engine.setAgentExecutor(async () => ({ status: 'success' }));
    engine.setBillingGate(() => true);
    engine.startRun(cleared);
    await engine.advanceOne(cleared);
    engine.startRun(stillArrears);
    await engine.advanceOne(stillArrears);
    expect(wdb.getRun(cleared).run_status).toBe('blocked');
    expect(wdb.getRun(stillArrears).run_status).toBe('blocked');

    // 只有 `cleared` 那家公司欠费清零了，`stillArrears` 还欠着
    engine.setBillingGate((companyId) => companyId === 'p4t5-still');

    const sweep = await engine.resumeBillingClearedRuns();
    expect(sweep.scanned).toBe(2);
    expect(sweep.resumed).toEqual([cleared]);
    expect(sweep.stillBlocked).toEqual([stillArrears]);

    expect(wdb.getRun(cleared).run_status).toBe('running'); // 恢复了，节点原地继续跑（手动模式停在 running）
    expect(wdb.getRun(cleared).block_reason).toBeNull();
    expect(wdb.getRun(stillArrears).run_status).toBe('blocked'); // 纹丝不动
    expect(wdb.getRun(stillArrears).block_reason).toBe('billing-arrears');
  });

  it('resumeBillingClearedRuns 支持按 companyId 只扫一家', async () => {
    const { runId: a } = buildSoloRun({ companyId: 'p4t5-scope-a' });
    const { runId: b } = buildSoloRun({ companyId: 'p4t5-scope-b' });
    engine.setAgentExecutor(async () => ({ status: 'success' }));
    engine.setBillingGate(() => true);
    engine.startRun(a); await engine.advanceOne(a);
    engine.startRun(b); await engine.advanceOne(b);

    engine.setBillingGate(null); // 两家都清零了，但这次只扫 a
    const sweep = await engine.resumeBillingClearedRuns({ companyId: 'p4t5-scope-a' });
    expect(sweep.scanned).toBe(1);
    expect(sweep.resumed).toEqual([a]);
    expect(wdb.getRun(a).run_status).toBe('running');
    expect(wdb.getRun(b).run_status).toBe('blocked'); // 没被这次扫描碰到
  });

  it('advanceBatch 排序回归：候选先查、budget/billing 后查——已收敛（零 ready 节点）'
    + '只是恰好预算/欠费也耗尽的 run，正确收敛到 completed，不会被误判成挂起', async () => {
    const { runId } = buildSoloRun();
    engine.startRun(runId);
    // 绕开引擎直接把节点写成 success（不经过 settleNode，run_status 就不会被顺带收敛），
    // 构造「run_status 还是 running，但已经没有 ready/running 节点」这个状态——
    // 同款手法 graph-engine-advance-batch.test.js 已经在用。
    const node = wdb.getRunState(runId).nodes[0];
    expect(wdb.updateNodeStatus(node.id, 'ready', 'running')).toBe(1);
    expect(wdb.updateNodeStatus(node.id, 'running', 'success')).toBe(1);

    // 预算与欠费都打成「耗尽/欠费」——如果 advanceBatch 在查候选之前就先查这两道门，
    // 会误判成 budget-exhausted 或 billing-arrears 并把这个其实已经跑完的 run 挂起。
    wdb.getDb().prepare(`UPDATE graph_runs SET budget_json = ? WHERE id = ?`)
      .run(JSON.stringify({ total_agent_calls: 0, spent: { agent_calls: 999 } }), runId);
    engine.setBillingGate(() => true);

    const res = await engine.advanceBatch(runId, { maxConcurrency: 3 });
    expect(res.ok).toBe(false);
    expect(res.reason).toBe('no-ready-node');
    expect(wdb.getRun(runId).run_status).toBe('completed');
    expect(wdb.getRun(runId).block_reason).toBeNull();
  });
});