opensource/sub2api
3
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-09-24 16:05:44 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #3614 from heathermhuang/codex/grok-media-group-gate-fix

Browse Source 
fix: enable Grok media generation groups
This commit is contained in:
Wesley Liddick
2026-07-02 17:34:09 +08:00
committed by GitHub
parent fafdd8342f 99a8d8ad28
commit 441b8ec2e2
6 changed files with 137 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/service/admin_service.go
+9 -1
View File
@@ -1814,6 +1814,12 @@ func defaultModelsListCandidateIDs(platform string) []string {
}
}
func defaultAllowImageGenerationForPlatform(platform string) bool {
// Grok image and video generation routes share the legacy image-generation gate.
// Older clients send the false zero value, so Grok groups must default enabled.
return platform == PlatformGrok
}
func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupInput) (*Group, error) {
if input.RateMultiplier <= 0 {
return nil, errors.New("rate_multiplier must be > 0")
@@ -1879,6 +1885,8 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
mcpXMLInject = *input.MCPXMLInject
}
allowImageGeneration := input.AllowImageGeneration || defaultAllowImageGenerationForPlatform(platform)
// 如果指定了复制账号的源分组，先获取账号 ID 列表
var accountIDsToCopy []int64
if len(input.CopyAccountsFromGroupIDs) > 0 {
@@ -1922,7 +1930,7 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
DailyLimitUSD: dailyLimit,
WeeklyLimitUSD: weeklyLimit,
MonthlyLimitUSD: monthlyLimit,
AllowImageGeneration: input.AllowImageGeneration,
AllowImageGeneration: allowImageGeneration,
ImageRateIndependent: input.ImageRateIndependent,
ImageRateMultiplier: imageRateMultiplier,
PeakRateEnabled: peakRateEnabled,
backend/internal/service/admin_service_group_test.go
+34
View File
@@ -198,6 +198,40 @@ func TestAdminService_CreateGroup_NilImagePricing(t *testing.T) {
require.Nil(t, repo.created.ImagePrice4K)
}
func TestAdminService_CreateGroup_DefaultsGrokMediaGenerationEnabled(t *testing.T) {
repo := &groupRepoStubForAdmin{}
svc := &adminServiceImpl{groupRepo: repo}
group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
Name: "grok-media",
Description: "Grok media group",
Platform: PlatformGrok,
RateMultiplier: 1.0,
})
require.NoError(t, err)
require.NotNil(t, group)
require.NotNil(t, repo.created)
require.True(t, repo.created.AllowImageGeneration)
require.True(t, group.AllowImageGeneration)
}
func TestAdminService_CreateGroup_PreservesNonGrokImageGenerationDisabled(t *testing.T) {
repo := &groupRepoStubForAdmin{}
svc := &adminServiceImpl{groupRepo: repo}
group, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
Name: "anthropic-text",
Description: "Anthropic text group",
Platform: PlatformAnthropic,
RateMultiplier: 1.0,
})
require.NoError(t, err)
require.NotNil(t, group)
require.NotNil(t, repo.created)
require.False(t, repo.created.AllowImageGeneration)
require.False(t, group.AllowImageGeneration)
}
// TestAdminService_UpdateGroup_WithImagePricing 测试更新分组时 ImagePrice 字段正确更新
func TestAdminService_UpdateGroup_WithImagePricing(t *testing.T) {
existingGroup := &Group{
backend/internal/service/grok_media.go
+32
View File
@@ -17,6 +17,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
type GrokMediaEndpoint string
@@ -296,6 +297,10 @@ func (s *OpenAIGatewayService) ForwardGrokMedia(
if err != nil {
return nil, err
}
body, contentType, err = normalizeGrokMediaForwardBody(endpoint, body, contentType)
if err != nil {
return nil, err
}
var bodyReader io.Reader
if endpoint.RequiresRequestBody() {
@@ -424,6 +429,33 @@ func prepareGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, conten
return out, "application/json", nil
}
func normalizeGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, contentType string) ([]byte, string, error) {
if !endpoint.RequiresRequestBody() || !gjson.ValidBytes(body) {
return body, contentType, nil
}
model := strings.TrimSpace(gjson.GetBytes(body, "model").String())
upstreamModel := normalizeGrokMediaModelForEndpoint(endpoint, model)
if upstreamModel == "" || upstreamModel == model {
return body, contentType, nil
}
out, err := sjson.SetBytes(body, "model", upstreamModel)
if err != nil {
return nil, "", fmt.Errorf("rewrite grok media model: %w", err)
}
return out, contentType, nil
}
func normalizeGrokMediaModelForEndpoint(endpoint GrokMediaEndpoint, model string) string {
model = strings.TrimSpace(model)
switch endpoint {
case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits:
if model == "grok-imagine" {
return "grok-imagine-image-quality"
}
}
return model
}
type grokMediaUsageMetadata struct {
ResponseID string
Usage OpenAIUsage
backend/internal/service/openai_gateway_grok_test.go
+44 -4
View File
@@ -146,6 +146,25 @@ func TestBuildGrokResponsesRequestRejectsUnsafeAccountBaseURL(t *testing.T) {
require.Contains(t, err.Error(), "invalid base url")
}
func TestGrokMediaGenerationGateCoversImagesAndVideo(t *testing.T) {
tests := []struct {
name string
endpoint GrokMediaEndpoint
want bool
}{
{name: "image generation", endpoint: GrokMediaEndpointImagesGenerations, want: true},
{name: "image edit", endpoint: GrokMediaEndpointImagesEdits, want: true},
{name: "video generation", endpoint: GrokMediaEndpointVideosGenerations, want: true},
{name: "video status", endpoint: GrokMediaEndpointVideoStatus, want: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.want, tt.endpoint.IsGenerationRequest())
})
}
}
func TestExtractGrokMediaModelSupportsJSONAndMultipart(t *testing.T) {
require.Equal(t, "grok-imagine", ExtractGrokMediaModel("application/json", []byte(`{"model":"grok-imagine"}`)))
@@ -182,7 +201,28 @@ func TestParseGrokMediaRequestBuildsMultipartModerationBody(t *testing.T) {
require.True(t, strings.HasPrefix(gjson.GetBytes(moderationBody, "images.0.image_url").String(), "data:image/"))
}
func TestForwardGrokMediaImagesGenerationPassthrough(t *testing.T) {
func TestNormalizeGrokMediaModelForEndpoint(t *testing.T) {
tests := []struct {
name string
endpoint GrokMediaEndpoint
model string
want string
}{
{name: "image generation alias", endpoint: GrokMediaEndpointImagesGenerations, model: "grok-imagine", want: "grok-imagine-image-quality"},
{name: "image edit alias", endpoint: GrokMediaEndpointImagesEdits, model: "grok-imagine", want: "grok-imagine-image-quality"},
{name: "image quality passthrough", endpoint: GrokMediaEndpointImagesGenerations, model: "grok-imagine-image-quality", want: "grok-imagine-image-quality"},
{name: "image fast passthrough", endpoint: GrokMediaEndpointImagesGenerations, model: "grok-imagine-image", want: "grok-imagine-image"},
{name: "video passthrough", endpoint: GrokMediaEndpointVideosGenerations, model: "grok-imagine-video", want: "grok-imagine-video"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
require.Equal(t, tt.want, normalizeGrokMediaModelForEndpoint(tt.endpoint, tt.model))
})
}
}
func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) {
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
gin.SetMode(gin.TestMode)
@@ -219,12 +259,12 @@ func TestForwardGrokMediaImagesGenerationPassthrough(t *testing.T) {
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.JSONEq(t, `{"model":"grok-imagine-image-quality","prompt":"draw a cat"}`, 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, "grok-imagine-image-quality", result.Model)
require.Equal(t, "grok-imagine-image-quality", result.BillingModel)
require.Equal(t, 1, result.ImageCount)
require.Equal(t, ImageBillingSize2K, result.ImageSize)
}
backend/migrations/158_enable_grok_media_generation_groups.sql
+7
View File
@@ -0,0 +1,7 @@
-- PR 3593 added Grok media routes for image generation, image edits, and video generation.
-- Existing Grok groups were created before the image-generation gate knew about
-- the Grok platform, so backfill them onto the same generation capability gate.
UPDATE groups
SET allow_image_generation = true
WHERE platform = 'grok'
AND allow_image_generation = false;
backend/migrations/auth_identity_payment_migrations_regression_test.go
+11
View File
@@ -169,6 +169,17 @@ func TestMigration151AddsAccountAutoPauseExpiryPartialIndex(t *testing.T) {
require.Contains(t, sql, "expires_at IS NOT NULL")
}
func TestMigration158BackfillsGrokMediaGenerationGroups(t *testing.T) {
content, err := FS.ReadFile("158_enable_grok_media_generation_groups.sql")
require.NoError(t, err)
sql := string(content)
require.Contains(t, sql, "UPDATE groups")
require.Contains(t, sql, "SET allow_image_generation = true")
require.Contains(t, sql, "WHERE platform = 'grok'")
require.Contains(t, sql, "AND allow_image_generation = false")
}
func TestMigration154AddsSparkShadowColumnsAndConstraintsWithoutHotIndexes(t *testing.T) {
content, err := FS.ReadFile("154_account_spark_shadow.sql")
require.NoError(t, err)