openclaw-voc-skill/claude-code/claude-code-voc-intelligence/mcp/src/core/batch-budget-guard.js

'use strict';

const { queryBilling } = require('./billing-client');
const { checkBudget, buildEstimateBands } = require('./budget-guard');

/**
 * Re-checks the authenticated balance and actual trace spend before each
 * collection batch. The query is read-only; NewAPI still charges each real
 * gateway call independently.
 */
function createBatchBudgetGuard({ input = {}, context = {}, operationFactory } = {}) {
  const budgetLimitCny = input.budgetCny ?? input.budgetLimitCny;
  const enabled = budgetLimitCny != null || input.enforceBudget === true || input.enforceBudget === 'true';
  let last = { enabled, status: 'disabled', allowed: true, checks: 0 };

  async function beforeBatch(batch = {}) {
    if (!enabled) return last;
    const operations = typeof operationFactory === 'function'
      ? operationFactory(batch)
      : [];
    if (!Array.isArray(operations) || operations.length === 0) return last;

    let result;
    try {
      const query = typeof input.queryBillingImpl === 'function' ? input.queryBillingImpl : queryBilling;
      result = await query({
        ...input,
        operations,
        include: ['balance', 'usage'],
        usageTraceId: context.usageTraceId,
        reportId: context.reportId,
        usageFilter: {
          ...(input.usageFilter || {}),
          usageTraceId: context.usageTraceId,
          ...(context.reportId ? { reportId: context.reportId } : {}),
          groupBy: ['usageTraceId'],
          page: 1,
          pageSize: 1
        },
        createPaymentQr: false
      });
    } catch (error) {
      last = {
        enabled,
        status: 'unknown',
        allowed: true,
        checks: last.checks + 1,
        warning: `批次预算检查暂时不可用：${error.message || String(error)}`
      };
      return last;
    }

    if (result?.status !== 'ok') {
      last = {
        enabled,
        status: 'unknown',
        allowed: true,
        checks: last.checks + 1,
        warning: result?.assistantMessage || '批次预算检查暂时不可用'
      };
      return last;
    }

    const estimate = result.data?.estimate || {};
    const usage = result.data?.usage?.totals || {};
    const bands = buildEstimateBands(estimate, { lowFactor: 1, highFactor: 1 });
    const budget = checkBudget({
      estimate,
      bands,
      balance: result.data?.balance || {},
      budgetLimitCny,
      actualCostCny: usage.costCny || 0,
      nextCostCny: estimate.amountCny || 0,
      rechargeBufferCny: input.rechargeBufferCny
    });
    last = {
      ...budget,
      enabled,
      checks: last.checks + 1,
      batchIndex: Number(batch.batchIndex || 0),
      keyword: batch.keyword || null,
      usageTraceId: context.usageTraceId,
      operations
    };
    return last;
  }

  beforeBatch.getState = () => ({ ...last });
  return { beforeBatch, snapshot: () => ({ ...last }) };
}

function findBudgetIssue(errors = []) {
  return errors.find(error => String(error?.kind || error?.errorKind || '') === 'budget') || null;
}

function budgetStoppedMessage(issue = {}) {
  const budget = issue.budget || issue;
  if (budget.status === 'needs_recharge') {
    return `批次预算检查发现当前余额不足，已停止新增调用。预计本批需 ¥${Number(budget.nextCostCny || 0).toFixed(2)}，当前缺口约 ¥${Number(budget.shortfallCny || 0).toFixed(2)}。`;
  }
  return `批次预算检查已阻止继续执行，避免超过本地预算上限 ¥${Number(budget.budgetLimitCny || 0).toFixed(2)}。`;
}

module.exports = { createBatchBudgetGuard, findBudgetIssue, budgetStoppedMessage };