From 844c5c19508ad08dad86caead838bfb72f92348a Mon Sep 17 00:00:00 2001 From: 3metaJun <251347867+3metaJun@users.noreply.github.com> Date: Sun, 30 Aug 2026 13:34:26 +0800 Subject: [PATCH] fix: normalize image billing inputs --- dto/openai_image.go | 51 +++++++++++++++++++ dto/openai_image_test.go | 23 +++++++++ relay/helper/billing_expr_request.go | 41 +++++++++++++++ relay/helper/billing_expr_request_test.go | 61 +++++++++++++++++++++++ 4 files changed, 176 insertions(+) diff --git a/dto/openai_image.go b/dto/openai_image.go index bd6f514cd32e..da07c098c9f1 100644 --- a/dto/openai_image.go +++ b/dto/openai_image.go @@ -3,6 +3,7 @@ package dto import ( "encoding/json" "reflect" + "strconv" "strings" "github.com/QuantumNous/new-api/common" @@ -173,6 +174,56 @@ func (i *ImageRequest) GetTokenCountMeta() *types.TokenCountMeta { } } +// BillingResolution returns the same 1K/2K/4K bucket used by sub2api image +// settlement. Unknown and automatic sizes default to 2K. +func (i *ImageRequest) BillingResolution() string { + switch strings.ToUpper(strings.TrimSpace(i.Resolution)) { + case "1K": + return "1K" + case "2K": + return "2K" + case "4K": + return "4K" + } + + parts := strings.Split(strings.ToLower(strings.TrimSpace(i.Size)), "x") + if len(parts) != 2 { + return "2K" + } + width, widthOK := imageBillingDimension(strings.TrimSpace(parts[0])) + height, heightOK := imageBillingDimension(strings.TrimSpace(parts[1])) + if !widthOK || !heightOK { + return "2K" + } + maxEdge := max(width, height) + switch { + case maxEdge <= 1024: + return "1K" + case maxEdge <= 2048: + return "2K" + default: + return "4K" + } +} + +func imageBillingDimension(value string) (int, bool) { + value = strings.TrimPrefix(value, "+") + value = strings.TrimLeft(value, "0") + if value == "" { + return 0, false + } + for _, char := range value { + if char < '0' || char > '9' { + return 0, false + } + } + if len(value) > 4 { + return 2049, true + } + dimension, err := strconv.Atoi(value) + return dimension, err == nil && dimension > 0 +} + func (i *ImageRequest) IsStream(c *gin.Context) bool { return i.Stream != nil && *i.Stream } diff --git a/dto/openai_image_test.go b/dto/openai_image_test.go index c02bf1f8cfcc..25ead5c20293 100644 --- a/dto/openai_image_test.go +++ b/dto/openai_image_test.go @@ -18,3 +18,26 @@ func TestImageRequestPreservesProviderResolutionFields(t *testing.T) { require.NoError(t, err) require.JSONEq(t, string(raw), string(encoded)) } + +func TestImageRequestBillingResolution(t *testing.T) { + tests := []struct { + name string + resolution string + size string + want string + }{ + {name: "explicit", resolution: "4k", size: "1024x1024", want: "4K"}, + {name: "one kilopixel", size: "1024x768", want: "1K"}, + {name: "two kilopixel", size: "2048x1152", want: "2K"}, + {name: "four kilopixel", size: "3840x2160", want: "4K"}, + {name: "overflowing dimension", size: "999999999999999999999999x1024", want: "4K"}, + {name: "automatic", size: "auto", want: "2K"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + request := ImageRequest{Resolution: test.resolution, Size: test.size} + require.Equal(t, test.want, request.BillingResolution()) + }) + } +} diff --git a/relay/helper/billing_expr_request.go b/relay/helper/billing_expr_request.go index 28a44bc8ae5a..c7e30f3849a7 100644 --- a/relay/helper/billing_expr_request.go +++ b/relay/helper/billing_expr_request.go @@ -1,6 +1,8 @@ package helper import ( + "encoding/json" + "errors" "strings" "github.com/QuantumNous/new-api/common" @@ -30,6 +32,14 @@ func ResolveIncomingBillingExprRequestInput(c *gin.Context, info *relaycommon.Re if err != nil { return billingexpr.RequestInput{}, err } + if info != nil { + if imageRequest, ok := info.Request.(*dto.ImageRequest); ok { + bodyBytes, err = mergeValidatedImageBillingFields(bodyBytes, imageRequest) + if err != nil { + return billingexpr.RequestInput{}, err + } + } + } input.Body = bodyBytes return input, nil } @@ -46,10 +56,41 @@ func BuildBillingExprRequestInputFromRequest(request dto.Request, headers map[st if err != nil { return billingexpr.RequestInput{}, err } + if imageRequest, ok := request.(*dto.ImageRequest); ok { + bodyBytes, err = mergeValidatedImageBillingFields(bodyBytes, imageRequest) + if err != nil { + return billingexpr.RequestInput{}, err + } + } input.Body = bodyBytes return input, nil } +func mergeValidatedImageBillingFields(body []byte, request *dto.ImageRequest) ([]byte, error) { + fields := make(map[string]json.RawMessage) + if len(body) > 0 { + if err := common.Unmarshal(body, &fields); err != nil { + return nil, err + } + if fields == nil { + return nil, errors.New("image billing request body must be a JSON object") + } + } + for key, value := range request.Extra { + if _, exists := fields[key]; !exists { + fields[key] = value + } + } + + imageN := uint(1) + if request.N != nil && *request.N > 0 { + imageN = *request.N + } + fields["n"], _ = common.Marshal(imageN) + fields["resolution"], _ = common.Marshal(request.BillingResolution()) + return common.Marshal(fields) +} + func readIncomingBillingExprBody(c *gin.Context) ([]byte, error) { if c == nil || c.Request == nil || !isJSONContentType(c.Request.Header.Get("Content-Type")) { return nil, nil diff --git a/relay/helper/billing_expr_request_test.go b/relay/helper/billing_expr_request_test.go index 9193f4b452e3..7abc2f48026c 100644 --- a/relay/helper/billing_expr_request_test.go +++ b/relay/helper/billing_expr_request_test.go @@ -2,6 +2,7 @@ package helper import ( "bytes" + "encoding/json" "io" "net/http" "net/http/httptest" @@ -37,6 +38,49 @@ func TestResolveIncomingBillingExprRequestInput(t *testing.T) { require.Equal(t, "application/json", input.Headers["Content-Type"]) } +func TestResolveIncomingBillingExprRequestInputUsesValidatedImageFields(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + ctx.Request.Header.Set("Content-Type", "application/json") + + body := []byte(`{"model":"gpt-image-2","prompt":"test","n":0,"size":"1024x1024","custom":"kept"}`) + ctx.Request.Body = io.NopCloser(bytes.NewReader(body)) + ctx.Set(common.KeyRequestBody, body) + n := uint(1) + info := &relaycommon.RelayInfo{ + Request: &dto.ImageRequest{ + Model: "gpt-image-2", + Prompt: "test", + N: &n, + Size: "1024x1024", + }, + } + + input, err := ResolveIncomingBillingExprRequestInput(ctx, info) + require.NoError(t, err) + require.Equal(t, float64(1), gjson.GetBytes(input.Body, "n").Float()) + require.Equal(t, "1K", gjson.GetBytes(input.Body, "resolution").String()) + require.Equal(t, "kept", gjson.GetBytes(input.Body, "custom").String()) +} + +func TestResolveIncomingBillingExprRequestInputRejectsNullImageBody(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil) + ctx.Request.Header.Set("Content-Type", "application/json") + + body := []byte(`null`) + ctx.Request.Body = io.NopCloser(bytes.NewReader(body)) + ctx.Set(common.KeyRequestBody, body) + info := &relaycommon.RelayInfo{Request: &dto.ImageRequest{Model: "gpt-image-2", Prompt: "test"}} + + _, err := ResolveIncomingBillingExprRequestInput(ctx, info) + require.EqualError(t, err, "image billing request body must be a JSON object") +} + func TestBuildBillingExprRequestInputFromRequest(t *testing.T) { request := &dto.GeneralOpenAIRequest{ Model: "gemini-3.1-pro-preview", @@ -61,3 +105,20 @@ func TestBuildBillingExprRequestInputFromRequest(t *testing.T) { require.Equal(t, "user", gjson.GetBytes(input.Body, "messages.0.role").String()) require.Equal(t, float64(3000), gjson.GetBytes(input.Body, "max_tokens").Float()) } + +func TestBuildBillingExprRequestInputFromImageRequestPreservesExtraFields(t *testing.T) { + n := uint(1) + request := &dto.ImageRequest{ + Model: "gpt-image-2", + Prompt: "test", + N: &n, + Size: "1024x1024", + Extra: map[string]json.RawMessage{"custom": json.RawMessage(`"kept"`)}, + } + + input, err := BuildBillingExprRequestInputFromRequest(request, nil) + require.NoError(t, err) + require.Equal(t, "kept", gjson.GetBytes(input.Body, "custom").String()) + require.Equal(t, float64(1), gjson.GetBytes(input.Body, "n").Float()) + require.Equal(t, "1K", gjson.GetBytes(input.Body, "resolution").String()) +}