feat(gateway): 添加租户余额管理功能

service/gateway/gateway_http_service.go

@@ -221,6 +221,66 @@ func GetVideoDuration(ctx context.Context, urls []string) (VideoDurationResp, er
 	return resp, nil
 }
 
+// DeductBalanceReq 扣减余额请求
+type DeductBalanceReq struct {
+	Id      uint64  `json:"id"`
+	Surplus float64 `json:"surplus"`
+}
+
+// DeductBalance 扣减租户余额
+func DeductBalance(ctx context.Context, tenantId uint64, amount float64) error {
+	apiURL := "admin-go/api/v1/system/tenant/edit"
+	headers := make(map[string]string)
+	if r := g.RequestFromCtx(ctx); r != nil {
+		for k, v := range r.Request.Header {
+			if len(v) > 0 {
+				headers[k] = v[0]
+			}
+		}
+	}
+
+	body := DeductBalanceReq{
+		Id:      tenantId,
+		Surplus: amount,
+	}
+	jsonData, _ := json.Marshal(body)
+
+	var resp struct{}
+	err := commonHttp.Put(ctx, apiURL, headers, &resp, jsonData)
+	if err != nil {
+		g.Log().Warningf(ctx, "[扣减余额] 失败 tenantId=%d amount=%.6f err=%v", tenantId, amount, err)
+		return err
+	}
+	g.Log().Infof(ctx, "[扣减余额] 成功 tenantId=%d amount=%.6f", tenantId, amount)
+	return nil
+}
+
+// TenantSurplusResp 租户余额返回
+type TenantSurplusResp struct {
+	Surplus float64 `json:"surplus"`
+}
+
+// GetTenantSurplus 获取租户余额
+func GetTenantSurplus(ctx context.Context, tenantId uint64) (float64, error) {
+	apiURL := fmt.Sprintf("admin-go/api/v1/system/tenant/getTenantDetails?tenantId=%d", tenantId)
+	headers := make(map[string]string)
+	if r := g.RequestFromCtx(ctx); r != nil {
+		for k, v := range r.Request.Header {
+			if len(v) > 0 {
+				headers[k] = v[0]
+			}
+		}
+	}
+
+	var resp TenantSurplusResp
+	err := commonHttp.Get(ctx, apiURL, headers, &resp, nil)
+	if err != nil {
+		g.Log().Warningf(ctx, "[获取余额] 失败 tenantId=%d err=%v", tenantId, err)
+		return 0, err
+	}
+	return resp.Surplus, nil
+}
+
 //// callback 向回调地址 POST 任务结果（与查询接口 GetTaskRes 出参一致）
 //func (s *audioTaskService) callback(ctx context.Context, taskID, status, errMsg, callbackURL string) {
 //	if callbackURL == "" {

service/task/worker.go

@@ -39,11 +39,23 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
 		result    map[string]any
 		err       error
 	)
-
 	g.Log().Infof(ctx, "[handleOne] 开始 taskId=%s model=%s", task.TaskID, task.ModelName)
 
 	// ============================================
-	// 1) 调用模型
+	// 1) 查询余额
+	// ============================================
+	var surplus float64
+	if len(model.BillingConfig) > 0 {
+		surplus, _ = gateway.GetTenantSurplus(ctx, task.TenantId)
+		if surplus <= 0 {
+			w.failTask(ctx, task, startTime, "租户余额不足")
+			return
+		}
+		g.Log().Infof(ctx, "[handleOne] 当前余额 tenantId=%d surplus=%.2f", task.TenantId, surplus)
+	}
+
+	// ============================================
+	// 2) 调用模型
 	// ============================================
 	for attempt := 0; ; attempt++ {
 		if attempt > 0 {
@@ -87,7 +99,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
 	}
 
 	// ============================================
-	// 2) 解析返回映射 + 存储 token 相关信息
+	// 3) 解析返回映射 + 存储 token 相关信息
 	// ============================================
 	mapped, err := util.MapResponsePayload(model.ResponseMapping, result)
 	if err != nil {
@@ -112,6 +124,14 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
 		if billingResult != nil {
 			task.BillingData[0] = billingResult
 		}
+
+		if billingResult != nil {
+			task.BillingData[0] = billingResult
+			totalFee := gconv.Float64(billingResult["total_fee"])
+			if totalFee > 0 {
+				_ = gateway.DeductBalance(util.AsyncCtx(ctx), task.TenantId, -totalFee)
+			}
+		}
 	}
 
 	task.ExpendTokens = gconv.Int64(mapped[entity.TotalTokens])
@@ -121,7 +141,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
 	}
 
 	// ============================================
-	// 3) 处理提示词相关数据解析涵盖重试
+	// 4) 处理提示词相关数据解析涵盖重试
 	// ============================================
 	if req.BuildType == public.BuildTypePrompt {
 		mapped, err = w.parseAndRetry(ctx, mapped, model, task, maxRetry)
@@ -133,7 +153,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
 	}
 
 	// ============================================
-	// 4) 上传 OSS（可重试）
+	// 5) 上传 OSS（可重试）
 	// ============================================
 	var oss *gateway.UploadFileResponse
 	for attempt := 0; attempt <= maxRetry; attempt++ {
@@ -152,7 +172,7 @@ func (w *asyncWorker) handleOne(ctx context.Context, task *entity.ModelGatewayTa
 	}
 
 	// ============================================
-	// 5) 成功收尾
+	// 6) 成功收尾
 	// ============================================
 	task.State = public.TaskStatusSuccess
 	task.DurationSeconds = int64(time.Since(startTime).Seconds())
@@ -318,7 +338,12 @@ func (w *asyncWorker) parseAndRetry(ctx context.Context, body map[string]any, mo
 			billingResult := util.CalculateBilling(model.BillingConfig, retryData)
 			if billingResult != nil {
 				task.BillingData = append(task.BillingData, billingResult)
+				totalFee := gconv.Float64(billingResult["total_fee"])
+				if totalFee > 0 {
+					_ = gateway.DeductBalance(util.AsyncCtx(ctx), task.TenantId, -totalFee)
+				}
 			}
+
 			task.ExpendTokens += gconv.Int64(mapped[entity.TotalTokens])
 			_, _ = dao.ModelGatewayTask.Update(ctx, &entity.ModelGatewayTask{
 				SQLBaseDO:    beans.SQLBaseDO{Id: task.Id},