Merge pull request #464 from touwaeriol/pr/antigravity-scope-ratelimit

feat(antigravity): 支持按配额域(scope)级别限流

backend/internal/handler/dto/mappers.go

@@ -204,6 +204,17 @@ func AccountFromServiceShallow(a *service.Account) *Account {
 		}
 	}
 
+	if scopeLimits := a.GetAntigravityScopeRateLimits(); len(scopeLimits) > 0 {
+		out.ScopeRateLimits = make(map[string]ScopeRateLimitInfo, len(scopeLimits))
+		now := time.Now()
+		for scope, remainingSec := range scopeLimits {
+			out.ScopeRateLimits[scope] = ScopeRateLimitInfo{
+				ResetAt:      now.Add(time.Duration(remainingSec) * time.Second),
+				RemainingSec: remainingSec,
+			}
+		}
+	}
+
 	return out
 }
 

backend/internal/handler/dto/types.go

@@ -2,6 +2,11 @@ package dto
 
 import "time"
 
+type ScopeRateLimitInfo struct {
+	ResetAt      time.Time `json:"reset_at"`
+	RemainingSec int64     `json:"remaining_sec"`
+}
+
 type User struct {
 	ID            int64     `json:"id"`
 	Email         string    `json:"email"`
@@ -108,6 +113,9 @@ type Account struct {
 	RateLimitResetAt *time.Time `json:"rate_limit_reset_at"`
 	OverloadUntil    *time.Time `json:"overload_until"`
 
+	// Antigravity scope 级限流状态（从 extra 提取）
+	ScopeRateLimits map[string]ScopeRateLimitInfo `json:"scope_rate_limits,omitempty"`
+
 	TempUnschedulableUntil  *time.Time `json:"temp_unschedulable_until"`
 	TempUnschedulableReason string     `json:"temp_unschedulable_reason"`
 

backend/internal/repository/account_repo.go

@@ -809,12 +809,21 @@ func (r *accountRepository) SetAntigravityQuotaScopeLimit(ctx context.Context, i
 		return err
 	}
 
-	path := "{antigravity_quota_scopes," + string(scope) + "}"
+	scopeKey := string(scope)
 	client := clientFromContext(ctx, r.client)
 	result, err := client.ExecContext(
 		ctx,
-		"UPDATE accounts SET extra = jsonb_set(COALESCE(extra, '{}'::jsonb), $1::text[], $2::jsonb, true), updated_at = NOW() WHERE id = $3 AND deleted_at IS NULL",
-		path,
+		`UPDATE accounts SET
+			extra = jsonb_set(
+				jsonb_set(COALESCE(extra, '{}'::jsonb), '{antigravity_quota_scopes}'::text[], COALESCE(extra->'antigravity_quota_scopes', '{}'::jsonb), true),
+				ARRAY['antigravity_quota_scopes', $1]::text[],
+				$2::jsonb,
+				true
+			),
+			updated_at = NOW(),
+			last_used_at = NOW()
+		WHERE id = $3 AND deleted_at IS NULL`,
+		scopeKey,
 		raw,
 		id,
 	)
@@ -829,6 +838,7 @@ func (r *accountRepository) SetAntigravityQuotaScopeLimit(ctx context.Context, i
 	if affected == 0 {
 		return service.ErrAccountNotFound
 	}
+
 	if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil {
 		log.Printf("[SchedulerOutbox] enqueue quota scope failed: account=%d err=%v", id, err)
 	}
@@ -849,12 +859,19 @@ func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, sco
 		return err
 	}
 
-	path := "{model_rate_limits," + scope + "}"
 	client := clientFromContext(ctx, r.client)
 	result, err := client.ExecContext(
 		ctx,
-		"UPDATE accounts SET extra = jsonb_set(COALESCE(extra, '{}'::jsonb), $1::text[], $2::jsonb, true), updated_at = NOW() WHERE id = $3 AND deleted_at IS NULL",
-		path,
+		`UPDATE accounts SET 
+			extra = jsonb_set(
+				jsonb_set(COALESCE(extra, '{}'::jsonb), '{model_rate_limits}'::text[], COALESCE(extra->'model_rate_limits', '{}'::jsonb), true),
+				ARRAY['model_rate_limits', $1]::text[],
+				$2::jsonb,
+				true
+			),
+			updated_at = NOW()
+		WHERE id = $3 AND deleted_at IS NULL`,
+		scope,
 		raw,
 		id,
 	)

backend/internal/service/antigravity_gateway_service.go

@@ -1537,7 +1537,11 @@ func sleepAntigravityBackoffWithContext(ctx context.Context, attempt int) bool {
 
 func antigravityUseScopeRateLimit() bool {
 	v := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityScopeRateLimitEnv)))
-	return v == "1" || v == "true" || v == "yes" || v == "on"
+	// 默认开启按配额域限流，只有明确设置为禁用值时才关闭
+	if v == "0" || v == "false" || v == "no" || v == "off" {
+		return false
+	}
+	return true
 }
 
 func (s *AntigravityGatewayService) handleUpstreamError(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope) {

backend/internal/service/antigravity_quota_scope.go

@@ -89,3 +89,30 @@ func (a *Account) antigravityQuotaScopeResetAt(scope AntigravityQuotaScope) *tim
 	}
 	return &resetAt
 }
+
+var antigravityAllScopes = []AntigravityQuotaScope{
+	AntigravityQuotaScopeClaude,
+	AntigravityQuotaScopeGeminiText,
+	AntigravityQuotaScopeGeminiImage,
+}
+
+func (a *Account) GetAntigravityScopeRateLimits() map[string]int64 {
+	if a == nil || a.Platform != PlatformAntigravity {
+		return nil
+	}
+	now := time.Now()
+	result := make(map[string]int64)
+	for _, scope := range antigravityAllScopes {
+		resetAt := a.antigravityQuotaScopeResetAt(scope)
+		if resetAt != nil && now.Before(*resetAt) {
+			remainingSec := int64(time.Until(*resetAt).Seconds())
+			if remainingSec > 0 {
+				result[string(scope)] = remainingSec
+			}
+		}
+	}
+	if len(result) == 0 {
+		return nil
+	}
+	return result
+}

backend/internal/service/ops_account_availability.go

@@ -67,6 +67,8 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
 
 		isAvailable := acc.Status == StatusActive && acc.Schedulable && !isRateLimited && !isOverloaded && !isTempUnsched
 
+		scopeRateLimits := acc.GetAntigravityScopeRateLimits()
+
 		if acc.Platform != "" {
 			if _, ok := platform[acc.Platform]; !ok {
 				platform[acc.Platform] = &PlatformAvailability{
@@ -84,6 +86,14 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
 			if hasError {
 				p.ErrorCount++
 			}
+			if len(scopeRateLimits) > 0 {
+				if p.ScopeRateLimitCount == nil {
+					p.ScopeRateLimitCount = make(map[string]int64)
+				}
+				for scope := range scopeRateLimits {
+					p.ScopeRateLimitCount[scope]++
+				}
+			}
 		}
 
 		for _, grp := range acc.Groups {
@@ -108,6 +118,14 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
 			if hasError {
 				g.ErrorCount++
 			}
+			if len(scopeRateLimits) > 0 {
+				if g.ScopeRateLimitCount == nil {
+					g.ScopeRateLimitCount = make(map[string]int64)
+				}
+				for scope := range scopeRateLimits {
+					g.ScopeRateLimitCount[scope]++
+				}
+			}
 		}
 
 		displayGroupID := int64(0)
@@ -140,6 +158,9 @@ func (s *OpsService) GetAccountAvailabilityStats(ctx context.Context, platformFi
 				item.RateLimitRemainingSec = &remainingSec
 			}
 		}
+		if len(scopeRateLimits) > 0 {
+			item.ScopeRateLimits = scopeRateLimits
+		}
 		if isOverloaded && acc.OverloadUntil != nil {
 			item.OverloadUntil = acc.OverloadUntil
 			remainingSec := int64(time.Until(*acc.OverloadUntil).Seconds())

backend/internal/service/ops_realtime_models.go

@@ -39,22 +39,24 @@ type AccountConcurrencyInfo struct {
 
 // PlatformAvailability aggregates account availability by platform.
 type PlatformAvailability struct {
-	Platform       string `json:"platform"`
-	TotalAccounts  int64  `json:"total_accounts"`
-	AvailableCount int64  `json:"available_count"`
-	RateLimitCount int64  `json:"rate_limit_count"`
-	ErrorCount     int64  `json:"error_count"`
+	Platform            string           `json:"platform"`
+	TotalAccounts       int64            `json:"total_accounts"`
+	AvailableCount      int64            `json:"available_count"`
+	RateLimitCount      int64            `json:"rate_limit_count"`
+	ScopeRateLimitCount map[string]int64 `json:"scope_rate_limit_count,omitempty"`
+	ErrorCount          int64            `json:"error_count"`
 }
 
 // GroupAvailability aggregates account availability by group.
 type GroupAvailability struct {
-	GroupID        int64  `json:"group_id"`
-	GroupName      string `json:"group_name"`
-	Platform       string `json:"platform"`
-	TotalAccounts  int64  `json:"total_accounts"`
-	AvailableCount int64  `json:"available_count"`
-	RateLimitCount int64  `json:"rate_limit_count"`
-	ErrorCount     int64  `json:"error_count"`
+	GroupID             int64            `json:"group_id"`
+	GroupName           string           `json:"group_name"`
+	Platform            string           `json:"platform"`
+	TotalAccounts       int64            `json:"total_accounts"`
+	AvailableCount      int64            `json:"available_count"`
+	RateLimitCount      int64            `json:"rate_limit_count"`
+	ScopeRateLimitCount map[string]int64 `json:"scope_rate_limit_count,omitempty"`
+	ErrorCount          int64            `json:"error_count"`
 }
 
 // AccountAvailability represents current availability for a single account.
@@ -72,10 +74,11 @@ type AccountAvailability struct {
 	IsOverloaded  bool `json:"is_overloaded"`
 	HasError      bool `json:"has_error"`
 
-	RateLimitResetAt       *time.Time `json:"rate_limit_reset_at"`
-	RateLimitRemainingSec  *int64     `json:"rate_limit_remaining_sec"`
-	OverloadUntil          *time.Time `json:"overload_until"`
-	OverloadRemainingSec   *int64     `json:"overload_remaining_sec"`
-	ErrorMessage           string     `json:"error_message"`
-	TempUnschedulableUntil *time.Time `json:"temp_unschedulable_until,omitempty"`
+	RateLimitResetAt       *time.Time       `json:"rate_limit_reset_at"`
+	RateLimitRemainingSec  *int64           `json:"rate_limit_remaining_sec"`
+	ScopeRateLimits        map[string]int64 `json:"scope_rate_limits,omitempty"`
+	OverloadUntil          *time.Time       `json:"overload_until"`
+	OverloadRemainingSec   *int64           `json:"overload_remaining_sec"`
+	ErrorMessage           string           `json:"error_message"`
+	TempUnschedulableUntil *time.Time       `json:"temp_unschedulable_until,omitempty"`
 }