feat(antigravity): 支持按模型类型配置重试次数

新增环境变量:
- GATEWAY_ANTIGRAVITY_MAX_RETRIES_CLAUDE
- GATEWAY_ANTIGRAVITY_MAX_RETRIES_GEMINI_TEXT
- GATEWAY_ANTIGRAVITY_MAX_RETRIES_GEMINI_IMAGE

未设置时回退到平台级 GATEWAY_ANTIGRAVITY_MAX_RETRIES
This commit is contained in:
song
2026-01-21 20:48:36 +08:00
parent 52c745bc62
commit 207e09500a

View File

@@ -31,10 +31,13 @@ const (
)
const (
antigravityMaxRetriesEnv = "GATEWAY_ANTIGRAVITY_MAX_RETRIES"
antigravityScopeRateLimitEnv = "GATEWAY_ANTIGRAVITY_429_SCOPE_LIMIT"
antigravityBillingModelEnv = "GATEWAY_ANTIGRAVITY_BILL_WITH_MAPPED_MODEL"
antigravityFallbackSecondsEnv = "GATEWAY_ANTIGRAVITY_FALLBACK_COOLDOWN_SECONDS"
antigravityMaxRetriesEnv = "GATEWAY_ANTIGRAVITY_MAX_RETRIES"
antigravityMaxRetriesClaudeEnv = "GATEWAY_ANTIGRAVITY_MAX_RETRIES_CLAUDE"
antigravityMaxRetriesGeminiTextEnv = "GATEWAY_ANTIGRAVITY_MAX_RETRIES_GEMINI_TEXT"
antigravityMaxRetriesGeminiImageEnv = "GATEWAY_ANTIGRAVITY_MAX_RETRIES_GEMINI_IMAGE"
antigravityScopeRateLimitEnv = "GATEWAY_ANTIGRAVITY_429_SCOPE_LIMIT"
antigravityBillingModelEnv = "GATEWAY_ANTIGRAVITY_BILL_WITH_MAPPED_MODEL"
antigravityFallbackSecondsEnv = "GATEWAY_ANTIGRAVITY_FALLBACK_COOLDOWN_SECONDS"
)
// antigravityRetryLoopParams 重试循环的参数
@@ -51,6 +54,7 @@ type antigravityRetryLoopParams struct {
httpUpstream HTTPUpstream
settingService *SettingService
handleError func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope)
maxRetries int // 可选0 表示使用平台级默认值
}
// antigravityRetryLoopResult 重试循环的结果
@@ -64,7 +68,10 @@ func antigravityRetryLoop(p antigravityRetryLoopParams) (*antigravityRetryLoopRe
if len(availableURLs) == 0 {
availableURLs = antigravity.BaseURLs
}
maxRetries := antigravityMaxRetries()
maxRetries := p.maxRetries
if maxRetries <= 0 {
maxRetries = antigravityMaxRetries()
}
var resp *http.Response
var usedBaseURL string
@@ -770,6 +777,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
httpUpstream: s.httpUpstream,
settingService: s.settingService,
handleError: s.handleUpstreamError,
maxRetries: antigravityMaxRetriesForModel(originalModel),
})
if err != nil {
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries")
@@ -846,6 +854,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
httpUpstream: s.httpUpstream,
settingService: s.settingService,
handleError: s.handleUpstreamError,
maxRetries: antigravityMaxRetriesForModel(originalModel),
})
if retryErr != nil {
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
@@ -1352,6 +1361,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
httpUpstream: s.httpUpstream,
settingService: s.settingService,
handleError: s.handleUpstreamError,
maxRetries: antigravityMaxRetriesForModel(originalModel),
})
if err != nil {
return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries")
@@ -1560,6 +1570,28 @@ func antigravityMaxRetries() int {
return value
}
// antigravityMaxRetriesForModel 根据模型类型获取重试次数
// 优先使用模型细分配置,未设置则回退到平台级配置
func antigravityMaxRetriesForModel(model string) int {
var envKey string
if strings.HasPrefix(model, "claude-") {
envKey = antigravityMaxRetriesClaudeEnv
} else if isImageGenerationModel(model) {
envKey = antigravityMaxRetriesGeminiImageEnv
} else if strings.HasPrefix(model, "gemini-") {
envKey = antigravityMaxRetriesGeminiTextEnv
}
if envKey != "" {
if raw := strings.TrimSpace(os.Getenv(envKey)); raw != "" {
if value, err := strconv.Atoi(raw); err == nil && value > 0 {
return value
}
}
}
return antigravityMaxRetries()
}
func antigravityUseMappedModelForBilling() bool {
v := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityBillingModelEnv)))
return v == "1" || v == "true" || v == "yes" || v == "on"