fix(openai): reject image models on chat completions

backend/internal/handler/chat_completions_image_model_test.go

@@ -0,0 +1,96 @@
+package handler
+
+import (
+	"bytes"
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+	"github.com/tidwall/gjson"
+)
+
+func TestChatCompletionsRejectsGPTImageModelsBeforeScheduling(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	for _, model := range []string{"gpt-image-1", "gpt-image-1.5", "gpt-image-2"} {
+		for _, tc := range []struct {
+			name string
+			call func(*gin.Context)
+		}{
+			{
+				name: "gateway",
+				call: (&GatewayHandler{}).ChatCompletions,
+			},
+			{
+				name: "openai_gateway",
+				call: newOpenAIImageChatRejectionHandler(t).ChatCompletions,
+			},
+		} {
+			t.Run(tc.name+"/"+model, func(t *testing.T) {
+				recorder := httptest.NewRecorder()
+				c, _ := gin.CreateTestContext(recorder)
+				body := []byte(`{"model":"` + model + `","messages":[{"role":"user","content":"draw"}]}`)
+				c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
+				setImageChatTestAuth(c)
+
+				tc.call(c)
+
+				require.Equal(t, http.StatusBadRequest, recorder.Code)
+				require.Equal(t, "invalid_request_error", gjson.Get(recorder.Body.String(), "error.type").String())
+				require.Contains(t, gjson.Get(recorder.Body.String(), "error.message").String(), "Chat Completions")
+				_, selected := c.Get(opsAccountIDKey)
+				require.False(t, selected, "rejection must happen before account selection")
+			})
+		}
+	}
+}
+
+func TestOpenAIChatCompletionsImageModelRejectionDoesNotAcquireConcurrency(t *testing.T) {
+	var acquireCalls atomic.Int64
+	cache := &concurrencyCacheMock{
+		acquireUserSlotFn: func(context.Context, int64, int, string) (bool, error) {
+			acquireCalls.Add(1)
+			return true, nil
+		},
+	}
+	h := newOpenAIImageChatRejectionHandlerWithCache(t, cache)
+	recorder := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(recorder)
+	c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewBufferString(
+		`{"model":"gpt-image-2","messages":[{"role":"user","content":"draw"}]}`,
+	))
+	setImageChatTestAuth(c)
+
+	h.ChatCompletions(c)
+
+	require.Equal(t, http.StatusBadRequest, recorder.Code)
+	require.Zero(t, acquireCalls.Load(), "rejection must happen before user/account concurrency and scheduling")
+}
+
+func newOpenAIImageChatRejectionHandler(t *testing.T) *OpenAIGatewayHandler {
+	t.Helper()
+	return newOpenAIImageChatRejectionHandlerWithCache(t, &concurrencyCacheMock{})
+}
+
+func newOpenAIImageChatRejectionHandlerWithCache(t *testing.T, cache *concurrencyCacheMock) *OpenAIGatewayHandler {
+	t.Helper()
+	return &OpenAIGatewayHandler{
+		gatewayService:      &service.OpenAIGatewayService{},
+		billingCacheService: &service.BillingCacheService{},
+		apiKeyService:       &service.APIKeyService{},
+		concurrencyHelper:   NewConcurrencyHelper(service.NewConcurrencyService(cache), SSEPingFormatNone, time.Second),
+	}
+}
+
+func setImageChatTestAuth(c *gin.Context) {
+	apiKey := &service.APIKey{ID: 4348, UserID: 4348, User: &service.User{ID: 4348}}
+	c.Set(string(middleware.ContextKeyAPIKey), apiKey)
+	c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 1})
+}

backend/internal/handler/gateway_handler_chat_completions.go

@@ -80,6 +80,10 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
 		h.chatCompletionsErrorResponse(c, http.StatusBadRequest, "invalid_request_error", invalidStreamFieldTypeMessage)
 		return
 	}
+	if service.IsGPTImageGenerationModel(reqModel) {
+		h.chatCompletionsErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "This model is not supported on the Chat Completions endpoint")
+		return
+	}
 	reqLog = reqLog.With(zap.String("model", reqModel), zap.Bool("stream", reqStream))
 
 	setOpsRequestContext(c, reqModel, reqStream)

backend/internal/handler/openai_chat_completions.go

@@ -80,6 +80,10 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
 		h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", invalidStreamFieldTypeMessage)
 		return
 	}
+	if service.IsGPTImageGenerationModel(reqModel) {
+		h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "This model is not supported on the Chat Completions endpoint")
+		return
+	}
 
 	reqLog = reqLog.With(zap.String("model", reqModel), zap.Bool("stream", reqStream))
 

backend/internal/service/openai_images.go

@@ -455,8 +455,13 @@ func applyOpenAIImagesDefaults(req *OpenAIImagesRequest) {
 }
 
 func isOpenAIImageGenerationModel(model string) bool {
+	return IsGPTImageGenerationModel(model) || isGrokImageGenerationModel(model)
+}
+
+// IsGPTImageGenerationModel identifies the GPT native image-generation model family.
+func IsGPTImageGenerationModel(model string) bool {
 	model = strings.ToLower(strings.TrimSpace(model))
-	return strings.HasPrefix(model, "gpt-image-") || isGrokImageGenerationModel(model)
+	return strings.HasPrefix(model, "gpt-image-")
 }
 
 func isGrokImageGenerationModel(model string) bool {