mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
Merge pull request #3593 from heathermhuang/codex/grok-media-routing
fix: route Grok media endpoints
This commit is contained in:
Wesley Liddick
@@ -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.
|
||||
|
||||
@@ -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},
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -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))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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"])
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user