Jaymr/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(billing): extend quantity validation and saturating conversions to remaining paths

Browse Source 
Bound max-tokens fields across all relay format validators, saturate
tiered-expression rounding and audio/tool/task token conversions, and
route legacy remix ratios through the guarded setter.
This commit is contained in:
CaIon
2026-07-07 01:08:52 +08:00
parent d0bd8aac74
commit c9943d37ad
13 changed files with 155 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
relay/channel/openai/audio.go
+3 -2
View File
@@ -103,8 +103,9 @@ func OpenaiTTSHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
usage.CompletionTokens = estimatedTokens
usage.CompletionTokenDetails.AudioTokens = estimatedTokens
} else if duration > 0 {
// 计算 token: ceil(duration) / 60.0 * 1000，即每分钟 1000 tokens
completionTokens := int(math.Round(math.Ceil(duration) / 60.0 * 1000))
// 计算 token: ceil(duration) / 60.0 * 1000，即每分钟 1000 tokens。
// duration 解析自上游返回的音频元数据，饱和转换防止 int 回绕。
completionTokens := common.QuotaFromFloat(math.Round(math.Ceil(duration) / 60.0 * 1000))
usage.CompletionTokens = completionTokens
usage.CompletionTokenDetails.AudioTokens = completionTokens
}
relay/channel/task/kling/adaptor.go
+2 -1
View File
@@ -358,7 +358,8 @@ func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, e
taskInfo.Url = video.Url
}
if tokens, err := strconv.ParseFloat(resPayload.Data.FinalUnitDeduction, 64); err == nil {
rounded := int(math.Ceil(tokens))
// 上游返回的扣费数值，饱和转换防止超大数值回绕成负数
rounded := common.QuotaFromFloat(math.Ceil(tokens))
if rounded > 0 {
taskInfo.CompletionTokens = rounded
taskInfo.TotalTokens = rounded
relay/helper/max_tokens_bounds_test.go
+72
View File
@@ -0,0 +1,72 @@
package helper
import (
"bytes"
"net/http"
"net/http/httptest"
"testing"
relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// TestMaxTokensBounds guards the billing invariant that user-supplied max
// token fields are bounded on every relay format. These values feed
// pre-consume quota math (preConsumedTokens * ratio); a huge or
// wrapped-negative value (e.g. 18446744073686646784 parsed into *uint) must
// be rejected at validation instead of corrupting the pre-charge.
func TestMaxTokensBounds(t *testing.T) {
gin.SetMode(gin.TestMode)
newJSONContext := func(t *testing.T, body string) *gin.Context {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, "/relay", bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
return c
}
const hugeN = "18446744073686646784"
t.Run("openai max_tokens overflow rejected", func(t *testing.T) {
c := newJSONContext(t, `{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"max_tokens":`+hugeN+`}`)
_, err := GetAndValidateTextRequest(c, relayconstant.RelayModeChatCompletions)
require.Error(t, err)
require.Contains(t, err.Error(), "max_tokens is invalid")
})
t.Run("openai max_completion_tokens overflow rejected", func(t *testing.T) {
c := newJSONContext(t, `{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}],"max_completion_tokens":`+hugeN+`}`)
_, err := GetAndValidateTextRequest(c, relayconstant.RelayModeChatCompletions)
require.Error(t, err)
require.Contains(t, err.Error(), "max_tokens is invalid")
})
t.Run("claude max_tokens overflow rejected", func(t *testing.T) {
c := newJSONContext(t, `{"model":"claude-sonnet-4","messages":[{"role":"user","content":"hi"}],"max_tokens":`+hugeN+`}`)
_, err := GetAndValidateClaudeRequest(c)
require.Error(t, err)
require.Contains(t, err.Error(), "max_tokens is invalid")
})
t.Run("claude normal max_tokens accepted", func(t *testing.T) {
c := newJSONContext(t, `{"model":"claude-sonnet-4","messages":[{"role":"user","content":"hi"}],"max_tokens":8192}`)
req, err := GetAndValidateClaudeRequest(c)
require.NoError(t, err)
require.EqualValues(t, 8192, *req.MaxTokens)
})
t.Run("gemini maxOutputTokens overflow rejected", func(t *testing.T) {
c := newJSONContext(t, `{"contents":[{"parts":[{"text":"hi"}]}],"generationConfig":{"maxOutputTokens":`+hugeN+`}}`)
_, err := GetAndValidateGeminiRequest(c)
require.Error(t, err)
require.Contains(t, err.Error(), "maxOutputTokens is invalid")
})
t.Run("responses max_output_tokens overflow rejected", func(t *testing.T) {
c := newJSONContext(t, `{"model":"gpt-4o","input":"hi","max_output_tokens":`+hugeN+`}`)
_, err := GetAndValidateResponsesRequest(c)
require.Error(t, err)
require.Contains(t, err.Error(), "max_output_tokens is invalid")
})
}
relay/helper/price.go
+4 -4
View File
@@ -112,12 +112,12 @@ 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 = int(float64(preConsumedTokens) * ratio)
preConsumedQuota = common.QuotaFromFloat(float64(preConsumedTokens) * ratio)
} else {
if meta.ImagePriceRatio != 0 {
modelPrice = modelPrice * meta.ImagePriceRatio
}
preConsumedQuota = int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
preConsumedQuota = common.QuotaFromFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
}
// check if free model pre-consume is disabled
@@ -194,7 +194,7 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types
freeModel := false
if usePrice {
quota = int(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
quota = common.QuotaFromFloat(modelPrice * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
if groupRatioInfo.GroupRatio == 0 || modelPrice == 0 {
quota = 0
@@ -203,7 +203,7 @@ func ModelPriceHelperPerCall(c *gin.Context, info *relaycommon.RelayInfo) (types
}
} else {
// 按量计费：以模型倍率的一半作为预扣额度
quota = int(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
quota = common.QuotaFromFloat(modelRatio / 2 * common.QuotaPerUnit * groupRatioInfo.GroupRatio)
modelPrice = -1
if !operation_setting.GetQuotaSetting().EnableFreeModelPreConsume {
if groupRatioInfo.GroupRatio == 0 || modelRatio == 0 {
relay/helper/valid_request.go
+24 -1
View File
@@ -114,6 +114,20 @@ func GetAndValidateEmbeddingRequest(c *gin.Context, relayMode int) (*dto.Embeddi
return embeddingRequest, nil
}
// maxTokensLimit bounds user-supplied max token fields. These values feed
// pre-consume quota math (preConsumedTokens * ratio); an unbounded value can
// overflow the conversion and corrupt billing.
const maxTokensLimit = math.MaxInt32 / 2
func exceedsMaxTokensLimit(values ...*uint) bool {
for _, v := range values {
if lo.FromPtrOr(v, uint(0)) > maxTokensLimit {
return true
}
}
return false
}
func GetAndValidateResponsesRequest(c *gin.Context) (*dto.OpenAIResponsesRequest, error) {
request := &dto.OpenAIResponsesRequest{}
err := common.UnmarshalBodyReusable(c, request)
@@ -126,6 +140,9 @@ func GetAndValidateResponsesRequest(c *gin.Context) (*dto.OpenAIResponsesRequest
if request.Input == nil {
return nil, errors.New("input is required")
}
if exceedsMaxTokensLimit(request.MaxOutputTokens) {
return nil, errors.New("max_output_tokens is invalid")
}
return request, nil
}
@@ -259,6 +276,9 @@ func GetAndValidateClaudeRequest(c *gin.Context) (textRequest *dto.ClaudeRequest
if textRequest.Model == "" {
return nil, errors.New("field model is required")
}
if exceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxTokensToSample) {
return nil, errors.New("max_tokens is invalid")
}
//if textRequest.Stream {
// relayInfo.IsStream = true
@@ -281,7 +301,7 @@ func GetAndValidateTextRequest(c *gin.Context, relayMode int) (*dto.GeneralOpenA
textRequest.Model = c.Param("model")
}
if lo.FromPtrOr(textRequest.MaxTokens, uint(0)) > math.MaxInt32/2 {
if exceedsMaxTokensLimit(textRequest.MaxTokens, textRequest.MaxCompletionTokens) {
return nil, errors.New("max_tokens is invalid")
}
if textRequest.Model == "" {
@@ -334,6 +354,9 @@ func GetAndValidateGeminiRequest(c *gin.Context) (*dto.GeminiChatRequest, error)
if len(request.Contents) == 0 && len(request.Requests) == 0 {
return nil, errors.New("contents is required")
}
if exceedsMaxTokensLimit(request.GenerationConfig.MaxOutputTokens) {
return nil, errors.New("maxOutputTokens is invalid")
}
//if c.Query("alt") == "sse" {
// relayInfo.IsStream = true
relay/relay_task.go
+7 -6
View File
@@ -121,14 +121,15 @@ func ResolveOriginTask(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskErr
if seconds <= 0 {
seconds = 4
}
sizeStr, _ := taskData["size"].(string)
if info.PriceData.OtherRatios == nil {
info.PriceData.OtherRatios = map[string]float64{}
// 历史任务数据可能包含未经校验的时长，作为计费乘数前必须钳制
if seconds > relaycommon.MaxTaskDurationSeconds {
seconds = relaycommon.MaxTaskDurationSeconds
}
info.PriceData.OtherRatios["seconds"] = float64(seconds)
info.PriceData.OtherRatios["size"] = 1
sizeStr, _ := taskData["size"].(string)
info.PriceData.AddOtherRatio("seconds", float64(seconds))
info.PriceData.AddOtherRatio("size", 1)
if sizeStr == "1792x1024" || sizeStr == "1024x1792" {
info.PriceData.OtherRatios["size"] = 1.666667
info.PriceData.AddOtherRatio("size", 1.666667)
}
}
}