feat(task-pricing): persist matrix billing snapshots

controller/relay.go

@@ -24,6 +24,7 @@ import (
 	"github.com/QuantumNous/new-api/service"
 	"github.com/QuantumNous/new-api/setting"
 	"github.com/QuantumNous/new-api/setting/operation_setting"
+	"github.com/QuantumNous/new-api/setting/ratio_setting"
 	"github.com/QuantumNous/new-api/types"
 
 	"github.com/bytedance/gopkg/util/gopool"
@@ -204,7 +205,6 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 			break
 		}
 
-
 		addUsedChannel(c, channel.Id)
 		bodyStorage, bodyErr := common.GetBodyStorage(c)
 		if bodyErr != nil {
@@ -488,6 +488,95 @@ func RelayTaskFetch(c *gin.Context) {
 	}
 }
 
+func preloadTaskPricingConfig(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError {
+	modelName := strings.TrimSpace(info.OriginModelName)
+	action := strings.TrimSpace(info.Action)
+	contentType := c.Request.Header.Get("Content-Type")
+
+	storage, err := common.GetBodyStorage(c)
+	if err != nil {
+		status := http.StatusBadRequest
+		if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
+			status = http.StatusRequestEntityTooLarge
+		}
+		return service.TaskErrorWrapperLocal(err, "read_request_body_failed", status)
+	}
+	defer func() {
+		_, _ = storage.Seek(0, io.SeekStart)
+		c.Request.Body = io.NopCloser(storage)
+	}()
+
+	switch {
+	case strings.HasPrefix(contentType, "application/json"):
+		body, err := storage.Bytes()
+		if err != nil {
+			return service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusBadRequest)
+		}
+		if strings.TrimSpace(string(body)) != "" {
+			var payload map[string]any
+			if err := common.Unmarshal(body, &payload); err != nil {
+				return service.TaskErrorWrapperLocal(err, "invalid_json", http.StatusBadRequest)
+			}
+			if raw, ok := payload["model"].(string); ok && strings.TrimSpace(raw) != "" {
+				modelName = strings.TrimSpace(raw)
+			}
+			if raw, ok := payload["action"].(string); ok && strings.TrimSpace(raw) != "" {
+				action = strings.TrimSpace(raw)
+			}
+		}
+	case strings.Contains(contentType, gin.MIMEMultipartPOSTForm):
+		form, err := common.ParseMultipartFormReusable(c)
+		if err != nil {
+			return service.TaskErrorWrapperLocal(err, "invalid_multipart_form", http.StatusBadRequest)
+		}
+		if vals := form.Value["model"]; len(vals) > 0 && strings.TrimSpace(vals[0]) != "" {
+			modelName = strings.TrimSpace(vals[0])
+		}
+		if vals := form.Value["action"]; len(vals) > 0 && strings.TrimSpace(vals[0]) != "" {
+			action = strings.TrimSpace(vals[0])
+		}
+	}
+
+	if modelName == "" && action != "" {
+		platform := constant.TaskPlatform(c.GetString("platform"))
+		modelName = service.CoverTaskActionToModelName(platform, action)
+	}
+	if modelName != "" {
+		info.OriginModelName = modelName
+	}
+	if action != "" {
+		info.Action = action
+	}
+
+	info.PricingConfigSnapshotLoaded = true
+	info.PricingConfigSnapshot = ratio_setting.GetPricingConfig(info.OriginModelName)
+	info.RequireMatrixUsageBilling = info.PricingConfigSnapshot != nil &&
+		info.PricingConfigSnapshot.BillingUnit == types.BillingUnitPer1MTokens
+	return nil
+}
+
+func buildTaskBillingContext(info *relaycommon.RelayInfo) *model.TaskBillingContext {
+	bc := &model.TaskBillingContext{
+		ModelPrice:      info.PriceData.ModelPrice,
+		GroupRatio:      info.PriceData.GroupRatioInfo.GroupRatio,
+		ModelRatio:      info.PriceData.ModelRatio,
+		OtherRatios:     info.PriceData.OtherRatios,
+		OriginModelName: info.OriginModelName,
+		PerCallBilling:  common.StringsContains(constant.TaskPricePatches, info.OriginModelName),
+	}
+	if decision := info.PricingDecisionFrozen; decision != nil && decision.BillingMode == types.BillingModeMatrix {
+		bc.BillingMode = decision.BillingMode
+		bc.PricingSnapshot = types.CloneMapAny(decision.Snapshot)
+		bc.BillingUnit = decision.BillingUnit
+		bc.TokenUnitPriceUSD = decision.TokenUnitPriceUSD
+		bc.PerCallBilling = decision.PerCallBilling
+		bc.ModelPrice = decision.PriceUSD
+		bc.GroupRatio = decision.GroupRatioInfo.GroupRatio
+		bc.OtherRatios = types.CloneRatios(decision.OtherRatios)
+	}
+	return bc
+}
+
 func RelayTask(c *gin.Context) {
 	relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
 	if err != nil {
@@ -503,6 +592,10 @@ func RelayTask(c *gin.Context) {
 		respondTaskError(c, taskErr)
 		return
 	}
+	if taskErr := preloadTaskPricingConfig(c, relayInfo); taskErr != nil {
+		respondTaskError(c, taskErr)
+		return
+	}
 
 	var result *relay.TaskSubmitResult
 	var taskErr *dto.TaskError
@@ -524,11 +617,9 @@ func RelayTask(c *gin.Context) {
 
 		if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
 			channel = lockedCh
-			if retryParam.GetRetry() > 0 {
-				if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
-					taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
-					break
-				}
+			if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
+				taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
+				break
 			}
 		} else {
 			var channelErr *types.NewAPIError
@@ -587,14 +678,7 @@ func RelayTask(c *gin.Context) {
 		task.PrivateData.BillingSource = relayInfo.BillingSource
 		task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
 		task.PrivateData.TokenId = relayInfo.TokenId
-		task.PrivateData.BillingContext = &model.TaskBillingContext{
-			ModelPrice:      relayInfo.PriceData.ModelPrice,
-			GroupRatio:      relayInfo.PriceData.GroupRatioInfo.GroupRatio,
-			ModelRatio:      relayInfo.PriceData.ModelRatio,
-			OtherRatios:     relayInfo.PriceData.OtherRatios,
-			OriginModelName: relayInfo.OriginModelName,
-			PerCallBilling:  common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName),
-		}
+		task.PrivateData.BillingContext = buildTaskBillingContext(relayInfo)
 		task.Quota = result.Quota
 		task.Data = result.TaskData
 		task.Action = relayInfo.Action

model/task.go

@@ -109,12 +109,16 @@ type TaskPrivateData struct {
 
 // TaskBillingContext 记录任务提交时的计费参数，以便轮询阶段可以重新计算额度。
 type TaskBillingContext struct {
-	ModelPrice      float64            `json:"model_price,omitempty"`       // 模型单价
-	GroupRatio      float64            `json:"group_ratio,omitempty"`       // 分组倍率
-	ModelRatio      float64            `json:"model_ratio,omitempty"`       // 模型倍率
-	OtherRatios     map[string]float64 `json:"other_ratios,omitempty"`      // 附加倍率（时长、分辨率等）
-	OriginModelName string             `json:"origin_model_name,omitempty"` // 模型名称，必须为OriginModelName
-	PerCallBilling  bool               `json:"per_call_billing,omitempty"`  // 按次计费：跳过轮询阶段的差额结算
+	ModelPrice        float64            `json:"model_price,omitempty"`       // 模型单价
+	GroupRatio        float64            `json:"group_ratio,omitempty"`       // 分组倍率
+	ModelRatio        float64            `json:"model_ratio,omitempty"`       // 模型倍率
+	OtherRatios       map[string]float64 `json:"other_ratios,omitempty"`      // 附加倍率（时长、分辨率等）
+	OriginModelName   string             `json:"origin_model_name,omitempty"` // 模型名称，必须为OriginModelName
+	PerCallBilling    bool               `json:"per_call_billing,omitempty"`  // 按次计费：跳过轮询阶段的差额结算
+	BillingMode       string             `json:"billing_mode,omitempty"`
+	PricingSnapshot   map[string]any     `json:"pricing_snapshot,omitempty"`
+	BillingUnit       string             `json:"billing_unit,omitempty"`
+	TokenUnitPriceUSD float64            `json:"token_unit_price_usd,omitempty"`
 }
 
 // GetUpstreamTaskID 获取上游真实 task ID（用于与 provider 通信）

relay/relay_task.go

@@ -19,6 +19,7 @@ import (
 	relayconstant "github.com/QuantumNous/new-api/relay/constant"
 	"github.com/QuantumNous/new-api/relay/helper"
 	"github.com/QuantumNous/new-api/service"
+	"github.com/QuantumNous/new-api/types"
 	"github.com/gin-gonic/gin"
 )
 
@@ -87,28 +88,19 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr
 	}
 	info.LockedChannel = ch
 
-	if originTask.ChannelId != info.ChannelId {
-		key, _, newAPIError := ch.GetNextEnabledKey()
-		if newAPIError != nil {
-			return service.TaskErrorWrapper(newAPIError, "channel_no_available_key", newAPIError.StatusCode)
-		}
-		common.SetContextKey(c, constant.ContextKeyChannelKey, key)
-		common.SetContextKey(c, constant.ContextKeyChannelType, ch.Type)
-		common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, ch.GetBaseURL())
-		common.SetContextKey(c, constant.ContextKeyChannelId, originTask.ChannelId)
-
-		info.ChannelBaseUrl = ch.GetBaseURL()
-		info.ChannelId = originTask.ChannelId
-		info.ChannelType = ch.Type
-		info.ApiKey = key
-	}
-
 	// 提取 remix 参数（时长、分辨率 → OtherRatios）
 	if info.Action == constant.TaskActionRemix {
 		if originTask.PrivateData.BillingContext != nil {
 			// 新的 remix 逻辑：直接从原始任务的 BillingContext 中提取 OtherRatios（如果存在）
-			for s, f := range originTask.PrivateData.BillingContext.OtherRatios {
-				info.PriceData.AddOtherRatio(s, f)
+			bc := originTask.PrivateData.BillingContext
+			info.OriginPricing = &types.OriginPricingSnapshot{
+				BillingMode:       bc.BillingMode,
+				BillingUnit:       bc.BillingUnit,
+				PriceUSD:          bc.ModelPrice,
+				TokenUnitPriceUSD: bc.TokenUnitPriceUSD,
+				Snapshot:          types.CloneMapAny(bc.PricingSnapshot),
+				OtherRatios:       types.CloneRatios(bc.OtherRatios),
+				PerCallBilling:    bc.PerCallBilling,
 			}
 		} else {
 			// 旧的 remix 逻辑：直接从 task data 解析 seconds 和 size（如果存在）
@@ -120,13 +112,16 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr
 				seconds = 4
 			}
 			sizeStr, _ := taskData["size"].(string)
-			if info.PriceData.OtherRatios == nil {
-				info.PriceData.OtherRatios = map[string]float64{}
+			otherRatios := map[string]float64{
+				"seconds": float64(seconds),
+				"size":    1,
 			}
-			info.PriceData.OtherRatios["seconds"] = float64(seconds)
-			info.PriceData.OtherRatios["size"] = 1
 			if sizeStr == "1792x1024" || sizeStr == "1024x1792" {
-				info.PriceData.OtherRatios["size"] = 1.666667
+				otherRatios["size"] = 1.666667
+			}
+			info.OriginPricing = &types.OriginPricingSnapshot{
+				OtherRatios:    otherRatios,
+				PerCallBilling: false,
 			}
 		}
 	}
@@ -181,18 +176,22 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
 		return nil, pricingErr
 	}
 	info.PriceData = priceData
+	matrixOrOriginPricing := info.OriginPricing != nil ||
+		(info.PricingDecisionFrozen != nil && info.PricingDecisionFrozen.BillingMode == types.BillingModeMatrix)
 
 	// 5. 计费估算：让适配器根据用户请求提供 OtherRatios（时长、分辨率等）
 	//    必须在 ModelPriceHelperPerCall 之后调用（它会重建 PriceData）。
 	//    ResolveOriginTask 可能已在 remix 路径中预设了 OtherRatios，此处合并。
-	if estimatedRatios := adaptor.EstimateBilling(c, info); len(estimatedRatios) > 0 {
-		for k, v := range estimatedRatios {
-			info.PriceData.AddOtherRatio(k, v)
+	if !matrixOrOriginPricing {
+		if estimatedRatios := adaptor.EstimateBilling(c, info); len(estimatedRatios) > 0 {
+			for k, v := range estimatedRatios {
+				info.PriceData.AddOtherRatio(k, v)
+			}
 		}
 	}
 
 	// 6. 将 OtherRatios 应用到基础额度
-	if !common.StringsContains(constant.TaskPricePatches, modelName) {
+	if !matrixOrOriginPricing && !common.StringsContains(constant.TaskPricePatches, modelName) {
 		for _, ra := range info.PriceData.OtherRatios {
 			if ra != 1.0 {
 				info.PriceData.Quota = int(float64(info.PriceData.Quota) * ra)
@@ -240,11 +239,13 @@ func RelayTaskSubmit(c *gin.Context, info *relaycommon.RelayInfo) (*TaskSubmitRe
 
 	// 11. 提交后计费调整：让适配器根据上游实际返回调整 OtherRatios
 	finalQuota := info.PriceData.Quota
-	if adjustedRatios := adaptor.AdjustBillingOnSubmit(info, taskData); len(adjustedRatios) > 0 {
-		// 基于调整后的 ratios 重新计算 quota
-		finalQuota = recalcQuotaFromRatios(info, adjustedRatios)
-		info.PriceData.OtherRatios = adjustedRatios
-		info.PriceData.Quota = finalQuota
+	if !matrixOrOriginPricing {
+		if adjustedRatios := adaptor.AdjustBillingOnSubmit(info, taskData); len(adjustedRatios) > 0 {
+			// 基于调整后的 ratios 重新计算 quota
+			finalQuota = recalcQuotaFromRatios(info, adjustedRatios)
+			info.PriceData.OtherRatios = adjustedRatios
+			info.PriceData.Quota = finalQuota
+		}
 	}
 
 	return &TaskSubmitResult{

service/task_billing.go

@@ -118,6 +118,12 @@ func taskBillingOther(task *model.Task) map[string]interface{} {
 	if bc := task.PrivateData.BillingContext; bc != nil {
 		other["model_price"] = bc.ModelPrice
 		other["group_ratio"] = bc.GroupRatio
+		if bc.BillingMode != "" {
+			other["billing_mode"] = bc.BillingMode
+			other["billing_unit"] = bc.BillingUnit
+			other["token_unit_price_usd"] = bc.TokenUnitPriceUSD
+			other["pricing_snapshot"] = bc.PricingSnapshot
+		}
 		if len(bc.OtherRatios) > 0 {
 			for k, v := range bc.OtherRatios {
 				other[k] = v
@@ -241,6 +247,42 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int
 // RecalculateTaskQuotaByTokens 根据实际 token 消耗重新计费（异步差额结算）。
 // 当任务成功且返回了 totalTokens 时，根据模型倍率和分组倍率重新计算实际扣费额度，
 // 与预扣费的差额进行补扣或退还。支持钱包和订阅计费来源。
+
+func RecalculateTaskQuotaWithOther(ctx context.Context, task *model.Task, actualQuota int, reason string, extra map[string]any) {
+	if len(extra) == 0 {
+		RecalculateTaskQuota(ctx, task, actualQuota, reason)
+		return
+	}
+	if task.PrivateData.BillingContext == nil {
+		task.PrivateData.BillingContext = &model.TaskBillingContext{}
+	}
+	if task.PrivateData.BillingContext.PricingSnapshot == nil {
+		task.PrivateData.BillingContext.PricingSnapshot = map[string]any{}
+	}
+	for k, v := range extra {
+		task.PrivateData.BillingContext.PricingSnapshot[k] = v
+	}
+	RecalculateTaskQuota(ctx, task, actualQuota, reason)
+}
+
+func RecalculateMatrixUsageTaskQuota(ctx context.Context, task *model.Task, totalTokens int) {
+	if totalTokens <= 0 {
+		return
+	}
+	bc := task.PrivateData.BillingContext
+	if bc == nil {
+		return
+	}
+	actualQuota := int(float64(totalTokens) * bc.TokenUnitPriceUSD / 1_000_000 * common.QuotaPerUnit * bc.GroupRatio)
+	if actualQuota <= 0 && totalTokens > 0 && bc.TokenUnitPriceUSD > 0 && bc.GroupRatio > 0 {
+		actualQuota = 1
+	}
+	reason := fmt.Sprintf("matrix usage billing: tokens=%d, token_unit_price_usd=%.6f, group_ratio=%.4f", totalTokens, bc.TokenUnitPriceUSD, bc.GroupRatio)
+	RecalculateTaskQuotaWithOther(ctx, task, actualQuota, reason, map[string]any{
+		"total_tokens": totalTokens,
+	})
+}
+
 func RecalculateTaskQuotaByTokens(ctx context.Context, task *model.Task, totalTokens int) {
 	if totalTokens <= 0 {
 		return

service/task_polling.go

@@ -17,6 +17,7 @@ import (
 	"github.com/QuantumNous/new-api/model"
 	"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
 	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/types"
 
 	"github.com/samber/lo"
 )
@@ -520,6 +521,24 @@ func truncateBase64(s string) string {
 //  2. taskResult.TotalTokens > 0 → 按 token 重算
 //  3. 都不满足 → 保持预扣额度不变
 func settleTaskBillingOnComplete(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, taskResult *relaycommon.TaskInfo) {
+	if bc := task.PrivateData.BillingContext; bc != nil && bc.BillingMode == types.BillingModeMatrix {
+		_, billingUnit, _, err := types.ResolveMatrixPricingSnapshot(bc.BillingUnit, bc.PricingSnapshot)
+		if err != nil {
+			logger.LogError(ctx, fmt.Sprintf("task %s matrix pricing snapshot invalid: %s", task.TaskID, err.Error()))
+			return
+		}
+		if billingUnit == types.BillingUnitPerCall {
+			logger.LogInfo(ctx, fmt.Sprintf("task %s matrix per_call billing, skip settlement", task.TaskID))
+			return
+		}
+		if billingUnit == types.BillingUnitPer1MTokens {
+			RecalculateMatrixUsageTaskQuota(ctx, task, taskResult.TotalTokens)
+			return
+		}
+		logger.LogError(ctx, fmt.Sprintf("task %s unknown matrix billing_unit: %s", task.TaskID, billingUnit))
+		return
+	}
+
 	// 0. 按次计费的任务不做差额结算
 	if bc := task.PrivateData.BillingContext; bc != nil && bc.PerCallBilling {
 		logger.LogInfo(ctx, fmt.Sprintf("任务 %s 按次计费，跳过差额结算", task.TaskID))

types/model_pricing.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"errors"
 	"fmt"
 
 	"github.com/QuantumNous/new-api/common"
@@ -131,3 +132,63 @@ func NormalizePricingConfigDefaults(cfg *PricingConfig) {
 		cfg.PreconsumeStrategy = PreconsumeStrategyExact
 	}
 }
+
+func ResolveMatrixPricingSnapshot(billingUnit string, snapshot map[string]any) (price float64, resolvedBillingUnit string, pricingMode string, err error) {
+	if snapshot == nil {
+		return 0, "", "", errors.New("pricing snapshot is required")
+	}
+	rawPrice, ok := snapshot["price"]
+	if !ok {
+		return 0, "", "", errors.New("pricing snapshot missing price")
+	}
+	price, ok = snapshotNumber(rawPrice)
+	if !ok {
+		return 0, "", "", errors.New("pricing snapshot price must be numeric")
+	}
+	rawMode, ok := snapshot["pricing_mode"].(string)
+	if !ok || rawMode == "" {
+		return 0, "", "", errors.New("pricing snapshot pricing_mode must be string")
+	}
+	switch PricingMode(rawMode) {
+	case PricingModeMatched, PricingModeFallbackMax, PricingModeFallbackDefault:
+		pricingMode = rawMode
+	default:
+		return 0, "", "", fmt.Errorf("unsupported pricing_mode %q", rawMode)
+	}
+
+	snapshotUnit, _ := snapshot["billing_unit"].(string)
+	if snapshotUnit != "" && snapshotUnit != BillingUnitPerCall && snapshotUnit != BillingUnitPer1MTokens {
+		return 0, "", "", fmt.Errorf("unsupported snapshot billing_unit %q", snapshotUnit)
+	}
+	if billingUnit != "" && billingUnit != BillingUnitPerCall && billingUnit != BillingUnitPer1MTokens {
+		return 0, "", "", fmt.Errorf("unsupported billing_unit %q", billingUnit)
+	}
+	switch {
+	case billingUnit != "" && snapshotUnit != "" && billingUnit != snapshotUnit:
+		return 0, "", "", fmt.Errorf("billing_unit %q does not match snapshot billing_unit %q", billingUnit, snapshotUnit)
+	case billingUnit != "":
+		resolvedBillingUnit = billingUnit
+	case snapshotUnit != "":
+		resolvedBillingUnit = snapshotUnit
+	default:
+		resolvedBillingUnit = BillingUnitPerCall
+	}
+	return price, resolvedBillingUnit, pricingMode, nil
+}
+
+func snapshotNumber(v any) (float64, bool) {
+	switch n := v.(type) {
+	case float64:
+		return n, true
+	case float32:
+		return float64(n), true
+	case int:
+		return float64(n), true
+	case int64:
+		return float64(n), true
+	case int32:
+		return float64(n), true
+	default:
+		return 0, false
+	}
+}