diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index 15821207..045902cb 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -29,7 +29,10 @@ const ( antigravityRetryMaxDelay = 16 * time.Second ) -const antigravityScopeRateLimitEnv = "GATEWAY_ANTIGRAVITY_429_SCOPE_LIMIT" +const ( + antigravityScopeRateLimitEnv = "GATEWAY_ANTIGRAVITY_429_SCOPE_LIMIT" + antigravityBillingModelEnv = "GATEWAY_ANTIGRAVITY_BILL_WITH_MAPPED_MODEL" +) // antigravityRetryLoopParams 重试循环的参数 type antigravityRetryLoopParams struct { @@ -711,6 +714,10 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, originalModel := claudeReq.Model mappedModel := s.getMappedModel(account, claudeReq.Model) quotaScope, _ := resolveAntigravityQuotaScope(originalModel) + billingModel := originalModel + if antigravityUseMappedModelForBilling() && strings.TrimSpace(mappedModel) != "" { + billingModel = mappedModel + } // 获取 access_token if s.tokenProvider == nil { @@ -971,7 +978,7 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, return &ForwardResult{ RequestID: requestID, Usage: *usage, - Model: originalModel, // 使用原始模型用于计费和日志 + Model: billingModel, Stream: claudeReq.Stream, Duration: time.Since(startTime), FirstTokenMs: firstTokenMs, @@ -1280,6 +1287,10 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co } mappedModel := s.getMappedModel(account, originalModel) + billingModel := originalModel + if antigravityUseMappedModelForBilling() && strings.TrimSpace(mappedModel) != "" { + billingModel = mappedModel + } // 获取 access_token if s.tokenProvider == nil { @@ -1478,7 +1489,7 @@ handleSuccess: return &ForwardResult{ RequestID: requestID, Usage: *usage, - Model: originalModel, + Model: billingModel, Stream: stream, Duration: time.Since(startTime), FirstTokenMs: firstTokenMs, @@ -1525,6 +1536,11 @@ func antigravityUseScopeRateLimit() bool { return v == "1" || v == "true" || v == "yes" || v == "on" } +func antigravityUseMappedModelForBilling() bool { + v := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityBillingModelEnv))) + return v == "1" || v == "true" || v == "yes" || v == "on" +} + func (s *AntigravityGatewayService) handleUpstreamError(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, quotaScope AntigravityQuotaScope) { // 429 使用 Gemini 格式解析(从 body 解析重置时间) if statusCode == 429 {