From 621927f710b46727223280fedbe32121bab59837 Mon Sep 17 00:00:00 2001 From: CaIon Date: Fri, 10 Jul 2026 21:53:51 +0800 Subject: [PATCH 1/3] fix(billing): reject saturated pre-consume quota --- pkg/billingexpr/round.go | 7 +++++++ relay/helper/price.go | 33 ++++++++++++++++++++++++----- relay/helper/price_test.go | 36 ++++++++++++++++++++++++++++++++ service/billing.go | 18 ++++++++++++++++ service/quota_saturation_test.go | 35 +++++++++++++++++++++++++++++++ 5 files changed, 124 insertions(+), 5 deletions(-) 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) +} From d9595831bf05b59198742dca7e779bbea3165141 Mon Sep 17 00:00:00 2001 From: CaIon Date: Fri, 10 Jul 2026 22:30:38 +0800 Subject: [PATCH 2/3] fix(billing): improve quota handling and error reporting for pre-consume operations --- common/quota_math.go | 57 +++++++++--- common/quota_math_test.go | 18 ++++ dto/openai_image.go | 13 ++- pkg/billingexpr/round.go | 8 +- relay/helper/openai_image_request_test.go | 11 +++ relay/helper/price.go | 50 ++++++----- relay/helper/price_test.go | 102 +++++++++++++++++++++- relay/image_handler.go | 10 --- service/billing.go | 3 +- service/quota_saturation_test.go | 4 + service/text_quota_test.go | 28 ++++++ types/request_meta.go | 3 +- 12 files changed, 246 insertions(+), 61 deletions(-) diff --git a/common/quota_math.go b/common/quota_math.go index 1a5a5ac7..86c72492 100644 --- a/common/quota_math.go +++ b/common/quota_math.go @@ -16,11 +16,14 @@ const ( MinQuota = math.MinInt32 ) +// QuotaClampKind identifies why a quota conversion had to be saturated. +type QuotaClampKind string + // Clamp kinds reported by QuotaClamp.Kind. const ( - QuotaClampOverflow = "overflow" - QuotaClampUnderflow = "underflow" - QuotaClampNaN = "nan" + QuotaClampOverflow QuotaClampKind = "overflow" + QuotaClampUnderflow QuotaClampKind = "underflow" + QuotaClampNaN QuotaClampKind = "nan" ) // QuotaClamp describes a single saturation event: a quota conversion whose @@ -28,10 +31,19 @@ const ( // therefore clamped. It is surfaced to billing callers so the event can be // recorded on the related consume/task log for admin auditing. type QuotaClamp struct { - Op string `json:"op"` // "QuotaFromFloat" | "QuotaRound" | "QuotaFromDecimal" - Kind string `json:"kind"` // "overflow" | "underflow" | "nan" - Original float64 `json:"original"` // best-effort pre-clamp value (decimal -> float64 approx) - Clamped int `json:"clamped"` // the saturated result actually used + Op string `json:"op"` // "QuotaFromFloat" | "QuotaRound" | "QuotaFromDecimal" + Kind QuotaClampKind `json:"kind"` // "overflow" | "underflow" | "nan" + Original float64 `json:"original"` // best-effort pre-clamp value (decimal -> float64 approx) + Clamped int `json:"clamped"` // the saturated result actually used +} + +// Error lets the same typed value serve both as the settlement audit marker +// and as the fail-fast error returned by strict pre-consume conversions. +func (c *QuotaClamp) Error() string { + if c == nil { + return "" + } + return fmt.Sprintf("quota conversion (%s) %s: original=%g, clamped=%d", c.Op, c.Kind, c.Original, c.Clamped) } // AuditMap renders the clamp as the marker stored under a log's @@ -58,19 +70,26 @@ func (c *QuotaClamp) AuditMap() map[string]interface{} { // record the event (e.g. on the consume log); the returned pointer is nil for // in-range values. func saturateQuota(value float64, op string) (int, *QuotaClamp) { + var clamp *QuotaClamp switch { case math.IsNaN(value): - SysError(fmt.Sprintf("quota conversion (%s) received NaN, falling back to 0", op)) - return 0, &QuotaClamp{Op: op, Kind: QuotaClampNaN, Original: value, Clamped: 0} + clamp = &QuotaClamp{Op: op, Kind: QuotaClampNaN, Original: value, Clamped: 0} case value >= MaxQuota: - SysError(fmt.Sprintf("quota conversion (%s) overflow: %g exceeds max quota, clamped to %d", op, value, MaxQuota)) - return MaxQuota, &QuotaClamp{Op: op, Kind: QuotaClampOverflow, Original: value, Clamped: MaxQuota} + clamp = &QuotaClamp{Op: op, Kind: QuotaClampOverflow, Original: value, Clamped: MaxQuota} case value <= MinQuota: - SysError(fmt.Sprintf("quota conversion (%s) underflow: %g below min quota, clamped to %d", op, value, MinQuota)) - return MinQuota, &QuotaClamp{Op: op, Kind: QuotaClampUnderflow, Original: value, Clamped: MinQuota} + clamp = &QuotaClamp{Op: op, Kind: QuotaClampUnderflow, Original: value, Clamped: MinQuota} default: return int(value), nil } + SysError(clamp.Error()) + return clamp.Clamped, clamp +} + +func strictQuota(quota int, clamp *QuotaClamp) (int, error) { + if clamp != nil { + return 0, clamp + } + return quota, nil } // QuotaFromFloat converts a computed quota value to int, truncating toward @@ -87,6 +106,12 @@ func QuotaFromFloatChecked(value float64) (int, *QuotaClamp) { return saturateQuota(value, "QuotaFromFloat") } +// QuotaFromFloatStrict converts an in-range value and returns a typed +// *QuotaClamp error instead of allowing a saturated result to reach billing. +func QuotaFromFloatStrict(value float64) (int, error) { + return strictQuota(QuotaFromFloatChecked(value)) +} + // QuotaRound converts a float64 quota value to int using half-away-from-zero // rounding, with saturation. Every tiered billing path (pre-consume, // settlement, breakdown validation, log fields) MUST use this to avoid +-1 @@ -102,6 +127,12 @@ func QuotaRoundChecked(value float64) (int, *QuotaClamp) { return saturateQuota(math.Round(value), "QuotaRound") } +// QuotaRoundStrict rounds an in-range value and returns a typed *QuotaClamp +// error instead of allowing a saturated result to reach billing. +func QuotaRoundStrict(value float64) (int, error) { + return strictQuota(QuotaRoundChecked(value)) +} + // QuotaFromDecimal converts a computed quota decimal to int with saturation. // The decimal is rounded (half away from zero) before conversion. func QuotaFromDecimal(d decimal.Decimal) int { diff --git a/common/quota_math_test.go b/common/quota_math_test.go index 63a1de29..2d8742e6 100644 --- a/common/quota_math_test.go +++ b/common/quota_math_test.go @@ -6,6 +6,7 @@ import ( "github.com/shopspring/decimal" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) // 2000 quota per call * n=18446744073686646784 overflows int64; the constant @@ -78,6 +79,23 @@ func TestQuotaFromFloatChecked(t *testing.T) { } } +func TestQuotaFromFloatStrictReturnsTypedClampError(t *testing.T) { + quota, err := QuotaFromFloatStrict(42.9) + require.NoError(t, err) + assert.Equal(t, 42, quota) + + quota, err = QuotaFromFloatStrict(overflowingProduct) + assert.Zero(t, quota) + var clamp *QuotaClamp + require.ErrorAs(t, err, &clamp) + assert.Equal(t, QuotaClampOverflow, clamp.Kind) + assert.Equal(t, MaxQuota, clamp.Clamped) + assert.ErrorContains(t, err, "QuotaFromFloat") + assert.ErrorContains(t, err, "overflow") + assert.ErrorContains(t, err, "original=") + assert.ErrorContains(t, err, "clamped=2147483647") +} + // TestQuotaRoundChecked verifies the rounding entry point reports clamps the // same way. func TestQuotaRoundChecked(t *testing.T) { diff --git a/dto/openai_image.go b/dto/openai_image.go index ef99c1c7..275fd555 100644 --- a/dto/openai_image.go +++ b/dto/openai_image.go @@ -155,14 +155,19 @@ func (i *ImageRequest) GetTokenCountMeta() *types.TokenCountMeta { } } - // n is NOT included here; it is handled via OtherRatio("n") in - // image_handler.go (default) or channel adaptors (actual count). - // Including n here caused double-counting for channels that also - // set OtherRatio("n") (e.g. Ali/Bailian). + imageN := uint(1) + if i.N != nil && *i.N > 0 { + imageN = *i.N + } + + // Keep n separate from ImagePriceRatio so size/quality and count remain + // independent billing dimensions. Fixed-price pre-consume stores this on + // PriceData, and image settlement reuses or replaces the same "n" ratio. return &types.TokenCountMeta{ CombineText: i.Prompt, MaxTokens: 1584, ImagePriceRatio: sizeRatio * qualityRatio, + BillingRatios: map[string]float64{"n": float64(imageN)}, } } diff --git a/pkg/billingexpr/round.go b/pkg/billingexpr/round.go index 6d7438d8..80d7d318 100644 --- a/pkg/billingexpr/round.go +++ b/pkg/billingexpr/round.go @@ -13,9 +13,7 @@ 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) +// QuotaRoundStrict rejects an unrepresentable pre-consume estimate. +func QuotaRoundStrict(f float64) (int, error) { + return common.QuotaRoundStrict(f) } diff --git a/relay/helper/openai_image_request_test.go b/relay/helper/openai_image_request_test.go index 7068b895..e9fb1b99 100644 --- a/relay/helper/openai_image_request_test.go +++ b/relay/helper/openai_image_request_test.go @@ -108,6 +108,16 @@ func TestGetAndValidOpenAIImageRequestNBounds(t *testing.T) { body: fmt.Sprintf(`{"model":"gpt-image-1","prompt":"a cat","n":%d}`, dto.MaxImageN), wantN: dto.MaxImageN, }, + { + name: "explicit n is accepted", + body: `{"model":"gpt-image-1","prompt":"a cat","n":3}`, + wantN: 3, + }, + { + name: "zero n defaults to 1", + body: `{"model":"gpt-image-1","prompt":"a cat","n":0}`, + wantN: 1, + }, { name: "absent n defaults to 1", body: `{"model":"gpt-image-1","prompt":"a cat"}`, @@ -127,6 +137,7 @@ func TestGetAndValidOpenAIImageRequestNBounds(t *testing.T) { require.NoError(t, err) require.NotNil(t, req.N) require.Equal(t, tt.wantN, *req.N) + require.Equal(t, float64(tt.wantN), req.GetTokenCountMeta().BillingRatios["n"]) }) } diff --git a/relay/helper/price.go b/relay/helper/price.go index 60fa043e..2e8ebb2d 100644 --- a/relay/helper/price.go +++ b/relay/helper/price.go @@ -32,10 +32,6 @@ 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 @@ -121,20 +117,15 @@ 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 - var clamp *common.QuotaClamp - preConsumedQuota, clamp = common.QuotaFromFloatChecked(float64(preConsumedTokens) * ratio) - if clamp != nil { - return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp) + quota, err := common.QuotaFromFloatStrict(float64(preConsumedTokens) * ratio) + if err != nil { + return types.PriceData{}, err } + preConsumedQuota = quota } else { if meta.ImagePriceRatio != 0 { modelPrice = modelPrice * meta.ImagePriceRatio } - 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 @@ -172,6 +163,17 @@ func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens CacheCreation1hRatio: cacheCreationRatio1h, QuotaToPreConsume: preConsumedQuota, } + if usePrice { + for name, ratio := range meta.BillingRatios { + priceData.AddOtherRatio(name, ratio) + } + quotaToPreConsume := priceData.ApplyOtherRatiosToFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + quota, err := common.QuotaFromFloatStrict(quotaToPreConsume) + if err != nil { + return types.PriceData{}, err + } + priceData.QuotaToPreConsume = quota + } if common.DebugEnabled { logger.LogDebug(c, "model_price_helper result: %s", priceData.ToSetting()) @@ -211,10 +213,10 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types freeModel := false if usePrice { - var clamp *common.QuotaClamp - quota, clamp = common.QuotaFromFloatChecked(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) - if clamp != nil { - return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp) + var err error + quota, err = common.QuotaFromFloatStrict(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + if err != nil { + return types.PriceData{}, err } if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 { @@ -224,10 +226,10 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types } } else { // 按量计费：以模型倍率的一半作为预扣额度 - var clamp *common.QuotaClamp - quota, clamp = common.QuotaFromFloatChecked(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio) - if clamp != nil { - return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp) + var err error + quota, err = common.QuotaFromFloatStrict(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio) + if err != nil { + return types.PriceData{}, err } modelPrice = -1 if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume { @@ -290,9 +292,9 @@ 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, clamp := billingexpr.QuotaRoundChecked(quotaBeforeGroup * groupRatioInfo.GroupRatio) - if clamp != nil { - return types.PriceData{}, preConsumeQuotaRangeError(info.OriginModelName, clamp) + preConsumedQuota, err := billingexpr.QuotaRoundStrict(quotaBeforeGroup * groupRatioInfo.GroupRatio) + if err != nil { + return types.PriceData{}, err } freeModel := false diff --git a/relay/helper/price_test.go b/relay/helper/price_test.go index c1a5416c..ef396146 100644 --- a/relay/helper/price_test.go +++ b/relay/helper/price_test.go @@ -10,6 +10,7 @@ import ( relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/setting/billing_setting" "github.com/QuantumNous/new-api/setting/config" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -52,7 +53,9 @@ func TestModelPriceHelperTieredUsesPreloadedRequestInput(t *testing.T) { }, } - priceData, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{}) + priceData, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{ + BillingRatios: map[string]float64{"n": 3}, + }) require.NoError(t, err) require.Equal(t, 1500, priceData.QuotaToPreConsume) require.NotNil(t, info.TieredBillingSnapshot) @@ -172,5 +175,100 @@ func TestModelPriceHelperTieredRejectsPreConsumeOverflow(t *testing.T) { _, err := ModelPriceHelper(ctx, info, 1000, &types.TokenCountMeta{}) - require.ErrorContains(t, err, "pre-consume quota is out of range") + var clamp *common.QuotaClamp + require.ErrorAs(t, err, &clamp) + require.Equal(t, "QuotaRound", clamp.Op) + require.Equal(t, common.QuotaClampOverflow, clamp.Kind) +} + +func TestModelPriceHelperRequestBillingRatiosOnlyApplyToFixedPrice(t *testing.T) { + gin.SetMode(gin.TestMode) + savedModelPrices := ratio_setting.ModelPrice2JSONString() + savedModelRatios := ratio_setting.ModelRatio2JSONString() + t.Cleanup(func() { + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(savedModelPrices)) + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(savedModelRatios)) + }) + + modelPrices, err := common.Marshal(map[string]float64{ + "fixed-image-price": 0.04, + "fractional-image-price": 0.0000012, + "overflow-image-price": float64(common.MaxQuota) / common.QuotaPerUnit / 2, + }) + require.NoError(t, err) + require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(string(modelPrices))) + modelRatios, err := common.Marshal(map[string]float64{"ratio-image-price": 15}) + require.NoError(t, err) + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(string(modelRatios))) + + tests := []struct { + name string + model string + wantQuota int + wantUsePrice bool + wantImageCount bool + }{ + { + name: "fixed price applies image count", + model: "fixed-image-price", + wantQuota: 180000, + wantUsePrice: true, + wantImageCount: true, + }, + { + name: "ratio price ignores request billing ratios", + model: "ratio-image-price", + wantQuota: 15000, + wantUsePrice: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set("group", "default") + info := &relaycommon.RelayInfo{ + OriginModelName: tt.model, + UserGroup: "default", + UsingGroup: "default", + } + meta := &types.TokenCountMeta{ + ImagePriceRatio: 3, + BillingRatios: map[string]float64{"n": 3}, + } + + priceData, err := ModelPriceHelper(ctx, info, 1000, meta) + + require.NoError(t, err) + require.Equal(t, tt.wantQuota, priceData.QuotaToPreConsume) + require.Equal(t, tt.wantUsePrice, priceData.UsePrice) + require.Equal(t, tt.wantImageCount, priceData.HasOtherRatio("n")) + require.Equal(t, priceData.OtherRatios(), info.PriceData.OtherRatios()) + }) + } + + newInfo := func(model string) (*gin.Context, *relaycommon.RelayInfo) { + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + ctx.Set("group", "default") + return ctx, &relaycommon.RelayInfo{ + OriginModelName: model, + UserGroup: "default", + UsingGroup: "default", + } + } + meta := &types.TokenCountMeta{BillingRatios: map[string]float64{"n": 3}} + + ctx, info := newInfo("fractional-image-price") + priceData, err := ModelPriceHelper(ctx, info, 0, meta) + require.NoError(t, err) + // 0.0000012 * 500000 * 3 = 1.8, then truncate once to 1. + require.Equal(t, 1, priceData.QuotaToPreConsume) + + ctx, info = newInfo("overflow-image-price") + _, err = ModelPriceHelper(ctx, info, 0, meta) + var clamp *common.QuotaClamp + require.ErrorAs(t, err, &clamp) + require.Equal(t, "QuotaFromFloat", clamp.Op) + require.Equal(t, common.QuotaClampOverflow, clamp.Kind) + require.Nil(t, info.Billing) } diff --git a/relay/image_handler.go b/relay/image_handler.go index 52aa3e78..0c21fabc 100644 --- a/relay/image_handler.go +++ b/relay/image_handler.go @@ -123,16 +123,6 @@ func ImageHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *type imageN = *request.N } - // n is handled via OtherRatio so it is applied exactly once in quota - // calculation (both price-based and ratio-based paths). - // Adaptors may have already set a more accurate count from the - // upstream response; only set the default when they haven't. - if info.PriceData.UsePrice { // only price model use N ratio - if !info.PriceData.HasOtherRatio("n") { - info.PriceData.AddOtherRatio("n", float64(imageN)) - } - } - if usage.(*dto.Usage).TotalTokens == 0 { usage.(*dto.Usage).TotalTokens = 1 } diff --git a/service/billing.go b/service/billing.go index 42cb68d3..02a093b6 100644 --- a/service/billing.go +++ b/service/billing.go @@ -19,9 +19,8 @@ const ( // 会话存储在 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), + relayInfo.QuotaClamp, types.ErrorCodeModelPriceError, http.StatusBadRequest, types.ErrOptionWithSkipRetry(), diff --git a/service/quota_saturation_test.go b/service/quota_saturation_test.go index bee96828..518bdde5 100644 --- a/service/quota_saturation_test.go +++ b/service/quota_saturation_test.go @@ -92,6 +92,10 @@ func TestPreConsumeBillingRejectsSaturatedQuotaBeforeDeduction(t *testing.T) { require.NotNil(t, apiErr) require.Equal(t, types.ErrorCodeModelPriceError, apiErr.GetErrorCode()) require.Equal(t, http.StatusBadRequest, apiErr.StatusCode) + require.Same(t, info.QuotaClamp, apiErr.Err) + var clamp *common.QuotaClamp + require.ErrorAs(t, apiErr, &clamp) + require.Same(t, info.QuotaClamp, clamp) require.Nil(t, info.Billing) } diff --git a/service/text_quota_test.go b/service/text_quota_test.go index c18fc47e..3b2bcb91 100644 --- a/service/text_quota_test.go +++ b/service/text_quota_test.go @@ -490,3 +490,31 @@ func TestTryTieredSettleNoClampInRange(t *testing.T) { require.NotNil(t, result) require.Nil(t, relayInfo.QuotaClamp, "in-range settlement must not record a clamp") } + +func TestCalculateTextQuotaSummaryFixedPriceAppliesImageCountOnceAndAllowsOverride(t *testing.T) { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + priceData := types.PriceData{ + ModelPrice: 0.12, + UsePrice: true, + GroupRatioInfo: types.GroupRatioInfo{ + GroupRatio: 1, + }, + } + priceData.AddOtherRatio("n", 3) + relayInfo := &relaycommon.RelayInfo{ + OriginModelName: "dall-e-3", + PriceData: priceData, + StartTime: time.Now(), + } + usage := &dto.Usage{PromptTokens: 1, TotalTokens: 1} + + summary := calculateTextQuotaSummary(ctx, relayInfo, usage) + require.Equal(t, 180000, summary.Quota) + + // An adaptor-reported actual count replaces the requested count rather + // than multiplying it a second time. + relayInfo.PriceData.AddOtherRatio("n", 2) + summary = calculateTextQuotaSummary(ctx, relayInfo, usage) + require.Equal(t, 120000, summary.Quota) +} diff --git a/types/request_meta.go b/types/request_meta.go index 476ea052..10f609b4 100644 --- a/types/request_meta.go +++ b/types/request_meta.go @@ -26,7 +26,8 @@ type TokenCountMeta struct { Files []*FileMeta `json:"files,omitempty"` // List of files, each with type and content MaxTokens int `json:"max_tokens,omitempty"` // Maximum tokens allowed in the request - ImagePriceRatio float64 `json:"image_ratio,omitempty"` // Ratio for image size, if applicable + ImagePriceRatio float64 `json:"image_ratio,omitempty"` // Ratio for image size, if applicable + BillingRatios map[string]float64 `json:"billing_ratios,omitempty"` // Validated request multipliers used by pre-consume billing //IsStreaming bool `json:"is_streaming,omitempty"` // Indicates if the request is streaming } From 269e4ff390594ef532a6587c48c6966fa617ce8e Mon Sep 17 00:00:00 2001 From: CaIon Date: Sat, 11 Jul 2026 13:13:56 +0800 Subject: [PATCH 3/3] feat(image): enhance image stream handling with client disconnect logic and billing adjustments --- relay/channel/openai/image_stream_test.go | 280 +++++++++++++++++++++- relay/channel/openai/relay_image.go | 145 +++++++---- 2 files changed, 377 insertions(+), 48 deletions(-) diff --git a/relay/channel/openai/image_stream_test.go b/relay/channel/openai/image_stream_test.go index a9b1e0b2..1adde5c9 100644 --- a/relay/channel/openai/image_stream_test.go +++ b/relay/channel/openai/image_stream_test.go @@ -1,14 +1,17 @@ package openai import ( + "context" "io" "net/http" "net/http/httptest" "strings" + "sync" "testing" "github.com/QuantumNous/new-api/constant" relaycommon "github.com/QuantumNous/new-api/relay/common" + relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -32,6 +35,37 @@ func newImageTestContext(t *testing.T, body, contentType string, isStream bool) return c, recorder, resp, info } +func TestOpenaiImageDoResponseUsesInfoIsStream(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + body := `{"created":1710000000,"data":[{"b64_json":"image"}]}` + + t.Run("non-stream response stays JSON", func(t *testing.T) { + c, recorder, resp, info := newImageTestContext(t, body, "application/json", false) + info.RelayMode = relayconstant.RelayModeImagesGenerations + + usage, err := (&Adaptor{}).DoResponse(c, resp, info) + + require.Nil(t, err) + require.NotNil(t, usage) + require.Equal(t, body, recorder.Body.String()) + }) + + t.Run("stream response converts JSON to SSE", func(t *testing.T) { + c, recorder, resp, info := newImageTestContext(t, body, "application/json", true) + info.RelayMode = relayconstant.RelayModeImagesGenerations + + usage, err := (&Adaptor{}).DoResponse(c, resp, info) + + require.Nil(t, err) + require.NotNil(t, usage) + require.Contains(t, recorder.Body.String(), `event: image_generation.completed`) + require.Contains(t, recorder.Body.String(), `data: [DONE]`) + }) +} + // TestOpenaiImageStreamHandlerForwardsSSEAndUsage covers the core SSE path: // chunks are forwarded with rebuilt event lines, usage is extracted and // normalized (input_tokens -> prompt_tokens with details), and [DONE] is @@ -56,6 +90,8 @@ func TestOpenaiImageStreamHandlerForwardsSSEAndUsage(t *testing.T) { }, "\n") c, recorder, resp, info := newImageTestContext(t, body, "text/event-stream", true) + info.PriceData.UsePrice = true + info.PriceData.AddOtherRatio("n", 3) usage, err := OpenaiImageStreamHandler(c, info, resp) require.Nil(t, err) @@ -69,6 +105,181 @@ func TestOpenaiImageStreamHandlerForwardsSSEAndUsage(t *testing.T) { require.Contains(t, recorder.Body.String(), `data: {"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7,"input_tokens_details":{"image_tokens":2,"text_tokens":1}}}`) require.Contains(t, recorder.Body.String(), `data: [DONE]`) require.Equal(t, "text/event-stream", recorder.Header().Get("Content-Type")) + require.Equal(t, 3.0, info.PriceData.OtherRatios()["n"], "streams without completed events keep the requested count") +} + +func TestOpenaiImageStreamHandlerUsesCompletedEventCount(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"type":"image_generation.partial_image","partial_image_index":0,"b64_json":"partial"}`, + ``, + `data: {"type":"image_generation.completed","b64_json":"first"}`, + ``, + `data: {"type":"image_edit.completed","b64_json":"second","usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7}}`, + ``, + `data: [DONE]`, + ``, + }, "\n") + + c, _, resp, info := newImageTestContext(t, body, "text/event-stream", true) + info.PriceData.UsePrice = true + info.PriceData.AddOtherRatio("n", 3) + + usage, err := OpenaiImageStreamHandler(c, info, resp) + + require.Nil(t, err) + require.Equal(t, 7, usage.TotalTokens) + require.Equal(t, 2.0, info.PriceData.OtherRatios()["n"]) +} + +// blockingBody serves one SSE chunk, then blocks until Close (the scanner's +// cleanup) and returns EOF — keeping the upstream "open" while the client-side +// disconnect is simulated elsewhere. +type blockingBody struct { + mu sync.Mutex + sent bool + chunk []byte + closed chan struct{} +} + +func (b *blockingBody) Read(p []byte) (int, error) { + b.mu.Lock() + if !b.sent { + b.sent = true + n := copy(p, b.chunk) + b.mu.Unlock() + return n, nil + } + b.mu.Unlock() + <-b.closed + return 0, io.EOF +} + +func (b *blockingBody) Close() error { + b.mu.Lock() + defer b.mu.Unlock() + select { + case <-b.closed: + default: + close(b.closed) + } + return nil +} + +// cancelAfterWriter cancels the request context right after the payload +// containing needle has been written to the client, simulating a client that +// disconnects after receiving that event. Cancelling from the write side (not +// the upstream read side) makes the abort deterministic: the handler has +// already processed and counted the event when the disconnect fires. +type cancelAfterWriter struct { + gin.ResponseWriter + needle string + cancel context.CancelFunc + once sync.Once +} + +func (w *cancelAfterWriter) Write(p []byte) (int, error) { + n, err := w.ResponseWriter.Write(p) + if strings.Contains(string(p), w.needle) { + w.once.Do(w.cancel) + } + return n, err +} + +func (w *cancelAfterWriter) WriteString(s string) (int, error) { + n, err := io.WriteString(w.ResponseWriter, s) + if strings.Contains(s, w.needle) { + w.once.Do(w.cancel) + } + return n, err +} + +func newDisconnectingImageStream(t *testing.T, sseBody, disconnectAfter string) (*gin.Context, *httptest.ResponseRecorder, *http.Response, *relaycommon.RelayInfo) { + t.Helper() + c, recorder, resp, info := newImageTestContext(t, "", "text/event-stream", true) + ctx, cancel := context.WithCancel(c.Request.Context()) + t.Cleanup(cancel) + c.Request = c.Request.WithContext(ctx) + c.Writer = &cancelAfterWriter{ResponseWriter: c.Writer, needle: disconnectAfter, cancel: cancel} + resp.Body = &blockingBody{ + chunk: []byte(sseBody), + closed: make(chan struct{}), + } + return c, recorder, resp, info +} + +// TestOpenaiImageStreamHandlerClientDisconnectKeepsRequestedCount guards the +// billing invariant: completed-event counting must not lower the charge when +// the client aborts the stream. Upstream already generated (and charged for) +// all requested images, so a disconnect after the first completed event keeps +// the requested n instead of dropping it to 1. +func TestOpenaiImageStreamHandlerClientDisconnectKeepsRequestedCount(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := "data: {\"type\":\"image_generation.completed\",\"b64_json\":\"first\"}\n\n" + c, recorder, resp, info := newDisconnectingImageStream(t, body, "first") + info.PriceData.UsePrice = true + info.PriceData.AddOtherRatio("n", 3) + + usage, err := OpenaiImageStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, usage) + require.NotNil(t, info.StreamStatus) + // A client abort surfaces as client_gone (main-loop ctx watch) or + // handler_stop (failed client write); both must be treated as untrusted. + require.Contains(t, + []relaycommon.StreamEndReason{relaycommon.StreamEndReasonClientGone, relaycommon.StreamEndReasonHandlerStop}, + info.StreamStatus.EndReason) + require.Contains(t, recorder.Body.String(), `"b64_json":"first"`) + require.Equal(t, 3.0, info.PriceData.OtherRatios()["n"], "client abort must not reduce the billed image count") +} + +// TestOpenaiImageStreamHandlerClientDisconnectRaisesCount covers the other +// direction of the abort guard: when completed events already exceed the +// recorded n, the higher actual count is billed even though the client aborted. +func TestOpenaiImageStreamHandlerClientDisconnectRaisesCount(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + + oldTimeout := constant.StreamingTimeout + constant.StreamingTimeout = 30 + t.Cleanup(func() { constant.StreamingTimeout = oldTimeout }) + + body := strings.Join([]string{ + `data: {"type":"image_generation.completed","b64_json":"first"}`, + ``, + `data: {"type":"image_generation.completed","b64_json":"second"}`, + ``, + ``, + }, "\n") + c, _, resp, info := newDisconnectingImageStream(t, body, "second") + info.PriceData.UsePrice = true + info.PriceData.AddOtherRatio("n", 1) + + usage, err := OpenaiImageStreamHandler(c, info, resp) + + require.Nil(t, err) + require.NotNil(t, usage) + require.NotNil(t, info.StreamStatus) + require.Contains(t, + []relaycommon.StreamEndReason{relaycommon.StreamEndReasonClientGone, relaycommon.StreamEndReasonHandlerStop}, + info.StreamStatus.EndReason) + require.Equal(t, 2.0, info.PriceData.OtherRatios()["n"], "completed events beyond the recorded n must raise the charge even on abort") } // TestOpenaiImageStreamHandlerWrapsJSONResponse covers the non-SSE fallback: @@ -78,9 +289,11 @@ func TestOpenaiImageStreamHandlerWrapsJSONResponse(t *testing.T) { gin.SetMode(gin.TestMode) t.Cleanup(func() { gin.SetMode(oldMode) }) - body := `{"created":1710000000,"data":[{"b64_json":"final","revised_prompt":"draw a cat"}],"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7,"input_tokens_details":{"image_tokens":2,"text_tokens":1}}}` + body := `{"created":1710000000,"data":[{"b64_json":"first","revised_prompt":"draw a cat"},{"b64_json":"second"}],"usage":{"input_tokens":3,"output_tokens":4,"total_tokens":7,"input_tokens_details":{"image_tokens":2,"text_tokens":1}}}` c, recorder, resp, info := newImageTestContext(t, body, "application/json", true) + info.PriceData.UsePrice = true + info.PriceData.AddOtherRatio("n", 3) usage, err := OpenaiImageStreamHandler(c, info, resp) require.Nil(t, err) @@ -93,9 +306,59 @@ func TestOpenaiImageStreamHandlerWrapsJSONResponse(t *testing.T) { require.Empty(t, recorder.Header().Get("Content-Length")) require.Contains(t, recorder.Body.String(), `event: image_generation.completed`) require.Contains(t, recorder.Body.String(), `"type":"image_generation.completed"`) - require.Contains(t, recorder.Body.String(), `"b64_json":"final"`) + require.Contains(t, recorder.Body.String(), `"b64_json":"first"`) + require.Contains(t, recorder.Body.String(), `"b64_json":"second"`) require.Contains(t, recorder.Body.String(), `"revised_prompt":"draw a cat"`) require.Contains(t, recorder.Body.String(), `data: [DONE]`) + require.Equal(t, 2, strings.Count(recorder.Body.String(), `event: image_generation.completed`)) + require.Equal(t, 2.0, info.PriceData.OtherRatios()["n"]) +} + +func TestOpenaiImageHandlerUsesPositiveActualCountForFixedPrice(t *testing.T) { + oldMode := gin.Mode() + gin.SetMode(gin.TestMode) + t.Cleanup(func() { gin.SetMode(oldMode) }) + longImage := strings.Repeat("a", 4096) + + tests := []struct { + name string + body string + usePrice bool + wantCount float64 + }{ + { + name: "fixed price uses data length", + body: `{"data":[{"b64_json":"` + longImage + `"},{"b64_json":"second"}]}`, + usePrice: true, + wantCount: 2, + }, + { + name: "empty data keeps requested count", + body: `{"data":[]}`, + usePrice: true, + wantCount: 3, + }, + { + name: "ratio billing ignores data length", + body: `{"data":[{"b64_json":"first"},{"b64_json":"second"}]}`, + usePrice: false, + wantCount: 3, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c, recorder, resp, info := newImageTestContext(t, tt.body, "application/json", false) + info.PriceData.UsePrice = tt.usePrice + info.PriceData.AddOtherRatio("n", 3) + + _, err := OpenaiImageHandler(c, info, resp) + + require.Nil(t, err) + require.Equal(t, tt.wantCount, info.PriceData.OtherRatios()["n"]) + require.Equal(t, tt.body, recorder.Body.String()) + }) + } } // TestOpenaiImageHandlersReturnJSONError covers JSON error responses for both @@ -132,6 +395,19 @@ func TestOpenaiImageHandlersReturnJSONError(t *testing.T) { require.Equal(t, "content moderation failed", err.ToOpenAIError().Message) require.Empty(t, recorder.Body.String()) }) + + t.Run("stream handler non-2xx stays JSON error", func(t *testing.T) { + c, recorder, resp, info := newImageTestContext(t, body, "application/json", true) + resp.StatusCode = http.StatusBadGateway + + usage, err := OpenaiImageStreamHandler(c, info, resp) + require.Nil(t, usage) + require.NotNil(t, err) + require.Equal(t, http.StatusBadGateway, err.StatusCode) + require.Equal(t, "content moderation failed", err.ToOpenAIError().Message) + require.Empty(t, recorder.Body.String()) + require.NotContains(t, recorder.Header().Get("Content-Type"), "text/event-stream") + }) } // TestOpenaiImageStreamHandlerRecordsUpstreamErrorEvent verifies that an error diff --git a/relay/channel/openai/relay_image.go b/relay/channel/openai/relay_image.go index 1eccbe2c..90fb8794 100644 --- a/relay/channel/openai/relay_image.go +++ b/relay/channel/openai/relay_image.go @@ -5,6 +5,7 @@ import ( "fmt" "io" "net/http" + "strconv" "strings" "time" @@ -17,8 +18,17 @@ import ( "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" ) +func updateOpenAIImageCount(info *relaycommon.RelayInfo, count int64) { + if info == nil || !info.PriceData.UsePrice || count <= 0 || count > int64(dto.MaxImageN) { + return + } + info.PriceData.AddOtherRatio("n", float64(count)) +} + // OpenaiImageHandler handles non-streaming OpenAI image responses // (generations/edits), returning the parsed usage for billing. func OpenaiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { @@ -39,6 +49,8 @@ func OpenaiImageHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http. return nil, types.WithOpenAIError(*oaiError, resp.StatusCode) } + updateOpenAIImageCount(info, gjson.GetBytes(responseBody, "data.#").Int()) + // 写入新的 response body service.IOCopyBytesGracefully(c, resp, responseBody) @@ -88,7 +100,7 @@ func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp return OpenaiImageHandler(c, info, resp) } if !strings.Contains(contentType, "text/event-stream") { - return OpenaiImageJSONAsStreamHandler(c, info, resp) + return openaiImageJSONAsStreamHandler(c, info, resp) } // Reuse the shared streaming engine (helper.StreamScannerHandler) so the // image streaming path gets the same ping keepalive, streaming-timeout @@ -98,6 +110,7 @@ func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp // field (real OpenAI image events keep event == type). usage := &dto.Usage{} var lastStreamData []byte + var completedImages int64 helper.StreamScannerHandler(c, resp, info, func(data string, sr *helper.StreamResult) { raw := common.StringToByteSlice(data) @@ -107,39 +120,64 @@ func OpenaiImageStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp // EndReason. HasErrors() flags the failure for logging/handling. sr.Error(fmt.Errorf("%s", extractOpenAIImageStreamErrorMessage(raw))) } - var usageResp dto.SimpleResponse - if err := common.Unmarshal(raw, &usageResp); err == nil { - normalizeOpenAIUsage(&usageResp.Usage) - if service.ValidUsage(&usageResp.Usage) { - usage = &usageResp.Usage + var chunk struct { + Type string `json:"type"` + Usage dto.Usage `json:"usage"` + } + if err := common.Unmarshal(raw, &chunk); err == nil { + normalizeOpenAIUsage(&chunk.Usage) + if service.ValidUsage(&chunk.Usage) { + usage = &chunk.Usage + } + if chunk.Type == "image_generation.completed" || chunk.Type == "image_edit.completed" { + completedImages++ } } - writeOpenaiImageStreamChunk(c, raw) + if err := writeOpenaiImageStreamChunk(c, raw); err != nil { + sr.Stop(err) + } }) // StreamScannerHandler consumes the upstream [DONE]; re-emit it so the // client still receives a terminal data: [DONE]. - if info != nil && info.StreamStatus != nil && info.StreamStatus.EndReason == relaycommon.StreamEndReasonDone { + if info.StreamStatus != nil && info.StreamStatus.EndReason == relaycommon.StreamEndReasonDone { helper.Done(c) } applyUsagePostProcessing(info, usage, lastStreamData) + // Only trust completedImages when upstream finished the stream (done/eof). + // On client-side aborts (client_gone, or handler_stop from a failed client + // write) the counter undercounts what upstream actually generated and + // charged, so keep the requested n — otherwise a client could pay for one + // image by disconnecting right after the first completed event. The abort + // guard only blocks lowering the charge: if completed events already + // exceed the recorded n, bill the higher actual count regardless. + if info.StreamStatus != nil { + upstreamFinished := info.StreamStatus.EndReason == relaycommon.StreamEndReasonDone || + info.StreamStatus.EndReason == relaycommon.StreamEndReasonEOF + requestedN := 1.0 + if n, ok := info.PriceData.OtherRatios()["n"]; ok { + requestedN = n + } + if upstreamFinished || float64(completedImages) > requestedN { + updateOpenAIImageCount(info, completedImages) + } + } return usage, nil } // writeOpenaiImageStreamChunk rebuilds the SSE frame for an image stream chunk: // it emits an "event:" line derived from the JSON "type" field (when present) // followed by the verbatim "data:" payload, mirroring helper.ResponseChunkData. -func writeOpenaiImageStreamChunk(c *gin.Context, data []byte) { +func writeOpenaiImageStreamChunk(c *gin.Context, data []byte) error { var payload struct { Type string `json:"type"` } _ = common.Unmarshal(data, &payload) if eventName := strings.TrimSpace(payload.Type); eventName != "" { - _ = helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventName}, string(data)) - return + return helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventName}, string(data)) } - _ = helper.StringData(c, string(data)) + return helper.StringData(c, string(data)) } // isOpenAIImageStreamErrorEvent detects upstream error chunks by JSON content @@ -192,7 +230,7 @@ func extractOpenAIImageStreamErrorMessage(data []byte) string { return "upstream image stream returned error event" } -func OpenaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { +func openaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) { defer service.CloseResponseBodyGracefully(resp) responseBody, err := io.ReadAll(resp.Body) @@ -200,49 +238,75 @@ func OpenaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, return nil, types.NewOpenAIError(err, types.ErrorCodeReadResponseBodyFailed, http.StatusInternalServerError) } - var imageResp dto.ImageResponse - if err := common.Unmarshal(responseBody, &imageResp); err != nil { + // Only decode usage/error. Do not Unmarshal data[] into dto.ImageResponse — + // b64_json values are large and would be copied into Go strings then + // re-marshaled for each SSE event. + var usageResp dto.SimpleResponse + if err := common.Unmarshal(responseBody, &usageResp); err != nil { return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } - - var usageResp dto.SimpleResponse - _ = common.Unmarshal(responseBody, &usageResp) if oaiError := usageResp.GetOpenAIError(); oaiError != nil && oaiError.Type != "" { return nil, types.WithOpenAIError(*oaiError, resp.StatusCode) } normalizeOpenAIUsage(&usageResp.Usage) applyUsagePostProcessing(info, &usageResp.Usage, responseBody) + imageCount := gjson.GetBytes(responseBody, "data.#").Int() + updateOpenAIImageCount(info, imageCount) + helper.SetEventStreamHeaders(c) c.Status(http.StatusOK) - created := imageResp.Created + created := gjson.GetBytes(responseBody, "created").Int() if created == 0 { created = time.Now().Unix() } if info != nil { info.SetFirstResponseTime() } - for _, image := range imageResp.Data { - payload := map[string]any{ - "type": "image_generation.completed", - "created_at": created, + + validUsage := service.ValidUsage(&usageResp.Usage) + var usageJSON []byte + if validUsage { + usageJSON, err = common.Marshal(usageResp.Usage) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } - if image.Url != "" { - payload["url"] = image.Url + } + + for i := int64(0); i < imageCount; i++ { + image := gjson.GetBytes(responseBody, "data."+strconv.FormatInt(i, 10)) + payload := []byte(`{"type":"image_generation.completed"}`) + payload, err = sjson.SetBytes(payload, "created_at", created) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } - if image.B64Json != "" { - payload["b64_json"] = image.B64Json + if validUsage { + payload, err = sjson.SetRawBytes(payload, "usage", usageJSON) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + } } - if image.RevisedPrompt != "" { - payload["revised_prompt"] = image.RevisedPrompt + // b64_json goes last: every sjson.Set* reallocates the whole payload, + // so inserting the large blob after all small fields avoids re-copying + // multi-MB buffers. + for _, field := range []string{"url", "revised_prompt", "b64_json"} { + value := image.Get(field) + if value.Type != gjson.String || value.Raw == `""` { + continue + } + raw := []byte(value.Raw) + if value.Index > 0 { + raw = responseBody[value.Index : value.Index+len(value.Raw)] + } + payload, err = sjson.SetRawBytes(payload, field, raw) + if err != nil { + return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + } } - if service.ValidUsage(&usageResp.Usage) { - payload["usage"] = usageResp.Usage - } - if err := writeOpenaiImageStreamPayload(c, "image_generation.completed", payload); err != nil { + if writeErr := helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: "image_generation.completed"}, string(payload)); writeErr != nil { if info != nil && info.StreamStatus != nil { - info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, err) + info.StreamStatus.SetEndReason(relaycommon.StreamEndReasonClientGone, writeErr) } return &usageResp.Usage, nil } @@ -254,7 +318,7 @@ func OpenaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, return &usageResp.Usage, nil } if info != nil { - info.ReceivedResponseCount += len(imageResp.Data) + info.ReceivedResponseCount += int(imageCount) if info.StreamStatus == nil { info.StreamStatus = relaycommon.NewStreamStatus() } @@ -263,17 +327,6 @@ func OpenaiImageJSONAsStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, return &usageResp.Usage, nil } -func writeOpenaiImageStreamPayload(c *gin.Context, eventName string, payload any) error { - data, err := common.Marshal(payload) - if err != nil { - return err - } - if eventName != "" { - return helper.ResponseChunkData(c, dto.ResponsesStreamResponse{Type: eventName}, string(data)) - } - return helper.StringData(c, string(data)) -} - func writeOpenaiImageStreamDone(c *gin.Context) error { return helper.StringData(c, "[DONE]") }