From 335edde9c80395eb684a9938736f754555a05558 Mon Sep 17 00:00:00 2001 From: Heatherm Huang Date: Sat, 18 Jul 2026 14:06:42 +0800 Subject: [PATCH] fix(grok): apply account model mapping to media --- backend/internal/handler/grok_media.go | 23 ++- backend/internal/handler/grok_media_test.go | 19 +++ backend/internal/service/grok_media.go | 20 ++- .../service/openai_gateway_grok_test.go | 132 +++++++++++++++++- 4 files changed, 178 insertions(+), 16 deletions(-) diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 6ae7129451..98f1ee69bf 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -95,6 +95,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. contentType := c.GetHeader("Content-Type") requestInfo := service.ParseGrokMediaRequest(contentType, body) requestModel := requestInfo.Model + routingModel := service.NormalizeGrokMediaModelForEndpoint(endpoint, requestModel, requestInfo.HasInputImage()) if endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) == "" { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required") return @@ -184,7 +185,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. apiKey.GroupID, "", sessionHash, - requestModel, + routingModel, failedAccountIDs, service.OpenAIUpstreamTransportHTTPSSE, requiredCapability, @@ -208,7 +209,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } if len(failedAccountIDs) == 0 { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok) + cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, service.PlatformGrok) if !cls.ModelNotFound { markOpsRoutingCapacityLimitedIfNoAvailable(c, err) } @@ -228,7 +229,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts") return } - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok) + cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, service.PlatformGrok) if !cls.ModelNotFound { markOpsRoutingCapacityLimited(c) } @@ -285,7 +286,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } if failoverErr.ShouldReportAccountScheduleFailure() { - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, nil), false, nil) } if c.Writer.Size() != writerSizeBeforeForward { h.handleFailoverExhausted(c, failoverErr, true) @@ -333,7 +334,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. ) continue } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, nil), false, nil) if c.Writer.Size() == writerSizeBeforeForward { h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed") } @@ -344,7 +345,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. return } - h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), true, nil) + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, result), true, nil) if endpoint.IsGenerationRequest() && strings.TrimSpace(result.ResponseID) != "" { if err := h.gatewayService.BindGrokMediaVideoRequestAccount(requestCtx, apiKey.GroupID, result.ResponseID, account.ID); err != nil { reqLog.Warn("grok_media.bind_video_request_account_failed", @@ -372,6 +373,16 @@ func grokMediaRequiredCapability(endpoint service.GrokMediaEndpoint) service.Ope return "" } +func grokMediaScheduleModel(account *service.Account, routingModel string, result *service.OpenAIForwardResult) string { + if result != nil && strings.TrimSpace(result.UpstreamModel) != "" { + return result.UpstreamModel + } + if account == nil { + return strings.TrimSpace(routingModel) + } + return account.GetMappedModel(routingModel) +} + func shouldRecordGrokMediaUsage(endpoint service.GrokMediaEndpoint, requestModel string) bool { return endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) != "" } diff --git a/backend/internal/handler/grok_media_test.go b/backend/internal/handler/grok_media_test.go index 3c58742c78..8c52bb3c83 100644 --- a/backend/internal/handler/grok_media_test.go +++ b/backend/internal/handler/grok_media_test.go @@ -73,3 +73,22 @@ func TestGrokMediaRequiredCapability(t *testing.T) { }) } } + +func TestGrokMediaScheduleModelUsesNormalizedMappedUpstream(t *testing.T) { + account := &service.Account{ + Platform: service.PlatformGrok, + Credentials: map[string]any{ + "model_mapping": map[string]any{ + "grok-imagine-video-1.5": "wrong-raw-model", + "grok-imagine-video": "mapped-video-model", + }, + }, + } + + require.Equal(t, "mapped-video-model", grokMediaScheduleModel(account, "grok-imagine-video", nil)) + require.Equal(t, "actual-upstream-model", grokMediaScheduleModel(account, "grok-imagine-video", &service.OpenAIForwardResult{ + UpstreamModel: "actual-upstream-model", + })) + require.Equal(t, "mapped-video-model", grokMediaScheduleModel(account, "grok-imagine-video", &service.OpenAIForwardResult{})) + require.Equal(t, "grok-imagine-video", grokMediaScheduleModel(nil, " grok-imagine-video ", nil)) +} diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index c75bd271f5..5aebc826f1 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -311,6 +311,18 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( return nil, err } requestInfo := ParseGrokMediaRequest(contentType, body) + upstreamModel := requestInfo.Model + if endpoint.RequiresRequestBody() && gjson.ValidBytes(body) { + if mappedModel := strings.TrimSpace(account.GetMappedModel(requestInfo.Model)); mappedModel != "" { + upstreamModel = mappedModel + } + if upstreamModel != requestInfo.Model { + body, err = sjson.SetBytes(body, "model", upstreamModel) + if err != nil { + return nil, fmt.Errorf("rewrite grok media account mapped model: %w", err) + } + } + } body, contentType, err = sanitizeGrokMediaForwardBody(endpoint, body, contentType) if err != nil { return nil, err @@ -372,7 +384,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( Usage: usage.Usage, Model: requestModel, BillingModel: requestModel, - UpstreamModel: requestModel, + UpstreamModel: upstreamModel, ResponseHeaders: resp.Header.Clone(), Duration: time.Since(startTime), ImageCount: usage.ImageCount, @@ -470,7 +482,7 @@ func normalizeGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, cont return nil, "", err } info := ParseGrokMediaRequest(contentType, body) - upstreamModel := normalizeGrokMediaModelForEndpoint(endpoint, info.Model, info.HasInputImage()) + upstreamModel := NormalizeGrokMediaModelForEndpoint(endpoint, info.Model, info.HasInputImage()) if upstreamModel == "" || upstreamModel == info.Model { return body, contentType, nil } @@ -552,7 +564,9 @@ func (r GrokMediaRequestInfo) HasInputImage() bool { return len(r.InputImageURLs) > 0 || len(r.Uploads) > 0 } -func normalizeGrokMediaModelForEndpoint(endpoint GrokMediaEndpoint, model string, hasInputImage bool) string { +// NormalizeGrokMediaModelForEndpoint resolves the built-in upstream model alias +// for a media endpoint before account-level model mapping and scheduling. +func NormalizeGrokMediaModelForEndpoint(endpoint GrokMediaEndpoint, model string, hasInputImage bool) string { model = strings.TrimSpace(model) switch endpoint { case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits: diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 67ad712147..b22e1bd1d8 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -519,7 +519,7 @@ func TestNormalizeGrokMediaModelForEndpoint(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - require.Equal(t, tt.want, normalizeGrokMediaModelForEndpoint(tt.endpoint, tt.model, tt.hasInputImage)) + require.Equal(t, tt.want, NormalizeGrokMediaModelForEndpoint(tt.endpoint, tt.model, tt.hasInputImage)) }) } } @@ -573,6 +573,115 @@ func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) { require.Equal(t, ImageBillingSize2K, result.ImageSize) } +func TestForwardGrokMediaAppliesAccountModelMappingAfterEndpointNormalization(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + endpoint GrokMediaEndpoint + path string + body string + modelMapping map[string]any + wantRequestModel string + wantUpstream string + wantBody string + responseBody string + }{ + { + name: "image generation maps normalized image alias", + endpoint: GrokMediaEndpointImagesGenerations, + path: "/v1/images/generations", + body: `{"model":"grok-imagine","prompt":"draw a cat"}`, + modelMapping: map[string]any{"grok-imagine-image-quality": "vendor-image-model"}, + wantRequestModel: "grok-imagine-image-quality", + wantUpstream: "vendor-image-model", + wantBody: `{"model":"vendor-image-model","prompt":"draw a cat"}`, + responseBody: `{"data":[]}`, + }, + { + name: "video generation maps text-only fallback model", + endpoint: GrokMediaEndpointVideosGenerations, + path: "/v1/videos/generations", + body: `{"model":"grok-imagine-video-1.5","prompt":"waves"}`, + modelMapping: map[string]any{"grok-imagine-video": "grok-image-video"}, + wantRequestModel: "grok-imagine-video", + wantUpstream: "grok-image-video", + wantBody: `{"model":"grok-image-video","prompt":"waves"}`, + responseBody: `{"request_id":"video-request-mapped"}`, + }, + { + name: "image-to-video preserves then maps the requested model", + endpoint: GrokMediaEndpointVideosGenerations, + path: "/v1/videos/generations", + body: `{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"url":"https://example.com/input.png"}}`, + modelMapping: map[string]any{"grok-imagine-video-1.5": "vendor-image-video"}, + wantRequestModel: "grok-imagine-video-1.5", + wantUpstream: "vendor-image-video", + wantBody: `{"model":"vendor-image-video","prompt":"animate","image":{"url":"https://example.com/input.png"}}`, + responseBody: `{"request_id":"image-video-request-mapped"}`, + }, + { + name: "mapping and image sanitization compose", + endpoint: GrokMediaEndpointImagesGenerations, + path: "/v1/images/generations", + body: `{"model":"grok-imagine","prompt":"draw","size":"1024x1024"}`, + modelMapping: map[string]any{"grok-imagine-image-quality": "vendor-image-model"}, + wantRequestModel: "grok-imagine-image-quality", + wantUpstream: "vendor-image-model", + wantBody: `{"model":"vendor-image-model","prompt":"draw"}`, + responseBody: `{"data":[]}`, + }, + { + name: "whitespace mapping target safely preserves normalized model", + endpoint: GrokMediaEndpointImagesGenerations, + path: "/v1/images/generations", + body: `{"model":"grok-imagine","prompt":"draw"}`, + modelMapping: map[string]any{"grok-imagine-image-quality": " "}, + wantRequestModel: "grok-imagine-image-quality", + wantUpstream: "grok-imagine-image-quality", + wantBody: `{"model":"grok-imagine-image-quality","prompt":"draw"}`, + responseBody: `{"data":[]}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, tt.path, strings.NewReader(tt.body)) + c.Request.Header.Set("Content-Type", "application/json") + + account := &Account{ + ID: 66, + Name: "grok-mapped", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "api-key", + "base_url": "https://xai.test/v1", + "model_mapping": tt.modelMapping, + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(tt.responseBody)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.ForwardGrokMedia(context.Background(), c, account, tt.endpoint, "", []byte(tt.body), "application/json") + + require.NoError(t, err) + require.JSONEq(t, tt.wantBody, string(upstream.lastBody)) + require.Equal(t, tt.wantRequestModel, result.Model) + require.Equal(t, tt.wantRequestModel, result.BillingModel) + require.Equal(t, tt.wantUpstream, result.UpstreamModel) + }) + } +} + func TestForwardGrokMediaImagesGenerationStripsUnsupportedSize(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") gin.SetMode(gin.TestMode) @@ -590,8 +699,9 @@ func TestForwardGrokMediaImagesGenerationStripsUnsupportedSize(t *testing.T) { Type: AccountTypeAPIKey, Concurrency: 1, Credentials: map[string]any{ - "api_key": "api-key", - "base_url": "https://xai.test/v1", + "api_key": "api-key", + "base_url": "https://xai.test/v1", + "model_mapping": map[string]any{"grok-imagine-edit": "vendor-image-edit"}, }, } upstream := &httpUpstreamRecorder{resp: &http.Response{ @@ -652,15 +762,17 @@ func TestForwardGrokMediaImagesEditMultipartConvertsToJSON(t *testing.T) { }} svc := &OpenAIGatewayService{httpUpstream: upstream} - _, err = svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesEdits, "", buf.Bytes(), writer.FormDataContentType()) + result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesEdits, "", buf.Bytes(), writer.FormDataContentType()) require.NoError(t, err) require.Equal(t, "https://xai.test/v1/images/edits", upstream.lastReq.URL.String()) require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type")) require.True(t, json.Valid(upstream.lastBody)) - require.Equal(t, "grok-imagine-edit", gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, "vendor-image-edit", gjson.GetBytes(upstream.lastBody, "model").String()) require.Equal(t, "edit this private image", gjson.GetBytes(upstream.lastBody, "prompt").String()) require.True(t, strings.HasPrefix(gjson.GetBytes(upstream.lastBody, "image.url").String(), "data:image/png;base64,")) require.False(t, gjson.GetBytes(upstream.lastBody, "image.image_url").Exists()) + require.Equal(t, "grok-imagine-edit", result.BillingModel) + require.Equal(t, "vendor-image-edit", result.UpstreamModel) } func TestForwardGrokMediaVideoGenerationReturnsUsageAndResponseID(t *testing.T) { @@ -857,7 +969,11 @@ func TestForwardGrokMediaVideoMutationEndpoints(t *testing.T) { account := &Account{ ID: 71, Name: "grok", Platform: PlatformGrok, Type: AccountTypeAPIKey, Concurrency: 1, - Credentials: map[string]any{"api_key": "api-key", "base_url": "https://xai.test/v1"}, + Credentials: map[string]any{ + "api_key": "api-key", + "base_url": "https://xai.test/v1", + "model_mapping": map[string]any{"grok-imagine-video": "vendor-video-mutation"}, + }, } upstream := &httpUpstreamRecorder{resp: &http.Response{ StatusCode: http.StatusOK, @@ -870,10 +986,12 @@ func TestForwardGrokMediaVideoMutationEndpoints(t *testing.T) { require.NoError(t, err) require.Equal(t, "https://xai.test/v1"+tt.path, upstream.lastReq.URL.String()) require.Equal(t, http.MethodPost, upstream.lastReq.Method) - require.JSONEq(t, string(body), string(upstream.lastBody)) + require.JSONEq(t, `{"model":"vendor-video-mutation","prompt":"continue","video":{"url":"https://example.com/in.mp4"},"duration":6}`, string(upstream.lastBody)) require.Equal(t, "video-mutation-123", result.ResponseID) require.Equal(t, 1, result.VideoCount) require.Equal(t, 6, result.VideoDurationSeconds) + require.Equal(t, "grok-imagine-video", result.BillingModel) + require.Equal(t, "vendor-video-mutation", result.UpstreamModel) }) } }