add modal price guard for multimodal apis in aigateway (#1254)

* add modal price guard for multimodal apis in aigateway

* Refactor resolution parsing functions

Removed redundant parseResolutionWxHFormat function and simplified logic in parseResolutionWxHPair.

---------

Co-authored-by: Jun Luo <luojun96@live.cn>
Co-authored-by: HaiHui <154984842+HaiHui886@users.noreply.github.com>

Author: 3 people

aigateway/handler/chat_trace_test.go

@@ -118,6 +118,38 @@ func TestPreflightTraceRecordsEarlyError(t *testing.T) {
 	requireSpanAttrValue(t, spans[0].Attributes, "error.category", "gateway_error")
 }
 
+func TestPreflightTraceRecordsModalPriceGuardError(t *testing.T) {
+	exporter := tracetest.NewInMemoryExporter()
+	provider := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter))
+	oldProvider := otel.GetTracerProvider()
+	otel.SetTracerProvider(provider)
+	t.Cleanup(func() {
+		otel.SetTracerProvider(oldProvider)
+		_ = provider.Shutdown(context.Background())
+	})
+
+	_, preflight := startPreflightTrace(context.Background(), preflightTraceStart{
+		API:       "/v1/images/generations",
+		RequestID: "req-1",
+		UserID:    "user-1",
+	})
+	preflight.SetTargetModel("model1", &resolvedModelTarget{
+		Model:     &types.Model{BaseModel: types.BaseModel{ID: "internal-model"}},
+		ModelName: "resolved-model",
+	})
+	preflight.RecordError(newInvalidRequestModelTargetError(
+		"invalid_request_error",
+		"resolution exceeds the maximum supported resolution",
+		modelTargetErrorOptions{},
+	), "modal_price_guard")
+
+	spans := exporter.GetSpans()
+	require.Len(t, spans, 1)
+	require.Equal(t, "Error", spans[0].Status.Code.String())
+	requireSpanAttrValue(t, spans[0].Attributes, "error.type", "modal_price_guard")
+	requireSpanAttrValue(t, spans[0].Attributes, "aigateway.request.model", "model1")
+}
+
 func TestPreflightTraceRecordsResolvedTargetModel(t *testing.T) {
 	exporter := tracetest.NewInMemoryExporter()
 	provider := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter))

aigateway/handler/model_price_guard_ce.go

@@ -7,3 +7,7 @@ import "opencsg.com/csghub-server/aigateway/types"
 func modelSKUPriceStatus(model *types.Model) (requiresSKUPrice bool, hasConfiguredSKUPrice bool) {
 	return false, false
 }
+
+func checkModalRequestAllowed(model *types.Model, size string) *modelTargetError {
+	return nil
+}

aigateway/handler/openai_image.go

@@ -70,6 +70,12 @@ func (h *OpenAIHandlerImpl) GenerateImage(c *gin.Context) {
 		return
 	}
 	preflight.SetTargetModel(modelID, modelTarget)
+
+	if guardErr := checkModalRequestAllowed(modelTarget.Model, string(req.Size)); guardErr != nil {
+		preflight.RecordError(guardErr, "modal_price_guard")
+		handleModelTargetError(c, ctx, modelID, "modal price guard rejected request", guardErr)
+		return
+	}
 	preflight.End()
 
 	traceCtx, generationRecorder := h.startModalGenerationTrace(ctx, modalTraceStartInput{

aigateway/handler/openai_image_edit.go

@@ -84,6 +84,12 @@ func (h *OpenAIHandlerImpl) EditImage(c *gin.Context) {
 		return
 	}
 	preflight.SetTargetModel(modelID, modelTarget)
+
+	if guardErr := checkModalRequestAllowed(modelTarget.Model, c.PostForm("size")); guardErr != nil {
+		preflight.RecordError(guardErr, "modal_price_guard")
+		handleModelTargetError(c, ctx, modelID, "modal price guard rejected request", guardErr)
+		return
+	}
 	preflight.End()
 
 	traceCtx, generationRecorder := h.startModalGenerationTrace(ctx, modalTraceStartInput{

aigateway/handler/openai_video.go

@@ -73,6 +73,12 @@ func (h *OpenAIHandlerImpl) CreateVideo(c *gin.Context) {
 		return
 	}
 	preflight.SetTargetModel(input.modelID, modelTarget)
+
+	if guardErr := checkModalRequestAllowed(modelTarget.Model, input.adapterReq.Size); guardErr != nil {
+		preflight.RecordError(guardErr, "modal_price_guard")
+		handleModelTargetError(c, ctx, input.modelID, "modal price guard rejected request", guardErr)
+		return
+	}
 	preflight.End()
 
 	traceCtx, generationRecorder := h.startModalGenerationTrace(ctx, modalTraceStartInput{

common/utils/common/resolution.go

@@ -7,12 +7,25 @@ import (
 	"strings"
 )
 
+// This file parses resolution strings in preset ("720p"), "WIDTHxHEIGHT",
+// or bare-integer format. ExtractEventResolutionMaxSide is lenient (used for
+// historical/billing data), while ParseResolutionMaxSide additionally
+// requires a strictly positive value (used for real-time request/SKU
+// validation), so gateway and accounting never interpret the same string
+// differently.
 var (
 	resolutionPFormat    = regexp.MustCompile(`^(\d+)\s*[pP]$`)
 	resolutionWxHFormat  = regexp.MustCompile(`^(\d+)\s*[xX×]\s*(\d+)$`)
 	resolutionPureFormat = regexp.MustCompile(`^(\d+)$`)
 )
 
+func longEdge(width, height int64) int64 {
+	if width > height {
+		return width
+	}
+	return height
+}
+
 func parseResolutionPFormat(s string) (int64, bool) {
 	matches := resolutionPFormat.FindStringSubmatch(s)
 	if matches == nil {
@@ -25,20 +38,25 @@ func parseResolutionPFormat(s string) (int64, bool) {
 	return v, true
 }
 
-func parseResolutionWxHFormat(s string) (int64, bool) {
+func parseResolutionWxHPair(s string) (width, height int64, ok bool) {
 	matches := resolutionWxHFormat.FindStringSubmatch(s)
 	if matches == nil {
-		return 0, false
+		return 0, 0, false
 	}
 	w, err1 := strconv.ParseInt(matches[1], 10, 64)
 	h, err2 := strconv.ParseInt(matches[2], 10, 64)
 	if err1 != nil || err2 != nil {
-		return 0, false
+		return 0, 0, false
 	}
-	if w > h {
-		return w, true
+	return w, h, true
+}
+
+func parseResolutionWxHFormat(s string) (int64, bool) {
+	w, h, ok := parseResolutionWxHPair(s)
+	if !ok {
+		return 0, false
 	}
-	return h, true
+	return longEdge(w, h), true
 }
 
 func parseResolutionPureFormat(s string) (int64, bool) {
@@ -71,3 +89,13 @@ func ExtractEventResolutionMaxSide(eventResolution string) (int64, error) {
 
 	return 0, fmt.Errorf("unsupported resolution format: %s, must be 1080P or 720 or 1920x1080", s)
 }
+
+// ParseResolutionMaxSide is like ExtractEventResolutionMaxSide but rejects
+// empty input and non-positive results.
+func ParseResolutionMaxSide(size string) (int64, bool) {
+	v, err := ExtractEventResolutionMaxSide(size)
+	if err != nil || v <= 0 {
+		return 0, false
+	}
+	return v, true
+}

common/utils/common/resolution_test.go

@@ -116,3 +116,39 @@ func TestExtractEventResolutionMaxSide(t *testing.T) {
 		}
 	}
 }
+
+func TestParseResolutionMaxSide(t *testing.T) {
+	tests := []struct {
+		name     string
+		size     string
+		wantSide int64
+		wantOK   bool
+	}{
+		{"square", "1024x1024", 1024, true},
+		{"landscape", "1280x720", 1280, true},
+		{"portrait", "720x1280", 1280, true},
+		{"whitespace wxh", "  1024  x  1536  ", 1536, true},
+		{"uppercase wxh", "1024X1024", 1024, true},
+		{"preset uppercase", "1080P", 1080, true},
+		{"preset lowercase", "720p", 720, true},
+		{"whitespace preset", "  1080p  ", 1080, true},
+		{"pure integer", "1024", 1024, true},
+		{"whitespace integer", "  2048  ", 2048, true},
+		{"empty", "", 0, false},
+		{"degenerate wxh one side zero", "0x1024", 1024, true}, // matches ExtractEventResolutionMaxSide leniency for a single zero dimension
+		{"wxh both sides zero", "0x0", 0, false},
+		{"negative wxh", "-1024x1024", 0, false},
+		{"zero preset", "0p", 0, false},
+		{"zero integer", "0", 0, false},
+		{"negative integer", "-1024", 0, false},
+		{"non numeric", "abcxdef", 0, false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			side, ok := ParseResolutionMaxSide(tt.size)
+			require.Equal(t, tt.wantOK, ok)
+			require.Equal(t, tt.wantSide, side)
+		})
+	}
+}