From caa8c47b68e4150c3aa6d73ac7f0caf373d4629a Mon Sep 17 00:00:00 2001
From: song <slovx2@gmail.com>
Date: Mon, 29 Dec 2025 21:28:28 +0800
Subject: [PATCH] =?UTF-8?q?fix(antigravity):=20=E4=BF=AE=E5=A4=8D=20429=20?=
 =?UTF-8?q?=E9=99=90=E6=B5=81=E5=A4=84=E7=90=86=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 只有 5 次重试全部失败后才标记账户限流
- 使用 Gemini 格式解析 429 响应中的重试时间
- Claude 模型无重试时间时默认 1 分钟，Gemini 默认 5 分钟
- 添加生图模型映射 gemini-3-pro-image-preview
---
 .../service/antigravity_gateway_service.go    | 30 ++++++++++++++++---
 .../service/gemini_messages_compat_service.go |  5 ++--
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go
index 8a5efa73..18a67fdf 100644
--- a/backend/internal/service/antigravity_gateway_service.go
+++ b/backend/internal/service/antigravity_gateway_service.go
@@ -51,23 +51,27 @@ var antigravityModelMapping = map[string]string{
 	"claude-haiku-4-5":           "gemini-3-flash",
 	"claude-3-haiku-20240307":    "gemini-3-flash",
 	"claude-haiku-4-5-20251001":  "gemini-3-flash",
+	// 生图模型：官方名 → Antigravity 内部名
+	"gemini-3-pro-image-preview": "gemini-3-pro-image",
 }
 
 // AntigravityGatewayService 处理 Antigravity 平台的 API 转发
 type AntigravityGatewayService struct {
+	accountRepo      AccountRepository
 	tokenProvider    *AntigravityTokenProvider
 	rateLimitService *RateLimitService
 	httpUpstream     HTTPUpstream
 }
 
 func NewAntigravityGatewayService(
-	_ AccountRepository,
+	accountRepo AccountRepository,
 	_ GatewayCache,
 	tokenProvider *AntigravityTokenProvider,
 	rateLimitService *RateLimitService,
 	httpUpstream HTTPUpstream,
 ) *AntigravityGatewayService {
 	return &AntigravityGatewayService{
+		accountRepo:      accountRepo,
 		tokenProvider:    tokenProvider,
 		rateLimitService: rateLimitService,
 		httpUpstream:     httpUpstream,
@@ -402,14 +406,15 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
 			respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
 			_ = resp.Body.Close()
 
-			if resp.StatusCode == 429 {
-				s.handleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
-			}
 			if attempt < antigravityMaxRetries {
 				log.Printf("Antigravity account %d: upstream status %d, retry %d/%d", account.ID, resp.StatusCode, attempt, antigravityMaxRetries)
 				sleepAntigravityBackoff(attempt)
 				continue
 			}
+			// 所有重试都失败，标记限流状态
+			if resp.StatusCode == 429 {
+				s.handleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
+			}
 			if action == "countTokens" {
 				estimated := estimateGeminiCountTokens(body)
 				c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
@@ -526,6 +531,23 @@ func sleepAntigravityBackoff(attempt int) {
 }
 
 func (s *AntigravityGatewayService) handleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, body []byte) {
+	// 429 使用 Gemini 格式解析（从 body 解析重置时间）
+	if statusCode == 429 {
+		resetAt := ParseGeminiRateLimitResetTime(body)
+		if resetAt == nil {
+			// 解析失败：Gemini 有重试时间用 5 分钟，Claude 没有用 1 分钟
+			defaultDur := 1 * time.Minute
+			if bytes.Contains(body, []byte("Please retry in")) || bytes.Contains(body, []byte("retryDelay")) {
+				defaultDur = 5 * time.Minute
+			}
+			ra := time.Now().Add(defaultDur)
+			_ = s.accountRepo.SetRateLimited(ctx, account.ID, ra)
+			return
+		}
+		_ = s.accountRepo.SetRateLimited(ctx, account.ID, time.Unix(*resetAt, 0))
+		return
+	}
+	// 其他错误码继续使用 rateLimitService
 	if s.rateLimitService == nil {
 		return
 	}
diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go
index c7374ad6..6ccbf8ec 100644
--- a/backend/internal/service/gemini_messages_compat_service.go
+++ b/backend/internal/service/gemini_messages_compat_service.go
@@ -1883,7 +1883,7 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont
 	if statusCode != 429 {
 		return
 	}
-	resetAt := parseGeminiRateLimitResetTime(body)
+	resetAt := ParseGeminiRateLimitResetTime(body)
 	if resetAt == nil {
 		ra := time.Now().Add(5 * time.Minute)
 		_ = s.accountRepo.SetRateLimited(ctx, account.ID, ra)
@@ -1892,7 +1892,8 @@ func (s *GeminiMessagesCompatService) handleGeminiUpstreamError(ctx context.Cont
 	_ = s.accountRepo.SetRateLimited(ctx, account.ID, time.Unix(*resetAt, 0))
 }
 
-func parseGeminiRateLimitResetTime(body []byte) *int64 {
+// ParseGeminiRateLimitResetTime 解析 Gemini 格式的 429 响应，返回重置时间的 Unix 时间戳
+func ParseGeminiRateLimitResetTime(body []byte) *int64 {
 	// Try to parse metadata.quotaResetDelay like "12.345s"
 	var parsed map[string]any
 	if err := json.Unmarshal(body, &parsed); err == nil {