package service import ( "context" "encoding/json" "io" "net/http" "strings" "sync" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" ) const antigravityCreditsOveragesKey = "antigravity_credits_overages" type antigravity429Category string const ( antigravity429Unknown antigravity429Category = "unknown" antigravity429RateLimited antigravity429Category = "rate_limited" antigravity429QuotaExhausted antigravity429Category = "quota_exhausted" ) var ( creditsExhaustedCache sync.Map antigravityQuotaExhaustedKeywords = []string{ "quota_exhausted", "quota exhausted", } creditsExhaustedKeywords = []string{ "google_one_ai", "insufficient credit", "insufficient credits", "not enough credit", "not enough credits", "credit exhausted", "credits exhausted", "credit balance", "minimumcreditamountforusage", "minimum credit amount for usage", "minimum credit", } ) // isCreditsExhausted 检查账号的 AI Credits 是否已被标记为耗尽。 func isCreditsExhausted(accountID int64) bool { v, ok := creditsExhaustedCache.Load(accountID) if !ok { return false } until, ok := v.(time.Time) if !ok || time.Now().After(until) { creditsExhaustedCache.Delete(accountID) return false } return true } // setCreditsExhausted 将账号标记为 AI Credits 已耗尽,直到指定时间。 func setCreditsExhausted(accountID int64, until time.Time) { creditsExhaustedCache.Store(accountID, until) } // clearCreditsExhausted 清除账号的 AI Credits 耗尽标记。 func clearCreditsExhausted(accountID int64) { creditsExhaustedCache.Delete(accountID) } // classifyAntigravity429 将 Antigravity 的 429 响应归类为配额耗尽、限流或未知。 func classifyAntigravity429(body []byte) antigravity429Category { if len(body) == 0 { return antigravity429Unknown } lowerBody := strings.ToLower(string(body)) for _, keyword := range antigravityQuotaExhaustedKeywords { if strings.Contains(lowerBody, keyword) { return antigravity429QuotaExhausted } } if info := parseAntigravitySmartRetryInfo(body); info != nil && !info.IsModelCapacityExhausted { return antigravity429RateLimited } return antigravity429Unknown } // injectEnabledCreditTypes 在已序列化的 v1internal JSON body 中注入 AI Credits 类型。 func injectEnabledCreditTypes(body []byte) []byte { var payload map[string]any if err := json.Unmarshal(body, &payload); err != nil { return nil } payload["enabledCreditTypes"] = []string{"GOOGLE_ONE_AI"} result, err := json.Marshal(payload) if err != nil { return nil } return result } // resolveCreditsOveragesModelKey 解析当前请求对应的 overages 状态模型 key。 func resolveCreditsOveragesModelKey(ctx context.Context, account *Account, upstreamModelName, requestedModel string) string { modelKey := strings.TrimSpace(upstreamModelName) if modelKey != "" { return modelKey } if account == nil { return "" } modelKey = resolveFinalAntigravityModelKey(ctx, account, requestedModel) if strings.TrimSpace(modelKey) != "" { return modelKey } return resolveAntigravityModelKey(requestedModel) } // canUseAntigravityCreditsOverages 判断当前请求是否应直接走已激活的 overages 链路。 func canUseAntigravityCreditsOverages(ctx context.Context, account *Account, requestedModel string) bool { if account == nil || account.Platform != PlatformAntigravity { return false } if !account.IsSchedulable() { return false } if !account.IsOveragesEnabled() || isCreditsExhausted(account.ID) { return false } return account.getAntigravityCreditsOveragesRemainingTimeWithContext(ctx, requestedModel) > 0 } func (a *Account) getAntigravityCreditsOveragesRemainingTimeWithContext(ctx context.Context, requestedModel string) time.Duration { if a == nil || a.Extra == nil { return 0 } modelKey := resolveFinalAntigravityModelKey(ctx, a, requestedModel) modelKey = strings.TrimSpace(modelKey) if modelKey == "" { return 0 } rawStates, ok := a.Extra[antigravityCreditsOveragesKey].(map[string]any) if !ok { return 0 } rawState, ok := rawStates[modelKey].(map[string]any) if !ok { return 0 } activeUntilRaw, ok := rawState["active_until"].(string) if !ok || strings.TrimSpace(activeUntilRaw) == "" { return 0 } activeUntil, err := time.Parse(time.RFC3339, activeUntilRaw) if err != nil { return 0 } remaining := time.Until(activeUntil) if remaining > 0 { return remaining } return 0 } func setAntigravityCreditsOveragesActive(ctx context.Context, repo AccountRepository, account *Account, modelKey string, activeUntil time.Time) { if repo == nil || account == nil || account.ID == 0 || strings.TrimSpace(modelKey) == "" { return } stateMap := copyAntigravityCreditsOveragesState(account) stateMap[modelKey] = map[string]any{ "activated_at": time.Now().UTC().Format(time.RFC3339), "active_until": activeUntil.UTC().Format(time.RFC3339), } ensureAccountExtra(account) account.Extra[antigravityCreditsOveragesKey] = stateMap if err := repo.UpdateExtra(ctx, account.ID, map[string]any{antigravityCreditsOveragesKey: stateMap}); err != nil { logger.LegacyPrintf("service.antigravity_gateway", "set overages state failed: account=%d model=%s err=%v", account.ID, modelKey, err) } } func clearAntigravityCreditsOveragesState(ctx context.Context, repo AccountRepository, accountID int64) error { if repo == nil || accountID == 0 { return nil } return repo.UpdateExtra(ctx, accountID, map[string]any{ antigravityCreditsOveragesKey: map[string]any{}, }) } func clearAntigravityCreditsOveragesStateForModel(ctx context.Context, repo AccountRepository, account *Account, modelKey string) { if repo == nil || account == nil || account.ID == 0 { return } stateMap := copyAntigravityCreditsOveragesState(account) delete(stateMap, modelKey) ensureAccountExtra(account) account.Extra[antigravityCreditsOveragesKey] = stateMap if err := repo.UpdateExtra(ctx, account.ID, map[string]any{antigravityCreditsOveragesKey: stateMap}); err != nil { logger.LegacyPrintf("service.antigravity_gateway", "clear overages state failed: account=%d model=%s err=%v", account.ID, modelKey, err) } } func copyAntigravityCreditsOveragesState(account *Account) map[string]any { result := make(map[string]any) if account == nil || account.Extra == nil { return result } rawState, ok := account.Extra[antigravityCreditsOveragesKey].(map[string]any) if !ok { return result } for key, value := range rawState { result[key] = value } return result } func ensureAccountExtra(account *Account) { if account != nil && account.Extra == nil { account.Extra = make(map[string]any) } } // shouldMarkCreditsExhausted 判断一次 credits 请求失败是否应标记为 credits 耗尽。 func shouldMarkCreditsExhausted(resp *http.Response, respBody []byte, reqErr error) bool { if reqErr != nil || resp == nil { return false } if resp.StatusCode >= 500 || resp.StatusCode == http.StatusRequestTimeout { return false } if isURLLevelRateLimit(respBody) { return false } if info := parseAntigravitySmartRetryInfo(respBody); info != nil { return false } bodyLower := strings.ToLower(string(respBody)) for _, keyword := range creditsExhaustedKeywords { if strings.Contains(bodyLower, keyword) { return true } } return false } type creditsOveragesRetryResult struct { handled bool resp *http.Response } // attemptCreditsOveragesRetry 在确认免费配额耗尽后,尝试注入 AI Credits 继续请求。 func (s *AntigravityGatewayService) attemptCreditsOveragesRetry( p antigravityRetryLoopParams, baseURL string, modelName string, waitDuration time.Duration, originalStatusCode int, respBody []byte, ) *creditsOveragesRetryResult { creditsBody := injectEnabledCreditTypes(p.body) if creditsBody == nil { return &creditsOveragesRetryResult{handled: false} } modelKey := resolveCreditsOveragesModelKey(p.ctx, p.account, modelName, p.requestedModel) logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 credit_overages_retry model=%s account=%d (injecting enabledCreditTypes)", p.prefix, modelKey, p.account.ID) creditsReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, creditsBody) if err != nil { logger.LegacyPrintf("service.antigravity_gateway", "%s credit_overages_failed model=%s account=%d build_request_err=%v", p.prefix, modelKey, p.account.ID, err) return &creditsOveragesRetryResult{handled: true} } creditsResp, err := p.httpUpstream.Do(creditsReq, p.proxyURL, p.account.ID, p.account.Concurrency) if err == nil && creditsResp != nil && creditsResp.StatusCode < 400 { clearCreditsExhausted(p.account.ID) activeUntil := s.resolveCreditsOveragesActiveUntil(respBody, waitDuration) setAntigravityCreditsOveragesActive(p.ctx, p.accountRepo, p.account, modelKey, activeUntil) logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d credit_overages_success model=%s account=%d active_until=%s", p.prefix, creditsResp.StatusCode, modelKey, p.account.ID, activeUntil.UTC().Format(time.RFC3339)) return &creditsOveragesRetryResult{handled: true, resp: creditsResp} } s.handleCreditsRetryFailure(p.prefix, modelKey, p.account, waitDuration, creditsResp, err) return &creditsOveragesRetryResult{handled: true} } func (s *AntigravityGatewayService) handleCreditsRetryFailure( prefix string, modelKey string, account *Account, waitDuration time.Duration, creditsResp *http.Response, reqErr error, ) { var creditsRespBody []byte creditsStatusCode := 0 if creditsResp != nil { creditsStatusCode = creditsResp.StatusCode if creditsResp.Body != nil { creditsRespBody, _ = io.ReadAll(io.LimitReader(creditsResp.Body, 64<<10)) _ = creditsResp.Body.Close() } } if shouldMarkCreditsExhausted(creditsResp, creditsRespBody, reqErr) && account != nil { exhaustedUntil := s.resolveCreditsOveragesActiveUntil(creditsRespBody, waitDuration) setCreditsExhausted(account.ID, exhaustedUntil) clearAntigravityCreditsOveragesStateForModel(context.Background(), s.accountRepo, account, modelKey) logger.LegacyPrintf("service.antigravity_gateway", "%s credit_overages_failed model=%s account=%d marked_exhausted=true status=%d exhausted_until=%v body=%s", prefix, modelKey, account.ID, creditsStatusCode, exhaustedUntil, truncateForLog(creditsRespBody, 200)) return } if account != nil { logger.LegacyPrintf("service.antigravity_gateway", "%s credit_overages_failed model=%s account=%d marked_exhausted=false status=%d err=%v body=%s", prefix, modelKey, account.ID, creditsStatusCode, reqErr, truncateForLog(creditsRespBody, 200)) } } func (s *AntigravityGatewayService) resolveCreditsOveragesActiveUntil(respBody []byte, waitDuration time.Duration) time.Time { resetAt := ParseGeminiRateLimitResetTime(respBody) defaultDur := waitDuration if defaultDur <= 0 { defaultDur = s.getDefaultRateLimitDuration() } return s.resolveResetTime(resetAt, defaultDur) }