mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-19 16:17:56 +00:00
fix(gemini): 完善gemini号池模式时retryable失效问题
此前 Gemini 三条转发路径对池模式账号命中 ErrorPolicySkipped 时直接把 上游错误体透传给客户端(强制标 500),既不同账号重试也不换号, pool_mode_retry_count 配置完全不生效;而 Anthropic/OpenAI 等路径 会构造 UpstreamFailoverError 交给 handler 层按池模式配置重试后换号。 - 新增 poolModeSkippedFailoverError:池模式 + 可 failover 状态码时 返回 UpstreamFailoverError,RetryableOnSameAccount 按 pool_mode_retry_status_codes(默认 401/403/429)判定 - 原生 v1beta(ForwardNative)与 Claude messages 兼容路径的 Skipped 分支接入;非池模式的自定义错误码透传行为不变 - chat completions 兼容路径的 failover 错误补上 RetryableOnSameAccount 池模式判定
This commit is contained in:
@@ -231,7 +231,11 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
})
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: evBody}
|
||||
return nil, &UpstreamFailoverError{
|
||||
StatusCode: resp.StatusCode,
|
||||
ResponseBody: evBody,
|
||||
RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode),
|
||||
}
|
||||
}
|
||||
|
||||
return nil, s.writeGeminiChatCompletionsMappedError(c, account, resp.StatusCode, requestID, evBody)
|
||||
|
||||
@@ -352,6 +352,68 @@ func TestGeminiErrorPolicyIntegration(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestPoolModeSkippedFailoverError — pool-mode accounts hitting
|
||||
// ErrorPolicySkipped must failover (align with other platform forwards)
|
||||
// instead of passing the upstream error through to the client.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
func TestPoolModeSkippedFailoverError(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := &GeminiMessagesCompatService{}
|
||||
|
||||
poolAccount := func(extra map[string]any) *Account {
|
||||
creds := map[string]any{"pool_mode": true}
|
||||
for k, v := range extra {
|
||||
creds[k] = v
|
||||
}
|
||||
return &Account{ID: 300, Type: AccountTypeAPIKey, Platform: PlatformGemini, Credentials: creds}
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
account *Account
|
||||
statusCode int
|
||||
expectFailover bool
|
||||
expectSameAccount bool
|
||||
}{
|
||||
{"pool_500_failover_no_same_account_retry", poolAccount(nil), 500, true, false},
|
||||
{"pool_429_failover_with_same_account_retry", poolAccount(nil), 429, true, true},
|
||||
{"pool_custom_retry_codes_500", poolAccount(map[string]any{
|
||||
"pool_mode_retry_status_codes": []any{float64(500)},
|
||||
}), 500, true, true},
|
||||
{"pool_400_not_failover_worthy", poolAccount(nil), 400, false, false},
|
||||
{"non_pool_account_keeps_passthrough", &Account{
|
||||
ID: 301, Type: AccountTypeAPIKey, Platform: PlatformGemini,
|
||||
Credentials: map[string]any{
|
||||
"custom_error_codes_enabled": true,
|
||||
"custom_error_codes": []any{float64(429)},
|
||||
},
|
||||
}, 500, false, false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
writer := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(writer)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
|
||||
body := []byte(`{"error":{"code":"bad_response_status_code","message":"openai_error"}}`)
|
||||
failoverErr := svc.poolModeSkippedFailoverError(c, tt.account, tt.statusCode, body, "req-1")
|
||||
|
||||
if !tt.expectFailover {
|
||||
require.Nil(t, failoverErr)
|
||||
return
|
||||
}
|
||||
require.NotNil(t, failoverErr)
|
||||
require.Equal(t, tt.statusCode, failoverErr.StatusCode)
|
||||
require.Equal(t, body, failoverErr.ResponseBody)
|
||||
require.Equal(t, tt.expectSameAccount, failoverErr.RetryableOnSameAccount)
|
||||
require.True(t, failoverErr.ShouldRetryNextAccount())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TestGeminiErrorPolicy_NilRateLimitService — verifies nil safety
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -945,6 +945,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
||||
if upstreamReqID == "" {
|
||||
upstreamReqID = resp.Header.Get("x-goog-request-id")
|
||||
}
|
||||
if failoverErr := s.poolModeSkippedFailoverError(c, account, resp.StatusCode, respBody, upstreamReqID); failoverErr != nil {
|
||||
return nil, failoverErr
|
||||
}
|
||||
return nil, s.writeGeminiMappedError(c, account, http.StatusInternalServerError, upstreamReqID, respBody)
|
||||
case ErrorPolicyMatched, ErrorPolicyTempUnscheduled:
|
||||
if policy == ErrorPolicyMatched {
|
||||
@@ -1452,6 +1455,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel)
|
||||
switch policy {
|
||||
case ErrorPolicySkipped:
|
||||
if failoverErr := s.poolModeSkippedFailoverError(c, account, resp.StatusCode, respBody, requestID); failoverErr != nil {
|
||||
return nil, failoverErr
|
||||
}
|
||||
respBody = unwrapIfNeeded(isOAuth, respBody)
|
||||
contentType := resp.Header.Get("Content-Type")
|
||||
if contentType == "" {
|
||||
@@ -1683,6 +1689,39 @@ func (s *GeminiMessagesCompatService) shouldFailoverGeminiUpstreamError(statusCo
|
||||
}
|
||||
}
|
||||
|
||||
// poolModeSkippedFailoverError 池模式账号命中 ErrorPolicySkipped 时构造 failover 错误:
|
||||
// 可 failover 的状态码返回 UpstreamFailoverError,交给 handler 层按 pool_mode_retry_count
|
||||
// 同账号重试后换号;返回 nil 表示不适用(非池模式或状态码不可 failover),由调用方透传。
|
||||
func (s *GeminiMessagesCompatService) poolModeSkippedFailoverError(c *gin.Context, account *Account, statusCode int, respBody []byte, upstreamRequestID string) *UpstreamFailoverError {
|
||||
if !account.IsPoolMode() || !s.shouldFailoverGeminiUpstreamError(statusCode) {
|
||||
return nil
|
||||
}
|
||||
upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody)))
|
||||
upstreamDetail := ""
|
||||
if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
|
||||
maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 2048
|
||||
}
|
||||
upstreamDetail = truncateString(string(respBody), maxBytes)
|
||||
}
|
||||
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
|
||||
Platform: account.Platform,
|
||||
AccountID: account.ID,
|
||||
AccountName: account.Name,
|
||||
UpstreamStatusCode: statusCode,
|
||||
UpstreamRequestID: upstreamRequestID,
|
||||
Kind: "failover",
|
||||
Message: upstreamMsg,
|
||||
Detail: upstreamDetail,
|
||||
})
|
||||
return &UpstreamFailoverError{
|
||||
StatusCode: statusCode,
|
||||
ResponseBody: respBody,
|
||||
RetryableOnSameAccount: account.IsPoolModeRetryableStatus(statusCode),
|
||||
}
|
||||
}
|
||||
|
||||
func sleepGeminiBackoff(attempt int) {
|
||||
delay := geminiRetryBaseDelay * time.Duration(1<<uint(attempt-1))
|
||||
if delay > geminiRetryMaxDelay {
|
||||
|
||||
Reference in New Issue
Block a user