refactor(antigravity): countTokens 端点直接返回空值

- Gemini 端点 countTokens 直接返回 {"totalTokens": 0}
- Claude 端点 countTokens 返回 {"input_tokens": 0}
- 移除透传上游和本地估算逻辑
This commit is contained in:
song
2026-01-02 12:38:03 +08:00
parent 9d3ec9e627
commit b2b842bf7a
2 changed files with 14 additions and 41 deletions

View File

@@ -467,8 +467,19 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
} }
switch action { switch action {
case "generateContent", "streamGenerateContent", "countTokens": case "generateContent", "streamGenerateContent":
// ok // 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: default:
return nil, s.writeGoogleError(c, http.StatusNotFound, "Unsupported action: "+action) 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) sleepAntigravityBackoff(attempt)
continue 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") 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 { if resp.StatusCode == 429 {
s.handleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) 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{ resp = &http.Response{
StatusCode: resp.StatusCode, StatusCode: resp.StatusCode,
Header: resp.Header.Clone(), 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)) respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
s.handleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) 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) { if s.shouldFailoverUpstreamError(resp.StatusCode) {
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode} return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode}
} }

View File

@@ -1750,10 +1750,9 @@ func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context,
body := parsed.Body body := parsed.Body
reqModel := parsed.Model reqModel := parsed.Model
// Antigravity 账户不支持 count_tokens 转发,返回估算 // Antigravity 账户不支持 count_tokens 转发,直接返回空
// 参考 Antigravity-Manager 和 proxycast 实现
if account.Platform == PlatformAntigravity { if account.Platform == PlatformAntigravity {
c.JSON(http.StatusOK, gin.H{"input_tokens": 100}) c.JSON(http.StatusOK, gin.H{"input_tokens": 0})
return nil return nil
} }