diff --git a/pkg/billingexpr/round.go b/pkg/billingexpr/round.go index 7e7e4196..6d7438d8 100644 --- a/pkg/billingexpr/round.go +++ b/pkg/billingexpr/round.go @@ -12,3 +12,10 @@ import "github.com/QuantumNous/new-api/common" func QuotaRound(f float64) int { return common.QuotaRound(f) } + +// QuotaRoundChecked is QuotaRound but also reports whether the result had to +// be saturated. Pre-consume callers use this to reject an unrepresentable +// estimate before any quota is deducted. +func QuotaRoundChecked(f float64) (int, *common.QuotaClamp) { + return common.QuotaRoundChecked(f) +} diff --git a/relay/helper/price.go b/relay/helper/price.go index e2bff583..60fa043e 100644 --- a/relay/helper/price.go +++ b/relay/helper/price.go @@ -32,6 +32,10 @@ func modelPriceNotConfiguredError(modelName string, userId int) error { ) } +func preConsumeQuotaRangeError(modelName string, clamp *common.QuotaClamp) error { + return fmt.Errorf("model %s pre-consume quota is out of range: operation=%s kind=%s value=%g", modelName, clamp.Op, clamp.Kind, clamp.Original) +} + // https://docs.claude.com/en/docs/build-with-claude/prompt-caching#1-hour-cache-duration const claudeCacheCreation1hMultiplier = 6 / 3.75 @@ -117,12 +121,20 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens audioRatio = ratio_setting.GetAudioRatio(info.OriginModelName) audioCompletionRatio = ratio_setting.GetAudioCompletionRatio(info.OriginModelName) ratio := modelRatio * groupRatioInfo.GroupRatio - preConsumedQuota = common.QuotaFromFloat(float64(preConsumedTokens) * ratio) + var clamp *common.QuotaClamp + preConsumedQuota, clamp = common.QuotaFromFloatChecked(float64(preConsumedTokens) * ratio) + if clamp != nil { + return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp) + } } else { if meta.ImagePriceRatio != 0 { modelPrice = modelPrice * meta.ImagePriceRatio } - preConsumedQuota = common.QuotaFromFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + var clamp *common.QuotaClamp + preConsumedQuota, clamp = common.QuotaFromFloatChecked(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + if clamp != nil { + return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp) + } } // check if free model pre-consume is disabled @@ -199,7 +211,11 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types freeModel := false if usePrice { - quota = common.QuotaFromFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + var clamp *common.QuotaClamp + quota, clamp = common.QuotaFromFloatChecked(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + if clamp != nil { + return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp) + } if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 { quota = 0 @@ -208,7 +224,11 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types } } else { // 按量计费：以模型倍率的一半作为预扣额度 - quota = common.QuotaFromFloat(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + var clamp *common.QuotaClamp + quota, clamp = common.QuotaFromFloatChecked(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + if clamp != nil { + return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp) + } modelPrice = -1 if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { if groupRatioInfo.GroupRatio == 0 || modelRatio == 0 { @@ -270,7 +290,10 @@ func modelPriceHelperTiered(c *gin.Context, info *relaycommon.RelayInfo, promptT // Expression coefficients are $/1M tokens prices; convert to quota the same way per-call billing does. quotaBeforeGroup := rawCost / 1_000_000 * common.QuotaPerUnit - preConsumedQuota := billingexpr.QuotaRound(quotaBeforeGroup * groupRatioInfo.GroupRatio) + preConsumedQuota, clamp := billingexpr.QuotaRoundChecked(quotaBeforeGroup * groupRatioInfo.GroupRatio) + if clamp != nil { + return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp) + } freeModel := false if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { diff --git a/relay/helper/price_test.go b/relay/helper/price_test.go index a8410e81..c1a5416c 100644 --- a/relay/helper/price_test.go +++ b/relay/helper/price_test.go @@ -138,3 +138,39 @@ func TestModelPriceHelperTieredPreConsumeMaxTokensFallback(t *testing.T) { }) } } + +func TestModelPriceHelperTieredRejectsPreConsumeOverflow(t *testing.T) { + gin.SetMode(gin.TestMode) + + saved := map[string]string{} + require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error { + saved[key] = value + return nil + })) + t.Cleanup(func() { + require.NoError(t, config.GlobalConfig.LoadFromDB(saved)) + }) + + require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{ + "billing_setting.billing_mode": `{"tiered-overflow-model":"tiered_expr"}`, + "billing_setting.billing_expr": `{"tiered-overflow-model":"tier(\"overflow\", p * 1000000000000000)"}`, + "group_ratio_setting.group_ratio": `{"default":1}`, + })) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + ctx.Set("group", "default") + info := &relaycommon.RelayInfo{ + OriginModelName: "tiered-overflow-model", + UserGroup: "default", + UsingGroup: "default", + BillingRequestInput: &billingexpr.RequestInput{ + Body: []byte(`{}`), + }, + } + + _, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{}) + + require.ErrorContains(t, err, "pre-consume quota is out of range") +} diff --git a/service/billing.go b/service/billing.go index 81daeed8..42cb68d3 100644 --- a/service/billing.go +++ b/service/billing.go @@ -2,6 +2,7 @@ package service import ( "fmt" + "net/http" "github.com/QuantumNous/new-api/logger" relaycommon "github.com/QuantumNous/new-api/relay/common" @@ -17,6 +18,23 @@ const ( // PreConsumeBilling 根据用户计费偏好创建 BillingSession 并执行预扣费。 // 会话存储在 relayInfo.Billing 上，供后续 Settle / Refund 使用。 func PreConsumeBilling(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) *types.NewAPIError { + if relayInfo != nil && relayInfo.QuotaClamp != nil { + clamp := relayInfo.QuotaClamp + return types.NewErrorWithStatusCode( + fmt.Errorf("pre-consume quota is out of range: operation=%s kind=%s value=%g", clamp.Op, clamp.Kind, clamp.Original), + types.ErrorCodeModelPriceError, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } + if preConsumedQuota < 0 { + return types.NewErrorWithStatusCode( + fmt.Errorf("pre-consume quota cannot be negative: %d", preConsumedQuota), + types.ErrorCodeModelPriceError, + http.StatusBadRequest, + types.ErrOptionWithSkipRetry(), + ) + } session, apiErr := NewBillingSession(c, relayInfo, preConsumedQuota) if apiErr != nil { return apiErr diff --git a/service/quota_saturation_test.go b/service/quota_saturation_test.go index ae8cdb6d..bee96828 100644 --- a/service/quota_saturation_test.go +++ b/service/quota_saturation_test.go @@ -1,10 +1,12 @@ package service import ( + "net/http" "testing" "github.com/QuantumNous/new-api/common" relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -72,3 +74,36 @@ func TestAttachQuotaSaturationNoClampNoMarker(t *testing.T) { _, hasAdmin := other["admin_info"] require.False(t, hasAdmin, "no admin_info should be added when there is no clamp") } + +func TestPreConsumeBillingRejectsSaturatedQuotaBeforeDeduction(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(nil) + info := &relaycommon.RelayInfo{ + QuotaClamp: &common.QuotaClamp{ + Op: "QuotaFromFloat", + Kind: common.QuotaClampOverflow, + Original: 1e30, + Clamped: common.MaxQuota, + }, + } + + apiErr := PreConsumeBilling(c, common.MaxQuota, info) + + require.NotNil(t, apiErr) + require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode()) + require.Equal(t, http.StatusBadRequest, apiErr.StatusCode) + require.Nil(t, info.Billing) +} + +func TestPreConsumeBillingRejectsNegativeQuotaBeforeDeduction(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(nil) + info := &relaycommon.RelayInfo{} + + apiErr := PreConsumeBilling(c, -1, info) + + require.NotNil(t, apiErr) + require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode()) + require.Equal(t, http.StatusBadRequest, apiErr.StatusCode) + require.Nil(t, info.Billing) +}