diff --git a/backend/internal/handler/endpoint.go b/backend/internal/handler/endpoint.go index e33d88241c..f8689a6e3a 100644 --- a/backend/internal/handler/endpoint.go +++ b/backend/internal/handler/endpoint.go @@ -21,6 +21,8 @@ const ( EndpointResponses = "/v1/responses" EndpointImagesGenerations = "/v1/images/generations" EndpointImagesEdits = "/v1/images/edits" + EndpointVideosGenerations = "/v1/videos/generations" + EndpointVideos = "/v1/videos" EndpointGeminiModels = "/v1beta/models" ) @@ -53,6 +55,10 @@ func NormalizeInboundEndpoint(path string) string { return EndpointImagesGenerations case strings.Contains(path, EndpointImagesEdits) || strings.Contains(path, "/images/edits"): return EndpointImagesEdits + case strings.Contains(path, EndpointVideosGenerations) || strings.Contains(path, "/videos/generations"): + return EndpointVideosGenerations + case strings.Contains(path, EndpointVideos) || strings.Contains(path, "/videos/"): + return EndpointVideos case strings.Contains(path, EndpointResponses): return EndpointResponses case strings.Contains(path, EndpointGeminiModels): @@ -78,7 +84,7 @@ func DeriveUpstreamEndpoint(inbound, rawRequestPath, platform string) string { switch platform { case service.PlatformOpenAI, service.PlatformGrok: - if inbound == EndpointEmbeddings || inbound == EndpointImagesGenerations || inbound == EndpointImagesEdits { + if inbound == EndpointEmbeddings || inbound == EndpointImagesGenerations || inbound == EndpointImagesEdits || inbound == EndpointVideosGenerations || inbound == EndpointVideos { return inbound } // OpenAI forwards everything to the Responses API. diff --git a/backend/internal/handler/endpoint_test.go b/backend/internal/handler/endpoint_test.go index 42b6d6e71b..55e3845ed4 100644 --- a/backend/internal/handler/endpoint_test.go +++ b/backend/internal/handler/endpoint_test.go @@ -28,6 +28,8 @@ func TestNormalizeInboundEndpoint(t *testing.T) { {"/v1/responses", EndpointResponses}, {"/v1/images/generations", EndpointImagesGenerations}, {"/v1/images/edits", EndpointImagesEdits}, + {"/v1/videos/generations", EndpointVideosGenerations}, + {"/v1/videos/req_123", EndpointVideos}, {"/v1beta/models", EndpointGeminiModels}, // Prefixed paths (antigravity, openai). @@ -81,6 +83,8 @@ func TestDeriveUpstreamEndpoint(t *testing.T) { {"openai embeddings", EndpointEmbeddings, "/v1/embeddings", service.PlatformOpenAI, EndpointEmbeddings}, {"openai image generations", EndpointImagesGenerations, "/v1/images/generations", service.PlatformOpenAI, EndpointImagesGenerations}, {"openai image edits", EndpointImagesEdits, "/openai/v1/images/edits", service.PlatformOpenAI, EndpointImagesEdits}, + {"grok video generations", EndpointVideosGenerations, "/v1/videos/generations", service.PlatformGrok, EndpointVideosGenerations}, + {"grok video status", EndpointVideos, "/videos/req_123", service.PlatformGrok, EndpointVideos}, // Antigravity — uses inbound to pick Claude vs Gemini upstream. {"antigravity claude", EndpointMessages, "/antigravity/v1/messages", service.PlatformAntigravity, EndpointMessages}, diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go new file mode 100644 index 0000000000..8e236ea49f --- /dev/null +++ b/backend/internal/handler/grok_media.go @@ -0,0 +1,377 @@ +package handler + +import ( + "context" + "errors" + "net/http" + "strconv" + "strings" + "time" + + pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" + "github.com/Wei-Shaw/sub2api/internal/pkg/ip" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "go.uber.org/zap" +) + +// GrokImages handles xAI image generation/editing through Grok groups. +func (h *OpenAIGatewayHandler) GrokImages(c *gin.Context) { + endpoint := service.GrokMediaEndpointImagesGenerations + if strings.Contains(c.Request.URL.Path, "/images/edits") { + endpoint = service.GrokMediaEndpointImagesEdits + } + h.handleGrokMedia(c, endpoint, "") +} + +// GrokVideoGeneration handles xAI video generation through Grok groups. +func (h *OpenAIGatewayHandler) GrokVideoGeneration(c *gin.Context) { + h.handleGrokMedia(c, service.GrokMediaEndpointVideosGenerations, "") +} + +// GrokVideoStatus handles xAI video status retrieval through Grok groups. +func (h *OpenAIGatewayHandler) GrokVideoStatus(c *gin.Context) { + h.handleGrokMedia(c, service.GrokMediaEndpointVideoStatus, c.Param("request_id")) +} + +func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.GrokMediaEndpoint, requestID string) { + streamStarted := false + defer h.recoverResponsesPanic(c, &streamStarted) + + requestStart := time.Now() + apiKey, ok := middleware2.GetAPIKeyFromContext(c) + if !ok { + h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key") + return + } + subject, ok := middleware2.GetAuthSubjectFromContext(c) + if !ok { + h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found") + return + } + + reqLog := requestLogger( + c, + "handler.openai_gateway.grok_media", + zap.Int64("user_id", subject.UserID), + zap.Int64("api_key_id", apiKey.ID), + zap.Any("group_id", apiKey.GroupID), + zap.String("endpoint", string(endpoint)), + ) + if !h.ensureResponsesDependencies(c, reqLog) { + return + } + + var body []byte + var err error + if endpoint.RequiresRequestBody() { + body, err = pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + if err != nil { + if maxErr, ok := extractMaxBytesError(err); ok { + h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) + return + } + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body") + return + } + if len(body) == 0 { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty") + return + } + } + + contentType := c.GetHeader("Content-Type") + requestInfo := service.ParseGrokMediaRequest(contentType, body) + requestModel := requestInfo.Model + if endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) == "" { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required") + return + } + if endpoint == service.GrokMediaEndpointVideoStatus && strings.TrimSpace(requestID) == "" { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "request_id is required") + return + } + + reqLog = reqLog.With(zap.String("model", requestModel)) + setOpsRequestContext(c, requestModel, false) + setOpsEndpointContext(c, "", int16(service.RequestTypeSync)) + + if endpoint.IsGenerationRequest() { + if !service.GroupAllowsImageGeneration(apiKey.Group) { + h.errorResponse(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage()) + return + } + if moderationBody := requestInfo.ModerationBody(); len(moderationBody) > 0 { + decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, requestModel, moderationBody) + if decision != nil && decision.Blocked { + h.errorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message) + return + } + } + imageReleaseFunc, acquired := h.acquireImageGenerationSlot(c, streamStarted) + if !acquired { + return + } + if imageReleaseFunc != nil { + defer imageReleaseFunc() + } + } + + if h.errorPassthroughService != nil { + service.BindErrorPassthroughService(c, h.errorPassthroughService) + } + + subscription, _ := middleware2.GetSubscriptionFromContext(c) + service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds()) + + userReleaseFunc, acquired := h.acquireResponsesUserSlot(c, subject.UserID, subject.Concurrency, false, &streamStarted, reqLog) + if !acquired { + return + } + if userReleaseFunc != nil { + defer userReleaseFunc() + } + + if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil { + reqLog.Info("grok_media.billing_eligibility_check_failed", zap.Error(err)) + status, code, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.errorResponse(c, status, code, message) + return + } + + sessionSeed := body + if len(sessionSeed) == 0 && strings.TrimSpace(requestID) != "" { + sessionSeed = []byte(requestID) + } + sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, sessionSeed) + if endpoint == service.GrokMediaEndpointVideoStatus { + sessionHash = service.GrokMediaVideoRequestSessionHash(requestID) + } + requestCtx := c.Request.Context() + failedAccountIDs := make(map[int64]struct{}) + sameAccountRetryCount := make(map[int64]int) + var lastFailoverErr *service.UpstreamFailoverError + switchCount := 0 + maxAccountSwitches := h.maxAccountSwitches + if maxAccountSwitches <= 0 { + maxAccountSwitches = 3 + } + routingStart := time.Now() + + for { + selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability( + requestCtx, + apiKey.GroupID, + "", + sessionHash, + requestModel, + failedAccountIDs, + service.OpenAIUpstreamTransportHTTPSSE, + "", + false, + service.PlatformGrok, + ) + if err != nil { + reqLog.Warn("grok_media.account_select_failed", + zap.Error(err), + zap.Int("excluded_account_count", len(failedAccountIDs)), + ) + if len(failedAccountIDs) == 0 { + cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok) + if !cls.ModelNotFound { + markOpsRoutingCapacityLimitedIfNoAvailable(c, err) + } + h.errorResponse(c, cls.Status, cls.ErrType, cls.Message) + return + } + if lastFailoverErr != nil { + h.handleFailoverExhausted(c, lastFailoverErr, false) + } else { + h.errorResponse(c, http.StatusBadGateway, "api_error", "Upstream request failed") + } + return + } + if selection == nil || selection.Account == nil { + cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok) + if !cls.ModelNotFound { + markOpsRoutingCapacityLimited(c) + } + h.errorResponse(c, cls.Status, cls.ErrType, cls.Message) + return + } + + reqLog.Debug("grok_media.account_schedule_decision", + zap.String("layer", scheduleDecision.Layer), + zap.Bool("sticky_session_hit", scheduleDecision.StickySessionHit), + zap.Int("candidate_count", scheduleDecision.CandidateCount), + zap.Int("top_k", scheduleDecision.TopK), + zap.Int64("latency_ms", scheduleDecision.LatencyMs), + zap.Float64("load_skew", scheduleDecision.LoadSkew), + ) + + account := selection.Account + sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account) + setOpsSelectedAccount(c, account.ID, account.Platform) + + accountReleaseFunc, accountAcquired := h.acquireResponsesAccountSlot(c, apiKey.GroupID, sessionHash, selection, false, &streamStarted, reqLog) + if !accountAcquired { + return + } + + service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds()) + forwardStart := time.Now() + writerSizeBeforeForward := c.Writer.Size() + result, err := func() (*service.OpenAIForwardResult, error) { + defer func() { + if accountReleaseFunc != nil { + accountReleaseFunc() + } + }() + return h.gatewayService.ForwardGrokMedia(requestCtx, c, account, endpoint, requestID, body, contentType) + }() + + forwardDurationMs := time.Since(forwardStart).Milliseconds() + upstreamLatencyMs, _ := getContextInt64(c, service.OpsUpstreamLatencyMsKey) + responseLatencyMs := forwardDurationMs + if upstreamLatencyMs > 0 && forwardDurationMs > upstreamLatencyMs { + responseLatencyMs = forwardDurationMs - upstreamLatencyMs + } + service.SetOpsLatencyMs(c, service.OpsResponseLatencyMsKey, responseLatencyMs) + + if err != nil { + var failoverErr *service.UpstreamFailoverError + if errors.As(err, &failoverErr) { + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + if c.Writer.Size() != writerSizeBeforeForward { + h.handleFailoverExhausted(c, failoverErr, true) + return + } + if failoverErr.RetryableOnSameAccount { + retryLimit := account.GetPoolModeRetryCount() + if sameAccountRetryCount[account.ID] < retryLimit { + sameAccountRetryCount[account.ID]++ + reqLog.Warn("grok_media.pool_mode_same_account_retry", + zap.Int64("account_id", account.ID), + zap.Int("upstream_status", failoverErr.StatusCode), + zap.Int("retry_limit", retryLimit), + zap.Int("retry_count", sameAccountRetryCount[account.ID]), + ) + select { + case <-requestCtx.Done(): + return + case <-time.After(sameAccountRetryDelay): + } + continue + } + } + h.gatewayService.RecordOpenAIAccountSwitch() + failedAccountIDs[account.ID] = struct{}{} + lastFailoverErr = failoverErr + if switchCount >= maxAccountSwitches { + h.handleFailoverExhausted(c, failoverErr, false) + return + } + switchCount++ + reqLog.Warn("grok_media.upstream_failover_switching", + zap.Int64("account_id", account.ID), + zap.Int("upstream_status", failoverErr.StatusCode), + zap.Int("switch_count", switchCount), + zap.Int("max_switches", maxAccountSwitches), + ) + continue + } + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + if c.Writer.Size() == writerSizeBeforeForward { + h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed") + } + reqLog.Warn("grok_media.forward_failed", + zap.Int64("account_id", account.ID), + zap.Error(err), + ) + return + } + + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + if endpoint == service.GrokMediaEndpointVideosGenerations && 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", + zap.Int64("account_id", account.ID), + zap.String("request_id", result.ResponseID), + zap.Error(err), + ) + } + } + if shouldRecordGrokMediaUsage(endpoint, requestModel) { + recordGrokMediaUsage(c, h, reqLog, apiKey, subject, subscription, account, result, requestModel, body, requestID) + } + reqLog.Debug("grok_media.request_completed", + zap.Int64("account_id", account.ID), + zap.Int("switch_count", switchCount), + ) + return + } +} + +func shouldRecordGrokMediaUsage(endpoint service.GrokMediaEndpoint, requestModel string) bool { + return endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) != "" +} + +func recordGrokMediaUsage( + c *gin.Context, + h *OpenAIGatewayHandler, + reqLog *zap.Logger, + apiKey *service.APIKey, + subject middleware2.AuthSubject, + subscription *service.UserSubscription, + account *service.Account, + result *service.OpenAIForwardResult, + requestModel string, + body []byte, + requestID string, +) { + userAgent := c.GetHeader("User-Agent") + clientIP := ip.GetClientIP(c) + payloadForHash := body + if len(payloadForHash) == 0 && strings.TrimSpace(requestID) != "" { + payloadForHash = []byte(requestID) + } + inboundEndpoint := GetInboundEndpoint(c) + upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + channelUsageFields := service.ChannelUsageFields{ + OriginalModel: requestModel, + ChannelMappedModel: requestModel, + } + h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) { + if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ + Result: result, + APIKey: apiKey, + User: apiKey.User, + Account: account, + Subscription: subscription, + InboundEndpoint: inboundEndpoint, + UpstreamEndpoint: upstreamEndpoint, + UserAgent: userAgent, + IPAddress: clientIP, + RequestPayloadHash: service.HashUsageRequestPayload(payloadForHash), + APIKeyService: h.apiKeyService, + QuotaPlatform: quotaPlatform, + ChannelUsageFields: channelUsageFields, + }); err != nil { + logger.L().With( + zap.String("component", "handler.openai_gateway.grok_media"), + zap.Int64("user_id", subject.UserID), + zap.Int64("api_key_id", apiKey.ID), + zap.Any("group_id", apiKey.GroupID), + zap.String("model", requestModel), + zap.Int64("account_id", account.ID), + ).Error("grok_media.record_usage_failed", zap.Error(err)) + reqLog.Debug("grok_media.record_usage_failed", zap.Error(err)) + } + }) +} diff --git a/backend/internal/handler/grok_media_test.go b/backend/internal/handler/grok_media_test.go new file mode 100644 index 0000000000..1b82f8fa6f --- /dev/null +++ b/backend/internal/handler/grok_media_test.go @@ -0,0 +1,54 @@ +package handler + +import ( + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestShouldRecordGrokMediaUsage(t *testing.T) { + tests := []struct { + name string + endpoint service.GrokMediaEndpoint + model string + want bool + }{ + { + name: "image generation records usage", + endpoint: service.GrokMediaEndpointImagesGenerations, + model: "grok-imagine", + want: true, + }, + { + name: "image edit records usage", + endpoint: service.GrokMediaEndpointImagesEdits, + model: "grok-imagine-edit", + want: true, + }, + { + name: "video generation records usage", + endpoint: service.GrokMediaEndpointVideosGenerations, + model: "grok-imagine-video-1.5", + want: true, + }, + { + name: "video status skips empty model usage", + endpoint: service.GrokMediaEndpointVideoStatus, + model: "", + want: false, + }, + { + name: "generation skips usage without model", + endpoint: service.GrokMediaEndpointImagesGenerations, + model: " ", + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.want, shouldRecordGrokMediaUsage(tt.endpoint, tt.model)) + }) + } +} diff --git a/backend/internal/pkg/xai/models.go b/backend/internal/pkg/xai/models.go index 0d289274fb..4902fcb94f 100644 --- a/backend/internal/pkg/xai/models.go +++ b/backend/internal/pkg/xai/models.go @@ -15,6 +15,12 @@ var defaultModels = []Model{ {ID: "grok-4.20-0309-reasoning", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Reasoning"}, {ID: "grok-4.20-0309-non-reasoning", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Non Reasoning"}, {ID: "grok-4.20-multi-agent-0309", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Multi Agent"}, + {ID: "grok-imagine", Object: "model", OwnedBy: "xai", DisplayName: "Grok Imagine"}, + {ID: "grok-imagine-image", Object: "model", OwnedBy: "xai", DisplayName: "Grok Imagine Image"}, + {ID: "grok-imagine-image-quality", Object: "model", OwnedBy: "xai", DisplayName: "Grok Imagine Image Quality"}, + {ID: "grok-imagine-edit", Object: "model", OwnedBy: "xai", DisplayName: "Grok Imagine Edit"}, + {ID: "grok-imagine-video", Object: "model", OwnedBy: "xai", DisplayName: "Grok Imagine Video"}, + {ID: "grok-imagine-video-1.5", Object: "model", OwnedBy: "xai", DisplayName: "Grok Imagine Video 1.5"}, } func DefaultModels() []Model { diff --git a/backend/internal/pkg/xai/oauth.go b/backend/internal/pkg/xai/oauth.go index 449b5cb865..1b3aadc08b 100644 --- a/backend/internal/pkg/xai/oauth.go +++ b/backend/internal/pkg/xai/oauth.go @@ -437,6 +437,42 @@ func BuildChatCompletionsURL(baseURL string) (string, error) { return validatedBaseURL + "/chat/completions", nil } +func BuildImagesGenerationsURL(baseURL string) (string, error) { + validatedBaseURL, err := ValidatedBaseURL(baseURL) + if err != nil { + return "", fmt.Errorf("invalid base url: %w", err) + } + return validatedBaseURL + "/images/generations", nil +} + +func BuildImagesEditsURL(baseURL string) (string, error) { + validatedBaseURL, err := ValidatedBaseURL(baseURL) + if err != nil { + return "", fmt.Errorf("invalid base url: %w", err) + } + return validatedBaseURL + "/images/edits", nil +} + +func BuildVideosGenerationsURL(baseURL string) (string, error) { + validatedBaseURL, err := ValidatedBaseURL(baseURL) + if err != nil { + return "", fmt.Errorf("invalid base url: %w", err) + } + return validatedBaseURL + "/videos/generations", nil +} + +func BuildVideoURL(baseURL, requestID string) (string, error) { + validatedBaseURL, err := ValidatedBaseURL(baseURL) + if err != nil { + return "", fmt.Errorf("invalid base url: %w", err) + } + requestID = strings.TrimSpace(requestID) + if requestID == "" { + return "", fmt.Errorf("request id is required") + } + return validatedBaseURL + "/videos/" + url.PathEscape(requestID), nil +} + // TokenResponse represents xAI OAuth token responses. type TokenResponse struct { AccessToken string `json:"access_token"` diff --git a/backend/internal/pkg/xai/oauth_test.go b/backend/internal/pkg/xai/oauth_test.go index fc48182fe7..68a4fea240 100644 --- a/backend/internal/pkg/xai/oauth_test.go +++ b/backend/internal/pkg/xai/oauth_test.go @@ -116,6 +116,27 @@ func TestValidateXAIURLsAllowOfficialOAuthAndGatewayHosts(t *testing.T) { require.Equal(t, DefaultCLIBaseURL+"/chat/completions", chatURL) } +func TestBuildGrokMediaURLs(t *testing.T) { + imagesURL, err := BuildImagesGenerationsURL(DefaultBaseURL + "/") + require.NoError(t, err) + require.Equal(t, DefaultBaseURL+"/images/generations", imagesURL) + + editsURL, err := BuildImagesEditsURL(DefaultBaseURL) + require.NoError(t, err) + require.Equal(t, DefaultBaseURL+"/images/edits", editsURL) + + videosURL, err := BuildVideosGenerationsURL(DefaultBaseURL) + require.NoError(t, err) + require.Equal(t, DefaultBaseURL+"/videos/generations", videosURL) + + videoURL, err := BuildVideoURL(DefaultBaseURL, "req 123") + require.NoError(t, err) + require.Equal(t, DefaultBaseURL+"/videos/req%20123", videoURL) + + _, err = BuildVideoURL(DefaultBaseURL, " ") + require.Error(t, err) +} + func TestValidateXAIURLsRejectArbitraryHostsByDefault(t *testing.T) { _, err := ValidateOAuthEndpointURL("https://auth.example.test/oauth2/token") require.Error(t, err) @@ -192,4 +213,10 @@ func TestDefaultModelMappingIncludesGrokAliases(t *testing.T) { require.Equal(t, "grok-4.20-0309-reasoning", mapping["grok-4.20-reasoning"]) require.Equal(t, "grok-4.20-0309-non-reasoning", mapping["grok-4.20-non-reasoning"]) require.Equal(t, "grok-4.20-multi-agent-0309", mapping["grok-4.20-multi-agent-0309"]) + require.Equal(t, "grok-imagine", mapping["grok-imagine"]) + require.Equal(t, "grok-imagine-image", mapping["grok-imagine-image"]) + require.Equal(t, "grok-imagine-image-quality", mapping["grok-imagine-image-quality"]) + require.Equal(t, "grok-imagine-edit", mapping["grok-imagine-edit"]) + require.Equal(t, "grok-imagine-video", mapping["grok-imagine-video"]) + require.Equal(t, "grok-imagine-video-1.5", mapping["grok-imagine-video-1.5"]) } diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index 54fb5c4d68..9522578051 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -42,6 +42,48 @@ func RegisterGatewayRoutes( isOpenAIGatewayPlatform := func(c *gin.Context) bool { return getGroupPlatform(c) == service.PlatformOpenAI } + imagesHandler := func(c *gin.Context) { + switch getGroupPlatform(c) { + case service.PlatformOpenAI: + h.OpenAIGateway.Images(c) + case service.PlatformGrok: + h.OpenAIGateway.GrokImages(c) + default: + service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) + c.JSON(http.StatusNotFound, gin.H{ + "error": gin.H{ + "type": "not_found_error", + "message": "Images API is not supported for this platform", + }, + }) + } + } + videoGenerationHandler := func(c *gin.Context) { + if getGroupPlatform(c) == service.PlatformGrok { + h.OpenAIGateway.GrokVideoGeneration(c) + return + } + service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) + c.JSON(http.StatusNotFound, gin.H{ + "error": gin.H{ + "type": "not_found_error", + "message": "Videos API is not supported for this platform", + }, + }) + } + videoStatusHandler := func(c *gin.Context) { + if getGroupPlatform(c) == service.PlatformGrok { + h.OpenAIGateway.GrokVideoStatus(c) + return + } + service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) + c.JSON(http.StatusNotFound, gin.H{ + "error": gin.H{ + "type": "not_found_error", + "message": "Videos API is not supported for this platform", + }, + }) + } // API网关（Claude API兼容） gateway := r.Group("/v1") gateway.Use(bodyLimit) @@ -120,32 +162,10 @@ func RegisterGatewayRoutes( } h.OpenAIGateway.Embeddings(c) }) - gateway.POST("/images/generations", func(c *gin.Context) { - if getGroupPlatform(c) != service.PlatformOpenAI { - service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) - c.JSON(http.StatusNotFound, gin.H{ - "error": gin.H{ - "type": "not_found_error", - "message": "Images API is not supported for this platform", - }, - }) - return - } - h.OpenAIGateway.Images(c) - }) - gateway.POST("/images/edits", func(c *gin.Context) { - if getGroupPlatform(c) != service.PlatformOpenAI { - service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) - c.JSON(http.StatusNotFound, gin.H{ - "error": gin.H{ - "type": "not_found_error", - "message": "Images API is not supported for this platform", - }, - }) - return - } - h.OpenAIGateway.Images(c) - }) + gateway.POST("/images/generations", imagesHandler) + gateway.POST("/images/edits", imagesHandler) + gateway.POST("/videos/generations", videoGenerationHandler) + gateway.GET("/videos/:request_id", videoStatusHandler) } // Gemini 原生 API 兼容层（Gemini SDK/CLI 直连） @@ -206,32 +226,10 @@ func RegisterGatewayRoutes( } h.OpenAIGateway.Embeddings(c) }) - r.POST("/images/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, func(c *gin.Context) { - if getGroupPlatform(c) != service.PlatformOpenAI { - service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) - c.JSON(http.StatusNotFound, gin.H{ - "error": gin.H{ - "type": "not_found_error", - "message": "Images API is not supported for this platform", - }, - }) - return - } - h.OpenAIGateway.Images(c) - }) - r.POST("/images/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, func(c *gin.Context) { - if getGroupPlatform(c) != service.PlatformOpenAI { - service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate) - c.JSON(http.StatusNotFound, gin.H{ - "error": gin.H{ - "type": "not_found_error", - "message": "Images API is not supported for this platform", - }, - }) - return - } - h.OpenAIGateway.Images(c) - }) + r.POST("/images/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, imagesHandler) + r.POST("/images/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, imagesHandler) + r.POST("/videos/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoGenerationHandler) + r.GET("/videos/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoStatusHandler) // Antigravity 模型列表 r.GET("/antigravity/models", gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.Gateway.AntigravityModels) diff --git a/backend/internal/server/routes/gateway_test.go b/backend/internal/server/routes/gateway_test.go index 0d8b6bd67a..2779dd8f01 100644 --- a/backend/internal/server/routes/gateway_test.go +++ b/backend/internal/server/routes/gateway_test.go @@ -83,6 +83,62 @@ func TestGatewayRoutesOpenAIImagesPathsAreRegistered(t *testing.T) { } } +func TestGatewayRoutesGrokImagesAndVideosPathsAreRegistered(t *testing.T) { + router := newGatewayRoutesTestRouter(service.PlatformGrok) + + for _, path := range []string{ + "/v1/images/generations", + "/v1/images/edits", + "/images/generations", + "/images/edits", + "/v1/videos/generations", + "/videos/generations", + } { + req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(`{"model":"grok-imagine","prompt":"draw a cat"}`)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + require.NotEqual(t, http.StatusNotFound, w.Code, "path=%s should hit Grok media handler", path) + require.NotContains(t, w.Body.String(), "not supported for this platform") + } + + for _, path := range []string{ + "/v1/videos/request-123", + "/videos/request-123", + } { + req := httptest.NewRequest(http.MethodGet, path, nil) + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + require.NotEqual(t, http.StatusNotFound, w.Code, "path=%s should hit Grok video handler", path) + require.NotContains(t, w.Body.String(), "not supported for this platform") + } +} + +func TestGatewayRoutesNonGrokVideosAreRejectedAtPlatformGate(t *testing.T) { + router := newGatewayRoutesTestRouter(service.PlatformOpenAI) + + for _, tc := range []struct { + method string + path string + body string + }{ + {http.MethodPost, "/v1/videos/generations", `{"model":"grok-imagine-video-1.5","prompt":"waves"}`}, + {http.MethodPost, "/videos/generations", `{"model":"grok-imagine-video-1.5","prompt":"waves"}`}, + {http.MethodGet, "/v1/videos/request-123", ""}, + {http.MethodGet, "/videos/request-123", ""}, + } { + req := httptest.NewRequest(tc.method, tc.path, strings.NewReader(tc.body)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + require.Equal(t, http.StatusNotFound, w.Code, "method=%s path=%s", tc.method, tc.path) + require.Contains(t, w.Body.String(), "Videos API is not supported for this platform") + } +} + func TestGatewayRoutesGrokAllowsCLICompatibilityEntrypoints(t *testing.T) { router := newGatewayRoutesTestRouter(service.PlatformGrok) diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go new file mode 100644 index 0000000000..fb05b64b01 --- /dev/null +++ b/backend/internal/service/grok_media.go @@ -0,0 +1,590 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "mime" + "mime/multipart" + "net/http" + "strconv" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +type GrokMediaEndpoint string + +const ( + GrokMediaEndpointImagesGenerations GrokMediaEndpoint = "images_generations" + GrokMediaEndpointImagesEdits GrokMediaEndpoint = "images_edits" + GrokMediaEndpointVideosGenerations GrokMediaEndpoint = "videos_generations" + GrokMediaEndpointVideoStatus GrokMediaEndpoint = "video_status" +) + +func (e GrokMediaEndpoint) RequiresRequestBody() bool { + return e != GrokMediaEndpointVideoStatus +} + +func (e GrokMediaEndpoint) IsGenerationRequest() bool { + switch e { + case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits, GrokMediaEndpointVideosGenerations: + return true + default: + return false + } +} + +type GrokMediaRequestInfo struct { + Model string + Prompt string + N int + Size string + SizeTier string + InputImageURLs []string + MaskImageURL string + Uploads []OpenAIImagesUpload + MaskUpload *OpenAIImagesUpload +} + +func (r GrokMediaRequestInfo) ModerationBody() []byte { + payload := map[string]any{} + if prompt := strings.TrimSpace(r.Prompt); prompt != "" { + payload["prompt"] = prompt + } + + images := make([]map[string]string, 0, len(r.InputImageURLs)+len(r.Uploads)+1) + for _, imageURL := range r.InputImageURLs { + if imageURL = strings.TrimSpace(imageURL); imageURL != "" { + images = append(images, map[string]string{"image_url": imageURL}) + } + } + for _, upload := range r.Uploads { + if dataURL := upload.ModerationDataURL(); dataURL != "" { + images = append(images, map[string]string{"image_url": dataURL}) + } + } + if maskURL := strings.TrimSpace(r.MaskImageURL); maskURL != "" { + images = append(images, map[string]string{"image_url": maskURL}) + } + if r.MaskUpload != nil { + if dataURL := r.MaskUpload.ModerationDataURL(); dataURL != "" { + images = append(images, map[string]string{"image_url": dataURL}) + } + } + if len(images) > 0 { + payload["images"] = images + } + if len(payload) == 0 { + return nil + } + body, err := json.Marshal(payload) + if err != nil { + return nil + } + return body +} + +func (e GrokMediaEndpoint) httpMethod() string { + if e == GrokMediaEndpointVideoStatus { + return http.MethodGet + } + return http.MethodPost +} + +func ExtractGrokMediaModel(contentType string, body []byte) string { + return ParseGrokMediaRequest(contentType, body).Model +} + +func ParseGrokMediaRequest(contentType string, body []byte) GrokMediaRequestInfo { + info := GrokMediaRequestInfo{N: 1} + if gjson.ValidBytes(body) { + parseGrokMediaJSONRequest(body, &info) + } else { + parseGrokMediaMultipartRequest(contentType, body, &info) + } + info.Model = strings.TrimSpace(info.Model) + info.Prompt = strings.TrimSpace(info.Prompt) + info.Size = strings.TrimSpace(info.Size) + info.SizeTier = NormalizeImageBillingTierOrDefault(info.Size) + if info.N <= 0 { + info.N = 1 + } + return info +} + +func parseGrokMediaJSONRequest(body []byte, info *GrokMediaRequestInfo) { + if info == nil { + return + } + info.Model = strings.TrimSpace(gjson.GetBytes(body, "model").String()) + info.Prompt = strings.TrimSpace(gjson.GetBytes(body, "prompt").String()) + info.Size = strings.TrimSpace(gjson.GetBytes(body, "size").String()) + if n := gjson.GetBytes(body, "n"); n.Exists() && n.Type == gjson.Number { + info.N = int(n.Int()) + } + appendJSONImageURLs := func(value gjson.Result) { + if !value.Exists() { + return + } + switch { + case value.IsArray(): + for _, item := range value.Array() { + if imageURL := strings.TrimSpace(item.Get("image_url").String()); imageURL != "" { + info.InputImageURLs = append(info.InputImageURLs, imageURL) + continue + } + if item.Type == gjson.String { + imageURL := strings.TrimSpace(item.String()) + if imageURL == "" { + continue + } + info.InputImageURLs = append(info.InputImageURLs, imageURL) + } + } + default: + if imageURL := strings.TrimSpace(value.Get("image_url").String()); imageURL != "" { + info.InputImageURLs = append(info.InputImageURLs, imageURL) + return + } + if value.Type == gjson.String { + imageURL := strings.TrimSpace(value.String()) + if imageURL == "" { + return + } + info.InputImageURLs = append(info.InputImageURLs, imageURL) + } + } + } + appendJSONImageURLs(gjson.GetBytes(body, "image")) + appendJSONImageURLs(gjson.GetBytes(body, "images")) + info.MaskImageURL = strings.TrimSpace(gjson.GetBytes(body, "mask.image_url").String()) +} + +func parseGrokMediaMultipartRequest(contentType string, body []byte, info *GrokMediaRequestInfo) { + if info == nil { + return + } + 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 err == io.EOF { + return + } + if err != nil { + return + } + name := strings.TrimSpace(part.FormName()) + if name == "" { + _ = part.Close() + continue + } + data, err := io.ReadAll(io.LimitReader(part, openAIImageMaxUploadPartSize)) + _ = part.Close() + if err != nil { + return + } + fileName := strings.TrimSpace(part.FileName()) + partContentType := strings.TrimSpace(part.Header.Get("Content-Type")) + if fileName != "" { + upload := OpenAIImagesUpload{ + FieldName: name, + FileName: fileName, + ContentType: partContentType, + Data: data, + } + if name == "mask" { + info.MaskUpload = &upload + continue + } + if name == "image" || strings.HasPrefix(name, "image[") { + info.Uploads = append(info.Uploads, upload) + } + continue + } + + value := strings.TrimSpace(string(data)) + switch name { + case "model": + info.Model = value + case "prompt": + info.Prompt = value + case "size": + info.Size = value + case "n": + if n, err := strconv.Atoi(value); err == nil { + info.N = n + } + case "image", "image_url": + if value != "" { + info.InputImageURLs = append(info.InputImageURLs, value) + } + case "mask", "mask_image_url": + info.MaskImageURL = value + } + } +} + +func GrokMediaVideoRequestSessionHash(requestID string) string { + requestID = strings.TrimSpace(requestID) + if requestID == "" { + return "" + } + return "grok-video:" + DeriveSessionHashFromSeed(requestID) +} + +func (s *OpenAIGatewayService) BindGrokMediaVideoRequestAccount(ctx context.Context, groupID *int64, requestID string, accountID int64) error { + return s.BindStickySession(ctx, groupID, GrokMediaVideoRequestSessionHash(requestID), accountID) +} + +func (e GrokMediaEndpoint) upstreamURL(baseURL, requestID string) (string, error) { + switch e { + case GrokMediaEndpointImagesGenerations: + return xai.BuildImagesGenerationsURL(baseURL) + case GrokMediaEndpointImagesEdits: + return xai.BuildImagesEditsURL(baseURL) + case GrokMediaEndpointVideosGenerations: + return xai.BuildVideosGenerationsURL(baseURL) + case GrokMediaEndpointVideoStatus: + return xai.BuildVideoURL(baseURL, requestID) + default: + return "", fmt.Errorf("unsupported grok media endpoint: %s", e) + } +} + +func (s *OpenAIGatewayService) ForwardGrokMedia( + ctx context.Context, + c *gin.Context, + account *Account, + endpoint GrokMediaEndpoint, + requestID string, + body []byte, + contentType string, +) (*OpenAIForwardResult, error) { + startTime := time.Now() + if account == nil { + return nil, fmt.Errorf("grok account is required") + } + if account.Platform != PlatformGrok { + return nil, fmt.Errorf("account platform %s is not supported for grok media", account.Platform) + } + + token, _, err := s.GetAccessToken(ctx, account) + if err != nil { + return nil, err + } + targetURL, err := endpoint.upstreamURL(account.GetGrokBaseURL(), requestID) + if err != nil { + return nil, err + } + + body, contentType, err = prepareGrokMediaForwardBody(endpoint, body, contentType) + if err != nil { + return nil, err + } + + var bodyReader io.Reader + if endpoint.RequiresRequestBody() { + bodyReader = bytes.NewReader(body) + } + upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) + defer releaseUpstreamCtx() + upstreamReq, err := http.NewRequestWithContext(upstreamCtx, endpoint.httpMethod(), targetURL, bodyReader) + if err != nil { + return nil, err + } + upstreamReq.Header.Set("Authorization", "Bearer "+token) + upstreamReq.Header.Set("Accept", "application/json") + upstreamReq.Header.Set("User-Agent", "sub2api-grok/1.0") + if endpoint.RequiresRequestBody() { + contentType = strings.TrimSpace(contentType) + if contentType == "" { + contentType = "application/json" + } + upstreamReq.Header.Set("Content-Type", contentType) + } + + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + upstreamStart := time.Now() + resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + if err != nil { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) + } + defer func() { _ = resp.Body.Close() }() + + requestIDHeader := firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")) + requestInfo := ParseGrokMediaRequest(contentType, body) + requestModel := requestInfo.Model + if resp.StatusCode >= 400 { + s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + return s.handleGrokMediaErrorResponse(ctx, resp, c, account, requestIDHeader, requestModel) + } + + s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError) + if err != nil { + return nil, err + } + writeGrokMediaResponse(c, resp, respBody, s.responseHeaderFilter) + usage := grokMediaUsageFromResponse(endpoint, requestInfo, respBody) + return &OpenAIForwardResult{ + RequestID: requestIDHeader, + ResponseID: usage.ResponseID, + Usage: usage.Usage, + Model: requestModel, + BillingModel: requestModel, + UpstreamModel: requestModel, + ResponseHeaders: resp.Header.Clone(), + Duration: time.Since(startTime), + ImageCount: usage.ImageCount, + ImageSize: usage.ImageSize, + ImageInputSize: usage.ImageInputSize, + ImageOutputSizes: usage.ImageOutputSizes, + }, nil +} + +func prepareGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, contentType string) ([]byte, string, error) { + if endpoint != GrokMediaEndpointImagesEdits || gjson.ValidBytes(body) { + return body, contentType, nil + } + mediaType, _, err := mime.ParseMediaType(strings.TrimSpace(contentType)) + if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") { + return body, contentType, nil + } + + info := ParseGrokMediaRequest(contentType, body) + payload := make(map[string]any) + if info.Model != "" { + payload["model"] = info.Model + } + if info.Prompt != "" { + payload["prompt"] = info.Prompt + } + if info.N > 1 { + payload["n"] = info.N + } + if info.Size != "" { + payload["size"] = info.Size + } + + images := make([]map[string]string, 0, len(info.InputImageURLs)+len(info.Uploads)) + for _, imageURL := range info.InputImageURLs { + if imageURL = strings.TrimSpace(imageURL); imageURL != "" { + images = append(images, map[string]string{"image_url": imageURL}) + } + } + for _, upload := range info.Uploads { + dataURL, err := openAIImageUploadToDataURL(upload) + if err != nil { + return nil, "", err + } + images = append(images, map[string]string{"image_url": dataURL}) + } + if len(images) > 0 { + payload["image"] = images[0] + if len(images) > 1 { + payload["images"] = images + } + } + + maskImageURL := strings.TrimSpace(info.MaskImageURL) + if info.MaskUpload != nil { + dataURL, err := openAIImageUploadToDataURL(*info.MaskUpload) + if err != nil { + return nil, "", err + } + maskImageURL = dataURL + } + if maskImageURL != "" { + payload["mask"] = map[string]string{"image_url": maskImageURL} + } + + out, err := marshalOpenAIUpstreamJSON(payload) + if err != nil { + return nil, "", err + } + return out, "application/json", nil +} + +type grokMediaUsageMetadata struct { + ResponseID string + Usage OpenAIUsage + ImageCount int + ImageSize string + ImageInputSize string + ImageOutputSizes []string +} + +func grokMediaUsageFromResponse(endpoint GrokMediaEndpoint, requestInfo GrokMediaRequestInfo, responseBody []byte) grokMediaUsageMetadata { + usage, _ := extractOpenAIUsageFromJSONBytes(responseBody) + meta := grokMediaUsageMetadata{Usage: usage} + switch endpoint { + case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits: + imageCount := countOpenAIResponseImageOutputsFromJSONBytes(responseBody) + if imageCount <= 0 { + imageCount = requestInfo.N + } + if imageCount <= 0 { + imageCount = 1 + } + meta.ImageCount = imageCount + meta.ImageSize = requestInfo.SizeTier + meta.ImageInputSize = requestInfo.Size + meta.ImageOutputSizes = collectOpenAIResponseImageOutputSizesFromJSONBytes(responseBody) + case GrokMediaEndpointVideosGenerations: + meta.ResponseID = extractGrokMediaVideoRequestID(responseBody) + meta.ImageCount = 1 + meta.ImageSize = requestInfo.SizeTier + meta.ImageInputSize = requestInfo.Size + } + return meta +} + +func extractGrokMediaVideoRequestID(body []byte) string { + if len(body) == 0 || !gjson.ValidBytes(body) { + return "" + } + for _, path := range []string{"request_id", "id", "data.request_id", "data.id", "video.request_id", "video.id"} { + if id := strings.TrimSpace(gjson.GetBytes(body, path).String()); id != "" { + return id + } + } + return "" +} + +func (s *OpenAIGatewayService) handleGrokMediaErrorResponse( + ctx context.Context, + resp *http.Response, + c *gin.Context, + account *Account, + requestIDHeader string, + requestedModel string, +) (*OpenAIForwardResult, error) { + body := s.readUpstreamErrorBody(resp) + upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(body))) + if upstreamMsg == "" { + upstreamMsg = fmt.Sprintf("xAI upstream returned status %d", resp.StatusCode) + } + + upstreamDetail := "" + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes + if maxBytes <= 0 { + maxBytes = 2048 + } + upstreamDetail = truncateString(string(body), maxBytes) + } + setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) + + if status, errType, errMsg, matched := applyErrorPassthroughRule( + c, + account.Platform, + resp.StatusCode, + body, + http.StatusBadGateway, + "upstream_error", + "Upstream request failed", + ); matched { + MarkResponseCommitted(c) + writeGrokMediaErrorResponse(c, status, errType, errMsg) + return nil, fmt.Errorf("upstream error: %d (passthrough rule matched) message=%s", resp.StatusCode, upstreamMsg) + } + + if !account.ShouldHandleErrorCode(resp.StatusCode) { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: requestIDHeader, + Kind: "http_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + MarkResponseCommitted(c) + writeGrokMediaErrorResponse(c, http.StatusInternalServerError, "upstream_error", "Upstream gateway error") + return nil, fmt.Errorf("upstream error: %d (not in custom error codes) message=%s", resp.StatusCode, upstreamMsg) + } + + s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body) + kind := "http_error" + if s.shouldFailoverUpstreamError(resp.StatusCode) { + kind = "failover" + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: requestIDHeader, + Kind: kind, + Message: upstreamMsg, + Detail: upstreamDetail, + }) + if kind == "failover" { + return nil, &UpstreamFailoverError{ + StatusCode: resp.StatusCode, + ResponseBody: body, + RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + } + } + + MarkResponseCommitted(c) + writeGrokMediaErrorResponse(c, resp.StatusCode, grokMediaErrorType(resp.StatusCode), upstreamMsg) + return nil, fmt.Errorf("upstream error: %d %s", resp.StatusCode, upstreamMsg) +} + +func grokMediaErrorType(statusCode int) string { + switch statusCode { + case http.StatusBadRequest: + return "invalid_request_error" + case http.StatusNotFound: + return "not_found_error" + case http.StatusTooManyRequests: + return "rate_limit_error" + default: + return "upstream_error" + } +} + +func writeGrokMediaErrorResponse(c *gin.Context, statusCode int, errType, message string) { + if c == nil || c.Writer == nil || c.Writer.Written() { + return + } + c.JSON(statusCode, gin.H{ + "error": gin.H{ + "type": strings.TrimSpace(errType), + "message": strings.TrimSpace(message), + }, + }) +} + +func writeGrokMediaResponse(c *gin.Context, resp *http.Response, body []byte, filter *responseheaders.CompiledHeaderFilter) { + if c == nil || resp == nil { + return + } + writeOpenAIPassthroughResponseHeaders(c.Writer.Header(), resp.Header, filter) + contentType := strings.TrimSpace(resp.Header.Get("Content-Type")) + if contentType == "" { + contentType = "application/json" + } + c.Data(resp.StatusCode, contentType, body) +} diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 4f6601e843..ad088be725 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -7,8 +7,10 @@ import ( "context" "encoding/json" "io" + "mime/multipart" "net/http" "net/http/httptest" + "net/textproto" "strings" "testing" "time" @@ -144,6 +146,279 @@ func TestBuildGrokResponsesRequestRejectsUnsafeAccountBaseURL(t *testing.T) { require.Contains(t, err.Error(), "invalid base url") } +func TestExtractGrokMediaModelSupportsJSONAndMultipart(t *testing.T) { + require.Equal(t, "grok-imagine", ExtractGrokMediaModel("application/json", []byte(`{"model":"grok-imagine"}`))) + + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + require.NoError(t, writer.WriteField("prompt", "draw a cat")) + require.NoError(t, writer.WriteField("model", "grok-imagine-edit")) + require.NoError(t, writer.Close()) + + require.Equal(t, "grok-imagine-edit", ExtractGrokMediaModel(writer.FormDataContentType(), buf.Bytes())) +} + +func TestParseGrokMediaRequestBuildsMultipartModerationBody(t *testing.T) { + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + require.NoError(t, writer.WriteField("prompt", "edit this private image")) + require.NoError(t, writer.WriteField("model", "grok-imagine-edit")) + partHeader := textproto.MIMEHeader{} + partHeader.Set("Content-Disposition", `form-data; name="image"; filename="input.png"`) + partHeader.Set("Content-Type", "image/png") + part, err := writer.CreatePart(partHeader) + require.NoError(t, err) + _, err = part.Write([]byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a}) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + info := ParseGrokMediaRequest(writer.FormDataContentType(), buf.Bytes()) + require.Equal(t, "grok-imagine-edit", info.Model) + require.Equal(t, "edit this private image", info.Prompt) + + moderationBody := info.ModerationBody() + require.NotEmpty(t, moderationBody) + require.Equal(t, "edit this private image", gjson.GetBytes(moderationBody, "prompt").String()) + require.True(t, strings.HasPrefix(gjson.GetBytes(moderationBody, "images.0.image_url").String(), "data:image/")) +} + +func TestForwardGrokMediaImagesGenerationPassthrough(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok-imagine","prompt":"draw a cat"}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + account := &Account{ + ID: 61, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "api-key", + "base_url": "https://xai.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{"xai-image-req"}, + }, + Body: io.NopCloser(strings.NewReader(`{"data":[]}`)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json") + require.NoError(t, err) + require.Equal(t, "https://xai.test/v1/images/generations", upstream.lastReq.URL.String()) + require.Equal(t, http.MethodPost, upstream.lastReq.Method) + require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization")) + require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type")) + require.JSONEq(t, string(body), string(upstream.lastBody)) + require.Equal(t, http.StatusOK, recorder.Code) + require.JSONEq(t, `{"data":[]}`, recorder.Body.String()) + require.Equal(t, "xai-image-req", result.RequestID) + require.Equal(t, "grok-imagine", result.Model) + require.Equal(t, "grok-imagine", result.BillingModel) + require.Equal(t, 1, result.ImageCount) + require.Equal(t, ImageBillingSize2K, result.ImageSize) +} + +func TestForwardGrokMediaImagesEditMultipartConvertsToJSON(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + gin.SetMode(gin.TestMode) + + var buf bytes.Buffer + writer := multipart.NewWriter(&buf) + require.NoError(t, writer.WriteField("model", "grok-imagine-edit")) + require.NoError(t, writer.WriteField("prompt", "edit this private image")) + partHeader := textproto.MIMEHeader{} + partHeader.Set("Content-Disposition", `form-data; name="image"; filename="input.png"`) + partHeader.Set("Content-Type", "image/png") + part, err := writer.CreatePart(partHeader) + require.NoError(t, err) + _, err = part.Write([]byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a}) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(buf.Bytes())) + c.Request.Header.Set("Content-Type", writer.FormDataContentType()) + + account := &Account{ + ID: 62, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "api-key", + "base_url": "https://xai.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + }, + Body: io.NopCloser(strings.NewReader(`{"data":[]}`)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + _, 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, "edit this private image", gjson.GetBytes(upstream.lastBody, "prompt").String()) + require.True(t, strings.HasPrefix(gjson.GetBytes(upstream.lastBody, "image.image_url").String(), "data:image/png;base64,")) +} + +func TestForwardGrokMediaVideoGenerationReturnsUsageAndResponseID(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok-imagine-video-1.5","prompt":"waves"}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + account := &Account{ + ID: 63, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "api-key", + "base_url": "https://xai.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{"xai-video-generate-req"}, + }, + Body: io.NopCloser(strings.NewReader(`{"request_id":"video-request-123","usage":{"prompt_tokens":3,"completion_tokens":4}}`)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json") + require.NoError(t, err) + require.Equal(t, "https://xai.test/v1/videos/generations", upstream.lastReq.URL.String()) + require.Equal(t, "video-request-123", result.ResponseID) + require.Equal(t, "grok-imagine-video-1.5", result.BillingModel) + require.Equal(t, 3, result.Usage.InputTokens) + require.Equal(t, 4, result.Usage.OutputTokens) + require.Equal(t, 1, result.ImageCount) +} + +func TestForwardGrokMediaVideoStatusUsesGETWithoutBody(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/videos/request-123", nil) + + account := &Account{ + ID: 62, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "api-key", + "base_url": "https://xai.test/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{"xai-video-req"}, + }, + Body: io.NopCloser(strings.NewReader(`{"id":"request-123","status":"completed"}`)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideoStatus, "request-123", nil, "") + require.NoError(t, err) + require.Equal(t, "https://xai.test/v1/videos/request-123", upstream.lastReq.URL.String()) + require.Equal(t, http.MethodGet, upstream.lastReq.Method) + require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization")) + require.Empty(t, upstream.lastReq.Header.Get("Content-Type")) + require.Empty(t, upstream.lastBody) + require.Equal(t, http.StatusOK, recorder.Code) + require.JSONEq(t, `{"id":"request-123","status":"completed"}`, recorder.Body.String()) + require.Equal(t, "xai-video-req", result.RequestID) +} + +func TestBindGrokMediaVideoRequestAccountUsesRequestIDStickyHash(t *testing.T) { + ctx := context.Background() + groupID := int64(7) + cache := &stubGatewayCache{} + svc := &OpenAIGatewayService{cache: cache} + + hash := GrokMediaVideoRequestSessionHash("video-request-123") + require.NotEmpty(t, hash) + require.NoError(t, svc.BindGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", 63)) + + accountID, err := svc.getStickySessionAccountID(ctx, &groupID, hash) + require.NoError(t, err) + require.Equal(t, int64(63), accountID) +} + +func TestForwardGrokMediaErrorHonorsCustomErrorCodes(t *testing.T) { + t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok-imagine","prompt":"draw a cat"}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + account := &Account{ + ID: 64, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "api-key", + "base_url": "https://xai.test/v1", + "custom_error_codes_enabled": true, + "custom_error_codes": []any{float64(http.StatusTooManyRequests)}, + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{"xai-error-req"}, + }, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"do not expose this upstream detail"}}`)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json") + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, http.StatusInternalServerError, recorder.Code) + require.Contains(t, recorder.Body.String(), "Upstream gateway error") + require.NotContains(t, recorder.Body.String(), "do not expose") +} + func TestForwardAsChatCompletionsForGrokUsesXAIChatCompletionsAndSnapshots(t *testing.T) { gin.SetMode(gin.TestMode)