fix: normalize image billing inputs

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
 }

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())
+		})
+	}
+}

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

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())
+}