mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 20:12:39 +08:00
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:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user