fix(gemini): stop cooling Vertex 429s until PST midnight, honor RetryInfo

Vertex AI reports 429s with a standard google.rpc.RetryInfo detail
(retryDelay like "39s"), but ParseGeminiRateLimitResetTime only knew the
AI Studio quotaResetDelay metadata, so parsing always failed for Vertex
bodies. The fallback then lumped service_account accounts in with
API-key/AI-Studio ones and rate limited them until the next PST
midnight, so a transient per-minute throttle could take a pay-as-you-go
account (and with it the whole group) offline for up to a day.

- Parse google.rpc.RetryInfo retryDelay from 429 bodies, capped at 15
  minutes to guard against absurd upstream values.
- When no reset time can be parsed, cool service_account (Vertex)
  accounts down for 60s instead of until PST midnight: pay-as-you-go
  quotas are short-window RPM/TPM with no daily reset.

Fallback behavior for Code Assist, Google One and API-key accounts is
unchanged, covered by new regression tests.
This commit is contained in:
Erik
2026-09-22 11:31:31 +08:00
parent 1c0a69c0ce
commit 7f1f85cb76
3 changed files with 138 additions and 9 deletions
@@ -606,6 +606,77 @@ func TestHandleGeminiUpstreamError_PoolMode429(t *testing.T) {
}
}
// Vertex(service_account)429 响应带 google.rpc.RetryInfo 时,应按 retryDelay 冷却,
// 而不是回退到 PST 午夜。
func TestHandleGeminiUpstreamError_VertexRetryInfoUsesParsedDelay(t *testing.T) {
repo := &rateLimit429AccountRepoStub{}
svc := &GeminiMessagesCompatService{accountRepo: repo}
account := &Account{
ID: 601,
Platform: PlatformGemini,
Type: AccountTypeServiceAccount,
Credentials: map[string]any{
"project_id": "my-vertex-project",
"location": "global",
},
}
body := []byte(`{"error":{"code":429,"message":"Resource exhausted. Please try again later.","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"39s"}]}}`)
before := time.Now()
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
require.Equal(t, 1, repo.rateLimitCalls)
require.Equal(t, int64(601), repo.lastRateLimitID)
require.WithinDuration(t, before.Add(39*time.Second), repo.lastRateLimitReset, 2*time.Second)
}
// Vertex(service_account)429 响应不带任何可解析的重置时间时,走短冷却兜底,
// 不再冷却到 PST 午夜(按量付费没有每日配额,长冷却会导致整组不可用)。
func TestHandleGeminiUpstreamError_VertexFallbackUsesShortCooldown(t *testing.T) {
repo := &rateLimit429AccountRepoStub{}
svc := &GeminiMessagesCompatService{accountRepo: repo}
account := &Account{
ID: 602,
Platform: PlatformGemini,
Type: AccountTypeServiceAccount,
Credentials: map[string]any{
"project_id": "my-vertex-project",
},
}
body := []byte(`{"error":{"code":429,"message":"Resource exhausted. Please try again later.","status":"RESOURCE_EXHAUSTED"}}`)
before := time.Now()
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
require.Equal(t, 1, repo.rateLimitCalls)
require.Equal(t, int64(602), repo.lastRateLimitID)
require.WithinDuration(t, before.Add(geminiVertexFallbackCooldown), repo.lastRateLimitReset, 2*time.Second)
// 显式守护:不应被冷却超过 15 分钟(更不应到次日 PST 午夜)
require.True(t, repo.lastRateLimitReset.Before(before.Add(15*time.Minute)))
}
// API Key(AI Studio)无法解析重置时间时仍冷却到 PST 午夜,行为保持不变。
func TestHandleGeminiUpstreamError_APIKeyFallbackStillPSTMidnight(t *testing.T) {
repo := &rateLimit429AccountRepoStub{}
svc := &GeminiMessagesCompatService{accountRepo: repo}
account := &Account{
ID: 603,
Platform: PlatformGemini,
Type: AccountTypeAPIKey,
}
body := []byte(`{"error":{"code":429,"message":"rate limit"}}`)
svc.handleGeminiUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{}, body)
require.Equal(t, 1, repo.rateLimitCalls)
require.Equal(t, int64(603), repo.lastRateLimitID)
expected := time.Unix(*nextGeminiDailyResetUnix(), 0)
require.WithinDuration(t, expected, repo.lastRateLimitReset, 5*time.Second)
}
type geminiErrorPolicyRepo struct {
mockAccountRepoForGemini
setErrorCalls int
@@ -40,6 +40,16 @@ const (
geminiRetryMaxDelay = 16 * time.Second
)
const (
// google.rpc.RetryInfo 的标准 @type(Vertex AI 429 响应携带,retryDelay 形如 "39s")
geminiRetryInfoTypeURL = "type.googleapis.com/google.rpc.RetryInfo"
// RetryInfo.retryDelay 的上限,防止异常值导致账号长时间不可用
geminiRetryInfoMaxDelay = 15 * time.Minute
// Vertex(service_account,按量付费)429 无法解析重置时间时的兜底冷却:
// 按量付费限速是短窗口 RPM/TPM,没有每日配额,不适用 PST 午夜重置
geminiVertexFallbackCooldown = time.Minute
)
// Gemini tool calling now requires `thoughtSignature` in parts that include `functionCall`.
// Many clients don't send it; we inject a known dummy signature to satisfy the validator.
// Ref: https://ai.google.dev/gemini-api/docs/thought-signatures
@@ -3140,6 +3150,10 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont
} 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 if account.IsVertexServiceAccount() {
// Vertex 按量付费:短窗限速,无每日配额,短冷却后即可恢复
ra = time.Now().Add(geminiVertexFallbackCooldown)
logger.LegacyPrintf("service.gemini_messages_compat", "[Gemini 429] Account %d (Vertex service account, project=%s) rate limited, cooldown=%v", account.ID, projectID, time.Until(ra).Truncate(time.Second))
} else {
// API Key / AI Studio OAuth: PST 午夜
if ts := nextGeminiDailyResetUnix(); ts != nil {
@@ -3172,19 +3186,29 @@ func ParseGeminiRateLimitResetTime(body []byte) *int64 {
}
}
// 遍历 error.details 查找 quotaResetDelay
// 遍历 error.details 查找 quotaResetDelay(AI Studio)或 RetryInfo.retryDelay(Vertex)
var found *int64
gjson.GetBytes(body, "error.details").ForEach(func(_, detail gjson.Result) bool {
v := detail.Get("metadata.quotaResetDelay").String()
if v == "" {
if v := detail.Get("metadata.quotaResetDelay").String(); v != "" {
if dur, err := time.ParseDuration(v); err == nil {
// Use ceil to avoid undercounting fractional seconds (e.g. 10.1s should not become 10s),
// which can affect scheduling decisions around thresholds (like 10s).
ts := time.Now().Unix() + int64(math.Ceil(dur.Seconds()))
found = &ts
return false
}
return true
}
if dur, err := time.ParseDuration(v); err == nil {
// Use ceil to avoid undercounting fractional seconds (e.g. 10.1s should not become 10s),
// which can affect scheduling decisions around thresholds (like 10s).
ts := time.Now().Unix() + int64(math.Ceil(dur.Seconds()))
found = &ts
return false
// Vertex 标准错误格式:google.rpc.RetryInfo.retryDelay("39s" / "1.5s")
if detail.Get("@type").String() == geminiRetryInfoTypeURL {
if dur, err := time.ParseDuration(detail.Get("retryDelay").String()); err == nil && dur > 0 {
if dur > geminiRetryInfoMaxDelay {
dur = geminiRetryInfoMaxDelay
}
ts := time.Now().Unix() + int64(math.Ceil(dur.Seconds()))
found = &ts
return false
}
}
return true
})
@@ -1017,6 +1017,40 @@ func TestParseGeminiRateLimitResetTime(t *testing.T) {
wantNil: false,
approxDelta: 10,
},
{
name: "Vertex RetryInfo retryDelay",
input: `{"error":{"code":429,"message":"Resource exhausted. Please try again later.","status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"39s"}]}}`,
wantNil: false,
approxDelta: 39,
},
{
name: "Vertex RetryInfo 小数秒向上取整",
input: `{"error":{"code":429,"status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"1.5s"}]}}`,
wantNil: false,
approxDelta: 2,
},
{
name: "Vertex RetryInfo 超过上限被截断到 15 分钟",
input: `{"error":{"code":429,"status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"3600s"}]}}`,
wantNil: false,
approxDelta: 900,
},
{
name: "retryDelay 存在但 @type 不是 RetryInfo 则忽略",
input: `{"error":{"code":429,"details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","retryDelay":"39s"}]}}`,
wantNil: true,
},
{
name: "RetryInfo retryDelay 非法值忽略",
input: `{"error":{"code":429,"details":[{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"soon"}]}}`,
wantNil: true,
},
{
name: "quotaResetDelay 与 RetryInfo 并存时先到者优先",
input: `{"error":{"code":429,"details":[{"metadata":{"quotaResetDelay":"120s"}},{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"39s"}]}}`,
wantNil: false,
approxDelta: 120,
},
}
for _, tt := range tests {