6e8eff9bb9
在运维监控的并发/排队卡片中，为 Antigravity 平台账号显示各 scope
(claude/gemini_text/gemini_image) 的限流数量统计，便于管理员了解
哪些 scope 正在被限流。
Cherry-picked from slovx2/sub2api: 08d6dc52
135 lines
3.5 KiB
Go
135 lines
3.5 KiB
Go
package service
|
|
|
|
import (
|
|
"slices"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
const antigravityQuotaScopesKey = "antigravity_quota_scopes"
|
|
|
|
// AntigravityQuotaScope 表示 Antigravity 的配额域
|
|
type AntigravityQuotaScope string
|
|
|
|
const (
|
|
AntigravityQuotaScopeClaude AntigravityQuotaScope = "claude"
|
|
AntigravityQuotaScopeGeminiText AntigravityQuotaScope = "gemini_text"
|
|
AntigravityQuotaScopeGeminiImage AntigravityQuotaScope = "gemini_image"
|
|
)
|
|
|
|
// IsScopeSupported 检查给定的 scope 是否在分组支持的 scope 列表中
|
|
func IsScopeSupported(supportedScopes []string, scope AntigravityQuotaScope) bool {
|
|
if len(supportedScopes) == 0 {
|
|
// 未配置时默认全部支持
|
|
return true
|
|
}
|
|
supported := slices.Contains(supportedScopes, string(scope))
|
|
return supported
|
|
}
|
|
|
|
// ResolveAntigravityQuotaScope 根据模型名称解析配额域（导出版本）
|
|
func ResolveAntigravityQuotaScope(requestedModel string) (AntigravityQuotaScope, bool) {
|
|
return resolveAntigravityQuotaScope(requestedModel)
|
|
}
|
|
|
|
// resolveAntigravityQuotaScope 根据模型名称解析配额域
|
|
func resolveAntigravityQuotaScope(requestedModel string) (AntigravityQuotaScope, bool) {
|
|
model := normalizeAntigravityModelName(requestedModel)
|
|
if model == "" {
|
|
return "", false
|
|
}
|
|
switch {
|
|
case strings.HasPrefix(model, "claude-"):
|
|
return AntigravityQuotaScopeClaude, true
|
|
case strings.HasPrefix(model, "gemini-"):
|
|
if isImageGenerationModel(model) {
|
|
return AntigravityQuotaScopeGeminiImage, true
|
|
}
|
|
return AntigravityQuotaScopeGeminiText, true
|
|
default:
|
|
return "", false
|
|
}
|
|
}
|
|
|
|
func normalizeAntigravityModelName(model string) string {
|
|
normalized := strings.ToLower(strings.TrimSpace(model))
|
|
normalized = strings.TrimPrefix(normalized, "models/")
|
|
return normalized
|
|
}
|
|
|
|
// IsSchedulableForModel 结合 Antigravity 配额域限流判断是否可调度
|
|
func (a *Account) IsSchedulableForModel(requestedModel string) bool {
|
|
if a == nil {
|
|
return false
|
|
}
|
|
if !a.IsSchedulable() {
|
|
return false
|
|
}
|
|
if a.isModelRateLimited(requestedModel) {
|
|
return false
|
|
}
|
|
if a.Platform != PlatformAntigravity {
|
|
return true
|
|
}
|
|
scope, ok := resolveAntigravityQuotaScope(requestedModel)
|
|
if !ok {
|
|
return true
|
|
}
|
|
resetAt := a.antigravityQuotaScopeResetAt(scope)
|
|
if resetAt == nil {
|
|
return true
|
|
}
|
|
now := time.Now()
|
|
return !now.Before(*resetAt)
|
|
}
|
|
|
|
func (a *Account) antigravityQuotaScopeResetAt(scope AntigravityQuotaScope) *time.Time {
|
|
if a == nil || a.Extra == nil || scope == "" {
|
|
return nil
|
|
}
|
|
rawScopes, ok := a.Extra[antigravityQuotaScopesKey].(map[string]any)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
rawScope, ok := rawScopes[string(scope)].(map[string]any)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
resetAtRaw, ok := rawScope["rate_limit_reset_at"].(string)
|
|
if !ok || strings.TrimSpace(resetAtRaw) == "" {
|
|
return nil
|
|
}
|
|
resetAt, err := time.Parse(time.RFC3339, resetAtRaw)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
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
|
|
}
|