antigravity: 区分切换后重试次数

backend/internal/service/antigravity_gateway_service.go

@@ -33,6 +33,7 @@ const (
 
 const (
 	antigravityMaxRetriesEnv            = "GATEWAY_ANTIGRAVITY_MAX_RETRIES"
+	antigravityMaxRetriesAfterSwitchEnv = "GATEWAY_ANTIGRAVITY_AFTER_SWITCHMAX_RETRIES"
 	antigravityMaxRetriesClaudeEnv      = "GATEWAY_ANTIGRAVITY_MAX_RETRIES_CLAUDE"
 	antigravityMaxRetriesGeminiTextEnv  = "GATEWAY_ANTIGRAVITY_MAX_RETRIES_GEMINI_TEXT"
 	antigravityMaxRetriesGeminiImageEnv = "GATEWAY_ANTIGRAVITY_MAX_RETRIES_GEMINI_IMAGE"
@@ -745,6 +746,8 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
 	if antigravityUseMappedModelForBilling() && strings.TrimSpace(mappedModel) != "" {
 		billingModel = mappedModel
 	}
+	afterSwitch := antigravityHasAccountSwitch(ctx)
+	maxRetries := antigravityMaxRetriesForModel(originalModel, afterSwitch)
 
 	// 获取 access_token
 	if s.tokenProvider == nil {
@@ -793,7 +796,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
 		httpUpstream:   s.httpUpstream,
 		settingService: s.settingService,
 		handleError:    s.handleUpstreamError,
-		maxRetries:     antigravityMaxRetriesForModel(originalModel),
+		maxRetries:     maxRetries,
 	})
 	if err != nil {
 		return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries")
@@ -870,7 +873,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
 					httpUpstream:   s.httpUpstream,
 					settingService: s.settingService,
 					handleError:    s.handleUpstreamError,
-					maxRetries:     antigravityMaxRetriesForModel(originalModel),
+					maxRetries:     maxRetries,
 				})
 				if retryErr != nil {
 					appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
@@ -1387,6 +1390,8 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
 	if antigravityUseMappedModelForBilling() && strings.TrimSpace(mappedModel) != "" {
 		billingModel = mappedModel
 	}
+	afterSwitch := antigravityHasAccountSwitch(ctx)
+	maxRetries := antigravityMaxRetriesForModel(originalModel, afterSwitch)
 
 	// 获取 access_token
 	if s.tokenProvider == nil {
@@ -1444,7 +1449,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
 		httpUpstream:   s.httpUpstream,
 		settingService: s.settingService,
 		handleError:    s.handleUpstreamError,
-		maxRetries:     antigravityMaxRetriesForModel(originalModel),
+		maxRetries:     maxRetries,
 	})
 	if err != nil {
 		return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries")
@@ -1641,6 +1646,16 @@ func antigravityUseScopeRateLimit() bool {
 	return v == "1" || v == "true" || v == "yes" || v == "on"
 }
 
+func antigravityHasAccountSwitch(ctx context.Context) bool {
+	if ctx == nil {
+		return false
+	}
+	if v, ok := ctx.Value(ctxkey.AccountSwitchCount).(int); ok {
+		return v > 0
+	}
+	return false
+}
+
 func antigravityMaxRetries() int {
 	raw := strings.TrimSpace(os.Getenv(antigravityMaxRetriesEnv))
 	if raw == "" {
@@ -1653,9 +1668,21 @@ func antigravityMaxRetries() int {
 	return value
 }
 
+func antigravityMaxRetriesAfterSwitch() int {
+	raw := strings.TrimSpace(os.Getenv(antigravityMaxRetriesAfterSwitchEnv))
+	if raw == "" {
+		return antigravityMaxRetries()
+	}
+	value, err := strconv.Atoi(raw)
+	if err != nil || value <= 0 {
+		return antigravityMaxRetries()
+	}
+	return value
+}
+
 // antigravityMaxRetriesForModel 根据模型类型获取重试次数
 // 优先使用模型细分配置，未设置则回退到平台级配置
-func antigravityMaxRetriesForModel(model string) int {
+func antigravityMaxRetriesForModel(model string, afterSwitch bool) int {
 	var envKey string
 	if strings.HasPrefix(model, "claude-") {
 		envKey = antigravityMaxRetriesClaudeEnv
@@ -1672,6 +1699,9 @@ func antigravityMaxRetriesForModel(model string) int {
 			}
 		}
 	}
+	if afterSwitch {
+		return antigravityMaxRetriesAfterSwitch()
+	}
 	return antigravityMaxRetries()
 }
 

backend/internal/service/antigravity_gateway_service_test.go

@@ -161,3 +161,28 @@ func TestAntigravityGatewayService_Forward_PromptTooLong(t *testing.T) {
 	require.Len(t, events, 1)
 	require.Equal(t, "prompt_too_long", events[0].Kind)
 }
+
+func TestAntigravityMaxRetriesForModel_AfterSwitch(t *testing.T) {
+	t.Setenv(antigravityMaxRetriesEnv, "4")
+	t.Setenv(antigravityMaxRetriesAfterSwitchEnv, "7")
+	t.Setenv(antigravityMaxRetriesClaudeEnv, "")
+	t.Setenv(antigravityMaxRetriesGeminiTextEnv, "")
+	t.Setenv(antigravityMaxRetriesGeminiImageEnv, "")
+
+	got := antigravityMaxRetriesForModel("claude-sonnet-4-5", false)
+	require.Equal(t, 4, got)
+
+	got = antigravityMaxRetriesForModel("claude-sonnet-4-5", true)
+	require.Equal(t, 7, got)
+}
+
+func TestAntigravityMaxRetriesForModel_AfterSwitchFallback(t *testing.T) {
+	t.Setenv(antigravityMaxRetriesEnv, "5")
+	t.Setenv(antigravityMaxRetriesAfterSwitchEnv, "")
+	t.Setenv(antigravityMaxRetriesClaudeEnv, "")
+	t.Setenv(antigravityMaxRetriesGeminiTextEnv, "")
+	t.Setenv(antigravityMaxRetriesGeminiImageEnv, "")
+
+	got := antigravityMaxRetriesForModel("gemini-2.5-flash", true)
+	require.Equal(t, 5, got)
+}

backend/internal/service/ops_retry.go

@@ -12,6 +12,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
 	infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
 	"github.com/gin-gonic/gin"
 	"github.com/lib/pq"
@@ -476,9 +477,13 @@ func (s *OpsService) executeClientRetry(ctx context.Context, reqType opsRetryReq
 			continue
 		}
 
+		attemptCtx := ctx
+		if switches > 0 {
+			attemptCtx = context.WithValue(attemptCtx, ctxkey.AccountSwitchCount, switches)
+		}
 		exec := func() *opsRetryExecution {
 			defer selection.ReleaseFunc()
-			return s.executeWithAccount(ctx, reqType, errorLog, body, account)
+			return s.executeWithAccount(attemptCtx, reqType, errorLog, body, account)
 		}()
 
 		if exec != nil {