ai-agent/workflow/service/flow/billing.go

package flow

import (
	"context"
	"errors"
	"fmt"

	commonHttp "gitea.redpowerfuture.com/red-future/common/http"
	"gitea.redpowerfuture.com/red-future/common/utils"
	"github.com/gogf/gf/v2/os/glog"
	"github.com/gogf/gf/v2/os/gtime"
	"github.com/gogf/gf/v2/util/gconv"
	"github.com/google/uuid"

	nodeDao "ai-agent/workflow/dao/node"
	sessionDao "ai-agent/workflow/dao/session"
	nodeDto "ai-agent/workflow/model/dto/node"
	"ai-agent/workflow/model/entity"
)

// ====================== 计费本地 DTO（shop-user-trade pricing，独立 module 不可 import，JSON 对齐） ======================

type pricingOpenOrderReq struct {
	UserId      int64  `json:"userId"`
	SubjectType string `json:"subjectType"`
	SubjectID   string `json:"subjectId"`
	ChargeMode  string `json:"chargeMode"`
	BizOrderNo  string `json:"bizOrderNo"`
}

type pricingGetConfigReq struct {
	SubjectType string `json:"subjectType"`
	SubjectID   string `json:"subjectId"`
}

type pricingGetConfigRes struct {
	Enabled int `json:"enabled"`
}

type pricingChargeOrderInfo struct {
	ID           int64   `json:"id"`
	Status       int     `json:"status"`       // 1已建单 2已结算 3已失败
	ActualAmount float64 `json:"actualAmount"` // 实收/实扣（元）：settle/cancel 响应回填，回写 exec_workflow.actual_amount
	CreatedAt    string  `json:"createdAt"`
	ChargeMode   string  `json:"chargeMode"` // per_item / per_token
}

// pricingSettleReq Settle 与 Cancel 同形状（OrderId + Usage）
type pricingSettleReq struct {
	OrderId int64          `json:"orderId"`
	Usage   map[string]any `json:"usage"`
}

type pricingFailReq struct {
	OrderId int64  `json:"orderId"`
	Reason  string `json:"reason"`
}

// 计价对象/模式常量（对齐 shop-user-trade consts/pricing）
const (
	pricingSubjectWorkflow     = "workflow"
	pricingChargeModePerItem   = "per_item"
	pricingChargeModePerSecond = "per_second"
	pricingChargeModePerToken  = "per_token"
	pricingOrderStatusCreated  = 1
)

// errBillingGateBlocked 计费门禁拦截（余额不足/钱包不可用/费率非法/未配置计价/用户缺失/per_second 无视频模型）：执行终局失败，
// 不进进程内重试、不落 recoverable（shouldRetry 与 handleExecute 分类均排除）。
var errBillingGateBlocked = errors.New("计费门禁拦截")

// pricingURL 组装 shop-user-trade 计价接口地址。
// 跨服务路由前缀 /pricing/controller/ 由 common http.RouteRegister 按 controller struct 名推导
// （pricingController → pricing/controller），GoFrame doSetHandler 恒 prefix+uri 拼接。
func pricingURL(sub string) string {
	return "shop-user-trade/pricing/controller/" + sub
}

// ====================== 建单（执行开始，不动钱） ======================

// openBillingOrder 工作流执行开始建计费单（不动钱）。
// 幂等键 bizOrderNo=wf:{execId}；复用终态 execId 重跑（原单已结算/失败）→ 开新单 wf:{execId}:{uuid8}，
// 新单 ID 落 exec_workflow.charge_order_id，结算据此定位。
// 门禁（余额>=min_balance、钱包须存在）失败 → errBillingGateBlocked，执行终局失败不重试。
// 计费是工作流执行的前置条件：用户缺失/计价未配置/预检失败/per_second 无视频模型 → 终局失败，不免费跑。
func openBillingOrder(ctx context.Context, execId int64, flowContent *entity.FlowInfo) error {
	user, err := utils.GetUserInfo(ctx)
	if err != nil || user == nil || user.Id == 0 {
		return fmt.Errorf("%w: 取不到用户 %v", errBillingGateBlocked, err)
	}
	// config/get 预检 = 服务存活探针 + 启用开关：
	// 预检失败（服务宕机/路由不通）→ 终局失败；
	// 预检通过（服务确认在、计价已开）后再调 open_order，其失败即为业务错误（余额/钱包/费率）→ 阻塞，
	// 确保 shop-user-trade 宕机时工作流执行也被阻断而非免费跑。
	var cfg pricingGetConfigRes
	if err = commonHttp.Get(ctx, pricingURL("config/get"), utils.HeadersFromCtx(ctx, utils.HeadersOptions{TokenFromQuery: true}), &cfg,
		"subjectType", pricingSubjectWorkflow, "subjectId", pricingSubjectWorkflow); err != nil {
		return fmt.Errorf("%w: 计价配置查询失败 %v", errBillingGateBlocked, err)
	}
	if cfg.Enabled != 1 {
		return fmt.Errorf("%w: 计价未启用", errBillingGateBlocked)
	}
	chargeMode := ""
	if flowContent != nil && flowContent.ChargeMode != "" {
		chargeMode = flowContent.ChargeMode
	} else {
		return fmt.Errorf("%w: 工作流执行未选择计费模式", errBillingGateBlocked)
	}
	// per_second 按秒计费以生成视频总时长为基础，工作流须含视频模型节点，否则配置非法 → 终局失败
	if chargeMode == pricingChargeModePerSecond && !flowHasVideoModel(ctx, flowContent) {
		return fmt.Errorf("%w: per_second 计费须工作流包含视频模型", errBillingGateBlocked)
	}
	info, err := openPricingOrder(ctx, int64(user.Id), chargeMode, fmt.Sprintf("wf:%d", execId))
	if err != nil {
		return err
	}
	orderId := info.ID
	if info.Status != pricingOrderStatusCreated {
		// 复用终态 execId：原单已结算/失败，开新单让本次运行独立计费
		info, err = openPricingOrder(ctx, int64(user.Id), chargeMode,
			fmt.Sprintf("wf:%d:%s", execId, uuid.NewString()[:8]))
		if err != nil {
			return err
		}
		orderId = info.ID
	}
	if err = sessionDao.ExecWorkflowDao.UpdateMap(ctx, execId, map[string]any{
		entity.ExecWorkflowCol.ChargeOrderId: orderId,
	}); err != nil {
		// 结算时按 bizOrderNo=wf:{execId} 兜底定位，不阻塞
		glog.Errorf(ctx, "工作流计费:记录 charge_order_id 失败 execId=%d: %v", execId, err)
	}
	return nil
}

// flowHasVideoModel 工作流是否包含视频模型节点（per_second 计费前置条件）。
// 遍历节点，按 modelId 去重后经 isVideoModel 查模型类型，任一为视频模型即 true；
// flowContent 缺失/查模型失败按无视频模型处理（per_second 被拦截，fail-closed）。
func flowHasVideoModel(ctx context.Context, flowContent *entity.FlowInfo) bool {
	if flowContent == nil {
		return false
	}
	seen := make(map[int64]struct{})
	for i := range flowContent.Nodes {
		modelId := flowContent.Nodes[i].ModelConfig.ModelId
		if modelId <= 0 {
			continue
		}
		if _, ok := seen[modelId]; ok {
			continue
		}
		seen[modelId] = struct{}{}
		if isVideoModel(ctx, modelId) {
			return true
		}
	}
	return false
}

// openPricingOrder 调 shop-user-trade 建单（幂等：同 subject+bizOrderNo 返回既有单）
func openPricingOrder(ctx context.Context, userId int64, chargeMode, bizOrderNo string) (*pricingChargeOrderInfo, error) {
	info := new(pricingChargeOrderInfo)
	err := commonHttp.Post(ctx, pricingURL("open_order"), utils.HeadersFromCtx(ctx, utils.HeadersOptions{TokenFromQuery: true}), info, &pricingOpenOrderReq{
		UserId:      userId,
		SubjectType: pricingSubjectWorkflow,
		SubjectID:   pricingSubjectWorkflow,
		ChargeMode:  chargeMode,
		BizOrderNo:  bizOrderNo,
	})
	if err != nil {
		return nil, fmt.Errorf("%w: %v", errBillingGateBlocked, err)
	}
	if info.ID == 0 {
		return nil, fmt.Errorf("%w: 未返回计费单ID", errBillingGateBlocked)
	}
	return info, nil
}

// ====================== 结算（终态） ======================

// settleBilling 工作流终态计费（recordWorkflow 汇聚全部执行路径后调用）：
// 成功→Settle 实收；用户取消→Cancel 按已消耗实收；永久失败（retryable=0/重试耗尽）→Fail 不扣费；
// 可恢复失败→跳过（订单留 CREATED，恢复续跑后结算/失败）。
// 全部计费调用错误仅记日志，不拖垮工作流终态落库。
func settleBilling(ctx context.Context, execId int64, runErr error, retryable, retryCount *int) {
	exec, err := sessionDao.ExecWorkflowDao.GetById(ctx, execId)
	if err != nil || exec == nil {
		glog.Errorf(ctx, "工作流计费:查询执行失败,跳过结算 execId=%d: %v", execId, err)
		return
	}
	orderId := exec.ChargeOrderId
	if orderId == 0 {
		// 兜底：charge_order_id 未落库（UpdateMap 失败）时按 bizOrderNo=wf:{execId} 定位；查不到→跳过
		info, e := getPricingOrder(ctx, "", fmt.Sprintf("wf:%d", execId))
		if e != nil || info == nil {
			return
		}
		orderId = info.ID
	}
	usage, full, err := workflowChargeUsage(ctx, exec, orderId, errors.Is(runErr, context.Canceled))
	if err != nil {
		glog.Errorf(ctx, "工作流计费:计算用量失败,跳过结算 execId=%d: %v", execId, err)
		return
	}
	// 终局实扣金额：settle/cancel 由 shop 结算响应回填（元）；fail / 可恢复=0（未扣费）
	var actual float64
	switch {
	case runErr == nil:
		actual, _ = callSettlePricing(ctx, orderId, usage, pricingURL("settle"))
	case errors.Is(runErr, context.Canceled):
		actual, _ = callSettlePricing(ctx, orderId, usage, pricingURL("cancel")) // Cancel 按已消耗实收
	case retryable != nil && *retryable == 0:
		callFailPricing(ctx, orderId, runErr.Error())
	case retryCount != nil && *retryCount >= execMaxRetryCount:
		callFailPricing(ctx, orderId, runErr.Error()) // 重试耗尽 → 永久失败
	default:
		// 可恢复失败（retryable=1 且未耗尽/关停中断）：订单留 CREATED，恢复续跑后结算
	}
	// 终局回填 exec_workflow（成功/取消/失败/可恢复统一落库，前端与对账均看此行）：
	// 模型消耗（total_tokens/total_fee，本次运行节点 token_info 按订单窗口聚合）+ 业务实扣
	// （actual_amount=settle/cancel 实收金额，失败/可恢复未结算=0）。失败/取消路径 SummaryLambda 不跑，
	// total 列此前恒为空，此处按同一窗口补齐（与结算口径一致，不混入上一运行残留）；
	// 可恢复失败也先落已消耗，续跑成功后同一订单收敛重算覆盖。错误仅记日志不拖垮终态落库。
	if full != nil {
		if err := sessionDao.ExecWorkflowDao.UpdateMap(ctx, execId, map[string]any{
			entity.ExecWorkflowCol.TotalTokens:  full.TotalTokens,
			entity.ExecWorkflowCol.TotalFee:     full.TotalFee,
			entity.ExecWorkflowCol.ActualAmount: actual,
		}); err != nil {
			glog.Errorf(ctx, "exec_workflow 回填消耗/实扣失败 execId=%d: %v", execId, err)
		}
	}
}

// workflowChargeUsage 计算工作流结算用量并返回全量聚合（full，含 TotalTokens/TotalFee 供终局回填 exec 行）：
// per_token → feeByModel（各模型按次已消耗费用，结算侧按此合计实收——每次模型调用在发生时已由
// shop /calc 计价，含「不足1分按1分」的按次兜底与调用时媒体/费率快照，不再按聚合 token 重算，
// 避免两笔 0.01 合并重算成 0.01）；
// per_item / per_second（其余模式）→ durationSec = 本次生成视频总时长
// （各视频模型节点记录里模型返回时长之和）。仅生成视频的工作流按时间计费，其余模型调用按条/token 计费。
// 用户取消（forCancel=true）时 per_item/per_second 在时长之外补收已消耗 token：把非视频节点的按次费用
// 一并上报（视频节点消耗已由时长计价覆盖，排除防双计）——中途取消通常无视频产出（durationSec=0），
// 但文本/分析节点可能已完成并消耗了 token，须按已消耗补收而非按 0 计。
func workflowChargeUsage(ctx context.Context, exec *entity.ExecWorkflow, orderId int64, forCancel bool) (usage map[string]any, full *nodeUsageAgg, err error) {
	info, err := getPricingOrder(ctx, gconv.String(orderId), "")
	if err != nil {
		return nil, nil, err
	}
	// 按订单收敛：只聚合订单创建后产生的节点记录。每次重跑（上单已终态）开新单，created_at 各自独立，
	// 隔离「上次已结算运行」与「本次运行」——否则重跑后取消会把上一次已扣费的 token 一起再扣。
	// shop 返回 gtime.String() 无时区本地墙钟，与 node_execution.created_at（timestamp without tz）同格式可比。
	var createdAtFrom *gtime.Time
	if info.CreatedAt != "" {
		createdAtFrom = gtime.NewFromStr(info.CreatedAt)
		if createdAtFrom == nil || createdAtFrom.IsZero() {
			return nil, nil, fmt.Errorf("计费单创建时间解析失败: %s", info.CreatedAt)
		}
	}
	full, nonVideo, err := workflowNodeUsage(ctx, exec, createdAtFrom)
	if err != nil {
		return nil, nil, err
	}
	switch {
	case info.ChargeMode == pricingChargeModePerToken:
		// per_token：按次已记录费用结算（含视频模型——per_token 无时长计价，视频模型按自身按次费用计收）
		return tokenUsageMap(full), full, nil
	case forCancel:
		// per_item/per_second 取消补收：时长（full，通常 0）+ 非视频节点按次已消耗费用
		// （nonVideo 排除视频节点：其消耗已由时长计价覆盖，避免双计）
		usage := tokenUsageMap(nonVideo)
		usage["durationSec"] = full.DurationSec
		return usage, full, nil
	default:
		// 正常结算走时长（视频产出按 per_item/per_second 档位/秒价），不附 token/费用拆分
		return map[string]any{"durationSec": full.DurationSec}, full, nil
	}
}

// nodeUsageAgg 本次执行聚合出的结算用量：FeeByModel（按生效（系统）模型 id 的按次已记录费用合计，
// shop 实收依据）+ DurationSec（时长，供 per_item/per_second 用）。逐模型 token/媒体明细不上报——
// 每调用一条留在 node_execution.token_info（model_id/total_tokens/prompt_tokens/completion_tokens/
// media_type/total_fee），订单层按需可从明细再聚合，不再冗余携带。
type nodeUsageAgg struct {
	// FeeByModel 各模型窗口内已消耗费用合计 = 节点 token_info.total_fee 求和。total_fee 本身是
	// 该节点内各次模型调用（每次经 shop /calc 计价，calculator.charge 内部已 ceilFen，「不足1分按1分」
	// 按调用次生效）费用之和 → 此处合计即「按次已记录费用」。结算侧按此实收，不再按聚合 token 重算。
	FeeByModel map[string]float64
	// DurationSec 各视频节点生成视频总时长（per_item/per_second 计价依据）
	DurationSec float64
	// TotalTokens / TotalFee 全量节点消耗合计（模型消耗 token / 模型按次费用），
	// 终局回填 exec_workflow.total_tokens/total_fee（区别于钱包实扣 ActualAmount）。
	TotalTokens int64
	TotalFee    float64
}

// workflowNodeUsage 聚合本次执行（FlowExecutionId + 订单创建时间下界）下各节点执行记录写入的用量。
// 返回两组聚合：
//   - full：全部节点记录——per_token 结算（视频模型也按自身按次费用计收，无时长计价）与时长累计用；
//   - nonVideo：排除 total_duration>0 的节点（视频节点生成时长，per_item/per_second 取消补收时其消耗
//     已由时长计价覆盖，不再按次补收，避免双计）。
//
// 聚合内容：feeByModel = 各模型 total_fee 求和（total_fee = 该节点内各次模型调用经 shop /calc 计价
// （已 ceilFen）的费用之和 → 按次已记录费用，shop 实收依据）；durationSec = 各视频节点生成视频总时长
// （模型返回，total_duration）累加。逐模型 token/媒体明细留在 node_execution.token_info，订单层不上报。
//
// 收敛到当前运行：重跑复用同一 exec 记录与节点组（检查点恢复的 SavedFlowInput 携带旧 node_group_id，
// exec_workflow 表也无该列持久化），node_group_id 无法区分运行；改按 created_at >= 订单创建时间过滤——
// 每次重跑开新单，各自 created_at 隔离本次运行消耗，避免把已结算的上一次运行用量一起计入。
func workflowNodeUsage(ctx context.Context, exec *entity.ExecWorkflow, createdAtFrom *gtime.Time) (full, nonVideo *nodeUsageAgg, err error) {
	records, _, err := nodeDao.NodeExecutionDao.ListByFlowExecutionId(ctx, &nodeDto.ListNodeExecutionByFlowReq{
		FlowExecutionId: exec.Id,
		CreatedAtFrom:   createdAtFrom,
	})
	if err != nil {
		return nil, nil, err
	}
	full = &nodeUsageAgg{FeeByModel: make(map[string]float64)}
	nonVideo = &nodeUsageAgg{FeeByModel: make(map[string]float64)}
	for _, rec := range records {
		for _, ti := range rec.TokenInfo {
			addNodeUsageEntry(full, ti)
			// total_duration>0 = 视频节点产出了生成时长（lambda 只有视频模型节点累加并落库该字段），
			// 该节点消耗由时长计价覆盖，排除出 nonVideo（per_item/per_second 取消补收按此计，防双计）
			if gconv.Float64(ti["total_duration"]) <= 0 {
				addNodeUsageEntry(nonVideo, ti)
			}
		}
	}
	return full, nonVideo, nil
}

// addNodeUsageEntry 把单条 token_info（节点一次模型调用或调用汇总）累加进聚合。
// gconv.String 兼容字符串与 JSONB 数字（float64）两种 model_id 写入，避免 (string) 断言丢弃条目。
func addNodeUsageEntry(agg *nodeUsageAgg, ti map[string]any) {
	modelID := gconv.String(ti["model_id"])
	agg.DurationSec += gconv.Float64(ti["total_duration"])
	agg.TotalTokens += gconv.Int64(ti["total_tokens"])
	agg.TotalFee += gconv.Float64(ti["total_fee"])
	if modelID == "" {
		return // 无模型 id（异常条目）不计费用（feeByModel 按模型键）
	}
	agg.FeeByModel[modelID] += gconv.Float64(ti["total_fee"])
}

// tokenUsageMap 按次费用口径组装结算用量（token 部分）：仅 feeByModel（shop 实收依据，
// per_token 无建单快照）。逐模型 token/媒体明细留在 node_execution.token_info，订单层不上报。
func tokenUsageMap(agg *nodeUsageAgg) map[string]any {
	return map[string]any{
		"feeByModel": agg.FeeByModel,
	}
}

// getPricingOrder 查计费单：id>0 按ID，否则按 subjectType+subjectId+bizOrderNo
func getPricingOrder(ctx context.Context, id, bizOrderNo string) (*pricingChargeOrderInfo, error) {
	info := new(pricingChargeOrderInfo)
	var data []any
	if id != "" {
		data = []any{"id", id}
	} else {
		data = []any{"subjectType", pricingSubjectWorkflow, "subjectId", pricingSubjectWorkflow, "bizOrderNo", bizOrderNo}
	}
	err := commonHttp.Get(ctx, pricingURL("order"), utils.HeadersFromCtx(ctx, utils.HeadersOptions{TokenFromQuery: true}), info, data...)
	if err != nil {
		return nil, err
	}
	if info.ID == 0 {
		return nil, errors.New("计费单不存在")
	}
	return info, nil
}

// callSettlePricing Settle（成功）或 Cancel（用户取消按已消耗实收），返回实收金额（元）。
// shop settle/cancel 响应体是结算后的 ChargeOrderInfo（actualAmount=本次实扣），幂等重复结算返回既有金额；
// 调用失败返回 0 并记日志（此时无法确知钱包是否已扣，actual_amount 置 0 保守，不臆造金额）。
func callSettlePricing(ctx context.Context, orderId int64, usage map[string]any, url string) (actual float64, err error) {
	info := new(pricingChargeOrderInfo)
	if err := commonHttp.Post(ctx, url, utils.HeadersFromCtx(ctx, utils.HeadersOptions{TokenFromQuery: true}), info,
		&pricingSettleReq{OrderId: orderId, Usage: usage}); err != nil {
		glog.Errorf(ctx, "工作流计费:结算失败 orderId=%d: %v", orderId, err)
		return 0, err
	}
	return info.ActualAmount, nil
}

// callFailPricing Fail（不扣费）
func callFailPricing(ctx context.Context, orderId int64, reason string) {
	if err := commonHttp.Post(ctx, pricingURL("fail"), utils.HeadersFromCtx(ctx, utils.HeadersOptions{TokenFromQuery: true}), &pricingChargeOrderInfo{},
		&pricingFailReq{OrderId: orderId, Reason: reason}); err != nil {
		glog.Errorf(ctx, "工作流计费:失败处理失败 orderId=%d: %v", orderId, err)
	}
}