a01686c637
Squash of 4 commits: - Fix Gemini rate limit scheduling - fix antigravity gemini rate limit scheduling - fix antigravity gemini limited account scheduling - fix antigravity test stubs for default lint
130 lines
3.6 KiB
Go
130 lines
3.6 KiB
Go
package service
|
||
|
||
import (
|
||
"context"
|
||
"strings"
|
||
"time"
|
||
)
|
||
|
||
const (
|
||
modelRateLimitsKey = "model_rate_limits"
|
||
antigravityGeminiModelRateLimitKey = "antigravity:gemini"
|
||
)
|
||
|
||
// isRateLimitActiveForKey 检查指定 key 的限流是否生效
|
||
func (a *Account) isRateLimitActiveForKey(key string) bool {
|
||
resetAt := a.modelRateLimitResetAt(key)
|
||
return resetAt != nil && time.Now().Before(*resetAt)
|
||
}
|
||
|
||
// getRateLimitRemainingForKey 获取指定 key 的限流剩余时间，0 表示未限流或已过期
|
||
func (a *Account) getRateLimitRemainingForKey(key string) time.Duration {
|
||
resetAt := a.modelRateLimitResetAt(key)
|
||
if resetAt == nil {
|
||
return 0
|
||
}
|
||
remaining := time.Until(*resetAt)
|
||
if remaining > 0 {
|
||
return remaining
|
||
}
|
||
return 0
|
||
}
|
||
|
||
func (a *Account) isModelRateLimitedWithContext(ctx context.Context, requestedModel string) bool {
|
||
if a == nil {
|
||
return false
|
||
}
|
||
|
||
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 == "" {
|
||
return false
|
||
}
|
||
return a.isRateLimitActiveForKey(modelKey)
|
||
}
|
||
|
||
// GetModelRateLimitRemainingTime 获取模型限流剩余时间
|
||
// 返回 0 表示未限流或已过期
|
||
func (a *Account) GetModelRateLimitRemainingTime(requestedModel string) time.Duration {
|
||
return a.GetModelRateLimitRemainingTimeWithContext(context.Background(), requestedModel)
|
||
}
|
||
|
||
func (a *Account) GetModelRateLimitRemainingTimeWithContext(ctx context.Context, requestedModel string) time.Duration {
|
||
if a == nil {
|
||
return 0
|
||
}
|
||
|
||
modelKey := a.GetMappedModel(requestedModel)
|
||
if a.Platform == PlatformAntigravity {
|
||
modelKey = resolveFinalAntigravityModelKey(ctx, a, requestedModel)
|
||
}
|
||
modelKey = strings.TrimSpace(modelKey)
|
||
if modelKey == "" {
|
||
return 0
|
||
}
|
||
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 {
|
||
modelKey := mapAntigravityModel(account, requestedModel)
|
||
if modelKey == "" {
|
||
return ""
|
||
}
|
||
// thinking 会影响 Antigravity 最终模型名（例如 claude-sonnet-4-5 -> claude-sonnet-4-5-thinking）
|
||
if enabled, ok := ThinkingEnabledFromContext(ctx); ok {
|
||
modelKey = applyThinkingModelSuffix(modelKey, enabled)
|
||
}
|
||
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
|
||
}
|
||
rawLimits, ok := a.Extra[modelRateLimitsKey].(map[string]any)
|
||
if !ok {
|
||
return nil
|
||
}
|
||
rawLimit, ok := rawLimits[scope].(map[string]any)
|
||
if !ok {
|
||
return nil
|
||
}
|
||
resetAtRaw, ok := rawLimit["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
|
||
}
|