fix: use tier cooldown for google one gemini 429

This commit is contained in:
yangzc2004-bit 2026-05-16 00:16:38 +08:00
parent 6e66edbb09
commit 87fac30459
2 changed files with 38 additions and 3 deletions

View File

@ -383,6 +383,37 @@ func TestGeminiErrorPolicy_NilRateLimitService(t *testing.T) {
// policy tests. Embeds mockAccountRepoForGemini and adds tracking.
// ---------------------------------------------------------------------------
func TestHandleGeminiUpstreamError_GoogleOneCapacityExhaustedUsesTierCooldown(t *testing.T) {
repo := &rateLimit429AccountRepoStub{}
quotaSvc := NewGeminiQuotaService(&config.Config{}, nil)
rlSvc := NewRateLimitService(repo, nil, &config.Config{}, quotaSvc, nil)
svc := &GeminiMessagesCompatService{
accountRepo: repo,
rateLimitService: rlSvc,
}
account := &Account{
ID: 511,
Platform: PlatformGemini,
Type: AccountTypeOAuth,
Credentials: map[string]any{
"oauth_type": "google_one",
"tier_id": "google_ai_pro",
},
}
body := []byte(`{"error":{"code":429,"details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","domain":"cloudcode-pa.googleapis.com","metadata":{"model":"gemini-3.1-pro-preview"},"reason":"MODEL_CAPACITY_EXHAUSTED"}],"message":"No capacity available for model gemini-3.1-pro-preview on the server","status":"RESOURCE_EXHAUSTED"}}`)
before := time.Now()
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
after := time.Now()
require.Equal(t, 1, repo.rateLimitCalls)
require.Equal(t, int64(511), repo.lastRateLimitID)
require.WithinDuration(t, before.Add(5*time.Minute), repo.lastRateLimitReset, 2*time.Second)
require.True(t, repo.lastRateLimitReset.After(before))
require.True(t, repo.lastRateLimitReset.Before(after.Add(5*time.Minute).Add(2*time.Second)))
}
type geminiErrorPolicyRepo struct {
mockAccountRepoForGemini
setErrorCalls int

View File

@ -2822,14 +2822,18 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont
if resetAt == nil {
// 根据账号类型使用不同的默认重置时间
var ra time.Time
if isCodeAssist {
// Code Assist: fallback cooldown by tier
if isCodeAssist || oauthType == "google_one" {
// Gemini CLI / Google One: fallback cooldown by tier
cooldown := geminiCooldownForTier(tierID)
if s.rateLimitService != nil {
cooldown = s.rateLimitService.GeminiCooldown(ctx, account)
}
ra = time.Now().Add(cooldown)
logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (Code Assist, tier=%s, project=%s) rate limited, cooldown=%v", account.ID, tierID, projectID, time.Until(ra).Truncate(time.Second))
if isCodeAssist {
logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (Code Assist, tier=%s, project=%s) rate limited, cooldown=%v", account.ID, tierID, projectID, time.Until(ra).Truncate(time.Second))
} else {
logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (Google One OAuth, tier=%s, project=%s) rate limited, cooldown=%v", account.ID, tierID, projectID, time.Until(ra).Truncate(time.Second))
}
} else {
// API Key / AI Studio OAuth: PST 午夜
if ts := nextGeminiDailyResetUnix(); ts != nil {