From b2b842bf7a579aa84f3a39a81fda3047987801e5 Mon Sep 17 00:00:00 2001
From: song <slovx2@gmail.com>
Date: Fri, 2 Jan 2026 12:38:03 +0800
Subject: [PATCH] =?UTF-8?q?refactor(antigravity):=20countTokens=20?=
 =?UTF-8?q?=E7=AB=AF=E7=82=B9=E7=9B=B4=E6=8E=A5=E8=BF=94=E5=9B=9E=E7=A9=BA?=
 =?UTF-8?q?=E5=80=BC?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Gemini 端点 countTokens 直接返回 {"totalTokens": 0}
- Claude 端点 countTokens 返回 {"input_tokens": 0}
- 移除透传上游和本地估算逻辑
---
 .../service/antigravity_gateway_service.go    | 50 +++++--------------
 backend/internal/service/gateway_service.go   |  5 +-
 2 files changed, 14 insertions(+), 41 deletions(-)

diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go
index 267d7548..0798f46d 100644
--- a/backend/internal/service/antigravity_gateway_service.go
+++ b/backend/internal/service/antigravity_gateway_service.go
@@ -467,8 +467,19 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
 	}
 
 	switch action {
-	case "generateContent", "streamGenerateContent", "countTokens":
+	case "generateContent", "streamGenerateContent":
 		// ok
+	case "countTokens":
+		// 直接返回空值，不透传上游
+		c.JSON(http.StatusOK, map[string]any{"totalTokens": 0})
+		return &ForwardResult{
+			RequestID:    "",
+			Usage:        ClaudeUsage{},
+			Model:        originalModel,
+			Stream:       false,
+			Duration:     time.Since(time.Now()),
+			FirstTokenMs: nil,
+		}, nil
 	default:
 		return nil, s.writeGoogleError(c, http.StatusNotFound, "Unsupported action: "+action)
 	}
@@ -523,18 +534,6 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
 				sleepAntigravityBackoff(attempt)
 				continue
 			}
-			if action == "countTokens" {
-				estimated := estimateGeminiCountTokens(body)
-				c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
-				return &ForwardResult{
-					RequestID:    "",
-					Usage:        ClaudeUsage{},
-					Model:        originalModel,
-					Stream:       false,
-					Duration:     time.Since(startTime),
-					FirstTokenMs: nil,
-				}, nil
-			}
 			return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries")
 		}
 
@@ -551,18 +550,6 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
 			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})
-				return &ForwardResult{
-					RequestID:    "",
-					Usage:        ClaudeUsage{},
-					Model:        originalModel,
-					Stream:       false,
-					Duration:     time.Since(startTime),
-					FirstTokenMs: nil,
-				}, nil
-			}
 			resp = &http.Response{
 				StatusCode: resp.StatusCode,
 				Header:     resp.Header.Clone(),
@@ -585,19 +572,6 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
 		respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
 		s.handleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody)
 
-		if action == "countTokens" {
-			estimated := estimateGeminiCountTokens(body)
-			c.JSON(http.StatusOK, map[string]any{"totalTokens": estimated})
-			return &ForwardResult{
-				RequestID:    requestID,
-				Usage:        ClaudeUsage{},
-				Model:        originalModel,
-				Stream:       false,
-				Duration:     time.Since(startTime),
-				FirstTokenMs: nil,
-			}, nil
-		}
-
 		if s.shouldFailoverUpstreamError(resp.StatusCode) {
 			return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode}
 		}
diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go
index ba1d5bb3..3932c35c 100644
--- a/backend/internal/service/gateway_service.go
+++ b/backend/internal/service/gateway_service.go
@@ -1750,10 +1750,9 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
 	body := parsed.Body
 	reqModel := parsed.Model
 
-	// Antigravity 账户不支持 count_tokens 转发，返回估算值
-	// 参考 Antigravity-Manager 和 proxycast 实现
+	// Antigravity 账户不支持 count_tokens 转发，直接返回空值
 	if account.Platform == PlatformAntigravity {
-		c.JSON(http.StatusOK, gin.H{"input_tokens": 100})
+		c.JSON(http.StatusOK, gin.H{"input_tokens": 0})
 		return nil
 	}