diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go index 6ae7129451..f51167c6e6 100644 --- a/backend/internal/handler/grok_media.go +++ b/backend/internal/handler/grok_media.go @@ -46,6 +46,11 @@ func (h *OpenAIGatewayHandler) GrokVideoStatus(c *gin.Context) { h.handleGrokMedia(c, service.GrokMediaEndpointVideoStatus, c.Param("request_id")) } +// GrokVideoContent proxies downloadable video content through the task's upstream account. +func (h *OpenAIGatewayHandler) GrokVideoContent(c *gin.Context) { + h.handleGrokMedia(c, service.GrokMediaEndpointVideoContent, c.Param("request_id")) +} + func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.GrokMediaEndpoint, requestID string) { streamStarted := false defer h.recoverResponsesPanic(c, &streamStarted) @@ -99,7 +104,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required") return } - if endpoint == service.GrokMediaEndpointVideoStatus && strings.TrimSpace(requestID) == "" { + if endpoint.IsVideoLookupRequest() && strings.TrimSpace(requestID) == "" { h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "request_id is required") return } @@ -159,7 +164,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. sessionSeed = []byte(requestID) } sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, sessionSeed) - if endpoint == service.GrokMediaEndpointVideoStatus { + if endpoint.IsVideoLookupRequest() { sessionHash = service.GrokMediaVideoRequestSessionHash(requestID) } requestCtx := c.Request.Context() @@ -334,7 +339,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service. continue } h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil) - if c.Writer.Size() == writerSizeBeforeForward { + if !service.IsResponseCommitted(c) && c.Writer.Size() == writerSizeBeforeForward { h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed") } reqLog.Warn("grok_media.forward_failed", diff --git a/backend/internal/handler/grok_media_test.go b/backend/internal/handler/grok_media_test.go index 3c58742c78..4e94665141 100644 --- a/backend/internal/handler/grok_media_test.go +++ b/backend/internal/handler/grok_media_test.go @@ -38,6 +38,12 @@ func TestShouldRecordGrokMediaUsage(t *testing.T) { model: "", want: false, }, + { + name: "video content skips usage", + endpoint: service.GrokMediaEndpointVideoContent, + model: "", + want: false, + }, { name: "generation skips usage without model", endpoint: service.GrokMediaEndpointImagesGenerations, @@ -65,6 +71,7 @@ func TestGrokMediaRequiredCapability(t *testing.T) { {name: "video edit", endpoint: service.GrokMediaEndpointVideosEdits, want: service.OpenAIEndpointCapabilityGrokMediaGeneration}, {name: "video extension", endpoint: service.GrokMediaEndpointVideosExtensions, want: service.OpenAIEndpointCapabilityGrokMediaGeneration}, {name: "video status preserves lookup", endpoint: service.GrokMediaEndpointVideoStatus, want: ""}, + {name: "video content preserves lookup", endpoint: service.GrokMediaEndpointVideoContent, want: ""}, } for _, tt := range tests { diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index a1e1de260d..6d091a079a 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -91,6 +91,19 @@ func RegisterGatewayRoutes( }, }) } + videoContentHandler := func(c *gin.Context) { + if getGroupPlatform(c) == service.PlatformGrok { + h.OpenAIGateway.GrokVideoContent(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", + }, + }) + } videoEditHandler := func(c *gin.Context) { if getGroupPlatform(c) == service.PlatformGrok { h.OpenAIGateway.GrokVideoEdit(c) @@ -209,6 +222,7 @@ func RegisterGatewayRoutes( gateway.POST("/videos/edits", videoEditHandler) gateway.POST("/videos/extensions", videoExtensionHandler) gateway.GET("/videos/:request_id", videoStatusHandler) + gateway.GET("/videos/:request_id/content", videoContentHandler) } // Gemini 原生 API 兼容层(Gemini SDK/CLI 直连) @@ -282,6 +296,7 @@ func RegisterGatewayRoutes( r.POST("/videos/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoEditHandler) r.POST("/videos/extensions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoExtensionHandler) r.GET("/videos/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoStatusHandler) + r.GET("/videos/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, videoContentHandler) // 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 f8abdb4062..80e319db76 100644 --- a/backend/internal/server/routes/gateway_test.go +++ b/backend/internal/server/routes/gateway_test.go @@ -160,6 +160,8 @@ func TestGatewayRoutesGrokImagesAndVideosPathsAreRegistered(t *testing.T) { for _, path := range []string{ "/v1/videos/request-123", "/videos/request-123", + "/v1/videos/request-123/content", + "/videos/request-123/content", } { req := httptest.NewRequest(http.MethodGet, path, nil) w := httptest.NewRecorder() @@ -186,6 +188,8 @@ func TestGatewayRoutesNonGrokVideosAreRejectedAtPlatformGate(t *testing.T) { {http.MethodPost, "/videos/extensions", `{"model":"grok-imagine-video","prompt":"waves","video":{"url":"https://example.com/in.mp4"}}`}, {http.MethodGet, "/v1/videos/request-123", ""}, {http.MethodGet, "/videos/request-123", ""}, + {http.MethodGet, "/v1/videos/request-123/content", ""}, + {http.MethodGet, "/videos/request-123/content", ""}, } { req := httptest.NewRequest(tc.method, tc.path, strings.NewReader(tc.body)) req.Header.Set("Content-Type", "application/json") diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index c75bd271f5..585fd50427 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -29,10 +29,15 @@ const ( GrokMediaEndpointVideosEdits GrokMediaEndpoint = "videos_edits" GrokMediaEndpointVideosExtensions GrokMediaEndpoint = "videos_extensions" GrokMediaEndpointVideoStatus GrokMediaEndpoint = "video_status" + GrokMediaEndpointVideoContent GrokMediaEndpoint = "video_content" ) func (e GrokMediaEndpoint) RequiresRequestBody() bool { - return e != GrokMediaEndpointVideoStatus + return !e.IsVideoLookupRequest() +} + +func (e GrokMediaEndpoint) IsVideoLookupRequest() bool { + return e == GrokMediaEndpointVideoStatus || e == GrokMediaEndpointVideoContent } func (e GrokMediaEndpoint) IsGenerationRequest() bool { @@ -97,7 +102,7 @@ func (r GrokMediaRequestInfo) ModerationBody() []byte { } func (e GrokMediaEndpoint) httpMethod() string { - if e == GrokMediaEndpointVideoStatus { + if e.IsVideoLookupRequest() { return http.MethodGet } return http.MethodPost @@ -327,7 +332,16 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( return nil, err } upstreamReq.Header.Set("Authorization", "Bearer "+token) - upstreamReq.Header.Set("Accept", "application/json") + if endpoint == GrokMediaEndpointVideoContent { + upstreamReq.Header.Set("Accept", "*/*") + if c != nil { + if rangeHeader := strings.TrimSpace(c.GetHeader("Range")); rangeHeader != "" { + upstreamReq.Header.Set("Range", rangeHeader) + } + } + } else { + upstreamReq.Header.Set("Accept", "application/json") + } if account.IsGrokOAuth() && isGrokCLIProxyTarget(targetURL) { applyGrokCLIHeaders(upstreamReq.Header) } @@ -360,10 +374,27 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( } s.updateGrokUsageFromResponse(ctx, account, resp.Header, resp.StatusCode) + if endpoint == GrokMediaEndpointVideoContent { + if err := writeGrokMediaContentResponse(c, resp); err != nil { + return nil, err + } + return &OpenAIForwardResult{ + RequestID: requestIDHeader, + ResponseHeaders: resp.Header.Clone(), + Duration: time.Since(startTime), + }, nil + } respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError) if err != nil { return nil, err } + if endpoint == GrokMediaEndpointVideoStatus { + respBody = rewriteGrokMediaVideoContentURLs( + respBody, + requestID, + grokMediaContentProxyURL(c, requestID), + ) + } writeGrokMediaResponse(c, resp, respBody, s.responseHeaderFilter) usage := grokMediaUsageFromResponse(endpoint, requestInfo, respBody) return &OpenAIForwardResult{ @@ -739,3 +770,144 @@ func writeGrokMediaResponse(c *gin.Context, resp *http.Response, body []byte, fi } c.Data(resp.StatusCode, contentType, body) } + +func writeGrokMediaContentResponse(c *gin.Context, resp *http.Response) error { + if c == nil || resp == nil || resp.Body == nil { + return fmt.Errorf("grok media content response is incomplete") + } + + for _, name := range []string{ + "Content-Type", + "Content-Length", + "Content-Range", + "Accept-Ranges", + "Content-Disposition", + } { + if value := strings.TrimSpace(resp.Header.Get(name)); value != "" { + c.Header(name, value) + } + } + if strings.TrimSpace(c.Writer.Header().Get("Content-Length")) == "" && resp.ContentLength >= 0 { + c.Header("Content-Length", strconv.FormatInt(resp.ContentLength, 10)) + } + if strings.TrimSpace(c.Writer.Header().Get("Content-Type")) == "" { + c.Header("Content-Type", "application/octet-stream") + } + c.Status(resp.StatusCode) + MarkResponseCommitted(c) + _, err := io.Copy(c.Writer, resp.Body) + return err +} + +func rewriteGrokMediaVideoContentURLs(body []byte, requestID, proxyURL string) []byte { + if len(body) == 0 || strings.TrimSpace(requestID) == "" || strings.TrimSpace(proxyURL) == "" || !gjson.ValidBytes(body) { + return body + } + + decoder := json.NewDecoder(bytes.NewReader(body)) + decoder.UseNumber() + var value any + if err := decoder.Decode(&value); err != nil { + return body + } + if !rewriteGrokMediaVideoContentURLValue(&value, requestID, proxyURL) { + return body + } + rewritten, err := json.Marshal(value) + if err != nil { + return body + } + return rewritten +} + +func rewriteGrokMediaVideoContentURLValue(value *any, requestID, proxyURL string) bool { + if value == nil { + return false + } + switch typed := (*value).(type) { + case map[string]any: + changed := false + for key, child := range typed { + childValue := child + if rewriteGrokMediaVideoContentURLValue(&childValue, requestID, proxyURL) { + typed[key] = childValue + changed = true + } + } + return changed + case []any: + changed := false + for index, child := range typed { + childValue := child + if rewriteGrokMediaVideoContentURLValue(&childValue, requestID, proxyURL) { + typed[index] = childValue + changed = true + } + } + return changed + case string: + if isGrokMediaVideoContentURL(typed, requestID) { + *value = proxyURL + return true + } + } + return false +} + +func isGrokMediaVideoContentURL(rawURL, requestID string) bool { + parsed, err := url.Parse(strings.TrimSpace(rawURL)) + if err != nil || parsed.Path == "" { + return false + } + segments := strings.Split(strings.Trim(parsed.Path, "/"), "/") + if len(segments) < 3 { + return false + } + requestID = strings.Trim(requestID, "/") + decodedID, err := url.PathUnescape(segments[len(segments)-2]) + if err != nil { + return false + } + return segments[len(segments)-3] == "videos" && + decodedID == requestID && + segments[len(segments)-1] == "content" +} + +func grokMediaContentProxyURL(c *gin.Context, requestID string) string { + if c == nil || c.Request == nil || c.Request.URL == nil || strings.TrimSpace(requestID) == "" { + return "" + } + pathPrefix := "" + if strings.HasPrefix(c.Request.URL.Path, "/v1/") { + pathPrefix = "/v1" + } + path := pathPrefix + "/videos/" + url.PathEscape(strings.Trim(requestID, "/")) + "/content" + host := strings.TrimSpace(c.GetHeader("X-Forwarded-Host")) + if comma := strings.IndexByte(host, ','); comma >= 0 { + host = strings.TrimSpace(host[:comma]) + } + if host == "" { + host = strings.TrimSpace(c.Request.Host) + } + if strings.ContainsAny(host, " \t\r\n") { + return "" + } + scheme := strings.ToLower(strings.TrimSpace(c.GetHeader("X-Forwarded-Proto"))) + if comma := strings.IndexByte(scheme, ','); comma >= 0 { + scheme = strings.TrimSpace(scheme[:comma]) + } + if scheme != "http" && scheme != "https" { + scheme = strings.ToLower(strings.TrimSpace(c.Request.URL.Scheme)) + } + if scheme != "http" && scheme != "https" { + if c.Request.TLS != nil { + scheme = "https" + } else { + scheme = "http" + } + } + if host == "" { + return path + } + return scheme + "://" + host + path +} diff --git a/backend/internal/service/grok_media_content_test.go b/backend/internal/service/grok_media_content_test.go new file mode 100644 index 0000000000..9668f35f7c --- /dev/null +++ b/backend/internal/service/grok_media_content_test.go @@ -0,0 +1,121 @@ +//go:build unit + +package service + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +type grokMediaContentUpstreamStub struct { + request *http.Request + response *http.Response +} + +func (s *grokMediaContentUpstreamStub) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + s.request = req + return s.response, nil +} + +func (s *grokMediaContentUpstreamStub) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) { + return s.Do(req, proxyURL, accountID, accountConcurrency) +} + +func grokMediaContentTestAccount() *Account { + return &Account{ + ID: 9, + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "upstream-key", + "base_url": "https://relay.example/v1", + }, + } +} + +func grokMediaContentTestContext(method, target string, headers map[string]string) (*gin.Context, *httptest.ResponseRecorder) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(method, target, nil) + for name, value := range headers { + c.Request.Header.Set(name, value) + } + return c, recorder +} + +func TestForwardGrokMediaContentUsesUpstreamCredentialAndStreamsRange(t *testing.T) { + upstream := &grokMediaContentUpstreamStub{ + response: &http.Response{ + StatusCode: http.StatusPartialContent, + Header: http.Header{ + "Content-Type": []string{"video/mp4"}, + "Content-Length": []string{"13"}, + "Content-Range": []string{"bytes 0-12/100"}, + "Accept-Ranges": []string{"bytes"}, + "Content-Disposition": []string{ + `attachment; filename="task-1.mp4"`, + }, + }, + Body: io.NopCloser(strings.NewReader("video-payload")), + }, + } + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + c, recorder := grokMediaContentTestContext(http.MethodGet, "https://api.example/v1/videos/task-1/content", map[string]string{ + "Range": "bytes=0-12", + }) + + result, err := svc.ForwardGrokMedia( + context.Background(), c, grokMediaContentTestAccount(), + GrokMediaEndpointVideoContent, "task-1", nil, "", + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, http.StatusPartialContent, recorder.Code) + require.Equal(t, "video-payload", recorder.Body.String()) + require.Equal(t, "https://relay.example/v1/videos/task-1/content", upstream.request.URL.String()) + require.Equal(t, "Bearer upstream-key", upstream.request.Header.Get("Authorization")) + require.Equal(t, "bytes=0-12", upstream.request.Header.Get("Range")) + require.Equal(t, "*/*", upstream.request.Header.Get("Accept")) + require.Equal(t, "video/mp4", recorder.Header().Get("Content-Type")) + require.Equal(t, "13", recorder.Header().Get("Content-Length")) + require.Equal(t, "bytes 0-12/100", recorder.Header().Get("Content-Range")) + require.Equal(t, "bytes", recorder.Header().Get("Accept-Ranges")) + require.Equal(t, `attachment; filename="task-1.mp4"`, recorder.Header().Get("Content-Disposition")) +} + +func TestForwardGrokVideoStatusRewritesOnlyProtectedContentURL(t *testing.T) { + statusBody := `{"id":"task-1","status":"completed","url":"https://relay.example/v1/videos/task-1/content","download_url":"/v1/videos/task-1/content","video_url":"https://vidgen.x.ai/task-1.mp4","counter":9007199254740993}` + upstream := &grokMediaContentUpstreamStub{ + response: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(statusBody)), + }, + } + svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + c, recorder := grokMediaContentTestContext(http.MethodGet, "https://api.example/v1/videos/task-1", nil) + + _, err := svc.ForwardGrokMedia( + context.Background(), c, grokMediaContentTestAccount(), + GrokMediaEndpointVideoStatus, "task-1", nil, "", + ) + + require.NoError(t, err) + require.Equal(t, http.StatusOK, recorder.Code) + require.Equal(t, "https://api.example/v1/videos/task-1/content", gjson.Get(recorder.Body.String(), "url").String()) + require.Equal(t, "https://api.example/v1/videos/task-1/content", gjson.Get(recorder.Body.String(), "download_url").String()) + require.Equal(t, "https://vidgen.x.ai/task-1.mp4", gjson.Get(recorder.Body.String(), "video_url").String()) + require.Equal(t, "9007199254740993", gjson.Get(recorder.Body.String(), "counter").String()) +} diff --git a/backend/internal/service/grok_upstream_url.go b/backend/internal/service/grok_upstream_url.go index e8d7fbf502..362e746e6f 100644 --- a/backend/internal/service/grok_upstream_url.go +++ b/backend/internal/service/grok_upstream_url.go @@ -112,6 +112,12 @@ func buildGrokMediaURL(account *Account, cfg *config.Config, endpoint GrokMediaE return xai.BuildVideosExtensionsURLWithValidator(baseURL, validator) case GrokMediaEndpointVideoStatus: return xai.BuildVideoURLWithValidator(baseURL, requestID, validator) + case GrokMediaEndpointVideoContent: + videoURL, err := xai.BuildVideoURLWithValidator(baseURL, requestID, validator) + if err != nil { + return "", err + } + return videoURL + "/content", nil default: return "", fmt.Errorf("unsupported grok media endpoint: %s", endpoint) } diff --git a/backend/internal/service/grok_upstream_url_test.go b/backend/internal/service/grok_upstream_url_test.go index 61c824e2ad..e2f6e8bce5 100644 --- a/backend/internal/service/grok_upstream_url_test.go +++ b/backend/internal/service/grok_upstream_url_test.go @@ -35,6 +35,10 @@ func TestGrokAPIKeyURLPolicyFollowsGlobalSecurityConfig(t *testing.T) { mediaURL, err := buildGrokMediaURL(account, cfg, GrokMediaEndpointImagesGenerations, "") require.NoError(t, err) require.Equal(t, "http://grok.example.test/v1/images/generations", mediaURL) + + contentURL, err := buildGrokMediaURL(account, cfg, GrokMediaEndpointVideoContent, "request 123") + require.NoError(t, err) + require.Equal(t, "http://grok.example.test/v1/videos/request%20123/content", contentURL) }) t.Run("insecure HTTP disabled", func(t *testing.T) {