diff --git a/backend/internal/service/antigravity_gateway_claude.go b/backend/internal/service/antigravity_gateway_claude.go new file mode 100644 index 0000000000..2011625984 --- /dev/null +++ b/backend/internal/service/antigravity_gateway_claude.go @@ -0,0 +1,754 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" +) + +// Forward 转发 Claude 协议请求（Claude → Gemini 转换） +// +// 限流处理流程: +// +// 请求 → antigravityRetryLoop → 预检查(remaining>0? → 切换账号) → 发送上游 +// ├─ 成功 → 正常返回 +// └─ 429/503 → handleSmartRetry +// ├─ retryDelay >= 7s → 设置模型限流 + 清除粘性绑定 → 切换账号 +// └─ retryDelay < 7s → 等待后重试 1 次 +// ├─ 成功 → 正常返回 +// └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号 +func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte, isStickySession bool) (*ForwardResult, error) { + // 上游透传账号直接转发，不走 OAuth token 刷新 + if account.Type == AccountTypeUpstream { + return s.ForwardUpstream(ctx, c, account, body) + } + + startTime := time.Now() + + sessionID := getSessionID(c) + prefix := logPrefix(sessionID, account.Name) + + // 解析 Claude 请求 + var claudeReq antigravity.ClaudeRequest + if err := json.Unmarshal(body, &claudeReq); err != nil { + return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request body") + } + if strings.TrimSpace(claudeReq.Model) == "" { + return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Missing model") + } + + originalModel := claudeReq.Model + mappedModel := s.getMappedModel(account, claudeReq.Model) + if mappedModel == "" { + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) + return nil, s.writeClaudeError(c, http.StatusForbidden, "permission_error", fmt.Sprintf("model %s not in whitelist", claudeReq.Model)) + } + // 应用 thinking 模式自动后缀：如果 thinking 开启且目标是 claude-sonnet-4-5，自动改为 thinking 版本 + thinkingEnabled := claudeReq.Thinking != nil && (claudeReq.Thinking.Type == "enabled" || claudeReq.Thinking.Type == "adaptive") + mappedModel = applyThinkingModelSuffix(mappedModel, thinkingEnabled) + billingModel := mappedModel + + // 获取 access_token + if s.tokenProvider == nil { + return nil, s.writeClaudeError(c, http.StatusBadGateway, "api_error", "Antigravity token provider not configured") + } + accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: []byte(`{"error":{"type":"authentication_error","message":"Failed to get upstream access token"},"type":"error"}`), + } + } + + projectID, err := resolveAntigravityProjectID(account) + if err != nil { + _ = s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", err.Error()) + return nil, err + } + + // 代理 URL + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + + // 获取转换选项 + // Antigravity 上游要求必须包含身份提示词，否则会返回 429 + transformOpts := s.getClaudeTransformOptions(ctx) + transformOpts.EnableIdentityPatch = true // 强制启用，Antigravity 上游必需 + + // 转换 Claude 请求为 Gemini 格式 + geminiBody, err := antigravity.TransformClaudeToGeminiWithOptions(&claudeReq, projectID, mappedModel, transformOpts) + if err != nil { + return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request") + } + + // Antigravity 上游只支持流式请求，统一使用 streamGenerateContent + // 如果客户端请求非流式，在响应处理阶段会收集完整流式响应后转换返回 + action := "streamGenerateContent" + + // 执行带重试的请求 + result, err := s.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: prefix, + account: account, + proxyURL: proxyURL, + accessToken: accessToken, + action: action, + body: geminiBody, + c: c, + httpUpstream: s.httpUpstream, + settingService: s.settingService, + accountRepo: s.accountRepo, + handleError: s.handleUpstreamError, + requestedModel: originalModel, + isStickySession: isStickySession, // Forward 由上层判断粘性会话 + groupID: 0, // Forward 方法没有 groupID，由上层处理粘性会话清除 + sessionHash: "", // Forward 方法没有 sessionHash，由上层处理粘性会话清除 + }) + if err != nil { + // 检查是否是账号切换信号，转换为 UpstreamFailoverError 让 Handler 切换账号 + if switchErr, ok := IsAntigravityAccountSwitchError(err); ok { + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusServiceUnavailable, + ForceCacheBilling: switchErr.IsStickySession, + } + } + // 区分客户端取消和真正的上游失败，返回更准确的错误消息 + if c.Request.Context().Err() != nil { + return nil, s.writeClaudeError(c, http.StatusBadGateway, "client_disconnected", "Client disconnected before upstream response") + } + return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries") + } + resp := result.resp + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode >= 400 { + respBody := s.readUpstreamErrorBody(resp) + + // 优先检测 thinking block 的 signature 相关错误（400）并重试一次： + // Antigravity /v1internal 链路在部分场景会对 thought/thinking signature 做严格校验， + // 当历史消息携带的 signature 不合法时会直接 400；去除 thinking 后可继续完成请求。 + if resp.StatusCode == http.StatusBadRequest && isSignatureRelatedError(respBody) && s.settingService.IsSignatureRectifierEnabled(ctx) { + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + logBody, maxBytes := s.getLogConfig() + upstreamDetail := s.getUpstreamErrorDetail(respBody) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "signature_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + // Conservative two-stage fallback: + // 1) Disable top-level thinking + thinking->text + // 2) Only if still signature-related 400: also downgrade tool_use/tool_result to text. + + retryStages := []struct { + name string + strip func(*antigravity.ClaudeRequest) (bool, error) + }{ + {name: "thinking-only", strip: stripThinkingFromClaudeRequest}, + {name: "thinking+tools", strip: stripSignatureSensitiveBlocksFromClaudeRequest}, + } + + for _, stage := range retryStages { + retryClaudeReq := claudeReq + retryClaudeReq.Messages = append([]antigravity.ClaudeMessage(nil), claudeReq.Messages...) + + stripped, stripErr := stage.strip(&retryClaudeReq) + if stripErr != nil || !stripped { + continue + } + + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: detected signature-related 400, retrying once (%s)", account.ID, stage.name) + + retryGeminiBody, txErr := antigravity.TransformClaudeToGeminiWithOptions(&retryClaudeReq, projectID, mappedModel, s.getClaudeTransformOptions(ctx)) + if txErr != nil { + continue + } + retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: prefix, + account: account, + proxyURL: proxyURL, + accessToken: accessToken, + action: action, + body: retryGeminiBody, + c: c, + httpUpstream: s.httpUpstream, + settingService: s.settingService, + accountRepo: s.accountRepo, + handleError: s.handleUpstreamError, + requestedModel: originalModel, + isStickySession: isStickySession, + groupID: 0, // Forward 方法没有 groupID，由上层处理粘性会话清除 + sessionHash: "", // Forward 方法没有 sessionHash，由上层处理粘性会话清除 + }) + if retryErr != nil { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: 0, + Kind: "signature_retry_request_error", + Message: sanitizeUpstreamErrorMessage(retryErr.Error()), + }) + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: signature retry request failed (%s): %v", account.ID, stage.name, retryErr) + continue + } + + retryResp := retryResult.resp + if retryResp.StatusCode < 400 { + _ = resp.Body.Close() + resp = retryResp + respBody = nil + break + } + + retryBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) + _ = retryResp.Body.Close() + if retryResp.StatusCode == http.StatusTooManyRequests { + retryBaseURL := "" + if retryResp.Request != nil && retryResp.Request.URL != nil { + retryBaseURL = retryResp.Request.URL.Scheme + "://" + retryResp.Request.URL.Host + } + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limited base_url=%s retry_stage=%s body=%s", prefix, retryBaseURL, stage.name, truncateForLog(retryBody, 200)) + } + kind := "signature_retry" + if strings.TrimSpace(stage.name) != "" { + kind = "signature_retry_" + strings.ReplaceAll(stage.name, "+", "_") + } + retryUpstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(retryBody)) + retryUpstreamMsg = sanitizeUpstreamErrorMessage(retryUpstreamMsg) + retryUpstreamDetail := "" + if logBody { + retryUpstreamDetail = truncateString(string(retryBody), maxBytes) + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: retryResp.StatusCode, + UpstreamRequestID: retryResp.Header.Get("x-request-id"), + Kind: kind, + Message: retryUpstreamMsg, + Detail: retryUpstreamDetail, + }) + + // If this stage fixed the signature issue, we stop; otherwise we may try the next stage. + if retryResp.StatusCode != http.StatusBadRequest || !isSignatureRelatedError(retryBody) { + respBody = retryBody + resp = &http.Response{ + StatusCode: retryResp.StatusCode, + Header: retryResp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + } + break + } + + // Still signature-related; capture context and allow next stage. + respBody = retryBody + resp = &http.Response{ + StatusCode: retryResp.StatusCode, + Header: retryResp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + } + } + } + + // Budget 整流：检测 budget_tokens 约束错误并自动修正重试 + if resp.StatusCode == http.StatusBadRequest && respBody != nil && !isSignatureRelatedError(respBody) { + errMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + if isThinkingBudgetConstraintError(errMsg) && s.settingService.IsBudgetRectifierEnabled(ctx) { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "budget_constraint_error", + Message: errMsg, + Detail: s.getUpstreamErrorDetail(respBody), + }) + + // 修正 claudeReq 的 thinking 参数（adaptive 模式不修正） + if claudeReq.Thinking == nil || claudeReq.Thinking.Type != "adaptive" { + retryClaudeReq := claudeReq + retryClaudeReq.Messages = append([]antigravity.ClaudeMessage(nil), claudeReq.Messages...) + // 创建新的 ThinkingConfig 避免修改原始 claudeReq.Thinking 指针 + retryClaudeReq.Thinking = &antigravity.ThinkingConfig{ + Type: "enabled", + BudgetTokens: BudgetRectifyBudgetTokens, + } + if retryClaudeReq.MaxTokens < BudgetRectifyMinMaxTokens { + retryClaudeReq.MaxTokens = BudgetRectifyMaxTokens + } + + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: detected budget_tokens constraint error, retrying with rectified budget (budget_tokens=%d, max_tokens=%d)", account.ID, BudgetRectifyBudgetTokens, BudgetRectifyMaxTokens) + + retryGeminiBody, txErr := antigravity.TransformClaudeToGeminiWithOptions(&retryClaudeReq, projectID, mappedModel, transformOpts) + if txErr == nil { + retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: prefix, + account: account, + proxyURL: proxyURL, + accessToken: accessToken, + action: action, + body: retryGeminiBody, + c: c, + httpUpstream: s.httpUpstream, + settingService: s.settingService, + accountRepo: s.accountRepo, + handleError: s.handleUpstreamError, + requestedModel: originalModel, + isStickySession: isStickySession, + groupID: 0, + sessionHash: "", + }) + if retryErr == nil { + retryResp := retryResult.resp + if retryResp.StatusCode < 400 { + _ = resp.Body.Close() + resp = retryResp + respBody = nil + } else { + retryBody := s.readUpstreamErrorBody(retryResp) + _ = retryResp.Body.Close() + respBody = retryBody + resp = &http.Response{ + StatusCode: retryResp.StatusCode, + Header: retryResp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + } + } + } else { + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: budget rectifier retry failed: %v", account.ID, retryErr) + } + } + } + } + } + + // 处理错误响应（重试后仍失败或不触发重试） + if resp.StatusCode >= 400 { + // 检测 prompt too long 错误，返回特殊错误类型供上层 fallback + if resp.StatusCode == http.StatusBadRequest && isPromptTooLongError(respBody) { + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + upstreamDetail := s.getUpstreamErrorDetail(respBody) + logBody, maxBytes := s.getLogConfig() + if logBody { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=400 prompt_too_long=true upstream_message=%q request_id=%s body=%s", prefix, upstreamMsg, resp.Header.Get("x-request-id"), truncateForLog(respBody, maxBytes)) + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "prompt_too_long", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return nil, &PromptTooLongError{ + StatusCode: resp.StatusCode, + RequestID: resp.Header.Get("x-request-id"), + Body: respBody, + } + } + + s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", isStickySession) + + // 精确匹配服务端配置类 400 错误，触发同账号重试 + failover + if resp.StatusCode == http.StatusBadRequest { + msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) + if isGoogleProjectConfigError(msg) { + upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) + upstreamDetail := s.getUpstreamErrorDetail(respBody) + log.Printf("%s status=400 google_config_error failover=true upstream_message=%q account=%d", prefix, upstreamMsg, account.ID) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "failover", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, RetryableOnSameAccount: true} + } + } + + if s.shouldFailoverUpstreamError(resp.StatusCode) { + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + upstreamDetail := s.getUpstreamErrorDetail(respBody) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "failover", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody} + } + + return nil, s.writeMappedClaudeError(c, account, resp.StatusCode, resp.Header.Get("x-request-id"), respBody) + } + } + + requestID := resp.Header.Get("x-request-id") + if requestID != "" { + c.Header("x-request-id", requestID) + } + + var usage *ClaudeUsage + var firstTokenMs *int + var clientDisconnect bool + if claudeReq.Stream { + // 客户端要求流式，直接透传转换 + streamRes, err := s.handleClaudeStreamingResponse(c, resp, startTime, originalModel) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_error error=%v", prefix, err) + return nil, err + } + usage = streamRes.usage + firstTokenMs = streamRes.firstTokenMs + clientDisconnect = streamRes.clientDisconnect + } else { + // 客户端要求非流式，收集流式响应后转换返回 + streamRes, err := s.handleClaudeStreamToNonStreaming(c, resp, startTime, originalModel) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_collect_error error=%v", prefix, err) + return nil, err + } + usage = streamRes.usage + firstTokenMs = streamRes.firstTokenMs + } + + return &ForwardResult{ + RequestID: requestID, + Usage: *usage, + Model: originalModel, + UpstreamModel: billingModel, + Stream: claudeReq.Stream, + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnect, + }, nil +} + +func isSignatureRelatedError(respBody []byte) bool { + msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) + if msg == "" { + // Fallback: best-effort scan of the raw payload. + msg = strings.ToLower(string(respBody)) + } + + // Keep this intentionally broad: different upstreams may use "signature" or "thought_signature". + if strings.Contains(msg, "thought_signature") || strings.Contains(msg, "signature") { + return true + } + + // Also detect thinking block structural errors: + // "Expected `thinking` or `redacted_thinking`, but found `text`" + if strings.Contains(msg, "expected") && (strings.Contains(msg, "thinking") || strings.Contains(msg, "redacted_thinking")) { + return true + } + + return false +} + +// isPromptTooLongError 检测是否为 prompt too long 错误 +func isPromptTooLongError(respBody []byte) bool { + msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) + if msg == "" { + msg = strings.ToLower(string(respBody)) + } + return strings.Contains(msg, "prompt is too long") || + strings.Contains(msg, "request is too long") || + strings.Contains(msg, "context length exceeded") || + strings.Contains(msg, "max_tokens") +} + +// isPassthroughErrorMessage 检查错误消息是否在透传白名单中 +func isPassthroughErrorMessage(msg string) bool { + lower := strings.ToLower(msg) + for _, pattern := range antigravityPassthroughErrorMessages { + if strings.Contains(lower, pattern) { + return true + } + } + return false +} + +// getPassthroughOrDefault 若消息在白名单内则返回原始消息，否则返回默认消息 +func getPassthroughOrDefault(upstreamMsg, defaultMsg string) string { + if isPassthroughErrorMessage(upstreamMsg) { + return upstreamMsg + } + return defaultMsg +} + +func extractAntigravityErrorMessage(body []byte) string { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + return "" + } + + // Google-style: {"error": {"message": "..."}} + if errObj, ok := payload["error"].(map[string]any); ok { + if msg, ok := errObj["message"].(string); ok && strings.TrimSpace(msg) != "" { + return msg + } + } + + // Fallback: top-level message + if msg, ok := payload["message"].(string); ok && strings.TrimSpace(msg) != "" { + return msg + } + + return "" +} + +// stripThinkingFromClaudeRequest converts thinking blocks to text blocks in a Claude Messages request. +// This preserves the thinking content while avoiding signature validation errors. +// Note: redacted_thinking blocks are removed because they cannot be converted to text. +// It also disables top-level `thinking` to avoid upstream structural constraints for thinking mode. +func stripThinkingFromClaudeRequest(req *antigravity.ClaudeRequest) (bool, error) { + if req == nil { + return false, nil + } + + changed := false + if req.Thinking != nil { + req.Thinking = nil + changed = true + } + + for i := range req.Messages { + raw := req.Messages[i].Content + if len(raw) == 0 { + continue + } + + // If content is a string, nothing to strip. + var str string + if json.Unmarshal(raw, &str) == nil { + continue + } + + // Otherwise treat as an array of blocks and convert thinking blocks to text. + var blocks []map[string]any + if err := json.Unmarshal(raw, &blocks); err != nil { + continue + } + + filtered := make([]map[string]any, 0, len(blocks)) + modifiedAny := false + for _, block := range blocks { + t, _ := block["type"].(string) + switch t { + case "thinking": + thinkingText, _ := block["thinking"].(string) + if thinkingText != "" { + filtered = append(filtered, map[string]any{ + "type": "text", + "text": thinkingText, + }) + } + modifiedAny = true + case "redacted_thinking": + modifiedAny = true + case "": + if thinkingText, hasThinking := block["thinking"].(string); hasThinking { + if thinkingText != "" { + filtered = append(filtered, map[string]any{ + "type": "text", + "text": thinkingText, + }) + } + modifiedAny = true + } else { + filtered = append(filtered, block) + } + default: + filtered = append(filtered, block) + } + } + + if !modifiedAny { + continue + } + + if len(filtered) == 0 { + filtered = append(filtered, map[string]any{ + "type": "text", + "text": "(content removed)", + }) + } + + newRaw, err := json.Marshal(filtered) + if err != nil { + return changed, err + } + req.Messages[i].Content = newRaw + changed = true + } + + return changed, nil +} + +// stripSignatureSensitiveBlocksFromClaudeRequest is a stronger retry degradation that additionally converts +// tool blocks to plain text. Use this only after a thinking-only retry still fails with signature errors. +func stripSignatureSensitiveBlocksFromClaudeRequest(req *antigravity.ClaudeRequest) (bool, error) { + if req == nil { + return false, nil + } + + changed := false + if req.Thinking != nil { + req.Thinking = nil + changed = true + } + + for i := range req.Messages { + raw := req.Messages[i].Content + if len(raw) == 0 { + continue + } + + // If content is a string, nothing to strip. + var str string + if json.Unmarshal(raw, &str) == nil { + continue + } + + // Otherwise treat as an array of blocks and convert signature-sensitive blocks to text. + var blocks []map[string]any + if err := json.Unmarshal(raw, &blocks); err != nil { + continue + } + + filtered := make([]map[string]any, 0, len(blocks)) + modifiedAny := false + for _, block := range blocks { + t, _ := block["type"].(string) + switch t { + case "thinking": + // Convert thinking to text, skip if empty + thinkingText, _ := block["thinking"].(string) + if thinkingText != "" { + filtered = append(filtered, map[string]any{ + "type": "text", + "text": thinkingText, + }) + } + modifiedAny = true + case "redacted_thinking": + // Remove redacted_thinking (cannot convert encrypted content) + modifiedAny = true + case "tool_use": + // Convert tool_use to text to avoid upstream signature/thought_signature validation errors. + // This is a retry-only degradation path, so we prioritise request validity over tool semantics. + name, _ := block["name"].(string) + id, _ := block["id"].(string) + input := block["input"] + inputJSON, _ := json.Marshal(input) + text := "(tool_use)" + if name != "" { + text += " name=" + name + } + if id != "" { + text += " id=" + id + } + if len(inputJSON) > 0 && string(inputJSON) != "null" { + text += " input=" + string(inputJSON) + } + filtered = append(filtered, map[string]any{ + "type": "text", + "text": text, + }) + modifiedAny = true + case "tool_result": + // Convert tool_result to text so it stays consistent when tool_use is downgraded. + toolUseID, _ := block["tool_use_id"].(string) + isError, _ := block["is_error"].(bool) + content := block["content"] + contentJSON, _ := json.Marshal(content) + text := "(tool_result)" + if toolUseID != "" { + text += " tool_use_id=" + toolUseID + } + if isError { + text += " is_error=true" + } + if len(contentJSON) > 0 && string(contentJSON) != "null" { + text += "\n" + string(contentJSON) + } + filtered = append(filtered, map[string]any{ + "type": "text", + "text": text, + }) + modifiedAny = true + case "": + // Handle untyped block with "thinking" field + if thinkingText, hasThinking := block["thinking"].(string); hasThinking { + if thinkingText != "" { + filtered = append(filtered, map[string]any{ + "type": "text", + "text": thinkingText, + }) + } + modifiedAny = true + } else { + filtered = append(filtered, block) + } + default: + filtered = append(filtered, block) + } + } + + if !modifiedAny { + continue + } + + if len(filtered) == 0 { + // Keep request valid: upstream rejects empty content arrays. + filtered = append(filtered, map[string]any{ + "type": "text", + "text": "(content removed)", + }) + } + + newRaw, err := json.Marshal(filtered) + if err != nil { + return changed, err + } + req.Messages[i].Content = newRaw + changed = true + } + + return changed, nil +} diff --git a/backend/internal/service/antigravity_gateway_gemini.go b/backend/internal/service/antigravity_gateway_gemini.go new file mode 100644 index 0000000000..f0d20244a8 --- /dev/null +++ b/backend/internal/service/antigravity_gateway_gemini.go @@ -0,0 +1,550 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" +) + +// ForwardGemini 转发 Gemini 协议请求 +// +// 限流处理流程: +// +// 请求 → antigravityRetryLoop → 预检查(remaining>0? → 切换账号) → 发送上游 +// ├─ 成功 → 正常返回 +// └─ 429/503 → handleSmartRetry +// ├─ retryDelay >= 7s → 设置模型限流 + 清除粘性绑定 → 切换账号 +// └─ retryDelay < 7s → 等待后重试 1 次 +// ├─ 成功 → 正常返回 +// └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号 +type ForwardGeminiOption func(*forwardGeminiOptions) + +type forwardGeminiOptions struct { + groupID int64 + sessionHash string +} + +func WithForwardGeminiSession(groupID int64, sessionHash string) ForwardGeminiOption { + return func(opts *forwardGeminiOptions) { + opts.groupID = groupID + opts.sessionHash = sessionHash + } +} + +func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool, options ...ForwardGeminiOption) (*ForwardResult, error) { + startTime := time.Now() + forwardOpts := forwardGeminiOptions{} + for _, apply := range options { + if apply != nil { + apply(&forwardOpts) + } + } + + sessionID := getSessionID(c) + prefix := logPrefix(sessionID, account.Name) + + if strings.TrimSpace(originalModel) == "" { + return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing model in URL") + } + if strings.TrimSpace(action) == "" { + return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing action in URL") + } + if len(body) == 0 { + return nil, s.writeGoogleError(c, http.StatusBadRequest, "Request body is empty") + } + + // 解析请求以获取 image_size（用于图片计费） + imageInputSize := s.extractImageInputSize(body) + imageSize := normalizeOpenAIImageSizeTier(imageInputSize) + + switch action { + 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(startTime), + FirstTokenMs: nil, + }, nil + default: + return nil, s.writeGoogleError(c, http.StatusNotFound, "Unsupported action: "+action) + } + + mappedModel := s.getMappedModel(account, originalModel) + if mappedModel == "" { + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) + return nil, s.writeGoogleError(c, http.StatusForbidden, fmt.Sprintf("model %s not in whitelist", originalModel)) + } + billingModel := mappedModel + + // 获取 access_token + if s.tokenProvider == nil { + return nil, s.writeGoogleError(c, http.StatusBadGateway, "Antigravity token provider not configured") + } + accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) + if err != nil { + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: []byte(`{"error":{"message":"Failed to get upstream access token","status":"UNAVAILABLE"}}`), + } + } + + projectID, err := resolveAntigravityProjectID(account) + if err != nil { + _ = s.writeGoogleError(c, http.StatusBadRequest, err.Error()) + return nil, err + } + + // 代理 URL + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + + // Antigravity 上游要求必须包含身份提示词，注入到请求中 + injectedBody, err := injectIdentityPatchToGeminiRequest(body) + if err != nil { + return nil, s.writeGoogleError(c, http.StatusBadRequest, "Invalid request body") + } + + // 清理 Schema + if cleanedBody, err := cleanGeminiRequest(injectedBody); err == nil { + injectedBody = cleanedBody + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Cleaned request schema in forwarded request for account %s", account.Name) + } else { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Failed to clean schema: %v", err) + } + + // 包装请求 + wrappedBody, err := s.wrapV1InternalRequest(projectID, mappedModel, injectedBody) + if err != nil { + return nil, s.writeGoogleError(c, http.StatusInternalServerError, "Failed to build upstream request") + } + + // Antigravity 上游只支持流式请求，统一使用 streamGenerateContent + // 如果客户端请求非流式，在响应处理阶段会收集完整流式响应后返回 + upstreamAction := "streamGenerateContent" + + // 执行带重试的请求 + result, err := s.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: prefix, + account: account, + proxyURL: proxyURL, + accessToken: accessToken, + action: upstreamAction, + body: wrappedBody, + c: c, + httpUpstream: s.httpUpstream, + settingService: s.settingService, + accountRepo: s.accountRepo, + handleError: s.handleUpstreamError, + requestedModel: originalModel, + isStickySession: isStickySession, // ForwardGemini 由上层判断粘性会话 + groupID: forwardOpts.groupID, + sessionHash: forwardOpts.sessionHash, + }) + if err != nil { + // 检查是否是账号切换信号，转换为 UpstreamFailoverError 让 Handler 切换账号 + if switchErr, ok := IsAntigravityAccountSwitchError(err); ok { + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusServiceUnavailable, + ForceCacheBilling: switchErr.IsStickySession, + } + } + // 区分客户端取消和真正的上游失败，返回更准确的错误消息 + if c.Request.Context().Err() != nil { + return nil, s.writeGoogleError(c, http.StatusBadGateway, "Client disconnected before upstream response") + } + return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries") + } + resp := result.resp + defer func() { + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + }() + + // 处理错误响应 + if resp.StatusCode >= 400 { + respBody := s.readUpstreamErrorBody(resp) + contentType := resp.Header.Get("Content-Type") + // 尽早关闭原始响应体，释放连接；后续逻辑仍可能需要读取 body，因此用内存副本重新包装。 + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + + // 模型兜底：模型不存在且开启 fallback 时，自动用 fallback 模型重试一次 + if s.settingService != nil && s.settingService.IsModelFallbackEnabled(ctx) && + isModelNotFoundError(resp.StatusCode, respBody) { + fallbackModel := s.settingService.GetFallbackModel(ctx, PlatformAntigravity) + if fallbackModel != "" && fallbackModel != mappedModel { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Model not found (%s), retrying with fallback model %s (account: %s)", mappedModel, fallbackModel, account.Name) + + fallbackWrapped, err := s.wrapV1InternalRequest(projectID, fallbackModel, injectedBody) + if err == nil { + fallbackReq, err := antigravity.NewAPIRequest(ctx, upstreamAction, accessToken, fallbackWrapped) + if err == nil { + fallbackResp, err := s.httpUpstream.Do(fallbackReq, proxyURL, account.ID, account.Concurrency) + if err == nil && fallbackResp.StatusCode < 400 { + _ = resp.Body.Close() + resp = fallbackResp + } else if fallbackResp != nil { + _ = fallbackResp.Body.Close() + } + } + } + } + } + + // Gemini 原生请求中的 thoughtSignature 可能来自旧上下文/旧账号，触发上游严格校验后返回 + // "Corrupted thought signature."。检测到此类 400 时，将 thoughtSignature 清理为 dummy 值后重试一次。 + signatureCheckBody := respBody + if unwrapped, unwrapErr := s.unwrapV1InternalResponse(respBody); unwrapErr == nil && len(unwrapped) > 0 { + signatureCheckBody = unwrapped + } + if resp.StatusCode == http.StatusBadRequest && + s.settingService != nil && + s.settingService.IsSignatureRectifierEnabled(ctx) && + isSignatureRelatedError(signatureCheckBody) && + bytes.Contains(injectedBody, []byte(`"thoughtSignature"`)) { + upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(signatureCheckBody))) + upstreamDetail := s.getUpstreamErrorDetail(signatureCheckBody) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "signature_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: detected signature-related 400, retrying with cleaned thought signatures", account.ID) + + cleanedInjectedBody := CleanGeminiNativeThoughtSignatures(injectedBody) + retryWrappedBody, wrapErr := s.wrapV1InternalRequest(projectID, mappedModel, cleanedInjectedBody) + if wrapErr == nil { + retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: ctx, + prefix: prefix, + account: account, + proxyURL: proxyURL, + accessToken: accessToken, + action: upstreamAction, + body: retryWrappedBody, + c: c, + httpUpstream: s.httpUpstream, + settingService: s.settingService, + accountRepo: s.accountRepo, + handleError: s.handleUpstreamError, + requestedModel: originalModel, + isStickySession: isStickySession, + groupID: forwardOpts.groupID, + sessionHash: forwardOpts.sessionHash, + }) + if retryErr == nil { + retryResp := retryResult.resp + if retryResp.StatusCode < 400 { + resp = retryResp + } else { + retryRespBody := s.readUpstreamErrorBody(retryResp) + _ = retryResp.Body.Close() + retryOpsBody := retryRespBody + if retryUnwrapped, unwrapErr := s.unwrapV1InternalResponse(retryRespBody); unwrapErr == nil && len(retryUnwrapped) > 0 { + retryOpsBody = retryUnwrapped + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: retryResp.StatusCode, + UpstreamRequestID: retryResp.Header.Get("x-request-id"), + Kind: "signature_retry", + Message: sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(retryOpsBody))), + Detail: s.getUpstreamErrorDetail(retryOpsBody), + }) + respBody = retryRespBody + resp = &http.Response{ + StatusCode: retryResp.StatusCode, + Header: retryResp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryRespBody)), + } + contentType = resp.Header.Get("Content-Type") + } + } else { + if switchErr, ok := IsAntigravityAccountSwitchError(retryErr); ok { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: http.StatusServiceUnavailable, + Kind: "failover", + Message: sanitizeUpstreamErrorMessage(retryErr.Error()), + }) + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusServiceUnavailable, + ForceCacheBilling: switchErr.IsStickySession, + } + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: 0, + Kind: "signature_retry_request_error", + Message: sanitizeUpstreamErrorMessage(retryErr.Error()), + }) + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: signature retry request failed: %v", account.ID, retryErr) + } + } else { + logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: signature retry wrap failed: %v", account.ID, wrapErr) + } + } + + // fallback 成功：继续按正常响应处理 + if resp.StatusCode < 400 { + goto handleSuccess + } + + requestID := resp.Header.Get("x-request-id") + if requestID != "" { + c.Header("x-request-id", requestID) + } + + unwrapped, unwrapErr := s.unwrapV1InternalResponse(respBody) + unwrappedForOps := unwrapped + if unwrapErr != nil || len(unwrappedForOps) == 0 { + unwrappedForOps = respBody + } + s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, forwardOpts.groupID, forwardOpts.sessionHash, isStickySession) + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(unwrappedForOps)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + upstreamDetail := s.getUpstreamErrorDetail(unwrappedForOps) + + // Always record upstream context for Ops error logs, even when we will failover. + setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) + + // 精确匹配服务端配置类 400 错误，触发同账号重试 + failover + if resp.StatusCode == http.StatusBadRequest && isGoogleProjectConfigError(strings.ToLower(upstreamMsg)) { + log.Printf("%s status=400 google_config_error failover=true upstream_message=%q account=%d", prefix, upstreamMsg, account.ID) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: requestID, + Kind: "failover", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: unwrappedForOps, RetryableOnSameAccount: true} + } + + if s.shouldFailoverUpstreamError(resp.StatusCode) { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: requestID, + Kind: "failover", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: unwrappedForOps} + } + if contentType == "" { + contentType = "application/json" + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: requestID, + Kind: "http_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] upstream error status=%d body=%s", resp.StatusCode, truncateForLog(unwrappedForOps, 500)) + MarkResponseCommitted(c) + c.Data(resp.StatusCode, contentType, unwrappedForOps) + return nil, fmt.Errorf("antigravity upstream error: %d", resp.StatusCode) + } + +handleSuccess: + requestID := resp.Header.Get("x-request-id") + if requestID != "" { + c.Header("x-request-id", requestID) + } + + var usage *ClaudeUsage + var firstTokenMs *int + var clientDisconnect bool + + if stream { + // 客户端要求流式，直接透传 + streamRes, err := s.handleGeminiStreamingResponse(c, resp, startTime) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_error error=%v", prefix, err) + return nil, err + } + usage = streamRes.usage + firstTokenMs = streamRes.firstTokenMs + clientDisconnect = streamRes.clientDisconnect + } else { + // 客户端要求非流式，收集流式响应后返回 + streamRes, err := s.handleGeminiStreamToNonStreaming(c, resp, startTime) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_collect_error error=%v", prefix, err) + return nil, err + } + usage = streamRes.usage + firstTokenMs = streamRes.firstTokenMs + } + + if usage == nil { + usage = &ClaudeUsage{} + } + + // 判断是否为图片生成模型 + imageCount := 0 + if isImageGenerationModel(mappedModel) { + // Gemini 图片生成 API 每次请求只生成一张图片（API 限制） + imageCount = 1 + } + + return &ForwardResult{ + RequestID: requestID, + Usage: *usage, + Model: originalModel, + UpstreamModel: billingModel, + Stream: stream, + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnect, + ImageCount: imageCount, + ImageSize: imageSize, + ImageInputSize: imageInputSize, + }, nil +} + +// cleanGeminiRequest 清理 Gemini 请求体中的 Schema +func cleanGeminiRequest(body []byte) ([]byte, error) { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + return nil, err + } + + modified := false + + // 1. 清理 Tools + if tools, ok := payload["tools"].([]any); ok && len(tools) > 0 { + for _, t := range tools { + toolMap, ok := t.(map[string]any) + if !ok { + continue + } + + // function_declarations (snake_case) or functionDeclarations (camelCase) + var funcs []any + if f, ok := toolMap["functionDeclarations"].([]any); ok { + funcs = f + } else if f, ok := toolMap["function_declarations"].([]any); ok { + funcs = f + } + + if len(funcs) == 0 { + continue + } + + for _, f := range funcs { + funcMap, ok := f.(map[string]any) + if !ok { + continue + } + + if params, ok := funcMap["parameters"].(map[string]any); ok { + antigravity.DeepCleanUndefined(params) + cleaned := antigravity.CleanJSONSchema(params) + funcMap["parameters"] = cleaned + modified = true + } + } + } + } + + if !modified { + return body, nil + } + + return json.Marshal(payload) +} + +// filterEmptyPartsFromGeminiRequest 过滤掉 parts 为空的消息 +// Gemini API 不接受空 parts，需要在请求前过滤 +func filterEmptyPartsFromGeminiRequest(body []byte) ([]byte, error) { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + return nil, err + } + + contents, ok := payload["contents"].([]any) + if !ok || len(contents) == 0 { + return body, nil + } + + filtered := make([]any, 0, len(contents)) + modified := false + + for _, c := range contents { + contentMap, ok := c.(map[string]any) + if !ok { + filtered = append(filtered, c) + continue + } + + parts, hasParts := contentMap["parts"] + if !hasParts { + filtered = append(filtered, c) + continue + } + + partsSlice, ok := parts.([]any) + if !ok { + filtered = append(filtered, c) + continue + } + + // 跳过 parts 为空数组的消息 + if len(partsSlice) == 0 { + modified = true + continue + } + + filtered = append(filtered, c) + } + + if !modified { + return body, nil + } + + payload["contents"] = filtered + return json.Marshal(payload) +} diff --git a/backend/internal/service/antigravity_gateway_retry.go b/backend/internal/service/antigravity_gateway_retry.go new file mode 100644 index 0000000000..f2f5d44f84 --- /dev/null +++ b/backend/internal/service/antigravity_gateway_retry.go @@ -0,0 +1,1279 @@ +package service + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "log" + mathrand "math/rand" + "net" + "net/http" + "os" + "strconv" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" +) + +// antigravityRetryLoopParams 重试循环的参数 +type antigravityRetryLoopParams struct { + ctx context.Context + prefix string + account *Account + proxyURL string + accessToken string + action string + body []byte + c *gin.Context + httpUpstream HTTPUpstream + settingService *SettingService + accountRepo AccountRepository // 用于智能重试的模型级别限流 + handleError func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult + requestedModel string // 用于限流检查的原始请求模型 + isStickySession bool // 是否为粘性会话（用于账号切换时的缓存计费判断） + groupID int64 // 用于模型级限流时清除粘性会话 + sessionHash string // 用于模型级限流时清除粘性会话 +} + +// antigravityRetryLoopResult 重试循环的结果 +type antigravityRetryLoopResult struct { + resp *http.Response +} + +// resolveAntigravityForwardBaseURL 解析转发用 base URL。 +// +// 默认使用生产端点 cloudcode-pa.googleapis.com（antigravity.BaseURLs 的首个地址， +// 与账号 OAuth 登录/测试连接所用的 antigravity.BaseURL 一致）。 +// +// 历史上这里改用 ForwardBaseURLs()（把 daily/sandbox 排到首位）并默认取首个地址， +// 导致网关把带生产 OAuth token 的请求发到 daily-cloudcode-pa.sandbox.googleapis.com， +// 上游拒绝 → 账号被 401「Invalid bearer token」/502 打入临时不可调度且无法恢复 +// （见 #3611 / #2962）。后台「测试连接」用的是生产端点，所以「测试成功但网关 401」。 +// +// daily/sandbox 端点仅供内部联调，需显式设置 +// GATEWAY_ANTIGRAVITY_FORWARD_BASE_URL=daily（或 sandbox）才启用。 +func resolveAntigravityForwardBaseURL() string { + baseURLs := antigravity.BaseURLs + if len(baseURLs) == 0 { + return "" + } + mode := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityForwardBaseURLEnv))) + if (mode == "daily" || mode == "sandbox") && len(baseURLs) > 1 { + return baseURLs[1] + } + return baseURLs[0] +} + +// smartRetryAction 智能重试的处理结果 +type smartRetryAction int + +const ( + smartRetryActionContinue smartRetryAction = iota // 继续默认重试逻辑 + smartRetryActionBreakWithResp // 结束循环并返回 resp + smartRetryActionContinueURL // 继续 URL fallback 循环 +) + +// smartRetryResult 智能重试的结果 +type smartRetryResult struct { + action smartRetryAction + resp *http.Response + err error + switchError *AntigravityAccountSwitchError // 模型限流时返回账号切换信号 +} + +// handleSmartRetry 处理 OAuth 账号的智能重试逻辑 +// 将 429/503 限流处理逻辑抽取为独立函数，减少 antigravityRetryLoop 的复杂度 +func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParams, resp *http.Response, respBody []byte, baseURL string, urlIdx int, availableURLs []string) *smartRetryResult { + // "Resource has been exhausted" 是 URL 级别限流，切换 URL（仅 429） + if resp.StatusCode == http.StatusTooManyRequests && isURLLevelRateLimit(respBody) && urlIdx < len(availableURLs)-1 { + logger.LegacyPrintf("service.antigravity_gateway", "%s URL fallback (429): %s -> %s", p.prefix, baseURL, availableURLs[urlIdx+1]) + return &smartRetryResult{action: smartRetryActionContinueURL} + } + + category := antigravity429Unknown + if resp.StatusCode == http.StatusTooManyRequests { + category = classifyAntigravity429(respBody) + } + + // 判断是否触发智能重试 + shouldSmartRetry, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, respBody) + + // AI Credits 超量请求： + // 仅在上游明确返回免费配额耗尽时才允许切换到 credits。 + if resp.StatusCode == http.StatusTooManyRequests && + category == antigravity429QuotaExhausted && + p.account.IsOveragesEnabled() && + !p.account.isCreditsExhausted() { + result := s.attemptCreditsOveragesRetry(p, baseURL, modelName, waitDuration, resp.StatusCode, respBody) + if result.handled && result.resp != nil { + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: result.resp, + } + } + } + + // 情况1: retryDelay >= 阈值，限流模型并切换账号 + if shouldRateLimitModel { + // 单账号 503 退避重试模式：不设限流、不切换账号，改为原地等待+重试 + // 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的，等几秒就能恢复。 + // 多账号场景下切换账号是最优选择，但单账号场景下设限流毫无意义（只会导致双重等待）。 + if resp.StatusCode == http.StatusServiceUnavailable && isSingleAccountRetry(p.ctx) { + return s.handleSingleAccountRetryInPlace(p, resp, respBody, baseURL, waitDuration, modelName) + } + + rateLimitDuration := waitDuration + if rateLimitDuration <= 0 { + rateLimitDuration = antigravityDefaultRateLimitDuration + } + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d oauth_long_delay model=%s account=%d upstream_retry_delay=%v body=%s (model rate limit, switch account)", + p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration, truncateForLog(respBody, 200)) + + resetAt := time.Now().Add(rateLimitDuration) + if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, false) { + p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited account=%d (no model mapping)", p.prefix, resp.StatusCode, p.account.ID) + } + s.clearStickySession(p.ctx, p.groupID, p.sessionHash) + + // 返回账号切换信号，让上层切换账号重试 + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + switchError: &AntigravityAccountSwitchError{ + OriginalAccountID: p.account.ID, + RateLimitedModel: modelName, + IsStickySession: p.isStickySession, + }, + } + } + + // 情况2: retryDelay < 阈值（或 MODEL_CAPACITY_EXHAUSTED），智能重试 + if shouldSmartRetry { + var lastRetryResp *http.Response + var lastRetryBody []byte + + // MODEL_CAPACITY_EXHAUSTED 使用独立的重试参数（60 次，固定 1s 间隔） + maxAttempts := antigravitySmartRetryMaxAttempts + if isModelCapacityExhausted { + maxAttempts = antigravityModelCapacityRetryMaxAttempts + waitDuration = antigravityModelCapacityRetryWait + + // 全局去重：如果其他 goroutine 已在重试同一模型且尚在 cooldown 中，直接返回 503 + if modelName != "" { + modelCapacityExhaustedMu.RLock() + cooldownUntil, exists := modelCapacityExhaustedUntil[modelName] + modelCapacityExhaustedMu.RUnlock() + if exists && time.Now().Before(cooldownUntil) { + log.Printf("%s status=%d model_capacity_exhausted_dedup model=%s account=%d cooldown_until=%v (skip retry)", + p.prefix, resp.StatusCode, modelName, p.account.ID, cooldownUntil.Format("15:04:05")) + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + }, + } + } + } + } + + for attempt := 1; attempt <= maxAttempts; attempt++ { + log.Printf("%s status=%d oauth_smart_retry attempt=%d/%d delay=%v model=%s account=%d", + p.prefix, resp.StatusCode, attempt, maxAttempts, waitDuration, modelName, p.account.ID) + + timer := time.NewTimer(waitDuration) + select { + case <-p.ctx.Done(): + timer.Stop() + log.Printf("%s status=context_canceled_during_smart_retry", p.prefix) + return &smartRetryResult{action: smartRetryActionBreakWithResp, err: p.ctx.Err()} + case <-timer.C: + } + + // 智能重试：创建新请求 + retryReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=smart_retry_request_build_failed error=%v", p.prefix, err) + p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + }, + } + } + + retryResp, retryErr := p.httpUpstream.Do(retryReq, p.proxyURL, p.account.ID, p.account.Concurrency) + if retryErr == nil && retryResp != nil && retryResp.StatusCode != http.StatusTooManyRequests && retryResp.StatusCode != http.StatusServiceUnavailable { + log.Printf("%s status=%d smart_retry_success attempt=%d/%d", p.prefix, retryResp.StatusCode, attempt, maxAttempts) + // 重试成功，清除 MODEL_CAPACITY_EXHAUSTED cooldown + if isModelCapacityExhausted && modelName != "" { + modelCapacityExhaustedMu.Lock() + delete(modelCapacityExhaustedUntil, modelName) + modelCapacityExhaustedMu.Unlock() + } + return &smartRetryResult{action: smartRetryActionBreakWithResp, resp: retryResp} + } + + // 网络错误时，继续重试 + if retryErr != nil || retryResp == nil { + log.Printf("%s status=smart_retry_network_error attempt=%d/%d error=%v", p.prefix, attempt, maxAttempts, retryErr) + continue + } + + // 重试失败，关闭之前的响应 + if lastRetryResp != nil { + _ = lastRetryResp.Body.Close() + } + lastRetryResp = retryResp + if retryResp != nil { + lastRetryBody, _ = io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) + _ = retryResp.Body.Close() + } + + // 解析新的重试信息，用于下次重试的等待时间（MODEL_CAPACITY_EXHAUSTED 使用固定循环，跳过） + if !isModelCapacityExhausted && attempt < maxAttempts && lastRetryBody != nil { + newShouldRetry, _, newWaitDuration, _, _ := shouldTriggerAntigravitySmartRetry(p.account, lastRetryBody) + if newShouldRetry && newWaitDuration > 0 { + waitDuration = newWaitDuration + } + } + } + + // 所有重试都失败 + rateLimitDuration := waitDuration + if rateLimitDuration <= 0 { + rateLimitDuration = antigravityDefaultRateLimitDuration + } + retryBody := lastRetryBody + if retryBody == nil { + retryBody = respBody + } + + // MODEL_CAPACITY_EXHAUSTED：模型容量不足，切换账号无意义 + // 直接返回上游错误响应，不设置模型限流，不切换账号 + if isModelCapacityExhausted { + // 设置 cooldown，让后续请求快速失败，避免重复重试 + if modelName != "" { + modelCapacityExhaustedMu.Lock() + modelCapacityExhaustedUntil[modelName] = time.Now().Add(antigravityModelCapacityCooldown) + modelCapacityExhaustedMu.Unlock() + } + log.Printf("%s status=%d smart_retry_exhausted_model_capacity attempts=%d model=%s account=%d body=%s (model capacity exhausted, not switching account)", + p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, truncateForLog(retryBody, 200)) + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + }, + } + } + + // 单账号 503 退避重试模式：智能重试耗尽后不设限流、不切换账号， + // 直接返回 503 让 Handler 层的单账号退避循环做最终处理。 + if resp.StatusCode == http.StatusServiceUnavailable && isSingleAccountRetry(p.ctx) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d smart_retry_exhausted_single_account attempts=%d model=%s account=%d body=%s (return 503 directly)", + p.prefix, resp.StatusCode, antigravitySmartRetryMaxAttempts, modelName, p.account.ID, truncateForLog(retryBody, 200)) + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + }, + } + } + + log.Printf("%s status=%d smart_retry_exhausted attempts=%d model=%s account=%d upstream_retry_delay=%v body=%s (switch account)", + p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, rateLimitDuration, truncateForLog(retryBody, 200)) + + resetAt := time.Now().Add(rateLimitDuration) + s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, true) + + // 清除粘性会话绑定，避免下次请求仍命中限流账号 + s.clearStickySession(p.ctx, p.groupID, p.sessionHash) + + // 返回账号切换信号，让上层切换账号重试 + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + switchError: &AntigravityAccountSwitchError{ + OriginalAccountID: p.account.ID, + RateLimitedModel: modelName, + IsStickySession: p.isStickySession, + }, + } + } + + // 未触发智能重试，继续默认重试逻辑 + return &smartRetryResult{action: smartRetryActionContinue} +} + +// handleSingleAccountRetryInPlace 单账号 503 退避重试的原地重试逻辑。 +// +// 在多账号场景下，收到 503 + 长 retryDelay（≥ 7s）时会设置模型限流 + 切换账号； +// 但在单账号场景下，设限流毫无意义（因为切换回来的还是同一个账号，还要等限流过期）。 +// 此方法改为在 Service 层原地等待 + 重试，避免双重等待问题： +// +// 旧流程：Service 设限流 → Handler 退避等待 → Service 等限流过期 → 再请求（总耗时 = 退避 + 限流） +// 新流程：Service 直接等 retryDelay → 重试 → 成功/再等 → 重试...（总耗时 ≈ 实际 retryDelay × 重试次数） +// +// 约束： +// - 单次等待不超过 antigravitySingleAccountSmartRetryMaxWait +// - 总累计等待不超过 antigravitySingleAccountSmartRetryTotalMaxWait +// - 最多重试 antigravitySingleAccountSmartRetryMaxAttempts 次 +func (s *AntigravityGatewayService) handleSingleAccountRetryInPlace( + p antigravityRetryLoopParams, + resp *http.Response, + respBody []byte, + baseURL string, + waitDuration time.Duration, + modelName string, +) *smartRetryResult { + // 限制单次等待时间 + if waitDuration > antigravitySingleAccountSmartRetryMaxWait { + waitDuration = antigravitySingleAccountSmartRetryMaxWait + } + if waitDuration < antigravitySmartRetryMinWait { + waitDuration = antigravitySmartRetryMinWait + } + + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_in_place model=%s account=%d upstream_retry_delay=%v (retrying in-place instead of rate-limiting)", + p.prefix, resp.StatusCode, modelName, p.account.ID, waitDuration) + + var lastRetryResp *http.Response + var lastRetryBody []byte + totalWaited := time.Duration(0) + + for attempt := 1; attempt <= antigravitySingleAccountSmartRetryMaxAttempts; attempt++ { + // 检查累计等待是否超限 + if totalWaited+waitDuration > antigravitySingleAccountSmartRetryTotalMaxWait { + remaining := antigravitySingleAccountSmartRetryTotalMaxWait - totalWaited + if remaining <= 0 { + logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: total_wait_exceeded total=%v max=%v, giving up", + p.prefix, totalWaited, antigravitySingleAccountSmartRetryTotalMaxWait) + break + } + waitDuration = remaining + } + + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry attempt=%d/%d delay=%v total_waited=%v model=%s account=%d", + p.prefix, resp.StatusCode, attempt, antigravitySingleAccountSmartRetryMaxAttempts, waitDuration, totalWaited, modelName, p.account.ID) + + timer := time.NewTimer(waitDuration) + select { + case <-p.ctx.Done(): + timer.Stop() + logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_single_account_retry", p.prefix) + return &smartRetryResult{action: smartRetryActionBreakWithResp, err: p.ctx.Err()} + case <-timer.C: + } + totalWaited += waitDuration + + // 创建新请求 + retryReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: request_build_failed error=%v", p.prefix, err) + break + } + + retryResp, retryErr := p.httpUpstream.Do(retryReq, p.proxyURL, p.account.ID, p.account.Concurrency) + if retryErr == nil && retryResp != nil && retryResp.StatusCode != http.StatusTooManyRequests && retryResp.StatusCode != http.StatusServiceUnavailable { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_success attempt=%d/%d total_waited=%v", + p.prefix, retryResp.StatusCode, attempt, antigravitySingleAccountSmartRetryMaxAttempts, totalWaited) + // 关闭之前的响应 + if lastRetryResp != nil { + _ = lastRetryResp.Body.Close() + } + return &smartRetryResult{action: smartRetryActionBreakWithResp, resp: retryResp} + } + + // 网络错误时继续重试 + if retryErr != nil || retryResp == nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: network_error attempt=%d/%d error=%v", + p.prefix, attempt, antigravitySingleAccountSmartRetryMaxAttempts, retryErr) + continue + } + + // 关闭之前的响应 + if lastRetryResp != nil { + _ = lastRetryResp.Body.Close() + } + lastRetryResp = retryResp + lastRetryBody, _ = io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) + _ = retryResp.Body.Close() + + // 解析新的重试信息，更新下次等待时间 + if attempt < antigravitySingleAccountSmartRetryMaxAttempts && lastRetryBody != nil { + _, _, newWaitDuration, _, _ := shouldTriggerAntigravitySmartRetry(p.account, lastRetryBody) + if newWaitDuration > 0 { + waitDuration = newWaitDuration + if waitDuration > antigravitySingleAccountSmartRetryMaxWait { + waitDuration = antigravitySingleAccountSmartRetryMaxWait + } + if waitDuration < antigravitySmartRetryMinWait { + waitDuration = antigravitySmartRetryMinWait + } + } + } + } + + // 所有重试都失败，不设限流，直接返回 503 + // Handler 层的单账号退避循环会做最终处理 + retryBody := lastRetryBody + if retryBody == nil { + retryBody = respBody + } + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_exhausted attempts=%d total_waited=%v model=%s account=%d body=%s (return 503 directly)", + p.prefix, resp.StatusCode, antigravitySingleAccountSmartRetryMaxAttempts, totalWaited, modelName, p.account.ID, truncateForLog(retryBody, 200)) + + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryBody)), + }, + } +} + +// antigravityRetryLoop 执行带 URL fallback 的重试循环 +func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopParams) (*antigravityRetryLoopResult, error) { + // 预检查：模型限流 + overages 启用 + 积分未耗尽 → 直接注入 AI Credits + overagesInjected := false + if p.requestedModel != "" && p.account.Platform == PlatformAntigravity && + p.account.IsOveragesEnabled() && !p.account.isCreditsExhausted() && + p.account.isModelRateLimitedWithContext(p.ctx, p.requestedModel) { + if creditsBody := injectEnabledCreditTypes(p.body); creditsBody != nil { + p.body = creditsBody + overagesInjected = true + logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: model_rate_limited_credits_inject model=%s account=%d (injecting enabledCreditTypes)", + p.prefix, p.requestedModel, p.account.ID) + } + } + + // 预检查：如果账号已限流，直接返回切换信号 + if p.requestedModel != "" { + if remaining := p.account.GetRateLimitRemainingTimeWithContext(p.ctx, p.requestedModel); remaining > 0 { + // 已注入积分的请求不再受普通模型限流预检查阻断。 + if overagesInjected { + logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: credits_injected_ignore_rate_limit remaining=%v model=%s account=%d", + p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) + } else if isSingleAccountRetry(p.ctx) { + // 单账号 503 退避重试模式：跳过限流预检查，直接发请求。 + // 首次请求设的限流是为了多账号调度器跳过该账号，在单账号模式下无意义。 + // 如果上游确实还不可用，handleSmartRetry → handleSingleAccountRetryInPlace + // 会在 Service 层原地等待+重试，不需要在预检查这里等。 + logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: single_account_retry skipping rate_limit remaining=%v model=%s account=%d (will retry in-place if 503)", + p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) + } else { + logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: rate_limit_switch remaining=%v model=%s account=%d", + p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) + return nil, &AntigravityAccountSwitchError{ + OriginalAccountID: p.account.ID, + RateLimitedModel: p.requestedModel, + IsStickySession: p.isStickySession, + } + } + } + } + + baseURL := resolveAntigravityForwardBaseURL() + if baseURL == "" { + return nil, errors.New("no antigravity forward base url configured") + } + availableURLs := []string{baseURL} + + var resp *http.Response + var usedBaseURL string + logBody := p.settingService != nil && p.settingService.cfg != nil && p.settingService.cfg.Gateway.LogUpstreamErrorBody + maxBytes := 2048 + if p.settingService != nil && p.settingService.cfg != nil && p.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > 0 { + maxBytes = p.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes + } + getUpstreamDetail := func(body []byte) string { + if !logBody { + return "" + } + return truncateString(string(body), maxBytes) + } + +urlFallbackLoop: + for urlIdx, baseURL := range availableURLs { + usedBaseURL = baseURL + allAttemptsInternal500 := true // 追踪本轮所有 attempt 是否全部命中 INTERNAL 500 + for attempt := 1; attempt <= antigravityMaxRetries; attempt++ { + select { + case <-p.ctx.Done(): + logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled error=%v", p.prefix, p.ctx.Err()) + return nil, p.ctx.Err() + default: + } + + upstreamReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) + if err != nil { + return nil, err + } + + resp, err = p.httpUpstream.Do(upstreamReq, p.proxyURL, p.account.ID, p.account.Concurrency) + if err == nil && resp == nil { + err = errors.New("upstream returned nil response") + } + if err != nil { + safeErr := sanitizeUpstreamErrorMessage(err.Error()) + appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ + Platform: p.account.Platform, + AccountID: p.account.ID, + AccountName: p.account.Name, + UpstreamStatusCode: 0, + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "request_error", + Message: safeErr, + }) + if shouldAntigravityFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 { + logger.LegacyPrintf("service.antigravity_gateway", "%s URL fallback (connection error): %s -> %s", p.prefix, baseURL, availableURLs[urlIdx+1]) + continue urlFallbackLoop + } + if attempt < antigravityMaxRetries { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=request_failed retry=%d/%d error=%v", p.prefix, attempt, antigravityMaxRetries, err) + if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) + return nil, p.ctx.Err() + } + continue + } + logger.LegacyPrintf("service.antigravity_gateway", "%s status=request_failed retries_exhausted error=%v", p.prefix, err) + setOpsUpstreamError(p.c, 0, safeErr, "") + return nil, fmt.Errorf("upstream request failed after retries: %w", err) + } + + // 统一处理错误响应 + if resp.StatusCode >= 400 { + respBody := s.readUpstreamErrorBody(resp) + _ = resp.Body.Close() + + if overagesInjected && shouldMarkCreditsExhausted(resp, respBody, nil) { + modelKey := resolveCreditsOveragesModelKey(p.ctx, p.account, "", p.requestedModel) + s.handleCreditsRetryFailure(p.ctx, p.prefix, modelKey, p.account, &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + }, nil) + } + + // ★ 统一入口：自定义错误码 + 临时不可调度 + if handled, outStatus, policyErr := s.applyErrorPolicy(p, resp.StatusCode, resp.Header, respBody); handled { + if policyErr != nil { + return nil, policyErr + } + resp = &http.Response{ + StatusCode: outStatus, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + } + break urlFallbackLoop + } + + // 429/503 限流处理：区分 URL 级别限流、智能重试和账户配额限流 + if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusServiceUnavailable { + // 尝试智能重试处理（OAuth 账号专用） + smartResult := s.handleSmartRetry(p, resp, respBody, baseURL, urlIdx, availableURLs) + switch smartResult.action { + case smartRetryActionContinueURL: + continue urlFallbackLoop + case smartRetryActionBreakWithResp: + if smartResult.err != nil { + return nil, smartResult.err + } + // 模型限流时返回切换账号信号 + if smartResult.switchError != nil { + return nil, smartResult.switchError + } + resp = smartResult.resp + break urlFallbackLoop + } + // smartRetryActionContinue: 继续默认重试逻辑 + + // 账户/模型配额限流，重试 3 次（指数退避）- 默认逻辑（非 OAuth 账号或解析失败） + if attempt < antigravityMaxRetries { + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ + Platform: p.account.Platform, + AccountID: p.account.ID, + AccountName: p.account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "retry", + Message: upstreamMsg, + Detail: getUpstreamDetail(respBody), + }) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d retry=%d/%d body=%s", p.prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 200)) + if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) + return nil, p.ctx.Err() + } + continue + } + + // 重试用尽，标记账户限流 + p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited base_url=%s body=%s", p.prefix, resp.StatusCode, baseURL, truncateForLog(respBody, 200)) + resp = &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + } + break urlFallbackLoop + } + + // 其他可重试错误（500/502/504/529，不包括 429 和 503） + if shouldRetryAntigravityError(resp.StatusCode) { + if attempt < antigravityMaxRetries { + upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ + Platform: p.account.Platform, + AccountID: p.account.ID, + AccountName: p.account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "retry", + Message: upstreamMsg, + Detail: getUpstreamDetail(respBody), + }) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d retry=%d/%d body=%s", p.prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 500)) + if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) + return nil, p.ctx.Err() + } + // 追踪 INTERNAL 500：非匹配的 attempt 清除标记 + if !isAntigravityInternalServerError(resp.StatusCode, respBody) { + allAttemptsInternal500 = false + } + continue + } + } + + // INTERNAL 500 渐进惩罚：3 次重试全部命中特定 500 时递增计数器并惩罚 + if allAttemptsInternal500 && isAntigravityInternalServerError(resp.StatusCode, respBody) { + s.handleInternal500RetryExhausted(p.ctx, p.prefix, p.account) + } + + // 其他 4xx 错误或重试用尽，直接返回 + resp = &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + } + break urlFallbackLoop + } + + // 成功响应（< 400） + break urlFallbackLoop + } + } + + if resp != nil && resp.StatusCode < 400 && usedBaseURL != "" { + antigravity.DefaultURLAvailability.MarkSuccess(usedBaseURL) + } + + // 成功响应时清零 INTERNAL 500 连续失败计数器（覆盖所有成功路径，含 smart retry） + if resp != nil && resp.StatusCode < 400 { + s.resetInternal500Counter(p.ctx, p.prefix, p.account.ID) + } + + return &antigravityRetryLoopResult{resp: resp}, nil +} + +// shouldRetryAntigravityError 判断是否应该重试 +func shouldRetryAntigravityError(statusCode int) bool { + switch statusCode { + case 429, 500, 502, 503, 504, 529: + return true + default: + return false + } +} + +// isURLLevelRateLimit 判断是否为 URL 级别的限流（应切换 URL 重试） +// "Resource has been exhausted" 是 URL/节点级别限流，切换 URL 可能成功 +// "exhausted your capacity on this model" 是账户/模型配额限流，切换 URL 无效 +func isURLLevelRateLimit(body []byte) bool { + // 快速检查：包含 "Resource has been exhausted" 且不包含 "capacity on this model" + bodyStr := string(body) + return strings.Contains(bodyStr, "Resource has been exhausted") && + !strings.Contains(bodyStr, "capacity on this model") +} + +// isAntigravityConnectionError 判断是否为连接错误（网络超时、DNS 失败、连接拒绝） +func isAntigravityConnectionError(err error) bool { + if err == nil { + return false + } + + // 检查超时错误 + var netErr net.Error + if errors.As(err, &netErr) && netErr.Timeout() { + return true + } + + // 检查连接错误（DNS 失败、连接拒绝） + var opErr *net.OpError + return errors.As(err, &opErr) +} + +// shouldAntigravityFallbackToNextURL 判断是否应切换到下一个 URL +// 仅连接错误和 HTTP 429 触发 URL 降级 +func shouldAntigravityFallbackToNextURL(err error, statusCode int) bool { + if isAntigravityConnectionError(err) { + return true + } + return statusCode == http.StatusTooManyRequests +} + +// getSessionID 从 gin.Context 获取 session_id（用于日志追踪） +func getSessionID(c *gin.Context) string { + if c == nil { + return "" + } + return c.GetHeader("session_id") +} + +// logPrefix 生成统一的日志前缀 +func logPrefix(sessionID, accountName string) string { + if sessionID != "" { + return fmt.Sprintf("[antigravity-Forward] session=%s account=%s", sessionID, accountName) + } + return fmt.Sprintf("[antigravity-Forward] account=%s", accountName) +} + +func (s *AntigravityGatewayService) shouldFailoverUpstreamError(statusCode int) bool { + switch statusCode { + case 401, 403, 429, 529: + return true + default: + return statusCode >= 500 + } +} + +// isGoogleProjectConfigError 判断（已提取的小写）错误消息是否属于 Google 服务端配置类问题。 +// 只精确匹配已知的服务端侧错误，避免对客户端请求错误做无意义重试。 +// 适用于所有走 Google 后端的平台（Antigravity、Gemini）。 +func isGoogleProjectConfigError(lowerMsg string) bool { + // Google 间歇性 Bug：Project ID 有效但被临时识别失败 + return strings.Contains(lowerMsg, "invalid project resource name") +} + +// googleConfigErrorCooldown 服务端配置类 400 错误的临时封禁时长 +const googleConfigErrorCooldown = 1 * time.Minute + +// tempUnscheduleGoogleConfigError 对服务端配置类 400 错误触发临时封禁， +// 避免短时间内反复调度到同一个有问题的账号。 +func tempUnscheduleGoogleConfigError(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) { + until := time.Now().Add(googleConfigErrorCooldown) + reason := "400: invalid project resource name (auto temp-unschedule 1m)" + if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil { + log.Printf("%s temp_unschedule_failed account=%d error=%v", logPrefix, accountID, err) + } else { + log.Printf("%s temp_unscheduled account=%d until=%v reason=%q", logPrefix, accountID, until.Format("15:04:05"), reason) + } +} + +// emptyResponseCooldown 空流式响应的临时封禁时长 +const emptyResponseCooldown = 1 * time.Minute + +// tempUnscheduleEmptyResponse 对空流式响应触发临时封禁， +// 避免短时间内反复调度到同一个返回空响应的账号。 +func tempUnscheduleEmptyResponse(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) { + until := time.Now().Add(emptyResponseCooldown) + reason := "empty stream response (auto temp-unschedule 1m)" + if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil { + log.Printf("%s temp_unschedule_failed account=%d error=%v", logPrefix, accountID, err) + } else { + log.Printf("%s temp_unscheduled account=%d until=%v reason=%q", logPrefix, accountID, until.Format("15:04:05"), reason) + } +} + +// sleepAntigravityBackoffWithContext 带 context 取消检查的退避等待 +// 返回 true 表示正常完成等待，false 表示 context 已取消 +func sleepAntigravityBackoffWithContext(ctx context.Context, attempt int) bool { + delay := antigravityRetryBaseDelay * time.Duration(1< antigravityRetryMaxDelay { + delay = antigravityRetryMaxDelay + } + + // +/- 20% jitter + r := mathrand.New(mathrand.NewSource(time.Now().UnixNano())) + jitter := time.Duration(float64(delay) * 0.2 * (r.Float64()*2 - 1)) + sleepFor := delay + jitter + if sleepFor < 0 { + sleepFor = 0 + } + + timer := time.NewTimer(sleepFor) + select { + case <-ctx.Done(): + timer.Stop() + return false + case <-timer.C: + return true + } +} + +// isSingleAccountRetry 检查 context 中是否设置了单账号退避重试标记 +func isSingleAccountRetry(ctx context.Context) bool { + v, _ := SingleAccountRetryFromContext(ctx) + return v +} + +// setModelRateLimitByModelName 使用官方模型 ID 设置模型级限流 +// 直接使用上游返回的模型 ID（如 claude-sonnet-4-5）作为限流 key +// 返回是否已成功设置（若模型名为空或 repo 为 nil 将返回 false） +func setModelRateLimitByModelName(ctx context.Context, repo AccountRepository, accountID int64, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { + if repo == nil || modelName == "" { + return false + } + // 直接使用官方模型 ID 作为 key，不再转换为 scope + if err := repo.SetModelRateLimit(ctx, accountID, modelName, resetAt); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_failed model=%s error=%v", prefix, statusCode, modelName, err) + return false + } + if afterSmartRetry { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_after_smart_retry model=%s account=%d reset_in=%v", prefix, statusCode, modelName, accountID, time.Until(resetAt).Truncate(time.Second)) + } else { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", prefix, statusCode, modelName, accountID, time.Until(resetAt).Truncate(time.Second)) + } + return true +} + +func (s *AntigravityGatewayService) setAntigravityModelRateLimits(ctx context.Context, repo AccountRepository, account *Account, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { + if account == nil || repo == nil { + return false + } + keys := antigravityModelRateLimitKeys(modelName) + if len(keys) == 0 { + return false + } + + success := false + for _, key := range keys { + if setModelRateLimitByModelName(ctx, repo, account.ID, key, prefix, statusCode, resetAt, afterSmartRetry) { + s.updateAccountModelRateLimitInCache(ctx, account, key, resetAt) + success = true + } + } + return success +} + +func (s *AntigravityGatewayService) clearStickySession(ctx context.Context, groupID int64, sessionHash string) { + if s == nil || s.cache == nil || strings.TrimSpace(sessionHash) == "" { + return + } + if err := s.cache.DeleteSessionAccountID(ctx, groupID, sessionHash); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] sticky_session_clear_failed group_id=%d session=%s err=%v", groupID, shortSessionHash(sessionHash), err) + } +} + +func antigravityFallbackCooldownSeconds() (time.Duration, bool) { + raw := strings.TrimSpace(os.Getenv(antigravityFallbackSecondsEnv)) + if raw == "" { + return 0, false + } + seconds, err := strconv.Atoi(raw) + if err != nil || seconds <= 0 { + return 0, false + } + return time.Duration(seconds) * time.Second, true +} + +// antigravitySmartRetryInfo 智能重试所需的信息 +type antigravitySmartRetryInfo struct { + RetryDelay time.Duration // 重试延迟时间 + ModelName string // 限流的模型名称（如 "claude-sonnet-4-5"） + IsModelCapacityExhausted bool // 是否为模型容量不足（MODEL_CAPACITY_EXHAUSTED） +} + +// parseAntigravitySmartRetryInfo 解析 Google RPC RetryInfo 和 ErrorInfo 信息 +// 返回解析结果，如果解析失败或不满足条件返回 nil +// +// 支持两种情况： +// 1. 429 RESOURCE_EXHAUSTED + RATE_LIMIT_EXCEEDED： +// - error.status == "RESOURCE_EXHAUSTED" +// - error.details[].reason == "RATE_LIMIT_EXCEEDED" +// +// 2. 503 UNAVAILABLE + MODEL_CAPACITY_EXHAUSTED： +// - error.status == "UNAVAILABLE" +// - error.details[].reason == "MODEL_CAPACITY_EXHAUSTED" +// +// 必须满足以下条件才会返回有效值： +// - error.details[] 中存在 @type == "type.googleapis.com/google.rpc.RetryInfo" 的元素 +// - 该元素包含 retryDelay 字段，格式为 "数字s"（如 "0.201506475s"） +func parseAntigravitySmartRetryInfo(body []byte) *antigravitySmartRetryInfo { + var parsed map[string]any + if err := json.Unmarshal(body, &parsed); err != nil { + return nil + } + + errObj, ok := parsed["error"].(map[string]any) + if !ok { + return nil + } + + // 检查 status 是否符合条件 + // 情况1: 429 RESOURCE_EXHAUSTED (需要进一步检查 reason == RATE_LIMIT_EXCEEDED) + // 情况2: 503 UNAVAILABLE (需要进一步检查 reason == MODEL_CAPACITY_EXHAUSTED) + status, _ := errObj["status"].(string) + isResourceExhausted := status == googleRPCStatusResourceExhausted + isUnavailable := status == googleRPCStatusUnavailable + + if !isResourceExhausted && !isUnavailable { + return nil + } + + details, ok := errObj["details"].([]any) + if !ok { + return nil + } + + var retryDelay time.Duration + var modelName string + var hasRateLimitExceeded bool // 429 需要此 reason + var hasModelCapacityExhausted bool // 503 需要此 reason + + for _, d := range details { + dm, ok := d.(map[string]any) + if !ok { + continue + } + + atType, _ := dm["@type"].(string) + + // 从 ErrorInfo 提取模型名称和 reason + if atType == googleRPCTypeErrorInfo { + if meta, ok := dm["metadata"].(map[string]any); ok { + if model, ok := meta["model"].(string); ok { + modelName = normalizeAntigravityModelName(model) + } + } + // 检查 reason + if reason, ok := dm["reason"].(string); ok { + if reason == googleRPCReasonModelCapacityExhausted { + hasModelCapacityExhausted = true + } + if reason == googleRPCReasonRateLimitExceeded { + hasRateLimitExceeded = true + } + } + continue + } + + // 从 RetryInfo 提取重试延迟 + if atType == googleRPCTypeRetryInfo { + delay, ok := dm["retryDelay"].(string) + if !ok || delay == "" { + continue + } + // 使用 time.ParseDuration 解析，支持所有 Go duration 格式 + // 例如: "0.5s", "10s", "4m50s", "1h30m", "200ms" 等 + dur, err := time.ParseDuration(delay) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] failed to parse retryDelay: %s error=%v", delay, err) + continue + } + retryDelay = dur + } + } + + // 验证条件 + // 情况1: RESOURCE_EXHAUSTED 需要有 RATE_LIMIT_EXCEEDED reason + // 情况2: UNAVAILABLE 需要有 MODEL_CAPACITY_EXHAUSTED reason + if isResourceExhausted && !hasRateLimitExceeded { + return nil + } + if isUnavailable && !hasModelCapacityExhausted { + return nil + } + + // 必须有模型名才返回有效结果 + if modelName == "" { + return nil + } + + // 如果上游未提供 retryDelay，使用默认限流时间 + if retryDelay <= 0 { + retryDelay = antigravityDefaultRateLimitDuration + } + + return &antigravitySmartRetryInfo{ + RetryDelay: retryDelay, + ModelName: modelName, + IsModelCapacityExhausted: hasModelCapacityExhausted, + } +} + +// shouldTriggerAntigravitySmartRetry 判断是否应该触发智能重试 +// 返回： +// - shouldRetry: 是否应该智能重试（retryDelay < antigravityRateLimitThreshold，或 MODEL_CAPACITY_EXHAUSTED） +// - shouldRateLimitModel: 是否应该限流模型并切换账号（仅 RATE_LIMIT_EXCEEDED 且 retryDelay >= 阈值） +// - waitDuration: 等待时间 +// - modelName: 限流的模型名称 +// - isModelCapacityExhausted: 是否为模型容量不足（MODEL_CAPACITY_EXHAUSTED） +func shouldTriggerAntigravitySmartRetry(account *Account, respBody []byte) (shouldRetry bool, shouldRateLimitModel bool, waitDuration time.Duration, modelName string, isModelCapacityExhausted bool) { + if account.Platform != PlatformAntigravity { + return false, false, 0, "", false + } + + info := parseAntigravitySmartRetryInfo(respBody) + if info == nil { + return false, false, 0, "", false + } + + // MODEL_CAPACITY_EXHAUSTED（模型容量不足）：所有账号共享同一模型容量池 + // 切换账号无意义，使用固定 1s 间隔重试 + if info.IsModelCapacityExhausted { + return true, false, antigravityModelCapacityRetryWait, info.ModelName, true + } + + // RATE_LIMIT_EXCEEDED（账号级限流）： + // retryDelay >= 阈值：直接限流模型，不重试 + // 注意：如果上游未提供 retryDelay，parseAntigravitySmartRetryInfo 已设置为默认 30s + if info.RetryDelay >= antigravityRateLimitThreshold { + return false, true, info.RetryDelay, info.ModelName, false + } + + // retryDelay < 阈值：智能重试 + waitDuration = info.RetryDelay + if waitDuration < antigravitySmartRetryMinWait { + waitDuration = antigravitySmartRetryMinWait + } + + return true, false, waitDuration, info.ModelName, false +} + +// handleModelRateLimitParams 模型级限流处理参数 +type handleModelRateLimitParams struct { + ctx context.Context + prefix string + account *Account + statusCode int + body []byte + cache GatewayCache + groupID int64 + sessionHash string + isStickySession bool +} + +// handleModelRateLimitResult 模型级限流处理结果 +type handleModelRateLimitResult struct { + Handled bool // 是否已处理 + ShouldRetry bool // 是否等待后重试 + WaitDuration time.Duration // 等待时间 + SwitchError *AntigravityAccountSwitchError // 账号切换错误 +} + +// handleModelRateLimit 处理模型级限流（在原有逻辑之前调用） +// 仅处理 429/503，解析模型名和 retryDelay +// - MODEL_CAPACITY_EXHAUSTED: 返回 Handled=true（实际重试由 handleSmartRetry 处理） +// - RATE_LIMIT_EXCEEDED + retryDelay < 阈值: 返回 ShouldRetry=true，由调用方等待后重试 +// - RATE_LIMIT_EXCEEDED + retryDelay >= 阈值: 设置模型限流 + 清除粘性会话 + 返回 SwitchError +func (s *AntigravityGatewayService) handleModelRateLimit(p *handleModelRateLimitParams) *handleModelRateLimitResult { + if p.statusCode != 429 && p.statusCode != 503 { + return &handleModelRateLimitResult{Handled: false} + } + + info := parseAntigravitySmartRetryInfo(p.body) + if info == nil || info.ModelName == "" { + return &handleModelRateLimitResult{Handled: false} + } + + // MODEL_CAPACITY_EXHAUSTED：模型容量不足，所有账号共享同一容量池 + // 切换账号无意义，不设置模型限流（实际重试由 handleSmartRetry 处理） + if info.IsModelCapacityExhausted { + log.Printf("%s status=%d model_capacity_exhausted model=%s (not switching account, retry handled by smart retry)", + p.prefix, p.statusCode, info.ModelName) + return &handleModelRateLimitResult{ + Handled: true, + } + } + + // RATE_LIMIT_EXCEEDED: < antigravityRateLimitThreshold: 等待后重试 + if info.RetryDelay < antigravityRateLimitThreshold { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_wait model=%s wait=%v", + p.prefix, p.statusCode, info.ModelName, info.RetryDelay) + return &handleModelRateLimitResult{ + Handled: true, + ShouldRetry: true, + WaitDuration: info.RetryDelay, + } + } + + // RATE_LIMIT_EXCEEDED: >= antigravityRateLimitThreshold: 设置限流 + 清除粘性会话 + 切换账号 + s.setModelRateLimitAndClearSession(p, info) + + return &handleModelRateLimitResult{ + Handled: true, + SwitchError: &AntigravityAccountSwitchError{ + OriginalAccountID: p.account.ID, + RateLimitedModel: info.ModelName, + IsStickySession: p.isStickySession, + }, + } +} + +// setModelRateLimitAndClearSession 设置模型限流并清除粘性会话 +func (s *AntigravityGatewayService) setModelRateLimitAndClearSession(p *handleModelRateLimitParams, info *antigravitySmartRetryInfo) { + resetAt := time.Now().Add(info.RetryDelay) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", + p.prefix, p.statusCode, info.ModelName, p.account.ID, info.RetryDelay) + + s.setAntigravityModelRateLimits(p.ctx, s.accountRepo, p.account, info.ModelName, p.prefix, p.statusCode, resetAt, false) + + // 清除粘性会话绑定 + if p.cache != nil && p.sessionHash != "" { + _ = p.cache.DeleteSessionAccountID(p.ctx, p.groupID, p.sessionHash) + } +} + +// updateAccountModelRateLimitInCache 立即更新 Redis 中账号的模型限流状态 +func (s *AntigravityGatewayService) updateAccountModelRateLimitInCache(ctx context.Context, account *Account, modelKey string, resetAt time.Time) { + if s.schedulerSnapshot == nil || account == nil || modelKey == "" { + return + } + + // 更新账号对象的 Extra 字段 + if account.Extra == nil { + account.Extra = make(map[string]any) + } + + limits, _ := account.Extra["model_rate_limits"].(map[string]any) + if limits == nil { + limits = make(map[string]any) + account.Extra["model_rate_limits"] = limits + } + + limits[modelKey] = map[string]any{ + "rate_limited_at": time.Now().UTC().Format(time.RFC3339), + "rate_limit_reset_at": resetAt.UTC().Format(time.RFC3339), + } + + // 更新 Redis 快照 + if err := s.schedulerSnapshot.UpdateAccountInCache(ctx, account); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] cache_update_failed account=%d model=%s err=%v", account.ID, modelKey, err) + } +} + +func (s *AntigravityGatewayService) handleUpstreamError( + ctx context.Context, prefix string, account *Account, + statusCode int, headers http.Header, body []byte, + requestedModel string, + groupID int64, sessionHash string, isStickySession bool, +) *handleModelRateLimitResult { + // 遵守自定义错误码策略：未命中则跳过所有限流处理 + if !account.ShouldHandleErrorCode(statusCode) { + return nil + } + // 模型级限流处理（优先） + result := s.handleModelRateLimit(&handleModelRateLimitParams{ + ctx: ctx, + prefix: prefix, + account: account, + statusCode: statusCode, + body: body, + cache: s.cache, + groupID: groupID, + sessionHash: sessionHash, + isStickySession: isStickySession, + }) + if result.Handled { + return result + } + + // 503 仅处理模型限流（MODEL_CAPACITY_EXHAUSTED），非模型限流不做额外处理 + // 避免将普通的 503 错误误判为账号问题 + if statusCode == 503 { + return nil + } + + // 429：尝试解析模型级限流，解析失败时兜底为账号级限流 + if statusCode == 429 { + if logBody, maxBytes := s.getLogConfig(); logBody { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity-Debug] 429 response body: %s", truncateString(string(body), maxBytes)) + } + + resetAt := ParseGeminiRateLimitResetTime(body) + defaultDur := s.getDefaultRateLimitDuration() + + // 尝试解析模型 key 并设置模型级限流 + // + // 注意：requestedModel 可能是"映射前"的请求模型名（例如 claude-opus-4-6）， + // 调度与限流判定使用的是 Antigravity 最终模型名（包含映射与 thinking 后缀）。 + // 因此这里必须写入最终模型 key，确保后续调度能正确避开已限流模型。 + modelKey := resolveFinalAntigravityModelKey(ctx, account, requestedModel) + if strings.TrimSpace(modelKey) == "" { + // 极少数情况下无法映射（理论上不应发生：能转发成功说明映射已通过）， + // 保持旧行为作为兜底，避免完全丢失模型级限流记录。 + modelKey = resolveAntigravityModelKey(requestedModel) + } + if modelKey != "" { + ra := s.resolveResetTime(resetAt, defaultDur) + if !s.setAntigravityModelRateLimits(ctx, s.accountRepo, account, modelKey, prefix, statusCode, ra, false) { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s", prefix, modelKey) + } else { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limited model=%s account=%d reset_at=%v reset_in=%v", + prefix, modelKey, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second)) + } + return nil + } + + // 无法解析模型 key，兜底为账号级限流 + ra := s.resolveResetTime(resetAt, defaultDur) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limited account=%d reset_at=%v reset_in=%v (fallback)", + prefix, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second)) + if err := s.accountRepo.SetRateLimited(ctx, account.ID, ra); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limit_set_failed account=%d error=%v", prefix, account.ID, err) + } + return nil + } + // 其他错误码继续使用 rateLimitService + if s.rateLimitService == nil { + return nil + } + shouldDisable := s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, headers, body) + if shouldDisable { + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d marked_error", prefix, statusCode) + } + return nil +} + +// getDefaultRateLimitDuration 获取默认限流时间 +func (s *AntigravityGatewayService) getDefaultRateLimitDuration() time.Duration { + defaultDur := antigravityDefaultRateLimitDuration + if s.settingService != nil && s.settingService.cfg != nil && s.settingService.cfg.Gateway.AntigravityFallbackCooldownMinutes > 0 { + defaultDur = time.Duration(s.settingService.cfg.Gateway.AntigravityFallbackCooldownMinutes) * time.Minute + } + if override, ok := antigravityFallbackCooldownSeconds(); ok { + defaultDur = override + } + return defaultDur +} + +// resolveResetTime 根据解析的重置时间或默认时长计算重置时间点 +func (s *AntigravityGatewayService) resolveResetTime(resetAt *int64, defaultDur time.Duration) time.Time { + if resetAt != nil { + return time.Unix(*resetAt, 0) + } + return time.Now().Add(defaultDur) +} diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index 585170438e..9e0cc804d5 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -1,28 +1,20 @@ package service import ( - "bufio" "bytes" "context" "encoding/json" "errors" "fmt" "io" - "log" "log/slog" - mathrand "math/rand" - "net" "net/http" - "os" - "strconv" "strings" "sync" - "sync/atomic" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - "github.com/gin-gonic/gin" "github.com/google/uuid" "github.com/tidwall/gjson" ) @@ -127,745 +119,6 @@ func (e *PromptTooLongError) Error() string { return fmt.Sprintf("prompt too long: status=%d", e.StatusCode) } -// antigravityRetryLoopParams 重试循环的参数 -type antigravityRetryLoopParams struct { - ctx context.Context - prefix string - account *Account - proxyURL string - accessToken string - action string - body []byte - c *gin.Context - httpUpstream HTTPUpstream - settingService *SettingService - accountRepo AccountRepository // 用于智能重试的模型级别限流 - handleError func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult - requestedModel string // 用于限流检查的原始请求模型 - isStickySession bool // 是否为粘性会话（用于账号切换时的缓存计费判断） - groupID int64 // 用于模型级限流时清除粘性会话 - sessionHash string // 用于模型级限流时清除粘性会话 -} - -// antigravityRetryLoopResult 重试循环的结果 -type antigravityRetryLoopResult struct { - resp *http.Response -} - -// resolveAntigravityForwardBaseURL 解析转发用 base URL。 -// -// 默认使用生产端点 cloudcode-pa.googleapis.com（antigravity.BaseURLs 的首个地址， -// 与账号 OAuth 登录/测试连接所用的 antigravity.BaseURL 一致）。 -// -// 历史上这里改用 ForwardBaseURLs()（把 daily/sandbox 排到首位）并默认取首个地址， -// 导致网关把带生产 OAuth token 的请求发到 daily-cloudcode-pa.sandbox.googleapis.com， -// 上游拒绝 → 账号被 401「Invalid bearer token」/502 打入临时不可调度且无法恢复 -// （见 #3611 / #2962）。后台「测试连接」用的是生产端点，所以「测试成功但网关 401」。 -// -// daily/sandbox 端点仅供内部联调，需显式设置 -// GATEWAY_ANTIGRAVITY_FORWARD_BASE_URL=daily（或 sandbox）才启用。 -func resolveAntigravityForwardBaseURL() string { - baseURLs := antigravity.BaseURLs - if len(baseURLs) == 0 { - return "" - } - mode := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityForwardBaseURLEnv))) - if (mode == "daily" || mode == "sandbox") && len(baseURLs) > 1 { - return baseURLs[1] - } - return baseURLs[0] -} - -// smartRetryAction 智能重试的处理结果 -type smartRetryAction int - -const ( - smartRetryActionContinue smartRetryAction = iota // 继续默认重试逻辑 - smartRetryActionBreakWithResp // 结束循环并返回 resp - smartRetryActionContinueURL // 继续 URL fallback 循环 -) - -// smartRetryResult 智能重试的结果 -type smartRetryResult struct { - action smartRetryAction - resp *http.Response - err error - switchError *AntigravityAccountSwitchError // 模型限流时返回账号切换信号 -} - -// handleSmartRetry 处理 OAuth 账号的智能重试逻辑 -// 将 429/503 限流处理逻辑抽取为独立函数，减少 antigravityRetryLoop 的复杂度 -func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParams, resp *http.Response, respBody []byte, baseURL string, urlIdx int, availableURLs []string) *smartRetryResult { - // "Resource has been exhausted" 是 URL 级别限流，切换 URL（仅 429） - if resp.StatusCode == http.StatusTooManyRequests && isURLLevelRateLimit(respBody) && urlIdx < len(availableURLs)-1 { - logger.LegacyPrintf("service.antigravity_gateway", "%s URL fallback (429): %s -> %s", p.prefix, baseURL, availableURLs[urlIdx+1]) - return &smartRetryResult{action: smartRetryActionContinueURL} - } - - category := antigravity429Unknown - if resp.StatusCode == http.StatusTooManyRequests { - category = classifyAntigravity429(respBody) - } - - // 判断是否触发智能重试 - shouldSmartRetry, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, respBody) - - // AI Credits 超量请求： - // 仅在上游明确返回免费配额耗尽时才允许切换到 credits。 - if resp.StatusCode == http.StatusTooManyRequests && - category == antigravity429QuotaExhausted && - p.account.IsOveragesEnabled() && - !p.account.isCreditsExhausted() { - result := s.attemptCreditsOveragesRetry(p, baseURL, modelName, waitDuration, resp.StatusCode, respBody) - if result.handled && result.resp != nil { - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: result.resp, - } - } - } - - // 情况1: retryDelay >= 阈值，限流模型并切换账号 - if shouldRateLimitModel { - // 单账号 503 退避重试模式：不设限流、不切换账号，改为原地等待+重试 - // 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的，等几秒就能恢复。 - // 多账号场景下切换账号是最优选择，但单账号场景下设限流毫无意义（只会导致双重等待）。 - if resp.StatusCode == http.StatusServiceUnavailable && isSingleAccountRetry(p.ctx) { - return s.handleSingleAccountRetryInPlace(p, resp, respBody, baseURL, waitDuration, modelName) - } - - rateLimitDuration := waitDuration - if rateLimitDuration <= 0 { - rateLimitDuration = antigravityDefaultRateLimitDuration - } - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d oauth_long_delay model=%s account=%d upstream_retry_delay=%v body=%s (model rate limit, switch account)", - p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration, truncateForLog(respBody, 200)) - - resetAt := time.Now().Add(rateLimitDuration) - if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, false) { - p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited account=%d (no model mapping)", p.prefix, resp.StatusCode, p.account.ID) - } - s.clearStickySession(p.ctx, p.groupID, p.sessionHash) - - // 返回账号切换信号，让上层切换账号重试 - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - switchError: &AntigravityAccountSwitchError{ - OriginalAccountID: p.account.ID, - RateLimitedModel: modelName, - IsStickySession: p.isStickySession, - }, - } - } - - // 情况2: retryDelay < 阈值（或 MODEL_CAPACITY_EXHAUSTED），智能重试 - if shouldSmartRetry { - var lastRetryResp *http.Response - var lastRetryBody []byte - - // MODEL_CAPACITY_EXHAUSTED 使用独立的重试参数（60 次，固定 1s 间隔） - maxAttempts := antigravitySmartRetryMaxAttempts - if isModelCapacityExhausted { - maxAttempts = antigravityModelCapacityRetryMaxAttempts - waitDuration = antigravityModelCapacityRetryWait - - // 全局去重：如果其他 goroutine 已在重试同一模型且尚在 cooldown 中，直接返回 503 - if modelName != "" { - modelCapacityExhaustedMu.RLock() - cooldownUntil, exists := modelCapacityExhaustedUntil[modelName] - modelCapacityExhaustedMu.RUnlock() - if exists && time.Now().Before(cooldownUntil) { - log.Printf("%s status=%d model_capacity_exhausted_dedup model=%s account=%d cooldown_until=%v (skip retry)", - p.prefix, resp.StatusCode, modelName, p.account.ID, cooldownUntil.Format("15:04:05")) - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - }, - } - } - } - } - - for attempt := 1; attempt <= maxAttempts; attempt++ { - log.Printf("%s status=%d oauth_smart_retry attempt=%d/%d delay=%v model=%s account=%d", - p.prefix, resp.StatusCode, attempt, maxAttempts, waitDuration, modelName, p.account.ID) - - timer := time.NewTimer(waitDuration) - select { - case <-p.ctx.Done(): - timer.Stop() - log.Printf("%s status=context_canceled_during_smart_retry", p.prefix) - return &smartRetryResult{action: smartRetryActionBreakWithResp, err: p.ctx.Err()} - case <-timer.C: - } - - // 智能重试：创建新请求 - retryReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=smart_retry_request_build_failed error=%v", p.prefix, err) - p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - }, - } - } - - retryResp, retryErr := p.httpUpstream.Do(retryReq, p.proxyURL, p.account.ID, p.account.Concurrency) - if retryErr == nil && retryResp != nil && retryResp.StatusCode != http.StatusTooManyRequests && retryResp.StatusCode != http.StatusServiceUnavailable { - log.Printf("%s status=%d smart_retry_success attempt=%d/%d", p.prefix, retryResp.StatusCode, attempt, maxAttempts) - // 重试成功，清除 MODEL_CAPACITY_EXHAUSTED cooldown - if isModelCapacityExhausted && modelName != "" { - modelCapacityExhaustedMu.Lock() - delete(modelCapacityExhaustedUntil, modelName) - modelCapacityExhaustedMu.Unlock() - } - return &smartRetryResult{action: smartRetryActionBreakWithResp, resp: retryResp} - } - - // 网络错误时，继续重试 - if retryErr != nil || retryResp == nil { - log.Printf("%s status=smart_retry_network_error attempt=%d/%d error=%v", p.prefix, attempt, maxAttempts, retryErr) - continue - } - - // 重试失败，关闭之前的响应 - if lastRetryResp != nil { - _ = lastRetryResp.Body.Close() - } - lastRetryResp = retryResp - if retryResp != nil { - lastRetryBody, _ = io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) - _ = retryResp.Body.Close() - } - - // 解析新的重试信息，用于下次重试的等待时间（MODEL_CAPACITY_EXHAUSTED 使用固定循环，跳过） - if !isModelCapacityExhausted && attempt < maxAttempts && lastRetryBody != nil { - newShouldRetry, _, newWaitDuration, _, _ := shouldTriggerAntigravitySmartRetry(p.account, lastRetryBody) - if newShouldRetry && newWaitDuration > 0 { - waitDuration = newWaitDuration - } - } - } - - // 所有重试都失败 - rateLimitDuration := waitDuration - if rateLimitDuration <= 0 { - rateLimitDuration = antigravityDefaultRateLimitDuration - } - retryBody := lastRetryBody - if retryBody == nil { - retryBody = respBody - } - - // MODEL_CAPACITY_EXHAUSTED：模型容量不足，切换账号无意义 - // 直接返回上游错误响应，不设置模型限流，不切换账号 - if isModelCapacityExhausted { - // 设置 cooldown，让后续请求快速失败，避免重复重试 - if modelName != "" { - modelCapacityExhaustedMu.Lock() - modelCapacityExhaustedUntil[modelName] = time.Now().Add(antigravityModelCapacityCooldown) - modelCapacityExhaustedMu.Unlock() - } - log.Printf("%s status=%d smart_retry_exhausted_model_capacity attempts=%d model=%s account=%d body=%s (model capacity exhausted, not switching account)", - p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, truncateForLog(retryBody, 200)) - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - }, - } - } - - // 单账号 503 退避重试模式：智能重试耗尽后不设限流、不切换账号， - // 直接返回 503 让 Handler 层的单账号退避循环做最终处理。 - if resp.StatusCode == http.StatusServiceUnavailable && isSingleAccountRetry(p.ctx) { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d smart_retry_exhausted_single_account attempts=%d model=%s account=%d body=%s (return 503 directly)", - p.prefix, resp.StatusCode, antigravitySmartRetryMaxAttempts, modelName, p.account.ID, truncateForLog(retryBody, 200)) - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - }, - } - } - - log.Printf("%s status=%d smart_retry_exhausted attempts=%d model=%s account=%d upstream_retry_delay=%v body=%s (switch account)", - p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, rateLimitDuration, truncateForLog(retryBody, 200)) - - resetAt := time.Now().Add(rateLimitDuration) - s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, true) - - // 清除粘性会话绑定，避免下次请求仍命中限流账号 - s.clearStickySession(p.ctx, p.groupID, p.sessionHash) - - // 返回账号切换信号，让上层切换账号重试 - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - switchError: &AntigravityAccountSwitchError{ - OriginalAccountID: p.account.ID, - RateLimitedModel: modelName, - IsStickySession: p.isStickySession, - }, - } - } - - // 未触发智能重试，继续默认重试逻辑 - return &smartRetryResult{action: smartRetryActionContinue} -} - -// handleSingleAccountRetryInPlace 单账号 503 退避重试的原地重试逻辑。 -// -// 在多账号场景下，收到 503 + 长 retryDelay（≥ 7s）时会设置模型限流 + 切换账号； -// 但在单账号场景下，设限流毫无意义（因为切换回来的还是同一个账号，还要等限流过期）。 -// 此方法改为在 Service 层原地等待 + 重试，避免双重等待问题： -// -// 旧流程：Service 设限流 → Handler 退避等待 → Service 等限流过期 → 再请求（总耗时 = 退避 + 限流） -// 新流程：Service 直接等 retryDelay → 重试 → 成功/再等 → 重试...（总耗时 ≈ 实际 retryDelay × 重试次数） -// -// 约束： -// - 单次等待不超过 antigravitySingleAccountSmartRetryMaxWait -// - 总累计等待不超过 antigravitySingleAccountSmartRetryTotalMaxWait -// - 最多重试 antigravitySingleAccountSmartRetryMaxAttempts 次 -func (s *AntigravityGatewayService) handleSingleAccountRetryInPlace( - p antigravityRetryLoopParams, - resp *http.Response, - respBody []byte, - baseURL string, - waitDuration time.Duration, - modelName string, -) *smartRetryResult { - // 限制单次等待时间 - if waitDuration > antigravitySingleAccountSmartRetryMaxWait { - waitDuration = antigravitySingleAccountSmartRetryMaxWait - } - if waitDuration < antigravitySmartRetryMinWait { - waitDuration = antigravitySmartRetryMinWait - } - - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_in_place model=%s account=%d upstream_retry_delay=%v (retrying in-place instead of rate-limiting)", - p.prefix, resp.StatusCode, modelName, p.account.ID, waitDuration) - - var lastRetryResp *http.Response - var lastRetryBody []byte - totalWaited := time.Duration(0) - - for attempt := 1; attempt <= antigravitySingleAccountSmartRetryMaxAttempts; attempt++ { - // 检查累计等待是否超限 - if totalWaited+waitDuration > antigravitySingleAccountSmartRetryTotalMaxWait { - remaining := antigravitySingleAccountSmartRetryTotalMaxWait - totalWaited - if remaining <= 0 { - logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: total_wait_exceeded total=%v max=%v, giving up", - p.prefix, totalWaited, antigravitySingleAccountSmartRetryTotalMaxWait) - break - } - waitDuration = remaining - } - - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry attempt=%d/%d delay=%v total_waited=%v model=%s account=%d", - p.prefix, resp.StatusCode, attempt, antigravitySingleAccountSmartRetryMaxAttempts, waitDuration, totalWaited, modelName, p.account.ID) - - timer := time.NewTimer(waitDuration) - select { - case <-p.ctx.Done(): - timer.Stop() - logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_single_account_retry", p.prefix) - return &smartRetryResult{action: smartRetryActionBreakWithResp, err: p.ctx.Err()} - case <-timer.C: - } - totalWaited += waitDuration - - // 创建新请求 - retryReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: request_build_failed error=%v", p.prefix, err) - break - } - - retryResp, retryErr := p.httpUpstream.Do(retryReq, p.proxyURL, p.account.ID, p.account.Concurrency) - if retryErr == nil && retryResp != nil && retryResp.StatusCode != http.StatusTooManyRequests && retryResp.StatusCode != http.StatusServiceUnavailable { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_success attempt=%d/%d total_waited=%v", - p.prefix, retryResp.StatusCode, attempt, antigravitySingleAccountSmartRetryMaxAttempts, totalWaited) - // 关闭之前的响应 - if lastRetryResp != nil { - _ = lastRetryResp.Body.Close() - } - return &smartRetryResult{action: smartRetryActionBreakWithResp, resp: retryResp} - } - - // 网络错误时继续重试 - if retryErr != nil || retryResp == nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: network_error attempt=%d/%d error=%v", - p.prefix, attempt, antigravitySingleAccountSmartRetryMaxAttempts, retryErr) - continue - } - - // 关闭之前的响应 - if lastRetryResp != nil { - _ = lastRetryResp.Body.Close() - } - lastRetryResp = retryResp - lastRetryBody, _ = io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) - _ = retryResp.Body.Close() - - // 解析新的重试信息，更新下次等待时间 - if attempt < antigravitySingleAccountSmartRetryMaxAttempts && lastRetryBody != nil { - _, _, newWaitDuration, _, _ := shouldTriggerAntigravitySmartRetry(p.account, lastRetryBody) - if newWaitDuration > 0 { - waitDuration = newWaitDuration - if waitDuration > antigravitySingleAccountSmartRetryMaxWait { - waitDuration = antigravitySingleAccountSmartRetryMaxWait - } - if waitDuration < antigravitySmartRetryMinWait { - waitDuration = antigravitySmartRetryMinWait - } - } - } - } - - // 所有重试都失败，不设限流，直接返回 503 - // Handler 层的单账号退避循环会做最终处理 - retryBody := lastRetryBody - if retryBody == nil { - retryBody = respBody - } - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_exhausted attempts=%d total_waited=%v model=%s account=%d body=%s (return 503 directly)", - p.prefix, resp.StatusCode, antigravitySingleAccountSmartRetryMaxAttempts, totalWaited, modelName, p.account.ID, truncateForLog(retryBody, 200)) - - return &smartRetryResult{ - action: smartRetryActionBreakWithResp, - resp: &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - }, - } -} - -// antigravityRetryLoop 执行带 URL fallback 的重试循环 -func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopParams) (*antigravityRetryLoopResult, error) { - // 预检查：模型限流 + overages 启用 + 积分未耗尽 → 直接注入 AI Credits - overagesInjected := false - if p.requestedModel != "" && p.account.Platform == PlatformAntigravity && - p.account.IsOveragesEnabled() && !p.account.isCreditsExhausted() && - p.account.isModelRateLimitedWithContext(p.ctx, p.requestedModel) { - if creditsBody := injectEnabledCreditTypes(p.body); creditsBody != nil { - p.body = creditsBody - overagesInjected = true - logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: model_rate_limited_credits_inject model=%s account=%d (injecting enabledCreditTypes)", - p.prefix, p.requestedModel, p.account.ID) - } - } - - // 预检查：如果账号已限流，直接返回切换信号 - if p.requestedModel != "" { - if remaining := p.account.GetRateLimitRemainingTimeWithContext(p.ctx, p.requestedModel); remaining > 0 { - // 已注入积分的请求不再受普通模型限流预检查阻断。 - if overagesInjected { - logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: credits_injected_ignore_rate_limit remaining=%v model=%s account=%d", - p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) - } else if isSingleAccountRetry(p.ctx) { - // 单账号 503 退避重试模式：跳过限流预检查，直接发请求。 - // 首次请求设的限流是为了多账号调度器跳过该账号，在单账号模式下无意义。 - // 如果上游确实还不可用，handleSmartRetry → handleSingleAccountRetryInPlace - // 会在 Service 层原地等待+重试，不需要在预检查这里等。 - logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: single_account_retry skipping rate_limit remaining=%v model=%s account=%d (will retry in-place if 503)", - p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) - } else { - logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: rate_limit_switch remaining=%v model=%s account=%d", - p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) - return nil, &AntigravityAccountSwitchError{ - OriginalAccountID: p.account.ID, - RateLimitedModel: p.requestedModel, - IsStickySession: p.isStickySession, - } - } - } - } - - baseURL := resolveAntigravityForwardBaseURL() - if baseURL == "" { - return nil, errors.New("no antigravity forward base url configured") - } - availableURLs := []string{baseURL} - - var resp *http.Response - var usedBaseURL string - logBody := p.settingService != nil && p.settingService.cfg != nil && p.settingService.cfg.Gateway.LogUpstreamErrorBody - maxBytes := 2048 - if p.settingService != nil && p.settingService.cfg != nil && p.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > 0 { - maxBytes = p.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes - } - getUpstreamDetail := func(body []byte) string { - if !logBody { - return "" - } - return truncateString(string(body), maxBytes) - } - -urlFallbackLoop: - for urlIdx, baseURL := range availableURLs { - usedBaseURL = baseURL - allAttemptsInternal500 := true // 追踪本轮所有 attempt 是否全部命中 INTERNAL 500 - for attempt := 1; attempt <= antigravityMaxRetries; attempt++ { - select { - case <-p.ctx.Done(): - logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled error=%v", p.prefix, p.ctx.Err()) - return nil, p.ctx.Err() - default: - } - - upstreamReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body) - if err != nil { - return nil, err - } - - resp, err = p.httpUpstream.Do(upstreamReq, p.proxyURL, p.account.ID, p.account.Concurrency) - if err == nil && resp == nil { - err = errors.New("upstream returned nil response") - } - if err != nil { - safeErr := sanitizeUpstreamErrorMessage(err.Error()) - appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ - Platform: p.account.Platform, - AccountID: p.account.ID, - AccountName: p.account.Name, - UpstreamStatusCode: 0, - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Kind: "request_error", - Message: safeErr, - }) - if shouldAntigravityFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 { - logger.LegacyPrintf("service.antigravity_gateway", "%s URL fallback (connection error): %s -> %s", p.prefix, baseURL, availableURLs[urlIdx+1]) - continue urlFallbackLoop - } - if attempt < antigravityMaxRetries { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=request_failed retry=%d/%d error=%v", p.prefix, attempt, antigravityMaxRetries, err) - if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) - return nil, p.ctx.Err() - } - continue - } - logger.LegacyPrintf("service.antigravity_gateway", "%s status=request_failed retries_exhausted error=%v", p.prefix, err) - setOpsUpstreamError(p.c, 0, safeErr, "") - return nil, fmt.Errorf("upstream request failed after retries: %w", err) - } - - // 统一处理错误响应 - if resp.StatusCode >= 400 { - respBody := s.readUpstreamErrorBody(resp) - _ = resp.Body.Close() - - if overagesInjected && shouldMarkCreditsExhausted(resp, respBody, nil) { - modelKey := resolveCreditsOveragesModelKey(p.ctx, p.account, "", p.requestedModel) - s.handleCreditsRetryFailure(p.ctx, p.prefix, modelKey, p.account, &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - }, nil) - } - - // ★ 统一入口：自定义错误码 + 临时不可调度 - if handled, outStatus, policyErr := s.applyErrorPolicy(p, resp.StatusCode, resp.Header, respBody); handled { - if policyErr != nil { - return nil, policyErr - } - resp = &http.Response{ - StatusCode: outStatus, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - } - break urlFallbackLoop - } - - // 429/503 限流处理：区分 URL 级别限流、智能重试和账户配额限流 - if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusServiceUnavailable { - // 尝试智能重试处理（OAuth 账号专用） - smartResult := s.handleSmartRetry(p, resp, respBody, baseURL, urlIdx, availableURLs) - switch smartResult.action { - case smartRetryActionContinueURL: - continue urlFallbackLoop - case smartRetryActionBreakWithResp: - if smartResult.err != nil { - return nil, smartResult.err - } - // 模型限流时返回切换账号信号 - if smartResult.switchError != nil { - return nil, smartResult.switchError - } - resp = smartResult.resp - break urlFallbackLoop - } - // smartRetryActionContinue: 继续默认重试逻辑 - - // 账户/模型配额限流，重试 3 次（指数退避）- 默认逻辑（非 OAuth 账号或解析失败） - if attempt < antigravityMaxRetries { - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ - Platform: p.account.Platform, - AccountID: p.account.ID, - AccountName: p.account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Kind: "retry", - Message: upstreamMsg, - Detail: getUpstreamDetail(respBody), - }) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d retry=%d/%d body=%s", p.prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 200)) - if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) - return nil, p.ctx.Err() - } - continue - } - - // 重试用尽，标记账户限流 - p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited base_url=%s body=%s", p.prefix, resp.StatusCode, baseURL, truncateForLog(respBody, 200)) - resp = &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - } - break urlFallbackLoop - } - - // 其他可重试错误（500/502/504/529，不包括 429 和 503） - if shouldRetryAntigravityError(resp.StatusCode) { - if attempt < antigravityMaxRetries { - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{ - Platform: p.account.Platform, - AccountID: p.account.ID, - AccountName: p.account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Kind: "retry", - Message: upstreamMsg, - Detail: getUpstreamDetail(respBody), - }) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d retry=%d/%d body=%s", p.prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 500)) - if !sleepAntigravityBackoffWithContext(p.ctx, attempt) { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix) - return nil, p.ctx.Err() - } - // 追踪 INTERNAL 500：非匹配的 attempt 清除标记 - if !isAntigravityInternalServerError(resp.StatusCode, respBody) { - allAttemptsInternal500 = false - } - continue - } - } - - // INTERNAL 500 渐进惩罚：3 次重试全部命中特定 500 时递增计数器并惩罚 - if allAttemptsInternal500 && isAntigravityInternalServerError(resp.StatusCode, respBody) { - s.handleInternal500RetryExhausted(p.ctx, p.prefix, p.account) - } - - // 其他 4xx 错误或重试用尽，直接返回 - resp = &http.Response{ - StatusCode: resp.StatusCode, - Header: resp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(respBody)), - } - break urlFallbackLoop - } - - // 成功响应（< 400） - break urlFallbackLoop - } - } - - if resp != nil && resp.StatusCode < 400 && usedBaseURL != "" { - antigravity.DefaultURLAvailability.MarkSuccess(usedBaseURL) - } - - // 成功响应时清零 INTERNAL 500 连续失败计数器（覆盖所有成功路径，含 smart retry） - if resp != nil && resp.StatusCode < 400 { - s.resetInternal500Counter(p.ctx, p.prefix, p.account.ID) - } - - return &antigravityRetryLoopResult{resp: resp}, nil -} - -// shouldRetryAntigravityError 判断是否应该重试 -func shouldRetryAntigravityError(statusCode int) bool { - switch statusCode { - case 429, 500, 502, 503, 504, 529: - return true - default: - return false - } -} - -// isURLLevelRateLimit 判断是否为 URL 级别的限流（应切换 URL 重试） -// "Resource has been exhausted" 是 URL/节点级别限流，切换 URL 可能成功 -// "exhausted your capacity on this model" 是账户/模型配额限流，切换 URL 无效 -func isURLLevelRateLimit(body []byte) bool { - // 快速检查：包含 "Resource has been exhausted" 且不包含 "capacity on this model" - bodyStr := string(body) - return strings.Contains(bodyStr, "Resource has been exhausted") && - !strings.Contains(bodyStr, "capacity on this model") -} - -// isAntigravityConnectionError 判断是否为连接错误（网络超时、DNS 失败、连接拒绝） -func isAntigravityConnectionError(err error) bool { - if err == nil { - return false - } - - // 检查超时错误 - var netErr net.Error - if errors.As(err, &netErr) && netErr.Timeout() { - return true - } - - // 检查连接错误（DNS 失败、连接拒绝） - var opErr *net.OpError - return errors.As(err, &opErr) -} - -// shouldAntigravityFallbackToNextURL 判断是否应切换到下一个 URL -// 仅连接错误和 HTTP 429 触发 URL 降级 -func shouldAntigravityFallbackToNextURL(err error, statusCode int) bool { - if isAntigravityConnectionError(err) { - return true - } - return statusCode == http.StatusTooManyRequests -} - -// getSessionID 从 gin.Context 获取 session_id（用于日志追踪） -func getSessionID(c *gin.Context) string { - if c == nil { - return "" - } - return c.GetHeader("session_id") -} - -// logPrefix 生成统一的日志前缀 -func logPrefix(sessionID, accountName string) string { - if sessionID != "" { - return fmt.Sprintf("[antigravity-Forward] session=%s account=%s", sessionID, accountName) - } - return fmt.Sprintf("[antigravity-Forward] account=%s", accountName) -} - // AntigravityGatewayService 处理 Antigravity 平台的 API 转发 type AntigravityGatewayService struct { accountRepo AccountRepository @@ -1384,3281 +637,3 @@ func (s *AntigravityGatewayService) unwrapV1InternalResponse(body []byte) ([]byt } return body, nil } - -// Forward 转发 Claude 协议请求（Claude → Gemini 转换） -// -// 限流处理流程: -// -// 请求 → antigravityRetryLoop → 预检查(remaining>0? → 切换账号) → 发送上游 -// ├─ 成功 → 正常返回 -// └─ 429/503 → handleSmartRetry -// ├─ retryDelay >= 7s → 设置模型限流 + 清除粘性绑定 → 切换账号 -// └─ retryDelay < 7s → 等待后重试 1 次 -// ├─ 成功 → 正常返回 -// └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号 -func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte, isStickySession bool) (*ForwardResult, error) { - // 上游透传账号直接转发，不走 OAuth token 刷新 - if account.Type == AccountTypeUpstream { - return s.ForwardUpstream(ctx, c, account, body) - } - - startTime := time.Now() - - sessionID := getSessionID(c) - prefix := logPrefix(sessionID, account.Name) - - // 解析 Claude 请求 - var claudeReq antigravity.ClaudeRequest - if err := json.Unmarshal(body, &claudeReq); err != nil { - return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request body") - } - if strings.TrimSpace(claudeReq.Model) == "" { - return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Missing model") - } - - originalModel := claudeReq.Model - mappedModel := s.getMappedModel(account, claudeReq.Model) - if mappedModel == "" { - MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) - return nil, s.writeClaudeError(c, http.StatusForbidden, "permission_error", fmt.Sprintf("model %s not in whitelist", claudeReq.Model)) - } - // 应用 thinking 模式自动后缀：如果 thinking 开启且目标是 claude-sonnet-4-5，自动改为 thinking 版本 - thinkingEnabled := claudeReq.Thinking != nil && (claudeReq.Thinking.Type == "enabled" || claudeReq.Thinking.Type == "adaptive") - mappedModel = applyThinkingModelSuffix(mappedModel, thinkingEnabled) - billingModel := mappedModel - - // 获取 access_token - if s.tokenProvider == nil { - return nil, s.writeClaudeError(c, http.StatusBadGateway, "api_error", "Antigravity token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: []byte(`{"error":{"type":"authentication_error","message":"Failed to get upstream access token"},"type":"error"}`), - } - } - - projectID, err := resolveAntigravityProjectID(account) - if err != nil { - _ = s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", err.Error()) - return nil, err - } - - // 代理 URL - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - // 获取转换选项 - // Antigravity 上游要求必须包含身份提示词，否则会返回 429 - transformOpts := s.getClaudeTransformOptions(ctx) - transformOpts.EnableIdentityPatch = true // 强制启用，Antigravity 上游必需 - - // 转换 Claude 请求为 Gemini 格式 - geminiBody, err := antigravity.TransformClaudeToGeminiWithOptions(&claudeReq, projectID, mappedModel, transformOpts) - if err != nil { - return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", "Invalid request") - } - - // Antigravity 上游只支持流式请求，统一使用 streamGenerateContent - // 如果客户端请求非流式，在响应处理阶段会收集完整流式响应后转换返回 - action := "streamGenerateContent" - - // 执行带重试的请求 - result, err := s.antigravityRetryLoop(antigravityRetryLoopParams{ - ctx: ctx, - prefix: prefix, - account: account, - proxyURL: proxyURL, - accessToken: accessToken, - action: action, - body: geminiBody, - c: c, - httpUpstream: s.httpUpstream, - settingService: s.settingService, - accountRepo: s.accountRepo, - handleError: s.handleUpstreamError, - requestedModel: originalModel, - isStickySession: isStickySession, // Forward 由上层判断粘性会话 - groupID: 0, // Forward 方法没有 groupID，由上层处理粘性会话清除 - sessionHash: "", // Forward 方法没有 sessionHash，由上层处理粘性会话清除 - }) - if err != nil { - // 检查是否是账号切换信号，转换为 UpstreamFailoverError 让 Handler 切换账号 - if switchErr, ok := IsAntigravityAccountSwitchError(err); ok { - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusServiceUnavailable, - ForceCacheBilling: switchErr.IsStickySession, - } - } - // 区分客户端取消和真正的上游失败，返回更准确的错误消息 - if c.Request.Context().Err() != nil { - return nil, s.writeClaudeError(c, http.StatusBadGateway, "client_disconnected", "Client disconnected before upstream response") - } - return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed after retries") - } - resp := result.resp - defer func() { _ = resp.Body.Close() }() - - if resp.StatusCode >= 400 { - respBody := s.readUpstreamErrorBody(resp) - - // 优先检测 thinking block 的 signature 相关错误（400）并重试一次： - // Antigravity /v1internal 链路在部分场景会对 thought/thinking signature 做严格校验， - // 当历史消息携带的 signature 不合法时会直接 400；去除 thinking 后可继续完成请求。 - if resp.StatusCode == http.StatusBadRequest && isSignatureRelatedError(respBody) && s.settingService.IsSignatureRectifierEnabled(ctx) { - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - logBody, maxBytes := s.getLogConfig() - upstreamDetail := s.getUpstreamErrorDetail(respBody) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "signature_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - // Conservative two-stage fallback: - // 1) Disable top-level thinking + thinking->text - // 2) Only if still signature-related 400: also downgrade tool_use/tool_result to text. - - retryStages := []struct { - name string - strip func(*antigravity.ClaudeRequest) (bool, error) - }{ - {name: "thinking-only", strip: stripThinkingFromClaudeRequest}, - {name: "thinking+tools", strip: stripSignatureSensitiveBlocksFromClaudeRequest}, - } - - for _, stage := range retryStages { - retryClaudeReq := claudeReq - retryClaudeReq.Messages = append([]antigravity.ClaudeMessage(nil), claudeReq.Messages...) - - stripped, stripErr := stage.strip(&retryClaudeReq) - if stripErr != nil || !stripped { - continue - } - - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: detected signature-related 400, retrying once (%s)", account.ID, stage.name) - - retryGeminiBody, txErr := antigravity.TransformClaudeToGeminiWithOptions(&retryClaudeReq, projectID, mappedModel, s.getClaudeTransformOptions(ctx)) - if txErr != nil { - continue - } - retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ - ctx: ctx, - prefix: prefix, - account: account, - proxyURL: proxyURL, - accessToken: accessToken, - action: action, - body: retryGeminiBody, - c: c, - httpUpstream: s.httpUpstream, - settingService: s.settingService, - accountRepo: s.accountRepo, - handleError: s.handleUpstreamError, - requestedModel: originalModel, - isStickySession: isStickySession, - groupID: 0, // Forward 方法没有 groupID，由上层处理粘性会话清除 - sessionHash: "", // Forward 方法没有 sessionHash，由上层处理粘性会话清除 - }) - if retryErr != nil { - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: 0, - Kind: "signature_retry_request_error", - Message: sanitizeUpstreamErrorMessage(retryErr.Error()), - }) - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: signature retry request failed (%s): %v", account.ID, stage.name, retryErr) - continue - } - - retryResp := retryResult.resp - if retryResp.StatusCode < 400 { - _ = resp.Body.Close() - resp = retryResp - respBody = nil - break - } - - retryBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 8<<10)) - _ = retryResp.Body.Close() - if retryResp.StatusCode == http.StatusTooManyRequests { - retryBaseURL := "" - if retryResp.Request != nil && retryResp.Request.URL != nil { - retryBaseURL = retryResp.Request.URL.Scheme + "://" + retryResp.Request.URL.Host - } - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limited base_url=%s retry_stage=%s body=%s", prefix, retryBaseURL, stage.name, truncateForLog(retryBody, 200)) - } - kind := "signature_retry" - if strings.TrimSpace(stage.name) != "" { - kind = "signature_retry_" + strings.ReplaceAll(stage.name, "+", "_") - } - retryUpstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(retryBody)) - retryUpstreamMsg = sanitizeUpstreamErrorMessage(retryUpstreamMsg) - retryUpstreamDetail := "" - if logBody { - retryUpstreamDetail = truncateString(string(retryBody), maxBytes) - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: retryResp.StatusCode, - UpstreamRequestID: retryResp.Header.Get("x-request-id"), - Kind: kind, - Message: retryUpstreamMsg, - Detail: retryUpstreamDetail, - }) - - // If this stage fixed the signature issue, we stop; otherwise we may try the next stage. - if retryResp.StatusCode != http.StatusBadRequest || !isSignatureRelatedError(retryBody) { - respBody = retryBody - resp = &http.Response{ - StatusCode: retryResp.StatusCode, - Header: retryResp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - } - break - } - - // Still signature-related; capture context and allow next stage. - respBody = retryBody - resp = &http.Response{ - StatusCode: retryResp.StatusCode, - Header: retryResp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - } - } - } - - // Budget 整流：检测 budget_tokens 约束错误并自动修正重试 - if resp.StatusCode == http.StatusBadRequest && respBody != nil && !isSignatureRelatedError(respBody) { - errMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - if isThinkingBudgetConstraintError(errMsg) && s.settingService.IsBudgetRectifierEnabled(ctx) { - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "budget_constraint_error", - Message: errMsg, - Detail: s.getUpstreamErrorDetail(respBody), - }) - - // 修正 claudeReq 的 thinking 参数（adaptive 模式不修正） - if claudeReq.Thinking == nil || claudeReq.Thinking.Type != "adaptive" { - retryClaudeReq := claudeReq - retryClaudeReq.Messages = append([]antigravity.ClaudeMessage(nil), claudeReq.Messages...) - // 创建新的 ThinkingConfig 避免修改原始 claudeReq.Thinking 指针 - retryClaudeReq.Thinking = &antigravity.ThinkingConfig{ - Type: "enabled", - BudgetTokens: BudgetRectifyBudgetTokens, - } - if retryClaudeReq.MaxTokens < BudgetRectifyMinMaxTokens { - retryClaudeReq.MaxTokens = BudgetRectifyMaxTokens - } - - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: detected budget_tokens constraint error, retrying with rectified budget (budget_tokens=%d, max_tokens=%d)", account.ID, BudgetRectifyBudgetTokens, BudgetRectifyMaxTokens) - - retryGeminiBody, txErr := antigravity.TransformClaudeToGeminiWithOptions(&retryClaudeReq, projectID, mappedModel, transformOpts) - if txErr == nil { - retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ - ctx: ctx, - prefix: prefix, - account: account, - proxyURL: proxyURL, - accessToken: accessToken, - action: action, - body: retryGeminiBody, - c: c, - httpUpstream: s.httpUpstream, - settingService: s.settingService, - accountRepo: s.accountRepo, - handleError: s.handleUpstreamError, - requestedModel: originalModel, - isStickySession: isStickySession, - groupID: 0, - sessionHash: "", - }) - if retryErr == nil { - retryResp := retryResult.resp - if retryResp.StatusCode < 400 { - _ = resp.Body.Close() - resp = retryResp - respBody = nil - } else { - retryBody := s.readUpstreamErrorBody(retryResp) - _ = retryResp.Body.Close() - respBody = retryBody - resp = &http.Response{ - StatusCode: retryResp.StatusCode, - Header: retryResp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryBody)), - } - } - } else { - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity account %d: budget rectifier retry failed: %v", account.ID, retryErr) - } - } - } - } - } - - // 处理错误响应（重试后仍失败或不触发重试） - if resp.StatusCode >= 400 { - // 检测 prompt too long 错误，返回特殊错误类型供上层 fallback - if resp.StatusCode == http.StatusBadRequest && isPromptTooLongError(respBody) { - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - upstreamDetail := s.getUpstreamErrorDetail(respBody) - logBody, maxBytes := s.getLogConfig() - if logBody { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=400 prompt_too_long=true upstream_message=%q request_id=%s body=%s", prefix, upstreamMsg, resp.Header.Get("x-request-id"), truncateForLog(respBody, maxBytes)) - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "prompt_too_long", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - return nil, &PromptTooLongError{ - StatusCode: resp.StatusCode, - RequestID: resp.Header.Get("x-request-id"), - Body: respBody, - } - } - - s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", isStickySession) - - // 精确匹配服务端配置类 400 错误，触发同账号重试 + failover - if resp.StatusCode == http.StatusBadRequest { - msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) - if isGoogleProjectConfigError(msg) { - upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) - upstreamDetail := s.getUpstreamErrorDetail(respBody) - log.Printf("%s status=400 google_config_error failover=true upstream_message=%q account=%d", prefix, upstreamMsg, account.ID) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "failover", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody, RetryableOnSameAccount: true} - } - } - - if s.shouldFailoverUpstreamError(resp.StatusCode) { - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - upstreamDetail := s.getUpstreamErrorDetail(respBody) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "failover", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody} - } - - return nil, s.writeMappedClaudeError(c, account, resp.StatusCode, resp.Header.Get("x-request-id"), respBody) - } - } - - requestID := resp.Header.Get("x-request-id") - if requestID != "" { - c.Header("x-request-id", requestID) - } - - var usage *ClaudeUsage - var firstTokenMs *int - var clientDisconnect bool - if claudeReq.Stream { - // 客户端要求流式，直接透传转换 - streamRes, err := s.handleClaudeStreamingResponse(c, resp, startTime, originalModel) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_error error=%v", prefix, err) - return nil, err - } - usage = streamRes.usage - firstTokenMs = streamRes.firstTokenMs - clientDisconnect = streamRes.clientDisconnect - } else { - // 客户端要求非流式，收集流式响应后转换返回 - streamRes, err := s.handleClaudeStreamToNonStreaming(c, resp, startTime, originalModel) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_collect_error error=%v", prefix, err) - return nil, err - } - usage = streamRes.usage - firstTokenMs = streamRes.firstTokenMs - } - - return &ForwardResult{ - RequestID: requestID, - Usage: *usage, - Model: originalModel, - UpstreamModel: billingModel, - Stream: claudeReq.Stream, - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, - ClientDisconnect: clientDisconnect, - }, nil -} - -func isSignatureRelatedError(respBody []byte) bool { - msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) - if msg == "" { - // Fallback: best-effort scan of the raw payload. - msg = strings.ToLower(string(respBody)) - } - - // Keep this intentionally broad: different upstreams may use "signature" or "thought_signature". - if strings.Contains(msg, "thought_signature") || strings.Contains(msg, "signature") { - return true - } - - // Also detect thinking block structural errors: - // "Expected `thinking` or `redacted_thinking`, but found `text`" - if strings.Contains(msg, "expected") && (strings.Contains(msg, "thinking") || strings.Contains(msg, "redacted_thinking")) { - return true - } - - return false -} - -// isPromptTooLongError 检测是否为 prompt too long 错误 -func isPromptTooLongError(respBody []byte) bool { - msg := strings.ToLower(strings.TrimSpace(extractAntigravityErrorMessage(respBody))) - if msg == "" { - msg = strings.ToLower(string(respBody)) - } - return strings.Contains(msg, "prompt is too long") || - strings.Contains(msg, "request is too long") || - strings.Contains(msg, "context length exceeded") || - strings.Contains(msg, "max_tokens") -} - -// isPassthroughErrorMessage 检查错误消息是否在透传白名单中 -func isPassthroughErrorMessage(msg string) bool { - lower := strings.ToLower(msg) - for _, pattern := range antigravityPassthroughErrorMessages { - if strings.Contains(lower, pattern) { - return true - } - } - return false -} - -// getPassthroughOrDefault 若消息在白名单内则返回原始消息，否则返回默认消息 -func getPassthroughOrDefault(upstreamMsg, defaultMsg string) string { - if isPassthroughErrorMessage(upstreamMsg) { - return upstreamMsg - } - return defaultMsg -} - -func extractAntigravityErrorMessage(body []byte) string { - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return "" - } - - // Google-style: {"error": {"message": "..."}} - if errObj, ok := payload["error"].(map[string]any); ok { - if msg, ok := errObj["message"].(string); ok && strings.TrimSpace(msg) != "" { - return msg - } - } - - // Fallback: top-level message - if msg, ok := payload["message"].(string); ok && strings.TrimSpace(msg) != "" { - return msg - } - - return "" -} - -// stripThinkingFromClaudeRequest converts thinking blocks to text blocks in a Claude Messages request. -// This preserves the thinking content while avoiding signature validation errors. -// Note: redacted_thinking blocks are removed because they cannot be converted to text. -// It also disables top-level `thinking` to avoid upstream structural constraints for thinking mode. -func stripThinkingFromClaudeRequest(req *antigravity.ClaudeRequest) (bool, error) { - if req == nil { - return false, nil - } - - changed := false - if req.Thinking != nil { - req.Thinking = nil - changed = true - } - - for i := range req.Messages { - raw := req.Messages[i].Content - if len(raw) == 0 { - continue - } - - // If content is a string, nothing to strip. - var str string - if json.Unmarshal(raw, &str) == nil { - continue - } - - // Otherwise treat as an array of blocks and convert thinking blocks to text. - var blocks []map[string]any - if err := json.Unmarshal(raw, &blocks); err != nil { - continue - } - - filtered := make([]map[string]any, 0, len(blocks)) - modifiedAny := false - for _, block := range blocks { - t, _ := block["type"].(string) - switch t { - case "thinking": - thinkingText, _ := block["thinking"].(string) - if thinkingText != "" { - filtered = append(filtered, map[string]any{ - "type": "text", - "text": thinkingText, - }) - } - modifiedAny = true - case "redacted_thinking": - modifiedAny = true - case "": - if thinkingText, hasThinking := block["thinking"].(string); hasThinking { - if thinkingText != "" { - filtered = append(filtered, map[string]any{ - "type": "text", - "text": thinkingText, - }) - } - modifiedAny = true - } else { - filtered = append(filtered, block) - } - default: - filtered = append(filtered, block) - } - } - - if !modifiedAny { - continue - } - - if len(filtered) == 0 { - filtered = append(filtered, map[string]any{ - "type": "text", - "text": "(content removed)", - }) - } - - newRaw, err := json.Marshal(filtered) - if err != nil { - return changed, err - } - req.Messages[i].Content = newRaw - changed = true - } - - return changed, nil -} - -// stripSignatureSensitiveBlocksFromClaudeRequest is a stronger retry degradation that additionally converts -// tool blocks to plain text. Use this only after a thinking-only retry still fails with signature errors. -func stripSignatureSensitiveBlocksFromClaudeRequest(req *antigravity.ClaudeRequest) (bool, error) { - if req == nil { - return false, nil - } - - changed := false - if req.Thinking != nil { - req.Thinking = nil - changed = true - } - - for i := range req.Messages { - raw := req.Messages[i].Content - if len(raw) == 0 { - continue - } - - // If content is a string, nothing to strip. - var str string - if json.Unmarshal(raw, &str) == nil { - continue - } - - // Otherwise treat as an array of blocks and convert signature-sensitive blocks to text. - var blocks []map[string]any - if err := json.Unmarshal(raw, &blocks); err != nil { - continue - } - - filtered := make([]map[string]any, 0, len(blocks)) - modifiedAny := false - for _, block := range blocks { - t, _ := block["type"].(string) - switch t { - case "thinking": - // Convert thinking to text, skip if empty - thinkingText, _ := block["thinking"].(string) - if thinkingText != "" { - filtered = append(filtered, map[string]any{ - "type": "text", - "text": thinkingText, - }) - } - modifiedAny = true - case "redacted_thinking": - // Remove redacted_thinking (cannot convert encrypted content) - modifiedAny = true - case "tool_use": - // Convert tool_use to text to avoid upstream signature/thought_signature validation errors. - // This is a retry-only degradation path, so we prioritise request validity over tool semantics. - name, _ := block["name"].(string) - id, _ := block["id"].(string) - input := block["input"] - inputJSON, _ := json.Marshal(input) - text := "(tool_use)" - if name != "" { - text += " name=" + name - } - if id != "" { - text += " id=" + id - } - if len(inputJSON) > 0 && string(inputJSON) != "null" { - text += " input=" + string(inputJSON) - } - filtered = append(filtered, map[string]any{ - "type": "text", - "text": text, - }) - modifiedAny = true - case "tool_result": - // Convert tool_result to text so it stays consistent when tool_use is downgraded. - toolUseID, _ := block["tool_use_id"].(string) - isError, _ := block["is_error"].(bool) - content := block["content"] - contentJSON, _ := json.Marshal(content) - text := "(tool_result)" - if toolUseID != "" { - text += " tool_use_id=" + toolUseID - } - if isError { - text += " is_error=true" - } - if len(contentJSON) > 0 && string(contentJSON) != "null" { - text += "\n" + string(contentJSON) - } - filtered = append(filtered, map[string]any{ - "type": "text", - "text": text, - }) - modifiedAny = true - case "": - // Handle untyped block with "thinking" field - if thinkingText, hasThinking := block["thinking"].(string); hasThinking { - if thinkingText != "" { - filtered = append(filtered, map[string]any{ - "type": "text", - "text": thinkingText, - }) - } - modifiedAny = true - } else { - filtered = append(filtered, block) - } - default: - filtered = append(filtered, block) - } - } - - if !modifiedAny { - continue - } - - if len(filtered) == 0 { - // Keep request valid: upstream rejects empty content arrays. - filtered = append(filtered, map[string]any{ - "type": "text", - "text": "(content removed)", - }) - } - - newRaw, err := json.Marshal(filtered) - if err != nil { - return changed, err - } - req.Messages[i].Content = newRaw - changed = true - } - - return changed, nil -} - -// ForwardGemini 转发 Gemini 协议请求 -// -// 限流处理流程: -// -// 请求 → antigravityRetryLoop → 预检查(remaining>0? → 切换账号) → 发送上游 -// ├─ 成功 → 正常返回 -// └─ 429/503 → handleSmartRetry -// ├─ retryDelay >= 7s → 设置模型限流 + 清除粘性绑定 → 切换账号 -// └─ retryDelay < 7s → 等待后重试 1 次 -// ├─ 成功 → 正常返回 -// └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号 -type ForwardGeminiOption func(*forwardGeminiOptions) - -type forwardGeminiOptions struct { - groupID int64 - sessionHash string -} - -func WithForwardGeminiSession(groupID int64, sessionHash string) ForwardGeminiOption { - return func(opts *forwardGeminiOptions) { - opts.groupID = groupID - opts.sessionHash = sessionHash - } -} - -func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool, options ...ForwardGeminiOption) (*ForwardResult, error) { - startTime := time.Now() - forwardOpts := forwardGeminiOptions{} - for _, apply := range options { - if apply != nil { - apply(&forwardOpts) - } - } - - sessionID := getSessionID(c) - prefix := logPrefix(sessionID, account.Name) - - if strings.TrimSpace(originalModel) == "" { - return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing model in URL") - } - if strings.TrimSpace(action) == "" { - return nil, s.writeGoogleError(c, http.StatusBadRequest, "Missing action in URL") - } - if len(body) == 0 { - return nil, s.writeGoogleError(c, http.StatusBadRequest, "Request body is empty") - } - - // 解析请求以获取 image_size（用于图片计费） - imageInputSize := s.extractImageInputSize(body) - imageSize := normalizeOpenAIImageSizeTier(imageInputSize) - - switch action { - 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(startTime), - FirstTokenMs: nil, - }, nil - default: - return nil, s.writeGoogleError(c, http.StatusNotFound, "Unsupported action: "+action) - } - - mappedModel := s.getMappedModel(account, originalModel) - if mappedModel == "" { - MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalFeatureGate) - return nil, s.writeGoogleError(c, http.StatusForbidden, fmt.Sprintf("model %s not in whitelist", originalModel)) - } - billingModel := mappedModel - - // 获取 access_token - if s.tokenProvider == nil { - return nil, s.writeGoogleError(c, http.StatusBadGateway, "Antigravity token provider not configured") - } - accessToken, err := s.tokenProvider.GetAccessToken(ctx, account) - if err != nil { - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: []byte(`{"error":{"message":"Failed to get upstream access token","status":"UNAVAILABLE"}}`), - } - } - - projectID, err := resolveAntigravityProjectID(account) - if err != nil { - _ = s.writeGoogleError(c, http.StatusBadRequest, err.Error()) - return nil, err - } - - // 代理 URL - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - // Antigravity 上游要求必须包含身份提示词，注入到请求中 - injectedBody, err := injectIdentityPatchToGeminiRequest(body) - if err != nil { - return nil, s.writeGoogleError(c, http.StatusBadRequest, "Invalid request body") - } - - // 清理 Schema - if cleanedBody, err := cleanGeminiRequest(injectedBody); err == nil { - injectedBody = cleanedBody - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Cleaned request schema in forwarded request for account %s", account.Name) - } else { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Failed to clean schema: %v", err) - } - - // 包装请求 - wrappedBody, err := s.wrapV1InternalRequest(projectID, mappedModel, injectedBody) - if err != nil { - return nil, s.writeGoogleError(c, http.StatusInternalServerError, "Failed to build upstream request") - } - - // Antigravity 上游只支持流式请求，统一使用 streamGenerateContent - // 如果客户端请求非流式，在响应处理阶段会收集完整流式响应后返回 - upstreamAction := "streamGenerateContent" - - // 执行带重试的请求 - result, err := s.antigravityRetryLoop(antigravityRetryLoopParams{ - ctx: ctx, - prefix: prefix, - account: account, - proxyURL: proxyURL, - accessToken: accessToken, - action: upstreamAction, - body: wrappedBody, - c: c, - httpUpstream: s.httpUpstream, - settingService: s.settingService, - accountRepo: s.accountRepo, - handleError: s.handleUpstreamError, - requestedModel: originalModel, - isStickySession: isStickySession, // ForwardGemini 由上层判断粘性会话 - groupID: forwardOpts.groupID, - sessionHash: forwardOpts.sessionHash, - }) - if err != nil { - // 检查是否是账号切换信号，转换为 UpstreamFailoverError 让 Handler 切换账号 - if switchErr, ok := IsAntigravityAccountSwitchError(err); ok { - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusServiceUnavailable, - ForceCacheBilling: switchErr.IsStickySession, - } - } - // 区分客户端取消和真正的上游失败，返回更准确的错误消息 - if c.Request.Context().Err() != nil { - return nil, s.writeGoogleError(c, http.StatusBadGateway, "Client disconnected before upstream response") - } - return nil, s.writeGoogleError(c, http.StatusBadGateway, "Upstream request failed after retries") - } - resp := result.resp - defer func() { - if resp != nil && resp.Body != nil { - _ = resp.Body.Close() - } - }() - - // 处理错误响应 - if resp.StatusCode >= 400 { - respBody := s.readUpstreamErrorBody(resp) - contentType := resp.Header.Get("Content-Type") - // 尽早关闭原始响应体，释放连接；后续逻辑仍可能需要读取 body，因此用内存副本重新包装。 - _ = resp.Body.Close() - resp.Body = io.NopCloser(bytes.NewReader(respBody)) - - // 模型兜底：模型不存在且开启 fallback 时，自动用 fallback 模型重试一次 - if s.settingService != nil && s.settingService.IsModelFallbackEnabled(ctx) && - isModelNotFoundError(resp.StatusCode, respBody) { - fallbackModel := s.settingService.GetFallbackModel(ctx, PlatformAntigravity) - if fallbackModel != "" && fallbackModel != mappedModel { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Model not found (%s), retrying with fallback model %s (account: %s)", mappedModel, fallbackModel, account.Name) - - fallbackWrapped, err := s.wrapV1InternalRequest(projectID, fallbackModel, injectedBody) - if err == nil { - fallbackReq, err := antigravity.NewAPIRequest(ctx, upstreamAction, accessToken, fallbackWrapped) - if err == nil { - fallbackResp, err := s.httpUpstream.Do(fallbackReq, proxyURL, account.ID, account.Concurrency) - if err == nil && fallbackResp.StatusCode < 400 { - _ = resp.Body.Close() - resp = fallbackResp - } else if fallbackResp != nil { - _ = fallbackResp.Body.Close() - } - } - } - } - } - - // Gemini 原生请求中的 thoughtSignature 可能来自旧上下文/旧账号，触发上游严格校验后返回 - // "Corrupted thought signature."。检测到此类 400 时，将 thoughtSignature 清理为 dummy 值后重试一次。 - signatureCheckBody := respBody - if unwrapped, unwrapErr := s.unwrapV1InternalResponse(respBody); unwrapErr == nil && len(unwrapped) > 0 { - signatureCheckBody = unwrapped - } - if resp.StatusCode == http.StatusBadRequest && - s.settingService != nil && - s.settingService.IsSignatureRectifierEnabled(ctx) && - isSignatureRelatedError(signatureCheckBody) && - bytes.Contains(injectedBody, []byte(`"thoughtSignature"`)) { - upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(signatureCheckBody))) - upstreamDetail := s.getUpstreamErrorDetail(signatureCheckBody) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "signature_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: detected signature-related 400, retrying with cleaned thought signatures", account.ID) - - cleanedInjectedBody := CleanGeminiNativeThoughtSignatures(injectedBody) - retryWrappedBody, wrapErr := s.wrapV1InternalRequest(projectID, mappedModel, cleanedInjectedBody) - if wrapErr == nil { - retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{ - ctx: ctx, - prefix: prefix, - account: account, - proxyURL: proxyURL, - accessToken: accessToken, - action: upstreamAction, - body: retryWrappedBody, - c: c, - httpUpstream: s.httpUpstream, - settingService: s.settingService, - accountRepo: s.accountRepo, - handleError: s.handleUpstreamError, - requestedModel: originalModel, - isStickySession: isStickySession, - groupID: forwardOpts.groupID, - sessionHash: forwardOpts.sessionHash, - }) - if retryErr == nil { - retryResp := retryResult.resp - if retryResp.StatusCode < 400 { - resp = retryResp - } else { - retryRespBody := s.readUpstreamErrorBody(retryResp) - _ = retryResp.Body.Close() - retryOpsBody := retryRespBody - if retryUnwrapped, unwrapErr := s.unwrapV1InternalResponse(retryRespBody); unwrapErr == nil && len(retryUnwrapped) > 0 { - retryOpsBody = retryUnwrapped - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: retryResp.StatusCode, - UpstreamRequestID: retryResp.Header.Get("x-request-id"), - Kind: "signature_retry", - Message: sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(retryOpsBody))), - Detail: s.getUpstreamErrorDetail(retryOpsBody), - }) - respBody = retryRespBody - resp = &http.Response{ - StatusCode: retryResp.StatusCode, - Header: retryResp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryRespBody)), - } - contentType = resp.Header.Get("Content-Type") - } - } else { - if switchErr, ok := IsAntigravityAccountSwitchError(retryErr); ok { - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: http.StatusServiceUnavailable, - Kind: "failover", - Message: sanitizeUpstreamErrorMessage(retryErr.Error()), - }) - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusServiceUnavailable, - ForceCacheBilling: switchErr.IsStickySession, - } - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: 0, - Kind: "signature_retry_request_error", - Message: sanitizeUpstreamErrorMessage(retryErr.Error()), - }) - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: signature retry request failed: %v", account.ID, retryErr) - } - } else { - logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: signature retry wrap failed: %v", account.ID, wrapErr) - } - } - - // fallback 成功：继续按正常响应处理 - if resp.StatusCode < 400 { - goto handleSuccess - } - - requestID := resp.Header.Get("x-request-id") - if requestID != "" { - c.Header("x-request-id", requestID) - } - - unwrapped, unwrapErr := s.unwrapV1InternalResponse(respBody) - unwrappedForOps := unwrapped - if unwrapErr != nil || len(unwrappedForOps) == 0 { - unwrappedForOps = respBody - } - s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, forwardOpts.groupID, forwardOpts.sessionHash, isStickySession) - upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(unwrappedForOps)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - upstreamDetail := s.getUpstreamErrorDetail(unwrappedForOps) - - // Always record upstream context for Ops error logs, even when we will failover. - setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) - - // 精确匹配服务端配置类 400 错误，触发同账号重试 + failover - if resp.StatusCode == http.StatusBadRequest && isGoogleProjectConfigError(strings.ToLower(upstreamMsg)) { - log.Printf("%s status=400 google_config_error failover=true upstream_message=%q account=%d", prefix, upstreamMsg, account.ID) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: requestID, - Kind: "failover", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: unwrappedForOps, RetryableOnSameAccount: true} - } - - if s.shouldFailoverUpstreamError(resp.StatusCode) { - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: requestID, - Kind: "failover", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: unwrappedForOps} - } - if contentType == "" { - contentType = "application/json" - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: requestID, - Kind: "http_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] upstream error status=%d body=%s", resp.StatusCode, truncateForLog(unwrappedForOps, 500)) - MarkResponseCommitted(c) - c.Data(resp.StatusCode, contentType, unwrappedForOps) - return nil, fmt.Errorf("antigravity upstream error: %d", resp.StatusCode) - } - -handleSuccess: - requestID := resp.Header.Get("x-request-id") - if requestID != "" { - c.Header("x-request-id", requestID) - } - - var usage *ClaudeUsage - var firstTokenMs *int - var clientDisconnect bool - - if stream { - // 客户端要求流式，直接透传 - streamRes, err := s.handleGeminiStreamingResponse(c, resp, startTime) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_error error=%v", prefix, err) - return nil, err - } - usage = streamRes.usage - firstTokenMs = streamRes.firstTokenMs - clientDisconnect = streamRes.clientDisconnect - } else { - // 客户端要求非流式，收集流式响应后返回 - streamRes, err := s.handleGeminiStreamToNonStreaming(c, resp, startTime) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=stream_collect_error error=%v", prefix, err) - return nil, err - } - usage = streamRes.usage - firstTokenMs = streamRes.firstTokenMs - } - - if usage == nil { - usage = &ClaudeUsage{} - } - - // 判断是否为图片生成模型 - imageCount := 0 - if isImageGenerationModel(mappedModel) { - // Gemini 图片生成 API 每次请求只生成一张图片（API 限制） - imageCount = 1 - } - - return &ForwardResult{ - RequestID: requestID, - Usage: *usage, - Model: originalModel, - UpstreamModel: billingModel, - Stream: stream, - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, - ClientDisconnect: clientDisconnect, - ImageCount: imageCount, - ImageSize: imageSize, - ImageInputSize: imageInputSize, - }, nil -} - -func (s *AntigravityGatewayService) shouldFailoverUpstreamError(statusCode int) bool { - switch statusCode { - case 401, 403, 429, 529: - return true - default: - return statusCode >= 500 - } -} - -// isGoogleProjectConfigError 判断（已提取的小写）错误消息是否属于 Google 服务端配置类问题。 -// 只精确匹配已知的服务端侧错误，避免对客户端请求错误做无意义重试。 -// 适用于所有走 Google 后端的平台（Antigravity、Gemini）。 -func isGoogleProjectConfigError(lowerMsg string) bool { - // Google 间歇性 Bug：Project ID 有效但被临时识别失败 - return strings.Contains(lowerMsg, "invalid project resource name") -} - -// googleConfigErrorCooldown 服务端配置类 400 错误的临时封禁时长 -const googleConfigErrorCooldown = 1 * time.Minute - -// tempUnscheduleGoogleConfigError 对服务端配置类 400 错误触发临时封禁， -// 避免短时间内反复调度到同一个有问题的账号。 -func tempUnscheduleGoogleConfigError(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) { - until := time.Now().Add(googleConfigErrorCooldown) - reason := "400: invalid project resource name (auto temp-unschedule 1m)" - if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil { - log.Printf("%s temp_unschedule_failed account=%d error=%v", logPrefix, accountID, err) - } else { - log.Printf("%s temp_unscheduled account=%d until=%v reason=%q", logPrefix, accountID, until.Format("15:04:05"), reason) - } -} - -// emptyResponseCooldown 空流式响应的临时封禁时长 -const emptyResponseCooldown = 1 * time.Minute - -// tempUnscheduleEmptyResponse 对空流式响应触发临时封禁， -// 避免短时间内反复调度到同一个返回空响应的账号。 -func tempUnscheduleEmptyResponse(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) { - until := time.Now().Add(emptyResponseCooldown) - reason := "empty stream response (auto temp-unschedule 1m)" - if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil { - log.Printf("%s temp_unschedule_failed account=%d error=%v", logPrefix, accountID, err) - } else { - log.Printf("%s temp_unscheduled account=%d until=%v reason=%q", logPrefix, accountID, until.Format("15:04:05"), reason) - } -} - -// sleepAntigravityBackoffWithContext 带 context 取消检查的退避等待 -// 返回 true 表示正常完成等待，false 表示 context 已取消 -func sleepAntigravityBackoffWithContext(ctx context.Context, attempt int) bool { - delay := antigravityRetryBaseDelay * time.Duration(1< antigravityRetryMaxDelay { - delay = antigravityRetryMaxDelay - } - - // +/- 20% jitter - r := mathrand.New(mathrand.NewSource(time.Now().UnixNano())) - jitter := time.Duration(float64(delay) * 0.2 * (r.Float64()*2 - 1)) - sleepFor := delay + jitter - if sleepFor < 0 { - sleepFor = 0 - } - - timer := time.NewTimer(sleepFor) - select { - case <-ctx.Done(): - timer.Stop() - return false - case <-timer.C: - return true - } -} - -// isSingleAccountRetry 检查 context 中是否设置了单账号退避重试标记 -func isSingleAccountRetry(ctx context.Context) bool { - v, _ := SingleAccountRetryFromContext(ctx) - return v -} - -// setModelRateLimitByModelName 使用官方模型 ID 设置模型级限流 -// 直接使用上游返回的模型 ID（如 claude-sonnet-4-5）作为限流 key -// 返回是否已成功设置（若模型名为空或 repo 为 nil 将返回 false） -func setModelRateLimitByModelName(ctx context.Context, repo AccountRepository, accountID int64, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { - if repo == nil || modelName == "" { - return false - } - // 直接使用官方模型 ID 作为 key，不再转换为 scope - if err := repo.SetModelRateLimit(ctx, accountID, modelName, resetAt); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_failed model=%s error=%v", prefix, statusCode, modelName, err) - return false - } - if afterSmartRetry { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_after_smart_retry model=%s account=%d reset_in=%v", prefix, statusCode, modelName, accountID, time.Until(resetAt).Truncate(time.Second)) - } else { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", prefix, statusCode, modelName, accountID, time.Until(resetAt).Truncate(time.Second)) - } - return true -} - -func (s *AntigravityGatewayService) setAntigravityModelRateLimits(ctx context.Context, repo AccountRepository, account *Account, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool { - if account == nil || repo == nil { - return false - } - keys := antigravityModelRateLimitKeys(modelName) - if len(keys) == 0 { - return false - } - - success := false - for _, key := range keys { - if setModelRateLimitByModelName(ctx, repo, account.ID, key, prefix, statusCode, resetAt, afterSmartRetry) { - s.updateAccountModelRateLimitInCache(ctx, account, key, resetAt) - success = true - } - } - return success -} - -func (s *AntigravityGatewayService) clearStickySession(ctx context.Context, groupID int64, sessionHash string) { - if s == nil || s.cache == nil || strings.TrimSpace(sessionHash) == "" { - return - } - if err := s.cache.DeleteSessionAccountID(ctx, groupID, sessionHash); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] sticky_session_clear_failed group_id=%d session=%s err=%v", groupID, shortSessionHash(sessionHash), err) - } -} - -func antigravityFallbackCooldownSeconds() (time.Duration, bool) { - raw := strings.TrimSpace(os.Getenv(antigravityFallbackSecondsEnv)) - if raw == "" { - return 0, false - } - seconds, err := strconv.Atoi(raw) - if err != nil || seconds <= 0 { - return 0, false - } - return time.Duration(seconds) * time.Second, true -} - -// antigravitySmartRetryInfo 智能重试所需的信息 -type antigravitySmartRetryInfo struct { - RetryDelay time.Duration // 重试延迟时间 - ModelName string // 限流的模型名称（如 "claude-sonnet-4-5"） - IsModelCapacityExhausted bool // 是否为模型容量不足（MODEL_CAPACITY_EXHAUSTED） -} - -// parseAntigravitySmartRetryInfo 解析 Google RPC RetryInfo 和 ErrorInfo 信息 -// 返回解析结果，如果解析失败或不满足条件返回 nil -// -// 支持两种情况： -// 1. 429 RESOURCE_EXHAUSTED + RATE_LIMIT_EXCEEDED： -// - error.status == "RESOURCE_EXHAUSTED" -// - error.details[].reason == "RATE_LIMIT_EXCEEDED" -// -// 2. 503 UNAVAILABLE + MODEL_CAPACITY_EXHAUSTED： -// - error.status == "UNAVAILABLE" -// - error.details[].reason == "MODEL_CAPACITY_EXHAUSTED" -// -// 必须满足以下条件才会返回有效值： -// - error.details[] 中存在 @type == "type.googleapis.com/google.rpc.RetryInfo" 的元素 -// - 该元素包含 retryDelay 字段，格式为 "数字s"（如 "0.201506475s"） -func parseAntigravitySmartRetryInfo(body []byte) *antigravitySmartRetryInfo { - var parsed map[string]any - if err := json.Unmarshal(body, &parsed); err != nil { - return nil - } - - errObj, ok := parsed["error"].(map[string]any) - if !ok { - return nil - } - - // 检查 status 是否符合条件 - // 情况1: 429 RESOURCE_EXHAUSTED (需要进一步检查 reason == RATE_LIMIT_EXCEEDED) - // 情况2: 503 UNAVAILABLE (需要进一步检查 reason == MODEL_CAPACITY_EXHAUSTED) - status, _ := errObj["status"].(string) - isResourceExhausted := status == googleRPCStatusResourceExhausted - isUnavailable := status == googleRPCStatusUnavailable - - if !isResourceExhausted && !isUnavailable { - return nil - } - - details, ok := errObj["details"].([]any) - if !ok { - return nil - } - - var retryDelay time.Duration - var modelName string - var hasRateLimitExceeded bool // 429 需要此 reason - var hasModelCapacityExhausted bool // 503 需要此 reason - - for _, d := range details { - dm, ok := d.(map[string]any) - if !ok { - continue - } - - atType, _ := dm["@type"].(string) - - // 从 ErrorInfo 提取模型名称和 reason - if atType == googleRPCTypeErrorInfo { - if meta, ok := dm["metadata"].(map[string]any); ok { - if model, ok := meta["model"].(string); ok { - modelName = normalizeAntigravityModelName(model) - } - } - // 检查 reason - if reason, ok := dm["reason"].(string); ok { - if reason == googleRPCReasonModelCapacityExhausted { - hasModelCapacityExhausted = true - } - if reason == googleRPCReasonRateLimitExceeded { - hasRateLimitExceeded = true - } - } - continue - } - - // 从 RetryInfo 提取重试延迟 - if atType == googleRPCTypeRetryInfo { - delay, ok := dm["retryDelay"].(string) - if !ok || delay == "" { - continue - } - // 使用 time.ParseDuration 解析，支持所有 Go duration 格式 - // 例如: "0.5s", "10s", "4m50s", "1h30m", "200ms" 等 - dur, err := time.ParseDuration(delay) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] failed to parse retryDelay: %s error=%v", delay, err) - continue - } - retryDelay = dur - } - } - - // 验证条件 - // 情况1: RESOURCE_EXHAUSTED 需要有 RATE_LIMIT_EXCEEDED reason - // 情况2: UNAVAILABLE 需要有 MODEL_CAPACITY_EXHAUSTED reason - if isResourceExhausted && !hasRateLimitExceeded { - return nil - } - if isUnavailable && !hasModelCapacityExhausted { - return nil - } - - // 必须有模型名才返回有效结果 - if modelName == "" { - return nil - } - - // 如果上游未提供 retryDelay，使用默认限流时间 - if retryDelay <= 0 { - retryDelay = antigravityDefaultRateLimitDuration - } - - return &antigravitySmartRetryInfo{ - RetryDelay: retryDelay, - ModelName: modelName, - IsModelCapacityExhausted: hasModelCapacityExhausted, - } -} - -// shouldTriggerAntigravitySmartRetry 判断是否应该触发智能重试 -// 返回： -// - shouldRetry: 是否应该智能重试（retryDelay < antigravityRateLimitThreshold，或 MODEL_CAPACITY_EXHAUSTED） -// - shouldRateLimitModel: 是否应该限流模型并切换账号（仅 RATE_LIMIT_EXCEEDED 且 retryDelay >= 阈值） -// - waitDuration: 等待时间 -// - modelName: 限流的模型名称 -// - isModelCapacityExhausted: 是否为模型容量不足（MODEL_CAPACITY_EXHAUSTED） -func shouldTriggerAntigravitySmartRetry(account *Account, respBody []byte) (shouldRetry bool, shouldRateLimitModel bool, waitDuration time.Duration, modelName string, isModelCapacityExhausted bool) { - if account.Platform != PlatformAntigravity { - return false, false, 0, "", false - } - - info := parseAntigravitySmartRetryInfo(respBody) - if info == nil { - return false, false, 0, "", false - } - - // MODEL_CAPACITY_EXHAUSTED（模型容量不足）：所有账号共享同一模型容量池 - // 切换账号无意义，使用固定 1s 间隔重试 - if info.IsModelCapacityExhausted { - return true, false, antigravityModelCapacityRetryWait, info.ModelName, true - } - - // RATE_LIMIT_EXCEEDED（账号级限流）： - // retryDelay >= 阈值：直接限流模型，不重试 - // 注意：如果上游未提供 retryDelay，parseAntigravitySmartRetryInfo 已设置为默认 30s - if info.RetryDelay >= antigravityRateLimitThreshold { - return false, true, info.RetryDelay, info.ModelName, false - } - - // retryDelay < 阈值：智能重试 - waitDuration = info.RetryDelay - if waitDuration < antigravitySmartRetryMinWait { - waitDuration = antigravitySmartRetryMinWait - } - - return true, false, waitDuration, info.ModelName, false -} - -// handleModelRateLimitParams 模型级限流处理参数 -type handleModelRateLimitParams struct { - ctx context.Context - prefix string - account *Account - statusCode int - body []byte - cache GatewayCache - groupID int64 - sessionHash string - isStickySession bool -} - -// handleModelRateLimitResult 模型级限流处理结果 -type handleModelRateLimitResult struct { - Handled bool // 是否已处理 - ShouldRetry bool // 是否等待后重试 - WaitDuration time.Duration // 等待时间 - SwitchError *AntigravityAccountSwitchError // 账号切换错误 -} - -// handleModelRateLimit 处理模型级限流（在原有逻辑之前调用） -// 仅处理 429/503，解析模型名和 retryDelay -// - MODEL_CAPACITY_EXHAUSTED: 返回 Handled=true（实际重试由 handleSmartRetry 处理） -// - RATE_LIMIT_EXCEEDED + retryDelay < 阈值: 返回 ShouldRetry=true，由调用方等待后重试 -// - RATE_LIMIT_EXCEEDED + retryDelay >= 阈值: 设置模型限流 + 清除粘性会话 + 返回 SwitchError -func (s *AntigravityGatewayService) handleModelRateLimit(p *handleModelRateLimitParams) *handleModelRateLimitResult { - if p.statusCode != 429 && p.statusCode != 503 { - return &handleModelRateLimitResult{Handled: false} - } - - info := parseAntigravitySmartRetryInfo(p.body) - if info == nil || info.ModelName == "" { - return &handleModelRateLimitResult{Handled: false} - } - - // MODEL_CAPACITY_EXHAUSTED：模型容量不足，所有账号共享同一容量池 - // 切换账号无意义，不设置模型限流（实际重试由 handleSmartRetry 处理） - if info.IsModelCapacityExhausted { - log.Printf("%s status=%d model_capacity_exhausted model=%s (not switching account, retry handled by smart retry)", - p.prefix, p.statusCode, info.ModelName) - return &handleModelRateLimitResult{ - Handled: true, - } - } - - // RATE_LIMIT_EXCEEDED: < antigravityRateLimitThreshold: 等待后重试 - if info.RetryDelay < antigravityRateLimitThreshold { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_wait model=%s wait=%v", - p.prefix, p.statusCode, info.ModelName, info.RetryDelay) - return &handleModelRateLimitResult{ - Handled: true, - ShouldRetry: true, - WaitDuration: info.RetryDelay, - } - } - - // RATE_LIMIT_EXCEEDED: >= antigravityRateLimitThreshold: 设置限流 + 清除粘性会话 + 切换账号 - s.setModelRateLimitAndClearSession(p, info) - - return &handleModelRateLimitResult{ - Handled: true, - SwitchError: &AntigravityAccountSwitchError{ - OriginalAccountID: p.account.ID, - RateLimitedModel: info.ModelName, - IsStickySession: p.isStickySession, - }, - } -} - -// setModelRateLimitAndClearSession 设置模型限流并清除粘性会话 -func (s *AntigravityGatewayService) setModelRateLimitAndClearSession(p *handleModelRateLimitParams, info *antigravitySmartRetryInfo) { - resetAt := time.Now().Add(info.RetryDelay) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", - p.prefix, p.statusCode, info.ModelName, p.account.ID, info.RetryDelay) - - s.setAntigravityModelRateLimits(p.ctx, s.accountRepo, p.account, info.ModelName, p.prefix, p.statusCode, resetAt, false) - - // 清除粘性会话绑定 - if p.cache != nil && p.sessionHash != "" { - _ = p.cache.DeleteSessionAccountID(p.ctx, p.groupID, p.sessionHash) - } -} - -// updateAccountModelRateLimitInCache 立即更新 Redis 中账号的模型限流状态 -func (s *AntigravityGatewayService) updateAccountModelRateLimitInCache(ctx context.Context, account *Account, modelKey string, resetAt time.Time) { - if s.schedulerSnapshot == nil || account == nil || modelKey == "" { - return - } - - // 更新账号对象的 Extra 字段 - if account.Extra == nil { - account.Extra = make(map[string]any) - } - - limits, _ := account.Extra["model_rate_limits"].(map[string]any) - if limits == nil { - limits = make(map[string]any) - account.Extra["model_rate_limits"] = limits - } - - limits[modelKey] = map[string]any{ - "rate_limited_at": time.Now().UTC().Format(time.RFC3339), - "rate_limit_reset_at": resetAt.UTC().Format(time.RFC3339), - } - - // 更新 Redis 快照 - if err := s.schedulerSnapshot.UpdateAccountInCache(ctx, account); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] cache_update_failed account=%d model=%s err=%v", account.ID, modelKey, err) - } -} - -func (s *AntigravityGatewayService) handleUpstreamError( - ctx context.Context, prefix string, account *Account, - statusCode int, headers http.Header, body []byte, - requestedModel string, - groupID int64, sessionHash string, isStickySession bool, -) *handleModelRateLimitResult { - // 遵守自定义错误码策略：未命中则跳过所有限流处理 - if !account.ShouldHandleErrorCode(statusCode) { - return nil - } - // 模型级限流处理（优先） - result := s.handleModelRateLimit(&handleModelRateLimitParams{ - ctx: ctx, - prefix: prefix, - account: account, - statusCode: statusCode, - body: body, - cache: s.cache, - groupID: groupID, - sessionHash: sessionHash, - isStickySession: isStickySession, - }) - if result.Handled { - return result - } - - // 503 仅处理模型限流（MODEL_CAPACITY_EXHAUSTED），非模型限流不做额外处理 - // 避免将普通的 503 错误误判为账号问题 - if statusCode == 503 { - return nil - } - - // 429：尝试解析模型级限流，解析失败时兜底为账号级限流 - if statusCode == 429 { - if logBody, maxBytes := s.getLogConfig(); logBody { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity-Debug] 429 response body: %s", truncateString(string(body), maxBytes)) - } - - resetAt := ParseGeminiRateLimitResetTime(body) - defaultDur := s.getDefaultRateLimitDuration() - - // 尝试解析模型 key 并设置模型级限流 - // - // 注意：requestedModel 可能是"映射前"的请求模型名（例如 claude-opus-4-6）， - // 调度与限流判定使用的是 Antigravity 最终模型名（包含映射与 thinking 后缀）。 - // 因此这里必须写入最终模型 key，确保后续调度能正确避开已限流模型。 - modelKey := resolveFinalAntigravityModelKey(ctx, account, requestedModel) - if strings.TrimSpace(modelKey) == "" { - // 极少数情况下无法映射（理论上不应发生：能转发成功说明映射已通过）， - // 保持旧行为作为兜底，避免完全丢失模型级限流记录。 - modelKey = resolveAntigravityModelKey(requestedModel) - } - if modelKey != "" { - ra := s.resolveResetTime(resetAt, defaultDur) - if !s.setAntigravityModelRateLimits(ctx, s.accountRepo, account, modelKey, prefix, statusCode, ra, false) { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s", prefix, modelKey) - } else { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limited model=%s account=%d reset_at=%v reset_in=%v", - prefix, modelKey, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second)) - } - return nil - } - - // 无法解析模型 key，兜底为账号级限流 - ra := s.resolveResetTime(resetAt, defaultDur) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limited account=%d reset_at=%v reset_in=%v (fallback)", - prefix, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second)) - if err := s.accountRepo.SetRateLimited(ctx, account.ID, ra); err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limit_set_failed account=%d error=%v", prefix, account.ID, err) - } - return nil - } - // 其他错误码继续使用 rateLimitService - if s.rateLimitService == nil { - return nil - } - shouldDisable := s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, headers, body) - if shouldDisable { - logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d marked_error", prefix, statusCode) - } - return nil -} - -// getDefaultRateLimitDuration 获取默认限流时间 -func (s *AntigravityGatewayService) getDefaultRateLimitDuration() time.Duration { - defaultDur := antigravityDefaultRateLimitDuration - if s.settingService != nil && s.settingService.cfg != nil && s.settingService.cfg.Gateway.AntigravityFallbackCooldownMinutes > 0 { - defaultDur = time.Duration(s.settingService.cfg.Gateway.AntigravityFallbackCooldownMinutes) * time.Minute - } - if override, ok := antigravityFallbackCooldownSeconds(); ok { - defaultDur = override - } - return defaultDur -} - -// resolveResetTime 根据解析的重置时间或默认时长计算重置时间点 -func (s *AntigravityGatewayService) resolveResetTime(resetAt *int64, defaultDur time.Duration) time.Time { - if resetAt != nil { - return time.Unix(*resetAt, 0) - } - return time.Now().Add(defaultDur) -} - -type antigravityStreamResult struct { - usage *ClaudeUsage - firstTokenMs *int - clientDisconnect bool // 客户端是否在流式传输过程中断开 -} - -// antigravityClientWriter 封装流式响应的客户端写入，自动检测断开并标记。 -// 断开后所有写入操作变为 no-op，调用方通过 Disconnected() 判断是否继续 drain 上游。 -type antigravityClientWriter struct { - w gin.ResponseWriter - flusher http.Flusher - disconnected bool - prefix string // 日志前缀，标识来源方法 -} - -func newAntigravityClientWriter(w gin.ResponseWriter, flusher http.Flusher, prefix string) *antigravityClientWriter { - return &antigravityClientWriter{w: w, flusher: flusher, prefix: prefix} -} - -// Write 写入数据到客户端，写入失败时标记断开并返回 false -func (cw *antigravityClientWriter) Write(p []byte) bool { - if cw.disconnected { - return false - } - if _, err := cw.w.Write(p); err != nil { - cw.markDisconnected() - return false - } - cw.flusher.Flush() - return true -} - -// Fprintf 格式化写入数据到客户端，写入失败时标记断开并返回 false -func (cw *antigravityClientWriter) Fprintf(format string, args ...any) bool { - if cw.disconnected { - return false - } - if _, err := fmt.Fprintf(cw.w, format, args...); err != nil { - cw.markDisconnected() - return false - } - cw.flusher.Flush() - return true -} - -func (cw *antigravityClientWriter) Disconnected() bool { return cw.disconnected } - -func (cw *antigravityClientWriter) markDisconnected() { - cw.disconnected = true - logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during streaming (%s), continuing to drain upstream for billing", cw.prefix) -} - -// handleStreamReadError 处理上游读取错误的通用逻辑。 -// 返回 (clientDisconnect, handled)：handled=true 表示错误已处理，调用方应返回已收集的 usage。 -func handleStreamReadError(err error, clientDisconnected bool, prefix string) (disconnect bool, handled bool) { - if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { - logger.LegacyPrintf("service.antigravity_gateway", "Context canceled during streaming (%s), returning collected usage", prefix) - return true, true - } - if clientDisconnected { - logger.LegacyPrintf("service.antigravity_gateway", "Upstream read error after client disconnect (%s): %v, returning collected usage", prefix, err) - return true, true - } - return false, false -} - -func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) { - c.Status(resp.StatusCode) - c.Header("Cache-Control", "no-cache") - c.Header("Connection", "keep-alive") - c.Header("X-Accel-Buffering", "no") - - contentType := resp.Header.Get("Content-Type") - if contentType == "" { - contentType = "text/event-stream; charset=utf-8" - } - c.Header("Content-Type", contentType) - - flusher, ok := c.Writer.(http.Flusher) - if !ok { - return nil, errors.New("streaming not supported") - } - - // 使用 Scanner 并限制单行大小，避免 ReadString 无上限导致 OOM - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.settingService.cfg.Gateway.MaxLineSize - } - scanBuf := getSSEScannerBuf64K() - scanner.Buffer(scanBuf[:0], maxLineSize) - usage := &ClaudeUsage{} - var firstTokenMs *int - - type scanEvent struct { - line string - err error - } - // 独立 goroutine 读取上游，避免读取阻塞影响超时处理 - events := make(chan scanEvent, 16) - done := make(chan struct{}) - sendEvent := func(ev scanEvent) bool { - select { - case events <- ev: - return true - case <-done: - return false - } - } - var lastReadAt int64 - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - go func(scanBuf *sseScannerBuf64K) { - defer putSSEScannerBuf64K(scanBuf) - defer close(events) - for scanner.Scan() { - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { - return - } - } - if err := scanner.Err(); err != nil { - _ = sendEvent(scanEvent{err: err}) - } - }(scanBuf) - defer close(done) - - // 上游数据间隔超时保护（防止上游挂起长期占用连接） - streamInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second - } - var intervalTicker *time.Ticker - if streamInterval > 0 { - intervalTicker = time.NewTicker(streamInterval) - defer intervalTicker.Stop() - } - var intervalCh <-chan time.Time - if intervalTicker != nil { - intervalCh = intervalTicker.C - } - - // 下游 keepalive：防止代理/Cloudflare Tunnel 因连接空闲而断开 - keepaliveInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { - keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second - } - var keepaliveTicker *time.Ticker - if keepaliveInterval > 0 { - keepaliveTicker = time.NewTicker(keepaliveInterval) - defer keepaliveTicker.Stop() - } - var keepaliveCh <-chan time.Time - if keepaliveTicker != nil { - keepaliveCh = keepaliveTicker.C - } - lastDataAt := time.Now() - - cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity gemini") - - // 仅发送一次错误事件，避免多次写入导致协议混乱 - errorEventSent := false - sendErrorEvent := func(reason string) { - if errorEventSent || cw.Disconnected() { - return - } - errorEventSent = true - _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"%s\"}\n\n", reason) - flusher.Flush() - } - - for { - select { - case ev, ok := <-events: - if !ok { - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}, nil - } - if ev.err != nil { - if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity gemini"); handled { - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: disconnect}, nil - } - if errors.Is(ev.err, bufio.ErrTooLong) { - logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err) - sendErrorEvent("response_too_large") - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, ev.err - } - sendErrorEvent("stream_read_error") - return nil, ev.err - } - - lastDataAt = time.Now() - - line := ev.line - trimmed := strings.TrimRight(line, "\r\n") - if strings.HasPrefix(trimmed, "data:") { - payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) - if payload == "" || payload == "[DONE]" { - cw.Fprintf("%s\n", line) - continue - } - - // 解包 v1internal 响应 - inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) - if parseErr == nil && inner != nil { - payload = string(inner) - } - - // 解析 usage - if u := extractGeminiUsage(inner); u != nil { - usage = u - } - var parsed map[string]any - if json.Unmarshal(inner, &parsed) == nil { - // Check for MALFORMED_FUNCTION_CALL - if candidates, ok := parsed["candidates"].([]any); ok && len(candidates) > 0 { - if cand, ok := candidates[0].(map[string]any); ok { - if fr, ok := cand["finishReason"].(string); ok && fr == "MALFORMED_FUNCTION_CALL" { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] MALFORMED_FUNCTION_CALL detected in forward stream") - if content, ok := cand["content"]; ok { - if b, err := json.Marshal(content); err == nil { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Malformed content: %s", string(b)) - } - } - } - } - } - } - - if firstTokenMs == nil { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - - cw.Fprintf("data: %s\n\n", payload) - continue - } - - cw.Fprintf("%s\n", line) - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - if cw.Disconnected() { - logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity gemini), returning collected usage") - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)") - sendErrorEvent("stream_timeout") - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") - - case <-keepaliveCh: - if cw.Disconnected() { - continue - } - if time.Since(lastDataAt) < keepaliveInterval { - continue - } - // SSE ping/keepalive：保持连接活跃防止 Cloudflare Tunnel 等代理断开 - if !cw.Fprintf(":\n\n") { - logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity gemini), continuing to drain upstream for billing") - continue - } - } - } -} - -// handleGeminiStreamToNonStreaming 读取上游流式响应，合并为非流式响应返回给客户端 -// Gemini 流式响应是增量的，需要累积所有 chunk 的内容 -func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) { - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.settingService.cfg.Gateway.MaxLineSize - } - scanBuf := getSSEScannerBuf64K() - scanner.Buffer(scanBuf[:0], maxLineSize) - - usage := &ClaudeUsage{} - var firstTokenMs *int - var last map[string]any - var lastWithParts map[string]any - var collectedImageParts []map[string]any // 收集所有包含图片的 parts - var collectedTextParts []string // 收集所有文本片段 - - type scanEvent struct { - line string - err error - } - - // 独立 goroutine 读取上游，避免读取阻塞影响超时处理 - events := make(chan scanEvent, 16) - done := make(chan struct{}) - sendEvent := func(ev scanEvent) bool { - select { - case events <- ev: - return true - case <-done: - return false - } - } - - var lastReadAt int64 - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - go func(scanBuf *sseScannerBuf64K) { - defer putSSEScannerBuf64K(scanBuf) - defer close(events) - for scanner.Scan() { - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { - return - } - } - if err := scanner.Err(); err != nil { - _ = sendEvent(scanEvent{err: err}) - } - }(scanBuf) - defer close(done) - - // 上游数据间隔超时保护（防止上游挂起长期占用连接） - streamInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second - } - var intervalTicker *time.Ticker - if streamInterval > 0 { - intervalTicker = time.NewTicker(streamInterval) - defer intervalTicker.Stop() - } - var intervalCh <-chan time.Time - if intervalTicker != nil { - intervalCh = intervalTicker.C - } - - for { - select { - case ev, ok := <-events: - if !ok { - // 流结束，返回收集的响应 - goto returnResponse - } - if ev.err != nil { - if errors.Is(ev.err, bufio.ErrTooLong) { - logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity non-stream): max_size=%d error=%v", maxLineSize, ev.err) - } - return nil, ev.err - } - - line := ev.line - trimmed := strings.TrimRight(line, "\r\n") - - if !strings.HasPrefix(trimmed, "data:") { - continue - } - - payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) - if payload == "" || payload == "[DONE]" { - continue - } - - // 解包 v1internal 响应 - inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) - if parseErr != nil { - continue - } - - var parsed map[string]any - if err := json.Unmarshal(inner, &parsed); err != nil { - continue - } - - // 记录首 token 时间 - if firstTokenMs == nil { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - - last = parsed - - // 提取 usage - if u := extractGeminiUsage(inner); u != nil { - usage = u - } - - // Check for MALFORMED_FUNCTION_CALL - if candidates, ok := parsed["candidates"].([]any); ok && len(candidates) > 0 { - if cand, ok := candidates[0].(map[string]any); ok { - if fr, ok := cand["finishReason"].(string); ok && fr == "MALFORMED_FUNCTION_CALL" { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] MALFORMED_FUNCTION_CALL detected in forward non-stream collect") - if content, ok := cand["content"]; ok { - if b, err := json.Marshal(content); err == nil { - logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Malformed content: %s", string(b)) - } - } - } - } - } - - // 保留最后一个有 parts 的响应 - if parts := extractGeminiParts(parsed); len(parts) > 0 { - lastWithParts = parsed - // 收集包含图片和文本的 parts - for _, part := range parts { - if inlineData, ok := part["inlineData"].(map[string]any); ok { - collectedImageParts = append(collectedImageParts, part) - _ = inlineData // 避免 unused 警告 - } - if text, ok := part["text"].(string); ok && text != "" { - collectedTextParts = append(collectedTextParts, text) - } - } - } - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity non-stream)") - return nil, fmt.Errorf("stream data interval timeout") - } - } - -returnResponse: - // 选择最后一个有效响应 - finalResponse := pickGeminiCollectResult(last, lastWithParts) - - // 处理空响应情况 — 触发同账号重试 + failover 切换账号 - if last == nil && lastWithParts == nil { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (gemini non-stream), triggering failover") - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), - RetryableOnSameAccount: true, - } - } - - // 如果收集到了图片 parts，需要合并到最终响应中 - if len(collectedImageParts) > 0 { - finalResponse = mergeImagePartsToResponse(finalResponse, collectedImageParts) - } - - // 如果收集到了文本，需要合并到最终响应中 - if len(collectedTextParts) > 0 { - finalResponse = mergeTextPartsToResponse(finalResponse, collectedTextParts) - } - - respBody, err := json.Marshal(finalResponse) - if err != nil { - return nil, fmt.Errorf("failed to marshal response: %w", err) - } - c.Data(http.StatusOK, "application/json", respBody) - - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil -} - -// getOrCreateGeminiParts 获取 Gemini 响应的 parts 结构，返回深拷贝和更新回调 -func getOrCreateGeminiParts(response map[string]any) (result map[string]any, existingParts []any, setParts func([]any)) { - // 深拷贝 response - result = make(map[string]any) - for k, v := range response { - result[k] = v - } - - // 获取或创建 candidates - candidates, ok := result["candidates"].([]any) - if !ok || len(candidates) == 0 { - candidates = []any{map[string]any{}} - } - - // 获取第一个 candidate - candidate, ok := candidates[0].(map[string]any) - if !ok { - candidate = make(map[string]any) - candidates[0] = candidate - } - - // 获取或创建 content - content, ok := candidate["content"].(map[string]any) - if !ok { - content = map[string]any{"role": "model"} - candidate["content"] = content - } - - // 获取现有 parts - existingParts, ok = content["parts"].([]any) - if !ok { - existingParts = []any{} - } - - // 返回更新回调 - setParts = func(newParts []any) { - content["parts"] = newParts - result["candidates"] = candidates - } - - return result, existingParts, setParts -} - -// mergeCollectedPartsToResponse 将收集的所有 parts 合并到 Gemini 响应中 -// 这个函数会合并所有类型的 parts：text、thinking、functionCall、inlineData 等 -// 保持原始顺序，只合并连续的普通 text parts -func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map[string]any) map[string]any { - if len(collectedParts) == 0 { - return response - } - - result, _, setParts := getOrCreateGeminiParts(response) - - // 合并策略： - // 1. 保持原始顺序 - // 2. 连续的普通 text parts 合并为一个 - // 3. thinking、functionCall、inlineData 等保持原样 - var mergedParts []any - var textBuffer strings.Builder - - flushTextBuffer := func() { - if textBuffer.Len() > 0 { - mergedParts = append(mergedParts, map[string]any{ - "text": textBuffer.String(), - }) - textBuffer.Reset() - } - } - - for _, part := range collectedParts { - // 检查是否是普通 text part - if text, ok := part["text"].(string); ok { - // 检查是否有 thought 标记 - if thought, _ := part["thought"].(bool); thought { - // thinking part，先刷新 text buffer，然后保留原样 - flushTextBuffer() - mergedParts = append(mergedParts, part) - } else { - // 普通 text，累积到 buffer - _, _ = textBuffer.WriteString(text) - } - } else { - // 非 text part（functionCall、inlineData 等），先刷新 text buffer，然后保留原样 - flushTextBuffer() - mergedParts = append(mergedParts, part) - } - } - - // 刷新剩余的 text - flushTextBuffer() - - setParts(mergedParts) - return result -} - -// mergeImagePartsToResponse 将收集到的图片 parts 合并到 Gemini 响应中 -func mergeImagePartsToResponse(response map[string]any, imageParts []map[string]any) map[string]any { - if len(imageParts) == 0 { - return response - } - - result, existingParts, setParts := getOrCreateGeminiParts(response) - - // 检查现有 parts 中是否已经有图片 - for _, p := range existingParts { - if pm, ok := p.(map[string]any); ok { - if _, hasInline := pm["inlineData"]; hasInline { - return result // 已有图片，不重复添加 - } - } - } - - // 添加收集到的图片 parts - for _, imgPart := range imageParts { - existingParts = append(existingParts, imgPart) - } - setParts(existingParts) - return result -} - -// mergeTextPartsToResponse 将收集到的文本合并到 Gemini 响应中 -func mergeTextPartsToResponse(response map[string]any, textParts []string) map[string]any { - if len(textParts) == 0 { - return response - } - - mergedText := strings.Join(textParts, "") - result, existingParts, setParts := getOrCreateGeminiParts(response) - - // 查找并更新第一个 text part，或创建新的 - newParts := make([]any, 0, len(existingParts)+1) - textUpdated := false - - for _, p := range existingParts { - pm, ok := p.(map[string]any) - if !ok { - newParts = append(newParts, p) - continue - } - if _, hasText := pm["text"]; hasText && !textUpdated { - // 用累积的文本替换 - newPart := make(map[string]any) - for k, v := range pm { - newPart[k] = v - } - newPart["text"] = mergedText - newParts = append(newParts, newPart) - textUpdated = true - } else { - newParts = append(newParts, pm) - } - } - - if !textUpdated { - newParts = append([]any{map[string]any{"text": mergedText}}, newParts...) - } - - setParts(newParts) - return result -} - -func (s *AntigravityGatewayService) writeClaudeError(c *gin.Context, status int, errType, message string) error { - MarkResponseCommitted(c) - c.JSON(status, gin.H{ - "type": "error", - "error": gin.H{"type": errType, "message": message}, - }) - return fmt.Errorf("%s", message) -} - -// WriteMappedClaudeError 导出版本，供 handler 层使用（如 fallback 错误处理） -func (s *AntigravityGatewayService) WriteMappedClaudeError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error { - return s.writeMappedClaudeError(c, account, upstreamStatus, upstreamRequestID, body) -} - -func (s *AntigravityGatewayService) writeMappedClaudeError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error { - MarkResponseCommitted(c) - upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - logBody, maxBytes := s.getLogConfig() - upstreamDetail := s.getUpstreamErrorDetail(body) - setOpsUpstreamError(c, upstreamStatus, upstreamMsg, upstreamDetail) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: upstreamStatus, - UpstreamRequestID: upstreamRequestID, - Kind: "http_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - // 记录上游错误详情便于排障（可选：由配置控制；不回显到客户端） - if logBody { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] upstream_error status=%d body=%s", upstreamStatus, truncateForLog(body, maxBytes)) - } - - // 检查错误透传规则 - if ptStatus, ptErrType, ptErrMsg, matched := applyErrorPassthroughRule( - c, account.Platform, upstreamStatus, body, - 0, "", "", - ); matched { - c.JSON(ptStatus, gin.H{ - "type": "error", - "error": gin.H{"type": ptErrType, "message": ptErrMsg}, - }) - if upstreamMsg == "" { - return fmt.Errorf("upstream error: %d", upstreamStatus) - } - return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg) - } - - var statusCode int - var errType, errMsg string - - switch upstreamStatus { - case 400: - statusCode = http.StatusBadRequest - errType = "invalid_request_error" - errMsg = getPassthroughOrDefault(upstreamMsg, "Invalid request") - case 401: - statusCode = http.StatusBadGateway - errType = "authentication_error" - errMsg = "Upstream authentication failed" - case 403: - statusCode = http.StatusBadGateway - errType = "permission_error" - errMsg = "Upstream access forbidden" - case 429: - statusCode = http.StatusTooManyRequests - errType = "rate_limit_error" - errMsg = "Upstream rate limit exceeded" - case 529: - statusCode = http.StatusServiceUnavailable - errType = "overloaded_error" - errMsg = "Upstream service overloaded" - default: - statusCode = http.StatusBadGateway - errType = "upstream_error" - errMsg = "Upstream request failed" - } - - c.JSON(statusCode, gin.H{ - "type": "error", - "error": gin.H{"type": errType, "message": errMsg}, - }) - if upstreamMsg == "" { - return fmt.Errorf("upstream error: %d", upstreamStatus) - } - return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg) -} - -func (s *AntigravityGatewayService) writeGoogleError(c *gin.Context, status int, message string) error { - MarkResponseCommitted(c) - statusStr := "UNKNOWN" - switch status { - case 400: - statusStr = "INVALID_ARGUMENT" - case 404: - statusStr = "NOT_FOUND" - case 429: - statusStr = "RESOURCE_EXHAUSTED" - case 500: - statusStr = "INTERNAL" - case 502, 503: - statusStr = "UNAVAILABLE" - } - - c.JSON(status, gin.H{ - "error": gin.H{ - "code": status, - "message": message, - "status": statusStr, - }, - }) - return fmt.Errorf("%s", message) -} - -// handleClaudeStreamToNonStreaming 收集上游流式响应，转换为 Claude 非流式格式返回 -// 用于处理客户端非流式请求但上游只支持流式的情况 -func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) { - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.settingService.cfg.Gateway.MaxLineSize - } - scanBuf := getSSEScannerBuf64K() - scanner.Buffer(scanBuf[:0], maxLineSize) - - var firstTokenMs *int - var last map[string]any - var lastWithParts map[string]any - var collectedParts []map[string]any // 收集所有 parts（包括 text、thinking、functionCall、inlineData 等） - - type scanEvent struct { - line string - err error - } - - // 独立 goroutine 读取上游，避免读取阻塞影响超时处理 - events := make(chan scanEvent, 16) - done := make(chan struct{}) - sendEvent := func(ev scanEvent) bool { - select { - case events <- ev: - return true - case <-done: - return false - } - } - - var lastReadAt int64 - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - go func(scanBuf *sseScannerBuf64K) { - defer putSSEScannerBuf64K(scanBuf) - defer close(events) - for scanner.Scan() { - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { - return - } - } - if err := scanner.Err(); err != nil { - _ = sendEvent(scanEvent{err: err}) - } - }(scanBuf) - defer close(done) - - // 上游数据间隔超时保护（防止上游挂起长期占用连接） - streamInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second - } - var intervalTicker *time.Ticker - if streamInterval > 0 { - intervalTicker = time.NewTicker(streamInterval) - defer intervalTicker.Stop() - } - var intervalCh <-chan time.Time - if intervalTicker != nil { - intervalCh = intervalTicker.C - } - - for { - select { - case ev, ok := <-events: - if !ok { - // 流结束，转换并返回响应 - goto returnResponse - } - if ev.err != nil { - if errors.Is(ev.err, bufio.ErrTooLong) { - logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity claude non-stream): max_size=%d error=%v", maxLineSize, ev.err) - } - return nil, ev.err - } - - line := ev.line - trimmed := strings.TrimRight(line, "\r\n") - - if !strings.HasPrefix(trimmed, "data:") { - continue - } - - payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) - if payload == "" || payload == "[DONE]" { - continue - } - - // 解包 v1internal 响应 - inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) - if parseErr != nil { - continue - } - - var parsed map[string]any - if err := json.Unmarshal(inner, &parsed); err != nil { - continue - } - - // 记录首 token 时间 - if firstTokenMs == nil { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - - last = parsed - - // 保留最后一个有 parts 的响应，并收集所有 parts - if parts := extractGeminiParts(parsed); len(parts) > 0 { - lastWithParts = parsed - - // 收集所有 parts（text、thinking、functionCall、inlineData 等） - collectedParts = append(collectedParts, parts...) - } - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity claude non-stream)") - return nil, fmt.Errorf("stream data interval timeout") - } - } - -returnResponse: - // 选择最后一个有效响应 - finalResponse := pickGeminiCollectResult(last, lastWithParts) - - // 处理空响应情况 — 触发同账号重试 + failover 切换账号 - if last == nil && lastWithParts == nil { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (claude non-stream), triggering failover") - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), - RetryableOnSameAccount: true, - } - } - - // 将收集的所有 parts 合并到最终响应中 - if len(collectedParts) > 0 { - finalResponse = mergeCollectedPartsToResponse(finalResponse, collectedParts) - } - - // 序列化为 JSON（Gemini 格式） - geminiBody, err := json.Marshal(finalResponse) - if err != nil { - return nil, fmt.Errorf("failed to marshal gemini response: %w", err) - } - - // 转换 Gemini 响应为 Claude 格式 - claudeResp, agUsage, err := antigravity.TransformGeminiToClaude(geminiBody, originalModel) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] transform_error error=%v body=%s", err, string(geminiBody)) - return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response") - } - - c.Data(http.StatusOK, "application/json", claudeResp) - - // 转换为 service.ClaudeUsage - usage := &ClaudeUsage{ - InputTokens: agUsage.InputTokens, - OutputTokens: agUsage.OutputTokens, - CacheCreationInputTokens: agUsage.CacheCreationInputTokens, - CacheReadInputTokens: agUsage.CacheReadInputTokens, - } - - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil -} - -// handleClaudeStreamingResponse 处理 Claude 流式响应（Gemini SSE → Claude SSE 转换） -func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) { - c.Header("Content-Type", "text/event-stream") - c.Header("Cache-Control", "no-cache") - c.Header("Connection", "keep-alive") - c.Header("X-Accel-Buffering", "no") - c.Status(http.StatusOK) - - flusher, ok := c.Writer.(http.Flusher) - if !ok { - return nil, errors.New("streaming not supported") - } - - processor := antigravity.NewStreamingProcessor(originalModel) - var firstTokenMs *int - // 使用 Scanner 并限制单行大小，避免 ReadString 无上限导致 OOM - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.settingService.cfg.Gateway.MaxLineSize - } - scanBuf := getSSEScannerBuf64K() - scanner.Buffer(scanBuf[:0], maxLineSize) - - // 辅助函数：转换 antigravity.ClaudeUsage 到 service.ClaudeUsage - convertUsage := func(agUsage *antigravity.ClaudeUsage) *ClaudeUsage { - if agUsage == nil { - return &ClaudeUsage{} - } - return &ClaudeUsage{ - InputTokens: agUsage.InputTokens, - OutputTokens: agUsage.OutputTokens, - CacheCreationInputTokens: agUsage.CacheCreationInputTokens, - CacheReadInputTokens: agUsage.CacheReadInputTokens, - } - } - - type scanEvent struct { - line string - err error - } - // 独立 goroutine 读取上游，避免读取阻塞影响超时处理 - events := make(chan scanEvent, 16) - done := make(chan struct{}) - sendEvent := func(ev scanEvent) bool { - select { - case events <- ev: - return true - case <-done: - return false - } - } - var lastReadAt int64 - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - go func(scanBuf *sseScannerBuf64K) { - defer putSSEScannerBuf64K(scanBuf) - defer close(events) - for scanner.Scan() { - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { - return - } - } - if err := scanner.Err(); err != nil { - _ = sendEvent(scanEvent{err: err}) - } - }(scanBuf) - defer close(done) - - streamInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second - } - var intervalTicker *time.Ticker - if streamInterval > 0 { - intervalTicker = time.NewTicker(streamInterval) - defer intervalTicker.Stop() - } - var intervalCh <-chan time.Time - if intervalTicker != nil { - intervalCh = intervalTicker.C - } - - // 下游 keepalive：防止代理/Cloudflare Tunnel 因连接空闲而断开 - keepaliveInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { - keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second - } - var keepaliveTicker *time.Ticker - if keepaliveInterval > 0 { - keepaliveTicker = time.NewTicker(keepaliveInterval) - defer keepaliveTicker.Stop() - } - var keepaliveCh <-chan time.Time - if keepaliveTicker != nil { - keepaliveCh = keepaliveTicker.C - } - lastDataAt := time.Now() - - cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity claude") - - // 仅发送一次错误事件，避免多次写入导致协议混乱 - errorEventSent := false - sendErrorEvent := func(reason string) { - if errorEventSent || cw.Disconnected() { - return - } - errorEventSent = true - _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"%s\"}\n\n", reason) - flusher.Flush() - } - - // finishUsage 是获取 processor 最终 usage 的辅助函数 - finishUsage := func() *ClaudeUsage { - _, agUsage := processor.Finish() - return convertUsage(agUsage) - } - - for { - select { - case ev, ok := <-events: - if !ok { - // 上游完成，发送结束事件 - finalEvents, agUsage := processor.Finish() - if len(finalEvents) > 0 { - cw.Write(finalEvents) - } else if !processor.MessageStartSent() && !cw.Disconnected() { - // 整个流未收到任何可解析的上游数据（全部 SSE 行均无法被 JSON 解析）， - // 触发 failover 在同账号重试，避免向客户端发出缺少 message_start 的残缺流 - logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Claude-Stream] empty stream response (no valid events parsed), triggering failover") - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), - RetryableOnSameAccount: true, - } - } - return &antigravityStreamResult{usage: convertUsage(agUsage), firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}, nil - } - if ev.err != nil { - if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity claude"); handled { - return &antigravityStreamResult{usage: finishUsage(), firstTokenMs: firstTokenMs, clientDisconnect: disconnect}, nil - } - if errors.Is(ev.err, bufio.ErrTooLong) { - logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err) - sendErrorEvent("response_too_large") - return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, ev.err - } - sendErrorEvent("stream_read_error") - return nil, fmt.Errorf("stream read error: %w", ev.err) - } - - lastDataAt = time.Now() - - // 处理 SSE 行，转换为 Claude 格式 - claudeEvents := processor.ProcessLine(strings.TrimRight(ev.line, "\r\n")) - if len(claudeEvents) > 0 { - if firstTokenMs == nil { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - cw.Write(claudeEvents) - } - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - if cw.Disconnected() { - logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity claude), returning collected usage") - return &antigravityStreamResult{usage: finishUsage(), firstTokenMs: firstTokenMs, clientDisconnect: true}, nil - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)") - sendErrorEvent("stream_timeout") - return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") - - case <-keepaliveCh: - if cw.Disconnected() { - continue - } - if time.Since(lastDataAt) < keepaliveInterval { - continue - } - // SSE ping 事件：Anthropic 原生格式，客户端会正确处理， - // 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开 - if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") { - logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity claude), continuing to drain upstream for billing") - continue - } - } - } -} - -func (s *AntigravityGatewayService) extractImageInputSize(body []byte) string { - var req antigravity.GeminiRequest - if err := json.Unmarshal(body, &req); err != nil { - return "" - } - - if req.GenerationConfig != nil && req.GenerationConfig.ImageConfig != nil { - return strings.TrimSpace(req.GenerationConfig.ImageConfig.ImageSize) - } - - return "" -} - -// isImageGenerationModel 判断模型是否为图片生成模型 -// 支持的模型：gemini-3.1-flash-image, gemini-3-pro-image, gemini-2.5-flash-image 等 -func isImageGenerationModel(model string) bool { - modelLower := strings.ToLower(model) - // 移除 models/ 前缀 - modelLower = strings.TrimPrefix(modelLower, "models/") - - // 精确匹配或前缀匹配 - return modelLower == "gemini-3.1-flash-image" || - modelLower == "gemini-3.1-flash-image-preview" || - strings.HasPrefix(modelLower, "gemini-3.1-flash-image-") || - modelLower == "gemini-3-pro-image" || - modelLower == "gemini-3-pro-image-preview" || - strings.HasPrefix(modelLower, "gemini-3-pro-image-") || - modelLower == "gemini-2.5-flash-image" || - modelLower == "gemini-2.5-flash-image-preview" || - strings.HasPrefix(modelLower, "gemini-2.5-flash-image-") -} - -// cleanGeminiRequest 清理 Gemini 请求体中的 Schema -func cleanGeminiRequest(body []byte) ([]byte, error) { - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return nil, err - } - - modified := false - - // 1. 清理 Tools - if tools, ok := payload["tools"].([]any); ok && len(tools) > 0 { - for _, t := range tools { - toolMap, ok := t.(map[string]any) - if !ok { - continue - } - - // function_declarations (snake_case) or functionDeclarations (camelCase) - var funcs []any - if f, ok := toolMap["functionDeclarations"].([]any); ok { - funcs = f - } else if f, ok := toolMap["function_declarations"].([]any); ok { - funcs = f - } - - if len(funcs) == 0 { - continue - } - - for _, f := range funcs { - funcMap, ok := f.(map[string]any) - if !ok { - continue - } - - if params, ok := funcMap["parameters"].(map[string]any); ok { - antigravity.DeepCleanUndefined(params) - cleaned := antigravity.CleanJSONSchema(params) - funcMap["parameters"] = cleaned - modified = true - } - } - } - } - - if !modified { - return body, nil - } - - return json.Marshal(payload) -} - -// filterEmptyPartsFromGeminiRequest 过滤掉 parts 为空的消息 -// Gemini API 不接受空 parts，需要在请求前过滤 -func filterEmptyPartsFromGeminiRequest(body []byte) ([]byte, error) { - var payload map[string]any - if err := json.Unmarshal(body, &payload); err != nil { - return nil, err - } - - contents, ok := payload["contents"].([]any) - if !ok || len(contents) == 0 { - return body, nil - } - - filtered := make([]any, 0, len(contents)) - modified := false - - for _, c := range contents { - contentMap, ok := c.(map[string]any) - if !ok { - filtered = append(filtered, c) - continue - } - - parts, hasParts := contentMap["parts"] - if !hasParts { - filtered = append(filtered, c) - continue - } - - partsSlice, ok := parts.([]any) - if !ok { - filtered = append(filtered, c) - continue - } - - // 跳过 parts 为空数组的消息 - if len(partsSlice) == 0 { - modified = true - continue - } - - filtered = append(filtered, c) - } - - if !modified { - return body, nil - } - - payload["contents"] = filtered - return json.Marshal(payload) -} - -// ForwardUpstream 使用 base_url + /v1/messages + 双 header 认证透传上游 Claude 请求 -func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) { - startTime := time.Now() - sessionID := getSessionID(c) - prefix := logPrefix(sessionID, account.Name) - - // 获取上游配置 - baseURL := strings.TrimSpace(account.GetCredential("base_url")) - apiKey := strings.TrimSpace(account.GetCredential("api_key")) - if baseURL == "" || apiKey == "" { - return nil, fmt.Errorf("upstream account missing base_url or api_key") - } - baseURL = strings.TrimSuffix(baseURL, "/") - - // 解析请求获取模型信息 - var claudeReq antigravity.ClaudeRequest - if err := json.Unmarshal(body, &claudeReq); err != nil { - return nil, fmt.Errorf("parse claude request: %w", err) - } - if strings.TrimSpace(claudeReq.Model) == "" { - return nil, fmt.Errorf("missing model") - } - originalModel := claudeReq.Model - - // 构建上游请求 URL - upstreamURL := baseURL + "/v1/messages" - - // 能力维度 sanitize：Anthropic-compatible 上游透传路径也需要保证 body↔beta header - // 对称。客户端 anthropic-beta header 不含 context-management-2025-06-27 但 body 带 - // context_management 时 strip，与 Anthropic 直连 / Bedrock / Vertex 路径保持一致。 - clientBeta := c.GetHeader("anthropic-beta") - if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, clientBeta); changed { - body = sanitized - } - - // 创建请求 - req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(body)) - if err != nil { - return nil, fmt.Errorf("create upstream request: %w", err) - } - - // 设置请求头 - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+apiKey) - req.Header.Set("x-api-key", apiKey) // Claude API 兼容 - - // 透传 Claude 相关 headers - if v := c.GetHeader("anthropic-version"); v != "" { - req.Header.Set("anthropic-version", v) - } - if v := clientBeta; v != "" { - req.Header.Set("anthropic-beta", v) - } - - // 代理 URL - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - // 发送请求 - resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) - if err != nil { - logger.LegacyPrintf("service.antigravity_gateway", "%s upstream request failed: %v", prefix, err) - return nil, fmt.Errorf("upstream request failed: %w", err) - } - defer func() { _ = resp.Body.Close() }() - - // 处理错误响应 - if resp.StatusCode >= 400 { - respBody := s.readUpstreamErrorBody(resp) - - // 429 错误时标记账号限流 - if resp.StatusCode == http.StatusTooManyRequests { - s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", false) - } - - // 透传上游错误 - c.Header("Content-Type", resp.Header.Get("Content-Type")) - c.Status(resp.StatusCode) - _, _ = c.Writer.Write(respBody) - - return &ForwardResult{ - Model: originalModel, - }, nil - } - - // 处理成功响应（流式/非流式） - var usage *ClaudeUsage - var firstTokenMs *int - var clientDisconnect bool - - if claudeReq.Stream { - // 流式响应：透传 - c.Header("Content-Type", "text/event-stream") - c.Header("Cache-Control", "no-cache") - c.Header("Connection", "keep-alive") - c.Header("X-Accel-Buffering", "no") - c.Status(http.StatusOK) - - streamRes := s.streamUpstreamResponse(c, resp, startTime) - usage = streamRes.usage - firstTokenMs = streamRes.firstTokenMs - clientDisconnect = streamRes.clientDisconnect - } else { - // 非流式响应：直接透传 - respBody, err := io.ReadAll(resp.Body) - if err != nil { - return nil, fmt.Errorf("read upstream response: %w", err) - } - - // 提取 usage - usage = s.extractClaudeUsage(respBody) - - c.Header("Content-Type", resp.Header.Get("Content-Type")) - c.Status(http.StatusOK) - _, _ = c.Writer.Write(respBody) - } - - // 构建计费结果 - duration := time.Since(startTime) - logger.LegacyPrintf("service.antigravity_gateway", "%s status=success duration_ms=%d", prefix, duration.Milliseconds()) - - return &ForwardResult{ - Model: originalModel, - Stream: claudeReq.Stream, - Duration: duration, - FirstTokenMs: firstTokenMs, - ClientDisconnect: clientDisconnect, - Usage: ClaudeUsage{ - InputTokens: usage.InputTokens, - OutputTokens: usage.OutputTokens, - CacheReadInputTokens: usage.CacheReadInputTokens, - CacheCreationInputTokens: usage.CacheCreationInputTokens, - }, - }, nil -} - -// streamUpstreamResponse 透传上游 SSE 流并提取 Claude usage -func (s *AntigravityGatewayService) streamUpstreamResponse(c *gin.Context, resp *http.Response, startTime time.Time) *antigravityStreamResult { - usage := &ClaudeUsage{} - var firstTokenMs *int - - scanner := bufio.NewScanner(resp.Body) - maxLineSize := defaultMaxLineSize - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.settingService.cfg.Gateway.MaxLineSize - } - scanner.Buffer(make([]byte, 64*1024), maxLineSize) - - type scanEvent struct { - line string - err error - } - events := make(chan scanEvent, 16) - done := make(chan struct{}) - sendEvent := func(ev scanEvent) bool { - select { - case events <- ev: - return true - case <-done: - return false - } - } - var lastReadAt int64 - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - go func() { - defer close(events) - for scanner.Scan() { - atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) - if !sendEvent(scanEvent{line: scanner.Text()}) { - return - } - } - if err := scanner.Err(); err != nil { - _ = sendEvent(scanEvent{err: err}) - } - }() - defer close(done) - - streamInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second - } - var intervalTicker *time.Ticker - if streamInterval > 0 { - intervalTicker = time.NewTicker(streamInterval) - defer intervalTicker.Stop() - } - var intervalCh <-chan time.Time - if intervalTicker != nil { - intervalCh = intervalTicker.C - } - - // 下游 keepalive：防止代理/Cloudflare Tunnel 因连接空闲而断开 - keepaliveInterval := time.Duration(0) - if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { - keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second - } - var keepaliveTicker *time.Ticker - if keepaliveInterval > 0 { - keepaliveTicker = time.NewTicker(keepaliveInterval) - defer keepaliveTicker.Stop() - } - var keepaliveCh <-chan time.Time - if keepaliveTicker != nil { - keepaliveCh = keepaliveTicker.C - } - lastDataAt := time.Now() - - flusher, _ := c.Writer.(http.Flusher) - cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity upstream") - - for { - select { - case ev, ok := <-events: - if !ok { - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()} - } - if ev.err != nil { - if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity upstream"); handled { - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: disconnect} - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream read error (antigravity upstream): %v", ev.err) - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs} - } - - lastDataAt = time.Now() - - line := ev.line - - // 记录首 token 时间 - if firstTokenMs == nil && len(line) > 0 { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - - // 尝试从 message_delta 或 message_stop 事件提取 usage - s.extractSSEUsage(line, usage) - - // 透传行 - cw.Fprintf("%s\n", line) - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - if cw.Disconnected() { - logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity upstream), returning collected usage") - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true} - } - logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity upstream)") - return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs} - - case <-keepaliveCh: - if cw.Disconnected() { - continue - } - if time.Since(lastDataAt) < keepaliveInterval { - continue - } - // SSE ping 事件：Anthropic 原生格式，客户端会正确处理， - // 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开 - if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") { - logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity upstream), continuing to drain upstream for billing") - continue - } - } - } -} - -// extractSSEUsage 从 SSE data 行中提取 Claude usage（用于流式透传场景） -// -// Anthropic streaming 的 usage 字段分布在两类事件中： -// - message_start：嵌套在 event.message.usage（input_tokens、cache_creation_input_tokens、 -// cache_read_input_tokens 等输入侧字段） -// - message_delta：位于顶层 event.usage（流结束时的最终 output_tokens） -// -// 仅读取顶层 event.usage 会漏掉 message_start 的输入侧字段，导致流式透传请求落库的 -// usage_logs 记录 input_tokens=0。 -func (s *AntigravityGatewayService) extractSSEUsage(line string, usage *ClaudeUsage) { - if !strings.HasPrefix(line, "data: ") { - return - } - dataStr := strings.TrimPrefix(line, "data: ") - var event map[string]any - if json.Unmarshal([]byte(dataStr), &event) != nil { - return - } - var u map[string]any - if eventType, _ := event["type"].(string); eventType == "message_start" { - if msg, ok := event["message"].(map[string]any); ok { - u, _ = msg["usage"].(map[string]any) - } - } else { - u, _ = event["usage"].(map[string]any) - } - if u == nil { - return - } - if v, ok := u["input_tokens"].(float64); ok && int(v) > 0 { - usage.InputTokens = int(v) - } - if v, ok := u["output_tokens"].(float64); ok && int(v) > 0 { - usage.OutputTokens = int(v) - } - if v, ok := u["cache_read_input_tokens"].(float64); ok && int(v) > 0 { - usage.CacheReadInputTokens = int(v) - } - if v, ok := u["cache_creation_input_tokens"].(float64); ok && int(v) > 0 { - usage.CacheCreationInputTokens = int(v) - } - // 解析嵌套的 cache_creation 对象中的 5m/1h 明细 - if cc, ok := u["cache_creation"].(map[string]any); ok { - if v, ok := cc["ephemeral_5m_input_tokens"].(float64); ok { - usage.CacheCreation5mTokens = int(v) - } - if v, ok := cc["ephemeral_1h_input_tokens"].(float64); ok { - usage.CacheCreation1hTokens = int(v) - } - } -} - -// extractClaudeUsage 从非流式 Claude 响应提取 usage -func (s *AntigravityGatewayService) extractClaudeUsage(body []byte) *ClaudeUsage { - usage := &ClaudeUsage{} - var resp map[string]any - if json.Unmarshal(body, &resp) != nil { - return usage - } - if u, ok := resp["usage"].(map[string]any); ok { - if v, ok := u["input_tokens"].(float64); ok { - usage.InputTokens = int(v) - } - if v, ok := u["output_tokens"].(float64); ok { - usage.OutputTokens = int(v) - } - if v, ok := u["cache_read_input_tokens"].(float64); ok { - usage.CacheReadInputTokens = int(v) - } - if v, ok := u["cache_creation_input_tokens"].(float64); ok { - usage.CacheCreationInputTokens = int(v) - } - // 解析嵌套的 cache_creation 对象中的 5m/1h 明细 - if cc, ok := u["cache_creation"].(map[string]any); ok { - if v, ok := cc["ephemeral_5m_input_tokens"].(float64); ok { - usage.CacheCreation5mTokens = int(v) - } - if v, ok := cc["ephemeral_1h_input_tokens"].(float64); ok { - usage.CacheCreation1hTokens = int(v) - } - } - } - return usage -} diff --git a/backend/internal/service/antigravity_gateway_streaming.go b/backend/internal/service/antigravity_gateway_streaming.go new file mode 100644 index 0000000000..1a6c59f617 --- /dev/null +++ b/backend/internal/service/antigravity_gateway_streaming.go @@ -0,0 +1,1150 @@ +package service + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "sync/atomic" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" +) + +type antigravityStreamResult struct { + usage *ClaudeUsage + firstTokenMs *int + clientDisconnect bool // 客户端是否在流式传输过程中断开 +} + +// antigravityClientWriter 封装流式响应的客户端写入，自动检测断开并标记。 +// 断开后所有写入操作变为 no-op，调用方通过 Disconnected() 判断是否继续 drain 上游。 +type antigravityClientWriter struct { + w gin.ResponseWriter + flusher http.Flusher + disconnected bool + prefix string // 日志前缀，标识来源方法 +} + +func newAntigravityClientWriter(w gin.ResponseWriter, flusher http.Flusher, prefix string) *antigravityClientWriter { + return &antigravityClientWriter{w: w, flusher: flusher, prefix: prefix} +} + +// Write 写入数据到客户端，写入失败时标记断开并返回 false +func (cw *antigravityClientWriter) Write(p []byte) bool { + if cw.disconnected { + return false + } + if _, err := cw.w.Write(p); err != nil { + cw.markDisconnected() + return false + } + cw.flusher.Flush() + return true +} + +// Fprintf 格式化写入数据到客户端，写入失败时标记断开并返回 false +func (cw *antigravityClientWriter) Fprintf(format string, args ...any) bool { + if cw.disconnected { + return false + } + if _, err := fmt.Fprintf(cw.w, format, args...); err != nil { + cw.markDisconnected() + return false + } + cw.flusher.Flush() + return true +} + +func (cw *antigravityClientWriter) Disconnected() bool { return cw.disconnected } + +func (cw *antigravityClientWriter) markDisconnected() { + cw.disconnected = true + logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during streaming (%s), continuing to drain upstream for billing", cw.prefix) +} + +// handleStreamReadError 处理上游读取错误的通用逻辑。 +// 返回 (clientDisconnect, handled)：handled=true 表示错误已处理，调用方应返回已收集的 usage。 +func handleStreamReadError(err error, clientDisconnected bool, prefix string) (disconnect bool, handled bool) { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + logger.LegacyPrintf("service.antigravity_gateway", "Context canceled during streaming (%s), returning collected usage", prefix) + return true, true + } + if clientDisconnected { + logger.LegacyPrintf("service.antigravity_gateway", "Upstream read error after client disconnect (%s): %v, returning collected usage", prefix, err) + return true, true + } + return false, false +} + +func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) { + c.Status(resp.StatusCode) + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + + contentType := resp.Header.Get("Content-Type") + if contentType == "" { + contentType = "text/event-stream; charset=utf-8" + } + c.Header("Content-Type", contentType) + + flusher, ok := c.Writer.(http.Flusher) + if !ok { + return nil, errors.New("streaming not supported") + } + + // 使用 Scanner 并限制单行大小，避免 ReadString 无上限导致 OOM + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.settingService.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + usage := &ClaudeUsage{} + var firstTokenMs *int + + type scanEvent struct { + line string + err error + } + // 独立 goroutine 读取上游，避免读取阻塞影响超时处理 + events := make(chan scanEvent, 16) + done := make(chan struct{}) + sendEvent := func(ev scanEvent) bool { + select { + case events <- ev: + return true + case <-done: + return false + } + } + var lastReadAt int64 + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + go func(scanBuf *sseScannerBuf64K) { + defer putSSEScannerBuf64K(scanBuf) + defer close(events) + for scanner.Scan() { + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + if !sendEvent(scanEvent{line: scanner.Text()}) { + return + } + } + if err := scanner.Err(); err != nil { + _ = sendEvent(scanEvent{err: err}) + } + }(scanBuf) + defer close(done) + + // 上游数据间隔超时保护（防止上游挂起长期占用连接） + streamInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var intervalTicker *time.Ticker + if streamInterval > 0 { + intervalTicker = time.NewTicker(streamInterval) + defer intervalTicker.Stop() + } + var intervalCh <-chan time.Time + if intervalTicker != nil { + intervalCh = intervalTicker.C + } + + // 下游 keepalive：防止代理/Cloudflare Tunnel 因连接空闲而断开 + keepaliveInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { + keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second + } + var keepaliveTicker *time.Ticker + if keepaliveInterval > 0 { + keepaliveTicker = time.NewTicker(keepaliveInterval) + defer keepaliveTicker.Stop() + } + var keepaliveCh <-chan time.Time + if keepaliveTicker != nil { + keepaliveCh = keepaliveTicker.C + } + lastDataAt := time.Now() + + cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity gemini") + + // 仅发送一次错误事件，避免多次写入导致协议混乱 + errorEventSent := false + sendErrorEvent := func(reason string) { + if errorEventSent || cw.Disconnected() { + return + } + errorEventSent = true + _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"%s\"}\n\n", reason) + flusher.Flush() + } + + for { + select { + case ev, ok := <-events: + if !ok { + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}, nil + } + if ev.err != nil { + if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity gemini"); handled { + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: disconnect}, nil + } + if errors.Is(ev.err, bufio.ErrTooLong) { + logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err) + sendErrorEvent("response_too_large") + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, ev.err + } + sendErrorEvent("stream_read_error") + return nil, ev.err + } + + lastDataAt = time.Now() + + line := ev.line + trimmed := strings.TrimRight(line, "\r\n") + if strings.HasPrefix(trimmed, "data:") { + payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + if payload == "" || payload == "[DONE]" { + cw.Fprintf("%s\n", line) + continue + } + + // 解包 v1internal 响应 + inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) + if parseErr == nil && inner != nil { + payload = string(inner) + } + + // 解析 usage + if u := extractGeminiUsage(inner); u != nil { + usage = u + } + var parsed map[string]any + if json.Unmarshal(inner, &parsed) == nil { + // Check for MALFORMED_FUNCTION_CALL + if candidates, ok := parsed["candidates"].([]any); ok && len(candidates) > 0 { + if cand, ok := candidates[0].(map[string]any); ok { + if fr, ok := cand["finishReason"].(string); ok && fr == "MALFORMED_FUNCTION_CALL" { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] MALFORMED_FUNCTION_CALL detected in forward stream") + if content, ok := cand["content"]; ok { + if b, err := json.Marshal(content); err == nil { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Malformed content: %s", string(b)) + } + } + } + } + } + } + + if firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + + cw.Fprintf("data: %s\n\n", payload) + continue + } + + cw.Fprintf("%s\n", line) + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + if cw.Disconnected() { + logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity gemini), returning collected usage") + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)") + sendErrorEvent("stream_timeout") + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") + + case <-keepaliveCh: + if cw.Disconnected() { + continue + } + if time.Since(lastDataAt) < keepaliveInterval { + continue + } + // SSE ping/keepalive：保持连接活跃防止 Cloudflare Tunnel 等代理断开 + if !cw.Fprintf(":\n\n") { + logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity gemini), continuing to drain upstream for billing") + continue + } + } + } +} + +// handleGeminiStreamToNonStreaming 读取上游流式响应，合并为非流式响应返回给客户端 +// Gemini 流式响应是增量的，需要累积所有 chunk 的内容 +func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) { + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.settingService.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + + usage := &ClaudeUsage{} + var firstTokenMs *int + var last map[string]any + var lastWithParts map[string]any + var collectedImageParts []map[string]any // 收集所有包含图片的 parts + var collectedTextParts []string // 收集所有文本片段 + + type scanEvent struct { + line string + err error + } + + // 独立 goroutine 读取上游，避免读取阻塞影响超时处理 + events := make(chan scanEvent, 16) + done := make(chan struct{}) + sendEvent := func(ev scanEvent) bool { + select { + case events <- ev: + return true + case <-done: + return false + } + } + + var lastReadAt int64 + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + go func(scanBuf *sseScannerBuf64K) { + defer putSSEScannerBuf64K(scanBuf) + defer close(events) + for scanner.Scan() { + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + if !sendEvent(scanEvent{line: scanner.Text()}) { + return + } + } + if err := scanner.Err(); err != nil { + _ = sendEvent(scanEvent{err: err}) + } + }(scanBuf) + defer close(done) + + // 上游数据间隔超时保护（防止上游挂起长期占用连接） + streamInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var intervalTicker *time.Ticker + if streamInterval > 0 { + intervalTicker = time.NewTicker(streamInterval) + defer intervalTicker.Stop() + } + var intervalCh <-chan time.Time + if intervalTicker != nil { + intervalCh = intervalTicker.C + } + + for { + select { + case ev, ok := <-events: + if !ok { + // 流结束，返回收集的响应 + goto returnResponse + } + if ev.err != nil { + if errors.Is(ev.err, bufio.ErrTooLong) { + logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity non-stream): max_size=%d error=%v", maxLineSize, ev.err) + } + return nil, ev.err + } + + line := ev.line + trimmed := strings.TrimRight(line, "\r\n") + + if !strings.HasPrefix(trimmed, "data:") { + continue + } + + payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + if payload == "" || payload == "[DONE]" { + continue + } + + // 解包 v1internal 响应 + inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) + if parseErr != nil { + continue + } + + var parsed map[string]any + if err := json.Unmarshal(inner, &parsed); err != nil { + continue + } + + // 记录首 token 时间 + if firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + + last = parsed + + // 提取 usage + if u := extractGeminiUsage(inner); u != nil { + usage = u + } + + // Check for MALFORMED_FUNCTION_CALL + if candidates, ok := parsed["candidates"].([]any); ok && len(candidates) > 0 { + if cand, ok := candidates[0].(map[string]any); ok { + if fr, ok := cand["finishReason"].(string); ok && fr == "MALFORMED_FUNCTION_CALL" { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] MALFORMED_FUNCTION_CALL detected in forward non-stream collect") + if content, ok := cand["content"]; ok { + if b, err := json.Marshal(content); err == nil { + logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] Malformed content: %s", string(b)) + } + } + } + } + } + + // 保留最后一个有 parts 的响应 + if parts := extractGeminiParts(parsed); len(parts) > 0 { + lastWithParts = parsed + // 收集包含图片和文本的 parts + for _, part := range parts { + if inlineData, ok := part["inlineData"].(map[string]any); ok { + collectedImageParts = append(collectedImageParts, part) + _ = inlineData // 避免 unused 警告 + } + if text, ok := part["text"].(string); ok && text != "" { + collectedTextParts = append(collectedTextParts, text) + } + } + } + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity non-stream)") + return nil, fmt.Errorf("stream data interval timeout") + } + } + +returnResponse: + // 选择最后一个有效响应 + finalResponse := pickGeminiCollectResult(last, lastWithParts) + + // 处理空响应情况 — 触发同账号重试 + failover 切换账号 + if last == nil && lastWithParts == nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (gemini non-stream), triggering failover") + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), + RetryableOnSameAccount: true, + } + } + + // 如果收集到了图片 parts，需要合并到最终响应中 + if len(collectedImageParts) > 0 { + finalResponse = mergeImagePartsToResponse(finalResponse, collectedImageParts) + } + + // 如果收集到了文本，需要合并到最终响应中 + if len(collectedTextParts) > 0 { + finalResponse = mergeTextPartsToResponse(finalResponse, collectedTextParts) + } + + respBody, err := json.Marshal(finalResponse) + if err != nil { + return nil, fmt.Errorf("failed to marshal response: %w", err) + } + c.Data(http.StatusOK, "application/json", respBody) + + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil +} + +// getOrCreateGeminiParts 获取 Gemini 响应的 parts 结构，返回深拷贝和更新回调 +func getOrCreateGeminiParts(response map[string]any) (result map[string]any, existingParts []any, setParts func([]any)) { + // 深拷贝 response + result = make(map[string]any) + for k, v := range response { + result[k] = v + } + + // 获取或创建 candidates + candidates, ok := result["candidates"].([]any) + if !ok || len(candidates) == 0 { + candidates = []any{map[string]any{}} + } + + // 获取第一个 candidate + candidate, ok := candidates[0].(map[string]any) + if !ok { + candidate = make(map[string]any) + candidates[0] = candidate + } + + // 获取或创建 content + content, ok := candidate["content"].(map[string]any) + if !ok { + content = map[string]any{"role": "model"} + candidate["content"] = content + } + + // 获取现有 parts + existingParts, ok = content["parts"].([]any) + if !ok { + existingParts = []any{} + } + + // 返回更新回调 + setParts = func(newParts []any) { + content["parts"] = newParts + result["candidates"] = candidates + } + + return result, existingParts, setParts +} + +// mergeCollectedPartsToResponse 将收集的所有 parts 合并到 Gemini 响应中 +// 这个函数会合并所有类型的 parts：text、thinking、functionCall、inlineData 等 +// 保持原始顺序，只合并连续的普通 text parts +func mergeCollectedPartsToResponse(response map[string]any, collectedParts []map[string]any) map[string]any { + if len(collectedParts) == 0 { + return response + } + + result, _, setParts := getOrCreateGeminiParts(response) + + // 合并策略： + // 1. 保持原始顺序 + // 2. 连续的普通 text parts 合并为一个 + // 3. thinking、functionCall、inlineData 等保持原样 + var mergedParts []any + var textBuffer strings.Builder + + flushTextBuffer := func() { + if textBuffer.Len() > 0 { + mergedParts = append(mergedParts, map[string]any{ + "text": textBuffer.String(), + }) + textBuffer.Reset() + } + } + + for _, part := range collectedParts { + // 检查是否是普通 text part + if text, ok := part["text"].(string); ok { + // 检查是否有 thought 标记 + if thought, _ := part["thought"].(bool); thought { + // thinking part，先刷新 text buffer，然后保留原样 + flushTextBuffer() + mergedParts = append(mergedParts, part) + } else { + // 普通 text，累积到 buffer + _, _ = textBuffer.WriteString(text) + } + } else { + // 非 text part（functionCall、inlineData 等），先刷新 text buffer，然后保留原样 + flushTextBuffer() + mergedParts = append(mergedParts, part) + } + } + + // 刷新剩余的 text + flushTextBuffer() + + setParts(mergedParts) + return result +} + +// mergeImagePartsToResponse 将收集到的图片 parts 合并到 Gemini 响应中 +func mergeImagePartsToResponse(response map[string]any, imageParts []map[string]any) map[string]any { + if len(imageParts) == 0 { + return response + } + + result, existingParts, setParts := getOrCreateGeminiParts(response) + + // 检查现有 parts 中是否已经有图片 + for _, p := range existingParts { + if pm, ok := p.(map[string]any); ok { + if _, hasInline := pm["inlineData"]; hasInline { + return result // 已有图片，不重复添加 + } + } + } + + // 添加收集到的图片 parts + for _, imgPart := range imageParts { + existingParts = append(existingParts, imgPart) + } + setParts(existingParts) + return result +} + +// mergeTextPartsToResponse 将收集到的文本合并到 Gemini 响应中 +func mergeTextPartsToResponse(response map[string]any, textParts []string) map[string]any { + if len(textParts) == 0 { + return response + } + + mergedText := strings.Join(textParts, "") + result, existingParts, setParts := getOrCreateGeminiParts(response) + + // 查找并更新第一个 text part，或创建新的 + newParts := make([]any, 0, len(existingParts)+1) + textUpdated := false + + for _, p := range existingParts { + pm, ok := p.(map[string]any) + if !ok { + newParts = append(newParts, p) + continue + } + if _, hasText := pm["text"]; hasText && !textUpdated { + // 用累积的文本替换 + newPart := make(map[string]any) + for k, v := range pm { + newPart[k] = v + } + newPart["text"] = mergedText + newParts = append(newParts, newPart) + textUpdated = true + } else { + newParts = append(newParts, pm) + } + } + + if !textUpdated { + newParts = append([]any{map[string]any{"text": mergedText}}, newParts...) + } + + setParts(newParts) + return result +} + +func (s *AntigravityGatewayService) writeClaudeError(c *gin.Context, status int, errType, message string) error { + MarkResponseCommitted(c) + c.JSON(status, gin.H{ + "type": "error", + "error": gin.H{"type": errType, "message": message}, + }) + return fmt.Errorf("%s", message) +} + +// WriteMappedClaudeError 导出版本，供 handler 层使用（如 fallback 错误处理） +func (s *AntigravityGatewayService) WriteMappedClaudeError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error { + return s.writeMappedClaudeError(c, account, upstreamStatus, upstreamRequestID, body) +} + +func (s *AntigravityGatewayService) writeMappedClaudeError(c *gin.Context, account *Account, upstreamStatus int, upstreamRequestID string, body []byte) error { + MarkResponseCommitted(c) + upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + logBody, maxBytes := s.getLogConfig() + upstreamDetail := s.getUpstreamErrorDetail(body) + setOpsUpstreamError(c, upstreamStatus, upstreamMsg, upstreamDetail) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: upstreamStatus, + UpstreamRequestID: upstreamRequestID, + Kind: "http_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + // 记录上游错误详情便于排障（可选：由配置控制；不回显到客户端） + if logBody { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] upstream_error status=%d body=%s", upstreamStatus, truncateForLog(body, maxBytes)) + } + + // 检查错误透传规则 + if ptStatus, ptErrType, ptErrMsg, matched := applyErrorPassthroughRule( + c, account.Platform, upstreamStatus, body, + 0, "", "", + ); matched { + c.JSON(ptStatus, gin.H{ + "type": "error", + "error": gin.H{"type": ptErrType, "message": ptErrMsg}, + }) + if upstreamMsg == "" { + return fmt.Errorf("upstream error: %d", upstreamStatus) + } + return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg) + } + + var statusCode int + var errType, errMsg string + + switch upstreamStatus { + case 400: + statusCode = http.StatusBadRequest + errType = "invalid_request_error" + errMsg = getPassthroughOrDefault(upstreamMsg, "Invalid request") + case 401: + statusCode = http.StatusBadGateway + errType = "authentication_error" + errMsg = "Upstream authentication failed" + case 403: + statusCode = http.StatusBadGateway + errType = "permission_error" + errMsg = "Upstream access forbidden" + case 429: + statusCode = http.StatusTooManyRequests + errType = "rate_limit_error" + errMsg = "Upstream rate limit exceeded" + case 529: + statusCode = http.StatusServiceUnavailable + errType = "overloaded_error" + errMsg = "Upstream service overloaded" + default: + statusCode = http.StatusBadGateway + errType = "upstream_error" + errMsg = "Upstream request failed" + } + + c.JSON(statusCode, gin.H{ + "type": "error", + "error": gin.H{"type": errType, "message": errMsg}, + }) + if upstreamMsg == "" { + return fmt.Errorf("upstream error: %d", upstreamStatus) + } + return fmt.Errorf("upstream error: %d message=%s", upstreamStatus, upstreamMsg) +} + +func (s *AntigravityGatewayService) writeGoogleError(c *gin.Context, status int, message string) error { + MarkResponseCommitted(c) + statusStr := "UNKNOWN" + switch status { + case 400: + statusStr = "INVALID_ARGUMENT" + case 404: + statusStr = "NOT_FOUND" + case 429: + statusStr = "RESOURCE_EXHAUSTED" + case 500: + statusStr = "INTERNAL" + case 502, 503: + statusStr = "UNAVAILABLE" + } + + c.JSON(status, gin.H{ + "error": gin.H{ + "code": status, + "message": message, + "status": statusStr, + }, + }) + return fmt.Errorf("%s", message) +} + +// handleClaudeStreamToNonStreaming 收集上游流式响应，转换为 Claude 非流式格式返回 +// 用于处理客户端非流式请求但上游只支持流式的情况 +func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) { + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.settingService.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + + var firstTokenMs *int + var last map[string]any + var lastWithParts map[string]any + var collectedParts []map[string]any // 收集所有 parts（包括 text、thinking、functionCall、inlineData 等） + + type scanEvent struct { + line string + err error + } + + // 独立 goroutine 读取上游，避免读取阻塞影响超时处理 + events := make(chan scanEvent, 16) + done := make(chan struct{}) + sendEvent := func(ev scanEvent) bool { + select { + case events <- ev: + return true + case <-done: + return false + } + } + + var lastReadAt int64 + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + go func(scanBuf *sseScannerBuf64K) { + defer putSSEScannerBuf64K(scanBuf) + defer close(events) + for scanner.Scan() { + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + if !sendEvent(scanEvent{line: scanner.Text()}) { + return + } + } + if err := scanner.Err(); err != nil { + _ = sendEvent(scanEvent{err: err}) + } + }(scanBuf) + defer close(done) + + // 上游数据间隔超时保护（防止上游挂起长期占用连接） + streamInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var intervalTicker *time.Ticker + if streamInterval > 0 { + intervalTicker = time.NewTicker(streamInterval) + defer intervalTicker.Stop() + } + var intervalCh <-chan time.Time + if intervalTicker != nil { + intervalCh = intervalTicker.C + } + + for { + select { + case ev, ok := <-events: + if !ok { + // 流结束，转换并返回响应 + goto returnResponse + } + if ev.err != nil { + if errors.Is(ev.err, bufio.ErrTooLong) { + logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity claude non-stream): max_size=%d error=%v", maxLineSize, ev.err) + } + return nil, ev.err + } + + line := ev.line + trimmed := strings.TrimRight(line, "\r\n") + + if !strings.HasPrefix(trimmed, "data:") { + continue + } + + payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:")) + if payload == "" || payload == "[DONE]" { + continue + } + + // 解包 v1internal 响应 + inner, parseErr := s.unwrapV1InternalResponse([]byte(payload)) + if parseErr != nil { + continue + } + + var parsed map[string]any + if err := json.Unmarshal(inner, &parsed); err != nil { + continue + } + + // 记录首 token 时间 + if firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + + last = parsed + + // 保留最后一个有 parts 的响应，并收集所有 parts + if parts := extractGeminiParts(parsed); len(parts) > 0 { + lastWithParts = parsed + + // 收集所有 parts（text、thinking、functionCall、inlineData 等） + collectedParts = append(collectedParts, parts...) + } + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity claude non-stream)") + return nil, fmt.Errorf("stream data interval timeout") + } + } + +returnResponse: + // 选择最后一个有效响应 + finalResponse := pickGeminiCollectResult(last, lastWithParts) + + // 处理空响应情况 — 触发同账号重试 + failover 切换账号 + if last == nil && lastWithParts == nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] warning: empty stream response (claude non-stream), triggering failover") + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), + RetryableOnSameAccount: true, + } + } + + // 将收集的所有 parts 合并到最终响应中 + if len(collectedParts) > 0 { + finalResponse = mergeCollectedPartsToResponse(finalResponse, collectedParts) + } + + // 序列化为 JSON（Gemini 格式） + geminiBody, err := json.Marshal(finalResponse) + if err != nil { + return nil, fmt.Errorf("failed to marshal gemini response: %w", err) + } + + // 转换 Gemini 响应为 Claude 格式 + claudeResp, agUsage, err := antigravity.TransformGeminiToClaude(geminiBody, originalModel) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] transform_error error=%v body=%s", err, string(geminiBody)) + return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response") + } + + c.Data(http.StatusOK, "application/json", claudeResp) + + // 转换为 service.ClaudeUsage + usage := &ClaudeUsage{ + InputTokens: agUsage.InputTokens, + OutputTokens: agUsage.OutputTokens, + CacheCreationInputTokens: agUsage.CacheCreationInputTokens, + CacheReadInputTokens: agUsage.CacheReadInputTokens, + } + + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs}, nil +} + +// handleClaudeStreamingResponse 处理 Claude 流式响应（Gemini SSE → Claude SSE 转换） +func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) { + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + c.Status(http.StatusOK) + + flusher, ok := c.Writer.(http.Flusher) + if !ok { + return nil, errors.New("streaming not supported") + } + + processor := antigravity.NewStreamingProcessor(originalModel) + var firstTokenMs *int + // 使用 Scanner 并限制单行大小，避免 ReadString 无上限导致 OOM + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.settingService.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + + // 辅助函数：转换 antigravity.ClaudeUsage 到 service.ClaudeUsage + convertUsage := func(agUsage *antigravity.ClaudeUsage) *ClaudeUsage { + if agUsage == nil { + return &ClaudeUsage{} + } + return &ClaudeUsage{ + InputTokens: agUsage.InputTokens, + OutputTokens: agUsage.OutputTokens, + CacheCreationInputTokens: agUsage.CacheCreationInputTokens, + CacheReadInputTokens: agUsage.CacheReadInputTokens, + } + } + + type scanEvent struct { + line string + err error + } + // 独立 goroutine 读取上游，避免读取阻塞影响超时处理 + events := make(chan scanEvent, 16) + done := make(chan struct{}) + sendEvent := func(ev scanEvent) bool { + select { + case events <- ev: + return true + case <-done: + return false + } + } + var lastReadAt int64 + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + go func(scanBuf *sseScannerBuf64K) { + defer putSSEScannerBuf64K(scanBuf) + defer close(events) + for scanner.Scan() { + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + if !sendEvent(scanEvent{line: scanner.Text()}) { + return + } + } + if err := scanner.Err(); err != nil { + _ = sendEvent(scanEvent{err: err}) + } + }(scanBuf) + defer close(done) + + streamInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var intervalTicker *time.Ticker + if streamInterval > 0 { + intervalTicker = time.NewTicker(streamInterval) + defer intervalTicker.Stop() + } + var intervalCh <-chan time.Time + if intervalTicker != nil { + intervalCh = intervalTicker.C + } + + // 下游 keepalive：防止代理/Cloudflare Tunnel 因连接空闲而断开 + keepaliveInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { + keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second + } + var keepaliveTicker *time.Ticker + if keepaliveInterval > 0 { + keepaliveTicker = time.NewTicker(keepaliveInterval) + defer keepaliveTicker.Stop() + } + var keepaliveCh <-chan time.Time + if keepaliveTicker != nil { + keepaliveCh = keepaliveTicker.C + } + lastDataAt := time.Now() + + cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity claude") + + // 仅发送一次错误事件，避免多次写入导致协议混乱 + errorEventSent := false + sendErrorEvent := func(reason string) { + if errorEventSent || cw.Disconnected() { + return + } + errorEventSent = true + _, _ = fmt.Fprintf(c.Writer, "event: error\ndata: {\"error\":\"%s\"}\n\n", reason) + flusher.Flush() + } + + // finishUsage 是获取 processor 最终 usage 的辅助函数 + finishUsage := func() *ClaudeUsage { + _, agUsage := processor.Finish() + return convertUsage(agUsage) + } + + for { + select { + case ev, ok := <-events: + if !ok { + // 上游完成，发送结束事件 + finalEvents, agUsage := processor.Finish() + if len(finalEvents) > 0 { + cw.Write(finalEvents) + } else if !processor.MessageStartSent() && !cw.Disconnected() { + // 整个流未收到任何可解析的上游数据（全部 SSE 行均无法被 JSON 解析）， + // 触发 failover 在同账号重试，避免向客户端发出缺少 message_start 的残缺流 + logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Claude-Stream] empty stream response (no valid events parsed), triggering failover") + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: []byte(`{"error":"empty stream response from upstream"}`), + RetryableOnSameAccount: true, + } + } + return &antigravityStreamResult{usage: convertUsage(agUsage), firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()}, nil + } + if ev.err != nil { + if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity claude"); handled { + return &antigravityStreamResult{usage: finishUsage(), firstTokenMs: firstTokenMs, clientDisconnect: disconnect}, nil + } + if errors.Is(ev.err, bufio.ErrTooLong) { + logger.LegacyPrintf("service.antigravity_gateway", "SSE line too long (antigravity): max_size=%d error=%v", maxLineSize, ev.err) + sendErrorEvent("response_too_large") + return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, ev.err + } + sendErrorEvent("stream_read_error") + return nil, fmt.Errorf("stream read error: %w", ev.err) + } + + lastDataAt = time.Now() + + // 处理 SSE 行，转换为 Claude 格式 + claudeEvents := processor.ProcessLine(strings.TrimRight(ev.line, "\r\n")) + if len(claudeEvents) > 0 { + if firstTokenMs == nil { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + cw.Write(claudeEvents) + } + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + if cw.Disconnected() { + logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity claude), returning collected usage") + return &antigravityStreamResult{usage: finishUsage(), firstTokenMs: firstTokenMs, clientDisconnect: true}, nil + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity)") + sendErrorEvent("stream_timeout") + return &antigravityStreamResult{usage: convertUsage(nil), firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") + + case <-keepaliveCh: + if cw.Disconnected() { + continue + } + if time.Since(lastDataAt) < keepaliveInterval { + continue + } + // SSE ping 事件：Anthropic 原生格式，客户端会正确处理， + // 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开 + if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") { + logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity claude), continuing to drain upstream for billing") + continue + } + } + } +} + +func (s *AntigravityGatewayService) extractImageInputSize(body []byte) string { + var req antigravity.GeminiRequest + if err := json.Unmarshal(body, &req); err != nil { + return "" + } + + if req.GenerationConfig != nil && req.GenerationConfig.ImageConfig != nil { + return strings.TrimSpace(req.GenerationConfig.ImageConfig.ImageSize) + } + + return "" +} + +// isImageGenerationModel 判断模型是否为图片生成模型 +// 支持的模型：gemini-3.1-flash-image, gemini-3-pro-image, gemini-2.5-flash-image 等 +func isImageGenerationModel(model string) bool { + modelLower := strings.ToLower(model) + // 移除 models/ 前缀 + modelLower = strings.TrimPrefix(modelLower, "models/") + + // 精确匹配或前缀匹配 + return modelLower == "gemini-3.1-flash-image" || + modelLower == "gemini-3.1-flash-image-preview" || + strings.HasPrefix(modelLower, "gemini-3.1-flash-image-") || + modelLower == "gemini-3-pro-image" || + modelLower == "gemini-3-pro-image-preview" || + strings.HasPrefix(modelLower, "gemini-3-pro-image-") || + modelLower == "gemini-2.5-flash-image" || + modelLower == "gemini-2.5-flash-image-preview" || + strings.HasPrefix(modelLower, "gemini-2.5-flash-image-") +} diff --git a/backend/internal/service/antigravity_gateway_upstream.go b/backend/internal/service/antigravity_gateway_upstream.go new file mode 100644 index 0000000000..2914863131 --- /dev/null +++ b/backend/internal/service/antigravity_gateway_upstream.go @@ -0,0 +1,375 @@ +package service + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "sync/atomic" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/gin-gonic/gin" +) + +// ForwardUpstream 使用 base_url + /v1/messages + 双 header 认证透传上游 Claude 请求 +func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) { + startTime := time.Now() + sessionID := getSessionID(c) + prefix := logPrefix(sessionID, account.Name) + + // 获取上游配置 + baseURL := strings.TrimSpace(account.GetCredential("base_url")) + apiKey := strings.TrimSpace(account.GetCredential("api_key")) + if baseURL == "" || apiKey == "" { + return nil, fmt.Errorf("upstream account missing base_url or api_key") + } + baseURL = strings.TrimSuffix(baseURL, "/") + + // 解析请求获取模型信息 + var claudeReq antigravity.ClaudeRequest + if err := json.Unmarshal(body, &claudeReq); err != nil { + return nil, fmt.Errorf("parse claude request: %w", err) + } + if strings.TrimSpace(claudeReq.Model) == "" { + return nil, fmt.Errorf("missing model") + } + originalModel := claudeReq.Model + + // 构建上游请求 URL + upstreamURL := baseURL + "/v1/messages" + + // 能力维度 sanitize：Anthropic-compatible 上游透传路径也需要保证 body↔beta header + // 对称。客户端 anthropic-beta header 不含 context-management-2025-06-27 但 body 带 + // context_management 时 strip，与 Anthropic 直连 / Bedrock / Vertex 路径保持一致。 + clientBeta := c.GetHeader("anthropic-beta") + if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, clientBeta); changed { + body = sanitized + } + + // 创建请求 + req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("create upstream request: %w", err) + } + + // 设置请求头 + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+apiKey) + req.Header.Set("x-api-key", apiKey) // Claude API 兼容 + + // 透传 Claude 相关 headers + if v := c.GetHeader("anthropic-version"); v != "" { + req.Header.Set("anthropic-version", v) + } + if v := clientBeta; v != "" { + req.Header.Set("anthropic-beta", v) + } + + // 代理 URL + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + + // 发送请求 + resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s upstream request failed: %v", prefix, err) + return nil, fmt.Errorf("upstream request failed: %w", err) + } + defer func() { _ = resp.Body.Close() }() + + // 处理错误响应 + if resp.StatusCode >= 400 { + respBody := s.readUpstreamErrorBody(resp) + + // 429 错误时标记账号限流 + if resp.StatusCode == http.StatusTooManyRequests { + s.handleUpstreamError(ctx, prefix, account, resp.StatusCode, resp.Header, respBody, originalModel, 0, "", false) + } + + // 透传上游错误 + c.Header("Content-Type", resp.Header.Get("Content-Type")) + c.Status(resp.StatusCode) + _, _ = c.Writer.Write(respBody) + + return &ForwardResult{ + Model: originalModel, + }, nil + } + + // 处理成功响应（流式/非流式） + var usage *ClaudeUsage + var firstTokenMs *int + var clientDisconnect bool + + if claudeReq.Stream { + // 流式响应：透传 + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + c.Status(http.StatusOK) + + streamRes := s.streamUpstreamResponse(c, resp, startTime) + usage = streamRes.usage + firstTokenMs = streamRes.firstTokenMs + clientDisconnect = streamRes.clientDisconnect + } else { + // 非流式响应：直接透传 + respBody, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("read upstream response: %w", err) + } + + // 提取 usage + usage = s.extractClaudeUsage(respBody) + + c.Header("Content-Type", resp.Header.Get("Content-Type")) + c.Status(http.StatusOK) + _, _ = c.Writer.Write(respBody) + } + + // 构建计费结果 + duration := time.Since(startTime) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=success duration_ms=%d", prefix, duration.Milliseconds()) + + return &ForwardResult{ + Model: originalModel, + Stream: claudeReq.Stream, + Duration: duration, + FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnect, + Usage: ClaudeUsage{ + InputTokens: usage.InputTokens, + OutputTokens: usage.OutputTokens, + CacheReadInputTokens: usage.CacheReadInputTokens, + CacheCreationInputTokens: usage.CacheCreationInputTokens, + }, + }, nil +} + +// streamUpstreamResponse 透传上游 SSE 流并提取 Claude usage +func (s *AntigravityGatewayService) streamUpstreamResponse(c *gin.Context, resp *http.Response, startTime time.Time) *antigravityStreamResult { + usage := &ClaudeUsage{} + var firstTokenMs *int + + scanner := bufio.NewScanner(resp.Body) + maxLineSize := defaultMaxLineSize + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.settingService.cfg.Gateway.MaxLineSize + } + scanner.Buffer(make([]byte, 64*1024), maxLineSize) + + type scanEvent struct { + line string + err error + } + events := make(chan scanEvent, 16) + done := make(chan struct{}) + sendEvent := func(ev scanEvent) bool { + select { + case events <- ev: + return true + case <-done: + return false + } + } + var lastReadAt int64 + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + go func() { + defer close(events) + for scanner.Scan() { + atomic.StoreInt64(&lastReadAt, time.Now().UnixNano()) + if !sendEvent(scanEvent{line: scanner.Text()}) { + return + } + } + if err := scanner.Err(); err != nil { + _ = sendEvent(scanEvent{err: err}) + } + }() + defer close(done) + + streamInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.settingService.cfg.Gateway.StreamDataIntervalTimeout) * time.Second + } + var intervalTicker *time.Ticker + if streamInterval > 0 { + intervalTicker = time.NewTicker(streamInterval) + defer intervalTicker.Stop() + } + var intervalCh <-chan time.Time + if intervalTicker != nil { + intervalCh = intervalTicker.C + } + + // 下游 keepalive：防止代理/Cloudflare Tunnel 因连接空闲而断开 + keepaliveInterval := time.Duration(0) + if s.settingService.cfg != nil && s.settingService.cfg.Gateway.StreamKeepaliveInterval > 0 { + keepaliveInterval = time.Duration(s.settingService.cfg.Gateway.StreamKeepaliveInterval) * time.Second + } + var keepaliveTicker *time.Ticker + if keepaliveInterval > 0 { + keepaliveTicker = time.NewTicker(keepaliveInterval) + defer keepaliveTicker.Stop() + } + var keepaliveCh <-chan time.Time + if keepaliveTicker != nil { + keepaliveCh = keepaliveTicker.C + } + lastDataAt := time.Now() + + flusher, _ := c.Writer.(http.Flusher) + cw := newAntigravityClientWriter(c.Writer, flusher, "antigravity upstream") + + for { + select { + case ev, ok := <-events: + if !ok { + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: cw.Disconnected()} + } + if ev.err != nil { + if disconnect, handled := handleStreamReadError(ev.err, cw.Disconnected(), "antigravity upstream"); handled { + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: disconnect} + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream read error (antigravity upstream): %v", ev.err) + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs} + } + + lastDataAt = time.Now() + + line := ev.line + + // 记录首 token 时间 + if firstTokenMs == nil && len(line) > 0 { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + + // 尝试从 message_delta 或 message_stop 事件提取 usage + s.extractSSEUsage(line, usage) + + // 透传行 + cw.Fprintf("%s\n", line) + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + if cw.Disconnected() { + logger.LegacyPrintf("service.antigravity_gateway", "Upstream timeout after client disconnect (antigravity upstream), returning collected usage") + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true} + } + logger.LegacyPrintf("service.antigravity_gateway", "Stream data interval timeout (antigravity upstream)") + return &antigravityStreamResult{usage: usage, firstTokenMs: firstTokenMs} + + case <-keepaliveCh: + if cw.Disconnected() { + continue + } + if time.Since(lastDataAt) < keepaliveInterval { + continue + } + // SSE ping 事件：Anthropic 原生格式，客户端会正确处理， + // 同时保持连接活跃防止 Cloudflare Tunnel 等代理断开 + if !cw.Fprintf("event: ping\ndata: {\"type\": \"ping\"}\n\n") { + logger.LegacyPrintf("service.antigravity_gateway", "Client disconnected during keepalive ping (antigravity upstream), continuing to drain upstream for billing") + continue + } + } + } +} + +// extractSSEUsage 从 SSE data 行中提取 Claude usage（用于流式透传场景） +// +// Anthropic streaming 的 usage 字段分布在两类事件中： +// - message_start：嵌套在 event.message.usage（input_tokens、cache_creation_input_tokens、 +// cache_read_input_tokens 等输入侧字段） +// - message_delta：位于顶层 event.usage（流结束时的最终 output_tokens） +// +// 仅读取顶层 event.usage 会漏掉 message_start 的输入侧字段，导致流式透传请求落库的 +// usage_logs 记录 input_tokens=0。 +func (s *AntigravityGatewayService) extractSSEUsage(line string, usage *ClaudeUsage) { + if !strings.HasPrefix(line, "data: ") { + return + } + dataStr := strings.TrimPrefix(line, "data: ") + var event map[string]any + if json.Unmarshal([]byte(dataStr), &event) != nil { + return + } + var u map[string]any + if eventType, _ := event["type"].(string); eventType == "message_start" { + if msg, ok := event["message"].(map[string]any); ok { + u, _ = msg["usage"].(map[string]any) + } + } else { + u, _ = event["usage"].(map[string]any) + } + if u == nil { + return + } + if v, ok := u["input_tokens"].(float64); ok && int(v) > 0 { + usage.InputTokens = int(v) + } + if v, ok := u["output_tokens"].(float64); ok && int(v) > 0 { + usage.OutputTokens = int(v) + } + if v, ok := u["cache_read_input_tokens"].(float64); ok && int(v) > 0 { + usage.CacheReadInputTokens = int(v) + } + if v, ok := u["cache_creation_input_tokens"].(float64); ok && int(v) > 0 { + usage.CacheCreationInputTokens = int(v) + } + // 解析嵌套的 cache_creation 对象中的 5m/1h 明细 + if cc, ok := u["cache_creation"].(map[string]any); ok { + if v, ok := cc["ephemeral_5m_input_tokens"].(float64); ok { + usage.CacheCreation5mTokens = int(v) + } + if v, ok := cc["ephemeral_1h_input_tokens"].(float64); ok { + usage.CacheCreation1hTokens = int(v) + } + } +} + +// extractClaudeUsage 从非流式 Claude 响应提取 usage +func (s *AntigravityGatewayService) extractClaudeUsage(body []byte) *ClaudeUsage { + usage := &ClaudeUsage{} + var resp map[string]any + if json.Unmarshal(body, &resp) != nil { + return usage + } + if u, ok := resp["usage"].(map[string]any); ok { + if v, ok := u["input_tokens"].(float64); ok { + usage.InputTokens = int(v) + } + if v, ok := u["output_tokens"].(float64); ok { + usage.OutputTokens = int(v) + } + if v, ok := u["cache_read_input_tokens"].(float64); ok { + usage.CacheReadInputTokens = int(v) + } + if v, ok := u["cache_creation_input_tokens"].(float64); ok { + usage.CacheCreationInputTokens = int(v) + } + // 解析嵌套的 cache_creation 对象中的 5m/1h 明细 + if cc, ok := u["cache_creation"].(map[string]any); ok { + if v, ok := cc["ephemeral_5m_input_tokens"].(float64); ok { + usage.CacheCreation5mTokens = int(v) + } + if v, ok := cc["ephemeral_1h_input_tokens"].(float64); ok { + usage.CacheCreation1hTokens = int(v) + } + } + } + return usage +}