zizi-al/sub2api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #2927 from moonagic/main

Browse Source 
fix antigravity gemini rate limit and account scheduling
This commit is contained in:
Wesley Liddick 2026-06-01 14:50:36 +08:00 committed by GitHub
commit 5f63fe1945
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
16 changed files with 747 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
backend/internal/handler/gateway_handler.go
View File

@ -441,7 +441,17 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
// 记录 Forward 前已写入字节数，Forward 后若增加则说明 SSE 内容已发，禁止 failover
writerSizeBeforeForward := c.Writer.Size()
if account.Platform == service.PlatformAntigravity {
result, err = h.antigravityGatewayService.ForwardGemini(requestCtx, c, account, reqModel, "generateContent", reqStream, body, hasBoundSession)
result, err = h.antigravityGatewayService.ForwardGemini(
requestCtx,
c,
account,
reqModel,
"generateContent",
reqStream,
body,
hasBoundSession,
service.WithForwardGeminiSession(derefGroupID(apiKey.GroupID), sessionKey),
)
} else {
result, err = h.geminiCompatService.Forward(requestCtx, c, account, body)
}

13
backend/internal/handler/gemini_v1beta_handler.go
View File

@ -477,8 +477,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
if fs.SwitchCount > 0 {
requestCtx = service.WithAccountSwitchCount(requestCtx, fs.SwitchCount, h.metadataBridgeEnabled())
}
sessionGroupID := derefGroupID(apiKey.GroupID)
if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey {
result, err = h.antigravityGatewayService.ForwardGemini(requestCtx, c, account, modelName, action, stream, body, hasBoundSession)
result, err = h.antigravityGatewayService.ForwardGemini(
requestCtx,
c,
account,
modelName,
action,
stream,
body,
hasBoundSession,
service.WithForwardGeminiSession(sessionGroupID, sessionKey),
)
} else {
result, err = h.geminiCompatService.ForwardNative(requestCtx, c, account, modelName, action, stream, body)
}

1
backend/internal/repository/scheduler_cache.go
View File

@ -559,6 +559,7 @@ func filterSchedulerExtra(extra map[string]any) map[string]any {
"auto_pause_7d_threshold",
"auto_pause_5h_disabled",
"auto_pause_7d_disabled",
"model_rate_limits",
}
filtered := make(map[string]any)
for _, key := range keys {

26
backend/internal/repository/scheduler_cache_unit_test.go
View File

@ -108,3 +108,29 @@ func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) {
require.Equal(t, true, got.Extra["auto_pause_5h_disabled"])
require.Equal(t, false, got.Extra["auto_pause_7d_disabled"])
}
func TestBuildSchedulerMetadataAccount_KeepsModelRateLimits(t *testing.T) {
account := service.Account{
ID: 90,
Platform: service.PlatformAntigravity,
Extra: map[string]any{
"model_rate_limits": map[string]any{
"gemini-3-flash": map[string]any{
"rate_limit_reset_at": "2026-05-30T10:10:00Z",
},
"antigravity:gemini": map[string]any{
"rate_limit_reset_at": "2026-05-30T10:10:00Z",
},
},
"unused_large_field": "drop-me",
},
}
got := buildSchedulerMetadataAccount(account)
limits, ok := got.Extra["model_rate_limits"].(map[string]any)
require.True(t, ok)
require.Contains(t, limits, "gemini-3-flash")
require.Contains(t, limits, "antigravity:gemini")
require.Nil(t, got.Extra["unused_large_field"])
}

61
backend/internal/service/antigravity_default_test_stubs_test.go Normal file
View File

@ -0,0 +1,61 @@
//go:build !unit
package service
import (
"context"
"time"
)
type defaultRateLimitCall struct {
accountID int64
resetAt time.Time
}
type defaultModelRateLimitCall struct {
accountID int64
modelKey string
resetAt time.Time
}
type defaultExtraUpdateCall struct {
accountID int64
updates map[string]any
}
type stubAntigravityAccountRepo struct {
AccountRepository
rateCalls []defaultRateLimitCall
modelRateLimitCalls []defaultModelRateLimitCall
extraUpdateCalls []defaultExtraUpdateCall
}
func (s *stubAntigravityAccountRepo) SetRateLimited(_ context.Context, id int64, resetAt time.Time) error {
s.rateCalls = append(s.rateCalls, defaultRateLimitCall{accountID: id, resetAt: resetAt})
return nil
}
func (s *stubAntigravityAccountRepo) SetModelRateLimit(_ context.Context, id int64, modelKey string, resetAt time.Time, _ ...string) error {
s.modelRateLimitCalls = append(s.modelRateLimitCalls, defaultModelRateLimitCall{accountID: id, modelKey: modelKey, resetAt: resetAt})
return nil
}
func (s *stubAntigravityAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error {
s.extraUpdateCalls = append(s.extraUpdateCalls, defaultExtraUpdateCall{accountID: id, updates: updates})
return nil
}
type defaultDeleteSessionCall struct {
groupID int64
sessionHash string
}
type stubSmartRetryCache struct {
GatewayCache
deleteCalls []defaultDeleteSessionCall
}
func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID int64, sessionHash string) error {
c.deleteCalls = append(c.deleteCalls, defaultDeleteSessionCall{groupID: groupID, sessionHash: sessionHash})
return nil
}

126
backend/internal/service/antigravity_gateway_service.go
View File

@ -228,12 +228,11 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration, truncateForLog(respBody, 200))
resetAt := time.Now().Add(rateLimitDuration)
if !setModelRateLimitByModelName(p.ctx, p.accountRepo, p.account.ID, modelName, p.prefix, resp.StatusCode, resetAt, false) {
if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, false) {
p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession)
logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited account=%d (no model mapping)", p.prefix, resp.StatusCode, p.account.ID)
} else {
s.updateAccountModelRateLimitInCache(p.ctx, p.account, modelName, resetAt)
}
s.clearStickySession(p.ctx, p.groupID, p.sessionHash)
// 返回账号切换信号，让上层切换账号重试
return &smartRetryResult{
@ -392,20 +391,10 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, rateLimitDuration, truncateForLog(retryBody, 200))
resetAt := time.Now().Add(rateLimitDuration)
if p.accountRepo != nil && modelName != "" {
if err := p.accountRepo.SetModelRateLimit(p.ctx, p.account.ID, modelName, resetAt); err != nil {
logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_failed model=%s error=%v", p.prefix, resp.StatusCode, modelName, err)
} else {
logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_after_smart_retry model=%s account=%d reset_in=%v",
p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration)
s.updateAccountModelRateLimitInCache(p.ctx, p.account, modelName, resetAt)
}
}
s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, true)
// 清除粘性会话绑定，避免下次请求仍命中限流账号
if s.cache != nil && p.sessionHash != "" {
_ = s.cache.DeleteSessionAccountID(p.ctx, p.groupID, p.sessionHash)
}
s.clearStickySession(p.ctx, p.groupID, p.sessionHash)
// 返回账号切换信号，让上层切换账号重试
return &smartRetryResult{
@ -954,8 +943,14 @@ func (s *AntigravityGatewayService) checkErrorPolicy(ctx context.Context, accoun
func (s *AntigravityGatewayService) applyErrorPolicy(p antigravityRetryLoopParams, statusCode int, headers http.Header, respBody []byte) (handled bool, outStatus int, retErr error) {
switch s.checkErrorPolicy(p.ctx, p.account, statusCode, respBody) {
case ErrorPolicySkipped:
if s.handleAntigravityModelRateLimitBeforePolicy(p, statusCode, headers, respBody) {
return true, statusCode, nil
}
return true, http.StatusInternalServerError, nil
case ErrorPolicyMatched:
if s.handleAntigravityModelRateLimitBeforePolicy(p, statusCode, headers, respBody) {
return true, statusCode, nil
}
_ = p.handleError(p.ctx, p.prefix, p.account, statusCode, headers, respBody,
p.requestedModel, p.groupID, p.sessionHash, p.isStickySession)
return true, statusCode, nil
@ -967,6 +962,31 @@ func (s *AntigravityGatewayService) applyErrorPolicy(p antigravityRetryLoopParam
return false, statusCode, nil
}
func (s *AntigravityGatewayService) handleAntigravityModelRateLimitBeforePolicy(p antigravityRetryLoopParams, statusCode int, headers http.Header, respBody []byte) bool {
if statusCode != http.StatusTooManyRequests && statusCode != http.StatusServiceUnavailable {
return false
}
if p.account == nil || p.account.Platform != PlatformAntigravity {
return false
}
_, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, respBody)
if isModelCapacityExhausted || !shouldRateLimitModel || strings.TrimSpace(modelName) == "" {
return false
}
rateLimitDuration := waitDuration
if rateLimitDuration <= 0 {
rateLimitDuration = antigravityDefaultRateLimitDuration
}
resetAt := time.Now().Add(rateLimitDuration)
if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, statusCode, resetAt, false) {
return false
}
s.clearStickySession(p.ctx, p.groupID, p.sessionHash)
logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_before_error_policy model=%s account=%d reset_in=%v",
p.prefix, statusCode, modelName, p.account.ID, rateLimitDuration)
return true
}
// mapAntigravityModel 获取映射后的模型名
// 完全依赖映射配置：账户映射（通配符）→ 默认映射兜底（DefaultAntigravityModelMapping）
// 注意：返回空字符串表示模型不被支持，调度时会过滤掉该账号
@ -974,6 +994,7 @@ func mapAntigravityModel(account *Account, requestedModel string) string {
if account == nil {
return ""
}
requestedModel = strings.TrimPrefix(requestedModel, "models/")
// 获取映射表（未配置时自动使用 DefaultAntigravityModelMapping）
mapping := account.GetModelMapping()
@ -2073,8 +2094,28 @@ func stripSignatureSensitiveBlocksFromClaudeRequest(req *antigravity.ClaudeReque
// └─ retryDelay < 7s → 等待后重试 1 次
// ├─ 成功 → 正常返回
// └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号
func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool) (*ForwardResult, error) {
type ForwardGeminiOption func(*forwardGeminiOptions)
type forwardGeminiOptions struct {
groupID int64
sessionHash string
}
func WithForwardGeminiSession(groupID int64, sessionHash string) ForwardGeminiOption {
return func(opts *forwardGeminiOptions) {
opts.groupID = groupID
opts.sessionHash = sessionHash
}
}
func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool, options ...ForwardGeminiOption) (*ForwardResult, error) {
startTime := time.Now()
forwardOpts := forwardGeminiOptions{}
for _, apply := range options {
if apply != nil {
apply(&forwardOpts)
}
}
sessionID := getSessionID(c)
prefix := logPrefix(sessionID, account.Name)
@ -2179,8 +2220,8 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
handleError: s.handleUpstreamError,
requestedModel: originalModel,
isStickySession: isStickySession, // ForwardGemini 由上层判断粘性会话
groupID: 0, // ForwardGemini 方法没有 groupID，由上层处理粘性会话清除
sessionHash: "", // ForwardGemini 方法没有 sessionHash，由上层处理粘性会话清除
groupID: forwardOpts.groupID,
sessionHash: forwardOpts.sessionHash,
})
if err != nil {
// 检查是否是账号切换信号，转换为 UpstreamFailoverError 让 Handler 切换账号
@ -2278,8 +2319,8 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
handleError: s.handleUpstreamError,
requestedModel: originalModel,
isStickySession: isStickySession,
groupID: 0,
sessionHash: "",
groupID: forwardOpts.groupID,
sessionHash: forwardOpts.sessionHash,
})
if retryErr == nil {
retryResp := retryResult.resp
@ -2355,7 +2396,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
if unwrapErr != nil || len(unwrappedForOps) == 0 {
unwrappedForOps = respBody
}
s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", isStickySession)
s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, forwardOpts.groupID, forwardOpts.sessionHash, isStickySession)
upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(unwrappedForOps))
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
upstreamDetail := s.getUpstreamErrorDetail(unwrappedForOps)
@ -2566,6 +2607,34 @@ func setModelRateLimitByModelName(ctx context.Context, repo AccountRepository, a
return true
}
func (s *AntigravityGatewayService) setAntigravityModelRateLimits(ctx context.Context, repo AccountRepository, account *Account, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool {
if account == nil || repo == nil {
return false
}
keys := antigravityModelRateLimitKeys(modelName)
if len(keys) == 0 {
return false
}
success := false
for _, key := range keys {
if setModelRateLimitByModelName(ctx, repo, account.ID, key, prefix, statusCode, resetAt, afterSmartRetry) {
s.updateAccountModelRateLimitInCache(ctx, account, key, resetAt)
success = true
}
}
return success
}
func (s *AntigravityGatewayService) clearStickySession(ctx context.Context, groupID int64, sessionHash string) {
if s == nil || s.cache == nil || strings.TrimSpace(sessionHash) == "" {
return
}
if err := s.cache.DeleteSessionAccountID(ctx, groupID, sessionHash); err != nil {
logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] sticky_session_clear_failed group_id=%d session=%s err=%v", groupID, shortSessionHash(sessionHash), err)
}
}
func antigravityFallbackCooldownSeconds() (time.Duration, bool) {
raw := strings.TrimSpace(os.Getenv(antigravityFallbackSecondsEnv))
if raw == "" {
@ -2644,7 +2713,7 @@ func parseAntigravitySmartRetryInfo(body []byte) *antigravitySmartRetryInfo {
if atType == googleRPCTypeErrorInfo {
if meta, ok := dm["metadata"].(map[string]any); ok {
if model, ok := meta["model"].(string); ok {
modelName = model
modelName = normalizeAntigravityModelName(model)
}
}
// 检查 reason
@ -2818,13 +2887,7 @@ func (s *AntigravityGatewayService) setModelRateLimitAndClearSession(p *handleMo
logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v",
p.prefix, p.statusCode, info.ModelName, p.account.ID, info.RetryDelay)
// 设置模型限流状态（数据库）
if err := s.accountRepo.SetModelRateLimit(p.ctx, p.account.ID, info.ModelName, resetAt); err != nil {
logger.LegacyPrintf("service.antigravity_gateway", "%s model_rate_limit_failed model=%s error=%v", p.prefix, info.ModelName, err)
}
// 立即更新 Redis 快照中账号的限流状态，避免并发请求重复选中
s.updateAccountModelRateLimitInCache(p.ctx, p.account, info.ModelName, resetAt)
s.setAntigravityModelRateLimits(p.ctx, s.accountRepo, p.account, info.ModelName, p.prefix, p.statusCode, resetAt, false)
// 清除粘性会话绑定
if p.cache != nil && p.sessionHash != "" {
@ -2914,12 +2977,11 @@ func (s *AntigravityGatewayService) handleUpstreamError(
}
if modelKey != "" {
ra := s.resolveResetTime(resetAt, defaultDur)
if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, modelKey, ra); err != nil {
logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s error=%v", prefix, modelKey, err)
if !s.setAntigravityModelRateLimits(ctx, s.accountRepo, account, modelKey, prefix, statusCode, ra, false) {
logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s", prefix, modelKey)
} else {
logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limited model=%s account=%d reset_at=%v reset_in=%v",
prefix, modelKey, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second))
s.updateAccountModelRateLimitInCache(ctx, account, modelKey, ra)
}
return nil
}

80
backend/internal/service/antigravity_gateway_service_test.go
View File

@ -500,6 +500,86 @@ func TestAntigravityGatewayService_ForwardGemini_StickySessionForceCacheBilling(
require.True(t, failoverErr.ForceCacheBilling, "ForceCacheBilling should be true for sticky session switch")
}
func TestAntigravityGatewayService_ForwardGemini_ClearsStickySessionOnGeminiRateLimit(t *testing.T) {
gin.SetMode(gin.TestMode)
writer := httptest.NewRecorder()
c, _ := gin.CreateTestContext(writer)
body, err := json.Marshal(map[string]any{
"contents": []map[string]any{
{"role": "user", "parts": []map[string]any{{"text": "hi"}}},
},
})
require.NoError(t, err)
req := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-3-flash-preview:generateContent", bytes.NewReader(body))
c.Request = req
respBody := []byte(`{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"}
]
}
}`)
upstream := &httpUpstreamStub{resp: &http.Response{
StatusCode: http.StatusTooManyRequests,
Header: http.Header{},
Body: io.NopCloser(bytes.NewReader(respBody)),
}}
repo := &stubAntigravityAccountRepo{}
cache := &stubSmartRetryCache{}
svc := &AntigravityGatewayService{
tokenProvider: &AntigravityTokenProvider{},
httpUpstream: upstream,
accountRepo: repo,
cache: cache,
}
account := &Account{
ID: 44,
Name: "acc-gemini-runtime-rate-limited",
Platform: PlatformAntigravity,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "token",
"expires_at": time.Now().Add(time.Hour).Format(time.RFC3339),
"project_id": "proj",
},
Extra: map[string]any{
"mixed_scheduling": true,
},
}
result, err := svc.ForwardGemini(
context.Background(),
c,
account,
"gemini-3-flash-preview",
"generateContent",
false,
body,
true,
WithForwardGeminiSession(77, "gemini:sticky-runtime"),
)
require.Nil(t, result)
var failoverErr *UpstreamFailoverError
require.ErrorAs(t, err, &failoverErr)
require.Equal(t, http.StatusServiceUnavailable, failoverErr.StatusCode)
require.Len(t, repo.modelRateLimitCalls, 2)
require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey)
require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
require.Len(t, cache.deleteCalls, 1)
require.Equal(t, int64(77), cache.deleteCalls[0].groupID)
require.Equal(t, "gemini:sticky-runtime", cache.deleteCalls[0].sessionHash)
}
// TestAntigravityGatewayService_Forward_BillsWithMappedModel
// 验证：Antigravity Claude 转发返回的计费模型使用映射后的模型
func TestAntigravityGatewayService_Forward_BillsWithMappedModel(t *testing.T) {

12
backend/internal/service/antigravity_quota_scope.go
View File

@ -8,7 +8,17 @@ import (
func normalizeAntigravityModelName(model string) string {
normalized := strings.ToLower(strings.TrimSpace(model))
normalized = strings.TrimPrefix(normalized, "models/")
if idx := strings.LastIndex(normalized, "/publishers/google/models/"); idx != -1 {
normalized = normalized[idx+len("/publishers/google/models/"):]
} else if idx := strings.LastIndex(normalized, "/publishers/anthropic/models/"); idx != -1 {
normalized = normalized[idx+len("/publishers/anthropic/models/"):]
} else if idx := strings.LastIndex(normalized, "/models/"); idx != -1 {
normalized = normalized[idx+len("/models/"):]
} else {
normalized = strings.TrimPrefix(normalized, "publishers/google/models/")
normalized = strings.TrimPrefix(normalized, "publishers/anthropic/models/")
normalized = strings.TrimPrefix(normalized, "models/")
}
return normalized
}

95
backend/internal/service/antigravity_rate_limit_test.go
View File

@ -821,6 +821,51 @@ func TestSetModelRateLimitByModelName_NotConvertToScope(t *testing.T) {
require.NotEqual(t, "claude_sonnet", call.modelKey, "should NOT be scope")
}
func TestSetAntigravityModelRateLimits_GeminiWritesFamilyScope(t *testing.T) {
repo := &stubAntigravityAccountRepo{}
svc := &AntigravityGatewayService{}
account := &Account{ID: 789, Platform: PlatformAntigravity}
resetAt := time.Now().Add(30 * time.Second)
success := svc.setAntigravityModelRateLimits(
context.Background(),
repo,
account,
"gemini-3-pro",
"[test]",
429,
resetAt,
false,
)
require.True(t, success)
require.Len(t, repo.modelRateLimitCalls, 2)
require.Equal(t, "gemini-3-pro", repo.modelRateLimitCalls[0].modelKey)
require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
}
func TestSetAntigravityModelRateLimits_ClaudeDoesNotWriteGeminiScope(t *testing.T) {
repo := &stubAntigravityAccountRepo{}
svc := &AntigravityGatewayService{}
account := &Account{ID: 790, Platform: PlatformAntigravity}
resetAt := time.Now().Add(30 * time.Second)
success := svc.setAntigravityModelRateLimits(
context.Background(),
repo,
account,
"claude-sonnet-4-5",
"[test]",
429,
resetAt,
false,
)
require.True(t, success)
require.Len(t, repo.modelRateLimitCalls, 1)
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
}
func TestAntigravityRetryLoop_PreCheck_SwitchesWhenRateLimited(t *testing.T) {
upstream := &recordingOKUpstream{}
account := &Account{
@ -1124,3 +1169,53 @@ func TestSchedulerSnapshotService_UpdateAccountInCache(t *testing.T) {
require.ErrorIs(t, err, expectedErr)
})
}
func TestNormalizeAntigravityModelName(t *testing.T) {
tests := []struct {
name string
model string
expected string
}{
{
name: "plain model name",
model: "gemini-1.5-pro",
expected: "gemini-1.5-pro",
},
{
name: "models/ prefix",
model: "models/gemini-1.5-pro",
expected: "gemini-1.5-pro",
},
{
name: "publishers/google/models/ prefix",
model: "publishers/google/models/gemini-1.5-pro",
expected: "gemini-1.5-pro",
},
{
name: "projects/.../publishers/google/models/ path",
model: "projects/my-proj/locations/us-central1/publishers/google/models/gemini-2.5-flash",
expected: "gemini-2.5-flash",
},
{
name: "publishers/anthropic/models/ prefix",
model: "publishers/anthropic/models/claude-sonnet-4-5",
expected: "claude-sonnet-4-5",
},
{
name: "projects/.../publishers/anthropic/models/ path",
model: "projects/my-proj/locations/global/publishers/anthropic/models/claude-sonnet-4-5",
expected: "claude-sonnet-4-5",
},
{
name: "mixed case and spaces",
model: " Models/Gemini-1.5-Pro ",
expected: "gemini-1.5-pro",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
actual := normalizeAntigravityModelName(tt.model)
require.Equal(t, tt.expected, actual)
})
}
}

8
backend/internal/service/antigravity_single_account_retry_test.go
View File

@ -196,8 +196,10 @@ func TestHandleSmartRetry_503_LongDelay_NoSingleAccountRetry_StillSwitches(t *te
require.Nil(t, result.resp, "should not return resp when switchError is set")
// 对照：多账号模式应设模型限流
require.Len(t, repo.modelRateLimitCalls, 1,
require.Len(t, repo.modelRateLimitCalls, 2,
"multi-account mode SHOULD set model rate limit")
require.Equal(t, "gemini-3-pro-high", repo.modelRateLimitCalls[0].modelKey)
require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
}
// TestHandleSmartRetry_429_LongDelay_SingleAccountRetry_StillSwitches
@ -412,8 +414,10 @@ func TestHandleSmartRetry_503_ShortDelay_NoSingleAccountRetry_SetsRateLimit(t *t
// 对照：多账号模式应返回 switchError
require.NotNil(t, result.switchError, "multi-account mode should return switchError for 503")
// 对照：多账号模式应设模型限流
require.Len(t, repo.modelRateLimitCalls, 1,
require.Len(t, repo.modelRateLimitCalls, 2,
"multi-account mode should set model rate limit")
require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey)
require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
}
// ---------------------------------------------------------------------------

28
backend/internal/service/antigravity_smart_retry_test.go
View File

@ -328,9 +328,10 @@ func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *test
require.Equal(t, "gemini-3-flash", result.switchError.RateLimitedModel)
require.False(t, result.switchError.IsStickySession)
// 验证模型限流已设置
require.Len(t, repo.modelRateLimitCalls, 1)
// 验证模型限流已设置：Gemini 同时写入精确模型和家族级 scope
require.Len(t, repo.modelRateLimitCalls, 2)
require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey)
require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
require.Len(t, upstream.calls, 1, "should have made one retry call (max attempts)")
}
@ -1104,10 +1105,9 @@ func TestHandleSmartRetry_ShortDelay_StickySession_SuccessRetry_NoDeleteSession(
require.Len(t, cache.deleteCalls, 0, "should NOT call DeleteSessionAccountID on successful retry")
}
// TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry
// 长延迟路径（情况1）在 handleSmartRetry 中不直接调用 DeleteSessionAccountID
// （清除由 handler 层的 shouldClearStickySession 在下次请求时处理）
func TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry(t *testing.T) {
// TestHandleSmartRetry_LongDelay_StickySession_ClearsSession
// 长延迟路径（情况1）应立即清除 sticky 绑定，避免下一次请求继续命中已限流账号。
func TestHandleSmartRetry_LongDelay_StickySession_ClearsSession(t *testing.T) {
repo := &stubAntigravityAccountRepo{}
cache := &stubSmartRetryCache{}
account := &Account{
@ -1159,10 +1159,9 @@ func TestHandleSmartRetry_LongDelay_StickySession_NoDeleteInHandleSmartRetry(t *
require.NotNil(t, result.switchError)
require.True(t, result.switchError.IsStickySession)
// 长延迟路径不在 handleSmartRetry 中调用 DeleteSessionAccountID
// （由上游 handler 的 shouldClearStickySession 处理）
require.Len(t, cache.deleteCalls, 0,
"long delay path should NOT call DeleteSessionAccountID in handleSmartRetry (handled by handler layer)")
require.Len(t, cache.deleteCalls, 1, "long delay path should clear sticky session in handleSmartRetry")
require.Equal(t, int64(42), cache.deleteCalls[0].groupID)
require.Equal(t, "sticky-hash-long-delay", cache.deleteCalls[0].sessionHash)
}
// TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession
@ -1227,6 +1226,10 @@ func TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession(t
require.Len(t, cache.deleteCalls, 1, "should call DeleteSessionAccountID after network error exhausts retry")
require.Equal(t, int64(99), cache.deleteCalls[0].groupID)
require.Equal(t, "sticky-net-error", cache.deleteCalls[0].sessionHash)
require.Len(t, repo.modelRateLimitCalls, 2)
require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey)
require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
}
// TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession
@ -1308,9 +1311,10 @@ func TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession
require.Equal(t, int64(77), cache.deleteCalls[0].groupID)
require.Equal(t, "sticky-503-short", cache.deleteCalls[0].sessionHash)
// 验证模型限流已设置
require.Len(t, repo.modelRateLimitCalls, 1)
// 验证模型限流已设置：Gemini 同时写入精确模型和家族级 scope
require.Len(t, repo.modelRateLimitCalls, 2)
require.Equal(t, "gemini-3-pro", repo.modelRateLimitCalls[0].modelKey)
require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
}
// TestAntigravityRetryLoop_SmartRetryFailed_StickySession_SwitchErrorPropagates

54
backend/internal/service/error_policy_test.go
View File

@ -389,6 +389,60 @@ func TestApplyErrorPolicy(t *testing.T) {
}
}
func TestApplyErrorPolicy_GeminiRateLimitBypassesCustomSkip(t *testing.T) {
repo := &stubAntigravityAccountRepo{}
cache := &stubSmartRetryCache{}
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
svc := &AntigravityGatewayService{
rateLimitService: rlSvc,
accountRepo: repo,
cache: cache,
}
account := &Account{
ID: 31,
Type: AccountTypeAPIKey,
Platform: PlatformAntigravity,
Credentials: map[string]any{
"custom_error_codes_enabled": true,
"custom_error_codes": []any{float64(500)},
},
}
body := []byte(`{
"error": {
"status": "RESOURCE_EXHAUSTED",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"}
]
}
}`)
p := antigravityRetryLoopParams{
ctx: context.Background(),
prefix: "[test]",
account: account,
accountRepo: repo,
groupID: 42,
sessionHash: "gemini:sticky",
handleError: func(context.Context, string, *Account, int, http.Header, []byte, string, int64, string, bool) *handleModelRateLimitResult {
t.Fatal("model rate limit should be handled before custom error fallback")
return nil
},
}
handled, outStatus, retErr := svc.applyErrorPolicy(p, http.StatusTooManyRequests, http.Header{}, body)
require.True(t, handled)
require.Equal(t, http.StatusTooManyRequests, outStatus)
require.NoError(t, retErr)
require.Len(t, repo.modelRateLimitCalls, 2)
require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey)
require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
require.Len(t, cache.deleteCalls, 1)
require.Equal(t, int64(42), cache.deleteCalls[0].groupID)
require.Equal(t, "gemini:sticky", cache.deleteCalls[0].sessionHash)
}
// ---------------------------------------------------------------------------
// errorPolicyRepoStub — minimal AccountRepository stub for error policy tests
// ---------------------------------------------------------------------------

100
backend/internal/service/gateway_multiplatform_test.go
View File

@ -1229,6 +1229,106 @@ func TestGatewayService_selectAccountWithMixedScheduling(t *testing.T) {
require.Equal(t, int64(2), acc.ID, "应选择优先级最高的账户（包含启用混合调度的antigravity）")
})
t.Run("混合调度-Gemini家族限流后跳过Antigravity账户", func(t *testing.T) {
resetAt := time.Now().Add(10 * time.Minute).Format(time.RFC3339)
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1,
Platform: PlatformAntigravity,
Priority: 1,
Status: StatusActive,
Schedulable: true,
Extra: map[string]any{
"mixed_scheduling": true,
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": resetAt,
},
},
},
},
{
ID: 2,
Platform: PlatformAntigravity,
Priority: 1,
Status: StatusActive,
Schedulable: true,
Extra: map[string]any{
"mixed_scheduling": true,
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": resetAt,
},
},
},
},
{
ID: 3,
Platform: PlatformAntigravity,
Priority: 2,
Status: StatusActive,
Schedulable: true,
Extra: map[string]any{"mixed_scheduling": true},
},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
svc := &GatewayService{
accountRepo: repo,
cache: &mockGatewayCacheForPlatform{},
cfg: testConfig(),
}
acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "", "gemini-3-pro-preview", nil, PlatformGemini)
require.NoError(t, err)
require.NotNil(t, acc)
require.Equal(t, int64(3), acc.ID)
})
t.Run("混合调度-Gemini家族限流不影响Claude调度", func(t *testing.T) {
resetAt := time.Now().Add(10 * time.Minute).Format(time.RFC3339)
repo := &mockAccountRepoForPlatform{
accounts: []Account{
{
ID: 1,
Platform: PlatformAntigravity,
Priority: 1,
Status: StatusActive,
Schedulable: true,
Extra: map[string]any{
"mixed_scheduling": true,
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": resetAt,
},
},
},
},
{ID: 2, Platform: PlatformAnthropic, Priority: 2, Status: StatusActive, Schedulable: true},
},
accountsByID: map[int64]*Account{},
}
for i := range repo.accounts {
repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i]
}
svc := &GatewayService{
accountRepo: repo,
cache: &mockGatewayCacheForPlatform{},
cfg: testConfig(),
}
acc, err := svc.selectAccountWithMixedScheduling(ctx, nil, "", "claude-sonnet-4-5", nil, PlatformAnthropic)
require.NoError(t, err)
require.NotNil(t, acc)
require.Equal(t, int64(1), acc.ID)
})
t.Run("混合调度-路由优先选择路由账号", func(t *testing.T) {
groupID := int64(30)
requestedModel := "claude-sonnet-4-5"

32
backend/internal/service/model_rate_limit.go
View File

@ -6,7 +6,10 @@ import (
"time"
)
const modelRateLimitsKey = "model_rate_limits"
const (
modelRateLimitsKey = "model_rate_limits"
antigravityGeminiModelRateLimitKey = "antigravity:gemini"
)
// isRateLimitActiveForKey 检查指定 key 的限流是否生效
func (a *Account) isRateLimitActiveForKey(key string) bool {
@ -35,6 +38,9 @@ func (a *Account) isModelRateLimitedWithContext(ctx context.Context, requestedMo
modelKey := a.GetMappedModel(requestedModel)
if a.Platform == PlatformAntigravity {
modelKey = resolveFinalAntigravityModelKey(ctx, a, requestedModel)
if isAntigravityGeminiModel(modelKey) && a.isRateLimitActiveForKey(antigravityGeminiModelRateLimitKey) {
return true
}
}
modelKey = strings.TrimSpace(modelKey)
if modelKey == "" {
@ -62,7 +68,13 @@ func (a *Account) GetModelRateLimitRemainingTimeWithContext(ctx context.Context,
if modelKey == "" {
return 0
}
return a.getRateLimitRemainingForKey(modelKey)
remaining := a.getRateLimitRemainingForKey(modelKey)
if a.Platform == PlatformAntigravity && isAntigravityGeminiModel(modelKey) {
if familyRemaining := a.getRateLimitRemainingForKey(antigravityGeminiModelRateLimitKey); familyRemaining > remaining {
return familyRemaining
}
}
return remaining
}
func resolveFinalAntigravityModelKey(ctx context.Context, account *Account, requestedModel string) string {
@ -77,6 +89,22 @@ func resolveFinalAntigravityModelKey(ctx context.Context, account *Account, requ
return modelKey
}
func isAntigravityGeminiModel(model string) bool {
return strings.HasPrefix(normalizeAntigravityModelName(model), "gemini-")
}
func antigravityModelRateLimitKeys(model string) []string {
model = strings.TrimSpace(model)
if model == "" {
return nil
}
keys := []string{model}
if isAntigravityGeminiModel(model) && model != antigravityGeminiModelRateLimitKey {
keys = append(keys, antigravityGeminiModelRateLimitKey)
}
return keys
}
func (a *Account) modelRateLimitResetAt(scope string) *time.Time {
if a == nil || a.Extra == nil || scope == "" {
return nil

62
backend/internal/service/model_rate_limit_test.go
View File

@ -121,6 +121,36 @@ func TestIsModelRateLimited(t *testing.T) {
requestedModel: "gemini-3-pro-preview",
expected: true,
},
{
name: "antigravity platform - gemini family rate limit blocks mapped preview",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "gemini-3-pro-preview",
expected: true,
},
{
name: "antigravity platform - gemini family rate limit does not block claude",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": future,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
expected: false,
},
{
name: "non-antigravity platform - gemini-3-pro-preview NOT mapped",
account: &Account{
@ -306,6 +336,38 @@ func TestGetModelRateLimitRemainingTime(t *testing.T) {
minExpected: 4 * time.Minute,
maxExpected: 6 * time.Minute,
},
{
name: "antigravity platform - gemini family rate limit remaining",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": future10m,
},
},
},
},
requestedModel: "gemini-3-pro-preview",
minExpected: 9 * time.Minute,
maxExpected: 11 * time.Minute,
},
{
name: "antigravity platform - gemini family remaining ignored for claude",
account: &Account{
Platform: PlatformAntigravity,
Extra: map[string]any{
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": future10m,
},
},
},
},
requestedModel: "claude-sonnet-4-5",
minExpected: 0,
maxExpected: 0,
},
}
for _, tt := range tests {

88
backend/internal/service/scheduler_snapshot_hydration_test.go
View File

@ -6,6 +6,8 @@ import (
"context"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
)
type snapshotHydrationCache struct {
@ -186,3 +188,89 @@ func TestGatewaySelectAccountWithLoadAwareness_HydratesSelectedAccountFromSchedu
t.Fatalf("expected hydrated api key, got %q", got)
}
}
func TestGatewaySelectAccountWithLoadAwareness_SkipsAntigravityGeminiFamilyRateLimitedSnapshot(t *testing.T) {
resetAt := time.Now().Add(10 * time.Minute).Format(time.RFC3339)
cache := &snapshotHydrationCache{
snapshot: []*Account{
{
ID: 1,
Platform: PlatformAntigravity,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Priority: 1,
AccountGroups: []AccountGroup{
{AccountID: 1, GroupID: 22},
},
GroupIDs: []int64{22},
Extra: map[string]any{
"mixed_scheduling": true,
modelRateLimitsKey: map[string]any{
antigravityGeminiModelRateLimitKey: map[string]any{
"rate_limit_reset_at": resetAt,
},
},
},
},
{
ID: 2,
Platform: PlatformAntigravity,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Priority: 2,
AccountGroups: []AccountGroup{
{AccountID: 2, GroupID: 22},
},
GroupIDs: []int64{22},
Extra: map[string]any{
"mixed_scheduling": true,
},
},
},
accounts: map[int64]*Account{
1: {ID: 1, Platform: PlatformAntigravity, Type: AccountTypeOAuth},
2: {ID: 2, Platform: PlatformAntigravity, Type: AccountTypeOAuth},
},
}
groupID := int64(22)
svc := &GatewayService{
schedulerSnapshot: NewSchedulerSnapshotService(cache, nil, nil, nil, nil),
groupRepo: &mockGroupRepoForGateway{
groups: map[int64]*Group{
groupID: {
ID: groupID,
Platform: PlatformGemini,
Status: StatusActive,
Hydrated: true,
},
},
},
concurrencyService: NewConcurrencyService(&mockConcurrencyCache{}),
cfg: &config.Config{
Gateway: config.GatewayConfig{
Scheduling: config.GatewaySchedulingConfig{
LoadBatchEnabled: true,
StickySessionMaxWaiting: 3,
StickySessionWaitTimeout: time.Second,
FallbackWaitTimeout: time.Second,
FallbackMaxWaiting: 10,
},
},
},
}
result, err := svc.SelectAccountWithLoadAwareness(context.Background(), &groupID, "", "gemini-3-flash-preview", nil, "", 0)
if err != nil {
t.Fatalf("SelectAccountWithLoadAwareness error: %v", err)
}
if result == nil || result.Account == nil {
t.Fatalf("expected selected account")
}
if result.Account.ID != 2 {
t.Fatalf("expected scheduler to skip Gemini-family limited antigravity account 1, got %d", result.Account.ID)
}
}