Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions aigateway/handler/chat_trace_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
4 changes: 4 additions & 0 deletions aigateway/handler/model_price_guard_ce.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
6 changes: 6 additions & 0 deletions aigateway/handler/openai_image.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down
6 changes: 6 additions & 0 deletions aigateway/handler/openai_image_edit.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down
6 changes: 6 additions & 0 deletions aigateway/handler/openai_video.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{
Expand Down
40 changes: 34 additions & 6 deletions common/utils/common/resolution.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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) {
Expand Down Expand Up @@ -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
}
36 changes: 36 additions & 0 deletions common/utils/common/resolution_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
})
}
}
Loading