diff --git a/backend/internal/handler/composite_platform_test.go b/backend/internal/handler/composite_platform_test.go index dc80ca3ff6..95dfa92d5f 100644 --- a/backend/internal/handler/composite_platform_test.go +++ b/backend/internal/handler/composite_platform_test.go @@ -4,6 +4,7 @@ import ( "net/http/httptest" "testing" + middleware2 "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" @@ -60,3 +61,25 @@ func TestCompositeTargetPlatformResolvedAllowsConcreteGroupWithoutResolution(t * require.True(t, compositeTargetPlatformResolved(c, apiKey, "llama-4-maverick")) } + +func TestClientRequestedModelUsesCompositePublicModel(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest("POST", "/v1/chat/completions", nil) + c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), service.CompositeRouteDecision{ + Matched: true, + Source: service.CompositeRouteSourceExplicit, + PublicModel: "public-alias", + TargetPlatform: service.PlatformOpenAI, + UpstreamModel: "gpt-5", + })) + + input := buildContentModerationInput(c, nil, middleware2.AuthSubject{UserID: 42}, service.ContentModerationProtocolOpenAIChat, "gpt-5", nil) + require.Equal(t, "public-alias", input.Model) + require.Equal(t, service.PlatformOpenAI, input.Provider) + + fields := clientRequestedUsageFields(c, service.ChannelMappingResult{MappedModel: "gpt-5"}, "gpt-5", "gpt-5") + require.Equal(t, "public-alias", fields.OriginalModel) + require.Equal(t, "public-alias", fields.ChannelMappedModel) + require.Equal(t, "public-alias\u2192gpt-5", fields.ModelMappingChain) +} diff --git a/backend/internal/handler/content_moderation_helper.go b/backend/internal/handler/content_moderation_helper.go index f91fd8a8f7..c7e7ff9b0f 100644 --- a/backend/internal/handler/content_moderation_helper.go +++ b/backend/internal/handler/content_moderation_helper.go @@ -23,6 +23,21 @@ func contentModerationErrorCode(decision *service.ContentModerationDecision) str return "content_policy_violation" } +func clientRequestedModel(c *gin.Context, fallback string) string { + fallback = strings.TrimSpace(fallback) + if c == nil || c.Request == nil { + return fallback + } + if model, ok := service.RequestedPublicModelFromContext(c.Request.Context()); ok { + return model + } + return fallback +} + +func clientRequestedUsageFields(c *gin.Context, mapping service.ChannelMappingResult, fallbackModel, upstreamModel string) service.ChannelUsageFields { + return mapping.ToUsageFields(clientRequestedModel(c, fallbackModel), upstreamModel) +} + func runContentModeration(c *gin.Context, reqLog *zap.Logger, svc *service.ContentModerationService, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol string, model string, body []byte) *service.ContentModerationDecision { if svc == nil || c == nil || c.Request == nil { return nil @@ -71,10 +86,13 @@ func buildContentModerationInput(c *gin.Context, apiKey *service.APIKey, subject UserID: subject.UserID, Endpoint: GetInboundEndpoint(c), Provider: contentModerationProvider(apiKey), - Model: strings.TrimSpace(model), + Model: clientRequestedModel(c, model), Protocol: protocol, Body: body, } + if resolvedPlatform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context()); ok { + input.Provider = resolvedPlatform + } if forcedPlatform, ok := middleware2.GetForcePlatformFromContext(c); ok { input.Provider = strings.TrimSpace(forcedPlatform) } diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 739b9e8e44..65fd059920 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -550,7 +550,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { RequestPayloadHash: requestPayloadHash, ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, - ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), }); err != nil { logger.L().With( zap.String("component", "handler.gateway.messages"), @@ -985,7 +985,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) { RequestPayloadHash: requestPayloadHash, ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, - ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), }); err != nil { logger.L().With( zap.String("component", "handler.gateway.messages"), @@ -1085,7 +1085,7 @@ func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID * } seen := make(map[string]struct{}) models := make([]string, 0) - for _, platform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformGrok} { + for _, platform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok} { for _, model := range h.gatewayService.GetAvailableModels(ctx, groupID, platform) { model = strings.TrimSpace(model) if model == "" { @@ -1304,7 +1304,7 @@ func defaultModelIDsForPlatform(platform string) []string { case service.PlatformComposite: ids := make([]string, 0) seen := make(map[string]struct{}) - for _, concretePlatform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformGrok} { + for _, concretePlatform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok} { for _, id := range defaultModelIDsForPlatform(concretePlatform) { if _, ok := seen[id]; ok { continue diff --git a/backend/internal/handler/gateway_handler_chat_completions.go b/backend/internal/handler/gateway_handler_chat_completions.go index 0f779441d0..3cfb9026e2 100644 --- a/backend/internal/handler/gateway_handler_chat_completions.go +++ b/backend/internal/handler/gateway_handler_chat_completions.go @@ -312,7 +312,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) { IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, - ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), }); err != nil { reqLog.Error("gateway.cc.record_usage_failed", zap.Int64("account_id", account.ID), diff --git a/backend/internal/handler/gateway_handler_responses.go b/backend/internal/handler/gateway_handler_responses.go index 475ad9e88f..bd545d2cf6 100644 --- a/backend/internal/handler/gateway_handler_responses.go +++ b/backend/internal/handler/gateway_handler_responses.go @@ -290,7 +290,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) { IPAddress: clientIP, RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, - ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), }); err != nil { reqLog.Error("gateway.responses.record_usage_failed", zap.Int64("account_id", account.ID), diff --git a/backend/internal/handler/gateway_models_test.go b/backend/internal/handler/gateway_models_test.go index af1bc5aea6..b9b419405a 100644 --- a/backend/internal/handler/gateway_models_test.go +++ b/backend/internal/handler/gateway_models_test.go @@ -61,6 +61,14 @@ func newGatewayModelsHandlerForTest(repo service.AccountRepository) *GatewayHand } } +func TestDefaultModelIDsForCompositeIncludesAntigravityDefaults(t *testing.T) { + antigravityIDs := defaultModelIDsForPlatform(service.PlatformAntigravity) + require.NotEmpty(t, antigravityIDs) + + compositeIDs := defaultModelIDsForPlatform(service.PlatformComposite) + require.Contains(t, compositeIDs, antigravityIDs[0]) +} + func TestGatewayModels_GeminiGroupFallsBackToGeminiModels(t *testing.T) { gin.SetMode(gin.TestMode) @@ -304,6 +312,15 @@ func TestGatewayModels_CompositeCustomModelsListFiltersAcrossConcretePlatforms(t }, }, }, + { + ID: 3, + Platform: service.PlatformAntigravity, + Credentials: map[string]any{ + "model_mapping": map[string]any{ + "ag-custom-model": "ag-custom-model", + }, + }, + }, }, }, }, @@ -318,7 +335,7 @@ func TestGatewayModels_CompositeCustomModelsListFiltersAcrossConcretePlatforms(t Platform: service.PlatformComposite, ModelsListConfig: service.GroupModelsListConfig{ Enabled: true, - Models: []string{"gemini-2.5-flash", "missing-model", "gpt-5.5"}, + Models: []string{"gemini-2.5-flash", "missing-model", "ag-custom-model", "gpt-5.5"}, }, }, }) @@ -329,7 +346,7 @@ func TestGatewayModels_CompositeCustomModelsListFiltersAcrossConcretePlatforms(t var got gatewayModelsResponseForTest require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) - require.Equal(t, []string{"gemini-2.5-flash", "gpt-5.5"}, modelIDsForTest(got.Data)) + require.Equal(t, []string{"gemini-2.5-flash", "ag-custom-model", "gpt-5.5"}, modelIDsForTest(got.Data)) } func TestGatewayModels_CustomModelsListKeepsConcreteModelAllowedByWildcardMapping(t *testing.T) { diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 557bec3f9f..7b349ae1db 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -549,7 +549,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { LongContextMultiplier: 2.0, // 超出部分双倍计费 ForceCacheBilling: forceCacheBilling, APIKeyService: h.apiKeyService, - ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), }); err != nil { logger.L().With( zap.String("component", "handler.gemini_v1beta.models"), diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index e90d7fcffb..8ccc4fb49e 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -232,7 +232,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { if service.GetOpsCyberPolicy(c) != nil { cyberBlockKeyChat = service.CyberSessionBlockKey(apiKey.ID, c, body) } - h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, reqModel, err != nil, cyberBlockKeyChat, channelMapping.ToUsageFields(reqModel, ""), service.HashUsageRequestPayload(body)) + h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, reqModel, err != nil, cyberBlockKeyChat, clientRequestedUsageFields(c, channelMapping, reqModel, ""), service.HashUsageRequestPayload(body)) forwardDurationMs := time.Since(forwardStart).Milliseconds() upstreamLatencyMs, _ := getContextInt64(c, service.OpsUpstreamLatencyMsKey) @@ -355,7 +355,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { IPAddress: clientIP, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, - ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { logger.L().With( diff --git a/backend/internal/handler/openai_embeddings.go b/backend/internal/handler/openai_embeddings.go index 0b3fd82ccc..b7bf64ce0e 100644 --- a/backend/internal/handler/openai_embeddings.go +++ b/backend/internal/handler/openai_embeddings.go @@ -257,7 +257,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) { IPAddress: clientIP, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, - ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), }); err != nil { logger.L().With( zap.String("component", "handler.openai_gateway.embeddings"), diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 8c24c85172..9fd18f6cb8 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -490,7 +490,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { if service.GetOpsCyberPolicy(c) != nil { cyberBlockKeyHTTP = service.CyberSessionBlockKey(apiKey.ID, c, sessionHashBody) } - h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, reqModel, err != nil, cyberBlockKeyHTTP, channelMapping.ToUsageFields(reqModel, ""), service.HashUsageRequestPayload(body)) + h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, reqModel, err != nil, cyberBlockKeyHTTP, clientRequestedUsageFields(c, channelMapping, reqModel, ""), service.HashUsageRequestPayload(body)) forwardDurationMs := time.Since(forwardStart).Milliseconds() upstreamLatencyMs, _ := getContextInt64(c, service.OpsUpstreamLatencyMsKey) responseLatencyMs := forwardDurationMs @@ -640,7 +640,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, - ChannelUsageFields: channelMapping.ToUsageFields(reqModel, result.UpstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { logger.L().With( @@ -1027,7 +1027,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { if service.GetOpsCyberPolicy(c) != nil { cyberBlockKeyMsg = service.CyberSessionBlockKey(apiKey.ID, c, body) } - h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, reqModel, err != nil, cyberBlockKeyMsg, channelMappingMsg.ToUsageFields(reqModel, ""), service.HashUsageRequestPayload(body)) + h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, reqModel, err != nil, cyberBlockKeyMsg, clientRequestedUsageFields(c, channelMappingMsg, reqModel, ""), service.HashUsageRequestPayload(body)) forwardDurationMs := time.Since(forwardStart).Milliseconds() upstreamLatencyMs, _ := getContextInt64(c, service.OpsUpstreamLatencyMsKey) responseLatencyMs := forwardDurationMs @@ -1150,7 +1150,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, - ChannelUsageFields: channelMappingMsg.ToUsageFields(reqModel, result.UpstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMappingMsg, reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { logger.L().With( @@ -1828,7 +1828,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { // 届时 defer 已清除标记）。 defer clearCyberPolicyTurnState(c) releaseTurnSlots() - h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, reqModel, turnErr != nil, cyberBlockKey, channelMappingWS.ToUsageFields(reqModel, ""), requestPayloadHash) + h.recordCyberPolicyIfMarked(c, apiKey, account, subscription, reqModel, turnErr != nil, cyberBlockKey, clientRequestedUsageFields(c, channelMappingWS, reqModel, ""), requestPayloadHash) if service.GetOpsCyberPolicy(c) != nil { cyberBlockedThisConn = true } @@ -1873,7 +1873,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, - ChannelUsageFields: channelMappingWS.ToUsageFields(reqModel, result.UpstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMappingWS, reqModel, result.UpstreamModel), CyberBlocked: cyberBlocked, }); err != nil { reqLog.Error("openai.websocket_record_usage_failed", @@ -2830,6 +2830,7 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey return } c.Set(cyberPolicyRecordedKey, true) + model = clientRequestedModel(c, model) requestID := c.Writer.Header().Get("X-Request-Id") var userID, apiKeyID int64 diff --git a/backend/internal/handler/openai_images.go b/backend/internal/handler/openai_images.go index 6509672eac..7ae070f6c6 100644 --- a/backend/internal/handler/openai_images.go +++ b/backend/internal/handler/openai_images.go @@ -75,13 +75,19 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { } requestModel := parsed.Model ensureCompositeTargetPlatform(c, apiKey, requestModel) + clientRequestModel := clientRequestedModel(c, requestModel) + routingModel := requestModel + if resolvedModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()); ok { + routingModel = resolvedModel + } if !compositeTargetPlatformAllowed(c, apiKey, requestModel, service.PlatformOpenAI) { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by this OpenAI-compatible endpoint for composite groups") return } reqLog = reqLog.With( - zap.String("model", requestModel), + zap.String("model", clientRequestModel), + zap.String("routing_model", routingModel), zap.Bool("stream", parsed.Stream), zap.Bool("multipart", parsed.Multipart), zap.String("capability", string(parsed.RequiredCapability)), @@ -103,14 +109,10 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { defer imageReleaseFunc() } - if parsed.Multipart { - setOpsRequestContext(c, requestModel, parsed.Stream) - } else { - setOpsRequestContext(c, requestModel, parsed.Stream) - } + setOpsRequestContext(c, clientRequestModel, parsed.Stream) setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(parsed.Stream, false))) - channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, requestModel) + channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, routingModel) if h.errorPassthroughService != nil { service.BindErrorPassthroughService(c, h.errorPassthroughService) @@ -158,7 +160,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { requestCtx, apiKey.GroupID, sessionHash, - requestModel, + routingModel, failedAccountIDs, parsed.RequiredCapability, ) @@ -172,7 +174,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { zap.Int("excluded_account_count", len(failedAccountIDs)), ) if len(failedAccountIDs) == 0 { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformOpenAI) + cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, clientRequestModel, routingModel, service.PlatformOpenAI) if !cls.ModelNotFound { markOpsRoutingCapacityLimitedIfNoAvailable(c, err) } @@ -191,7 +193,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { return } if selection == nil || selection.Account == nil { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformOpenAI) + cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, clientRequestModel, routingModel, service.PlatformOpenAI) if !cls.ModelNotFound { markOpsRoutingCapacityLimited(c) } @@ -386,14 +388,14 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) { RequestPayloadHash: requestPayloadHash, APIKeyService: h.apiKeyService, QuotaPlatform: quotaPlatform, - ChannelUsageFields: channelMapping.ToUsageFields(requestModel, upstreamModel), + ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, requestModel, upstreamModel), }); err != nil { logger.L().With( zap.String("component", "handler.openai_gateway.images"), zap.Int64("user_id", subject.UserID), zap.Int64("api_key_id", apiKey.ID), zap.Any("group_id", apiKey.GroupID), - zap.String("model", requestModel), + zap.String("model", clientRequestModel), zap.Int64("account_id", account.ID), ).Error("openai.images.record_usage_failed", zap.Error(err)) } diff --git a/backend/internal/server/routes/composite_platform_test.go b/backend/internal/server/routes/composite_platform_test.go index 7b7e5f8c8d..682d129ebc 100644 --- a/backend/internal/server/routes/composite_platform_test.go +++ b/backend/internal/server/routes/composite_platform_test.go @@ -1,8 +1,10 @@ package routes import ( + "bytes" "context" "io" + "mime/multipart" "net/http" "net/http/httptest" "strings" @@ -131,6 +133,67 @@ func TestCompositeTargetPlatformMiddlewareUsesExplicitRouteAndRewritesBody(t *te require.Equal(t, http.StatusNoContent, w.Code) } +func TestCompositeTargetPlatformMiddlewareUsesExplicitRouteForMultipartImages(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{ + routes: []service.CompositeModelRoute{ + { + ID: 1, + GroupID: 1, + PublicModel: "image-alias", + MatchType: service.CompositeRouteMatchExact, + TargetPlatform: service.PlatformOpenAI, + UpstreamModel: "gpt-image-1", + Endpoint: service.CompositeRouteEndpointImages, + Priority: 100, + Enabled: true, + }, + }, + }) + router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) { + groupID := int64(1) + c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{ + GroupID: &groupID, + Group: &service.Group{ID: groupID, Platform: service.PlatformComposite}, + }) + c.Next() + }))) + router.Use(compositeTargetPlatformMiddleware(resolver)) + router.POST("/v1/images/edits", func(c *gin.Context) { + platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context()) + require.True(t, ok) + require.Equal(t, service.PlatformOpenAI, platform) + + upstreamModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()) + require.True(t, ok) + require.Equal(t, "gpt-image-1", upstreamModel) + + publicModel, ok := service.RequestedPublicModelFromContext(c.Request.Context()) + require.True(t, ok) + require.Equal(t, "image-alias", publicModel) + + body, err := io.ReadAll(c.Request.Body) + require.NoError(t, err) + require.Contains(t, string(body), "image-alias") + c.Status(http.StatusNoContent) + }) + + var body bytes.Buffer + writer := multipart.NewWriter(&body) + require.NoError(t, writer.WriteField("model", "image-alias")) + require.NoError(t, writer.WriteField("prompt", "draw")) + require.NoError(t, writer.Close()) + + req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body.Bytes())) + req.Header.Set("Content-Type", writer.FormDataContentType()) + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusNoContent, w.Code) +} + func TestCompositeGeminiTargetPlatformMiddlewareUsesPathRoute(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index b61d5d5e3f..c12bddcd0c 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -4,6 +4,8 @@ import ( "bytes" "errors" "io" + "mime" + "mime/multipart" "net/http" "strconv" "strings" @@ -387,7 +389,7 @@ func compositeTargetPlatformMiddleware(resolver *service.CompositeRouteResolver) return } - model := strings.TrimSpace(gjson.GetBytes(body, "model").String()) + model := compositeRequestModelFromBody(c.GetHeader("Content-Type"), body) if model != "" { decision, err := resolver.Resolve(c.Request.Context(), apiKey.Group.ID, model, compositeRouteEndpointForPath(c.Request.URL.Path)) if err != nil { @@ -409,6 +411,42 @@ func compositeTargetPlatformMiddleware(resolver *service.CompositeRouteResolver) } } +func compositeRequestModelFromBody(contentType string, body []byte) string { + if model := strings.TrimSpace(gjson.GetBytes(body, "model").String()); model != "" { + return model + } + return compositeMultipartModelFromBody(contentType, body) +} + +func compositeMultipartModelFromBody(contentType string, body []byte) string { + mediaType, params, err := mime.ParseMediaType(strings.TrimSpace(contentType)) + if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") { + return "" + } + boundary := strings.TrimSpace(params["boundary"]) + if boundary == "" { + return "" + } + reader := multipart.NewReader(bytes.NewReader(body), boundary) + for { + part, err := reader.NextPart() + if errors.Is(err, io.EOF) { + return "" + } + if err != nil { + return "" + } + if part.FormName() != "model" || part.FileName() != "" { + continue + } + data, err := io.ReadAll(part) + if err != nil { + return "" + } + return strings.TrimSpace(string(data)) + } +} + func compositeGeminiTargetPlatformMiddleware(resolver *service.CompositeRouteResolver) gin.HandlerFunc { if resolver == nil { resolver = service.NewCompositeRouteResolver(nil)