From 9fedf899e2d26fa33f295211ca9393b54ad05c14 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Sat, 12 Sep 2026 14:34:31 +0800 Subject: [PATCH] fix(images): route compatible Gemini image models through API keys --- backend/internal/handler/openai_images.go | 2 +- .../composite_images_compatible_test.go | 160 ++++++++++++++++++ backend/internal/service/account.go | 2 + backend/internal/service/openai_images.go | 42 ++++- .../service/openai_images_compatible_test.go | 3 + 5 files changed, 205 insertions(+), 4 deletions(-) create mode 100644 backend/internal/server/routes/composite_images_compatible_test.go diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 18e1d59195..9900861715 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -165,7 +165,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { sessionHash, routingModel, failedAccountIDs, - parsed.RequiredCapability, + parsed.RequiredCapabilityForModel(channelMapping.MappedModel), ) if err != nil { if failoverClientGone(c) { diff --git a/backend/internal/server/routes/composite_images_compatible_test.go b/backend/internal/server/routes/composite_images_compatible_test.go new file mode 100644 index 0000000000..1e0b9406ba --- /dev/null +++ b/backend/internal/server/routes/composite_images_compatible_test.go @@ -0,0 +1,160 @@ +//go:build unit + +package routes + +import ( + "bytes" + "context" + "fmt" + "io" + "mime/multipart" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/handler" + servermiddleware "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" +) + +type compatibleImagesAccounts struct { + service.AccountRepository + accounts []service.Account +} + +func (r compatibleImagesAccounts) GetByID(_ context.Context, id int64) (*service.Account, error) { + for _, account := range r.accounts { + if account.ID == id { + return &account, nil + } + } + return nil, service.ErrNoAvailableAccounts +} + +func (r compatibleImagesAccounts) ListSchedulableByPlatform(context.Context, string) ([]service.Account, error) { + return r.accounts, nil +} + +func (r compatibleImagesAccounts) ListSchedulableUngroupedByPlatform(context.Context, string) ([]service.Account, error) { + return r.accounts, nil +} + +func (r compatibleImagesAccounts) ListSchedulableByGroupIDAndPlatform(context.Context, int64, string) ([]service.Account, error) { + return r.accounts, nil +} + +type compatibleImagesUpstream struct { + service.HTTPUpstream + accountIDs []int64 + body []byte +} + +func (u *compatibleImagesUpstream) Do(req *http.Request, _ string, id int64, _ int) (*http.Response, error) { + u.accountIDs = append(u.accountIDs, id) + var err error + u.body, err = io.ReadAll(req.Body) + if err != nil { + return nil, err + } + return &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": {"application/json"}, "X-Request-Id": {"compatible-image-test"}}, Body: io.NopCloser(strings.NewReader(`{"data":[{"b64_json":"aW1hZ2U="}]}`))}, nil +} + +type compatibleImagesUsage struct { + service.UsageLogRepository + logs []*service.UsageLog +} + +func (r *compatibleImagesUsage) Create(_ context.Context, log *service.UsageLog) (bool, error) { + r.logs = append(r.logs, log) + return true, nil +} + +func TestCompositeCompatibleImagesEndToEnd(t *testing.T) { + gin.SetMode(gin.TestMode) + for _, scenario := range []string{"generation", "json_edit", "multipart_alias", "disabled", "native_only", "restricted"} { + t.Run(scenario, func(t *testing.T) { + const model = "gemini-3.1-flash-image" + groupID := int64(101) + price := 0.17 + group := &service.Group{ID: groupID, Platform: service.PlatformComposite, AllowImageGeneration: scenario != "disabled", RateMultiplier: 1, ImagePrice1K: &price, ImagePrice2K: &price, ImagePrice4K: &price} + accounts := []service.Account{} + // Native accounts deliberately advertise the same model and have higher + // priority, proving the image capability fence controls selection. + for i, typ := range []string{service.AccountTypeOAuth, service.AccountTypeSetupToken, service.AccountTypeAPIKey} { + if scenario == "native_only" && typ == service.AccountTypeAPIKey { + continue + } + mapping := map[string]any{model: model} + if scenario == "restricted" && typ == service.AccountTypeAPIKey { + mapping = map[string]any{"gpt-image-2": "gpt-image-2"} + } + accounts = append(accounts, service.Account{ID: int64(i + 1), Platform: service.PlatformOpenAI, Type: typ, Status: service.StatusActive, Schedulable: true, Priority: i, Credentials: map[string]any{"api_key": "fixture-key", "access_token": "unused-native-token", "base_url": "https://compatible.example/v1", "model_mapping": mapping}}) + } + cfg := &config.Config{RunMode: config.RunModeSimple} + repo := compatibleImagesAccounts{accounts: accounts} + upstream, usage := &compatibleImagesUpstream{}, &compatibleImagesUsage{} + billingCache := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billingCache.Stop) + gateway := service.NewOpenAIGatewayService(repo, usage, nil, nil, nil, nil, nil, cfg, nil, nil, service.NewBillingService(cfg, nil), nil, billingCache, upstream, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil) + imagesHandler := handler.NewOpenAIGatewayHandler(gateway, service.NewConcurrencyService(nil), billingCache, service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg), nil, nil, nil, nil, cfg) + publicModel := model + if scenario == "multipart_alias" { + publicModel = "public-image" + } + resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{routes: []service.CompositeModelRoute{{ID: 1, GroupID: groupID, PublicModel: publicModel, MatchType: service.CompositeRouteMatchExact, TargetPlatform: service.PlatformOpenAI, UpstreamModel: model, Endpoint: service.CompositeRouteEndpointImages, Enabled: true}}}) + router := gin.New() + router.Use(func(c *gin.Context) { + c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{ID: 202, GroupID: &groupID, Group: group, User: &service.User{ID: 303}}) + c.Set(string(servermiddleware.ContextKeyUser), servermiddleware.AuthSubject{UserID: 303}) + c.Next() + }) + router.Use(compositeTargetPlatformMiddleware(resolver)) + router.POST("/v1/images/generations", imagesHandler.Images) + router.POST("/v1/images/edits", imagesHandler.Images) + endpoint, contentType := "/v1/images/generations", "application/json" + body := []byte(fmt.Sprintf(`{"model":%q,"prompt":"draw","size":"1024x1024"}`, publicModel)) + if scenario == "json_edit" { + endpoint = "/v1/images/edits" + body = []byte(fmt.Sprintf(`{"model":%q,"prompt":"draw","images":[{"image_url":"https://source.example/input.png"}]}`, publicModel)) + } + if scenario == "multipart_alias" { + endpoint = "/v1/images/edits" + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + require.NoError(t, writer.WriteField("model", publicModel)) + require.NoError(t, writer.WriteField("prompt", "draw")) + part, err := writer.CreateFormFile("image", "input.png") + require.NoError(t, err) + _, err = part.Write([]byte("fixture-image")) + require.NoError(t, err) + require.NoError(t, writer.Close()) + body, contentType = buf.Bytes(), writer.FormDataContentType() + } + req := httptest.NewRequest(http.MethodPost, endpoint, bytes.NewReader(body)) + req.Header.Set("Content-Type", contentType) + rec := httptest.NewRecorder() + router.ServeHTTP(rec, req) + if scenario == "disabled" || scenario == "native_only" || scenario == "restricted" { + require.GreaterOrEqual(t, rec.Code, 400, rec.Body.String()) + if scenario == "disabled" { + require.Equal(t, http.StatusForbidden, rec.Code) + } + require.Empty(t, upstream.accountIDs) + require.Empty(t, usage.logs) + return + } + require.Equal(t, http.StatusOK, rec.Code, rec.Body.String()) + require.Equal(t, []int64{3}, upstream.accountIDs) + require.Contains(t, string(upstream.body), model) + require.Contains(t, rec.Body.String(), "aW1hZ2U=") + require.Len(t, usage.logs, 1, "the image must reach usage recording, not just return HTTP 200") + require.Equal(t, 1, usage.logs[0].ImageCount) + require.InDelta(t, price, usage.logs[0].ActualCost, 1e-9) + require.Equal(t, publicModel, usage.logs[0].OriginalModel) + }) + } +} diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index b5268e9102..bee3e8946a 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -2008,6 +2008,8 @@ func (a *Account) SupportsOpenAIImageCapability(capability OpenAIImagesCapabilit return false } switch capability { + case OpenAIImagesCapabilityAPIKey: + return a.Type == AccountTypeAPIKey case OpenAIImagesCapabilityBasic, OpenAIImagesCapabilityNative: return a.Type == AccountTypeOAuth || a.Type == AccountTypeSetupToken || a.Type == AccountTypeAPIKey default: diff --git a/backend/internal/service/openai_images.go b/backend/internal/service/openai_images.go index 0fe06a334b..c71a9ea05d 100644 --- a/backend/internal/service/openai_images.go +++ b/backend/internal/service/openai_images.go @@ -59,6 +59,8 @@ type OpenAIImagesCapability string const ( OpenAIImagesCapabilityBasic OpenAIImagesCapability = "images-basic" OpenAIImagesCapabilityNative OpenAIImagesCapability = "images-native" + // Compatible provider models require the API-key Images passthrough path. + OpenAIImagesCapabilityAPIKey OpenAIImagesCapability = "images-apikey" ) type OpenAIImagesUpload struct { @@ -227,7 +229,14 @@ func (s *OpenAIGatewayService) ParseOpenAIImagesRequest(c *gin.Context, body []b } applyOpenAIImagesDefaults(req) - if err := validateOpenAIImagesModel(req.Model); err != nil { + // Composite middleware preserves multipart bodies, including their public + // alias. Validate and forward the resolved model without altering uploads. + if platform, _ := ResolvedTargetPlatformFromContext(c.Request.Context()); platform == PlatformOpenAI { + if model, ok := ResolvedUpstreamModelFromContext(c.Request.Context()); ok { + req.Model = model + } + } + if err := validateCompatibleImagesModel(req.Model); err != nil { return nil, err } req.SizeTier = normalizeOpenAIImageSizeTier(req.Size) @@ -494,6 +503,30 @@ func validateOpenAIImagesModel(model string) error { return fmt.Errorf("images endpoint requires an image model, got %q", model) } +// Keep this separate from isOpenAIImageGenerationModel: that predicate also +// drives native Responses tool conversion, pricing and rate-limit policy. +func isGeminiCompatibleImageModel(model string) bool { + model = strings.ToLower(strings.TrimSpace(model)) + return strings.HasPrefix(model, "gemini-") && + (strings.HasSuffix(model, "-image") || strings.Contains(model, "-image-")) +} + +func validateCompatibleImagesModel(model string) error { + if isGeminiCompatibleImageModel(model) { + return nil + } + return validateOpenAIImagesModel(model) +} + +// RequiredCapabilityForModel also applies the API-key-only fence when channel +// mapping introduces a compatible provider model after request parsing. +func (req *OpenAIImagesRequest) RequiredCapabilityForModel(model string) OpenAIImagesCapability { + if isGeminiCompatibleImageModel(model) { + return OpenAIImagesCapabilityAPIKey + } + return req.RequiredCapability +} + func normalizeOpenAIImagesEndpointPath(path string) string { trimmed := strings.TrimSpace(path) switch { @@ -510,6 +543,9 @@ func classifyOpenAIImagesCapability(req *OpenAIImagesRequest) OpenAIImagesCapabi if req == nil { return OpenAIImagesCapabilityNative } + if isGeminiCompatibleImageModel(req.Model) { + return OpenAIImagesCapabilityAPIKey + } if req.ExplicitModel || req.ExplicitSize { return OpenAIImagesCapabilityNative } @@ -594,11 +630,11 @@ func (s *OpenAIGatewayService) forwardOpenAIImagesAPIKey( if mapped := strings.TrimSpace(channelMappedModel); mapped != "" { requestModel = mapped } - if err := validateOpenAIImagesModel(requestModel); err != nil { + if err := validateCompatibleImagesModel(requestModel); err != nil { return nil, err } upstreamModel := account.GetMappedModel(requestModel) - if err := validateOpenAIImagesModel(upstreamModel); err != nil { + if err := validateCompatibleImagesModel(upstreamModel); err != nil { return nil, err } SetOpsUpstreamModel(c, upstreamModel) diff --git a/backend/internal/service/openai_images_compatible_test.go b/backend/internal/service/openai_images_compatible_test.go index 4adc7c208a..10e50f8104 100644 --- a/backend/internal/service/openai_images_compatible_test.go +++ b/backend/internal/service/openai_images_compatible_test.go @@ -98,6 +98,9 @@ func TestCompatibleImagesForwardGemini(t *testing.T) { svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} parsed, err := svc.ParseOpenAIImagesRequest(c, body) require.NoError(t, err) + if kind == "channel_mapping" { + require.Equal(t, OpenAIImagesCapabilityAPIKey, parsed.RequiredCapabilityForModel(channelModel)) + } result, err := svc.ForwardImages(c.Request.Context(), c, &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Credentials: credentials}, body, parsed, channelModel) require.NoError(t, err) require.Equal(t, 1, result.ImageCount)