From 084d26cbd2c9ba865222501fd6001a388a76b83c Mon Sep 17 00:00:00 2001 From: shaw Date: Tue, 7 Jul 2026 23:32:57 +0800 Subject: [PATCH] =?UTF-8?q?refactor(service):=20=E7=BA=AF=E7=A7=BB?= =?UTF-8?q?=E5=8A=A8=E6=8B=86=E5=88=86=20gateway=5Fservice.go=EF=BC=887294?= =?UTF-8?q?=E2=86=921289=E8=A1=8C=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../service/gateway_claude_oauth_body.go | 1210 ++++ .../internal/service/gateway_count_tokens.go | 606 ++ backend/internal/service/gateway_forward.go | 959 +++ backend/internal/service/gateway_service.go | 6005 ----------------- .../service/gateway_upstream_request.go | 923 +++ .../service/gateway_upstream_response.go | 1426 ++++ .../internal/service/gateway_usage_billing.go | 978 +++ 7 files changed, 6102 insertions(+), 6005 deletions(-) create mode 100644 backend/internal/service/gateway_claude_oauth_body.go create mode 100644 backend/internal/service/gateway_count_tokens.go create mode 100644 backend/internal/service/gateway_forward.go create mode 100644 backend/internal/service/gateway_upstream_request.go create mode 100644 backend/internal/service/gateway_upstream_response.go create mode 100644 backend/internal/service/gateway_usage_billing.go diff --git a/backend/internal/service/gateway_claude_oauth_body.go b/backend/internal/service/gateway_claude_oauth_body.go new file mode 100644 index 0000000000..80604dd8ac --- /dev/null +++ b/backend/internal/service/gateway_claude_oauth_body.go @@ -0,0 +1,1210 @@ +package service + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/json" + "fmt" + "strconv" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/pkg/anthropicfp" + "github.com/Wei-Shaw/sub2api/internal/pkg/claude" + infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/google/uuid" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + + "github.com/gin-gonic/gin" +) + +type anthropicCacheControlPayload struct { + Type string `json:"type"` + TTL string `json:"ttl,omitempty"` +} + +type anthropicSystemTextBlockPayload struct { + Type string `json:"type"` + Text string `json:"text"` + CacheControl *anthropicCacheControlPayload `json:"cache_control,omitempty"` +} + +type anthropicMetadataPayload struct { + UserID string `json:"user_id"` +} + +// replaceModelInBody 替换请求体中的model字段 +// 优先使用定点修改，尽量保持客户端原始字段顺序。 +func (s *GatewayService) replaceModelInBody(body []byte, newModel string) []byte { + return ReplaceModelInBody(body, newModel) +} + +type claudeOAuthNormalizeOptions struct { + injectMetadata bool + metadataUserID string + stripSystemCacheControl bool +} + +// sanitizeSystemText rewrites only the fixed OpenCode identity sentence (if present). +// We intentionally avoid broad keyword replacement in system prompts to prevent +// accidentally changing user-provided instructions. +func sanitizeSystemText(text string) string { + if text == "" { + return text + } + // Some clients include a fixed OpenCode identity sentence. Anthropic may treat + // this as a non-Claude-Code fingerprint, so rewrite it to the canonical + // Claude Code banner before generic "OpenCode"/"opencode" replacements. + text = strings.ReplaceAll( + text, + "You are OpenCode, the best coding agent on the planet.", + strings.TrimSpace(claudeCodeSystemPrompt), + ) + return text +} + +func marshalAnthropicSystemTextBlock(text string, includeCacheControl bool) ([]byte, error) { + block := anthropicSystemTextBlockPayload{ + Type: "text", + Text: text, + } + if includeCacheControl { + block.CacheControl = &anthropicCacheControlPayload{ + Type: "ephemeral", + TTL: claude.DefaultCacheControlTTL, + } + } + return json.Marshal(block) +} + +func marshalAnthropicSystemTextBlockWithCacheControl(text string, cacheControl any) ([]byte, error) { + block := map[string]any{ + "type": "text", + "text": text, + } + if cacheControl != nil { + block["cache_control"] = cacheControl + } + return json.Marshal(block) +} + +func marshalAnthropicMetadata(userID string) ([]byte, error) { + return json.Marshal(anthropicMetadataPayload{UserID: userID}) +} + +func buildJSONArrayRaw(items [][]byte) []byte { + if len(items) == 0 { + return []byte("[]") + } + + total := 2 + for _, item := range items { + total += len(item) + } + total += len(items) - 1 + + buf := make([]byte, 0, total) + buf = append(buf, '[') + for i, item := range items { + if i > 0 { + buf = append(buf, ',') + } + buf = append(buf, item...) + } + buf = append(buf, ']') + return buf +} + +func setJSONValueBytes(body []byte, path string, value any) ([]byte, bool) { + next, err := sjson.SetBytes(body, path, value) + if err != nil { + return body, false + } + return next, true +} + +func setJSONRawBytes(body []byte, path string, raw []byte) ([]byte, bool) { + next, err := sjson.SetRawBytes(body, path, raw) + if err != nil { + return body, false + } + return next, true +} + +func deleteJSONPathBytes(body []byte, path string) ([]byte, bool) { + next, err := sjson.DeleteBytes(body, path) + if err != nil { + return body, false + } + return next, true +} + +func normalizeClaudeOAuthSystemBody(body []byte, opts claudeOAuthNormalizeOptions) ([]byte, bool) { + sys := gjson.GetBytes(body, "system") + if !sys.Exists() { + return body, false + } + + out := body + modified := false + + switch { + case sys.Type == gjson.String: + sanitized := sanitizeSystemText(sys.String()) + if sanitized != sys.String() { + if next, ok := setJSONValueBytes(out, "system", sanitized); ok { + out = next + modified = true + } + } + case sys.IsArray(): + index := 0 + sys.ForEach(func(_, item gjson.Result) bool { + if item.Get("type").String() == "text" { + textResult := item.Get("text") + if textResult.Exists() && textResult.Type == gjson.String { + text := textResult.String() + sanitized := sanitizeSystemText(text) + if sanitized != text { + if next, ok := setJSONValueBytes(out, fmt.Sprintf("system.%d.text", index), sanitized); ok { + out = next + modified = true + } + } + } + } + + if opts.stripSystemCacheControl && item.Get("cache_control").Exists() { + if next, ok := deleteJSONPathBytes(out, fmt.Sprintf("system.%d.cache_control", index)); ok { + out = next + modified = true + } + } + + index++ + return true + }) + } + + return out, modified +} + +func ensureClaudeOAuthMetadataUserID(body []byte, userID string) ([]byte, bool) { + if strings.TrimSpace(userID) == "" { + return body, false + } + + metadata := gjson.GetBytes(body, "metadata") + if !metadata.Exists() || metadata.Type == gjson.Null { + raw, err := marshalAnthropicMetadata(userID) + if err != nil { + return body, false + } + return setJSONRawBytes(body, "metadata", raw) + } + + trimmedRaw := strings.TrimSpace(metadata.Raw) + if strings.HasPrefix(trimmedRaw, "{") { + existing := metadata.Get("user_id") + if existing.Exists() && existing.Type == gjson.String && existing.String() != "" { + return body, false + } + return setJSONValueBytes(body, "metadata.user_id", userID) + } + + raw, err := marshalAnthropicMetadata(userID) + if err != nil { + return body, false + } + return setJSONRawBytes(body, "metadata", raw) +} + +func normalizeClaudeOAuthRequestBody(body []byte, modelID string, opts claudeOAuthNormalizeOptions) ([]byte, string) { + if len(body) == 0 { + return body, modelID + } + + out := body + modified := false + + if next, changed := normalizeClaudeOAuthSystemBody(out, opts); changed { + out = next + modified = true + } + + rawModel := gjson.GetBytes(out, "model") + if rawModel.Exists() && rawModel.Type == gjson.String { + normalized := claude.NormalizeModelID(rawModel.String()) + if normalized != rawModel.String() { + if next, ok := setJSONValueBytes(out, "model", normalized); ok { + out = next + modified = true + } + modelID = normalized + } + } + + // 确保 tools 字段存在（即使为空数组） + if !gjson.GetBytes(out, "tools").Exists() { + if next, ok := setJSONRawBytes(out, "tools", []byte("[]")); ok { + out = next + modified = true + } + } + + if opts.injectMetadata && opts.metadataUserID != "" { + if next, changed := ensureClaudeOAuthMetadataUserID(out, opts.metadataUserID); changed { + out = next + modified = true + } + } + + // temperature：真实 Claude Code CLI 总是发送 temperature（默认 1，客户端可覆盖）。 + // 之前的实现直接 delete 会导致 payload 缺字段，与真实 CLI 字节级不一致。 + // 策略：客户端传了什么就透传；没传则补默认 1。 + if !gjson.GetBytes(out, "temperature").Exists() { + if next, ok := setJSONValueBytes(out, "temperature", 1); ok { + out = next + modified = true + } + } + + // max_tokens：真实 CLI 的默认值是 128000。缺失时补齐以对齐指纹。 + if !gjson.GetBytes(out, "max_tokens").Exists() { + if next, ok := setJSONValueBytes(out, "max_tokens", 128000); ok { + out = next + modified = true + } + } + + // context_management：thinking.type 为 enabled/adaptive 时，真实 CLI 会自动 + // 附带 {"edits":[{"type":"clear_thinking_20251015","keep":"all"}]}。 + // 客户端显式传了就透传；否则按 CLI 行为补齐。 + // + // 注：本函数不按 model 名决定是否保留 context_management。“最终 beta + // header 不含 context-management-2025-06-27 时 strip 字段”的能力维度 + // 对称约束由 sanitizeAnthropicBodyForBetaTokens 在 buildUpstreamRequest / + // buildCountTokensRequest 层统一执行，与 Bedrock 路径的 + // sanitizeBedrockFieldsForBetaTokens 对称。 + if !gjson.GetBytes(out, "context_management").Exists() { + thinkingType := gjson.GetBytes(out, "thinking.type").String() + if thinkingType == "enabled" || thinkingType == "adaptive" { + const cmDefault = `{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]}` + if next, ok := setJSONRawBytes(out, "context_management", []byte(cmDefault)); ok { + out = next + modified = true + } + } + } + + // tool_choice：与 Parrot 对齐，不再无条件删除。 + // - 客户端传了 {"type":"tool","name":"X"} → 保留结构，name 由 + // applyToolNameRewriteToBody 同步映射为假名 + // - 其他形态（auto/any/none）原样透传 + // 如果 body 里完全没有 tools（空数组），tool_choice 没意义时才删除 + if !gjson.GetBytes(out, "tools").IsArray() || len(gjson.GetBytes(out, "tools").Array()) == 0 { + if gjson.GetBytes(out, "tool_choice").Exists() { + if next, ok := deleteJSONPathBytes(out, "tool_choice"); ok { + out = next + modified = true + } + } + } + + if !modified { + return body, modelID + } + + return out, modelID +} + +func (s *GatewayService) buildOAuthMetadataUserID(parsed *ParsedRequest, account *Account, fp *Fingerprint) string { + if parsed == nil || account == nil { + return "" + } + if parsed.MetadataUserID != "" { + return "" + } + + userID := strings.TrimSpace(account.GetClaudeUserID()) + if userID == "" && fp != nil { + userID = fp.ClientID + } + if userID == "" { + // Fall back to a random, well-formed client id so we can still satisfy + // Claude Code OAuth requirements when account metadata is incomplete. + userID = generateClientID() + } + + // session_id 用"会话级稳定种子"派生（账号 + 客户端区分因子 + 首条 user 文本）： + // 随对话在尾部追加 messages 时保持不变，贴近真实 CC 进程级稳定的 session_id。 + // 不复用 GenerateSessionHash —— 后者是粘性路由键、按设计逐轮变化（见其测试）。 + var firstUserText string + if parsed.Body != nil { + firstUserText = extractFirstUserText(parsed.Body.Bytes()) + } + seed := buildStableSessionSeed(account.ID, sessionContextDiscriminator(parsed.SessionContext), firstUserText) + sessionID := generateSessionUUID(seed) + + // 根据指纹 UA 版本选择输出格式 + var uaVersion string + if fp != nil { + uaVersion = ExtractCLIVersion(fp.UserAgent) + } + accountUUID := strings.TrimSpace(account.GetExtraString("account_uuid")) + return FormatMetadataUserID(userID, accountUUID, sessionID, uaVersion) +} + +// applyClaudeCodeOAuthMimicryToBody 将"非 Claude Code 客户端 + Claude OAuth 账号" +// 路径上原本只在 /v1/messages 里做的完整伪装应用到任意 body 上。 +// +// 这是 /v1/messages 主路径上 rewriteSystemForNonClaudeCode + +// normalizeClaudeOAuthRequestBody 流程的通用版，供 OpenAI 协议兼容层 +// (ForwardAsChatCompletions / ForwardAsResponses) 复用。 +// +// 未抽离之前，OpenAI 协议兼容层仅做 injectClaudeCodePrompt（前置追加）， +// 而仓内 /v1/messages 路径自己的注释明确说过"仅前置追加无法通过 Anthropic +// 第三方检测"；那条注释就是本函数存在的根因。 +// +// 参数： +// - ctx / c：用于读取指纹和 gateway settings；c 可为 nil（如 count_tokens）。 +// - account：必须是 OAuth 账号，且调用方已判断不是 Claude Code 客户端。 +// - body：已经 marshal 成 Anthropic /v1/messages 格式的请求体。 +// - systemRaw：body 中原始 system 字段（用于判断是否需要 rewrite）。 +// - model：最终会发给上游的模型 ID（用于 haiku 旁路 + metadata 版本选择）。 +// +// 返回：改写后的 body。即使中间任何一步失败，也会退化成原 body（不会 panic）。 +func (s *GatewayService) applyClaudeCodeOAuthMimicryToBody( + ctx context.Context, + c *gin.Context, + account *Account, + body []byte, + systemRaw any, + model string, +) []byte { + if account == nil || !account.IsOAuth() || len(body) == 0 { + return body + } + + systemPromptInjectionEnabled, systemPrompt, systemPromptBlocks := s.claudeOAuthSystemPromptInjectionSettings(ctx) + systemRewritten := false + if systemPromptInjectionEnabled && !strings.Contains(strings.ToLower(model), "haiku") { + body = rewriteSystemForNonClaudeCodeWithPromptBlocks(body, normalizeSystemParam(systemRaw), systemPrompt, systemPromptBlocks) + systemRewritten = true + } + + normalizeOpts := claudeOAuthNormalizeOptions{stripSystemCacheControl: !systemRewritten} + + if s.identityService != nil && c != nil && c.Request != nil { + if fp, err := s.identityService.GetOrCreateFingerprint(ctx, account.ID, c.Request.Header); err == nil && fp != nil { + mimicMPT := false + if s.settingService != nil { + _, mimicMPT, _ = s.settingService.GetGatewayForwardingSettings(ctx) + } + if !mimicMPT { + if uid := s.buildOAuthMetadataUserIDFromBody(ctx, account, fp, body); uid != "" { + normalizeOpts.injectMetadata = true + normalizeOpts.metadataUserID = uid + } + } + } + } + + body, _ = normalizeClaudeOAuthRequestBody(body, model, normalizeOpts) + + // Phase D+E+F: messages cache 策略 + 工具名混淆 + tools[-1] 断点 + // 对齐 Parrot transform_request 里剩余的字段级改写。顺序有语义约束： + // 1) messages cache：仅在配置开启时清除客户端断点并注入代理断点 + // 2) tool rewrite：最后改 tools[*].name / tool_choice.name 并在 tools[-1] + // 上打断点；mapping 存入 gin.Context 供响应侧 bytes.Replace 还原。 + body = s.rewriteMessageCacheControlIfEnabled(ctx, body) + + if rw := buildToolNameRewriteFromBody(body); rw != nil { + body = applyToolNameRewriteToBody(body, rw) + if c != nil { + c.Set(toolNameRewriteKey, rw) + } + } else { + body = applyToolsLastCacheBreakpoint(body) + } + + return body +} + +// buildOAuthMetadataUserIDFromBody 是 buildOAuthMetadataUserID 的变体， +// 适用于调用方手上没有 ParsedRequest 的场景（如 OpenAI 协议兼容层）。 +// +// 与 buildOAuthMetadataUserID 的唯一区别： +// - session hash 从 body 本体按同样规则重算，而不是读取 ParsedRequest 缓存值。 +// - 如果 body 里已经存在 metadata.user_id，则返回空（由 ensureClaudeOAuthMetadataUserID +// 自行决定是否覆盖）。 +func (s *GatewayService) buildOAuthMetadataUserIDFromBody( + ctx context.Context, + account *Account, + fp *Fingerprint, + body []byte, +) string { + _ = ctx + if account == nil { + return "" + } + if existing := gjson.GetBytes(body, "metadata.user_id").String(); existing != "" { + return "" + } + + userID := strings.TrimSpace(account.GetClaudeUserID()) + if userID == "" && fp != nil { + userID = fp.ClientID + } + if userID == "" { + userID = generateClientID() + } + + // 与 buildOAuthMetadataUserID 一致：用会话级稳定种子，避免整 body 哈希导致 + // 每轮（甚至每个 token 变化）都重算出不同的 session_id。 + var clientDiscriminator string + if fp != nil { + clientDiscriminator = fp.ClientID + } + seed := buildStableSessionSeed(account.ID, clientDiscriminator, extractFirstUserText(body)) + sessionID := generateSessionUUID(seed) + + var uaVersion string + if fp != nil { + uaVersion = ExtractCLIVersion(fp.UserAgent) + } + accountUUID := strings.TrimSpace(account.GetExtraString("account_uuid")) + return FormatMetadataUserID(userID, accountUUID, sessionID, uaVersion) +} + +// buildStableSessionSeed 为伪装路径合成的 metadata.user_id session_id 生成"会话级稳定"种子。 +// +// 真实 Claude Code 的 session_id 是进程级随机 UUID，在一段会话内跨请求保持不变。无状态代理 +// 无法恢复该值，这里用"会话内不变的锚点"近似：账号 ID + 客户端区分因子 + 首条 user 消息文本。 +// 对话在尾部追加 messages 时这三者都不变，因此 generateSessionUUID(seed) 跨轮稳定。 +// +// 注意：粘性路由键 GenerateSessionHash 按设计逐轮变化（见其测试），本函数与之独立、互不影响。 +// accountID 恒存在，故 seed 永不为空 —— 输出始终是确定性 UUID，而非随机值。 +func buildStableSessionSeed(accountID int64, clientDiscriminator, firstUserText string) string { + var b strings.Builder + _, _ = b.WriteString(strconv.FormatInt(accountID, 10)) + _, _ = b.WriteString("::") + _, _ = b.WriteString(clientDiscriminator) + _, _ = b.WriteString("::") + _, _ = b.WriteString(firstUserText) + return b.String() +} + +// sessionContextDiscriminator 把请求上下文（客户端 IP / 归一化 UA / API Key ID）拼成 +// 一个跨客户端的区分因子，避免不同用户的相同首条消息派生出相同 session_id。 +func sessionContextDiscriminator(sc *SessionContext) string { + if sc == nil { + return "" + } + return sc.ClientIP + ":" + NormalizeSessionUserAgent(sc.UserAgent) + ":" + strconv.FormatInt(sc.APIKeyID, 10) +} + +// GenerateSessionUUID creates a deterministic UUID4 from a seed string. +func GenerateSessionUUID(seed string) string { + return generateSessionUUID(seed) +} + +func generateSessionUUID(seed string) string { + if seed == "" { + return uuid.NewString() + } + hash := sha256.Sum256([]byte(seed)) + bytes := hash[:16] + bytes[6] = (bytes[6] & 0x0f) | 0x40 + bytes[8] = (bytes[8] & 0x3f) | 0x80 + return fmt.Sprintf("%x-%x-%x-%x-%x", + bytes[0:4], bytes[4:6], bytes[6:8], bytes[8:10], bytes[10:16]) +} + +// normalizeSystemParam 将 json.RawMessage 类型的 system 参数转为标准 Go 类型（string / []any / nil）， +// 避免 type switch 中 json.RawMessage（底层 []byte）无法匹配 case string / case []any / case nil 的问题。 +// 这是 Go 的 typed nil 陷阱：(json.RawMessage, nil) ≠ (nil, nil)。 +func normalizeSystemParam(system any) any { + raw, ok := system.(json.RawMessage) + if !ok { + return system + } + if len(raw) == 0 { + return nil + } + var parsed any + if err := json.Unmarshal(raw, &parsed); err != nil { + return nil + } + return parsed +} + +// systemIncludesClaudeCodePrompt 检查 system 中是否已包含 Claude Code 提示词 +// 使用前缀匹配支持多种变体（标准版、Agent SDK 版等） +func systemIncludesClaudeCodePrompt(system any) bool { + system = normalizeSystemParam(system) + switch v := system.(type) { + case string: + return hasClaudeCodePrefix(v) + case []any: + for _, item := range v { + if m, ok := item.(map[string]any); ok { + if text, ok := m["text"].(string); ok && hasClaudeCodePrefix(text) { + return true + } + } + } + } + return false +} + +// hasClaudeCodePrefix 检查文本是否以 Claude Code 提示词的特征前缀开头 +func hasClaudeCodePrefix(text string) bool { + for _, prefix := range claudeCodePromptPrefixes { + if strings.HasPrefix(text, prefix) { + return true + } + } + return false +} + +// injectClaudeCodePrompt 在 system 开头注入 Claude Code 提示词 +// 处理 null、字符串、数组三种格式 +func injectClaudeCodePrompt(body []byte, system any) []byte { + system = normalizeSystemParam(system) + claudeCodeBlock, err := marshalAnthropicSystemTextBlock(claudeCodeSystemPrompt, true) + if err != nil { + logger.LegacyPrintf("service.gateway", "Warning: failed to build Claude Code prompt block: %v", err) + return body + } + // Opencode plugin applies an extra safeguard: it not only prepends the Claude Code + // banner, it also prefixes the next system instruction with the same banner plus + // a blank line. This helps when upstream concatenates system instructions. + claudeCodePrefix := strings.TrimSpace(claudeCodeSystemPrompt) + + var items [][]byte + + switch v := system.(type) { + case nil: + items = [][]byte{claudeCodeBlock} + case string: + // Be tolerant of older/newer clients that may differ only by trailing whitespace/newlines. + if strings.TrimSpace(v) == "" || strings.TrimSpace(v) == strings.TrimSpace(claudeCodeSystemPrompt) { + items = [][]byte{claudeCodeBlock} + } else { + // Mirror opencode behavior: keep the banner as a separate system entry, + // but also prefix the next system text with the banner. + merged := v + if !strings.HasPrefix(v, claudeCodePrefix) { + merged = claudeCodePrefix + "\n\n" + v + } + nextBlock, buildErr := marshalAnthropicSystemTextBlock(merged, false) + if buildErr != nil { + logger.LegacyPrintf("service.gateway", "Warning: failed to build prefixed Claude Code system block: %v", buildErr) + return body + } + items = [][]byte{claudeCodeBlock, nextBlock} + } + case []any: + items = make([][]byte, 0, len(v)+1) + items = append(items, claudeCodeBlock) + prefixedNext := false + systemResult := gjson.GetBytes(body, "system") + if systemResult.IsArray() { + systemResult.ForEach(func(_, item gjson.Result) bool { + textResult := item.Get("text") + if textResult.Exists() && textResult.Type == gjson.String && + strings.TrimSpace(textResult.String()) == strings.TrimSpace(claudeCodeSystemPrompt) { + return true + } + + raw := []byte(item.Raw) + // Prefix the first subsequent text system block once. + if !prefixedNext && item.Get("type").String() == "text" && textResult.Exists() && textResult.Type == gjson.String { + text := textResult.String() + if strings.TrimSpace(text) != "" && !strings.HasPrefix(text, claudeCodePrefix) { + next, setErr := sjson.SetBytes(raw, "text", claudeCodePrefix+"\n\n"+text) + if setErr == nil { + raw = next + prefixedNext = true + } + } + } + items = append(items, raw) + return true + }) + } else { + for _, item := range v { + m, ok := item.(map[string]any) + if !ok { + raw, marshalErr := json.Marshal(item) + if marshalErr == nil { + items = append(items, raw) + } + continue + } + if text, ok := m["text"].(string); ok && strings.TrimSpace(text) == strings.TrimSpace(claudeCodeSystemPrompt) { + continue + } + if !prefixedNext { + if blockType, _ := m["type"].(string); blockType == "text" { + if text, ok := m["text"].(string); ok && strings.TrimSpace(text) != "" && !strings.HasPrefix(text, claudeCodePrefix) { + m["text"] = claudeCodePrefix + "\n\n" + text + prefixedNext = true + } + } + } + raw, marshalErr := json.Marshal(m) + if marshalErr == nil { + items = append(items, raw) + } + } + } + default: + items = [][]byte{claudeCodeBlock} + } + + result, ok := setJSONRawBytes(body, "system", buildJSONArrayRaw(items)) + if !ok { + logger.LegacyPrintf("service.gateway", "Warning: failed to inject Claude Code prompt") + return body + } + return result +} + +// rewriteSystemForNonClaudeCode 将非 Claude Code 客户端的 system prompt 迁移至 messages， +// system 字段仅保留 Claude Code 标识提示词。 +// Anthropic 基于 system 参数内容检测第三方应用，仅前置追加 Claude Code 提示词 +// 无法通过检测，因为后续内容仍为非 Claude Code 格式。 +// 策略：将原始 system prompt 提取并注入为 user/assistant 消息对，system 仅保留 Claude Code 标识。 +func rewriteSystemForNonClaudeCode(body []byte, system any) []byte { + return rewriteSystemForNonClaudeCodeWithPromptBlocks(body, system, "", "") +} + +func rewriteSystemForNonClaudeCodeWithPrompt(body []byte, system any, expansionPrompt string) []byte { + return rewriteSystemForNonClaudeCodeWithPromptBlocks(body, system, expansionPrompt, "") +} + +type claudeOAuthSystemPromptBlockConfig struct { + Enabled *bool `json:"enabled,omitempty"` + Type string `json:"type,omitempty"` + Text string `json:"text,omitempty"` + CacheControl json.RawMessage `json:"cache_control,omitempty"` +} + +type claudeOAuthSystemPromptBlocksEnvelope struct { + Blocks []claudeOAuthSystemPromptBlockConfig `json:"blocks"` +} + +func defaultClaudeOAuthExpansionPrompt(expansionPrompt string) string { + expansionPrompt = strings.TrimSpace(expansionPrompt) + if expansionPrompt == "" { + return claudeCodeSystemPromptExpansion + } + return expansionPrompt +} + +func parseClaudeOAuthSystemPromptBlocksConfig(raw string) ([]claudeOAuthSystemPromptBlockConfig, error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return nil, nil + } + if strings.HasPrefix(raw, "[") { + var blocks []claudeOAuthSystemPromptBlockConfig + if err := json.Unmarshal([]byte(raw), &blocks); err != nil { + return nil, err + } + return blocks, nil + } + var envelope claudeOAuthSystemPromptBlocksEnvelope + if err := json.Unmarshal([]byte(raw), &envelope); err != nil { + return nil, err + } + return envelope.Blocks, nil +} + +func decodeClaudeOAuthSystemPromptCacheControl(raw json.RawMessage) (any, error) { + trimmed := bytes.TrimSpace(raw) + if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) || bytes.Equal(trimmed, []byte("false")) { + return nil, nil + } + if bytes.Equal(trimmed, []byte("true")) { + return map[string]string{ + "type": "ephemeral", + "ttl": claude.DefaultCacheControlTTL, + }, nil + } + var value any + if err := json.Unmarshal(trimmed, &value); err != nil { + return nil, err + } + if _, ok := value.(map[string]any); !ok { + return nil, fmt.Errorf("cache_control must be boolean, null, or object") + } + return value, nil +} + +func expandClaudeOAuthSystemPromptTextTemplate(body []byte, text string, expansionPrompt string) (string, error) { + if text == "" { + return "", nil + } + expansionPrompt = defaultClaudeOAuthExpansionPrompt(expansionPrompt) + billingText, err := buildBillingAttributionText(body, claude.CLICurrentVersion) + if err != nil { + return "", err + } + fp := computeClaudeCodeFingerprint(body, claude.CLICurrentVersion) + replacer := strings.NewReplacer( + "{billing_header}", billingText, + "{cc_version}", claude.CLICurrentVersion, + "{fp}", fp, + "{claude_code_system_prompt}", claudeCodeSystemPrompt, + "{claude_code_expansion_prompt}", expansionPrompt, + ) + return replacer.Replace(text), nil +} + +func defaultClaudeOAuthSystemPromptBlockConfig() []claudeOAuthSystemPromptBlockConfig { + enabled := true + return []claudeOAuthSystemPromptBlockConfig{ + { + Enabled: &enabled, + Type: "text", + Text: "{billing_header}", + }, + { + Enabled: &enabled, + Type: "text", + Text: "{claude_code_system_prompt}", + }, + { + Enabled: &enabled, + Type: "text", + Text: "{claude_code_expansion_prompt}", + CacheControl: json.RawMessage( + fmt.Sprintf(`{"type":"ephemeral","ttl":%q}`, claude.DefaultCacheControlTTL), + ), + }, + } +} + +func buildClaudeOAuthSystemPromptBlocksJSON(body []byte, expansionPrompt string, blocksConfig string) ([][]byte, error) { + blocks, err := parseClaudeOAuthSystemPromptBlocksConfig(blocksConfig) + if err != nil { + return nil, err + } + if len(blocks) == 0 { + blocks = defaultClaudeOAuthSystemPromptBlockConfig() + } + + items := make([][]byte, 0, len(blocks)) + for i, block := range blocks { + if block.Enabled != nil && !*block.Enabled { + continue + } + blockType := strings.TrimSpace(block.Type) + if blockType == "" { + blockType = "text" + } + if blockType != "text" { + return nil, fmt.Errorf("system block %d type %q is not supported", i, block.Type) + } + text, err := expandClaudeOAuthSystemPromptTextTemplate(body, block.Text, expansionPrompt) + if err != nil { + return nil, err + } + if strings.TrimSpace(text) == "" { + continue + } + cacheControl, err := decodeClaudeOAuthSystemPromptCacheControl(block.CacheControl) + if err != nil { + return nil, fmt.Errorf("system block %d cache_control: %w", i, err) + } + raw, err := marshalAnthropicSystemTextBlockWithCacheControl(text, cacheControl) + if err != nil { + return nil, err + } + items = append(items, raw) + } + return items, nil +} + +func ValidateClaudeOAuthSystemPromptBlocksConfig(raw string) error { + if strings.TrimSpace(raw) == "" { + return nil + } + blocks, err := parseClaudeOAuthSystemPromptBlocksConfig(raw) + if err != nil { + return infraerrors.BadRequest("INVALID_CLAUDE_OAUTH_SYSTEM_PROMPT_BLOCKS", "claude oauth system prompt blocks must be valid JSON") + } + for i, block := range blocks { + blockType := strings.TrimSpace(block.Type) + if blockType == "" { + blockType = "text" + } + if blockType != "text" { + return infraerrors.BadRequest("INVALID_CLAUDE_OAUTH_SYSTEM_PROMPT_BLOCKS", fmt.Sprintf("system block %d type must be text", i)) + } + if _, err := decodeClaudeOAuthSystemPromptCacheControl(block.CacheControl); err != nil { + return infraerrors.BadRequest("INVALID_CLAUDE_OAUTH_SYSTEM_PROMPT_BLOCKS", fmt.Sprintf("system block %d cache_control is invalid", i)) + } + } + return nil +} + +func rewriteSystemForNonClaudeCodeWithPromptBlocks(body []byte, system any, expansionPrompt string, blocksConfig string) []byte { + system = normalizeSystemParam(system) + expansionPrompt = defaultClaudeOAuthExpansionPrompt(expansionPrompt) + + // 1. 提取原始 system prompt 文本 + var originalSystemText string + switch v := system.(type) { + case string: + originalSystemText = strings.TrimSpace(v) + case []any: + var parts []string + for _, item := range v { + if m, ok := item.(map[string]any); ok { + if text, ok := m["text"].(string); ok && strings.TrimSpace(text) != "" { + parts = append(parts, text) + } + } + } + originalSystemText = strings.Join(parts, "\n\n") + } + + // 2. 构造 system 数组，对齐真实 Claude Code CLI 的 3-block 形态： + // [0] billing attribution block（cc_version={cliVer}.{fp}; cc_entrypoint=cli;） + // [1] "You are Claude Code..." 身份前缀 block（默认不带 cache_control） + // [2] 工具无关的通用提示词扩充 block（带 cache_control 作为稳定缓存断点） + // + // 真实 CC 的 system 在身份前缀之后还有大段提示词，仅有 2 块会在块数/体量上明显 + // 区别于真实 CLI。这里注入 claudeCodeSystemPromptExpansion（中性段落）把形态做到 + // 接近真实，同时不注入会污染被代理用户行为的工具专属指令。 + // + // 缺失 billing block 的系统 payload 是 Anthropic 判定第三方的关键信号之一 + // （真实 CLI 每个请求都带）。新版 CLI 已取消 cch=... 签名字段，故 block 不再注入 + // cch（见 buildBillingAttributionText）。 + systemBlocks, blockErr := buildClaudeOAuthSystemPromptBlocksJSON(body, expansionPrompt, blocksConfig) + if blockErr != nil { + logger.LegacyPrintf("service.gateway", "Warning: failed to build configured Claude OAuth system blocks: %v", blockErr) + systemBlocks, blockErr = buildClaudeOAuthSystemPromptBlocksJSON(body, expansionPrompt, "") + } + if blockErr != nil { + logger.LegacyPrintf("service.gateway", "Warning: failed to build default Claude OAuth system blocks: %v", blockErr) + return body + } + out, ok := setJSONRawBytes(body, "system", buildJSONArrayRaw(systemBlocks)) + if !ok { + logger.LegacyPrintf("service.gateway", "Warning: failed to set Claude Code system prompt") + return body + } + + // 3. 将原始 system prompt 作为 user/assistant 消息对注入到 messages 开头 + // 模型仍通过 messages 接收完整指令，保留客户端功能 + ccPromptTrimmed := strings.TrimSpace(claudeCodeSystemPrompt) + if originalSystemText != "" && originalSystemText != ccPromptTrimmed && !hasClaudeCodePrefix(originalSystemText) { + instrMsg, err1 := json.Marshal(map[string]any{ + "role": "user", + "content": []map[string]any{ + {"type": "text", "text": "[System Instructions]\n" + originalSystemText}, + }, + }) + ackMsg, err2 := json.Marshal(map[string]any{ + "role": "assistant", + "content": []map[string]any{ + {"type": "text", "text": "Understood. I will follow these instructions."}, + }, + }) + if err1 != nil || err2 != nil { + logger.LegacyPrintf("service.gateway", "Warning: failed to marshal system-to-messages injection") + return out + } + + // 重建 messages 数组：[instruction, ack, ...originalMessages] + items := [][]byte{instrMsg, ackMsg} + messagesResult := gjson.GetBytes(out, "messages") + if messagesResult.IsArray() { + messagesResult.ForEach(func(_, msg gjson.Result) bool { + items = append(items, []byte(msg.Raw)) + return true + }) + } + + if next, setOk := setJSONRawBytes(out, "messages", buildJSONArrayRaw(items)); setOk { + out = next + } + } + + return out +} + +type cacheControlPath struct { + path string + log string +} + +func collectCacheControlPaths(body []byte) (invalidThinking []cacheControlPath, messagePaths []string, toolPaths []string, systemPaths []string) { + system := gjson.GetBytes(body, "system") + if system.IsArray() { + sysIndex := 0 + system.ForEach(func(_, item gjson.Result) bool { + if item.Get("cache_control").Exists() { + path := fmt.Sprintf("system.%d.cache_control", sysIndex) + if item.Get("type").String() == "thinking" { + invalidThinking = append(invalidThinking, cacheControlPath{ + path: path, + log: "[Warning] Removed illegal cache_control from thinking block in system", + }) + } else { + systemPaths = append(systemPaths, path) + } + } + sysIndex++ + return true + }) + } + + messages := gjson.GetBytes(body, "messages") + if messages.IsArray() { + msgIndex := 0 + messages.ForEach(func(_, msg gjson.Result) bool { + content := msg.Get("content") + if content.IsArray() { + contentIndex := 0 + content.ForEach(func(_, item gjson.Result) bool { + if item.Get("cache_control").Exists() { + path := fmt.Sprintf("messages.%d.content.%d.cache_control", msgIndex, contentIndex) + if item.Get("type").String() == "thinking" { + invalidThinking = append(invalidThinking, cacheControlPath{ + path: path, + log: fmt.Sprintf("[Warning] Removed illegal cache_control from thinking block in messages[%d].content[%d]", msgIndex, contentIndex), + }) + } else { + messagePaths = append(messagePaths, path) + } + } + contentIndex++ + return true + }) + } + msgIndex++ + return true + }) + } + + tools := gjson.GetBytes(body, "tools") + if tools.IsArray() { + toolIndex := 0 + tools.ForEach(func(_, tool gjson.Result) bool { + if tool.Get("cache_control").Exists() { + toolPaths = append(toolPaths, fmt.Sprintf("tools.%d.cache_control", toolIndex)) + } + toolIndex++ + return true + }) + } + + return invalidThinking, messagePaths, toolPaths, systemPaths +} + +// enforceCacheControlLimit 强制执行 cache_control 块数量限制（最多 4 个） +// 超限时优先移除工具断点，再移除 messages 断点，最后才移除 system 断点。 +func enforceCacheControlLimit(body []byte) []byte { + if len(body) == 0 { + return body + } + + invalidThinking, messagePaths, toolPaths, systemPaths := collectCacheControlPaths(body) + out := body + modified := false + + // 先清理 thinking 块中的非法 cache_control（thinking 块不支持该字段） + for _, item := range invalidThinking { + if !gjson.GetBytes(out, item.path).Exists() { + continue + } + next, ok := deleteJSONPathBytes(out, item.path) + if !ok { + continue + } + out = next + modified = true + logger.LegacyPrintf("service.gateway", "%s", item.log) + } + + count := len(messagePaths) + len(toolPaths) + len(systemPaths) + if count <= maxCacheControlBlocks { + if modified { + return out + } + return body + } + + // 超限：优先从 tools 中移除，再从 messages 中移除，最后才从 system 中移除。 + remaining := count - maxCacheControlBlocks + for i := len(toolPaths) - 1; i >= 0 && remaining > 0; i-- { + path := toolPaths[i] + if !gjson.GetBytes(out, path).Exists() { + continue + } + next, ok := deleteJSONPathBytes(out, path) + if !ok { + continue + } + out = next + modified = true + remaining-- + } + + for _, path := range messagePaths { + if remaining <= 0 { + break + } + if !gjson.GetBytes(out, path).Exists() { + continue + } + next, ok := deleteJSONPathBytes(out, path) + if !ok { + continue + } + out = next + modified = true + remaining-- + } + + for i := len(systemPaths) - 1; i >= 0 && remaining > 0; i-- { + path := systemPaths[i] + if !gjson.GetBytes(out, path).Exists() { + continue + } + next, ok := deleteJSONPathBytes(out, path) + if !ok { + continue + } + out = next + modified = true + remaining-- + } + + if modified { + return out + } + return body +} + +// injectAnthropicCacheControlTTL1h 将已有 ephemeral cache_control 块的 ttl 强制写为 1h。 +// 仅修改已经存在的 cache_control，不新增缓存断点。 +func injectAnthropicCacheControlTTL1h(body []byte) []byte { + return forceEphemeralCacheControlTTL(body, cacheTTLTarget1h) +} + +func forceEphemeralCacheControlTTL(body []byte, ttl string) []byte { + if len(body) == 0 || ttl == "" { + return body + } + out := body + var paths []string + addPath := func(path string, value gjson.Result) { + cc := value.Get("cache_control") + if !cc.Exists() || cc.Get("type").String() != "ephemeral" { + return + } + if cc.Get("ttl").String() == ttl { + return + } + paths = append(paths, path+".cache_control.ttl") + } + + if topCC := gjson.GetBytes(body, "cache_control"); topCC.Exists() && topCC.Get("type").String() == "ephemeral" && topCC.Get("ttl").String() != ttl { + paths = append(paths, "cache_control.ttl") + } + + system := gjson.GetBytes(body, "system") + if system.IsArray() { + idx := -1 + system.ForEach(func(_, block gjson.Result) bool { + idx++ + addPath(fmt.Sprintf("system.%d", idx), block) + return true + }) + } + + messages := gjson.GetBytes(body, "messages") + if messages.IsArray() { + msgIdx := -1 + messages.ForEach(func(_, msg gjson.Result) bool { + msgIdx++ + content := msg.Get("content") + if !content.IsArray() { + return true + } + contentIdx := -1 + content.ForEach(func(_, block gjson.Result) bool { + contentIdx++ + addPath(fmt.Sprintf("messages.%d.content.%d", msgIdx, contentIdx), block) + return true + }) + return true + }) + } + + tools := gjson.GetBytes(body, "tools") + if tools.IsArray() { + idx := -1 + tools.ForEach(func(_, tool gjson.Result) bool { + idx++ + addPath(fmt.Sprintf("tools.%d", idx), tool) + return true + }) + } + + for _, path := range paths { + if next, err := sjson.SetBytes(out, path, ttl); err == nil { + out = next + } + } + return out +} + +func (s *GatewayService) shouldInjectAnthropicCacheTTL1h(ctx context.Context, account *Account) bool { + if account == nil || !account.IsAnthropicOAuthOrSetupToken() || s == nil || s.settingService == nil { + return false + } + return s.settingService.IsAnthropicCacheTTL1hInjectionEnabled(ctx) +} + +// shouldNormalizeClientDateline reports whether the request body's client +// dateline should be normalized before forwarding to Anthropic. The switch is +// scoped to Anthropic OAuth/SetupToken accounts only; API-Key accounts and +// non-Anthropic platforms bypass this step entirely. +func (s *GatewayService) shouldNormalizeClientDateline(ctx context.Context, account *Account) bool { + if account == nil || !account.IsAnthropicOAuthOrSetupToken() || s == nil || s.settingService == nil { + return false + } + return s.settingService.IsClientDatelineNormalizationEnabled(ctx) +} + +// normalizeClientDatelineIfEnabled applies dateline normalization to body when +// the switch is on and the account qualifies. Returns (nextBody, true) only +// when the body actually changed; otherwise returns (nil, false) so callers +// can skip the writeback. +func (s *GatewayService) normalizeClientDatelineIfEnabled(ctx context.Context, account *Account, body []byte) ([]byte, bool) { + if !s.shouldNormalizeClientDateline(ctx, account) { + return nil, false + } + next, _, changed := anthropicfp.NormalizeDateline(body) + if !changed { + return nil, false + } + return next, true +} + +func (s *GatewayService) claudeOAuthSystemPromptInjectionSettings(ctx context.Context) (bool, string, string) { + if s == nil || s.settingService == nil { + return true, "", "" + } + return s.settingService.GetClaudeOAuthSystemPromptInjectionSettings(ctx) +} diff --git a/backend/internal/service/gateway_count_tokens.go b/backend/internal/service/gateway_count_tokens.go new file mode 100644 index 0000000000..500f655660 --- /dev/null +++ b/backend/internal/service/gateway_count_tokens.go @@ -0,0 +1,606 @@ +package service + +import ( + "bytes" + "context" + "errors" + "fmt" + "net/http" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/pkg/claude" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/tidwall/gjson" + + "github.com/gin-gonic/gin" +) + +// ForwardCountTokens 转发 count_tokens 请求到上游 API +// 特点：不记录使用量、仅支持非流式响应 +func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, account *Account, parsed *ParsedRequest) error { + if parsed == nil { + s.countTokensError(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty") + return fmt.Errorf("parse request: empty request") + } + + if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() { + passthroughBody := parsed.Body.Bytes() + if reqModel := parsed.Model; reqModel != "" { + if mappedModel := account.GetMappedModel(reqModel); mappedModel != reqModel { + passthroughBody = s.replaceModelInBody(passthroughBody, mappedModel) + logger.LegacyPrintf("service.gateway", "CountTokens passthrough model mapping: %s -> %s (account: %s)", reqModel, mappedModel, account.Name) + } + } + return s.forwardCountTokensAnthropicAPIKeyPassthrough(ctx, c, account, passthroughBody) + } + + // Bedrock 不支持 count_tokens 端点 + if account != nil && account.IsBedrock() { + s.countTokensError(c, http.StatusNotFound, "not_found_error", "count_tokens endpoint is not supported for Bedrock") + return nil + } + + body := parsed.Body.Bytes() + replaceBody := func(next []byte) error { + if err := parsed.ReplaceBody(next); err != nil { + return fmt.Errorf("rewrite count_tokens body: %w", err) + } + body = parsed.Body.Bytes() + return nil + } + reqModel := parsed.Model + + // Pre-filter: strip empty text blocks to prevent upstream 400. + if err := replaceBody(StripEmptyTextBlocks(body)); err != nil { + return err + } + + isClaudeCodeCT := IsClaudeCodeClient(ctx) || isClaudeCodeClient(c.GetHeader("User-Agent"), parsed.MetadataUserID) + shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCodeCT + + if shouldMimicClaudeCode { + normalizeOpts := claudeOAuthNormalizeOptions{stripSystemCacheControl: true} + var normalizedBody []byte + normalizedBody, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) + if err := replaceBody(normalizedBody); err != nil { + return err + } + + if err := replaceBody(s.rewriteMessageCacheControlIfEnabled(ctx, body)); err != nil { + return err + } + if rw := buildToolNameRewriteFromBody(body); rw != nil { + if err := replaceBody(applyToolNameRewriteToBody(body, rw)); err != nil { + return err + } + } else { + if err := replaceBody(applyToolsLastCacheBreakpoint(body)); err != nil { + return err + } + } + } + + // Antigravity 账户不支持 count_tokens，返回 404 让客户端 fallback 到本地估算。 + // 返回 nil 避免 handler 层记录为错误，也不设置 ops 上游错误上下文。 + if account.Platform == PlatformAntigravity { + s.countTokensError(c, http.StatusNotFound, "not_found_error", "count_tokens endpoint is not supported for this platform") + return nil + } + + // 应用模型映射： + // - APIKey 账号：使用账号级别的显式映射（如果配置），否则透传原始模型名 + // - OAuth/SetupToken 账号：使用 Anthropic 标准映射（短ID → 长ID） + if reqModel != "" { + mappedModel := reqModel + mappingSource := "" + if account.Type == AccountTypeAPIKey { + mappedModel = account.GetMappedModel(reqModel) + if mappedModel != reqModel { + mappingSource = "account" + } + } + if mappingSource == "" && account.Platform == PlatformAnthropic && account.Type != AccountTypeAPIKey { + normalized := claude.NormalizeModelID(reqModel) + if normalized != reqModel { + mappedModel = normalized + mappingSource = "prefix" + } + } + if mappedModel != reqModel { + originalReqModel := reqModel + if err := replaceBody(s.replaceModelInBody(body, mappedModel)); err != nil { + return err + } + reqModel = mappedModel + parsed.Model = mappedModel + logger.LegacyPrintf("service.gateway", "CountTokens model mapping applied: %s -> %s (account: %s, source=%s)", originalReqModel, mappedModel, account.Name, mappingSource) + } + } + + // 获取凭证 + token, tokenType, err := s.GetAccessToken(ctx, account) + if err != nil { + s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to get access token") + return err + } + + // 构建上游请求 + upstreamReq, wireBody, err := s.buildCountTokensRequest(ctx, c, account, body, token, tokenType, reqModel, shouldMimicClaudeCode) + if err != nil { + s.countTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request") + return err + } + // 先记录首发 wire body；如果后面进入 400 retry，retry 会基于未签名的逻辑 body 重新构建。 + acceptedWireBody := wireBody + + // 获取代理URL（自定义 base URL 模式下，proxy 通过 buildCustomRelayURL 作为查询参数传递） + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + if !account.IsCustomBaseURLEnabled() || account.GetCustomBaseURL() == "" { + proxyURL = account.Proxy.URL() + } + } + + // 发送请求 + resp, err := s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account)) + if err != nil { + setOpsUpstreamError(c, 0, sanitizeUpstreamErrorMessage(err.Error()), "") + s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Request failed") + return fmt.Errorf("upstream request failed: %w", err) + } + + // 读取响应体 + countTokensTooLarge := func(c *gin.Context) { + s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream response too large") + } + respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, countTokensTooLarge) + _ = resp.Body.Close() + if err != nil { + if !errors.Is(err, ErrUpstreamResponseBodyTooLarge) { + s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response") + } + return err + } + + // 检测 thinking block 签名错误（400）并重试一次（过滤 thinking blocks） + if resp.StatusCode == 400 && s.shouldRectifySignatureError(ctx, account, respBody, reqModel) { + logger.LegacyPrintf("service.gateway", "Account %d: detected thinking block signature error on count_tokens, retrying with filtered thinking blocks", account.ID) + + filteredBody := FilterThinkingBlocksForRetry(body, reqModel) + retryReq, retryWireBody, buildErr := s.buildCountTokensRequest(ctx, c, account, filteredBody, token, tokenType, reqModel, shouldMimicClaudeCode) + if buildErr == nil { + retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account)) + if retryErr == nil { + if retryResp.StatusCode < 400 { + // count_tokens 签名重试成功后记录最终 wire body，错误响应仍保留原 body 便于后续处理。 + acceptedWireBody = retryWireBody + } + resp = retryResp + respBody, err = ReadUpstreamResponseBody(resp.Body, s.cfg, c, countTokensTooLarge) + _ = resp.Body.Close() + if err != nil { + if !errors.Is(err, ErrUpstreamResponseBodyTooLarge) { + s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response") + } + return err + } + } + } + } + + if resp.StatusCode < 400 && !bytes.Equal(acceptedWireBody, body) { + // count_tokens 成功后再同步最终 wire body，避免 retry 从已签名 body 派生。 + if err := replaceBody(acceptedWireBody); err != nil { + return err + } + } + + // 处理错误响应 + if resp.StatusCode >= 400 { + // 标记账号状态（429/529等） + s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + + upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + upstreamDetail := "" + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes + if maxBytes <= 0 { + maxBytes = 2048 + } + upstreamDetail = truncateString(string(respBody), maxBytes) + } + setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) + + // 记录上游错误摘要便于排障（不回显请求内容） + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + logger.LegacyPrintf("service.gateway", + "count_tokens upstream error %d (account=%d platform=%s type=%s): %s", + resp.StatusCode, + account.ID, + account.Platform, + account.Type, + truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), + ) + } + + // 返回简化的错误响应 + errMsg := "Upstream request failed" + switch resp.StatusCode { + case 429: + errMsg = "Rate limit exceeded" + case 529: + errMsg = "Service overloaded" + } + s.countTokensError(c, resp.StatusCode, "upstream_error", errMsg) + if upstreamMsg == "" { + return fmt.Errorf("upstream error: %d", resp.StatusCode) + } + return fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, upstreamMsg) + } + + // 透传成功响应 + c.Data(resp.StatusCode, "application/json", respBody) + return nil +} + +func (s *GatewayService) forwardCountTokensAnthropicAPIKeyPassthrough(ctx context.Context, c *gin.Context, account *Account, body []byte) error { + token, tokenType, err := s.GetAccessToken(ctx, account) + if err != nil { + s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to get access token") + return err + } + if tokenType != "apikey" { + s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Invalid account token type") + return fmt.Errorf("anthropic api key passthrough requires apikey token, got: %s", tokenType) + } + + upstreamReq, err := s.buildCountTokensRequestAnthropicAPIKeyPassthrough(ctx, c, account, body, token) + if err != nil { + s.countTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request") + return err + } + + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + + resp, err := s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account)) + if err != nil { + setOpsUpstreamError(c, 0, sanitizeUpstreamErrorMessage(err.Error()), "") + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: 0, + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Passthrough: true, + Kind: "request_error", + Message: sanitizeUpstreamErrorMessage(err.Error()), + }) + s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Request failed") + return fmt.Errorf("upstream request failed: %w", err) + } + + countTokensTooLarge := func(c *gin.Context) { + s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream response too large") + } + respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, countTokensTooLarge) + _ = resp.Body.Close() + if err != nil { + if !errors.Is(err, ErrUpstreamResponseBodyTooLarge) { + s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response") + } + return err + } + + if resp.StatusCode >= 400 { + if s.rateLimitService != nil { + s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + } + + upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + + // 中转站不支持 count_tokens 端点时（404），返回 404 让客户端 fallback 到本地估算。 + // 仅在错误消息明确指向 count_tokens endpoint 不存在时生效，避免误吞其他 404（如错误 base_url）。 + // 返回 nil 避免 handler 层记录为错误，也不设置 ops 上游错误上下文。 + if isCountTokensUnsupported404(resp.StatusCode, respBody) { + logger.LegacyPrintf("service.gateway", + "[count_tokens] Upstream does not support count_tokens (404), returning 404: account=%d name=%s msg=%s", + account.ID, account.Name, truncateString(upstreamMsg, 512)) + s.countTokensError(c, http.StatusNotFound, "not_found_error", "count_tokens endpoint is not supported by upstream") + return nil + } + + upstreamDetail := "" + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes + if maxBytes <= 0 { + maxBytes = 2048 + } + upstreamDetail = truncateString(string(respBody), maxBytes) + } + setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Passthrough: true, + Kind: "http_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + errMsg := "Upstream request failed" + switch resp.StatusCode { + case 429: + errMsg = "Rate limit exceeded" + case 529: + errMsg = "Service overloaded" + } + s.countTokensError(c, resp.StatusCode, "upstream_error", errMsg) + if upstreamMsg == "" { + return fmt.Errorf("upstream error: %d", resp.StatusCode) + } + return fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, upstreamMsg) + } + + writeAnthropicPassthroughResponseHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) + contentType := strings.TrimSpace(resp.Header.Get("Content-Type")) + if contentType == "" { + contentType = "application/json" + } + c.Data(resp.StatusCode, contentType, respBody) + return nil +} + +func (s *GatewayService) buildCountTokensRequestAnthropicAPIKeyPassthrough( + ctx context.Context, + c *gin.Context, + account *Account, + body []byte, + token string, +) (*http.Request, error) { + targetURL := claudeAPICountTokensURL + baseURL := account.GetBaseURL() + if baseURL != "" { + validatedURL, err := s.validateUpstreamBaseURL(baseURL) + if err != nil { + return nil, err + } + targetURL = validatedURL + "/v1/messages/count_tokens?beta=true" + } + body = sanitizeCountTokensRequestBody(body) + + // 同 buildUpstreamRequestAnthropicAPIKeyPassthrough：能力维度 sanitize。 + clientBeta := "" + if c != nil && c.Request != nil { + clientBeta = getHeaderRaw(c.Request.Header, "anthropic-beta") + } + // 账号覆写了 anthropic-beta 时，覆写值即最终上游值：净化以覆写值为准 + if beta, ok := account.HeaderOverrideValue("anthropic-beta"); ok { + clientBeta = beta + } + if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, clientBeta); changed { + body = sanitized + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body)) + if err != nil { + return nil, err + } + + if c != nil && c.Request != nil { + for key, values := range c.Request.Header { + lowerKey := strings.ToLower(strings.TrimSpace(key)) + if !allowedHeaders[lowerKey] { + continue + } + wireKey := resolveWireCasing(key) + for _, v := range values { + addHeaderRaw(req.Header, wireKey, v) + } + } + } + + req.Header.Del("authorization") + req.Header.Del("x-api-key") + req.Header.Del("x-goog-api-key") + req.Header.Del("cookie") + setAnthropicAPIKeyAuthHeader(req.Header, account, token) + + if req.Header.Get("content-type") == "" { + req.Header.Set("content-type", "application/json") + } + if req.Header.Get("anthropic-version") == "" { + req.Header.Set("anthropic-version", "2023-06-01") + } + + // 账号级请求头覆写（最终生效，覆盖上面所有来源的同名头） + account.ApplyHeaderOverrides(req.Header) + + return req, nil +} + +// buildCountTokensRequest 构建 count_tokens 上游请求 +func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, mimicClaudeCode bool) (*http.Request, []byte, error) { + // 确定目标 URL + targetURL := claudeAPICountTokensURL + if account.Type == AccountTypeAPIKey { + baseURL := account.GetBaseURL() + if baseURL != "" { + validatedURL, err := s.validateUpstreamBaseURL(baseURL) + if err != nil { + return nil, nil, err + } + targetURL = validatedURL + "/v1/messages/count_tokens?beta=true" + } + } else if account.IsCustomBaseURLEnabled() { + customURL := account.GetCustomBaseURL() + if customURL == "" { + return nil, nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) + } + validatedURL, err := s.validateUpstreamBaseURL(customURL) + if err != nil { + return nil, nil, err + } + targetURL = s.buildCustomRelayURL(validatedURL, "/v1/messages/count_tokens", account) + } + + clientHeaders := http.Header{} + if c != nil && c.Request != nil { + clientHeaders = c.Request.Header + } + + // OAuth 账号：应用统一指纹和重写 userID（受设置开关控制） + // 如果启用了会话ID伪装，会在重写后替换 session 部分为固定值 + ctEnableFP, ctEnableMPT := true, false + if s.settingService != nil { + ctEnableFP, ctEnableMPT, _ = s.settingService.GetGatewayForwardingSettings(ctx) + } + var ctFingerprint *Fingerprint + if account.IsOAuth() && s.identityService != nil { + fp, err := s.identityService.GetOrCreateFingerprint(ctx, account.ID, clientHeaders) + if err == nil { + ctFingerprint = fp + if !ctEnableMPT { + accountUUID := account.GetExtraString("account_uuid") + if accountUUID != "" && fp.ClientID != "" { + if newBody, err := s.identityService.RewriteUserIDWithMasking(ctx, body, account, accountUUID, fp.ClientID, fp.UserAgent); err == nil && len(newBody) > 0 { + body = newBody + } + } + } + } + } + + // 同步 billing header cc_version 与实际发送的 User-Agent 版本 + if ctFingerprint != nil && ctEnableFP { + body = syncBillingHeaderVersion(body, ctFingerprint.UserAgent) + } + + // === 计算最终 anthropic-beta header（先于 body sanitize 与 CCH 签名）=== + // 顺序约束同 buildUpstreamRequest。 + ctEffectiveDropSet := mergeDropSets(s.getBetaPolicyFilterSet(ctx, c, account, modelID)) + finalBetaHeader, finalBetaShouldSet := s.computeFinalCountTokensAnthropicBeta( + tokenType, mimicClaudeCode, modelID, clientHeaders, body, ctEffectiveDropSet, + ) + + // 账号覆写了 anthropic-beta 时，覆写值即最终上游值：净化以覆写值为准 + if beta, ok := account.HeaderOverrideValue("anthropic-beta"); ok { + finalBetaHeader, finalBetaShouldSet = beta, true + } + + // 能力维度 body sanitize：与最终 anthropic-beta header 对称 + if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, finalBetaHeader); changed { + body = sanitized + } + + body = sanitizeCountTokensRequestBody(body) + + req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body)) + if err != nil { + return nil, nil, err + } + + // 设置认证头（保持原始大小写） + if tokenType == "oauth" { + setHeaderRaw(req.Header, "authorization", "Bearer "+token) + } else { + setAnthropicAPIKeyAuthHeader(req.Header, account, token) + } + + // 白名单透传 headers（恢复真实 wire casing） + for key, values := range clientHeaders { + lowerKey := strings.ToLower(key) + if allowedHeaders[lowerKey] { + wireKey := resolveWireCasing(key) + for _, v := range values { + addHeaderRaw(req.Header, wireKey, v) + } + } + } + + // OAuth 账号：应用指纹到请求头（受设置开关控制） + if ctEnableFP && ctFingerprint != nil { + s.identityService.ApplyFingerprint(req, ctFingerprint) + } + + // 确保必要的 headers 存在（保持原始大小写） + if getHeaderRaw(req.Header, "content-type") == "" { + setHeaderRaw(req.Header, "content-type", "application/json") + } + if getHeaderRaw(req.Header, "anthropic-version") == "" { + setHeaderRaw(req.Header, "anthropic-version", "2023-06-01") + } + if tokenType == "oauth" { + applyClaudeOAuthHeaderDefaults(req) + } + + // OAuth + mimic Claude Code：强制注入 CLI 指纹 header + if tokenType == "oauth" && mimicClaudeCode { + applyClaudeCodeMimicHeaders(req, false) + } + + // 写入最终 anthropic-beta header（Del 一次避免白名单透传值残留） + deleteHeaderAllForms(req.Header, "anthropic-beta") + if finalBetaShouldSet { + setHeaderRaw(req.Header, "anthropic-beta", finalBetaHeader) + } + + // 同步 X-Claude-Code-Session-Id 头：取 body 中已处理的 metadata.user_id 的 session_id 覆盖 + if sessionHeader := getHeaderRaw(req.Header, "X-Claude-Code-Session-Id"); sessionHeader != "" { + if uid := gjson.GetBytes(body, "metadata.user_id").String(); uid != "" { + if parsed := ParseMetadataUserID(uid); parsed != nil { + setHeaderRaw(req.Header, "X-Claude-Code-Session-Id", parsed.SessionID) + } + } + } + + // 账号级请求头覆写（仅 anthropic/openai api_key 账号启用时生效；OAuth 路径 no-op） + account.ApplyHeaderOverrides(req.Header) + + if c != nil && tokenType == "oauth" { + c.Set(claudeMimicDebugInfoKey, buildClaudeMimicDebugLine(req, body, account, tokenType, mimicClaudeCode)) + } + if s.debugClaudeMimicEnabled() { + logClaudeMimicDebug(req, body, account, tokenType, mimicClaudeCode) + } + + return req, body, nil +} + +func sanitizeCountTokensRequestBody(body []byte) []byte { + out := body + for _, path := range []string{ + "temperature", + "top_p", + "top_k", + "stream", + "stop_sequences", + "stop", + } { + if gjson.GetBytes(out, path).Exists() { + if next, ok := deleteJSONPathBytes(out, path); ok { + out = next + } + } + } + return out +} + +// countTokensError 返回 count_tokens 错误响应 +func (s *GatewayService) countTokensError(c *gin.Context, status int, errType, message string) { + c.JSON(status, gin.H{ + "type": "error", + "error": gin.H{ + "type": errType, + "message": message, + }, + }) +} diff --git a/backend/internal/service/gateway_forward.go b/backend/internal/service/gateway_forward.go new file mode 100644 index 0000000000..e690aaab22 --- /dev/null +++ b/backend/internal/service/gateway_forward.go @@ -0,0 +1,959 @@ +package service + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "log/slog" + "net/http" + "strconv" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/claude" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + + "github.com/gin-gonic/gin" +) + +// 重试相关常量 +const ( + // 最大尝试次数（包含首次请求）。过多重试会导致请求堆积与资源耗尽。 + maxRetryAttempts = 5 + + // 指数退避：第 N 次失败后的等待 = retryBaseDelay * 2^(N-1)，并且上限为 retryMaxDelay。 + retryBaseDelay = 300 * time.Millisecond + retryMaxDelay = 3 * time.Second + + // 最大重试耗时（包含请求本身耗时 + 退避等待时间）。 + // 用于防止极端情况下 goroutine 长时间堆积导致资源耗尽。 + maxRetryElapsed = 10 * time.Second +) + +func (s *GatewayService) shouldRetryUpstreamError(account *Account, statusCode int) bool { + // OAuth/Setup Token 账号：仅 403 重试 + if account.IsOAuth() { + return statusCode == 403 + } + + // API Key 账号：未配置的错误码重试 + return !account.ShouldHandleErrorCode(statusCode) +} + +// shouldFailoverUpstreamError determines whether an upstream error should trigger account failover. +func (s *GatewayService) shouldFailoverUpstreamError(statusCode int) bool { + switch statusCode { + case 401, 403, 429, 529: + return true + default: + return statusCode >= 500 + } +} + +func retryBackoffDelay(attempt int) time.Duration { + // attempt 从 1 开始，表示第 attempt 次请求刚失败，需要等待后进行第 attempt+1 次请求。 + if attempt <= 0 { + return retryBaseDelay + } + delay := retryBaseDelay * time.Duration(1<<(attempt-1)) + if delay > retryMaxDelay { + return retryMaxDelay + } + return delay +} + +func sleepWithContext(ctx context.Context, d time.Duration) error { + if d <= 0 { + return nil + } + timer := time.NewTimer(d) + defer func() { + if !timer.Stop() { + select { + case <-timer.C: + default: + } + } + }() + + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +// Forward 转发请求到Claude API +func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, parsed *ParsedRequest) (*ForwardResult, error) { + startTime := time.Now() + if parsed == nil { + return nil, fmt.Errorf("parse request: empty request") + } + + // Web Search 模拟：纯 web_search 请求时，直接调用搜索 API 构造响应 + if account != nil && s.shouldEmulateWebSearch(ctx, account, parsed.GroupID, parsed.Body.Bytes()) { + return s.handleWebSearchEmulation(ctx, c, account, parsed) + } + + if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() { + passthroughBody := parsed.Body.Bytes() + passthroughModel := parsed.Model + if passthroughModel != "" { + if mappedModel := account.GetMappedModel(passthroughModel); mappedModel != passthroughModel { + passthroughBody = s.replaceModelInBody(passthroughBody, mappedModel) + logger.LegacyPrintf("service.gateway", "Passthrough model mapping: %s -> %s (account: %s)", parsed.Model, mappedModel, account.Name) + passthroughModel = mappedModel + } + } + return s.forwardAnthropicAPIKeyPassthroughWithInput(ctx, c, account, anthropicPassthroughForwardInput{ + Body: passthroughBody, + Parsed: parsed, + RequestModel: passthroughModel, + OriginalModel: parsed.Model, + RequestStream: parsed.Stream, + StartTime: startTime, + }) + } + + if account != nil && account.IsBedrock() { + return s.forwardBedrock(ctx, c, account, parsed, startTime) + } + + // Beta policy: evaluate once; block check + cache filter set for buildUpstreamRequest. + // Always overwrite the cache to prevent stale values from a previous retry with a different account. + if account.Platform == PlatformAnthropic && c != nil { + policy := s.evaluateBetaPolicy(ctx, c.GetHeader("anthropic-beta"), account, parsed.Model) + if policy.blockErr != nil { + return nil, policy.blockErr + } + filterSet := policy.filterSet + if filterSet == nil { + filterSet = map[string]struct{}{} + } + c.Set(betaPolicyFilterSetKey, filterSet) + } + + body := parsed.Body.Bytes() + replaceBody := func(next []byte) error { + if err := parsed.ReplaceBody(next); err != nil { + return fmt.Errorf("rewrite request body: %w", err) + } + body = parsed.Body.Bytes() + return nil + } + reqModel := parsed.Model + reqStream := parsed.Stream + originalModel := reqModel + + // === DEBUG: 打印客户端原始请求（headers + body 摘要）=== + if c != nil { + s.debugLogGatewaySnapshot("CLIENT_ORIGINAL", c.Request.Header, body, map[string]string{ + "account": fmt.Sprintf("%d(%s)", account.ID, account.Name), + "account_type": string(account.Type), + "model": reqModel, + "stream": strconv.FormatBool(reqStream), + }) + } + + // Claude Code 客户端判定：UA 匹配 claude-cli/* 且携带 metadata.user_id。 + // 真正的 Claude Code 客户端自带完整的 system prompt、cache_control 断点和 header， + // 不需要代理做任何 body 级别的 mimicry；强行替换反而会破坏客户端的缓存策略 + // （长 system prompt 被替换为 ~45 tokens 的短 prompt，低于 Anthropic 1024 token + // 最低缓存门槛，导致系统级缓存失效）。 + // + // 对于非 Claude Code 的第三方客户端（opencode 等），仍然走完整 mimicry。 + isClaudeCode := IsClaudeCodeClient(ctx) || isClaudeCodeClient(c.GetHeader("User-Agent"), parsed.MetadataUserID) + shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCode + + if shouldMimicClaudeCode { + // 与 Parrot 对齐：OAuth 账号无条件重写 system（即使客户端已发了 Claude Code + // 风格的 system prompt）。原因：第三方工具（opencode 等）会发 "You are Claude + // Code..." system prompt 但缺少 billing attribution block，导致 Anthropic + // 检测到"有 CC prompt 但无 billing block"的不一致而判为 third-party。 + // Parrot 的 transform_request 从不检查客户端 system 内容，直接覆盖。 + systemRewritten := false + if !strings.Contains(strings.ToLower(reqModel), "haiku") { + systemRaw, _ := parsed.SystemValue() + systemPromptInjectionEnabled, systemPrompt, systemPromptBlocks := s.claudeOAuthSystemPromptInjectionSettings(ctx) + if systemPromptInjectionEnabled { + if err := replaceBody(rewriteSystemForNonClaudeCodeWithPromptBlocks(body, systemRaw, systemPrompt, systemPromptBlocks)); err != nil { + return nil, err + } + systemRewritten = true + } + } + + // system 被重写时保留 CC prompt 的 cache_control: ephemeral（匹配真实 Claude Code 行为）； + // 未重写时（haiku / 注入开关关闭）剥离客户端 cache_control，与原有行为一致。 + // 两种情况下 enforceCacheControlLimit 都会兜底处理上限。 + normalizeOpts := claudeOAuthNormalizeOptions{stripSystemCacheControl: !systemRewritten} + if s.identityService != nil { + fp, err := s.identityService.GetOrCreateFingerprint(ctx, account.ID, c.Request.Header) + if err == nil && fp != nil { + // metadata 透传开启时跳过 metadata 注入 + _, mimicMPT, _ := s.settingService.GetGatewayForwardingSettings(ctx) + if !mimicMPT { + if metadataUserID := s.buildOAuthMetadataUserID(parsed, account, fp); metadataUserID != "" { + normalizeOpts.injectMetadata = true + normalizeOpts.metadataUserID = metadataUserID + } + } + } + } + + var normalizedBody []byte + normalizedBody, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) + if err := replaceBody(normalizedBody); err != nil { + return nil, err + } + + // D/E/F: 可选 messages cache 策略 + 工具名混淆 + tools[-1] 断点 + // 与 forward_as_chat_completions / forward_as_responses 路径对齐， + // 原生 /v1/messages 路径也走同一套可配置字段级改写。 + if err := replaceBody(s.rewriteMessageCacheControlIfEnabled(ctx, body)); err != nil { + return nil, err + } + if rw := buildToolNameRewriteFromBody(body); rw != nil { + if err := replaceBody(applyToolNameRewriteToBody(body, rw)); err != nil { + return nil, err + } + c.Set(toolNameRewriteKey, rw) + } else { + if err := replaceBody(applyToolsLastCacheBreakpoint(body)); err != nil { + return nil, err + } + } + } + + // 客户端 dateline 归一化：仅对 Anthropic OAuth/SetupToken 账号生效。 + // 抹除 "Today's date is …" 语句里可能被注入的隐写指纹（4 种撇号 × 2 种日期 + // 分隔符），还原为 ASCII 撇号 + "-" 分隔符。运行在 mimicry 分支之外， + // 保证真实 Claude Code 客户端注入的指纹同样被清洗。 + if next, ok := s.normalizeClientDatelineIfEnabled(ctx, account, body); ok { + if err := replaceBody(next); err != nil { + return nil, err + } + } + + // 强制执行 cache_control 块数量限制（最多 4 个） + if err := replaceBody(enforceCacheControlLimit(body)); err != nil { + return nil, err + } + + // 应用模型映射： + // - APIKey 账号：使用账号级别的显式映射（如果配置），否则透传原始模型名 + // - OAuth/SetupToken 账号：使用 Anthropic 标准映射（短ID → 长ID） + mappedModel := reqModel + mappingSource := "" + if account.Type == AccountTypeAPIKey { + mappedModel = account.GetMappedModel(reqModel) + if mappedModel != reqModel { + mappingSource = "account" + } + } + if mappingSource == "" && account.Platform == PlatformAnthropic && account.Type == AccountTypeServiceAccount { + if candidate, matched := account.ResolveMappedModel(reqModel); matched { + mappedModel = candidate + mappingSource = "account" + } else { + normalized := normalizeVertexAnthropicModelID(claude.NormalizeModelID(reqModel)) + if normalized != reqModel { + mappedModel = normalized + mappingSource = "vertex" + } + } + } + if mappingSource == "" && account.Platform == PlatformAnthropic && account.Type != AccountTypeAPIKey { + normalized := claude.NormalizeModelID(reqModel) + if normalized != reqModel { + mappedModel = normalized + mappingSource = "prefix" + } + } + if mappedModel != reqModel { + // 替换请求体中的模型名 + if err := replaceBody(s.replaceModelInBody(body, mappedModel)); err != nil { + return nil, err + } + reqModel = mappedModel + parsed.Model = mappedModel + logger.LegacyPrintf("service.gateway", "Model mapping applied: %s -> %s (account: %s, source=%s)", originalModel, mappedModel, account.Name, mappingSource) + } + + if s.shouldInjectAnthropicCacheTTL1h(ctx, account) { + if err := replaceBody(injectAnthropicCacheControlTTL1h(body)); err != nil { + return nil, err + } + } + + // 获取凭证 + token, tokenType, err := s.GetAccessToken(ctx, account) + if err != nil { + return nil, err + } + + // 获取代理URL（自定义 base URL 模式下，proxy 通过 buildCustomRelayURL 作为查询参数传递） + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + if !account.IsCustomBaseURLEnabled() || account.GetCustomBaseURL() == "" { + proxyURL = account.Proxy.URL() + } + } + + // 解析 TLS 指纹 profile（同一请求生命周期内不变，避免重试循环中重复解析） + tlsProfile := s.tlsFPProfileService.ResolveTLSProfile(account) + + // 调试日志：记录即将转发的账号信息 + logger.LegacyPrintf("service.gateway", "[Forward] Using account: ID=%d Name=%s Platform=%s Type=%s TLSFingerprint=%v Proxy=%s", + account.ID, account.Name, account.Platform, account.Type, tlsProfile, proxyURL) + // Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400. + if err := replaceBody(StripEmptyTextBlocks(body)); err != nil { + return nil, err + } + // Pre-filter: strip web-search history blocks the upstream cannot accept + // (emulation-synthesized server_tool_use / web_search_tool_result always; + // genuine ones additionally for passback-required upstreams). See + // FilterWebSearchHistoryBlocks. reqModel 此时已是映射后的模型 ID。 + if err := replaceBody(FilterWebSearchHistoryBlocks(body, reqModel)); err != nil { + return nil, err + } + // Pre-filter: remove thinking blocks with missing/invalid signatures before forwarding. + // Clients (e.g. Claude Code) sometimes send multi-turn conversations where a historical + // assistant message contains a thinking block that is missing the required "signature" field, + // causing upstream to reject the request with 400 "thinking.signature: Field required". + // FilterThinkingBlocks removes only the invalid blocks; thinking blocks with valid signatures + // are preserved. This avoids relying solely on the post-error retry path, which can time out + // (maxRetryElapsed = 10s) for long conversations before the retry budget is exhausted. + // + // 仅 anthropic-strict 模型族执行此过滤；passback-required 上游 (DeepSeek/Kimi/GLM 等) + // 要求历史 thinking block 原样回传，过滤反而制造 400。reqModel 此时已是映射后的模型 ID。 + if err := replaceBody(FilterThinkingBlocks(body, reqModel)); err != nil { + return nil, err + } + // Chinese LLM thinking.type 协议差异补正（如 MiniMax 只接受 adaptive；Anthropic-SDK + // 客户端默认发 enabled）。仅对 passback-required 上游生效（claude-* 不会进来）。 + if ResolveThinkingProtocol(reqModel) == ThinkingProtocolPassbackRequired { + if rewritten, applied := NormalizeChineseLLMThinking(body, reqModel); applied { + if err := replaceBody(rewritten); err != nil { + return nil, err + } + logger.LegacyPrintf("service.gateway", "Account %d: rewrote thinking.type for %s (Anthropic-SDK default 'enabled' -> vendor-specific)", account.ID, reqModel) + } + } + + // 重试循环 + var resp *http.Response + lastWireBody := body + retryStart := time.Now() + for attempt := 1; attempt <= maxRetryAttempts; attempt++ { + // 构建上游请求（每次重试需要重新构建，因为请求体需要重新读取） + upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream) + upstreamReq, wireBody, err := s.buildUpstreamRequest(upstreamCtx, c, account, body, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + releaseUpstreamCtx() + if err != nil { + return nil, err + } + // 记录本次实际发送的 wire body；只有请求成功后才写回 ParsedRequest，避免 400 retry 基于已签名 CCH 再改写。 + lastWireBody = wireBody + + // 发送请求 + resp, err = s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, tlsProfile) + if err != nil { + if resp != nil && resp.Body != nil { + _ = resp.Body.Close() + } + // Ensure the client receives an error response (handlers assume Forward writes on non-failover errors). + safeErr := sanitizeUpstreamErrorMessage(err.Error()) + setOpsUpstreamError(c, 0, safeErr, "") + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: 0, + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "request_error", + Message: safeErr, + }) + c.JSON(http.StatusBadGateway, gin.H{ + "type": "error", + "error": gin.H{ + "type": "upstream_error", + "message": "Upstream request failed", + }, + }) + return nil, fmt.Errorf("upstream request failed: %s", safeErr) + } + + // 优先检测thinking block签名错误（400）并重试一次 + if resp.StatusCode == 400 { + respBody, readErr := s.readUpstreamErrorBody(resp) + if readErr == nil { + _ = resp.Body.Close() + + if s.shouldRectifySignatureError(ctx, account, respBody, reqModel) { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "signature_error", + Message: extractUpstreamErrorMessage(respBody), + Detail: func() string { + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return "" + }(), + }) + + looksLikeToolSignatureError := func(msg string) bool { + m := strings.ToLower(msg) + return strings.Contains(m, "tool_use") || + strings.Contains(m, "tool_result") || + strings.Contains(m, "functioncall") || + strings.Contains(m, "function_call") || + strings.Contains(m, "functionresponse") || + strings.Contains(m, "function_response") + } + + // 避免在重试预算已耗尽时再发起额外请求 + if time.Since(retryStart) >= maxRetryElapsed { + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + break + } + logger.LegacyPrintf("service.gateway", "[warn] Account %d: thinking blocks have invalid signature, retrying with filtered blocks", account.ID) + + // Conservative two-stage fallback: + // 1) Disable thinking + thinking->text (preserve content) + // 2) Only if upstream still errors AND error message points to tool/function signature issues: + // also downgrade tool_use/tool_result blocks to text. + + filteredBody := FilterThinkingBlocksForRetry(body, reqModel) + retryCtx, releaseRetryCtx := detachStreamUpstreamContext(ctx, reqStream) + retryReq, retryWireBody, buildErr := s.buildUpstreamRequest(retryCtx, c, account, filteredBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + releaseRetryCtx() + if buildErr == nil { + retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, tlsProfile) + if retryErr == nil { + if retryResp.StatusCode < 400 { + // 重试请求被上游接受后同步 ParsedRequest，保证 usage/日志看到真实请求体。 + lastWireBody = retryWireBody + if err := replaceBody(retryWireBody); err != nil { + _ = retryResp.Body.Close() + return nil, err + } + logger.LegacyPrintf("service.gateway", "Account %d: thinking block retry succeeded (blocks downgraded)", account.ID) + resp = retryResp + break + } + + retryRespBody, retryReadErr := s.readUpstreamErrorBody(retryResp) + _ = retryResp.Body.Close() + if retryReadErr == nil && retryResp.StatusCode == 400 && s.isSignatureErrorPattern(ctx, account, retryRespBody) { + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: retryResp.StatusCode, + UpstreamRequestID: retryResp.Header.Get("x-request-id"), + UpstreamURL: safeUpstreamURL(retryReq.URL.String()), + Kind: "signature_retry_thinking", + Message: extractUpstreamErrorMessage(retryRespBody), + Detail: func() string { + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + return truncateString(string(retryRespBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return "" + }(), + }) + msg2 := extractUpstreamErrorMessage(retryRespBody) + if looksLikeToolSignatureError(msg2) && time.Since(retryStart) < maxRetryElapsed { + logger.LegacyPrintf("service.gateway", "Account %d: signature retry still failing and looks tool-related, retrying with tool blocks downgraded", account.ID) + filteredBody2 := FilterSignatureSensitiveBlocksForRetry(body, reqModel) + retryCtx2, releaseRetryCtx2 := detachStreamUpstreamContext(ctx, reqStream) + retryReq2, retryWireBody2, buildErr2 := s.buildUpstreamRequest(retryCtx2, c, account, filteredBody2, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + releaseRetryCtx2() + if buildErr2 == nil { + retryResp2, retryErr2 := s.httpUpstream.DoWithTLS(retryReq2, proxyURL, account.ID, account.Concurrency, tlsProfile) + if retryErr2 == nil { + if retryResp2.StatusCode < 400 { + // 二阶段工具块降级成功时也必须更新当前 body。 + lastWireBody = retryWireBody2 + if err := replaceBody(retryWireBody2); err != nil { + _ = retryResp2.Body.Close() + return nil, err + } + } + resp = retryResp2 + break + } + if retryResp2 != nil && retryResp2.Body != nil { + _ = retryResp2.Body.Close() + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: 0, + UpstreamURL: safeUpstreamURL(retryReq2.URL.String()), + Kind: "signature_retry_tools_request_error", + Message: sanitizeUpstreamErrorMessage(retryErr2.Error()), + }) + logger.LegacyPrintf("service.gateway", "Account %d: tool-downgrade signature retry failed: %v", account.ID, retryErr2) + } else { + logger.LegacyPrintf("service.gateway", "Account %d: tool-downgrade signature retry build failed: %v", account.ID, buildErr2) + } + } + } + + // Fall back to the original retry response context. + resp = &http.Response{ + StatusCode: retryResp.StatusCode, + Header: retryResp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(retryRespBody)), + } + break + } + if retryResp != nil && retryResp.Body != nil { + _ = retryResp.Body.Close() + } + logger.LegacyPrintf("service.gateway", "Account %d: signature error retry failed: %v", account.ID, retryErr) + } else { + logger.LegacyPrintf("service.gateway", "Account %d: signature error retry build request failed: %v", account.ID, buildErr) + } + + // Retry failed: restore original response body and continue handling. + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + break + } + // 不是签名错误（或整流器已关闭），继续检查 budget 约束 + errMsg := extractUpstreamErrorMessage(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"), + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "budget_constraint_error", + Message: errMsg, + Detail: func() string { + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return "" + }(), + }) + + rectifiedBody, applied := RectifyThinkingBudget(body) + if applied && time.Since(retryStart) < maxRetryElapsed { + logger.LegacyPrintf("service.gateway", "Account %d: detected budget_tokens constraint error, retrying with rectified budget (budget_tokens=%d, max_tokens=%d)", account.ID, BudgetRectifyBudgetTokens, BudgetRectifyMaxTokens) + budgetRetryCtx, releaseBudgetRetryCtx := detachStreamUpstreamContext(ctx, reqStream) + budgetRetryReq, budgetWireBody, buildErr := s.buildUpstreamRequest(budgetRetryCtx, c, account, rectifiedBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) + releaseBudgetRetryCtx() + if buildErr == nil { + budgetRetryResp, retryErr := s.httpUpstream.DoWithTLS(budgetRetryReq, proxyURL, account.ID, account.Concurrency, tlsProfile) + if retryErr == nil { + if budgetRetryResp.StatusCode < 400 { + // budget 修正请求成功后，ParsedRequest 也要描述被接受的修正版。 + lastWireBody = budgetWireBody + if err := replaceBody(budgetWireBody); err != nil { + _ = budgetRetryResp.Body.Close() + return nil, err + } + } + resp = budgetRetryResp + break + } + if budgetRetryResp != nil && budgetRetryResp.Body != nil { + _ = budgetRetryResp.Body.Close() + } + logger.LegacyPrintf("service.gateway", "Account %d: budget rectifier retry failed: %v", account.ID, retryErr) + } else { + logger.LegacyPrintf("service.gateway", "Account %d: budget rectifier retry build failed: %v", account.ID, buildErr) + } + } + } + + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + } + } + + // 检查是否需要通用重试（排除400，因为400已经在上面特殊处理过了） + if resp.StatusCode >= 400 && resp.StatusCode != 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) { + if attempt < maxRetryAttempts { + elapsed := time.Since(retryStart) + if elapsed >= maxRetryElapsed { + break + } + + delay := retryBackoffDelay(attempt) + remaining := maxRetryElapsed - elapsed + if delay > remaining { + delay = remaining + } + if delay <= 0 { + break + } + + respBody, _ := s.readUpstreamErrorBody(resp) + _ = resp.Body.Close() + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), + Kind: "retry", + Message: extractUpstreamErrorMessage(respBody), + Detail: func() string { + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return "" + }(), + }) + logger.LegacyPrintf("service.gateway", "Account %d: upstream error %d, retry %d/%d after %v (elapsed=%v/%v)", + account.ID, resp.StatusCode, attempt, maxRetryAttempts, delay, elapsed, maxRetryElapsed) + if err := sleepWithContext(ctx, delay); err != nil { + return nil, err + } + continue + } + // 最后一次尝试也失败，跳出循环处理重试耗尽 + break + } + + // 不需要重试（成功或不可重试的错误），跳出循环 + // DEBUG: 输出响应 headers（用于检测 rate limit 信息） + if account.Platform == PlatformGemini && resp.StatusCode < 400 && s.cfg != nil && s.cfg.Gateway.GeminiDebugResponseHeaders { + logger.LegacyPrintf("service.gateway", "[DEBUG] Gemini API Response Headers for account %d:", account.ID) + for k, v := range resp.Header { + logger.LegacyPrintf("service.gateway", "[DEBUG] %s: %v", k, v) + } + } + break + } + if resp == nil || resp.Body == nil { + return nil, errors.New("upstream request failed: empty response") + } + defer func() { _ = resp.Body.Close() }() + + // 处理重试耗尽的情况 + if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) { + if s.shouldFailoverUpstreamError(resp.StatusCode) { + respBody, _ := s.readUpstreamErrorBody(resp) + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + + // 调试日志：打印重试耗尽后的错误响应 + logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d RequestID=%s Body=%s", + account.ID, account.Name, resp.StatusCode, resp.Header.Get("x-request-id"), truncateString(string(respBody), 1000)) + + s.handleRetryExhaustedSideEffects(ctx, resp, account) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "retry_exhausted_failover", + Message: extractUpstreamErrorMessage(respBody), + Detail: func() string { + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return "" + }(), + }) + return nil, &UpstreamFailoverError{ + StatusCode: resp.StatusCode, + ResponseBody: respBody, + RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + } + } + return s.handleRetryExhaustedError(ctx, resp, c, account) + } + + // 处理可切换账号的错误 + if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) { + respBody, _ := s.readUpstreamErrorBody(resp) + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + + // 调试日志：打印上游错误响应 + logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (failover): Account=%d(%s) Status=%d RequestID=%s Body=%s", + account.ID, account.Name, resp.StatusCode, resp.Header.Get("x-request-id"), truncateString(string(respBody), 1000)) + + s.handleFailoverSideEffects(ctx, resp, account, reqModel) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "failover", + Message: extractUpstreamErrorMessage(respBody), + Detail: func() string { + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return "" + }(), + }) + return nil, &UpstreamFailoverError{ + StatusCode: resp.StatusCode, + ResponseBody: respBody, + RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + } + } + if resp.StatusCode >= 400 { + // 可选：对部分 400 触发 failover（默认关闭以保持语义） + if resp.StatusCode == 400 && s.cfg != nil && s.cfg.Gateway.FailoverOn400 { + respBody, readErr := s.readUpstreamErrorBody(resp) + if readErr != nil { + // ReadAll failed, fall back to normal error handling without consuming the stream + return s.handleErrorResponse(ctx, resp, c, account, reqModel) + } + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + + if s.shouldFailoverOn400(respBody) { + upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + upstreamDetail := "" + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes + if maxBytes <= 0 { + maxBytes = 2048 + } + upstreamDetail = truncateString(string(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: "failover_on_400", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + if s.cfg.Gateway.LogUpstreamErrorBody { + logger.LegacyPrintf("service.gateway", + "Account %d: 400 error, attempting failover: %s", + account.ID, + truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), + ) + } else { + logger.LegacyPrintf("service.gateway", "Account %d: 400 error, attempting failover", account.ID) + } + s.handleFailoverSideEffects(ctx, resp, account, reqModel) + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody} + } + } + return s.handleErrorResponse(ctx, resp, c, account, reqModel) + } + + // 处理正常响应 + + if !bytes.Equal(lastWireBody, body) { + // 成功后再同步最终 wire body，避免失败重试从已签名 CCH 的 body 继续派生。 + if err := replaceBody(lastWireBody); err != nil { + return nil, err + } + } + + // 触发上游接受回调（提前释放串行锁，不等流完成） + if parsed.OnUpstreamAccepted != nil { + parsed.OnUpstreamAccepted() + } + + var usage *ClaudeUsage + var firstTokenMs *int + var clientDisconnect bool + if reqStream { + streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, reqModel, shouldMimicClaudeCode) + if err != nil { + var sseErr *sseStreamErrorEventError + if errors.As(err, &sseErr) { + // 上游 HTTP 200 + SSE 流体内出现 event:error 帧。 + // 保留 StatusCode=403 以兼容既有 failover/客户端响应语义， + // 但补全 ResponseBody 与 ops 上下文，让运维日志能反映上游真实错误。 + body := []byte(sseErr.RawData) + + upstreamMsg := sanitizeUpstreamErrorMessage( + strings.TrimSpace(extractUpstreamErrorMessage(body)), + ) + + upstreamDetail := "" + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes + if maxBytes <= 0 { + maxBytes = 2048 + } + upstreamDetail = truncateString(sseErr.RawData, maxBytes) + } + + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: 403, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "stream_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + logger.LegacyPrintf("service.gateway", + "[Forward] SSE error event in stream: Account=%d(%s) RequestID=%s Body=%s", + account.ID, account.Name, resp.Header.Get("x-request-id"), + truncateString(sseErr.RawData, 1000), + ) + + return nil, &UpstreamFailoverError{ + StatusCode: 403, + ResponseBody: body, + } + } + return nil, err + } + usage = streamResult.usage + firstTokenMs = streamResult.firstTokenMs + clientDisconnect = streamResult.clientDisconnect + } else { + usage, err = s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, reqModel) + if err != nil { + return nil, err + } + } + + return &ForwardResult{ + RequestID: resp.Header.Get("x-request-id"), + Usage: *usage, + Model: originalModel, // 使用原始模型用于计费和日志 + UpstreamModel: mappedModel, + Stream: reqStream, + Duration: time.Since(startTime), + FirstTokenMs: firstTokenMs, + ClientDisconnect: clientDisconnect, + }, nil +} + +// ResolveChannelMapping 委托渠道服务解析模型映射 +func (s *GatewayService) ResolveChannelMapping(ctx context.Context, groupID int64, model string) ChannelMappingResult { + if s.channelService == nil { + return ChannelMappingResult{MappedModel: model} + } + return s.channelService.ResolveChannelMapping(ctx, groupID, model) +} + +// ReplaceModelInBody 替换请求体中的模型名（导出供 handler 使用） +func (s *GatewayService) ReplaceModelInBody(body []byte, newModel string) []byte { + return ReplaceModelInBody(body, newModel) +} + +// IsModelRestricted 检查模型是否被渠道限制 +func (s *GatewayService) IsModelRestricted(ctx context.Context, groupID int64, model string) bool { + if s.channelService == nil { + return false + } + return s.channelService.IsModelRestricted(ctx, groupID, model) +} + +// ResolveChannelMappingAndRestrict 解析渠道映射。 +// 模型限制检查已移至调度阶段（checkChannelPricingRestriction），restricted 始终返回 false。 +func (s *GatewayService) ResolveChannelMappingAndRestrict(ctx context.Context, groupID *int64, model string) (ChannelMappingResult, bool) { + if s.channelService == nil { + return ChannelMappingResult{MappedModel: model}, false + } + return s.channelService.ResolveChannelMappingAndRestrict(ctx, groupID, model) +} + +// checkChannelPricingRestriction 根据渠道计费基准检查模型是否受定价列表限制。 +// 供调度阶段预检查（requested / channel_mapped）。 +// upstream 需逐账号检查，此处返回 false。 +func (s *GatewayService) checkChannelPricingRestriction(ctx context.Context, groupID *int64, requestedModel string) bool { + if groupID == nil || s.channelService == nil || requestedModel == "" { + return false + } + mapping := s.channelService.ResolveChannelMapping(ctx, *groupID, requestedModel) + billingModel := billingModelForRestriction(mapping.BillingModelSource, requestedModel, mapping.MappedModel) + if billingModel == "" { + return false + } + return s.channelService.IsModelRestricted(ctx, *groupID, billingModel) +} + +// billingModelForRestriction 根据计费基准确定限制检查使用的模型。 +// upstream 返回空（需逐账号检查）。 +func billingModelForRestriction(source, requestedModel, channelMappedModel string) string { + switch source { + case BillingModelSourceRequested: + return requestedModel + case BillingModelSourceUpstream: + return "" + case BillingModelSourceChannelMapped: + return channelMappedModel + default: + return channelMappedModel + } +} + +// isUpstreamModelRestrictedByChannel 检查账号映射后的上游模型是否受渠道定价限制。 +// 仅在 BillingModelSource="upstream" 且 RestrictModels=true 时由调度循环调用。 +func (s *GatewayService) isUpstreamModelRestrictedByChannel(ctx context.Context, groupID int64, account *Account, requestedModel string) bool { + if s.channelService == nil { + return false + } + upstreamModel := resolveAccountUpstreamModel(account, requestedModel) + if upstreamModel == "" { + return false + } + return s.channelService.IsModelRestricted(ctx, groupID, upstreamModel) +} + +// resolveAccountUpstreamModel 确定账号将请求模型映射为什么上游模型。 +func resolveAccountUpstreamModel(account *Account, requestedModel string) string { + if account.Platform == PlatformAntigravity { + return mapAntigravityModel(account, requestedModel) + } + return account.GetMappedModel(requestedModel) +} + +// needsUpstreamChannelRestrictionCheck 判断是否需要在调度循环中逐账号检查上游模型的渠道限制。 +func (s *GatewayService) needsUpstreamChannelRestrictionCheck(ctx context.Context, groupID *int64) bool { + if groupID == nil || s.channelService == nil { + return false + } + ch, err := s.channelService.GetChannelForGroup(ctx, *groupID) + if err != nil { + slog.Warn("failed to check channel upstream restriction", "group_id", *groupID, "error", err) + return false + } + if ch == nil || !ch.RestrictModels { + return false + } + return ch.BillingModelSource == BillingModelSourceUpstream +} + +// isStickyAccountUpstreamRestricted 检查粘性会话命中的账号是否受 upstream 渠道限制。 +// 合并 needsUpstreamChannelRestrictionCheck + isUpstreamModelRestrictedByChannel 两步调用， +// 供 sticky session 条件链使用，避免内联多个函数调用导致行过长。 +func (s *GatewayService) isStickyAccountUpstreamRestricted(ctx context.Context, groupID *int64, account *Account, requestedModel string) bool { + if groupID == nil { + return false + } + if !s.needsUpstreamChannelRestrictionCheck(ctx, groupID) { + return false + } + return s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel) +} diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 79150f5598..c151375623 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -1,18 +1,13 @@ package service import ( - "bufio" "bytes" "context" - "crypto/sha256" "encoding/json" "errors" "fmt" - "io" "log/slog" - "net" "net/http" - "net/url" "os" "path/filepath" "regexp" @@ -20,27 +15,16 @@ import ( "strconv" "strings" "sync/atomic" - "syscall" "time" "unsafe" "github.com/Wei-Shaw/sub2api/internal/config" - "github.com/Wei-Shaw/sub2api/internal/pkg/anthropicfp" - "github.com/Wei-Shaw/sub2api/internal/pkg/claude" - "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" - infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors" "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" - "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" "github.com/cespare/xxhash/v2" - "github.com/google/uuid" gocache "github.com/patrickmn/go-cache" "github.com/tidwall/gjson" - "github.com/tidwall/sjson" "golang.org/x/sync/singleflight" - - "github.com/gin-gonic/gin" ) const ( @@ -1055,509 +1039,6 @@ func (s *GatewayService) hashContent(content string) string { return strconv.FormatUint(h, 36) } -type anthropicCacheControlPayload struct { - Type string `json:"type"` - TTL string `json:"ttl,omitempty"` -} - -type anthropicSystemTextBlockPayload struct { - Type string `json:"type"` - Text string `json:"text"` - CacheControl *anthropicCacheControlPayload `json:"cache_control,omitempty"` -} - -type anthropicMetadataPayload struct { - UserID string `json:"user_id"` -} - -// replaceModelInBody 替换请求体中的model字段 -// 优先使用定点修改，尽量保持客户端原始字段顺序。 -func (s *GatewayService) replaceModelInBody(body []byte, newModel string) []byte { - return ReplaceModelInBody(body, newModel) -} - -type claudeOAuthNormalizeOptions struct { - injectMetadata bool - metadataUserID string - stripSystemCacheControl bool -} - -// sanitizeSystemText rewrites only the fixed OpenCode identity sentence (if present). -// We intentionally avoid broad keyword replacement in system prompts to prevent -// accidentally changing user-provided instructions. -func sanitizeSystemText(text string) string { - if text == "" { - return text - } - // Some clients include a fixed OpenCode identity sentence. Anthropic may treat - // this as a non-Claude-Code fingerprint, so rewrite it to the canonical - // Claude Code banner before generic "OpenCode"/"opencode" replacements. - text = strings.ReplaceAll( - text, - "You are OpenCode, the best coding agent on the planet.", - strings.TrimSpace(claudeCodeSystemPrompt), - ) - return text -} - -func marshalAnthropicSystemTextBlock(text string, includeCacheControl bool) ([]byte, error) { - block := anthropicSystemTextBlockPayload{ - Type: "text", - Text: text, - } - if includeCacheControl { - block.CacheControl = &anthropicCacheControlPayload{ - Type: "ephemeral", - TTL: claude.DefaultCacheControlTTL, - } - } - return json.Marshal(block) -} - -func marshalAnthropicSystemTextBlockWithCacheControl(text string, cacheControl any) ([]byte, error) { - block := map[string]any{ - "type": "text", - "text": text, - } - if cacheControl != nil { - block["cache_control"] = cacheControl - } - return json.Marshal(block) -} - -func marshalAnthropicMetadata(userID string) ([]byte, error) { - return json.Marshal(anthropicMetadataPayload{UserID: userID}) -} - -func buildJSONArrayRaw(items [][]byte) []byte { - if len(items) == 0 { - return []byte("[]") - } - - total := 2 - for _, item := range items { - total += len(item) - } - total += len(items) - 1 - - buf := make([]byte, 0, total) - buf = append(buf, '[') - for i, item := range items { - if i > 0 { - buf = append(buf, ',') - } - buf = append(buf, item...) - } - buf = append(buf, ']') - return buf -} - -func setJSONValueBytes(body []byte, path string, value any) ([]byte, bool) { - next, err := sjson.SetBytes(body, path, value) - if err != nil { - return body, false - } - return next, true -} - -func setJSONRawBytes(body []byte, path string, raw []byte) ([]byte, bool) { - next, err := sjson.SetRawBytes(body, path, raw) - if err != nil { - return body, false - } - return next, true -} - -func deleteJSONPathBytes(body []byte, path string) ([]byte, bool) { - next, err := sjson.DeleteBytes(body, path) - if err != nil { - return body, false - } - return next, true -} - -func normalizeClaudeOAuthSystemBody(body []byte, opts claudeOAuthNormalizeOptions) ([]byte, bool) { - sys := gjson.GetBytes(body, "system") - if !sys.Exists() { - return body, false - } - - out := body - modified := false - - switch { - case sys.Type == gjson.String: - sanitized := sanitizeSystemText(sys.String()) - if sanitized != sys.String() { - if next, ok := setJSONValueBytes(out, "system", sanitized); ok { - out = next - modified = true - } - } - case sys.IsArray(): - index := 0 - sys.ForEach(func(_, item gjson.Result) bool { - if item.Get("type").String() == "text" { - textResult := item.Get("text") - if textResult.Exists() && textResult.Type == gjson.String { - text := textResult.String() - sanitized := sanitizeSystemText(text) - if sanitized != text { - if next, ok := setJSONValueBytes(out, fmt.Sprintf("system.%d.text", index), sanitized); ok { - out = next - modified = true - } - } - } - } - - if opts.stripSystemCacheControl && item.Get("cache_control").Exists() { - if next, ok := deleteJSONPathBytes(out, fmt.Sprintf("system.%d.cache_control", index)); ok { - out = next - modified = true - } - } - - index++ - return true - }) - } - - return out, modified -} - -func ensureClaudeOAuthMetadataUserID(body []byte, userID string) ([]byte, bool) { - if strings.TrimSpace(userID) == "" { - return body, false - } - - metadata := gjson.GetBytes(body, "metadata") - if !metadata.Exists() || metadata.Type == gjson.Null { - raw, err := marshalAnthropicMetadata(userID) - if err != nil { - return body, false - } - return setJSONRawBytes(body, "metadata", raw) - } - - trimmedRaw := strings.TrimSpace(metadata.Raw) - if strings.HasPrefix(trimmedRaw, "{") { - existing := metadata.Get("user_id") - if existing.Exists() && existing.Type == gjson.String && existing.String() != "" { - return body, false - } - return setJSONValueBytes(body, "metadata.user_id", userID) - } - - raw, err := marshalAnthropicMetadata(userID) - if err != nil { - return body, false - } - return setJSONRawBytes(body, "metadata", raw) -} - -func normalizeClaudeOAuthRequestBody(body []byte, modelID string, opts claudeOAuthNormalizeOptions) ([]byte, string) { - if len(body) == 0 { - return body, modelID - } - - out := body - modified := false - - if next, changed := normalizeClaudeOAuthSystemBody(out, opts); changed { - out = next - modified = true - } - - rawModel := gjson.GetBytes(out, "model") - if rawModel.Exists() && rawModel.Type == gjson.String { - normalized := claude.NormalizeModelID(rawModel.String()) - if normalized != rawModel.String() { - if next, ok := setJSONValueBytes(out, "model", normalized); ok { - out = next - modified = true - } - modelID = normalized - } - } - - // 确保 tools 字段存在（即使为空数组） - if !gjson.GetBytes(out, "tools").Exists() { - if next, ok := setJSONRawBytes(out, "tools", []byte("[]")); ok { - out = next - modified = true - } - } - - if opts.injectMetadata && opts.metadataUserID != "" { - if next, changed := ensureClaudeOAuthMetadataUserID(out, opts.metadataUserID); changed { - out = next - modified = true - } - } - - // temperature：真实 Claude Code CLI 总是发送 temperature（默认 1，客户端可覆盖）。 - // 之前的实现直接 delete 会导致 payload 缺字段，与真实 CLI 字节级不一致。 - // 策略：客户端传了什么就透传；没传则补默认 1。 - if !gjson.GetBytes(out, "temperature").Exists() { - if next, ok := setJSONValueBytes(out, "temperature", 1); ok { - out = next - modified = true - } - } - - // max_tokens：真实 CLI 的默认值是 128000。缺失时补齐以对齐指纹。 - if !gjson.GetBytes(out, "max_tokens").Exists() { - if next, ok := setJSONValueBytes(out, "max_tokens", 128000); ok { - out = next - modified = true - } - } - - // context_management：thinking.type 为 enabled/adaptive 时，真实 CLI 会自动 - // 附带 {"edits":[{"type":"clear_thinking_20251015","keep":"all"}]}。 - // 客户端显式传了就透传；否则按 CLI 行为补齐。 - // - // 注：本函数不按 model 名决定是否保留 context_management。“最终 beta - // header 不含 context-management-2025-06-27 时 strip 字段”的能力维度 - // 对称约束由 sanitizeAnthropicBodyForBetaTokens 在 buildUpstreamRequest / - // buildCountTokensRequest 层统一执行，与 Bedrock 路径的 - // sanitizeBedrockFieldsForBetaTokens 对称。 - if !gjson.GetBytes(out, "context_management").Exists() { - thinkingType := gjson.GetBytes(out, "thinking.type").String() - if thinkingType == "enabled" || thinkingType == "adaptive" { - const cmDefault = `{"edits":[{"type":"clear_thinking_20251015","keep":"all"}]}` - if next, ok := setJSONRawBytes(out, "context_management", []byte(cmDefault)); ok { - out = next - modified = true - } - } - } - - // tool_choice：与 Parrot 对齐，不再无条件删除。 - // - 客户端传了 {"type":"tool","name":"X"} → 保留结构，name 由 - // applyToolNameRewriteToBody 同步映射为假名 - // - 其他形态（auto/any/none）原样透传 - // 如果 body 里完全没有 tools（空数组），tool_choice 没意义时才删除 - if !gjson.GetBytes(out, "tools").IsArray() || len(gjson.GetBytes(out, "tools").Array()) == 0 { - if gjson.GetBytes(out, "tool_choice").Exists() { - if next, ok := deleteJSONPathBytes(out, "tool_choice"); ok { - out = next - modified = true - } - } - } - - if !modified { - return body, modelID - } - - return out, modelID -} - -func (s *GatewayService) buildOAuthMetadataUserID(parsed *ParsedRequest, account *Account, fp *Fingerprint) string { - if parsed == nil || account == nil { - return "" - } - if parsed.MetadataUserID != "" { - return "" - } - - userID := strings.TrimSpace(account.GetClaudeUserID()) - if userID == "" && fp != nil { - userID = fp.ClientID - } - if userID == "" { - // Fall back to a random, well-formed client id so we can still satisfy - // Claude Code OAuth requirements when account metadata is incomplete. - userID = generateClientID() - } - - // session_id 用"会话级稳定种子"派生（账号 + 客户端区分因子 + 首条 user 文本）： - // 随对话在尾部追加 messages 时保持不变，贴近真实 CC 进程级稳定的 session_id。 - // 不复用 GenerateSessionHash —— 后者是粘性路由键、按设计逐轮变化（见其测试）。 - var firstUserText string - if parsed.Body != nil { - firstUserText = extractFirstUserText(parsed.Body.Bytes()) - } - seed := buildStableSessionSeed(account.ID, sessionContextDiscriminator(parsed.SessionContext), firstUserText) - sessionID := generateSessionUUID(seed) - - // 根据指纹 UA 版本选择输出格式 - var uaVersion string - if fp != nil { - uaVersion = ExtractCLIVersion(fp.UserAgent) - } - accountUUID := strings.TrimSpace(account.GetExtraString("account_uuid")) - return FormatMetadataUserID(userID, accountUUID, sessionID, uaVersion) -} - -// applyClaudeCodeOAuthMimicryToBody 将"非 Claude Code 客户端 + Claude OAuth 账号" -// 路径上原本只在 /v1/messages 里做的完整伪装应用到任意 body 上。 -// -// 这是 /v1/messages 主路径上 rewriteSystemForNonClaudeCode + -// normalizeClaudeOAuthRequestBody 流程的通用版，供 OpenAI 协议兼容层 -// (ForwardAsChatCompletions / ForwardAsResponses) 复用。 -// -// 未抽离之前，OpenAI 协议兼容层仅做 injectClaudeCodePrompt（前置追加）， -// 而仓内 /v1/messages 路径自己的注释明确说过"仅前置追加无法通过 Anthropic -// 第三方检测"；那条注释就是本函数存在的根因。 -// -// 参数： -// - ctx / c：用于读取指纹和 gateway settings；c 可为 nil（如 count_tokens）。 -// - account：必须是 OAuth 账号，且调用方已判断不是 Claude Code 客户端。 -// - body：已经 marshal 成 Anthropic /v1/messages 格式的请求体。 -// - systemRaw：body 中原始 system 字段（用于判断是否需要 rewrite）。 -// - model：最终会发给上游的模型 ID（用于 haiku 旁路 + metadata 版本选择）。 -// -// 返回：改写后的 body。即使中间任何一步失败，也会退化成原 body（不会 panic）。 -func (s *GatewayService) applyClaudeCodeOAuthMimicryToBody( - ctx context.Context, - c *gin.Context, - account *Account, - body []byte, - systemRaw any, - model string, -) []byte { - if account == nil || !account.IsOAuth() || len(body) == 0 { - return body - } - - systemPromptInjectionEnabled, systemPrompt, systemPromptBlocks := s.claudeOAuthSystemPromptInjectionSettings(ctx) - systemRewritten := false - if systemPromptInjectionEnabled && !strings.Contains(strings.ToLower(model), "haiku") { - body = rewriteSystemForNonClaudeCodeWithPromptBlocks(body, normalizeSystemParam(systemRaw), systemPrompt, systemPromptBlocks) - systemRewritten = true - } - - normalizeOpts := claudeOAuthNormalizeOptions{stripSystemCacheControl: !systemRewritten} - - if s.identityService != nil && c != nil && c.Request != nil { - if fp, err := s.identityService.GetOrCreateFingerprint(ctx, account.ID, c.Request.Header); err == nil && fp != nil { - mimicMPT := false - if s.settingService != nil { - _, mimicMPT, _ = s.settingService.GetGatewayForwardingSettings(ctx) - } - if !mimicMPT { - if uid := s.buildOAuthMetadataUserIDFromBody(ctx, account, fp, body); uid != "" { - normalizeOpts.injectMetadata = true - normalizeOpts.metadataUserID = uid - } - } - } - } - - body, _ = normalizeClaudeOAuthRequestBody(body, model, normalizeOpts) - - // Phase D+E+F: messages cache 策略 + 工具名混淆 + tools[-1] 断点 - // 对齐 Parrot transform_request 里剩余的字段级改写。顺序有语义约束： - // 1) messages cache：仅在配置开启时清除客户端断点并注入代理断点 - // 2) tool rewrite：最后改 tools[*].name / tool_choice.name 并在 tools[-1] - // 上打断点；mapping 存入 gin.Context 供响应侧 bytes.Replace 还原。 - body = s.rewriteMessageCacheControlIfEnabled(ctx, body) - - if rw := buildToolNameRewriteFromBody(body); rw != nil { - body = applyToolNameRewriteToBody(body, rw) - if c != nil { - c.Set(toolNameRewriteKey, rw) - } - } else { - body = applyToolsLastCacheBreakpoint(body) - } - - return body -} - -// buildOAuthMetadataUserIDFromBody 是 buildOAuthMetadataUserID 的变体， -// 适用于调用方手上没有 ParsedRequest 的场景（如 OpenAI 协议兼容层）。 -// -// 与 buildOAuthMetadataUserID 的唯一区别： -// - session hash 从 body 本体按同样规则重算，而不是读取 ParsedRequest 缓存值。 -// - 如果 body 里已经存在 metadata.user_id，则返回空（由 ensureClaudeOAuthMetadataUserID -// 自行决定是否覆盖）。 -func (s *GatewayService) buildOAuthMetadataUserIDFromBody( - ctx context.Context, - account *Account, - fp *Fingerprint, - body []byte, -) string { - _ = ctx - if account == nil { - return "" - } - if existing := gjson.GetBytes(body, "metadata.user_id").String(); existing != "" { - return "" - } - - userID := strings.TrimSpace(account.GetClaudeUserID()) - if userID == "" && fp != nil { - userID = fp.ClientID - } - if userID == "" { - userID = generateClientID() - } - - // 与 buildOAuthMetadataUserID 一致：用会话级稳定种子，避免整 body 哈希导致 - // 每轮（甚至每个 token 变化）都重算出不同的 session_id。 - var clientDiscriminator string - if fp != nil { - clientDiscriminator = fp.ClientID - } - seed := buildStableSessionSeed(account.ID, clientDiscriminator, extractFirstUserText(body)) - sessionID := generateSessionUUID(seed) - - var uaVersion string - if fp != nil { - uaVersion = ExtractCLIVersion(fp.UserAgent) - } - accountUUID := strings.TrimSpace(account.GetExtraString("account_uuid")) - return FormatMetadataUserID(userID, accountUUID, sessionID, uaVersion) -} - -// buildStableSessionSeed 为伪装路径合成的 metadata.user_id session_id 生成"会话级稳定"种子。 -// -// 真实 Claude Code 的 session_id 是进程级随机 UUID，在一段会话内跨请求保持不变。无状态代理 -// 无法恢复该值，这里用"会话内不变的锚点"近似：账号 ID + 客户端区分因子 + 首条 user 消息文本。 -// 对话在尾部追加 messages 时这三者都不变，因此 generateSessionUUID(seed) 跨轮稳定。 -// -// 注意：粘性路由键 GenerateSessionHash 按设计逐轮变化（见其测试），本函数与之独立、互不影响。 -// accountID 恒存在，故 seed 永不为空 —— 输出始终是确定性 UUID，而非随机值。 -func buildStableSessionSeed(accountID int64, clientDiscriminator, firstUserText string) string { - var b strings.Builder - _, _ = b.WriteString(strconv.FormatInt(accountID, 10)) - _, _ = b.WriteString("::") - _, _ = b.WriteString(clientDiscriminator) - _, _ = b.WriteString("::") - _, _ = b.WriteString(firstUserText) - return b.String() -} - -// sessionContextDiscriminator 把请求上下文（客户端 IP / 归一化 UA / API Key ID）拼成 -// 一个跨客户端的区分因子，避免不同用户的相同首条消息派生出相同 session_id。 -func sessionContextDiscriminator(sc *SessionContext) string { - if sc == nil { - return "" - } - return sc.ClientIP + ":" + NormalizeSessionUserAgent(sc.UserAgent) + ":" + strconv.FormatInt(sc.APIKeyID, 10) -} - -// GenerateSessionUUID creates a deterministic UUID4 from a seed string. -func GenerateSessionUUID(seed string) string { - return generateSessionUUID(seed) -} - -func generateSessionUUID(seed string) string { - if seed == "" { - return uuid.NewString() - } - hash := sha256.Sum256([]byte(seed)) - bytes := hash[:16] - bytes[6] = (bytes[6] & 0x0f) | 0x40 - bytes[8] = (bytes[8] & 0x3f) | 0x80 - return fmt.Sprintf("%x-%x-%x-%x-%x", - bytes[0:4], bytes[4:6], bytes[6:8], bytes[8:10], bytes[10:16]) -} - // GetAccessToken 获取账号凭证 func (s *GatewayService) GetAccessToken(ctx context.Context, account *Account) (string, string, error) { switch account.Type { @@ -1608,5474 +1089,6 @@ func (s *GatewayService) getOAuthToken(ctx context.Context, account *Account) (s return accessToken, "oauth", nil } -// 重试相关常量 -const ( - // 最大尝试次数（包含首次请求）。过多重试会导致请求堆积与资源耗尽。 - maxRetryAttempts = 5 - - // 指数退避：第 N 次失败后的等待 = retryBaseDelay * 2^(N-1)，并且上限为 retryMaxDelay。 - retryBaseDelay = 300 * time.Millisecond - retryMaxDelay = 3 * time.Second - - // 最大重试耗时（包含请求本身耗时 + 退避等待时间）。 - // 用于防止极端情况下 goroutine 长时间堆积导致资源耗尽。 - maxRetryElapsed = 10 * time.Second -) - -func (s *GatewayService) shouldRetryUpstreamError(account *Account, statusCode int) bool { - // OAuth/Setup Token 账号：仅 403 重试 - if account.IsOAuth() { - return statusCode == 403 - } - - // API Key 账号：未配置的错误码重试 - return !account.ShouldHandleErrorCode(statusCode) -} - -// shouldFailoverUpstreamError determines whether an upstream error should trigger account failover. -func (s *GatewayService) shouldFailoverUpstreamError(statusCode int) bool { - switch statusCode { - case 401, 403, 429, 529: - return true - default: - return statusCode >= 500 - } -} - -func retryBackoffDelay(attempt int) time.Duration { - // attempt 从 1 开始，表示第 attempt 次请求刚失败，需要等待后进行第 attempt+1 次请求。 - if attempt <= 0 { - return retryBaseDelay - } - delay := retryBaseDelay * time.Duration(1<<(attempt-1)) - if delay > retryMaxDelay { - return retryMaxDelay - } - return delay -} - -func sleepWithContext(ctx context.Context, d time.Duration) error { - if d <= 0 { - return nil - } - timer := time.NewTimer(d) - defer func() { - if !timer.Stop() { - select { - case <-timer.C: - default: - } - } - }() - - select { - case <-ctx.Done(): - return ctx.Err() - case <-timer.C: - return nil - } -} - -// isClaudeCodeClient 判断请求是否来自真正的 Claude Code 客户端。 -// 判定条件： -// 1. User-Agent 匹配 claude-cli/X.Y.Z（大小写不敏感） -// 2. metadata.user_id 符合 Claude Code 格式（legacy 或 JSON 格式） -// -// 只检查 metadata.user_id 非空不够严格：第三方工具（opencode 等）可能伪造 UA -// 并附带任意 metadata.user_id 字符串，从而绕过 mimicry。必须通过 ParseMetadataUserID -// 验证格式才能确认是真正的 Claude Code 客户端。 -func isClaudeCodeClient(userAgent string, metadataUserID string) bool { - if !claudeCliUserAgentRe.MatchString(userAgent) { - return false - } - return ParseMetadataUserID(metadataUserID) != nil -} - -func shouldUseClaudeCodeNoopDeltaKeepalive(userAgent string) bool { - version := ExtractCLIVersion(userAgent) - if version == "" { - return false - } - return CompareVersions(version, claudeCodeNoopDeltaKeepaliveMinVersion) >= 0 -} - -func claudeCodeKeepaliveDeltaTypeForContentBlock(blockType string) string { - switch blockType { - case "text": - return "text_delta" - case "tool_use": - return "input_json_delta" - case "thinking": - return "thinking_delta" - default: - return "" - } -} - -func claudeCodeKeepaliveFieldForDeltaType(deltaType string) string { - switch deltaType { - case "text_delta": - return "text" - case "input_json_delta": - return "partial_json" - case "thinking_delta": - return "thinking" - default: - return "" - } -} - -func buildClaudeCodeNoopDeltaKeepalive(index int, deltaType string) (string, bool) { - fieldName := claudeCodeKeepaliveFieldForDeltaType(deltaType) - if fieldName == "" { - return "", false - } - return fmt.Sprintf("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":%d,\"delta\":{\"type\":\"%s\",\"%s\":\"\"}}\n\n", index, deltaType, fieldName), true -} - -func sseEventIndex(event map[string]any) (int, bool) { - switch v := event["index"].(type) { - case float64: - return int(v), true - case int: - return v, true - case int64: - return int(v), true - case json.Number: - i, err := v.Int64() - if err != nil { - return 0, false - } - return int(i), true - default: - return 0, false - } -} - -// normalizeSystemParam 将 json.RawMessage 类型的 system 参数转为标准 Go 类型（string / []any / nil）， -// 避免 type switch 中 json.RawMessage（底层 []byte）无法匹配 case string / case []any / case nil 的问题。 -// 这是 Go 的 typed nil 陷阱：(json.RawMessage, nil) ≠ (nil, nil)。 -func normalizeSystemParam(system any) any { - raw, ok := system.(json.RawMessage) - if !ok { - return system - } - if len(raw) == 0 { - return nil - } - var parsed any - if err := json.Unmarshal(raw, &parsed); err != nil { - return nil - } - return parsed -} - -// systemIncludesClaudeCodePrompt 检查 system 中是否已包含 Claude Code 提示词 -// 使用前缀匹配支持多种变体（标准版、Agent SDK 版等） -func systemIncludesClaudeCodePrompt(system any) bool { - system = normalizeSystemParam(system) - switch v := system.(type) { - case string: - return hasClaudeCodePrefix(v) - case []any: - for _, item := range v { - if m, ok := item.(map[string]any); ok { - if text, ok := m["text"].(string); ok && hasClaudeCodePrefix(text) { - return true - } - } - } - } - return false -} - -// hasClaudeCodePrefix 检查文本是否以 Claude Code 提示词的特征前缀开头 -func hasClaudeCodePrefix(text string) bool { - for _, prefix := range claudeCodePromptPrefixes { - if strings.HasPrefix(text, prefix) { - return true - } - } - return false -} - -// injectClaudeCodePrompt 在 system 开头注入 Claude Code 提示词 -// 处理 null、字符串、数组三种格式 -func injectClaudeCodePrompt(body []byte, system any) []byte { - system = normalizeSystemParam(system) - claudeCodeBlock, err := marshalAnthropicSystemTextBlock(claudeCodeSystemPrompt, true) - if err != nil { - logger.LegacyPrintf("service.gateway", "Warning: failed to build Claude Code prompt block: %v", err) - return body - } - // Opencode plugin applies an extra safeguard: it not only prepends the Claude Code - // banner, it also prefixes the next system instruction with the same banner plus - // a blank line. This helps when upstream concatenates system instructions. - claudeCodePrefix := strings.TrimSpace(claudeCodeSystemPrompt) - - var items [][]byte - - switch v := system.(type) { - case nil: - items = [][]byte{claudeCodeBlock} - case string: - // Be tolerant of older/newer clients that may differ only by trailing whitespace/newlines. - if strings.TrimSpace(v) == "" || strings.TrimSpace(v) == strings.TrimSpace(claudeCodeSystemPrompt) { - items = [][]byte{claudeCodeBlock} - } else { - // Mirror opencode behavior: keep the banner as a separate system entry, - // but also prefix the next system text with the banner. - merged := v - if !strings.HasPrefix(v, claudeCodePrefix) { - merged = claudeCodePrefix + "\n\n" + v - } - nextBlock, buildErr := marshalAnthropicSystemTextBlock(merged, false) - if buildErr != nil { - logger.LegacyPrintf("service.gateway", "Warning: failed to build prefixed Claude Code system block: %v", buildErr) - return body - } - items = [][]byte{claudeCodeBlock, nextBlock} - } - case []any: - items = make([][]byte, 0, len(v)+1) - items = append(items, claudeCodeBlock) - prefixedNext := false - systemResult := gjson.GetBytes(body, "system") - if systemResult.IsArray() { - systemResult.ForEach(func(_, item gjson.Result) bool { - textResult := item.Get("text") - if textResult.Exists() && textResult.Type == gjson.String && - strings.TrimSpace(textResult.String()) == strings.TrimSpace(claudeCodeSystemPrompt) { - return true - } - - raw := []byte(item.Raw) - // Prefix the first subsequent text system block once. - if !prefixedNext && item.Get("type").String() == "text" && textResult.Exists() && textResult.Type == gjson.String { - text := textResult.String() - if strings.TrimSpace(text) != "" && !strings.HasPrefix(text, claudeCodePrefix) { - next, setErr := sjson.SetBytes(raw, "text", claudeCodePrefix+"\n\n"+text) - if setErr == nil { - raw = next - prefixedNext = true - } - } - } - items = append(items, raw) - return true - }) - } else { - for _, item := range v { - m, ok := item.(map[string]any) - if !ok { - raw, marshalErr := json.Marshal(item) - if marshalErr == nil { - items = append(items, raw) - } - continue - } - if text, ok := m["text"].(string); ok && strings.TrimSpace(text) == strings.TrimSpace(claudeCodeSystemPrompt) { - continue - } - if !prefixedNext { - if blockType, _ := m["type"].(string); blockType == "text" { - if text, ok := m["text"].(string); ok && strings.TrimSpace(text) != "" && !strings.HasPrefix(text, claudeCodePrefix) { - m["text"] = claudeCodePrefix + "\n\n" + text - prefixedNext = true - } - } - } - raw, marshalErr := json.Marshal(m) - if marshalErr == nil { - items = append(items, raw) - } - } - } - default: - items = [][]byte{claudeCodeBlock} - } - - result, ok := setJSONRawBytes(body, "system", buildJSONArrayRaw(items)) - if !ok { - logger.LegacyPrintf("service.gateway", "Warning: failed to inject Claude Code prompt") - return body - } - return result -} - -// rewriteSystemForNonClaudeCode 将非 Claude Code 客户端的 system prompt 迁移至 messages， -// system 字段仅保留 Claude Code 标识提示词。 -// Anthropic 基于 system 参数内容检测第三方应用，仅前置追加 Claude Code 提示词 -// 无法通过检测，因为后续内容仍为非 Claude Code 格式。 -// 策略：将原始 system prompt 提取并注入为 user/assistant 消息对，system 仅保留 Claude Code 标识。 -func rewriteSystemForNonClaudeCode(body []byte, system any) []byte { - return rewriteSystemForNonClaudeCodeWithPromptBlocks(body, system, "", "") -} - -func rewriteSystemForNonClaudeCodeWithPrompt(body []byte, system any, expansionPrompt string) []byte { - return rewriteSystemForNonClaudeCodeWithPromptBlocks(body, system, expansionPrompt, "") -} - -type claudeOAuthSystemPromptBlockConfig struct { - Enabled *bool `json:"enabled,omitempty"` - Type string `json:"type,omitempty"` - Text string `json:"text,omitempty"` - CacheControl json.RawMessage `json:"cache_control,omitempty"` -} - -type claudeOAuthSystemPromptBlocksEnvelope struct { - Blocks []claudeOAuthSystemPromptBlockConfig `json:"blocks"` -} - -func defaultClaudeOAuthExpansionPrompt(expansionPrompt string) string { - expansionPrompt = strings.TrimSpace(expansionPrompt) - if expansionPrompt == "" { - return claudeCodeSystemPromptExpansion - } - return expansionPrompt -} - -func parseClaudeOAuthSystemPromptBlocksConfig(raw string) ([]claudeOAuthSystemPromptBlockConfig, error) { - raw = strings.TrimSpace(raw) - if raw == "" { - return nil, nil - } - if strings.HasPrefix(raw, "[") { - var blocks []claudeOAuthSystemPromptBlockConfig - if err := json.Unmarshal([]byte(raw), &blocks); err != nil { - return nil, err - } - return blocks, nil - } - var envelope claudeOAuthSystemPromptBlocksEnvelope - if err := json.Unmarshal([]byte(raw), &envelope); err != nil { - return nil, err - } - return envelope.Blocks, nil -} - -func decodeClaudeOAuthSystemPromptCacheControl(raw json.RawMessage) (any, error) { - trimmed := bytes.TrimSpace(raw) - if len(trimmed) == 0 || bytes.Equal(trimmed, []byte("null")) || bytes.Equal(trimmed, []byte("false")) { - return nil, nil - } - if bytes.Equal(trimmed, []byte("true")) { - return map[string]string{ - "type": "ephemeral", - "ttl": claude.DefaultCacheControlTTL, - }, nil - } - var value any - if err := json.Unmarshal(trimmed, &value); err != nil { - return nil, err - } - if _, ok := value.(map[string]any); !ok { - return nil, fmt.Errorf("cache_control must be boolean, null, or object") - } - return value, nil -} - -func expandClaudeOAuthSystemPromptTextTemplate(body []byte, text string, expansionPrompt string) (string, error) { - if text == "" { - return "", nil - } - expansionPrompt = defaultClaudeOAuthExpansionPrompt(expansionPrompt) - billingText, err := buildBillingAttributionText(body, claude.CLICurrentVersion) - if err != nil { - return "", err - } - fp := computeClaudeCodeFingerprint(body, claude.CLICurrentVersion) - replacer := strings.NewReplacer( - "{billing_header}", billingText, - "{cc_version}", claude.CLICurrentVersion, - "{fp}", fp, - "{claude_code_system_prompt}", claudeCodeSystemPrompt, - "{claude_code_expansion_prompt}", expansionPrompt, - ) - return replacer.Replace(text), nil -} - -func defaultClaudeOAuthSystemPromptBlockConfig() []claudeOAuthSystemPromptBlockConfig { - enabled := true - return []claudeOAuthSystemPromptBlockConfig{ - { - Enabled: &enabled, - Type: "text", - Text: "{billing_header}", - }, - { - Enabled: &enabled, - Type: "text", - Text: "{claude_code_system_prompt}", - }, - { - Enabled: &enabled, - Type: "text", - Text: "{claude_code_expansion_prompt}", - CacheControl: json.RawMessage( - fmt.Sprintf(`{"type":"ephemeral","ttl":%q}`, claude.DefaultCacheControlTTL), - ), - }, - } -} - -func buildClaudeOAuthSystemPromptBlocksJSON(body []byte, expansionPrompt string, blocksConfig string) ([][]byte, error) { - blocks, err := parseClaudeOAuthSystemPromptBlocksConfig(blocksConfig) - if err != nil { - return nil, err - } - if len(blocks) == 0 { - blocks = defaultClaudeOAuthSystemPromptBlockConfig() - } - - items := make([][]byte, 0, len(blocks)) - for i, block := range blocks { - if block.Enabled != nil && !*block.Enabled { - continue - } - blockType := strings.TrimSpace(block.Type) - if blockType == "" { - blockType = "text" - } - if blockType != "text" { - return nil, fmt.Errorf("system block %d type %q is not supported", i, block.Type) - } - text, err := expandClaudeOAuthSystemPromptTextTemplate(body, block.Text, expansionPrompt) - if err != nil { - return nil, err - } - if strings.TrimSpace(text) == "" { - continue - } - cacheControl, err := decodeClaudeOAuthSystemPromptCacheControl(block.CacheControl) - if err != nil { - return nil, fmt.Errorf("system block %d cache_control: %w", i, err) - } - raw, err := marshalAnthropicSystemTextBlockWithCacheControl(text, cacheControl) - if err != nil { - return nil, err - } - items = append(items, raw) - } - return items, nil -} - -func ValidateClaudeOAuthSystemPromptBlocksConfig(raw string) error { - if strings.TrimSpace(raw) == "" { - return nil - } - blocks, err := parseClaudeOAuthSystemPromptBlocksConfig(raw) - if err != nil { - return infraerrors.BadRequest("INVALID_CLAUDE_OAUTH_SYSTEM_PROMPT_BLOCKS", "claude oauth system prompt blocks must be valid JSON") - } - for i, block := range blocks { - blockType := strings.TrimSpace(block.Type) - if blockType == "" { - blockType = "text" - } - if blockType != "text" { - return infraerrors.BadRequest("INVALID_CLAUDE_OAUTH_SYSTEM_PROMPT_BLOCKS", fmt.Sprintf("system block %d type must be text", i)) - } - if _, err := decodeClaudeOAuthSystemPromptCacheControl(block.CacheControl); err != nil { - return infraerrors.BadRequest("INVALID_CLAUDE_OAUTH_SYSTEM_PROMPT_BLOCKS", fmt.Sprintf("system block %d cache_control is invalid", i)) - } - } - return nil -} - -func rewriteSystemForNonClaudeCodeWithPromptBlocks(body []byte, system any, expansionPrompt string, blocksConfig string) []byte { - system = normalizeSystemParam(system) - expansionPrompt = defaultClaudeOAuthExpansionPrompt(expansionPrompt) - - // 1. 提取原始 system prompt 文本 - var originalSystemText string - switch v := system.(type) { - case string: - originalSystemText = strings.TrimSpace(v) - case []any: - var parts []string - for _, item := range v { - if m, ok := item.(map[string]any); ok { - if text, ok := m["text"].(string); ok && strings.TrimSpace(text) != "" { - parts = append(parts, text) - } - } - } - originalSystemText = strings.Join(parts, "\n\n") - } - - // 2. 构造 system 数组，对齐真实 Claude Code CLI 的 3-block 形态： - // [0] billing attribution block（cc_version={cliVer}.{fp}; cc_entrypoint=cli;） - // [1] "You are Claude Code..." 身份前缀 block（默认不带 cache_control） - // [2] 工具无关的通用提示词扩充 block（带 cache_control 作为稳定缓存断点） - // - // 真实 CC 的 system 在身份前缀之后还有大段提示词，仅有 2 块会在块数/体量上明显 - // 区别于真实 CLI。这里注入 claudeCodeSystemPromptExpansion（中性段落）把形态做到 - // 接近真实，同时不注入会污染被代理用户行为的工具专属指令。 - // - // 缺失 billing block 的系统 payload 是 Anthropic 判定第三方的关键信号之一 - // （真实 CLI 每个请求都带）。新版 CLI 已取消 cch=... 签名字段，故 block 不再注入 - // cch（见 buildBillingAttributionText）。 - systemBlocks, blockErr := buildClaudeOAuthSystemPromptBlocksJSON(body, expansionPrompt, blocksConfig) - if blockErr != nil { - logger.LegacyPrintf("service.gateway", "Warning: failed to build configured Claude OAuth system blocks: %v", blockErr) - systemBlocks, blockErr = buildClaudeOAuthSystemPromptBlocksJSON(body, expansionPrompt, "") - } - if blockErr != nil { - logger.LegacyPrintf("service.gateway", "Warning: failed to build default Claude OAuth system blocks: %v", blockErr) - return body - } - out, ok := setJSONRawBytes(body, "system", buildJSONArrayRaw(systemBlocks)) - if !ok { - logger.LegacyPrintf("service.gateway", "Warning: failed to set Claude Code system prompt") - return body - } - - // 3. 将原始 system prompt 作为 user/assistant 消息对注入到 messages 开头 - // 模型仍通过 messages 接收完整指令，保留客户端功能 - ccPromptTrimmed := strings.TrimSpace(claudeCodeSystemPrompt) - if originalSystemText != "" && originalSystemText != ccPromptTrimmed && !hasClaudeCodePrefix(originalSystemText) { - instrMsg, err1 := json.Marshal(map[string]any{ - "role": "user", - "content": []map[string]any{ - {"type": "text", "text": "[System Instructions]\n" + originalSystemText}, - }, - }) - ackMsg, err2 := json.Marshal(map[string]any{ - "role": "assistant", - "content": []map[string]any{ - {"type": "text", "text": "Understood. I will follow these instructions."}, - }, - }) - if err1 != nil || err2 != nil { - logger.LegacyPrintf("service.gateway", "Warning: failed to marshal system-to-messages injection") - return out - } - - // 重建 messages 数组：[instruction, ack, ...originalMessages] - items := [][]byte{instrMsg, ackMsg} - messagesResult := gjson.GetBytes(out, "messages") - if messagesResult.IsArray() { - messagesResult.ForEach(func(_, msg gjson.Result) bool { - items = append(items, []byte(msg.Raw)) - return true - }) - } - - if next, setOk := setJSONRawBytes(out, "messages", buildJSONArrayRaw(items)); setOk { - out = next - } - } - - return out -} - -type cacheControlPath struct { - path string - log string -} - -func collectCacheControlPaths(body []byte) (invalidThinking []cacheControlPath, messagePaths []string, toolPaths []string, systemPaths []string) { - system := gjson.GetBytes(body, "system") - if system.IsArray() { - sysIndex := 0 - system.ForEach(func(_, item gjson.Result) bool { - if item.Get("cache_control").Exists() { - path := fmt.Sprintf("system.%d.cache_control", sysIndex) - if item.Get("type").String() == "thinking" { - invalidThinking = append(invalidThinking, cacheControlPath{ - path: path, - log: "[Warning] Removed illegal cache_control from thinking block in system", - }) - } else { - systemPaths = append(systemPaths, path) - } - } - sysIndex++ - return true - }) - } - - messages := gjson.GetBytes(body, "messages") - if messages.IsArray() { - msgIndex := 0 - messages.ForEach(func(_, msg gjson.Result) bool { - content := msg.Get("content") - if content.IsArray() { - contentIndex := 0 - content.ForEach(func(_, item gjson.Result) bool { - if item.Get("cache_control").Exists() { - path := fmt.Sprintf("messages.%d.content.%d.cache_control", msgIndex, contentIndex) - if item.Get("type").String() == "thinking" { - invalidThinking = append(invalidThinking, cacheControlPath{ - path: path, - log: fmt.Sprintf("[Warning] Removed illegal cache_control from thinking block in messages[%d].content[%d]", msgIndex, contentIndex), - }) - } else { - messagePaths = append(messagePaths, path) - } - } - contentIndex++ - return true - }) - } - msgIndex++ - return true - }) - } - - tools := gjson.GetBytes(body, "tools") - if tools.IsArray() { - toolIndex := 0 - tools.ForEach(func(_, tool gjson.Result) bool { - if tool.Get("cache_control").Exists() { - toolPaths = append(toolPaths, fmt.Sprintf("tools.%d.cache_control", toolIndex)) - } - toolIndex++ - return true - }) - } - - return invalidThinking, messagePaths, toolPaths, systemPaths -} - -// enforceCacheControlLimit 强制执行 cache_control 块数量限制（最多 4 个） -// 超限时优先移除工具断点，再移除 messages 断点，最后才移除 system 断点。 -func enforceCacheControlLimit(body []byte) []byte { - if len(body) == 0 { - return body - } - - invalidThinking, messagePaths, toolPaths, systemPaths := collectCacheControlPaths(body) - out := body - modified := false - - // 先清理 thinking 块中的非法 cache_control（thinking 块不支持该字段） - for _, item := range invalidThinking { - if !gjson.GetBytes(out, item.path).Exists() { - continue - } - next, ok := deleteJSONPathBytes(out, item.path) - if !ok { - continue - } - out = next - modified = true - logger.LegacyPrintf("service.gateway", "%s", item.log) - } - - count := len(messagePaths) + len(toolPaths) + len(systemPaths) - if count <= maxCacheControlBlocks { - if modified { - return out - } - return body - } - - // 超限：优先从 tools 中移除，再从 messages 中移除，最后才从 system 中移除。 - remaining := count - maxCacheControlBlocks - for i := len(toolPaths) - 1; i >= 0 && remaining > 0; i-- { - path := toolPaths[i] - if !gjson.GetBytes(out, path).Exists() { - continue - } - next, ok := deleteJSONPathBytes(out, path) - if !ok { - continue - } - out = next - modified = true - remaining-- - } - - for _, path := range messagePaths { - if remaining <= 0 { - break - } - if !gjson.GetBytes(out, path).Exists() { - continue - } - next, ok := deleteJSONPathBytes(out, path) - if !ok { - continue - } - out = next - modified = true - remaining-- - } - - for i := len(systemPaths) - 1; i >= 0 && remaining > 0; i-- { - path := systemPaths[i] - if !gjson.GetBytes(out, path).Exists() { - continue - } - next, ok := deleteJSONPathBytes(out, path) - if !ok { - continue - } - out = next - modified = true - remaining-- - } - - if modified { - return out - } - return body -} - -// injectAnthropicCacheControlTTL1h 将已有 ephemeral cache_control 块的 ttl 强制写为 1h。 -// 仅修改已经存在的 cache_control，不新增缓存断点。 -func injectAnthropicCacheControlTTL1h(body []byte) []byte { - return forceEphemeralCacheControlTTL(body, cacheTTLTarget1h) -} - -func forceEphemeralCacheControlTTL(body []byte, ttl string) []byte { - if len(body) == 0 || ttl == "" { - return body - } - out := body - var paths []string - addPath := func(path string, value gjson.Result) { - cc := value.Get("cache_control") - if !cc.Exists() || cc.Get("type").String() != "ephemeral" { - return - } - if cc.Get("ttl").String() == ttl { - return - } - paths = append(paths, path+".cache_control.ttl") - } - - if topCC := gjson.GetBytes(body, "cache_control"); topCC.Exists() && topCC.Get("type").String() == "ephemeral" && topCC.Get("ttl").String() != ttl { - paths = append(paths, "cache_control.ttl") - } - - system := gjson.GetBytes(body, "system") - if system.IsArray() { - idx := -1 - system.ForEach(func(_, block gjson.Result) bool { - idx++ - addPath(fmt.Sprintf("system.%d", idx), block) - return true - }) - } - - messages := gjson.GetBytes(body, "messages") - if messages.IsArray() { - msgIdx := -1 - messages.ForEach(func(_, msg gjson.Result) bool { - msgIdx++ - content := msg.Get("content") - if !content.IsArray() { - return true - } - contentIdx := -1 - content.ForEach(func(_, block gjson.Result) bool { - contentIdx++ - addPath(fmt.Sprintf("messages.%d.content.%d", msgIdx, contentIdx), block) - return true - }) - return true - }) - } - - tools := gjson.GetBytes(body, "tools") - if tools.IsArray() { - idx := -1 - tools.ForEach(func(_, tool gjson.Result) bool { - idx++ - addPath(fmt.Sprintf("tools.%d", idx), tool) - return true - }) - } - - for _, path := range paths { - if next, err := sjson.SetBytes(out, path, ttl); err == nil { - out = next - } - } - return out -} - -func (s *GatewayService) shouldInjectAnthropicCacheTTL1h(ctx context.Context, account *Account) bool { - if account == nil || !account.IsAnthropicOAuthOrSetupToken() || s == nil || s.settingService == nil { - return false - } - return s.settingService.IsAnthropicCacheTTL1hInjectionEnabled(ctx) -} - -// shouldNormalizeClientDateline reports whether the request body's client -// dateline should be normalized before forwarding to Anthropic. The switch is -// scoped to Anthropic OAuth/SetupToken accounts only; API-Key accounts and -// non-Anthropic platforms bypass this step entirely. -func (s *GatewayService) shouldNormalizeClientDateline(ctx context.Context, account *Account) bool { - if account == nil || !account.IsAnthropicOAuthOrSetupToken() || s == nil || s.settingService == nil { - return false - } - return s.settingService.IsClientDatelineNormalizationEnabled(ctx) -} - -// normalizeClientDatelineIfEnabled applies dateline normalization to body when -// the switch is on and the account qualifies. Returns (nextBody, true) only -// when the body actually changed; otherwise returns (nil, false) so callers -// can skip the writeback. -func (s *GatewayService) normalizeClientDatelineIfEnabled(ctx context.Context, account *Account, body []byte) ([]byte, bool) { - if !s.shouldNormalizeClientDateline(ctx, account) { - return nil, false - } - next, _, changed := anthropicfp.NormalizeDateline(body) - if !changed { - return nil, false - } - return next, true -} - -func (s *GatewayService) claudeOAuthSystemPromptInjectionSettings(ctx context.Context) (bool, string, string) { - if s == nil || s.settingService == nil { - return true, "", "" - } - return s.settingService.GetClaudeOAuthSystemPromptInjectionSettings(ctx) -} - -// Forward 转发请求到Claude API -func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, parsed *ParsedRequest) (*ForwardResult, error) { - startTime := time.Now() - if parsed == nil { - return nil, fmt.Errorf("parse request: empty request") - } - - // Web Search 模拟：纯 web_search 请求时，直接调用搜索 API 构造响应 - if account != nil && s.shouldEmulateWebSearch(ctx, account, parsed.GroupID, parsed.Body.Bytes()) { - return s.handleWebSearchEmulation(ctx, c, account, parsed) - } - - if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() { - passthroughBody := parsed.Body.Bytes() - passthroughModel := parsed.Model - if passthroughModel != "" { - if mappedModel := account.GetMappedModel(passthroughModel); mappedModel != passthroughModel { - passthroughBody = s.replaceModelInBody(passthroughBody, mappedModel) - logger.LegacyPrintf("service.gateway", "Passthrough model mapping: %s -> %s (account: %s)", parsed.Model, mappedModel, account.Name) - passthroughModel = mappedModel - } - } - return s.forwardAnthropicAPIKeyPassthroughWithInput(ctx, c, account, anthropicPassthroughForwardInput{ - Body: passthroughBody, - Parsed: parsed, - RequestModel: passthroughModel, - OriginalModel: parsed.Model, - RequestStream: parsed.Stream, - StartTime: startTime, - }) - } - - if account != nil && account.IsBedrock() { - return s.forwardBedrock(ctx, c, account, parsed, startTime) - } - - // Beta policy: evaluate once; block check + cache filter set for buildUpstreamRequest. - // Always overwrite the cache to prevent stale values from a previous retry with a different account. - if account.Platform == PlatformAnthropic && c != nil { - policy := s.evaluateBetaPolicy(ctx, c.GetHeader("anthropic-beta"), account, parsed.Model) - if policy.blockErr != nil { - return nil, policy.blockErr - } - filterSet := policy.filterSet - if filterSet == nil { - filterSet = map[string]struct{}{} - } - c.Set(betaPolicyFilterSetKey, filterSet) - } - - body := parsed.Body.Bytes() - replaceBody := func(next []byte) error { - if err := parsed.ReplaceBody(next); err != nil { - return fmt.Errorf("rewrite request body: %w", err) - } - body = parsed.Body.Bytes() - return nil - } - reqModel := parsed.Model - reqStream := parsed.Stream - originalModel := reqModel - - // === DEBUG: 打印客户端原始请求（headers + body 摘要）=== - if c != nil { - s.debugLogGatewaySnapshot("CLIENT_ORIGINAL", c.Request.Header, body, map[string]string{ - "account": fmt.Sprintf("%d(%s)", account.ID, account.Name), - "account_type": string(account.Type), - "model": reqModel, - "stream": strconv.FormatBool(reqStream), - }) - } - - // Claude Code 客户端判定：UA 匹配 claude-cli/* 且携带 metadata.user_id。 - // 真正的 Claude Code 客户端自带完整的 system prompt、cache_control 断点和 header， - // 不需要代理做任何 body 级别的 mimicry；强行替换反而会破坏客户端的缓存策略 - // （长 system prompt 被替换为 ~45 tokens 的短 prompt，低于 Anthropic 1024 token - // 最低缓存门槛，导致系统级缓存失效）。 - // - // 对于非 Claude Code 的第三方客户端（opencode 等），仍然走完整 mimicry。 - isClaudeCode := IsClaudeCodeClient(ctx) || isClaudeCodeClient(c.GetHeader("User-Agent"), parsed.MetadataUserID) - shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCode - - if shouldMimicClaudeCode { - // 与 Parrot 对齐：OAuth 账号无条件重写 system（即使客户端已发了 Claude Code - // 风格的 system prompt）。原因：第三方工具（opencode 等）会发 "You are Claude - // Code..." system prompt 但缺少 billing attribution block，导致 Anthropic - // 检测到"有 CC prompt 但无 billing block"的不一致而判为 third-party。 - // Parrot 的 transform_request 从不检查客户端 system 内容，直接覆盖。 - systemRewritten := false - if !strings.Contains(strings.ToLower(reqModel), "haiku") { - systemRaw, _ := parsed.SystemValue() - systemPromptInjectionEnabled, systemPrompt, systemPromptBlocks := s.claudeOAuthSystemPromptInjectionSettings(ctx) - if systemPromptInjectionEnabled { - if err := replaceBody(rewriteSystemForNonClaudeCodeWithPromptBlocks(body, systemRaw, systemPrompt, systemPromptBlocks)); err != nil { - return nil, err - } - systemRewritten = true - } - } - - // system 被重写时保留 CC prompt 的 cache_control: ephemeral（匹配真实 Claude Code 行为）； - // 未重写时（haiku / 注入开关关闭）剥离客户端 cache_control，与原有行为一致。 - // 两种情况下 enforceCacheControlLimit 都会兜底处理上限。 - normalizeOpts := claudeOAuthNormalizeOptions{stripSystemCacheControl: !systemRewritten} - if s.identityService != nil { - fp, err := s.identityService.GetOrCreateFingerprint(ctx, account.ID, c.Request.Header) - if err == nil && fp != nil { - // metadata 透传开启时跳过 metadata 注入 - _, mimicMPT, _ := s.settingService.GetGatewayForwardingSettings(ctx) - if !mimicMPT { - if metadataUserID := s.buildOAuthMetadataUserID(parsed, account, fp); metadataUserID != "" { - normalizeOpts.injectMetadata = true - normalizeOpts.metadataUserID = metadataUserID - } - } - } - } - - var normalizedBody []byte - normalizedBody, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) - if err := replaceBody(normalizedBody); err != nil { - return nil, err - } - - // D/E/F: 可选 messages cache 策略 + 工具名混淆 + tools[-1] 断点 - // 与 forward_as_chat_completions / forward_as_responses 路径对齐， - // 原生 /v1/messages 路径也走同一套可配置字段级改写。 - if err := replaceBody(s.rewriteMessageCacheControlIfEnabled(ctx, body)); err != nil { - return nil, err - } - if rw := buildToolNameRewriteFromBody(body); rw != nil { - if err := replaceBody(applyToolNameRewriteToBody(body, rw)); err != nil { - return nil, err - } - c.Set(toolNameRewriteKey, rw) - } else { - if err := replaceBody(applyToolsLastCacheBreakpoint(body)); err != nil { - return nil, err - } - } - } - - // 客户端 dateline 归一化：仅对 Anthropic OAuth/SetupToken 账号生效。 - // 抹除 "Today's date is …" 语句里可能被注入的隐写指纹（4 种撇号 × 2 种日期 - // 分隔符），还原为 ASCII 撇号 + "-" 分隔符。运行在 mimicry 分支之外， - // 保证真实 Claude Code 客户端注入的指纹同样被清洗。 - if next, ok := s.normalizeClientDatelineIfEnabled(ctx, account, body); ok { - if err := replaceBody(next); err != nil { - return nil, err - } - } - - // 强制执行 cache_control 块数量限制（最多 4 个） - if err := replaceBody(enforceCacheControlLimit(body)); err != nil { - return nil, err - } - - // 应用模型映射： - // - APIKey 账号：使用账号级别的显式映射（如果配置），否则透传原始模型名 - // - OAuth/SetupToken 账号：使用 Anthropic 标准映射（短ID → 长ID） - mappedModel := reqModel - mappingSource := "" - if account.Type == AccountTypeAPIKey { - mappedModel = account.GetMappedModel(reqModel) - if mappedModel != reqModel { - mappingSource = "account" - } - } - if mappingSource == "" && account.Platform == PlatformAnthropic && account.Type == AccountTypeServiceAccount { - if candidate, matched := account.ResolveMappedModel(reqModel); matched { - mappedModel = candidate - mappingSource = "account" - } else { - normalized := normalizeVertexAnthropicModelID(claude.NormalizeModelID(reqModel)) - if normalized != reqModel { - mappedModel = normalized - mappingSource = "vertex" - } - } - } - if mappingSource == "" && account.Platform == PlatformAnthropic && account.Type != AccountTypeAPIKey { - normalized := claude.NormalizeModelID(reqModel) - if normalized != reqModel { - mappedModel = normalized - mappingSource = "prefix" - } - } - if mappedModel != reqModel { - // 替换请求体中的模型名 - if err := replaceBody(s.replaceModelInBody(body, mappedModel)); err != nil { - return nil, err - } - reqModel = mappedModel - parsed.Model = mappedModel - logger.LegacyPrintf("service.gateway", "Model mapping applied: %s -> %s (account: %s, source=%s)", originalModel, mappedModel, account.Name, mappingSource) - } - - if s.shouldInjectAnthropicCacheTTL1h(ctx, account) { - if err := replaceBody(injectAnthropicCacheControlTTL1h(body)); err != nil { - return nil, err - } - } - - // 获取凭证 - token, tokenType, err := s.GetAccessToken(ctx, account) - if err != nil { - return nil, err - } - - // 获取代理URL（自定义 base URL 模式下，proxy 通过 buildCustomRelayURL 作为查询参数传递） - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - if !account.IsCustomBaseURLEnabled() || account.GetCustomBaseURL() == "" { - proxyURL = account.Proxy.URL() - } - } - - // 解析 TLS 指纹 profile（同一请求生命周期内不变，避免重试循环中重复解析） - tlsProfile := s.tlsFPProfileService.ResolveTLSProfile(account) - - // 调试日志：记录即将转发的账号信息 - logger.LegacyPrintf("service.gateway", "[Forward] Using account: ID=%d Name=%s Platform=%s Type=%s TLSFingerprint=%v Proxy=%s", - account.ID, account.Name, account.Platform, account.Type, tlsProfile, proxyURL) - // Pre-filter: strip empty text blocks (including nested in tool_result) to prevent upstream 400. - if err := replaceBody(StripEmptyTextBlocks(body)); err != nil { - return nil, err - } - // Pre-filter: strip web-search history blocks the upstream cannot accept - // (emulation-synthesized server_tool_use / web_search_tool_result always; - // genuine ones additionally for passback-required upstreams). See - // FilterWebSearchHistoryBlocks. reqModel 此时已是映射后的模型 ID。 - if err := replaceBody(FilterWebSearchHistoryBlocks(body, reqModel)); err != nil { - return nil, err - } - // Pre-filter: remove thinking blocks with missing/invalid signatures before forwarding. - // Clients (e.g. Claude Code) sometimes send multi-turn conversations where a historical - // assistant message contains a thinking block that is missing the required "signature" field, - // causing upstream to reject the request with 400 "thinking.signature: Field required". - // FilterThinkingBlocks removes only the invalid blocks; thinking blocks with valid signatures - // are preserved. This avoids relying solely on the post-error retry path, which can time out - // (maxRetryElapsed = 10s) for long conversations before the retry budget is exhausted. - // - // 仅 anthropic-strict 模型族执行此过滤；passback-required 上游 (DeepSeek/Kimi/GLM 等) - // 要求历史 thinking block 原样回传，过滤反而制造 400。reqModel 此时已是映射后的模型 ID。 - if err := replaceBody(FilterThinkingBlocks(body, reqModel)); err != nil { - return nil, err - } - // Chinese LLM thinking.type 协议差异补正（如 MiniMax 只接受 adaptive；Anthropic-SDK - // 客户端默认发 enabled）。仅对 passback-required 上游生效（claude-* 不会进来）。 - if ResolveThinkingProtocol(reqModel) == ThinkingProtocolPassbackRequired { - if rewritten, applied := NormalizeChineseLLMThinking(body, reqModel); applied { - if err := replaceBody(rewritten); err != nil { - return nil, err - } - logger.LegacyPrintf("service.gateway", "Account %d: rewrote thinking.type for %s (Anthropic-SDK default 'enabled' -> vendor-specific)", account.ID, reqModel) - } - } - - // 重试循环 - var resp *http.Response - lastWireBody := body - retryStart := time.Now() - for attempt := 1; attempt <= maxRetryAttempts; attempt++ { - // 构建上游请求（每次重试需要重新构建，因为请求体需要重新读取） - upstreamCtx, releaseUpstreamCtx := detachStreamUpstreamContext(ctx, reqStream) - upstreamReq, wireBody, err := s.buildUpstreamRequest(upstreamCtx, c, account, body, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) - releaseUpstreamCtx() - if err != nil { - return nil, err - } - // 记录本次实际发送的 wire body；只有请求成功后才写回 ParsedRequest，避免 400 retry 基于已签名 CCH 再改写。 - lastWireBody = wireBody - - // 发送请求 - resp, err = s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, tlsProfile) - if err != nil { - if resp != nil && resp.Body != nil { - _ = resp.Body.Close() - } - // Ensure the client receives an error response (handlers assume Forward writes on non-failover errors). - safeErr := sanitizeUpstreamErrorMessage(err.Error()) - setOpsUpstreamError(c, 0, safeErr, "") - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: 0, - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Kind: "request_error", - Message: safeErr, - }) - c.JSON(http.StatusBadGateway, gin.H{ - "type": "error", - "error": gin.H{ - "type": "upstream_error", - "message": "Upstream request failed", - }, - }) - return nil, fmt.Errorf("upstream request failed: %s", safeErr) - } - - // 优先检测thinking block签名错误（400）并重试一次 - if resp.StatusCode == 400 { - respBody, readErr := s.readUpstreamErrorBody(resp) - if readErr == nil { - _ = resp.Body.Close() - - if s.shouldRectifySignatureError(ctx, account, respBody, reqModel) { - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Kind: "signature_error", - Message: extractUpstreamErrorMessage(respBody), - Detail: func() string { - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) - } - return "" - }(), - }) - - looksLikeToolSignatureError := func(msg string) bool { - m := strings.ToLower(msg) - return strings.Contains(m, "tool_use") || - strings.Contains(m, "tool_result") || - strings.Contains(m, "functioncall") || - strings.Contains(m, "function_call") || - strings.Contains(m, "functionresponse") || - strings.Contains(m, "function_response") - } - - // 避免在重试预算已耗尽时再发起额外请求 - if time.Since(retryStart) >= maxRetryElapsed { - resp.Body = io.NopCloser(bytes.NewReader(respBody)) - break - } - logger.LegacyPrintf("service.gateway", "[warn] Account %d: thinking blocks have invalid signature, retrying with filtered blocks", account.ID) - - // Conservative two-stage fallback: - // 1) Disable thinking + thinking->text (preserve content) - // 2) Only if upstream still errors AND error message points to tool/function signature issues: - // also downgrade tool_use/tool_result blocks to text. - - filteredBody := FilterThinkingBlocksForRetry(body, reqModel) - retryCtx, releaseRetryCtx := detachStreamUpstreamContext(ctx, reqStream) - retryReq, retryWireBody, buildErr := s.buildUpstreamRequest(retryCtx, c, account, filteredBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) - releaseRetryCtx() - if buildErr == nil { - retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, tlsProfile) - if retryErr == nil { - if retryResp.StatusCode < 400 { - // 重试请求被上游接受后同步 ParsedRequest，保证 usage/日志看到真实请求体。 - lastWireBody = retryWireBody - if err := replaceBody(retryWireBody); err != nil { - _ = retryResp.Body.Close() - return nil, err - } - logger.LegacyPrintf("service.gateway", "Account %d: thinking block retry succeeded (blocks downgraded)", account.ID) - resp = retryResp - break - } - - retryRespBody, retryReadErr := s.readUpstreamErrorBody(retryResp) - _ = retryResp.Body.Close() - if retryReadErr == nil && retryResp.StatusCode == 400 && s.isSignatureErrorPattern(ctx, account, retryRespBody) { - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: retryResp.StatusCode, - UpstreamRequestID: retryResp.Header.Get("x-request-id"), - UpstreamURL: safeUpstreamURL(retryReq.URL.String()), - Kind: "signature_retry_thinking", - Message: extractUpstreamErrorMessage(retryRespBody), - Detail: func() string { - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - return truncateString(string(retryRespBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) - } - return "" - }(), - }) - msg2 := extractUpstreamErrorMessage(retryRespBody) - if looksLikeToolSignatureError(msg2) && time.Since(retryStart) < maxRetryElapsed { - logger.LegacyPrintf("service.gateway", "Account %d: signature retry still failing and looks tool-related, retrying with tool blocks downgraded", account.ID) - filteredBody2 := FilterSignatureSensitiveBlocksForRetry(body, reqModel) - retryCtx2, releaseRetryCtx2 := detachStreamUpstreamContext(ctx, reqStream) - retryReq2, retryWireBody2, buildErr2 := s.buildUpstreamRequest(retryCtx2, c, account, filteredBody2, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) - releaseRetryCtx2() - if buildErr2 == nil { - retryResp2, retryErr2 := s.httpUpstream.DoWithTLS(retryReq2, proxyURL, account.ID, account.Concurrency, tlsProfile) - if retryErr2 == nil { - if retryResp2.StatusCode < 400 { - // 二阶段工具块降级成功时也必须更新当前 body。 - lastWireBody = retryWireBody2 - if err := replaceBody(retryWireBody2); err != nil { - _ = retryResp2.Body.Close() - return nil, err - } - } - resp = retryResp2 - break - } - if retryResp2 != nil && retryResp2.Body != nil { - _ = retryResp2.Body.Close() - } - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: 0, - UpstreamURL: safeUpstreamURL(retryReq2.URL.String()), - Kind: "signature_retry_tools_request_error", - Message: sanitizeUpstreamErrorMessage(retryErr2.Error()), - }) - logger.LegacyPrintf("service.gateway", "Account %d: tool-downgrade signature retry failed: %v", account.ID, retryErr2) - } else { - logger.LegacyPrintf("service.gateway", "Account %d: tool-downgrade signature retry build failed: %v", account.ID, buildErr2) - } - } - } - - // Fall back to the original retry response context. - resp = &http.Response{ - StatusCode: retryResp.StatusCode, - Header: retryResp.Header.Clone(), - Body: io.NopCloser(bytes.NewReader(retryRespBody)), - } - break - } - if retryResp != nil && retryResp.Body != nil { - _ = retryResp.Body.Close() - } - logger.LegacyPrintf("service.gateway", "Account %d: signature error retry failed: %v", account.ID, retryErr) - } else { - logger.LegacyPrintf("service.gateway", "Account %d: signature error retry build request failed: %v", account.ID, buildErr) - } - - // Retry failed: restore original response body and continue handling. - resp.Body = io.NopCloser(bytes.NewReader(respBody)) - break - } - // 不是签名错误（或整流器已关闭），继续检查 budget 约束 - errMsg := extractUpstreamErrorMessage(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"), - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Kind: "budget_constraint_error", - Message: errMsg, - Detail: func() string { - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) - } - return "" - }(), - }) - - rectifiedBody, applied := RectifyThinkingBudget(body) - if applied && time.Since(retryStart) < maxRetryElapsed { - logger.LegacyPrintf("service.gateway", "Account %d: detected budget_tokens constraint error, retrying with rectified budget (budget_tokens=%d, max_tokens=%d)", account.ID, BudgetRectifyBudgetTokens, BudgetRectifyMaxTokens) - budgetRetryCtx, releaseBudgetRetryCtx := detachStreamUpstreamContext(ctx, reqStream) - budgetRetryReq, budgetWireBody, buildErr := s.buildUpstreamRequest(budgetRetryCtx, c, account, rectifiedBody, token, tokenType, reqModel, reqStream, shouldMimicClaudeCode) - releaseBudgetRetryCtx() - if buildErr == nil { - budgetRetryResp, retryErr := s.httpUpstream.DoWithTLS(budgetRetryReq, proxyURL, account.ID, account.Concurrency, tlsProfile) - if retryErr == nil { - if budgetRetryResp.StatusCode < 400 { - // budget 修正请求成功后，ParsedRequest 也要描述被接受的修正版。 - lastWireBody = budgetWireBody - if err := replaceBody(budgetWireBody); err != nil { - _ = budgetRetryResp.Body.Close() - return nil, err - } - } - resp = budgetRetryResp - break - } - if budgetRetryResp != nil && budgetRetryResp.Body != nil { - _ = budgetRetryResp.Body.Close() - } - logger.LegacyPrintf("service.gateway", "Account %d: budget rectifier retry failed: %v", account.ID, retryErr) - } else { - logger.LegacyPrintf("service.gateway", "Account %d: budget rectifier retry build failed: %v", account.ID, buildErr) - } - } - } - - resp.Body = io.NopCloser(bytes.NewReader(respBody)) - } - } - - // 检查是否需要通用重试（排除400，因为400已经在上面特殊处理过了） - if resp.StatusCode >= 400 && resp.StatusCode != 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) { - if attempt < maxRetryAttempts { - elapsed := time.Since(retryStart) - if elapsed >= maxRetryElapsed { - break - } - - delay := retryBackoffDelay(attempt) - remaining := maxRetryElapsed - elapsed - if delay > remaining { - delay = remaining - } - if delay <= 0 { - break - } - - respBody, _ := s.readUpstreamErrorBody(resp) - _ = resp.Body.Close() - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Kind: "retry", - Message: extractUpstreamErrorMessage(respBody), - Detail: func() string { - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) - } - return "" - }(), - }) - logger.LegacyPrintf("service.gateway", "Account %d: upstream error %d, retry %d/%d after %v (elapsed=%v/%v)", - account.ID, resp.StatusCode, attempt, maxRetryAttempts, delay, elapsed, maxRetryElapsed) - if err := sleepWithContext(ctx, delay); err != nil { - return nil, err - } - continue - } - // 最后一次尝试也失败，跳出循环处理重试耗尽 - break - } - - // 不需要重试（成功或不可重试的错误），跳出循环 - // DEBUG: 输出响应 headers（用于检测 rate limit 信息） - if account.Platform == PlatformGemini && resp.StatusCode < 400 && s.cfg != nil && s.cfg.Gateway.GeminiDebugResponseHeaders { - logger.LegacyPrintf("service.gateway", "[DEBUG] Gemini API Response Headers for account %d:", account.ID) - for k, v := range resp.Header { - logger.LegacyPrintf("service.gateway", "[DEBUG] %s: %v", k, v) - } - } - break - } - if resp == nil || resp.Body == nil { - return nil, errors.New("upstream request failed: empty response") - } - defer func() { _ = resp.Body.Close() }() - - // 处理重试耗尽的情况 - if resp.StatusCode >= 400 && s.shouldRetryUpstreamError(account, resp.StatusCode) { - if s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := s.readUpstreamErrorBody(resp) - _ = resp.Body.Close() - resp.Body = io.NopCloser(bytes.NewReader(respBody)) - - // 调试日志：打印重试耗尽后的错误响应 - logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (retry exhausted, failover): Account=%d(%s) Status=%d RequestID=%s Body=%s", - account.ID, account.Name, resp.StatusCode, resp.Header.Get("x-request-id"), truncateString(string(respBody), 1000)) - - s.handleRetryExhaustedSideEffects(ctx, resp, account) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "retry_exhausted_failover", - Message: extractUpstreamErrorMessage(respBody), - Detail: func() string { - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) - } - return "" - }(), - }) - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), - } - } - return s.handleRetryExhaustedError(ctx, resp, c, account) - } - - // 处理可切换账号的错误 - if resp.StatusCode >= 400 && s.shouldFailoverUpstreamError(resp.StatusCode) { - respBody, _ := s.readUpstreamErrorBody(resp) - _ = resp.Body.Close() - resp.Body = io.NopCloser(bytes.NewReader(respBody)) - - // 调试日志：打印上游错误响应 - logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (failover): Account=%d(%s) Status=%d RequestID=%s Body=%s", - account.ID, account.Name, resp.StatusCode, resp.Header.Get("x-request-id"), truncateString(string(respBody), 1000)) - - s.handleFailoverSideEffects(ctx, resp, account, reqModel) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "failover", - Message: extractUpstreamErrorMessage(respBody), - Detail: func() string { - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - return truncateString(string(respBody), s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) - } - return "" - }(), - }) - return nil, &UpstreamFailoverError{ - StatusCode: resp.StatusCode, - ResponseBody: respBody, - RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), - } - } - if resp.StatusCode >= 400 { - // 可选：对部分 400 触发 failover（默认关闭以保持语义） - if resp.StatusCode == 400 && s.cfg != nil && s.cfg.Gateway.FailoverOn400 { - respBody, readErr := s.readUpstreamErrorBody(resp) - if readErr != nil { - // ReadAll failed, fall back to normal error handling without consuming the stream - return s.handleErrorResponse(ctx, resp, c, account, reqModel) - } - _ = resp.Body.Close() - resp.Body = io.NopCloser(bytes.NewReader(respBody)) - - if s.shouldFailoverOn400(respBody) { - upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - upstreamDetail := "" - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes - if maxBytes <= 0 { - maxBytes = 2048 - } - upstreamDetail = truncateString(string(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: "failover_on_400", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - if s.cfg.Gateway.LogUpstreamErrorBody { - logger.LegacyPrintf("service.gateway", - "Account %d: 400 error, attempting failover: %s", - account.ID, - truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), - ) - } else { - logger.LegacyPrintf("service.gateway", "Account %d: 400 error, attempting failover", account.ID) - } - s.handleFailoverSideEffects(ctx, resp, account, reqModel) - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBody} - } - } - return s.handleErrorResponse(ctx, resp, c, account, reqModel) - } - - // 处理正常响应 - - if !bytes.Equal(lastWireBody, body) { - // 成功后再同步最终 wire body，避免失败重试从已签名 CCH 的 body 继续派生。 - if err := replaceBody(lastWireBody); err != nil { - return nil, err - } - } - - // 触发上游接受回调（提前释放串行锁，不等流完成） - if parsed.OnUpstreamAccepted != nil { - parsed.OnUpstreamAccepted() - } - - var usage *ClaudeUsage - var firstTokenMs *int - var clientDisconnect bool - if reqStream { - streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, reqModel, shouldMimicClaudeCode) - if err != nil { - var sseErr *sseStreamErrorEventError - if errors.As(err, &sseErr) { - // 上游 HTTP 200 + SSE 流体内出现 event:error 帧。 - // 保留 StatusCode=403 以兼容既有 failover/客户端响应语义， - // 但补全 ResponseBody 与 ops 上下文，让运维日志能反映上游真实错误。 - body := []byte(sseErr.RawData) - - upstreamMsg := sanitizeUpstreamErrorMessage( - strings.TrimSpace(extractUpstreamErrorMessage(body)), - ) - - upstreamDetail := "" - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes - if maxBytes <= 0 { - maxBytes = 2048 - } - upstreamDetail = truncateString(sseErr.RawData, maxBytes) - } - - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: 403, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "stream_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - logger.LegacyPrintf("service.gateway", - "[Forward] SSE error event in stream: Account=%d(%s) RequestID=%s Body=%s", - account.ID, account.Name, resp.Header.Get("x-request-id"), - truncateString(sseErr.RawData, 1000), - ) - - return nil, &UpstreamFailoverError{ - StatusCode: 403, - ResponseBody: body, - } - } - return nil, err - } - usage = streamResult.usage - firstTokenMs = streamResult.firstTokenMs - clientDisconnect = streamResult.clientDisconnect - } else { - usage, err = s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, reqModel) - if err != nil { - return nil, err - } - } - - return &ForwardResult{ - RequestID: resp.Header.Get("x-request-id"), - Usage: *usage, - Model: originalModel, // 使用原始模型用于计费和日志 - UpstreamModel: mappedModel, - Stream: reqStream, - Duration: time.Since(startTime), - FirstTokenMs: firstTokenMs, - ClientDisconnect: clientDisconnect, - }, nil -} - -func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, reqStream bool, mimicClaudeCode bool) (*http.Request, []byte, error) { - if account.Platform == PlatformAnthropic && account.Type == AccountTypeServiceAccount { - req, err := s.buildUpstreamRequestAnthropicVertex(ctx, c, account, body, token, modelID, reqStream) - return req, body, err - } - - // 确定目标URL - targetURL := claudeAPIURL - if account.Type == AccountTypeAPIKey { - baseURL := account.GetBaseURL() - if baseURL != "" { - validatedURL, err := s.validateUpstreamBaseURL(baseURL) - if err != nil { - return nil, nil, err - } - targetURL = validatedURL + "/v1/messages?beta=true" - } - } else if account.IsCustomBaseURLEnabled() { - customURL := account.GetCustomBaseURL() - if customURL == "" { - return nil, nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) - } - validatedURL, err := s.validateUpstreamBaseURL(customURL) - if err != nil { - return nil, nil, err - } - targetURL = s.buildCustomRelayURL(validatedURL, "/v1/messages", account) - } - - clientHeaders := http.Header{} - if c != nil && c.Request != nil { - clientHeaders = c.Request.Header - } - - // OAuth账号：应用统一指纹和metadata重写（受设置开关控制） - var fingerprint *Fingerprint - enableFP, enableMPT := true, false - if s.settingService != nil { - enableFP, enableMPT, _ = s.settingService.GetGatewayForwardingSettings(ctx) - } - if account.IsOAuth() && s.identityService != nil { - // 1. 获取或创建指纹（包含随机生成的ClientID） - fp, err := s.identityService.GetOrCreateFingerprint(ctx, account.ID, clientHeaders) - if err != nil { - logger.LegacyPrintf("service.gateway", "Warning: failed to get fingerprint for account %d: %v", account.ID, err) - // 失败时降级为透传原始headers - } else { - if enableFP { - fingerprint = fp - } - - // 2. 重写metadata.user_id（需要指纹中的ClientID和账号的account_uuid） - // 如果启用了会话ID伪装，会在重写后替换 session 部分为固定值 - // 当 metadata 透传开启时跳过重写 - if !enableMPT { - accountUUID := account.GetExtraString("account_uuid") - if accountUUID != "" && fp.ClientID != "" { - if newBody, err := s.identityService.RewriteUserIDWithMasking(ctx, body, account, accountUUID, fp.ClientID, fp.UserAgent); err == nil && len(newBody) > 0 { - body = newBody - } - } - } - } - } - - // 同步 billing header cc_version 与实际发送的 User-Agent 版本 - if fingerprint != nil { - body = syncBillingHeaderVersion(body, fingerprint.UserAgent) - } - - // === 计算最终 anthropic-beta header（先于 body sanitize 与 CCH 签名）=== - // - // 顺序约束： - // 1) 算 finalBeta（纯函数，不依赖 req.Header；mimicry 路径会忽略客户端 beta， - // 与原“OAuth + mimicClaudeCode 跳过白名单透传”行为对齐） - // 2) 按 finalBeta 做能力维度 body sanitize（如 context-management beta 缺失 → - // strip body.context_management，与 Bedrock 路径对称） - // 3) CCH 签名（必须使用 strip 后的 body，否则 hash 与最终 body 不一致 → - // 被 Anthropic 判 third-party） - // 4) NewRequest（body 至此最终敲定） - // 5) 透传白名单 / fingerprint / mimic header / 写入 finalBeta - policyFilterSet := s.getBetaPolicyFilterSet(ctx, c, account, modelID) - effectiveDropSet := mergeDropSets(policyFilterSet) - finalBetaHeader, finalBetaShouldSet := s.computeFinalAnthropicBeta( - tokenType, mimicClaudeCode, modelID, clientHeaders, body, effectiveDropSet, - ) - - // 账号覆写了 anthropic-beta 时，覆写值即最终上游值（由下方 ApplyHeaderOverrides 写入）： - // body 能力净化必须以覆写值为准，否则 header/body 不对称会被上游 400。 - if beta, ok := account.HeaderOverrideValue("anthropic-beta"); ok { - finalBetaHeader, finalBetaShouldSet = beta, true - } - - // 能力维度 body sanitize：与最终 anthropic-beta header 对称 - if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, finalBetaHeader); changed { - body = sanitized - } - - req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body)) - if err != nil { - return nil, nil, err - } - - // 设置认证头（保持原始大小写） - if tokenType == "oauth" { - setHeaderRaw(req.Header, "authorization", "Bearer "+token) - } else { - setAnthropicAPIKeyAuthHeader(req.Header, account, token) - } - - // 白名单透传 headers - // OAuth mimicry 路径：跳过客户端 header 透传，与 Parrot 对齐。 - // Parrot 的 build_upstream_headers 只发 9 个精确 header，不透传任何客户端 header。 - // 透传客户端 header 会引入不一致的 x-stainless-* / anthropic-beta / user-agent / - // x-claude-code-session-id 等值，和我们注入的伪装 header 冲突，被 Anthropic 判 third-party。 - if tokenType != "oauth" || !mimicClaudeCode { - for key, values := range clientHeaders { - lowerKey := strings.ToLower(key) - if allowedHeaders[lowerKey] { - wireKey := resolveWireCasing(key) - for _, v := range values { - addHeaderRaw(req.Header, wireKey, v) - } - } - } - } - - // OAuth账号：应用缓存的指纹到请求头（覆盖白名单透传的头） - if fingerprint != nil { - s.identityService.ApplyFingerprint(req, fingerprint) - } - - // 确保必要的headers存在（保持原始大小写） - if getHeaderRaw(req.Header, "content-type") == "" { - setHeaderRaw(req.Header, "content-type", "application/json") - } - if getHeaderRaw(req.Header, "anthropic-version") == "" { - setHeaderRaw(req.Header, "anthropic-version", "2023-06-01") - } - if tokenType == "oauth" { - applyClaudeOAuthHeaderDefaults(req) - } - - // OAuth + mimic Claude Code：强制注入 CLI 指纹相关 header - // （user-agent/x-stainless-*/x-app/Accept/x-stainless-helper-method/x-client-request-id） - if tokenType == "oauth" && mimicClaudeCode { - applyClaudeCodeMimicHeaders(req, reqStream) - } - - // 写入最终 anthropic-beta header - // 注：透传分支白名单可能写入了客户端 anthropic-beta，无条件 Del 一次再按 finalBeta - // 决定是否 set，确保 dropSet 过滤后的结果一定覆盖客户端原始值。 - deleteHeaderAllForms(req.Header, "anthropic-beta") - if finalBetaShouldSet { - setHeaderRaw(req.Header, "anthropic-beta", finalBetaHeader) - } - - // 同步 X-Claude-Code-Session-Id 头：取 body 中已处理的 metadata.user_id 的 session_id 覆盖 - if sessionHeader := getHeaderRaw(req.Header, "X-Claude-Code-Session-Id"); sessionHeader != "" { - if uid := gjson.GetBytes(body, "metadata.user_id").String(); uid != "" { - if parsed := ParseMetadataUserID(uid); parsed != nil { - setHeaderRaw(req.Header, "X-Claude-Code-Session-Id", parsed.SessionID) - } - } - } - - // 账号级请求头覆写（仅 anthropic/openai api_key 账号启用时生效；OAuth 路径 no-op）。 - // 放在所有 header 逻辑之后，确保配置值对同名头拥有最终决定权。 - account.ApplyHeaderOverrides(req.Header) - - // === DEBUG: 打印上游转发请求（headers + body 摘要），与 CLIENT_ORIGINAL 对比 === - s.debugLogGatewaySnapshot("UPSTREAM_FORWARD", req.Header, body, map[string]string{ - "url": req.URL.String(), - "token_type": tokenType, - "mimic_claude_code": strconv.FormatBool(mimicClaudeCode), - "fingerprint_applied": strconv.FormatBool(fingerprint != nil), - "enable_fp": strconv.FormatBool(enableFP), - "enable_mpt": strconv.FormatBool(enableMPT), - }) - - // Always capture a compact fingerprint line for later error diagnostics. - // We only print it when needed (or when the explicit debug flag is enabled). - if c != nil && tokenType == "oauth" { - c.Set(claudeMimicDebugInfoKey, buildClaudeMimicDebugLine(req, body, account, tokenType, mimicClaudeCode)) - } - if s.debugClaudeMimicEnabled() { - logClaudeMimicDebug(req, body, account, tokenType, mimicClaudeCode) - } - - return req, body, nil -} - -// vertexSupportedBetaTokens 是 Vertex AI 的 Anthropic 端点接受的 anthropic-beta -// 白名单。Vertex 对任何未知 token 直接 HTTP 400，故采用白名单（与 Bedrock 的 -// bedrockSupportedBetaTokens 同思路）而非黑名单：未来 Claude Code 新增的、Vertex 尚未 -// 支持的 token 天然被剥离。当 Vertex 新增支持某 beta 时在此补充。 -// -// 明确排除（issue #3358 中 Vertex 报 400 的 token）：advisor-tool-2026-03-01、 -// prompt-caching-scope-2026-01-05、redact-thinking-2026-02-12、 -// thinking-token-count-2026-05-13；以及 claude-code-20250219 / oauth-2025-04-20 等 -// 客户端身份 beta——Vertex service_account 走 Bearer 鉴权，不需要它们。 -var vertexSupportedBetaTokens = map[string]bool{ - "context-1m-2025-08-07": true, - "context-management-2025-06-27": true, - "fine-grained-tool-streaming-2025-05-14": true, - "interleaved-thinking-2025-05-14": true, -} - -// filterVertexBetaTokens 解析 client 的 anthropic-beta header，先剔除 drop 集合中的 -// token（BetaPolicy filter + 默认 drop），再只保留 Vertex 支持的 token，去重后逗号拼接。 -// 返回最终 header（可能为空字符串）。 -func filterVertexBetaTokens(header string, drop map[string]struct{}) string { - tokens := parseAnthropicBetaHeader(header) - if len(tokens) == 0 { - return "" - } - out := make([]string, 0, len(tokens)) - seen := make(map[string]bool, len(tokens)) - for _, t := range tokens { - if _, dropped := drop[t]; dropped { - continue - } - if !vertexSupportedBetaTokens[t] { - continue - } - if seen[t] { - continue - } - seen[t] = true - out = append(out, t) - } - return strings.Join(out, ",") -} - -func (s *GatewayService) buildUpstreamRequestAnthropicVertex( - ctx context.Context, - c *gin.Context, - account *Account, - body []byte, - token string, - modelID string, - reqStream bool, -) (*http.Request, error) { - vertexBody, err := buildVertexAnthropicRequestBody(body) - if err != nil { - return nil, err - } - - // 计算最终 outgoing anthropic-beta。Vertex AI 的 Anthropic 端点只接受一小撮 - // beta token，未知 token 会直接 HTTP 400——近期 Claude Code CLI 透传的 - // advisor-tool-2026-03-01 / prompt-caching-scope-2026-01-05 / - // redact-thinking-2026-02-12 / thinking-token-count-2026-05-13 都不被 Vertex 接受 - // （issue #3358）。这里复用 BetaPolicy 的 block 检查（与 Bedrock 的 - // resolveBedrockBetaTokensForRequest 对称），再按 vertexSupportedBetaTokens 白名单 - // 剥离其余 token，使该路径与 Anthropic 直连 / Bedrock 路径行为一致。 - clientBeta := "" - if c != nil && c.Request != nil { - clientBeta = getHeaderRaw(c.Request.Header, "anthropic-beta") - } - policy := s.evaluateBetaPolicy(ctx, clientBeta, account, modelID) - if policy.blockErr != nil { - return nil, policy.blockErr - } - finalBeta := filterVertexBetaTokens(clientBeta, mergeDropSets(policy.filterSet)) - - // 能力维度 sanitize：基于最终 beta（而非原始 client 值）决定是否保留 body 中的 - // context_management，与 Anthropic 直连 / Bedrock 路径对称。 - if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(vertexBody, finalBeta); changed { - vertexBody = sanitized - } - fullURL, err := buildVertexAnthropicURL(account.VertexProjectID(), account.VertexLocation(modelID), modelID, reqStream) - if err != nil { - return nil, err - } - req, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(vertexBody)) - if err != nil { - return nil, err - } - - if c != nil && c.Request != nil { - for key, values := range c.Request.Header { - lowerKey := strings.ToLower(strings.TrimSpace(key)) - if !allowedHeaders[lowerKey] || lowerKey == "anthropic-version" { - continue - } - wireKey := resolveWireCasing(key) - for _, v := range values { - addHeaderRaw(req.Header, wireKey, v) - } - } - } - - req.Header.Del("authorization") - req.Header.Del("x-api-key") - req.Header.Del("x-goog-api-key") - req.Header.Del("cookie") - req.Header.Del("anthropic-version") - setHeaderRaw(req.Header, "authorization", "Bearer "+token) - setHeaderRaw(req.Header, "content-type", "application/json") - - // 覆盖上面白名单 loop 写入的原始 client anthropic-beta，使用过滤后的最终值。 - // finalBeta 为空（全部被剥离）时不下发该 header，与 Vertex 无 beta 请求一致。 - deleteHeaderAllForms(req.Header, "anthropic-beta") - if finalBeta != "" { - setHeaderRaw(req.Header, "anthropic-beta", finalBeta) - } - - s.debugLogGatewaySnapshot("UPSTREAM_FORWARD_VERTEX_ANTHROPIC", req.Header, vertexBody, map[string]string{ - "url": req.URL.String(), - "token_type": "service_account", - "model": modelID, - "stream": strconv.FormatBool(reqStream), - }) - - return req, nil -} - -// getBetaHeader 处理anthropic-beta header -// 对于OAuth账号，需要确保包含oauth-2025-04-20 -func (s *GatewayService) getBetaHeader(modelID string, clientBetaHeader string) string { - // 如果客户端传了anthropic-beta - if clientBetaHeader != "" { - // 已包含oauth beta则直接返回 - if strings.Contains(clientBetaHeader, claude.BetaOAuth) { - return clientBetaHeader - } - - // 需要添加oauth beta - parts := strings.Split(clientBetaHeader, ",") - for i, p := range parts { - parts[i] = strings.TrimSpace(p) - } - - // 在claude-code-20250219后面插入oauth beta - claudeCodeIdx := -1 - for i, p := range parts { - if p == claude.BetaClaudeCode { - claudeCodeIdx = i - break - } - } - - if claudeCodeIdx >= 0 { - // 在claude-code后面插入 - newParts := make([]string, 0, len(parts)+1) - newParts = append(newParts, parts[:claudeCodeIdx+1]...) - newParts = append(newParts, claude.BetaOAuth) - newParts = append(newParts, parts[claudeCodeIdx+1:]...) - return strings.Join(newParts, ",") - } - - // 没有claude-code，放在第一位 - return claude.BetaOAuth + "," + clientBetaHeader - } - - // 客户端没传，根据模型生成 - // haiku 模型不需要 claude-code beta - if strings.Contains(strings.ToLower(modelID), "haiku") { - return claude.HaikuBetaHeader - } - - return claude.DefaultBetaHeader -} - -func requestNeedsBetaFeatures(body []byte) bool { - tools := gjson.GetBytes(body, "tools") - if tools.Exists() && tools.IsArray() && len(tools.Array()) > 0 { - return true - } - thinkingType := gjson.GetBytes(body, "thinking.type").String() - if strings.EqualFold(thinkingType, "enabled") || strings.EqualFold(thinkingType, "adaptive") { - return true - } - return false -} - -func defaultAPIKeyBetaHeader(body []byte) string { - modelID := gjson.GetBytes(body, "model").String() - if strings.Contains(strings.ToLower(modelID), "haiku") { - return claude.APIKeyHaikuBetaHeader - } - return claude.APIKeyBetaHeader -} - -func applyClaudeOAuthHeaderDefaults(req *http.Request) { - if req == nil { - return - } - if getHeaderRaw(req.Header, "Accept") == "" { - setHeaderRaw(req.Header, "Accept", "application/json") - } - for key, value := range claude.DefaultHeaders { - if value == "" { - continue - } - if getHeaderRaw(req.Header, key) == "" { - setHeaderRaw(req.Header, resolveWireCasing(key), value) - } - } -} - -func mergeAnthropicBeta(required []string, incoming string) string { - seen := make(map[string]struct{}, len(required)+8) - out := make([]string, 0, len(required)+8) - - add := func(v string) { - v = strings.TrimSpace(v) - if v == "" { - return - } - if _, ok := seen[v]; ok { - return - } - seen[v] = struct{}{} - out = append(out, v) - } - - for _, r := range required { - add(r) - } - for _, p := range strings.Split(incoming, ",") { - add(p) - } - return strings.Join(out, ",") -} - -func mergeAnthropicBetaDropping(required []string, incoming string, drop map[string]struct{}) string { - merged := mergeAnthropicBeta(required, incoming) - if merged == "" || len(drop) == 0 { - return merged - } - out := make([]string, 0, 8) - for _, p := range strings.Split(merged, ",") { - p = strings.TrimSpace(p) - if p == "" { - continue - } - if _, ok := drop[p]; ok { - continue - } - out = append(out, p) - } - return strings.Join(out, ",") -} - -// computeFinalAnthropicBeta 计算发往上游的最终 anthropic-beta header 值。 -// -// 设计动机：将原本在 buildUpstreamRequest 内联在一起、依赖 req.Header 的 -// anthropic-beta 计算逻辑抽成纯函数。这样调用方可以在 NewRequest 之前 -// 就提前拿到最终 beta header，进而能按它对 body 做能力维度 sanitize 后再做 -// CCH 签名——一举修复了以下之前由顺序依赖导致的能力维度 sanitize -// 无法部署的问题（签名与最终 body 不一致可以被判 third-party）。 -// -// 返回 (value, shouldSet)： -// - shouldSet=false 意为“不主动设置 anthropic-beta header”，与原代码“ -// API-key 账号 + 客户端未传 anthropic-beta + InjectBetaForAPIKey 未开启或 -// requestNeedsBetaFeatures=false”的行为对齐。 -// - shouldSet=true 时 value 可能为空字符串（例如客户端透传的 beta 被 dropSet -// 全部过滤掉），这与原代码中 setHeaderRaw 的结果一致。 -// -// clientHeaders 是客户端原始 HTTP header（通常为 c.Request.Header）；nil 时按“客户端 -// 未传”处理。body 是已经 metadata 重写 / billing version sync 之后但未 sanitize 上游 -// 不兼容字段之前的版本。 -func (s *GatewayService) computeFinalAnthropicBeta( - tokenType string, - mimicClaudeCode bool, - modelID string, - clientHeaders http.Header, - body []byte, - effectiveDropSet map[string]struct{}, -) (string, bool) { - clientBeta := "" - if clientHeaders != nil { - clientBeta = getHeaderRaw(clientHeaders, "anthropic-beta") - } - - if tokenType == "oauth" { - if mimicClaudeCode { - // mimic 路径：原代码跳过白名单透传，incomingBeta 总是空字符串。 - // 这里传空 string 以严格对齐原行为。 - requiredBetas := []string{claude.BetaOAuth, claude.BetaInterleavedThinking} - if !strings.Contains(strings.ToLower(modelID), "haiku") { - requiredBetas = claude.FullClaudeCodeMimicryBetas() - } - return mergeAnthropicBetaDropping(requiredBetas, "", effectiveDropSet), true - } - // 真 Claude Code 客户端透传路径 - return stripBetaTokensWithSet(s.getBetaHeader(modelID, clientBeta), effectiveDropSet), true - } - - // API-key accounts - if clientBeta != "" { - return stripBetaTokensWithSet(clientBeta, effectiveDropSet), true - } - if s.cfg != nil && s.cfg.Gateway.InjectBetaForAPIKey { - if requestNeedsBetaFeatures(body) { - if beta := defaultAPIKeyBetaHeader(body); beta != "" { - return beta, true - } - } - } - return "", false -} - -// computeFinalCountTokensAnthropicBeta 是 count_tokens 路径上 anthropic-beta header 的 -// 计算纯函数。语义与 computeFinalAnthropicBeta 对齐，但备份了 count_tokens 独有的 -// 两条特殊规则： -// -// - OAuth mimic：requiredBetas 为 FullClaudeCodeMimicryBetas + BetaTokenCounting -// （与 messages 不同的是：不按 haiku 排除；count_tokens 始终携带 token-counting beta） -// - OAuth 透传 + 客户端未传 anthropic-beta：补齐 CountTokensBetaHeader -// - OAuth 透传 + 客户端传了：补齐 BetaTokenCounting（如果未含） -// -// 返回语义同 computeFinalAnthropicBeta。 -func (s *GatewayService) computeFinalCountTokensAnthropicBeta( - tokenType string, - mimicClaudeCode bool, - modelID string, - clientHeaders http.Header, - body []byte, - effectiveDropSet map[string]struct{}, -) (string, bool) { - clientBeta := "" - if clientHeaders != nil { - clientBeta = getHeaderRaw(clientHeaders, "anthropic-beta") - } - - if tokenType == "oauth" { - if mimicClaudeCode { - // 与原代码严格等价：original buildCountTokensRequest 在 count_tokens mimic - // 分支上**不**会跳过白名单透传（与 messages mimic 路径不同），所以 - // incomingBeta = req.Header[anthropic-beta] = 客户端透传过来的 client beta。 - // 重构后直接从 clientHeaders 拿同一个值，保持行为一致。 - requiredBetas := append(claude.FullClaudeCodeMimicryBetas(), claude.BetaTokenCounting) - return mergeAnthropicBetaDropping(requiredBetas, clientBeta, effectiveDropSet), true - } - if clientBeta == "" { - return claude.CountTokensBetaHeader, true - } - beta := s.getBetaHeader(modelID, clientBeta) - if !strings.Contains(beta, claude.BetaTokenCounting) { - beta = beta + "," + claude.BetaTokenCounting - } - return stripBetaTokensWithSet(beta, effectiveDropSet), true - } - - // API-key accounts - if clientBeta != "" { - return stripBetaTokensWithSet(clientBeta, effectiveDropSet), true - } - if s.cfg != nil && s.cfg.Gateway.InjectBetaForAPIKey { - if requestNeedsBetaFeatures(body) { - if beta := defaultAPIKeyBetaHeader(body); beta != "" { - return beta, true - } - } - } - return "", false -} - -// stripBetaTokens removes the given beta tokens from a comma-separated header value. -func stripBetaTokens(header string, tokens []string) string { - if header == "" || len(tokens) == 0 { - return header - } - return stripBetaTokensWithSet(header, buildBetaTokenSet(tokens)) -} - -func stripBetaTokensWithSet(header string, drop map[string]struct{}) string { - if header == "" || len(drop) == 0 { - return header - } - parts := strings.Split(header, ",") - out := make([]string, 0, len(parts)) - for _, p := range parts { - p = strings.TrimSpace(p) - if p == "" { - continue - } - if _, ok := drop[p]; ok { - continue - } - out = append(out, p) - } - if len(out) == len(parts) { - return header // no change, avoid allocation - } - return strings.Join(out, ",") -} - -// BetaBlockedError indicates a request was blocked by a beta policy rule. -type BetaBlockedError struct { - Message string -} - -func (e *BetaBlockedError) Error() string { return e.Message } - -// betaPolicyResult holds the evaluated result of beta policy rules for a single request. -type betaPolicyResult struct { - blockErr *BetaBlockedError // non-nil if a block rule matched - filterSet map[string]struct{} // tokens to filter (may be nil) -} - -// evaluateBetaPolicy loads settings once and evaluates all rules against the given request. -func (s *GatewayService) evaluateBetaPolicy(ctx context.Context, betaHeader string, account *Account, model string) betaPolicyResult { - if s.settingService == nil { - return betaPolicyResult{} - } - settings, err := s.settingService.GetBetaPolicySettings(ctx) - if err != nil || settings == nil { - return betaPolicyResult{} - } - isOAuth := account.IsOAuth() - isBedrock := account.IsBedrock() - var result betaPolicyResult - for _, rule := range settings.Rules { - if !betaPolicyScopeMatches(rule.Scope, isOAuth, isBedrock) { - continue - } - effectiveAction, effectiveErrMsg := resolveRuleAction(rule, model) - switch effectiveAction { - case BetaPolicyActionBlock: - if result.blockErr == nil && betaHeader != "" && containsBetaToken(betaHeader, rule.BetaToken) { - msg := effectiveErrMsg - if msg == "" { - msg = "beta feature " + rule.BetaToken + " is not allowed" - } - result.blockErr = &BetaBlockedError{Message: msg} - } - case BetaPolicyActionFilter: - if result.filterSet == nil { - result.filterSet = make(map[string]struct{}) - } - result.filterSet[rule.BetaToken] = struct{}{} - } - } - return result -} - -// mergeDropSets merges the static defaultDroppedBetasSet with dynamic policy filter tokens. -// Returns defaultDroppedBetasSet directly when policySet is empty (zero allocation). -func mergeDropSets(policySet map[string]struct{}, extra ...string) map[string]struct{} { - if len(policySet) == 0 && len(extra) == 0 { - return defaultDroppedBetasSet - } - m := make(map[string]struct{}, len(defaultDroppedBetasSet)+len(policySet)+len(extra)) - for t := range defaultDroppedBetasSet { - m[t] = struct{}{} - } - for t := range policySet { - m[t] = struct{}{} - } - for _, t := range extra { - m[t] = struct{}{} - } - return m -} - -// betaPolicyFilterSetKey is the gin.Context key for caching the policy filter set within a request. -const betaPolicyFilterSetKey = "betaPolicyFilterSet" - -// getBetaPolicyFilterSet returns the beta policy filter set, using the gin context cache if available. -// In the /v1/messages path, Forward() evaluates the policy first and caches the result; -// buildUpstreamRequest reuses it (zero extra DB calls). In the count_tokens path, this -// evaluates on demand (one DB call). -func (s *GatewayService) getBetaPolicyFilterSet(ctx context.Context, c *gin.Context, account *Account, model string) map[string]struct{} { - if c != nil { - if v, ok := c.Get(betaPolicyFilterSetKey); ok { - if fs, ok := v.(map[string]struct{}); ok { - return fs - } - } - } - return s.evaluateBetaPolicy(ctx, "", account, model).filterSet -} - -// betaPolicyScopeMatches checks whether a rule's scope matches the current account type. -func betaPolicyScopeMatches(scope string, isOAuth bool, isBedrock bool) bool { - switch scope { - case BetaPolicyScopeAll: - return true - case BetaPolicyScopeOAuth: - return isOAuth - case BetaPolicyScopeAPIKey: - return !isOAuth && !isBedrock - case BetaPolicyScopeBedrock: - return isBedrock - default: - return true // unknown scope → match all (fail-open) - } -} - -// matchModelWhitelist checks if a model matches any pattern in the whitelist. -// Reuses matchModelPattern from group.go which supports exact and wildcard prefix matching. -func matchModelWhitelist(model string, whitelist []string) bool { - for _, pattern := range whitelist { - if matchModelPattern(pattern, model) { - return true - } - } - return false -} - -// resolveRuleAction determines the effective action and error message for a rule given the request model. -// When ModelWhitelist is empty, the rule's primary Action/ErrorMessage applies unconditionally. -// When non-empty, Action applies to matching models; FallbackAction/FallbackErrorMessage applies to others. -func resolveRuleAction(rule BetaPolicyRule, model string) (action, errorMessage string) { - if len(rule.ModelWhitelist) == 0 { - return rule.Action, rule.ErrorMessage - } - if matchModelWhitelist(model, rule.ModelWhitelist) { - return rule.Action, rule.ErrorMessage - } - if rule.FallbackAction != "" { - return rule.FallbackAction, rule.FallbackErrorMessage - } - return BetaPolicyActionPass, "" // default fallback: pass (fail-open) -} - -// droppedBetaSet returns claude.DroppedBetas as a set, with optional extra tokens. -func droppedBetaSet(extra ...string) map[string]struct{} { - m := make(map[string]struct{}, len(defaultDroppedBetasSet)+len(extra)) - for t := range defaultDroppedBetasSet { - m[t] = struct{}{} - } - for _, t := range extra { - m[t] = struct{}{} - } - return m -} - -// containsBetaToken checks if a comma-separated header value contains the given token. -func containsBetaToken(header, token string) bool { - if header == "" || token == "" { - return false - } - for _, p := range strings.Split(header, ",") { - if strings.TrimSpace(p) == token { - return true - } - } - return false -} - -func filterBetaTokens(tokens []string, filterSet map[string]struct{}) []string { - if len(tokens) == 0 || len(filterSet) == 0 { - return tokens - } - kept := make([]string, 0, len(tokens)) - for _, token := range tokens { - if _, filtered := filterSet[token]; !filtered { - kept = append(kept, token) - } - } - return kept -} - -func (s *GatewayService) resolveBedrockBetaTokensForRequest( - ctx context.Context, - account *Account, - betaHeader string, - body []byte, - modelID string, -) ([]string, error) { - // 1. 对原始 header 中的 beta token 做 block 检查（快速失败） - policy := s.evaluateBetaPolicy(ctx, betaHeader, account, modelID) - if policy.blockErr != nil { - return nil, policy.blockErr - } - - // 2. 解析 header + body 自动注入 + Bedrock 转换/过滤 - betaTokens := ResolveBedrockBetaTokens(betaHeader, body, modelID) - - // 3. 对最终 token 列表再做 block 检查，捕获通过 body 自动注入绕过 header block 的情况。 - // 例如：管理员 block 了 interleaved-thinking，客户端不在 header 中带该 token， - // 但请求体中包含 thinking 字段 → autoInjectBedrockBetaTokens 会自动补齐 → - // 如果不做此检查，block 规则会被绕过。 - if blockErr := s.checkBetaPolicyBlockForTokens(ctx, betaTokens, account, modelID); blockErr != nil { - return nil, blockErr - } - - return filterBetaTokens(betaTokens, policy.filterSet), nil -} - -// checkBetaPolicyBlockForTokens 检查 token 列表中是否有被管理员 block 规则命中的 token。 -// 用于补充 evaluateBetaPolicy 对 header 的检查，覆盖 body 自动注入的 token。 -func (s *GatewayService) checkBetaPolicyBlockForTokens(ctx context.Context, tokens []string, account *Account, model string) *BetaBlockedError { - if s.settingService == nil || len(tokens) == 0 { - return nil - } - settings, err := s.settingService.GetBetaPolicySettings(ctx) - if err != nil || settings == nil { - return nil - } - isOAuth := account.IsOAuth() - isBedrock := account.IsBedrock() - tokenSet := buildBetaTokenSet(tokens) - for _, rule := range settings.Rules { - effectiveAction, effectiveErrMsg := resolveRuleAction(rule, model) - if effectiveAction != BetaPolicyActionBlock { - continue - } - if !betaPolicyScopeMatches(rule.Scope, isOAuth, isBedrock) { - continue - } - if _, present := tokenSet[rule.BetaToken]; present { - msg := effectiveErrMsg - if msg == "" { - msg = "beta feature " + rule.BetaToken + " is not allowed" - } - return &BetaBlockedError{Message: msg} - } - } - return nil -} - -func buildBetaTokenSet(tokens []string) map[string]struct{} { - m := make(map[string]struct{}, len(tokens)) - for _, t := range tokens { - if t == "" { - continue - } - m[t] = struct{}{} - } - return m -} - -var defaultDroppedBetasSet = buildBetaTokenSet(claude.DroppedBetas) - -// applyClaudeCodeMimicHeaders forces "Claude Code-like" request headers. -// This mirrors opencode-anthropic-auth behavior: do not trust downstream -// headers when using Claude Code-scoped OAuth credentials. -func applyClaudeCodeMimicHeaders(req *http.Request, isStream bool) { - if req == nil { - return - } - // Start with the standard defaults (fill missing). - applyClaudeOAuthHeaderDefaults(req) - // Then force key headers to match Claude Code fingerprint regardless of what the client sent. - // 使用 resolveWireCasing 确保 key 与真实 wire format 一致（如 "x-app" 而非 "X-App"） - for key, value := range claude.DefaultHeaders { - if value == "" { - continue - } - setHeaderRaw(req.Header, resolveWireCasing(key), value) - } - // Real Claude CLI uses Accept: application/json (even for streaming). - setHeaderRaw(req.Header, "Accept", "application/json") - if isStream { - setHeaderRaw(req.Header, "x-stainless-helper-method", "stream") - } - // Real Claude CLI 每个请求都会生成一个新的 UUID 放在 x-client-request-id。 - // 上游会以此作为会话/请求指纹的一部分，缺失或重复都可能触发第三方判定。 - if getHeaderRaw(req.Header, "x-client-request-id") == "" { - setHeaderRaw(req.Header, "x-client-request-id", uuid.NewString()) - } -} - -func truncateForLog(b []byte, maxBytes int) string { - if maxBytes <= 0 { - maxBytes = 2048 - } - if len(b) > maxBytes { - b = b[:maxBytes] - } - s := string(b) - // 保持一行，避免污染日志格式 - s = strings.ReplaceAll(s, "\n", "\\n") - s = strings.ReplaceAll(s, "\r", "\\r") - return s -} - -// shouldRectifySignatureError 统一判断是否应触发签名整流（strip thinking blocks 并重试）。 -// 根据账号类型检查对应的开关和匹配模式。 -// -// mappedModel 用于按 thinking 协议族分流：passback-required (DeepSeek/Kimi/GLM 等) 上游 -// 的 400 不是签名缺失问题，retry 任何 thinking 变形都会破坏「原样回传」契约——直接透传 -// 错误给客户端。详见 thinking_protocol.go。 -func (s *GatewayService) shouldRectifySignatureError(ctx context.Context, account *Account, respBody []byte, mappedModel string) bool { - if !ShouldRectifyThinkingSignatureError(mappedModel) { - return false - } - if account.Type == AccountTypeAPIKey { - // API Key 账号：独立开关，一次读取配置 - settings, err := s.settingService.GetRectifierSettings(ctx) - if err != nil || !settings.Enabled || !settings.APIKeySignatureEnabled { - return false - } - // 先检查内置模式（同 OAuth），再检查自定义关键词 - if s.isThinkingBlockSignatureError(respBody) { - return true - } - return matchSignaturePatterns(respBody, settings.APIKeySignaturePatterns) - } - // OAuth/SetupToken/Upstream/Bedrock 等：保持原有行为（内置模式 + 原开关） - return s.isThinkingBlockSignatureError(respBody) && s.settingService.IsSignatureRectifierEnabled(ctx) -} - -// isSignatureErrorPattern 仅做模式匹配，不检查开关。 -// 用于已进入重试流程后的二阶段检测（此时开关已在首次调用时验证过）。 -func (s *GatewayService) isSignatureErrorPattern(ctx context.Context, account *Account, respBody []byte) bool { - if s.isThinkingBlockSignatureError(respBody) { - return true - } - if account.Type == AccountTypeAPIKey { - settings, err := s.settingService.GetRectifierSettings(ctx) - if err != nil { - return false - } - return matchSignaturePatterns(respBody, settings.APIKeySignaturePatterns) - } - return false -} - -// matchSignaturePatterns 检查响应体是否匹配自定义关键词列表（不区分大小写）。 -func matchSignaturePatterns(respBody []byte, patterns []string) bool { - if len(patterns) == 0 { - return false - } - bodyLower := strings.ToLower(string(respBody)) - for _, p := range patterns { - p = strings.TrimSpace(p) - if p == "" { - continue - } - if strings.Contains(bodyLower, strings.ToLower(p)) { - return true - } - } - return false -} - -// isThinkingBlockSignatureError 检测是否是thinking block相关错误 -// 这类错误可以通过过滤thinking blocks并重试来解决 -func (s *GatewayService) isThinkingBlockSignatureError(respBody []byte) bool { - msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) - if msg == "" { - return false - } - - // 检测signature相关的错误（更宽松的匹配） - // 例如: "Invalid `signature` in `thinking` block", "***.signature" 等 - if strings.Contains(msg, "signature") { - return true - } - - // 检测 thinking block 顺序/类型错误 - // 例如: "Expected `thinking` or `redacted_thinking`, but found `text`" - if strings.Contains(msg, "expected") && (strings.Contains(msg, "thinking") || strings.Contains(msg, "redacted_thinking")) { - logger.LegacyPrintf("service.gateway", "[SignatureCheck] Detected thinking block type error") - return true - } - - // 检测 thinking block 被修改的错误 - // 例如: "thinking or redacted_thinking blocks in the latest assistant message cannot be modified" - if strings.Contains(msg, "cannot be modified") && (strings.Contains(msg, "thinking") || strings.Contains(msg, "redacted_thinking")) { - logger.LegacyPrintf("service.gateway", "[SignatureCheck] Detected thinking block modification error") - return true - } - - // 检测空消息内容错误（可能是过滤 thinking blocks 后导致的，或客户端发送了空 text block） - // 例如: "all messages must have non-empty content" - // "messages: text content blocks must be non-empty" - if strings.Contains(msg, "non-empty content") || strings.Contains(msg, "empty content") || - strings.Contains(msg, "content blocks must be non-empty") { - logger.LegacyPrintf("service.gateway", "[SignatureCheck] Detected empty content error") - return true - } - - // 检测 thinking block 缺少 thinking 字段的错误（跨模型切换时常见： - // 其他模型回过的 assistant 历史里有 type=thinking 但没有 thinking 文本， - // 喂给开启 extended thinking 的 claude 时会被拒） - // 例如: "messages.1.content.0.thinking: each thinking block must contain thinking" - if strings.Contains(msg, "thinking block must contain") { - logger.LegacyPrintf("service.gateway", "[SignatureCheck] Detected thinking block missing content error") - return true - } - - return false -} - -func (s *GatewayService) shouldFailoverOn400(respBody []byte) bool { - // 只对"可能是兼容性差异导致"的 400 允许切换，避免无意义重试。 - // 默认保守：无法识别则不切换。 - msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) - if msg == "" { - return false - } - - // 缺少/错误的 beta header：换账号/链路可能成功（尤其是混合调度时）。 - // 更精确匹配 beta 相关的兼容性问题，避免误触发切换。 - if strings.Contains(msg, "anthropic-beta") || - strings.Contains(msg, "beta feature") || - strings.Contains(msg, "requires beta") { - return true - } - - // thinking/tool streaming 等兼容性约束（常见于中间转换链路） - if strings.Contains(msg, "thinking") || strings.Contains(msg, "thought_signature") || strings.Contains(msg, "signature") { - return true - } - if strings.Contains(msg, "tool_use") || strings.Contains(msg, "tool_result") || strings.Contains(msg, "tools") { - return true - } - - return false -} - -// sanitizeStreamError 返回不含网络地址的客户端可见错误描述。 -// 默认 (*net.OpError).Error() 会拼接 Source/Addr 字段，泄露内部 IP/端口与上游 -// 服务器地址（例如 "read tcp 10.0.0.1:54321->52.1.2.3:443: read: connection -// reset by peer"）。该函数只保留可识别的错误类别，原始 err 仍在调用点写入日志。 -func sanitizeStreamError(err error) string { - if err == nil { - return "" - } - switch { - case errors.Is(err, io.ErrUnexpectedEOF): - return "unexpected EOF" - case errors.Is(err, io.EOF): - return "EOF" - case errors.Is(err, context.Canceled): - return "canceled" - case errors.Is(err, context.DeadlineExceeded): - return "deadline exceeded" - case errors.Is(err, syscall.ECONNRESET): - return "connection reset by peer" - case errors.Is(err, syscall.ECONNABORTED): - return "connection aborted" - case errors.Is(err, syscall.ETIMEDOUT): - return "connection timed out" - case errors.Is(err, syscall.EPIPE): - return "broken pipe" - case errors.Is(err, syscall.ECONNREFUSED): - return "connection refused" - } - var netErr *net.OpError - if errors.As(err, &netErr) { - if netErr.Timeout() { - if netErr.Op != "" { - return netErr.Op + " timeout" - } - return "i/o timeout" - } - if netErr.Op != "" { - return netErr.Op + " network error" - } - } - return "upstream connection error" -} - -// ExtractUpstreamErrorMessage 从上游响应体中提取错误消息 -// 支持 Claude 风格的错误格式：{"type":"error","error":{"type":"...","message":"..."}} -func ExtractUpstreamErrorMessage(body []byte) string { - return extractUpstreamErrorMessage(body) -} - -func extractUpstreamErrorMessage(body []byte) string { - // Claude 风格：{"type":"error","error":{"type":"...","message":"..."}} - if m := gjson.GetBytes(body, "error.message").String(); strings.TrimSpace(m) != "" { - inner := strings.TrimSpace(m) - // 有些上游会把完整 JSON 作为字符串塞进 message - if strings.HasPrefix(inner, "{") { - if innerMsg := gjson.Get(inner, "error.message").String(); strings.TrimSpace(innerMsg) != "" { - return innerMsg - } - } - return m - } - - // ChatGPT 内部 API 风格：{"detail":"..."} - if d := gjson.GetBytes(body, "detail").String(); strings.TrimSpace(d) != "" { - return d - } - - // 兜底：尝试顶层 message - return gjson.GetBytes(body, "message").String() -} - -func extractUpstreamErrorCode(body []byte) string { - if code := strings.TrimSpace(gjson.GetBytes(body, "error.code").String()); code != "" { - return code - } - - inner := strings.TrimSpace(gjson.GetBytes(body, "error.message").String()) - if !strings.HasPrefix(inner, "{") { - return "" - } - - if code := strings.TrimSpace(gjson.Get(inner, "error.code").String()); code != "" { - return code - } - - if lastBrace := strings.LastIndex(inner, "}"); lastBrace >= 0 { - if code := strings.TrimSpace(gjson.Get(inner[:lastBrace+1], "error.code").String()); code != "" { - return code - } - } - - return "" -} - -func isCountTokensUnsupported404(statusCode int, body []byte) bool { - if statusCode != http.StatusNotFound { - return false - } - msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(body))) - if msg == "" { - return false - } - if strings.Contains(msg, "/v1/messages/count_tokens") { - return true - } - return strings.Contains(msg, "count_tokens") && strings.Contains(msg, "not found") -} - -func (s *GatewayService) readUpstreamErrorBody(resp *http.Response) ([]byte, error) { - if resp == nil || resp.Body == nil { - return nil, nil - } - limit := gatewayUpstreamErrorBodyReadLimit - if s != nil && s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody && s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) { - limit = int64(s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) - } - return io.ReadAll(io.LimitReader(resp.Body, limit)) -} - -func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, requestedModel ...string) (*ForwardResult, error) { - body, _ := s.readUpstreamErrorBody(resp) - - // 调试日志：打印上游错误响应 - logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (non-retryable): Account=%d(%s) Status=%d RequestID=%s Body=%s", - account.ID, account.Name, resp.StatusCode, resp.Header.Get("x-request-id"), truncateString(string(body), 1000)) - - upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - - // Print a compact upstream request fingerprint when we hit the Claude Code OAuth - // credential scope error. This avoids requiring env-var tweaks in a fixed deploy. - if isClaudeCodeCredentialScopeError(upstreamMsg) && c != nil { - if v, ok := c.Get(claudeMimicDebugInfoKey); ok { - if line, ok := v.(string); ok && strings.TrimSpace(line) != "" { - logger.LegacyPrintf("service.gateway", "[ClaudeMimicDebugOnError] status=%d request_id=%s %s", - resp.StatusCode, - resp.Header.Get("x-request-id"), - line, - ) - } - } - } - - // Enrich Ops error logs with upstream status + message, and optionally a truncated body snippet. - upstreamDetail := "" - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes - if maxBytes <= 0 { - maxBytes = 2048 - } - upstreamDetail = truncateString(string(body), maxBytes) - } - setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "http_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - // 处理上游错误，标记账号状态 - shouldDisable := false - if s.rateLimitService != nil { - if len(requestedModel) > 0 { - shouldDisable = s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel[0]) - } else { - shouldDisable = s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body) - } - } - if shouldDisable { - return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: body} - } - - MarkResponseCommitted(c) - - // 记录上游错误响应体摘要便于排障（可选：由配置控制；不回显到客户端） - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - logger.LegacyPrintf("service.gateway", - "Upstream error %d (account=%d platform=%s type=%s): %s", - resp.StatusCode, - account.ID, - account.Platform, - account.Type, - truncateForLog(body, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), - ) - } - - // 非 failover 错误也支持错误透传规则匹配。 - if status, errType, errMsg, matched := applyErrorPassthroughRule( - c, - account.Platform, - resp.StatusCode, - body, - http.StatusBadGateway, - "upstream_error", - "Upstream request failed", - ); matched { - c.JSON(status, gin.H{ - "type": "error", - "error": gin.H{ - "type": errType, - "message": errMsg, - }, - }) - - summary := upstreamMsg - if summary == "" { - summary = errMsg - } - if summary == "" { - return nil, fmt.Errorf("upstream error: %d (passthrough rule matched)", resp.StatusCode) - } - return nil, fmt.Errorf("upstream error: %d (passthrough rule matched) message=%s", resp.StatusCode, summary) - } - - // 根据状态码返回适当的自定义错误响应（不透传上游详细信息） - var errType, errMsg string - var statusCode int - - switch resp.StatusCode { - case 400: - c.Data(http.StatusBadRequest, "application/json", body) - summary := upstreamMsg - if summary == "" { - summary = truncateForLog(body, 512) - } - if summary == "" { - return nil, fmt.Errorf("upstream error: %d", resp.StatusCode) - } - return nil, fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, summary) - case 401: - statusCode = http.StatusBadGateway - errType = "upstream_error" - errMsg = "Upstream authentication failed, please contact administrator" - case 403: - statusCode = http.StatusBadGateway - errType = "upstream_error" - errMsg = "Upstream access forbidden, please contact administrator" - case 429: - statusCode = http.StatusTooManyRequests - errType = "rate_limit_error" - errMsg = "Upstream rate limit exceeded, please retry later" - case 529: - statusCode = http.StatusServiceUnavailable - errType = "overloaded_error" - errMsg = "Upstream service overloaded, please retry later" - case 500, 502, 503, 504: - statusCode = http.StatusBadGateway - errType = "upstream_error" - errMsg = "Upstream service temporarily unavailable" - 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 nil, fmt.Errorf("upstream error: %d", resp.StatusCode) - } - return nil, fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, upstreamMsg) -} - -func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, resp *http.Response, account *Account) { - body, _ := s.readUpstreamErrorBody(resp) - statusCode := resp.StatusCode - - // OAuth/Setup Token 账号的 403：标记账号异常 - if account.IsOAuth() && statusCode == 403 { - s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, resp.Header, body) - logger.LegacyPrintf("service.gateway", "Account %d: marked as error after %d retries for status %d", account.ID, maxRetryAttempts, statusCode) - } else { - // API Key 未配置错误码：不标记账号状态 - logger.LegacyPrintf("service.gateway", "Account %d: upstream error %d after %d retries (not marking account)", account.ID, statusCode, maxRetryAttempts) - } -} - -func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) { - body, _ := s.readUpstreamErrorBody(resp) - if len(requestedModel) > 0 { - s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel[0]) - return - } - s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body) -} - -// handleRetryExhaustedError 处理重试耗尽后的错误 -// OAuth 403：标记账号异常 -// API Key 未配置错误码：仅返回错误，不标记账号 -func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account) (*ForwardResult, error) { - MarkResponseCommitted(c) - // Capture upstream error body before side-effects consume the stream. - respBody, _ := s.readUpstreamErrorBody(resp) - _ = resp.Body.Close() - resp.Body = io.NopCloser(bytes.NewReader(respBody)) - - s.handleRetryExhaustedSideEffects(ctx, resp, account) - - upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - - if isClaudeCodeCredentialScopeError(upstreamMsg) && c != nil { - if v, ok := c.Get(claudeMimicDebugInfoKey); ok { - if line, ok := v.(string); ok && strings.TrimSpace(line) != "" { - logger.LegacyPrintf("service.gateway", "[ClaudeMimicDebugOnError] status=%d request_id=%s %s", - resp.StatusCode, - resp.Header.Get("x-request-id"), - line, - ) - } - } - } - - upstreamDetail := "" - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes - if maxBytes <= 0 { - maxBytes = 2048 - } - upstreamDetail = truncateString(string(respBody), maxBytes) - } - setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - Kind: "retry_exhausted", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - logger.LegacyPrintf("service.gateway", - "Upstream error %d retries_exhausted (account=%d platform=%s type=%s): %s", - resp.StatusCode, - account.ID, - account.Platform, - account.Type, - truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), - ) - } - - if status, errType, errMsg, matched := applyErrorPassthroughRule( - c, - account.Platform, - resp.StatusCode, - respBody, - http.StatusBadGateway, - "upstream_error", - "Upstream request failed after retries", - ); matched { - c.JSON(status, gin.H{ - "type": "error", - "error": gin.H{ - "type": errType, - "message": errMsg, - }, - }) - - summary := upstreamMsg - if summary == "" { - summary = errMsg - } - if summary == "" { - return nil, fmt.Errorf("upstream error: %d (retries exhausted, passthrough rule matched)", resp.StatusCode) - } - return nil, fmt.Errorf("upstream error: %d (retries exhausted, passthrough rule matched) message=%s", resp.StatusCode, summary) - } - - // 返回统一的重试耗尽错误响应 - c.JSON(http.StatusBadGateway, gin.H{ - "type": "error", - "error": gin.H{ - "type": "upstream_error", - "message": "Upstream request failed after retries", - }, - }) - - if upstreamMsg == "" { - return nil, fmt.Errorf("upstream error: %d (retries exhausted)", resp.StatusCode) - } - return nil, fmt.Errorf("upstream error: %d (retries exhausted) message=%s", resp.StatusCode, upstreamMsg) -} - -// streamingResult 流式响应结果 -type streamingResult struct { - usage *ClaudeUsage - firstTokenMs *int - clientDisconnect bool // 客户端是否在流式传输过程中断开 -} - -func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel string, mimicClaudeCode bool) (*streamingResult, error) { - // 更新5h窗口状态 - s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header) - - if s.responseHeaderFilter != nil { - responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) - } - - // 设置SSE响应头 - c.Header("Content-Type", "text/event-stream") - c.Header("Cache-Control", "no-cache") - c.Header("Connection", "keep-alive") - c.Header("X-Accel-Buffering", "no") - - // 透传其他响应头 - if v := resp.Header.Get("x-request-id"); v != "" { - c.Header("x-request-id", v) - } - - w := c.Writer - flusher, ok := w.(http.Flusher) - if !ok { - return nil, errors.New("streaming not supported") - } - - usage := &ClaudeUsage{} - var firstTokenMs *int - scanner := bufio.NewScanner(resp.Body) - // 设置更大的buffer以处理长行 - maxLineSize := defaultMaxLineSize - if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 { - maxLineSize = s.cfg.Gateway.MaxLineSize - } - scanBuf := getSSEScannerBuf64K() - scanner.Buffer(scanBuf[:0], maxLineSize) - - type scanEvent struct { - line string - err error - } - // 独立 goroutine 读取上游，避免读取阻塞导致超时/keepalive无法处理 - 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.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { - streamInterval = time.Duration(s.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.cfg != nil && s.cfg.Gateway.StreamKeepaliveInterval > 0 { - keepaliveInterval = time.Duration(s.cfg.Gateway.StreamKeepaliveInterval) * time.Second - } - var keepaliveTimer *time.Timer - if keepaliveInterval > 0 { - keepaliveTimer = time.NewTimer(keepaliveInterval) - defer keepaliveTimer.Stop() - } - var keepaliveCh <-chan time.Time - if keepaliveTimer != nil { - keepaliveCh = keepaliveTimer.C - } - lastDataAt := time.Now() - resetKeepaliveTimer := func() { - if keepaliveTimer == nil { - return - } - if !keepaliveTimer.Stop() { - select { - case <-keepaliveTimer.C: - default: - } - } - keepaliveTimer.Reset(keepaliveInterval) - } - - // 仅发送一次错误事件，避免多次写入导致协议混乱（写失败时尽力通知客户端）。 - // 事件格式遵循 Anthropic SSE 标准：{"type":"error","error":{"type":,"message":}} - // 这样 Anthropic SDK / Claude Code 等客户端能按标准 error 类型解析，UI 能显示具体错误文案， - // 服务端 ExtractUpstreamErrorMessage 也能从透传的 body 中提取 message。 - errorEventSent := false - sendErrorEvent := func(reason, message string) { - if errorEventSent { - return - } - errorEventSent = true - if message == "" { - message = reason - } - body, err := json.Marshal(map[string]any{ - "type": "error", - "error": map[string]string{ - "type": reason, - "message": message, - }, - }) - if err != nil { - // json.Marshal 不可能在已知 string-only 输入上失败，保守 fallback - body = []byte(fmt.Sprintf(`{"type":"error","error":{"type":%q,"message":%q}}`, reason, message)) - } - _, _ = fmt.Fprintf(w, "event: error\ndata: %s\n\n", body) - flusher.Flush() - } - - needModelReplace := originalModel != mappedModel - clientDisconnected := false // 客户端断开标志，断开后继续读取上游以获取完整usage - sawTerminalEvent := false - useNoopDeltaKeepalive := c != nil && c.Request != nil && shouldUseClaudeCodeNoopDeltaKeepalive(c.GetHeader("User-Agent")) - noopDeltaKeepaliveBlockIndex := -1 - noopDeltaKeepaliveDeltaType := "" - - pendingEventLines := make([]string, 0, 4) - - processSSEEvent := func(lines []string) ([]string, string, *sseUsagePatch, error) { - if len(lines) == 0 { - return nil, "", nil, nil - } - - eventName := "" - dataLine := "" - for _, line := range lines { - trimmed := strings.TrimSpace(line) - if strings.HasPrefix(trimmed, "event:") { - eventName = strings.TrimSpace(strings.TrimPrefix(trimmed, "event:")) - continue - } - if dataLine == "" && sseDataRe.MatchString(trimmed) { - dataLine = sseDataRe.ReplaceAllString(trimmed, "") - } - } - - if eventName == "error" { - return nil, dataLine, nil, &sseStreamErrorEventError{RawData: dataLine} - } - - if dataLine == "" { - return []string{strings.Join(lines, "\n") + "\n\n"}, "", nil, nil - } - - if dataLine == "[DONE]" { - sawTerminalEvent = true - block := "" - if eventName != "" { - block = "event: " + eventName + "\n" - } - block += "data: " + dataLine + "\n\n" - return []string{block}, dataLine, nil, nil - } - - var event map[string]any - if err := json.Unmarshal([]byte(dataLine), &event); err != nil { - // JSON 解析失败，直接透传原始数据 - block := "" - if eventName != "" { - block = "event: " + eventName + "\n" - } - block += "data: " + dataLine + "\n\n" - return []string{block}, dataLine, nil, nil - } - - eventType, _ := event["type"].(string) - if eventName == "" { - eventName = eventType - } - eventChanged := false - - if useNoopDeltaKeepalive { - switch eventType { - case "content_block_start": - if idx, ok := sseEventIndex(event); ok { - noopDeltaKeepaliveBlockIndex = -1 - noopDeltaKeepaliveDeltaType = "" - if contentBlock, ok := event["content_block"].(map[string]any); ok { - blockType, _ := contentBlock["type"].(string) - if deltaType := claudeCodeKeepaliveDeltaTypeForContentBlock(blockType); deltaType != "" { - noopDeltaKeepaliveBlockIndex = idx - noopDeltaKeepaliveDeltaType = deltaType - } - } - } - case "content_block_delta": - if idx, ok := sseEventIndex(event); ok { - if delta, ok := event["delta"].(map[string]any); ok { - deltaType, _ := delta["type"].(string) - if claudeCodeKeepaliveFieldForDeltaType(deltaType) != "" { - noopDeltaKeepaliveBlockIndex = idx - noopDeltaKeepaliveDeltaType = deltaType - } - } - } - case "content_block_stop": - if idx, ok := sseEventIndex(event); ok && idx == noopDeltaKeepaliveBlockIndex { - noopDeltaKeepaliveBlockIndex = -1 - noopDeltaKeepaliveDeltaType = "" - } - case "message_stop": - noopDeltaKeepaliveBlockIndex = -1 - noopDeltaKeepaliveDeltaType = "" - } - } - - // 兼容 Kimi cached_tokens → cache_read_input_tokens - if eventType == "message_start" { - if msg, ok := event["message"].(map[string]any); ok { - if u, ok := msg["usage"].(map[string]any); ok { - eventChanged = reconcileCachedTokens(u) || eventChanged - } - } - } - if eventType == "message_delta" { - if u, ok := event["usage"].(map[string]any); ok { - eventChanged = reconcileCachedTokens(u) || eventChanged - } - } - - // Cache TTL Override: 重写 SSE 事件中的 cache_creation 分类。 - // 账号级设置优先；全局 1h 请求注入开启时，默认把 usage 计费归回 5m。 - if overrideTarget, ok := s.resolveCacheTTLUsageOverrideTarget(ctx, account); ok { - if eventType == "message_start" { - if msg, ok := event["message"].(map[string]any); ok { - if u, ok := msg["usage"].(map[string]any); ok { - eventChanged = rewriteCacheCreationJSON(u, overrideTarget) || eventChanged - } - } - } - if eventType == "message_delta" { - if u, ok := event["usage"].(map[string]any); ok { - eventChanged = rewriteCacheCreationJSON(u, overrideTarget) || eventChanged - } - } - } - - if needModelReplace { - if msg, ok := event["message"].(map[string]any); ok { - if model, ok := msg["model"].(string); ok && model == mappedModel { - msg["model"] = originalModel - eventChanged = true - } - } - } - - usagePatch := s.extractSSEUsagePatch(event) - if anthropicStreamEventIsTerminal(eventName, dataLine) { - sawTerminalEvent = true - } - if !eventChanged { - block := "" - if eventName != "" { - block = "event: " + eventName + "\n" - } - block += "data: " + dataLine + "\n\n" - return []string{block}, dataLine, usagePatch, nil - } - - newData, err := json.Marshal(event) - if err != nil { - // 序列化失败，直接透传原始数据 - block := "" - if eventName != "" { - block = "event: " + eventName + "\n" - } - block += "data: " + dataLine + "\n\n" - return []string{block}, dataLine, usagePatch, nil - } - - block := "" - if eventName != "" { - block = "event: " + eventName + "\n" - } - block += "data: " + string(newData) + "\n\n" - return []string{block}, string(newData), usagePatch, nil - } - - for { - select { - case ev, ok := <-events: - if !ok { - // 上游完成，返回结果 - if !sawTerminalEvent { - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, fmt.Errorf("stream usage incomplete: missing terminal event") - } - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, nil - } - if ev.err != nil { - if sawTerminalEvent { - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, nil - } - // 检测 context 取消（客户端断开会导致 context 取消，进而影响上游读取） - if errors.Is(ev.err, context.Canceled) || errors.Is(ev.err, context.DeadlineExceeded) { - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, fmt.Errorf("stream usage incomplete: %w", ev.err) - } - // 客户端已通过写入失败检测到断开，上游也出错了，返回已收集的 usage - if clientDisconnected { - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, fmt.Errorf("stream usage incomplete after disconnect: %w", ev.err) - } - // 客户端未断开，正常的错误处理 - if errors.Is(ev.err, bufio.ErrTooLong) { - logger.LegacyPrintf("service.gateway", "SSE line too long: account=%d max_size=%d error=%v", account.ID, maxLineSize, ev.err) - sendErrorEvent("response_too_large", fmt.Sprintf("upstream SSE line exceeded %d bytes", maxLineSize)) - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, ev.err - } - // 上游中途读错误（unexpected EOF / connection reset 等，常见于 HTTP/2 GOAWAY）： - // 若尚未向客户端写过任何字节，包成 UpstreamFailoverError 让 handler 层走 failover/重试。 - // 已经开始写流时 SSE 协议无 resume，只能透传错误事件给客户端。 - // 注意:面向客户端的 disconnectMsg 必须用 sanitizeStreamError 剥离地址, - // 默认 *net.OpError 的 Error() 会泄露内部 IP/端口和上游地址。完整 ev.err - // 仅在下方 LegacyPrintf 内部日志中保留供运维诊断。 - disconnectMsg := "upstream stream disconnected: " + sanitizeStreamError(ev.err) - if !c.Writer.Written() { - logger.LegacyPrintf("service.gateway", "Upstream stream read error before any client output (account=%d), failing over: %v", account.ID, ev.err) - body, _ := json.Marshal(map[string]any{ - "type": "error", - "error": map[string]string{ - "type": "upstream_disconnected", - "message": disconnectMsg, - }, - }) - return nil, &UpstreamFailoverError{ - StatusCode: http.StatusBadGateway, - ResponseBody: body, - RetryableOnSameAccount: true, - } - } - sendErrorEvent("stream_read_error", disconnectMsg) - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream read error: %w", ev.err) - } - line := ev.line - trimmed := strings.TrimSpace(line) - - if trimmed == "" { - if len(pendingEventLines) == 0 { - continue - } - - outputBlocks, data, usagePatch, err := processSSEEvent(pendingEventLines) - pendingEventLines = pendingEventLines[:0] - if err != nil { - if clientDisconnected { - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil - } - return nil, err - } - - for _, block := range outputBlocks { - if !clientDisconnected { - restored := reverseToolNamesIfPresent(c, []byte(block)) - if _, werr := fmt.Fprint(w, string(restored)); werr != nil { - clientDisconnected = true - logger.LegacyPrintf("service.gateway", "Client disconnected during streaming, continuing to drain upstream for billing") - break - } - flusher.Flush() - lastDataAt = time.Now() - resetKeepaliveTimer() - } - if data != "" { - if firstTokenMs == nil && data != "[DONE]" { - ms := int(time.Since(startTime).Milliseconds()) - firstTokenMs = &ms - } - if usagePatch != nil { - mergeSSEUsagePatch(usage, usagePatch) - } - } - } - continue - } - - pendingEventLines = append(pendingEventLines, line) - - case <-intervalCh: - lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) - if time.Since(lastRead) < streamInterval { - continue - } - if clientDisconnected { - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, fmt.Errorf("stream usage incomplete after timeout") - } - logger.LegacyPrintf("service.gateway", "Stream data interval timeout: account=%d model=%s interval=%s", account.ID, originalModel, streamInterval) - // 处理流超时，可能标记账户为临时不可调度或错误状态 - if s.rateLimitService != nil { - s.rateLimitService.HandleStreamTimeout(ctx, account, originalModel) - } - sendErrorEvent("stream_timeout", fmt.Sprintf("upstream stream idle for %s", streamInterval)) - return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") - - case <-keepaliveCh: - if clientDisconnected { - continue - } - if time.Since(lastDataAt) < keepaliveInterval { - resetKeepaliveTimer() - continue - } - keepaliveBlock := "event: ping\ndata: {\"type\": \"ping\"}\n\n" - if useNoopDeltaKeepalive && noopDeltaKeepaliveBlockIndex >= 0 { - if block, ok := buildClaudeCodeNoopDeltaKeepalive(noopDeltaKeepaliveBlockIndex, noopDeltaKeepaliveDeltaType); ok { - keepaliveBlock = block - } - } - if _, werr := fmt.Fprint(w, keepaliveBlock); werr != nil { - clientDisconnected = true - logger.LegacyPrintf("service.gateway", "Client disconnected during keepalive ping, continuing to drain upstream for billing") - continue - } - flusher.Flush() - lastDataAt = time.Now() - resetKeepaliveTimer() - } - } - -} - -func (s *GatewayService) parseSSEUsage(data string, usage *ClaudeUsage) { - if usage == nil { - return - } - - var event map[string]any - if err := json.Unmarshal([]byte(data), &event); err != nil { - return - } - - if patch := s.extractSSEUsagePatch(event); patch != nil { - mergeSSEUsagePatch(usage, patch) - } -} - -type sseUsagePatch struct { - inputTokens int - hasInputTokens bool - outputTokens int - hasOutputTokens bool - cacheCreationInputTokens int - hasCacheCreationInput bool - cacheReadInputTokens int - hasCacheReadInput bool - cacheCreation5mTokens int - hasCacheCreation5m bool - cacheCreation1hTokens int - hasCacheCreation1h bool -} - -func (s *GatewayService) extractSSEUsagePatch(event map[string]any) *sseUsagePatch { - if len(event) == 0 { - return nil - } - - eventType, _ := event["type"].(string) - switch eventType { - case "message_start": - msg, _ := event["message"].(map[string]any) - usageObj, _ := msg["usage"].(map[string]any) - if len(usageObj) == 0 { - return nil - } - - patch := &sseUsagePatch{} - patch.hasInputTokens = true - if v, ok := parseSSEUsageInt(usageObj["input_tokens"]); ok { - patch.inputTokens = v - } - patch.hasCacheCreationInput = true - if v, ok := parseSSEUsageInt(usageObj["cache_creation_input_tokens"]); ok { - patch.cacheCreationInputTokens = v - } - patch.hasCacheReadInput = true - if v, ok := parseSSEUsageInt(usageObj["cache_read_input_tokens"]); ok { - patch.cacheReadInputTokens = v - } - if cc, ok := usageObj["cache_creation"].(map[string]any); ok { - if v, exists := parseSSEUsageInt(cc["ephemeral_5m_input_tokens"]); exists { - patch.cacheCreation5mTokens = v - patch.hasCacheCreation5m = true - } - if v, exists := parseSSEUsageInt(cc["ephemeral_1h_input_tokens"]); exists { - patch.cacheCreation1hTokens = v - patch.hasCacheCreation1h = true - } - } - return patch - - case "message_delta": - usageObj, _ := event["usage"].(map[string]any) - if len(usageObj) == 0 { - return nil - } - - patch := &sseUsagePatch{} - if v, ok := parseSSEUsageInt(usageObj["input_tokens"]); ok && v > 0 { - patch.inputTokens = v - patch.hasInputTokens = true - } - if v, ok := parseSSEUsageInt(usageObj["output_tokens"]); ok && v > 0 { - patch.outputTokens = v - patch.hasOutputTokens = true - } - if v, ok := parseSSEUsageInt(usageObj["cache_creation_input_tokens"]); ok && v > 0 { - patch.cacheCreationInputTokens = v - patch.hasCacheCreationInput = true - } - if v, ok := parseSSEUsageInt(usageObj["cache_read_input_tokens"]); ok && v > 0 { - patch.cacheReadInputTokens = v - patch.hasCacheReadInput = true - } - if cc, ok := usageObj["cache_creation"].(map[string]any); ok { - if v, exists := parseSSEUsageInt(cc["ephemeral_5m_input_tokens"]); exists && v > 0 { - patch.cacheCreation5mTokens = v - patch.hasCacheCreation5m = true - } - if v, exists := parseSSEUsageInt(cc["ephemeral_1h_input_tokens"]); exists && v > 0 { - patch.cacheCreation1hTokens = v - patch.hasCacheCreation1h = true - } - } - return patch - } - - return nil -} - -func mergeSSEUsagePatch(usage *ClaudeUsage, patch *sseUsagePatch) { - if usage == nil || patch == nil { - return - } - - if patch.hasInputTokens { - usage.InputTokens = patch.inputTokens - } - if patch.hasCacheCreationInput { - usage.CacheCreationInputTokens = patch.cacheCreationInputTokens - } - if patch.hasCacheReadInput { - usage.CacheReadInputTokens = patch.cacheReadInputTokens - } - if patch.hasOutputTokens { - usage.OutputTokens = patch.outputTokens - } - if patch.hasCacheCreation5m { - usage.CacheCreation5mTokens = patch.cacheCreation5mTokens - } - if patch.hasCacheCreation1h { - usage.CacheCreation1hTokens = patch.cacheCreation1hTokens - } -} - -func parseSSEUsageInt(value any) (int, bool) { - switch v := value.(type) { - case float64: - return int(v), true - case float32: - return int(v), true - case int: - return v, true - case int64: - return int(v), true - case int32: - return int(v), true - case json.Number: - if i, err := v.Int64(); err == nil { - return int(i), true - } - if f, err := v.Float64(); err == nil { - return int(f), true - } - case string: - if parsed, err := strconv.Atoi(strings.TrimSpace(v)); err == nil { - return parsed, true - } - } - return 0, false -} - -// applyCacheTTLOverride 将所有 cache creation tokens 归入指定的 TTL 类型。 -// target 为 "5m" 或 "1h"。返回 true 表示发生了变更。 -func applyCacheTTLOverride(usage *ClaudeUsage, target string) bool { - // Fallback: 如果只有聚合字段但无 5m/1h 明细，将聚合字段归入 5m 默认类别 - if usage.CacheCreation5mTokens == 0 && usage.CacheCreation1hTokens == 0 && usage.CacheCreationInputTokens > 0 { - usage.CacheCreation5mTokens = usage.CacheCreationInputTokens - } - - total := usage.CacheCreation5mTokens + usage.CacheCreation1hTokens - if total == 0 { - return false - } - switch target { - case "1h": - if usage.CacheCreation1hTokens == total { - return false // 已经全是 1h - } - usage.CacheCreation1hTokens = total - usage.CacheCreation5mTokens = 0 - default: // "5m" - if usage.CacheCreation5mTokens == total { - return false // 已经全是 5m - } - usage.CacheCreation5mTokens = total - usage.CacheCreation1hTokens = 0 - } - return true -} - -// rewriteCacheCreationJSON 在 JSON usage 对象中重写 cache_creation 嵌套对象的 TTL 分类。 -// usageObj 是 usage JSON 对象（map[string]any）。 -func rewriteCacheCreationJSON(usageObj map[string]any, target string) bool { - ccObj, ok := usageObj["cache_creation"].(map[string]any) - if !ok { - return false - } - v5m, _ := parseSSEUsageInt(ccObj["ephemeral_5m_input_tokens"]) - v1h, _ := parseSSEUsageInt(ccObj["ephemeral_1h_input_tokens"]) - total := v5m + v1h - if total == 0 { - return false - } - switch target { - case "1h": - if v1h == total { - return false - } - ccObj["ephemeral_1h_input_tokens"] = float64(total) - ccObj["ephemeral_5m_input_tokens"] = float64(0) - default: // "5m" - if v5m == total { - return false - } - ccObj["ephemeral_5m_input_tokens"] = float64(total) - ccObj["ephemeral_1h_input_tokens"] = float64(0) - } - return true -} - -func (s *GatewayService) resolveCacheTTLUsageOverrideTarget(ctx context.Context, account *Account) (string, bool) { - if account == nil { - return "", false - } - if account.IsCacheTTLOverrideEnabled() { - return account.GetCacheTTLOverrideTarget(), true - } - if account.IsAnthropicOAuthOrSetupToken() && s != nil && s.settingService != nil && s.settingService.IsAnthropicCacheTTL1hInjectionEnabled(ctx) { - return cacheTTLTarget5m, true - } - return "", false -} - -func (s *GatewayService) handleNonStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, originalModel, mappedModel string) (*ClaudeUsage, error) { - // 更新5h窗口状态 - s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header) - - body, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, anthropicTooLargeError) - if err != nil { - return nil, err - } - - // 解析usage - var response struct { - Usage ClaudeUsage `json:"usage"` - } - if err := json.Unmarshal(body, &response); err != nil { - if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices { - return nil, s.invalidNonStreamingJSONFailoverError(ctx, resp, account, body, err, mappedModel) - } - return nil, fmt.Errorf("parse response: %w", err) - } - - // 解析嵌套的 cache_creation 对象中的 5m/1h 明细 - cc5m := gjson.GetBytes(body, "usage.cache_creation.ephemeral_5m_input_tokens") - cc1h := gjson.GetBytes(body, "usage.cache_creation.ephemeral_1h_input_tokens") - if cc5m.Exists() || cc1h.Exists() { - response.Usage.CacheCreation5mTokens = int(cc5m.Int()) - response.Usage.CacheCreation1hTokens = int(cc1h.Int()) - } - - // 兼容 Kimi cached_tokens → cache_read_input_tokens - if response.Usage.CacheReadInputTokens == 0 { - cachedTokens := gjson.GetBytes(body, "usage.cached_tokens").Int() - if cachedTokens > 0 { - response.Usage.CacheReadInputTokens = int(cachedTokens) - if newBody, err := sjson.SetBytes(body, "usage.cache_read_input_tokens", cachedTokens); err == nil { - body = newBody - } - } - } - - // Cache TTL Override: 重写 non-streaming 响应中的 cache_creation 分类。 - // 账号级设置优先；全局 1h 请求注入开启时，默认把 usage 计费归回 5m。 - if overrideTarget, ok := s.resolveCacheTTLUsageOverrideTarget(ctx, account); ok { - if applyCacheTTLOverride(&response.Usage, overrideTarget) { - // 同步更新 body JSON 中的嵌套 cache_creation 对象 - if newBody, err := sjson.SetBytes(body, "usage.cache_creation.ephemeral_5m_input_tokens", response.Usage.CacheCreation5mTokens); err == nil { - body = newBody - } - if newBody, err := sjson.SetBytes(body, "usage.cache_creation.ephemeral_1h_input_tokens", response.Usage.CacheCreation1hTokens); err == nil { - body = newBody - } - } - } - - // 如果有模型映射，替换响应中的model字段 - if originalModel != mappedModel { - body = s.replaceModelInResponseBody(body, mappedModel, originalModel) - } - - responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) - - contentType := "application/json" - if s.cfg != nil && !s.cfg.Security.ResponseHeaders.Enabled { - if upstreamType := resp.Header.Get("Content-Type"); upstreamType != "" { - contentType = upstreamType - } - } - - body = reverseToolNamesIfPresent(c, body) - - // 写入响应 - c.Data(resp.StatusCode, contentType, body) - - return &response.Usage, nil -} - -// replaceModelInResponseBody 替换响应体中的model字段 -// 使用 gjson/sjson 精确替换，避免全量 JSON 反序列化 -func (s *GatewayService) replaceModelInResponseBody(body []byte, fromModel, toModel string) []byte { - if m := gjson.GetBytes(body, "model"); m.Exists() && m.Str == fromModel { - newBody, err := sjson.SetBytes(body, "model", toModel) - if err != nil { - return body - } - return newBody - } - return body -} - -func (s *GatewayService) getUserGroupRateMultiplier(ctx context.Context, userID, groupID int64, groupDefaultMultiplier float64) float64 { - if s == nil { - return groupDefaultMultiplier - } - resolver := s.userGroupRateResolver - if resolver == nil { - resolver = newUserGroupRateResolver( - s.userGroupRateRepo, - s.userGroupRateCache, - resolveUserGroupRateCacheTTL(s.cfg), - &s.userGroupRateSF, - "service.gateway", - ) - } - return resolver.Resolve(ctx, userID, groupID, groupDefaultMultiplier) -} - -// RecordUsageInput 记录使用量的输入参数。 -// 异步 worker 只接收计费所需快照，不能持有 ParsedRequest/RequestBodyRef 这类大请求体引用。 -type RecordUsageInput struct { - Result *ForwardResult - APIKey *APIKey - User *User - Account *Account - Subscription *UserSubscription // 可选：订阅信息 - InboundEndpoint string // 入站端点（客户端请求路径） - UpstreamEndpoint string // 上游端点（标准化后的上游路径） - UserAgent string // 请求的 User-Agent - IPAddress string // 请求的客户端 IP 地址 - RequestPayloadHash string // 请求体语义哈希，用于降低 request_id 误复用时的静默误去重风险 - ForceCacheBilling bool // 强制缓存计费：将 input_tokens 转为 cache_read 计费（用于粘性会话切换） - APIKeyService APIKeyQuotaUpdater // 可选：用于更新API Key配额 - QuotaPlatform string // user×platform 配额计量平台：handler 在请求 ctx 内经 QuotaPlatform() 算定后传入（后扣运行在 worker 池 background ctx 上，取不到 ForcePlatform） - - ChannelUsageFields // 渠道映射信息（由 handler 在 Forward 前解析） -} - -// APIKeyQuotaUpdater defines the interface for updating API Key quota and rate limit usage -type APIKeyQuotaUpdater interface { - UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error - UpdateRateLimitUsage(ctx context.Context, apiKeyID int64, cost float64) error -} - -type apiKeyAuthCacheInvalidator interface { - InvalidateAuthCacheByKey(ctx context.Context, key string) -} - -type usageLogBestEffortWriter interface { - CreateBestEffort(ctx context.Context, log *UsageLog) error -} - -// postUsageBillingParams 统一扣费所需的参数 -type postUsageBillingParams struct { - Cost *CostBreakdown - User *User - APIKey *APIKey - Account *Account - Subscription *UserSubscription - RequestPayloadHash string - IsSubscriptionBill bool - AccountRateMultiplier float64 - APIKeyService APIKeyQuotaUpdater - Platform string // 来自 APIKey 关联 Group 的平台标识 -} - -// PlatformFromAPIKey 从 APIKey 关联的 Group 推导 platform 名称。 -// apiKey 为 nil 或 Group 信息缺失时返回空串（调用方据此 short-circuit quota 累加）。 -// 导出供 handler 层调用。 -func PlatformFromAPIKey(apiKey *APIKey) string { - if apiKey == nil || apiKey.Group == nil { - return "" - } - return apiKey.Group.Platform -} - -// QuotaPlatform 返回 user×platform 配额计量使用的平台标识。 -// 强制平台路由（如 /antigravity）优先按 ctx 中的 ForcePlatform 计量，否则回退到 -// APIKey 关联 Group 的平台。 -// -// 注意：必须用带 ForcePlatform 的请求 context 调用（如 handler 的 c.Request.Context()）。 -// 后扣运行在 worker 池的 background ctx 上没有 ForcePlatform，因此后扣平台由 handler -// 预先算定、经 RecordUsageInput.QuotaPlatform 传入，不要在后扣链路用 worker ctx 调用本函数。 -func QuotaPlatform(ctx context.Context, apiKey *APIKey) string { - if fp, ok := ctx.Value(ctxkey.ForcePlatform).(string); ok && fp != "" { - return fp - } - return PlatformFromAPIKey(apiKey) -} - -func (p *postUsageBillingParams) shouldDeductAPIKeyQuota() bool { - return p.Cost.ActualCost > 0 && p.APIKey.Quota > 0 && p.APIKeyService != nil -} - -func (p *postUsageBillingParams) shouldUpdateRateLimits() bool { - return p.Cost.ActualCost > 0 && p.APIKey.HasRateLimits() && p.APIKeyService != nil -} - -func (p *postUsageBillingParams) shouldUpdateAccountQuota() bool { - return p.Cost.TotalCost > 0 && p.Account.IsAPIKeyOrBedrock() && p.Account.HasAnyQuotaLimit() -} - -// postUsageBilling is the legacy fallback billing path used when the unified -// billing repo is unavailable (nil). Production uses applyUsageBilling → repo.Apply -// for atomic billing. This path only runs in tests or degraded mode. -func postUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *billingDeps) { - billingCtx, cancel := detachedBillingContext(ctx) - defer cancel() - - cost := p.Cost - - if p.IsSubscriptionBill { - // Subscription usage tracked by ActualCost so group rate multiplier - // consumes the quota at the expected speed. - if cost.ActualCost > 0 { - if err := deps.userSubRepo.IncrementUsage(billingCtx, p.Subscription.ID, cost.ActualCost); err != nil { - slog.Error("increment subscription usage failed", "subscription_id", p.Subscription.ID, "error", err) - } - } - } else { - if cost.ActualCost > 0 { - if err := deps.userRepo.DeductBalance(billingCtx, p.User.ID, cost.ActualCost); err != nil { - slog.Error("deduct balance failed", "user_id", p.User.ID, "error", err) - } else if deps.billingCacheService != nil { - if err := deps.billingCacheService.InvalidateUserBalance(billingCtx, p.User.ID); err != nil { - slog.Warn("invalidate balance cache after legacy deduction failed", "user_id", p.User.ID, "error", err) - } - } - } - } - - if p.shouldDeductAPIKeyQuota() { - if err := p.APIKeyService.UpdateQuotaUsed(billingCtx, p.APIKey.ID, cost.ActualCost); err != nil { - slog.Error("update api key quota failed", "api_key_id", p.APIKey.ID, "error", err) - } - } - - if p.shouldUpdateRateLimits() { - if err := p.APIKeyService.UpdateRateLimitUsage(billingCtx, p.APIKey.ID, cost.ActualCost); err != nil { - slog.Error("update api key rate limit usage failed", "api_key_id", p.APIKey.ID, "error", err) - } - } - - if p.shouldUpdateAccountQuota() { - accountCost := cost.TotalCost * p.AccountRateMultiplier - if err := deps.accountRepo.IncrementQuotaUsed(billingCtx, p.Account.ID, accountCost); err != nil { - slog.Error("increment account quota used failed", "account_id", p.Account.ID, "cost", accountCost, "error", err) - } - } - - // Platform quota 累加（legacy 兜底路径）：仅对 standard（余额）模式生效；订阅模式豁免；仅对有 limit 的用户写 - // - HasUserPlatformQuotaLimit 守卫:与正常路径对齐，无 limit 公司跳过 - // - 新增 Redis 同步写:enforcement 走 Redis，legacy 路径也必须同步写，否则 preflight 看不到消费 - // - flusher_enabled=false（降级）:保留原有同步直写 DB - // - flusher_enabled=true:跳过直写 DB，由 flusher 异步批量刷（markDirty 在 IncrementUserPlatformQuotaUsage 内部完成） - // - 失败仅记 ALERT log + counter，不阻断主扣费流程 - if !p.IsSubscriptionBill && p.Platform != "" && cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil { - if deps.billingCacheService.HasUserPlatformQuotaLimit(billingCtx, p.User.ID, p.Platform) { - deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, cost.ActualCost) - if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled { - // 降级路径:flusher 未启用时保留原有同步直写 DB - if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(billingCtx, p.User.ID, p.Platform, cost.ActualCost, time.Now().UTC()); err != nil { - userPlatformQuotaDBIncrLegacyErrorTotal.Add(1) - logger.LegacyPrintf("service.gateway", "ALERT: legacy incr user platform quota DB failed user=%d platform=%s cost=%f: %v", p.User.ID, p.Platform, cost.ActualCost, err) - } - } - // flusher_enabled=true:不直写 DB，flusher 异步批量刷 - } - } - - // NOTE: finalizePostUsageBilling is NOT called here to avoid double-queuing - // cache updates. The legacy path does DB writes directly; the finalize path - // does cache queue + notifications. Notifications are dispatched separately - // by the caller after recording the usage log. -} - -func resolveUsageBillingRequestID(ctx context.Context, upstreamRequestID string) string { - if ctx != nil { - if clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(clientRequestID) != "" { - return "client:" + strings.TrimSpace(clientRequestID) - } - if requestID, _ := ctx.Value(ctxkey.RequestID).(string); strings.TrimSpace(requestID) != "" { - return "local:" + strings.TrimSpace(requestID) - } - } - if requestID := strings.TrimSpace(upstreamRequestID); requestID != "" { - return requestID - } - return "generated:" + generateRequestID() -} - -func resolveUsageBillingPayloadFingerprint(ctx context.Context, requestPayloadHash string) string { - if payloadHash := strings.TrimSpace(requestPayloadHash); payloadHash != "" { - return payloadHash - } - if ctx != nil { - if clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(clientRequestID) != "" { - return "client:" + strings.TrimSpace(clientRequestID) - } - if requestID, _ := ctx.Value(ctxkey.RequestID).(string); strings.TrimSpace(requestID) != "" { - return "local:" + strings.TrimSpace(requestID) - } - } - return "" -} - -func buildUsageBillingCommand(requestID string, usageLog *UsageLog, p *postUsageBillingParams) *UsageBillingCommand { - if p == nil || p.Cost == nil || p.APIKey == nil || p.User == nil || p.Account == nil { - return nil - } - - cmd := &UsageBillingCommand{ - RequestID: requestID, - APIKeyID: p.APIKey.ID, - UserID: p.User.ID, - AccountID: p.Account.ID, - AccountType: p.Account.Type, - RequestPayloadHash: strings.TrimSpace(p.RequestPayloadHash), - } - if usageLog != nil { - cmd.Model = usageLog.Model - cmd.BillingType = usageLog.BillingType - cmd.InputTokens = usageLog.InputTokens - cmd.OutputTokens = usageLog.OutputTokens - cmd.CacheCreationTokens = usageLog.CacheCreationTokens - cmd.CacheReadTokens = usageLog.CacheReadTokens - cmd.ImageCount = usageLog.ImageCount - if usageLog.ServiceTier != nil { - cmd.ServiceTier = *usageLog.ServiceTier - } - if usageLog.ReasoningEffort != nil { - cmd.ReasoningEffort = *usageLog.ReasoningEffort - } - if usageLog.SubscriptionID != nil { - cmd.SubscriptionID = usageLog.SubscriptionID - } - } - - // Record subscription / balance cost using ActualCost so the group (and any - // user-specific) rate multiplier consumes subscription quota at the expected - // speed. TotalCost remains the raw (pre-multiplier) value; downstream guards - // on "> 0" still correctly skip free subscriptions (RateMultiplier == 0). - if p.IsSubscriptionBill && p.Subscription != nil && p.Cost.TotalCost > 0 { - cmd.SubscriptionID = &p.Subscription.ID - cmd.SubscriptionCost = p.Cost.ActualCost - } else if p.Cost.ActualCost > 0 { - cmd.BalanceCost = p.Cost.ActualCost - } - - if p.shouldDeductAPIKeyQuota() { - cmd.APIKeyQuotaCost = p.Cost.ActualCost - } - if p.shouldUpdateRateLimits() { - cmd.APIKeyRateLimitCost = p.Cost.ActualCost - } - if p.shouldUpdateAccountQuota() { - cmd.AccountQuotaCost = p.Cost.TotalCost * p.AccountRateMultiplier - } - - cmd.Normalize() - return cmd -} - -func applyUsageBilling(ctx context.Context, requestID string, usageLog *UsageLog, p *postUsageBillingParams, deps *billingDeps, repo UsageBillingRepository) (bool, error) { - if p == nil || deps == nil { - return false, nil - } - - cmd := buildUsageBillingCommand(requestID, usageLog, p) - if cmd == nil || cmd.RequestID == "" || repo == nil { - postUsageBilling(ctx, p, deps) - return true, nil - } - - billingCtx, cancel := detachedBillingContext(ctx) - defer cancel() - - result, err := repo.Apply(billingCtx, cmd) - if err != nil { - return false, err - } - - if result == nil || !result.Applied { - deps.deferredService.ScheduleLastUsedUpdate(p.Account.ID) - return false, nil - } - - if result.APIKeyQuotaExhausted { - if invalidator, ok := p.APIKeyService.(apiKeyAuthCacheInvalidator); ok && p.APIKey != nil && p.APIKey.Key != "" { - invalidator.InvalidateAuthCacheByKey(billingCtx, p.APIKey.Key) - } - } - - finalizePostUsageBilling(billingCtx, p, deps, result) - return true, nil -} - -func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) { - if p == nil || p.Cost == nil || deps == nil { - return - } - - if p.IsSubscriptionBill { - if p.Cost.ActualCost > 0 && p.User != nil && p.APIKey != nil && p.APIKey.GroupID != nil { - deps.billingCacheService.QueueUpdateSubscriptionUsage(p.User.ID, *p.APIKey.GroupID, p.Cost.ActualCost) - } - } else if p.Cost.ActualCost > 0 && p.User != nil { - syncBalanceCacheAfterDeduction(ctx, p, deps, result) - } - - if p.Cost.ActualCost > 0 && p.APIKey != nil && p.APIKey.HasRateLimits() { - deps.billingCacheService.QueueUpdateAPIKeyRateLimitUsage(p.APIKey.ID, p.Cost.ActualCost) - } - - deps.deferredService.ScheduleLastUsedUpdate(p.Account.ID) - - // Platform quota 累加：仅在 standard（余额）模式生效；订阅模式豁免；仅对有 limit 的用户写 - // Redis 同步写 + DB 异步持久化（flag=false 降级）或 flusher 异步刷（flag=true）: - // - HasUserPlatformQuotaLimit 守卫:无 limit 的公司跳过,避免无效写入 + 浪费 Redis 容量 - // - Redis 同步:确保下次 preflight 立即看到最新 usage,把 TOCTOU 超支窗口 - // 限制在并发 in-flight 请求数量内（旧实现的异步入队会让超支无限累积直到 worker 处理） - // - DB 异步(flusher_enabled=false):在独立 goroutine 中走 detached context,失败用 ALERT log 触发 oncall 对账 - // - flusher_enabled=true:不直写 DB,由 flusher 异步批量刷（markDirty 已在 IncrementUserPlatformQuotaUsage 内部完成） - if !p.IsSubscriptionBill && p.Platform != "" && p.Cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil { - if deps.billingCacheService.HasUserPlatformQuotaLimit(ctx, p.User.ID, p.Platform) { - deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, p.Cost.ActualCost) - if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled { - // 降级路径:flusher 未启用时保留原有异步直写 DB - dbCtx, dbCancel := detachUpstreamContext(ctx) - userID, platform, cost := p.User.ID, p.Platform, p.Cost.ActualCost - go func() { - defer func() { - if r := recover(); r != nil { - logger.LegacyPrintf("service.gateway", "ALERT: panic in user platform quota incr goroutine user=%d platform=%s: %v", userID, platform, r) - } - }() - defer dbCancel() - if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(dbCtx, userID, platform, cost, time.Now().UTC()); err != nil { - // 失败计数器:暴露给 GatewayUserPlatformQuotaIncrStats(),由 ops 面板做斜率告警。 - userPlatformQuotaDBIncrErrorTotal.Add(1) - // ALERT 级别:DB 持久化失败意味着 Redis cache 失效后该笔 cost 永久丢失, - // 用户配额视图与实际消费会偏差,oncall 需要据此对账或人工补录。 - logger.LegacyPrintf("service.gateway", "ALERT: incr user platform quota DB failed user=%d platform=%s cost=%f: %v", userID, platform, cost, err) - } - }() - } - // flusher_enabled=true:不直写 DB,flusher 异步批量刷 - } - } - - // Notification checks run async — all parameters are already captured, - // no dependency on the request context or upstream connection. - go notifyBalanceLow(p, deps, result) - go notifyAccountQuota(p, deps, result) -} - -func syncBalanceCacheAfterDeduction(ctx context.Context, p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) { - if p == nil || p.Cost == nil || p.User == nil || deps == nil || deps.billingCacheService == nil { - return - } - if result != nil && result.NewBalance != nil && deps.billingCacheService.balanceBelowEligibilityThreshold(*result.NewBalance) { - if err := deps.billingCacheService.InvalidateUserBalance(ctx, p.User.ID); err != nil { - slog.Warn("invalidate balance cache after exhausted deduction failed", - "user_id", p.User.ID, - "new_balance", *result.NewBalance, - "balance_overdrafted", result.BalanceOverdrafted, - "error", err, - ) - } - return - } - deps.billingCacheService.QueueDeductBalance(p.User.ID, p.Cost.ActualCost) -} - -// notifyBalanceLow sends balance low notification after deduction. -// When result.NewBalance is available (from DB transaction RETURNING), it is used directly -// to reconstruct oldBalance, avoiding stale Redis reads and concurrent-deduction races. -func notifyBalanceLow(p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) { - defer func() { - if r := recover(); r != nil { - slog.Error("panic in notifyBalanceLow", "recover", r) - } - }() - if p.IsSubscriptionBill || p.Cost.ActualCost <= 0 || p.User == nil || deps.balanceNotifyService == nil { - slog.Debug("notifyBalanceLow: skipped", - "is_subscription", p.IsSubscriptionBill, - "actual_cost", p.Cost.ActualCost, - "user_nil", p.User == nil, - "service_nil", deps.balanceNotifyService == nil, - ) - return - } - - oldBalance := resolveOldBalance(p, result) - slog.Debug("notifyBalanceLow: calling CheckBalanceAfterDeduction", - "user_id", p.User.ID, - "old_balance", oldBalance, - "cost", p.Cost.ActualCost, - "notify_enabled", p.User.BalanceNotifyEnabled, - "threshold", p.User.BalanceNotifyThreshold, - "result_has_new_balance", result != nil && result.NewBalance != nil, - ) - deps.balanceNotifyService.CheckBalanceAfterDeduction(context.Background(), p.User, oldBalance, p.Cost.ActualCost) -} - -// resolveOldBalance returns the pre-deduction balance. -// Prefers the DB transaction result (newBalance + cost) over snapshot. -func resolveOldBalance(p *postUsageBillingParams, result *UsageBillingApplyResult) float64 { - if result != nil && result.NewBalance != nil { - return *result.NewBalance + p.Cost.ActualCost - } - // Legacy fallback: snapshot balance from request context - return p.User.Balance -} - -// notifyAccountQuota sends account quota threshold notification after increment. -// When result.QuotaState is available (from DB transaction RETURNING), it is passed directly -// to avoid a separate DB read that may see stale or concurrently-modified data. -func notifyAccountQuota(p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) { - defer func() { - if r := recover(); r != nil { - slog.Error("panic in notifyAccountQuota", "recover", r) - } - }() - if p.Cost.TotalCost <= 0 || p.Account == nil || !p.Account.IsAPIKeyOrBedrock() || deps.balanceNotifyService == nil { - slog.Debug("notifyAccountQuota: skipped", - "total_cost", p.Cost.TotalCost, - "account_nil", p.Account == nil, - "is_apikey_or_bedrock", p.Account != nil && p.Account.IsAPIKeyOrBedrock(), - "service_nil", deps.balanceNotifyService == nil, - ) - return - } - accountCost := p.Cost.TotalCost * p.AccountRateMultiplier - var quotaState *AccountQuotaState - if result != nil { - quotaState = result.QuotaState - } - slog.Debug("notifyAccountQuota: calling CheckAccountQuotaAfterIncrement", - "account_id", p.Account.ID, - "account_cost", accountCost, - "has_quota_state", quotaState != nil, - ) - deps.balanceNotifyService.CheckAccountQuotaAfterIncrement(context.Background(), p.Account, accountCost, quotaState) -} - -func detachedBillingContext(ctx context.Context) (context.Context, context.CancelFunc) { - base := context.Background() - if ctx != nil { - base = context.WithoutCancel(ctx) - } - return context.WithTimeout(base, postUsageBillingTimeout) -} - -func detachStreamUpstreamContext(ctx context.Context, stream bool) (context.Context, context.CancelFunc) { - if ctx == nil { - return context.Background(), func() {} - } - if !stream { - return ctx, func() {} - } - return context.WithoutCancel(ctx), func() {} -} - -func detachUpstreamContext(ctx context.Context) (context.Context, context.CancelFunc) { - if ctx == nil { - return context.Background(), func() {} - } - return context.WithoutCancel(ctx), func() {} -} - -// billingDeps 扣费逻辑依赖的服务（由各 gateway service 提供） -type billingDeps struct { - accountRepo AccountRepository - userRepo UserRepository - userSubRepo UserSubscriptionRepository - billingCacheService *BillingCacheService - deferredService *DeferredService - balanceNotifyService *BalanceNotifyService - userPlatformQuotaRepo UserPlatformQuotaRepository - cfg *config.Config -} - -func (s *GatewayService) billingDeps() *billingDeps { - return &billingDeps{ - accountRepo: s.accountRepo, - userRepo: s.userRepo, - userSubRepo: s.userSubRepo, - billingCacheService: s.billingCacheService, - deferredService: s.deferredService, - balanceNotifyService: s.balanceNotifyService, - userPlatformQuotaRepo: s.userPlatformQuotaRepo, - cfg: s.cfg, - } -} - -func writeUsageLogBestEffort(ctx context.Context, repo UsageLogRepository, usageLog *UsageLog, logKey string) { - if repo == nil || usageLog == nil { - return - } - usageCtx, cancel := detachedBillingContext(ctx) - defer cancel() - - if writer, ok := repo.(usageLogBestEffortWriter); ok { - if err := writer.CreateBestEffort(usageCtx, usageLog); err != nil { - logger.LegacyPrintf(logKey, "Create usage log failed: %v", err) - // 计费已在此前完成，日志必须落库：dropped（批处理队列超时）同样走同步兜底， - // 否则会出现“已扣费但无 usage_log”的对账缺口（issue #3656）。 - // 重复写入由 usage_logs 的 ON CONFLICT (request_id, api_key_id) DO NOTHING 防护。 - fallbackCtx := usageCtx - if usageCtx.Err() != nil { - // usageCtx 已耗尽（best-effort 入队阻塞到期限）：换新的 detached 窗口，避免兜底必然失败。 - var fallbackCancel context.CancelFunc - fallbackCtx, fallbackCancel = detachedBillingContext(context.Background()) - defer fallbackCancel() - } - if _, syncErr := repo.Create(fallbackCtx, usageLog); syncErr != nil { - logger.LegacyPrintf(logKey, "Create usage log sync fallback failed: %v", syncErr) - } - } - return - } - - if _, err := repo.Create(usageCtx, usageLog); err != nil { - logger.LegacyPrintf(logKey, "Create usage log failed: %v", err) - } -} - -// recordUsageOpts 内部选项，参数化普通计费与长上下文计费的差异点。 -type recordUsageOpts struct { - // 长上下文计费（仅 Gemini 路径需要） - LongContextThreshold int - LongContextMultiplier float64 -} - -// RecordUsage 记录使用量并扣费（或更新订阅用量） -func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInput) error { - return s.recordUsageCore(ctx, &recordUsageCoreInput{ - Result: input.Result, - APIKey: input.APIKey, - User: input.User, - Account: input.Account, - Subscription: input.Subscription, - InboundEndpoint: input.InboundEndpoint, - UpstreamEndpoint: input.UpstreamEndpoint, - UserAgent: input.UserAgent, - IPAddress: input.IPAddress, - RequestPayloadHash: input.RequestPayloadHash, - ForceCacheBilling: input.ForceCacheBilling, - APIKeyService: input.APIKeyService, - QuotaPlatform: input.QuotaPlatform, - ChannelUsageFields: input.ChannelUsageFields, - }, &recordUsageOpts{}) -} - -// RecordUsageLongContextInput 记录使用量的输入参数（支持长上下文双倍计费） -type RecordUsageLongContextInput struct { - Result *ForwardResult - APIKey *APIKey - User *User - Account *Account - Subscription *UserSubscription // 可选：订阅信息 - InboundEndpoint string // 入站端点（客户端请求路径） - UpstreamEndpoint string // 上游端点（标准化后的上游路径） - UserAgent string // 请求的 User-Agent - IPAddress string // 请求的客户端 IP 地址 - RequestPayloadHash string // 请求体语义哈希，用于降低 request_id 误复用时的静默误去重风险 - LongContextThreshold int // 长上下文阈值（如 200000） - LongContextMultiplier float64 // 超出阈值部分的倍率（如 2.0） - ForceCacheBilling bool // 强制缓存计费：将 input_tokens 转为 cache_read 计费（用于粘性会话切换） - APIKeyService APIKeyQuotaUpdater // API Key 配额服务（可选） - QuotaPlatform string // user×platform 配额计量平台：handler 在请求 ctx 内经 QuotaPlatform() 算定后传入（后扣运行在 worker 池 background ctx 上，取不到 ForcePlatform） - - ChannelUsageFields // 渠道映射信息（由 handler 在 Forward 前解析） -} - -// RecordUsageWithLongContext 记录使用量并扣费，支持长上下文双倍计费（用于 Gemini） -func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input *RecordUsageLongContextInput) error { - return s.recordUsageCore(ctx, &recordUsageCoreInput{ - Result: input.Result, - APIKey: input.APIKey, - User: input.User, - Account: input.Account, - Subscription: input.Subscription, - InboundEndpoint: input.InboundEndpoint, - UpstreamEndpoint: input.UpstreamEndpoint, - UserAgent: input.UserAgent, - IPAddress: input.IPAddress, - RequestPayloadHash: input.RequestPayloadHash, - ForceCacheBilling: input.ForceCacheBilling, - APIKeyService: input.APIKeyService, - QuotaPlatform: input.QuotaPlatform, - ChannelUsageFields: input.ChannelUsageFields, - }, &recordUsageOpts{ - LongContextThreshold: input.LongContextThreshold, - LongContextMultiplier: input.LongContextMultiplier, - }) -} - -// recordUsageCoreInput 是 recordUsageCore 的公共输入字段，从两种输入结构体中提取。 -type recordUsageCoreInput struct { - Result *ForwardResult - APIKey *APIKey - User *User - Account *Account - Subscription *UserSubscription - InboundEndpoint string - UpstreamEndpoint string - UserAgent string - IPAddress string - RequestPayloadHash string - ForceCacheBilling bool - APIKeyService APIKeyQuotaUpdater - QuotaPlatform string - ChannelUsageFields -} - -// recordUsageCore 是 RecordUsage 和 RecordUsageWithLongContext 的统一实现。 -// LongContextThreshold > 0 时 Token 计费回退走 CalculateCostWithLongContext。 -func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsageCoreInput, opts *recordUsageOpts) error { - result := input.Result - apiKey := input.APIKey - user := input.User - account := input.Account - subscription := input.Subscription - ApplyForwardImageBillingResolution(result) - - // 强制缓存计费：将 input_tokens 转为 cache_read_input_tokens - // 用于粘性会话切换时的特殊计费处理 - if input.ForceCacheBilling && result.Usage.InputTokens > 0 { - logger.LegacyPrintf("service.gateway", "force_cache_billing: %d input_tokens → cache_read_input_tokens (account=%d)", - result.Usage.InputTokens, account.ID) - result.Usage.CacheReadInputTokens += result.Usage.InputTokens - result.Usage.InputTokens = 0 - } - - // Cache TTL Override: 确保计费时 token 分类与账号设置一致。 - // 账号级设置优先；全局 1h 请求注入开启时，默认把 usage 计费归回 5m。 - cacheTTLOverridden := false - if overrideTarget, ok := s.resolveCacheTTLUsageOverrideTarget(ctx, account); ok { - applyCacheTTLOverride(&result.Usage, overrideTarget) - cacheTTLOverridden = (result.Usage.CacheCreation5mTokens + result.Usage.CacheCreation1hTokens) > 0 - } - - // 获取费率倍数（优先级：用户专属 > 分组默认 > 系统默认） - multiplier := 1.0 - if s.cfg != nil { - multiplier = s.cfg.Default.RateMultiplier - } - if apiKey.GroupID != nil && apiKey.Group != nil { - groupDefault := apiKey.Group.RateMultiplier - multiplier = s.getUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, groupDefault) - } - // token 倍率叠加高峰因子（token 计费含图片 token，图片按次倍率不受影响）。高峰因子按请求时刻现算， - // 不并入上面的 getUserGroupRateMultiplier，以免污染 user:group 倍率缓存。 - multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, timezone.Now()) - - // 确定计费模型 - billingModel := forwardResultBillingModel(result.Model, result.UpstreamModel) - if input.BillingModelSource == BillingModelSourceChannelMapped && input.ChannelMappedModel != "" { - billingModel = input.ChannelMappedModel - } - if input.BillingModelSource == BillingModelSourceRequested && input.OriginalModel != "" { - billingModel = input.OriginalModel - } - - // 确定 RequestedModel（渠道映射前的原始模型） - requestedModel := result.Model - if input.OriginalModel != "" { - requestedModel = input.OriginalModel - } - - // 计算费用 - cost := s.calculateRecordUsageCost(ctx, result, apiKey, billingModel, multiplier, imageMultiplier, opts) - - // 判断计费方式：订阅模式 vs 余额模式 - isSubscriptionBilling := subscription != nil && apiKey.Group != nil && apiKey.Group.IsSubscriptionType() - billingType := BillingTypeBalance - if isSubscriptionBilling { - billingType = BillingTypeSubscription - } - - // 创建使用日志 - accountRateMultiplier := account.BillingRateMultiplier() - usageLog := s.buildRecordUsageLog(ctx, input, result, apiKey, user, account, subscription, - requestedModel, multiplier, imageMultiplier, accountRateMultiplier, billingType, cacheTTLOverridden, cost, opts) - - // 计算账号统计定价费用（使用最终上游模型匹配自定义规则） - if apiKey.GroupID != nil { - applyAccountStatsCost(ctx, usageLog, s.channelService, s.billingService, - account.ID, *apiKey.GroupID, result.UpstreamModel, result.Model, - // Anthropic's input_tokens excludes cache_read and cache_creation (billed separately); - // OpenAI gateway uses actualInputTokens which also excludes cache_read for the same reason. - UsageTokens{ - InputTokens: result.Usage.InputTokens, - OutputTokens: result.Usage.OutputTokens, - CacheCreationTokens: result.Usage.CacheCreationInputTokens, - CacheReadTokens: result.Usage.CacheReadInputTokens, - ImageOutputTokens: result.Usage.ImageOutputTokens, - }, - cost.TotalCost, - ) - } - - if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { - writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway") - logger.LegacyPrintf("service.gateway", "[SIMPLE MODE] Usage recorded (not billed): user=%d, tokens=%d", usageLog.UserID, usageLog.TotalTokens()) - s.deferredService.ScheduleLastUsedUpdate(account.ID) - return nil - } - - // 配额平台由 handler 在请求 ctx 内经 QuotaPlatform() 算定并通过 input 传入； - // 后扣运行在 worker 池的 background ctx 上，无法再从 ctx 取 ForcePlatform。 - // 缺省（未设置）时回退到分组平台，保持对其它调用方的兼容。 - quotaPlatform := input.QuotaPlatform - if quotaPlatform == "" { - quotaPlatform = PlatformFromAPIKey(apiKey) - } - requestID := usageLog.RequestID - _, billingErr := applyUsageBilling(ctx, requestID, usageLog, &postUsageBillingParams{ - Cost: cost, - User: user, - APIKey: apiKey, - Account: account, - Subscription: subscription, - RequestPayloadHash: resolveUsageBillingPayloadFingerprint(ctx, input.RequestPayloadHash), - IsSubscriptionBill: isSubscriptionBilling, - AccountRateMultiplier: accountRateMultiplier, - APIKeyService: input.APIKeyService, - Platform: quotaPlatform, - }, s.billingDeps(), s.usageBillingRepo) - - if billingErr != nil { - return billingErr - } - writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway") - - return nil -} - -// calculateRecordUsageCost 根据请求类型和选项计算费用。 -func (s *GatewayService) calculateRecordUsageCost( - ctx context.Context, - result *ForwardResult, - apiKey *APIKey, - billingModel string, - multiplier float64, - imageMultiplier float64, - opts *recordUsageOpts, -) *CostBreakdown { - // 图片生成：渠道定价为 token 计费时走 token 路径，否则走图片计费 - if result.ImageCount > 0 { - if resolved := s.resolveChannelPricing(ctx, billingModel, apiKey); resolved != nil && resolved.Mode == BillingModeToken { - return s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts) - } - return s.calculateImageCost(ctx, result, apiKey, billingModel, imageMultiplier) - } - - // Token 计费 - return s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts) -} - -// resolveChannelPricing 检查指定模型是否存在渠道级别定价。 -// 返回非 nil 的 ResolvedPricing 表示有渠道定价，nil 表示走默认定价路径。 -func (s *GatewayService) resolveChannelPricing(ctx context.Context, billingModel string, apiKey *APIKey) *ResolvedPricing { - if s.resolver == nil || apiKey.Group == nil { - return nil - } - gid := apiKey.Group.ID - resolved := s.resolver.Resolve(ctx, PricingInput{Model: billingModel, GroupID: &gid}) - if resolved.Source == PricingSourceChannel { - return resolved - } - return nil -} - -// calculateImageCost 计算图片生成费用：渠道级别定价优先，否则走按次计费。 -func (s *GatewayService) calculateImageCost( - ctx context.Context, - result *ForwardResult, - apiKey *APIKey, - billingModel string, - multiplier float64, -) *CostBreakdown { - sizeTier := NormalizeImageBillingTierOrDefault(result.ImageSize) - if resolved := s.resolveChannelPricing(ctx, billingModel, apiKey); resolved != nil { - tokens := UsageTokens{ - InputTokens: result.Usage.InputTokens, - OutputTokens: result.Usage.OutputTokens, - ImageOutputTokens: result.Usage.ImageOutputTokens, - } - gid := apiKey.Group.ID - cost, err := s.billingService.CalculateCostUnified(CostInput{ - Ctx: ctx, - Model: billingModel, - GroupID: &gid, - Tokens: tokens, - RequestCount: result.ImageCount, - SizeTier: sizeTier, - RateMultiplier: multiplier, - Resolver: s.resolver, - Resolved: resolved, - }) - if err != nil { - logger.LegacyPrintf("service.gateway", "Calculate image token cost failed: %v", err) - return &CostBreakdown{ActualCost: 0} - } - return cost - } - - var groupConfig *ImagePriceConfig - if apiKey.Group != nil { - groupConfig = &ImagePriceConfig{ - Price1K: apiKey.Group.ImagePrice1K, - Price2K: apiKey.Group.ImagePrice2K, - Price4K: apiKey.Group.ImagePrice4K, - } - } - return s.billingService.CalculateImageCost(billingModel, sizeTier, result.ImageCount, groupConfig, multiplier) -} - -// calculateTokenCost 计算 Token 计费：根据 opts 决定走普通/长上下文/渠道统一计费。 -func (s *GatewayService) calculateTokenCost( - ctx context.Context, - result *ForwardResult, - apiKey *APIKey, - billingModel string, - multiplier float64, - opts *recordUsageOpts, -) *CostBreakdown { - tokens := UsageTokens{ - InputTokens: result.Usage.InputTokens, - OutputTokens: result.Usage.OutputTokens, - CacheCreationTokens: result.Usage.CacheCreationInputTokens, - CacheReadTokens: result.Usage.CacheReadInputTokens, - CacheCreation5mTokens: result.Usage.CacheCreation5mTokens, - CacheCreation1hTokens: result.Usage.CacheCreation1hTokens, - ImageOutputTokens: result.Usage.ImageOutputTokens, - } - - var cost *CostBreakdown - var err error - - // 优先尝试渠道定价 → CalculateCostUnified - if resolved := s.resolveChannelPricing(ctx, billingModel, apiKey); resolved != nil { - gid := apiKey.Group.ID - cost, err = s.billingService.CalculateCostUnified(CostInput{ - Ctx: ctx, - Model: billingModel, - GroupID: &gid, - Tokens: tokens, - RequestCount: 1, - RateMultiplier: multiplier, - Resolver: s.resolver, - Resolved: resolved, - }) - } else if opts.LongContextThreshold > 0 { - // 长上下文双倍计费（如 Gemini 200K 阈值） - cost, err = s.billingService.CalculateCostWithLongContext(billingModel, tokens, multiplier, opts.LongContextThreshold, opts.LongContextMultiplier) - } else { - cost, err = s.billingService.CalculateCost(billingModel, tokens, multiplier) - } - if err != nil { - logger.LegacyPrintf("service.gateway", "Calculate cost failed: %v", err) - return &CostBreakdown{ActualCost: 0} - } - return cost -} - -// buildRecordUsageLog 构建使用日志并设置计费模式。 -func (s *GatewayService) buildRecordUsageLog( - ctx context.Context, - input *recordUsageCoreInput, - result *ForwardResult, - apiKey *APIKey, - user *User, - account *Account, - subscription *UserSubscription, - requestedModel string, - multiplier float64, - imageMultiplier float64, - accountRateMultiplier float64, - billingType int8, - cacheTTLOverridden bool, - cost *CostBreakdown, - opts *recordUsageOpts, -) *UsageLog { - durationMs := int(result.Duration.Milliseconds()) - requestID := resolveUsageBillingRequestID(ctx, result.RequestID) - usageLog := &UsageLog{ - UserID: user.ID, - APIKeyID: apiKey.ID, - AccountID: account.ID, - RequestID: requestID, - Model: result.Model, - RequestedModel: requestedModel, - UpstreamModel: optionalNonEqualStringPtr(result.UpstreamModel, result.Model), - ReasoningEffort: result.ReasoningEffort, - InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint), - UpstreamEndpoint: optionalTrimmedStringPtr(input.UpstreamEndpoint), - InputTokens: result.Usage.InputTokens, - OutputTokens: result.Usage.OutputTokens, - CacheCreationTokens: result.Usage.CacheCreationInputTokens, - CacheReadTokens: result.Usage.CacheReadInputTokens, - CacheCreation5mTokens: result.Usage.CacheCreation5mTokens, - CacheCreation1hTokens: result.Usage.CacheCreation1hTokens, - ImageOutputTokens: result.Usage.ImageOutputTokens, - RateMultiplier: multiplier, - AccountRateMultiplier: &accountRateMultiplier, - BillingType: billingType, - BillingMode: resolveBillingMode(result, cost), - Stream: result.Stream, - DurationMs: &durationMs, - FirstTokenMs: result.FirstTokenMs, - ImageCount: result.ImageCount, - ImageSize: optionalTrimmedStringPtr(result.ImageSize), - ImageInputSize: optionalTrimmedStringPtr(result.ImageInputSize), - ImageOutputSize: optionalTrimmedStringPtr(result.ImageOutputSize), - ImageSizeSource: optionalTrimmedStringPtr(result.ImageSizeSource), - ImageSizeBreakdown: result.ImageSizeBreakdown, - CacheTTLOverridden: cacheTTLOverridden, - ChannelID: optionalInt64Ptr(input.ChannelID), - ModelMappingChain: optionalTrimmedStringPtr(input.ModelMappingChain), - UserAgent: optionalTrimmedStringPtr(input.UserAgent), - IPAddress: optionalTrimmedStringPtr(input.IPAddress), - GroupID: apiKey.GroupID, - SubscriptionID: optionalSubscriptionID(subscription), - CreatedAt: time.Now(), - } - if result.ImageCount > 0 && (cost == nil || cost.BillingMode != string(BillingModeToken)) { - usageLog.RateMultiplier = imageMultiplier - } - if cost != nil { - usageLog.InputCost = cost.InputCost - usageLog.OutputCost = cost.OutputCost - usageLog.ImageOutputCost = cost.ImageOutputCost - usageLog.CacheCreationCost = cost.CacheCreationCost - usageLog.CacheReadCost = cost.CacheReadCost - usageLog.TotalCost = cost.TotalCost - usageLog.ActualCost = cost.ActualCost - } - - return usageLog -} - -// resolveBillingMode 根据计费结果和请求类型确定计费模式。 -func resolveBillingMode(result *ForwardResult, cost *CostBreakdown) *string { - var mode string - switch { - case cost != nil && cost.BillingMode != "": - mode = cost.BillingMode - case result.ImageCount > 0: - mode = string(BillingModeImage) - default: - mode = string(BillingModeToken) - } - return &mode -} - -func optionalSubscriptionID(subscription *UserSubscription) *int64 { - if subscription != nil { - return &subscription.ID - } - return nil -} - -// ResolveChannelMapping 委托渠道服务解析模型映射 -func (s *GatewayService) ResolveChannelMapping(ctx context.Context, groupID int64, model string) ChannelMappingResult { - if s.channelService == nil { - return ChannelMappingResult{MappedModel: model} - } - return s.channelService.ResolveChannelMapping(ctx, groupID, model) -} - -// ReplaceModelInBody 替换请求体中的模型名（导出供 handler 使用） -func (s *GatewayService) ReplaceModelInBody(body []byte, newModel string) []byte { - return ReplaceModelInBody(body, newModel) -} - -// IsModelRestricted 检查模型是否被渠道限制 -func (s *GatewayService) IsModelRestricted(ctx context.Context, groupID int64, model string) bool { - if s.channelService == nil { - return false - } - return s.channelService.IsModelRestricted(ctx, groupID, model) -} - -// ResolveChannelMappingAndRestrict 解析渠道映射。 -// 模型限制检查已移至调度阶段（checkChannelPricingRestriction），restricted 始终返回 false。 -func (s *GatewayService) ResolveChannelMappingAndRestrict(ctx context.Context, groupID *int64, model string) (ChannelMappingResult, bool) { - if s.channelService == nil { - return ChannelMappingResult{MappedModel: model}, false - } - return s.channelService.ResolveChannelMappingAndRestrict(ctx, groupID, model) -} - -// checkChannelPricingRestriction 根据渠道计费基准检查模型是否受定价列表限制。 -// 供调度阶段预检查（requested / channel_mapped）。 -// upstream 需逐账号检查，此处返回 false。 -func (s *GatewayService) checkChannelPricingRestriction(ctx context.Context, groupID *int64, requestedModel string) bool { - if groupID == nil || s.channelService == nil || requestedModel == "" { - return false - } - mapping := s.channelService.ResolveChannelMapping(ctx, *groupID, requestedModel) - billingModel := billingModelForRestriction(mapping.BillingModelSource, requestedModel, mapping.MappedModel) - if billingModel == "" { - return false - } - return s.channelService.IsModelRestricted(ctx, *groupID, billingModel) -} - -// billingModelForRestriction 根据计费基准确定限制检查使用的模型。 -// upstream 返回空（需逐账号检查）。 -func billingModelForRestriction(source, requestedModel, channelMappedModel string) string { - switch source { - case BillingModelSourceRequested: - return requestedModel - case BillingModelSourceUpstream: - return "" - case BillingModelSourceChannelMapped: - return channelMappedModel - default: - return channelMappedModel - } -} - -// isUpstreamModelRestrictedByChannel 检查账号映射后的上游模型是否受渠道定价限制。 -// 仅在 BillingModelSource="upstream" 且 RestrictModels=true 时由调度循环调用。 -func (s *GatewayService) isUpstreamModelRestrictedByChannel(ctx context.Context, groupID int64, account *Account, requestedModel string) bool { - if s.channelService == nil { - return false - } - upstreamModel := resolveAccountUpstreamModel(account, requestedModel) - if upstreamModel == "" { - return false - } - return s.channelService.IsModelRestricted(ctx, groupID, upstreamModel) -} - -// resolveAccountUpstreamModel 确定账号将请求模型映射为什么上游模型。 -func resolveAccountUpstreamModel(account *Account, requestedModel string) string { - if account.Platform == PlatformAntigravity { - return mapAntigravityModel(account, requestedModel) - } - return account.GetMappedModel(requestedModel) -} - -// needsUpstreamChannelRestrictionCheck 判断是否需要在调度循环中逐账号检查上游模型的渠道限制。 -func (s *GatewayService) needsUpstreamChannelRestrictionCheck(ctx context.Context, groupID *int64) bool { - if groupID == nil || s.channelService == nil { - return false - } - ch, err := s.channelService.GetChannelForGroup(ctx, *groupID) - if err != nil { - slog.Warn("failed to check channel upstream restriction", "group_id", *groupID, "error", err) - return false - } - if ch == nil || !ch.RestrictModels { - return false - } - return ch.BillingModelSource == BillingModelSourceUpstream -} - -// isStickyAccountUpstreamRestricted 检查粘性会话命中的账号是否受 upstream 渠道限制。 -// 合并 needsUpstreamChannelRestrictionCheck + isUpstreamModelRestrictedByChannel 两步调用， -// 供 sticky session 条件链使用，避免内联多个函数调用导致行过长。 -func (s *GatewayService) isStickyAccountUpstreamRestricted(ctx context.Context, groupID *int64, account *Account, requestedModel string) bool { - if groupID == nil { - return false - } - if !s.needsUpstreamChannelRestrictionCheck(ctx, groupID) { - return false - } - return s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel) -} - -// ForwardCountTokens 转发 count_tokens 请求到上游 API -// 特点：不记录使用量、仅支持非流式响应 -func (s *GatewayService) ForwardCountTokens(ctx context.Context, c *gin.Context, account *Account, parsed *ParsedRequest) error { - if parsed == nil { - s.countTokensError(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty") - return fmt.Errorf("parse request: empty request") - } - - if account != nil && account.IsAnthropicAPIKeyPassthroughEnabled() { - passthroughBody := parsed.Body.Bytes() - if reqModel := parsed.Model; reqModel != "" { - if mappedModel := account.GetMappedModel(reqModel); mappedModel != reqModel { - passthroughBody = s.replaceModelInBody(passthroughBody, mappedModel) - logger.LegacyPrintf("service.gateway", "CountTokens passthrough model mapping: %s -> %s (account: %s)", reqModel, mappedModel, account.Name) - } - } - return s.forwardCountTokensAnthropicAPIKeyPassthrough(ctx, c, account, passthroughBody) - } - - // Bedrock 不支持 count_tokens 端点 - if account != nil && account.IsBedrock() { - s.countTokensError(c, http.StatusNotFound, "not_found_error", "count_tokens endpoint is not supported for Bedrock") - return nil - } - - body := parsed.Body.Bytes() - replaceBody := func(next []byte) error { - if err := parsed.ReplaceBody(next); err != nil { - return fmt.Errorf("rewrite count_tokens body: %w", err) - } - body = parsed.Body.Bytes() - return nil - } - reqModel := parsed.Model - - // Pre-filter: strip empty text blocks to prevent upstream 400. - if err := replaceBody(StripEmptyTextBlocks(body)); err != nil { - return err - } - - isClaudeCodeCT := IsClaudeCodeClient(ctx) || isClaudeCodeClient(c.GetHeader("User-Agent"), parsed.MetadataUserID) - shouldMimicClaudeCode := account.IsOAuth() && !isClaudeCodeCT - - if shouldMimicClaudeCode { - normalizeOpts := claudeOAuthNormalizeOptions{stripSystemCacheControl: true} - var normalizedBody []byte - normalizedBody, reqModel = normalizeClaudeOAuthRequestBody(body, reqModel, normalizeOpts) - if err := replaceBody(normalizedBody); err != nil { - return err - } - - if err := replaceBody(s.rewriteMessageCacheControlIfEnabled(ctx, body)); err != nil { - return err - } - if rw := buildToolNameRewriteFromBody(body); rw != nil { - if err := replaceBody(applyToolNameRewriteToBody(body, rw)); err != nil { - return err - } - } else { - if err := replaceBody(applyToolsLastCacheBreakpoint(body)); err != nil { - return err - } - } - } - - // Antigravity 账户不支持 count_tokens，返回 404 让客户端 fallback 到本地估算。 - // 返回 nil 避免 handler 层记录为错误，也不设置 ops 上游错误上下文。 - if account.Platform == PlatformAntigravity { - s.countTokensError(c, http.StatusNotFound, "not_found_error", "count_tokens endpoint is not supported for this platform") - return nil - } - - // 应用模型映射： - // - APIKey 账号：使用账号级别的显式映射（如果配置），否则透传原始模型名 - // - OAuth/SetupToken 账号：使用 Anthropic 标准映射（短ID → 长ID） - if reqModel != "" { - mappedModel := reqModel - mappingSource := "" - if account.Type == AccountTypeAPIKey { - mappedModel = account.GetMappedModel(reqModel) - if mappedModel != reqModel { - mappingSource = "account" - } - } - if mappingSource == "" && account.Platform == PlatformAnthropic && account.Type != AccountTypeAPIKey { - normalized := claude.NormalizeModelID(reqModel) - if normalized != reqModel { - mappedModel = normalized - mappingSource = "prefix" - } - } - if mappedModel != reqModel { - originalReqModel := reqModel - if err := replaceBody(s.replaceModelInBody(body, mappedModel)); err != nil { - return err - } - reqModel = mappedModel - parsed.Model = mappedModel - logger.LegacyPrintf("service.gateway", "CountTokens model mapping applied: %s -> %s (account: %s, source=%s)", originalReqModel, mappedModel, account.Name, mappingSource) - } - } - - // 获取凭证 - token, tokenType, err := s.GetAccessToken(ctx, account) - if err != nil { - s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to get access token") - return err - } - - // 构建上游请求 - upstreamReq, wireBody, err := s.buildCountTokensRequest(ctx, c, account, body, token, tokenType, reqModel, shouldMimicClaudeCode) - if err != nil { - s.countTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request") - return err - } - // 先记录首发 wire body；如果后面进入 400 retry，retry 会基于未签名的逻辑 body 重新构建。 - acceptedWireBody := wireBody - - // 获取代理URL（自定义 base URL 模式下，proxy 通过 buildCustomRelayURL 作为查询参数传递） - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - if !account.IsCustomBaseURLEnabled() || account.GetCustomBaseURL() == "" { - proxyURL = account.Proxy.URL() - } - } - - // 发送请求 - resp, err := s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account)) - if err != nil { - setOpsUpstreamError(c, 0, sanitizeUpstreamErrorMessage(err.Error()), "") - s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Request failed") - return fmt.Errorf("upstream request failed: %w", err) - } - - // 读取响应体 - countTokensTooLarge := func(c *gin.Context) { - s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream response too large") - } - respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, countTokensTooLarge) - _ = resp.Body.Close() - if err != nil { - if !errors.Is(err, ErrUpstreamResponseBodyTooLarge) { - s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response") - } - return err - } - - // 检测 thinking block 签名错误（400）并重试一次（过滤 thinking blocks） - if resp.StatusCode == 400 && s.shouldRectifySignatureError(ctx, account, respBody, reqModel) { - logger.LegacyPrintf("service.gateway", "Account %d: detected thinking block signature error on count_tokens, retrying with filtered thinking blocks", account.ID) - - filteredBody := FilterThinkingBlocksForRetry(body, reqModel) - retryReq, retryWireBody, buildErr := s.buildCountTokensRequest(ctx, c, account, filteredBody, token, tokenType, reqModel, shouldMimicClaudeCode) - if buildErr == nil { - retryResp, retryErr := s.httpUpstream.DoWithTLS(retryReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account)) - if retryErr == nil { - if retryResp.StatusCode < 400 { - // count_tokens 签名重试成功后记录最终 wire body，错误响应仍保留原 body 便于后续处理。 - acceptedWireBody = retryWireBody - } - resp = retryResp - respBody, err = ReadUpstreamResponseBody(resp.Body, s.cfg, c, countTokensTooLarge) - _ = resp.Body.Close() - if err != nil { - if !errors.Is(err, ErrUpstreamResponseBodyTooLarge) { - s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response") - } - return err - } - } - } - } - - if resp.StatusCode < 400 && !bytes.Equal(acceptedWireBody, body) { - // count_tokens 成功后再同步最终 wire body，避免 retry 从已签名 body 派生。 - if err := replaceBody(acceptedWireBody); err != nil { - return err - } - } - - // 处理错误响应 - if resp.StatusCode >= 400 { - // 标记账号状态（429/529等） - s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) - - upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - upstreamDetail := "" - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes - if maxBytes <= 0 { - maxBytes = 2048 - } - upstreamDetail = truncateString(string(respBody), maxBytes) - } - setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) - - // 记录上游错误摘要便于排障（不回显请求内容） - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - logger.LegacyPrintf("service.gateway", - "count_tokens upstream error %d (account=%d platform=%s type=%s): %s", - resp.StatusCode, - account.ID, - account.Platform, - account.Type, - truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), - ) - } - - // 返回简化的错误响应 - errMsg := "Upstream request failed" - switch resp.StatusCode { - case 429: - errMsg = "Rate limit exceeded" - case 529: - errMsg = "Service overloaded" - } - s.countTokensError(c, resp.StatusCode, "upstream_error", errMsg) - if upstreamMsg == "" { - return fmt.Errorf("upstream error: %d", resp.StatusCode) - } - return fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, upstreamMsg) - } - - // 透传成功响应 - c.Data(resp.StatusCode, "application/json", respBody) - return nil -} - -func (s *GatewayService) forwardCountTokensAnthropicAPIKeyPassthrough(ctx context.Context, c *gin.Context, account *Account, body []byte) error { - token, tokenType, err := s.GetAccessToken(ctx, account) - if err != nil { - s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to get access token") - return err - } - if tokenType != "apikey" { - s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Invalid account token type") - return fmt.Errorf("anthropic api key passthrough requires apikey token, got: %s", tokenType) - } - - upstreamReq, err := s.buildCountTokensRequestAnthropicAPIKeyPassthrough(ctx, c, account, body, token) - if err != nil { - s.countTokensError(c, http.StatusInternalServerError, "api_error", "Failed to build request") - return err - } - - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - resp, err := s.httpUpstream.DoWithTLS(upstreamReq, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account)) - if err != nil { - setOpsUpstreamError(c, 0, sanitizeUpstreamErrorMessage(err.Error()), "") - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: 0, - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Passthrough: true, - Kind: "request_error", - Message: sanitizeUpstreamErrorMessage(err.Error()), - }) - s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Request failed") - return fmt.Errorf("upstream request failed: %w", err) - } - - countTokensTooLarge := func(c *gin.Context) { - s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Upstream response too large") - } - respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, countTokensTooLarge) - _ = resp.Body.Close() - if err != nil { - if !errors.Is(err, ErrUpstreamResponseBodyTooLarge) { - s.countTokensError(c, http.StatusBadGateway, "upstream_error", "Failed to read response") - } - return err - } - - if resp.StatusCode >= 400 { - if s.rateLimitService != nil { - s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) - } - - upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) - upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) - - // 中转站不支持 count_tokens 端点时（404），返回 404 让客户端 fallback 到本地估算。 - // 仅在错误消息明确指向 count_tokens endpoint 不存在时生效，避免误吞其他 404（如错误 base_url）。 - // 返回 nil 避免 handler 层记录为错误，也不设置 ops 上游错误上下文。 - if isCountTokensUnsupported404(resp.StatusCode, respBody) { - logger.LegacyPrintf("service.gateway", - "[count_tokens] Upstream does not support count_tokens (404), returning 404: account=%d name=%s msg=%s", - account.ID, account.Name, truncateString(upstreamMsg, 512)) - s.countTokensError(c, http.StatusNotFound, "not_found_error", "count_tokens endpoint is not supported by upstream") - return nil - } - - upstreamDetail := "" - if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { - maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes - if maxBytes <= 0 { - maxBytes = 2048 - } - upstreamDetail = truncateString(string(respBody), maxBytes) - } - setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) - appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ - Platform: account.Platform, - AccountID: account.ID, - AccountName: account.Name, - UpstreamStatusCode: resp.StatusCode, - UpstreamRequestID: resp.Header.Get("x-request-id"), - UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()), - Passthrough: true, - Kind: "http_error", - Message: upstreamMsg, - Detail: upstreamDetail, - }) - - errMsg := "Upstream request failed" - switch resp.StatusCode { - case 429: - errMsg = "Rate limit exceeded" - case 529: - errMsg = "Service overloaded" - } - s.countTokensError(c, resp.StatusCode, "upstream_error", errMsg) - if upstreamMsg == "" { - return fmt.Errorf("upstream error: %d", resp.StatusCode) - } - return fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, upstreamMsg) - } - - writeAnthropicPassthroughResponseHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) - contentType := strings.TrimSpace(resp.Header.Get("Content-Type")) - if contentType == "" { - contentType = "application/json" - } - c.Data(resp.StatusCode, contentType, respBody) - return nil -} - -func (s *GatewayService) buildCountTokensRequestAnthropicAPIKeyPassthrough( - ctx context.Context, - c *gin.Context, - account *Account, - body []byte, - token string, -) (*http.Request, error) { - targetURL := claudeAPICountTokensURL - baseURL := account.GetBaseURL() - if baseURL != "" { - validatedURL, err := s.validateUpstreamBaseURL(baseURL) - if err != nil { - return nil, err - } - targetURL = validatedURL + "/v1/messages/count_tokens?beta=true" - } - body = sanitizeCountTokensRequestBody(body) - - // 同 buildUpstreamRequestAnthropicAPIKeyPassthrough：能力维度 sanitize。 - clientBeta := "" - if c != nil && c.Request != nil { - clientBeta = getHeaderRaw(c.Request.Header, "anthropic-beta") - } - // 账号覆写了 anthropic-beta 时，覆写值即最终上游值：净化以覆写值为准 - if beta, ok := account.HeaderOverrideValue("anthropic-beta"); ok { - clientBeta = beta - } - if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, clientBeta); changed { - body = sanitized - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body)) - if err != nil { - return nil, err - } - - if c != nil && c.Request != nil { - for key, values := range c.Request.Header { - lowerKey := strings.ToLower(strings.TrimSpace(key)) - if !allowedHeaders[lowerKey] { - continue - } - wireKey := resolveWireCasing(key) - for _, v := range values { - addHeaderRaw(req.Header, wireKey, v) - } - } - } - - req.Header.Del("authorization") - req.Header.Del("x-api-key") - req.Header.Del("x-goog-api-key") - req.Header.Del("cookie") - setAnthropicAPIKeyAuthHeader(req.Header, account, token) - - if req.Header.Get("content-type") == "" { - req.Header.Set("content-type", "application/json") - } - if req.Header.Get("anthropic-version") == "" { - req.Header.Set("anthropic-version", "2023-06-01") - } - - // 账号级请求头覆写（最终生效，覆盖上面所有来源的同名头） - account.ApplyHeaderOverrides(req.Header) - - return req, nil -} - -// buildCountTokensRequest 构建 count_tokens 上游请求 -func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, mimicClaudeCode bool) (*http.Request, []byte, error) { - // 确定目标 URL - targetURL := claudeAPICountTokensURL - if account.Type == AccountTypeAPIKey { - baseURL := account.GetBaseURL() - if baseURL != "" { - validatedURL, err := s.validateUpstreamBaseURL(baseURL) - if err != nil { - return nil, nil, err - } - targetURL = validatedURL + "/v1/messages/count_tokens?beta=true" - } - } else if account.IsCustomBaseURLEnabled() { - customURL := account.GetCustomBaseURL() - if customURL == "" { - return nil, nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) - } - validatedURL, err := s.validateUpstreamBaseURL(customURL) - if err != nil { - return nil, nil, err - } - targetURL = s.buildCustomRelayURL(validatedURL, "/v1/messages/count_tokens", account) - } - - clientHeaders := http.Header{} - if c != nil && c.Request != nil { - clientHeaders = c.Request.Header - } - - // OAuth 账号：应用统一指纹和重写 userID（受设置开关控制） - // 如果启用了会话ID伪装，会在重写后替换 session 部分为固定值 - ctEnableFP, ctEnableMPT := true, false - if s.settingService != nil { - ctEnableFP, ctEnableMPT, _ = s.settingService.GetGatewayForwardingSettings(ctx) - } - var ctFingerprint *Fingerprint - if account.IsOAuth() && s.identityService != nil { - fp, err := s.identityService.GetOrCreateFingerprint(ctx, account.ID, clientHeaders) - if err == nil { - ctFingerprint = fp - if !ctEnableMPT { - accountUUID := account.GetExtraString("account_uuid") - if accountUUID != "" && fp.ClientID != "" { - if newBody, err := s.identityService.RewriteUserIDWithMasking(ctx, body, account, accountUUID, fp.ClientID, fp.UserAgent); err == nil && len(newBody) > 0 { - body = newBody - } - } - } - } - } - - // 同步 billing header cc_version 与实际发送的 User-Agent 版本 - if ctFingerprint != nil && ctEnableFP { - body = syncBillingHeaderVersion(body, ctFingerprint.UserAgent) - } - - // === 计算最终 anthropic-beta header（先于 body sanitize 与 CCH 签名）=== - // 顺序约束同 buildUpstreamRequest。 - ctEffectiveDropSet := mergeDropSets(s.getBetaPolicyFilterSet(ctx, c, account, modelID)) - finalBetaHeader, finalBetaShouldSet := s.computeFinalCountTokensAnthropicBeta( - tokenType, mimicClaudeCode, modelID, clientHeaders, body, ctEffectiveDropSet, - ) - - // 账号覆写了 anthropic-beta 时，覆写值即最终上游值：净化以覆写值为准 - if beta, ok := account.HeaderOverrideValue("anthropic-beta"); ok { - finalBetaHeader, finalBetaShouldSet = beta, true - } - - // 能力维度 body sanitize：与最终 anthropic-beta header 对称 - if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, finalBetaHeader); changed { - body = sanitized - } - - body = sanitizeCountTokensRequestBody(body) - - req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body)) - if err != nil { - return nil, nil, err - } - - // 设置认证头（保持原始大小写） - if tokenType == "oauth" { - setHeaderRaw(req.Header, "authorization", "Bearer "+token) - } else { - setAnthropicAPIKeyAuthHeader(req.Header, account, token) - } - - // 白名单透传 headers（恢复真实 wire casing） - for key, values := range clientHeaders { - lowerKey := strings.ToLower(key) - if allowedHeaders[lowerKey] { - wireKey := resolveWireCasing(key) - for _, v := range values { - addHeaderRaw(req.Header, wireKey, v) - } - } - } - - // OAuth 账号：应用指纹到请求头（受设置开关控制） - if ctEnableFP && ctFingerprint != nil { - s.identityService.ApplyFingerprint(req, ctFingerprint) - } - - // 确保必要的 headers 存在（保持原始大小写） - if getHeaderRaw(req.Header, "content-type") == "" { - setHeaderRaw(req.Header, "content-type", "application/json") - } - if getHeaderRaw(req.Header, "anthropic-version") == "" { - setHeaderRaw(req.Header, "anthropic-version", "2023-06-01") - } - if tokenType == "oauth" { - applyClaudeOAuthHeaderDefaults(req) - } - - // OAuth + mimic Claude Code：强制注入 CLI 指纹 header - if tokenType == "oauth" && mimicClaudeCode { - applyClaudeCodeMimicHeaders(req, false) - } - - // 写入最终 anthropic-beta header（Del 一次避免白名单透传值残留） - deleteHeaderAllForms(req.Header, "anthropic-beta") - if finalBetaShouldSet { - setHeaderRaw(req.Header, "anthropic-beta", finalBetaHeader) - } - - // 同步 X-Claude-Code-Session-Id 头：取 body 中已处理的 metadata.user_id 的 session_id 覆盖 - if sessionHeader := getHeaderRaw(req.Header, "X-Claude-Code-Session-Id"); sessionHeader != "" { - if uid := gjson.GetBytes(body, "metadata.user_id").String(); uid != "" { - if parsed := ParseMetadataUserID(uid); parsed != nil { - setHeaderRaw(req.Header, "X-Claude-Code-Session-Id", parsed.SessionID) - } - } - } - - // 账号级请求头覆写（仅 anthropic/openai api_key 账号启用时生效；OAuth 路径 no-op） - account.ApplyHeaderOverrides(req.Header) - - if c != nil && tokenType == "oauth" { - c.Set(claudeMimicDebugInfoKey, buildClaudeMimicDebugLine(req, body, account, tokenType, mimicClaudeCode)) - } - if s.debugClaudeMimicEnabled() { - logClaudeMimicDebug(req, body, account, tokenType, mimicClaudeCode) - } - - return req, body, nil -} - -func sanitizeCountTokensRequestBody(body []byte) []byte { - out := body - for _, path := range []string{ - "temperature", - "top_p", - "top_k", - "stream", - "stop_sequences", - "stop", - } { - if gjson.GetBytes(out, path).Exists() { - if next, ok := deleteJSONPathBytes(out, path); ok { - out = next - } - } - } - return out -} - -// countTokensError 返回 count_tokens 错误响应 -func (s *GatewayService) countTokensError(c *gin.Context, status int, errType, message string) { - c.JSON(status, gin.H{ - "type": "error", - "error": gin.H{ - "type": errType, - "message": message, - }, - }) -} - -// buildCustomRelayURL 构建自定义中继转发 URL -// 在 path 后附加 beta=true 和可选的 proxy 查询参数 -func (s *GatewayService) buildCustomRelayURL(baseURL, path string, account *Account) string { - u := strings.TrimRight(baseURL, "/") + path + "?beta=true" - if account.ProxyID != nil && account.Proxy != nil { - proxyURL := account.Proxy.URL() - if proxyURL != "" { - u += "&proxy=" + url.QueryEscape(proxyURL) - } - } - return u -} - -func (s *GatewayService) validateUpstreamBaseURL(raw string) (string, error) { - if s.cfg != nil && !s.cfg.Security.URLAllowlist.Enabled { - normalized, err := urlvalidator.ValidateURLFormat(raw, s.cfg.Security.URLAllowlist.AllowInsecureHTTP) - if err != nil { - return "", fmt.Errorf("invalid base_url: %w", err) - } - return normalized, nil - } - normalized, err := urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{ - AllowedHosts: s.cfg.Security.URLAllowlist.UpstreamHosts, - RequireAllowlist: true, - AllowPrivate: s.cfg.Security.URLAllowlist.AllowPrivateHosts, - }) - if err != nil { - return "", fmt.Errorf("invalid base_url: %w", err) - } - return normalized, nil -} - // GetAvailableModels returns the list of models available for a group // It aggregates model_mapping keys from all schedulable accounts in the group func (s *GatewayService) GetAvailableModels(ctx context.Context, groupID *int64, platform string) []string { @@ -7183,24 +1196,6 @@ func (s *GatewayService) InvalidateAvailableModelsCache(groupID *int64, platform } } -// reconcileCachedTokens 兼容 Kimi 等上游： -// 将 OpenAI 风格的 cached_tokens 映射到 Claude 标准的 cache_read_input_tokens -func reconcileCachedTokens(usage map[string]any) bool { - if usage == nil { - return false - } - cacheRead, _ := usage["cache_read_input_tokens"].(float64) - if cacheRead > 0 { - return false // 已有标准字段，无需处理 - } - cached, _ := usage["cached_tokens"].(float64) - if cached <= 0 { - return false - } - usage["cache_read_input_tokens"] = cached - return true -} - const debugGatewayBodyDefaultFilename = "gateway_debug.log" // initDebugGatewayBodyFile 初始化网关调试日志文件。 diff --git a/backend/internal/service/gateway_upstream_request.go b/backend/internal/service/gateway_upstream_request.go new file mode 100644 index 0000000000..ac13f0be59 --- /dev/null +++ b/backend/internal/service/gateway_upstream_request.go @@ -0,0 +1,923 @@ +package service + +import ( + "bytes" + "context" + "fmt" + "net/http" + "net/url" + "strconv" + "strings" + + "github.com/Wei-Shaw/sub2api/internal/pkg/claude" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" + "github.com/google/uuid" + "github.com/tidwall/gjson" + + "github.com/gin-gonic/gin" +) + +func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, tokenType, modelID string, reqStream bool, mimicClaudeCode bool) (*http.Request, []byte, error) { + if account.Platform == PlatformAnthropic && account.Type == AccountTypeServiceAccount { + req, err := s.buildUpstreamRequestAnthropicVertex(ctx, c, account, body, token, modelID, reqStream) + return req, body, err + } + + // 确定目标URL + targetURL := claudeAPIURL + if account.Type == AccountTypeAPIKey { + baseURL := account.GetBaseURL() + if baseURL != "" { + validatedURL, err := s.validateUpstreamBaseURL(baseURL) + if err != nil { + return nil, nil, err + } + targetURL = validatedURL + "/v1/messages?beta=true" + } + } else if account.IsCustomBaseURLEnabled() { + customURL := account.GetCustomBaseURL() + if customURL == "" { + return nil, nil, fmt.Errorf("custom_base_url is enabled but not configured for account %d", account.ID) + } + validatedURL, err := s.validateUpstreamBaseURL(customURL) + if err != nil { + return nil, nil, err + } + targetURL = s.buildCustomRelayURL(validatedURL, "/v1/messages", account) + } + + clientHeaders := http.Header{} + if c != nil && c.Request != nil { + clientHeaders = c.Request.Header + } + + // OAuth账号：应用统一指纹和metadata重写（受设置开关控制） + var fingerprint *Fingerprint + enableFP, enableMPT := true, false + if s.settingService != nil { + enableFP, enableMPT, _ = s.settingService.GetGatewayForwardingSettings(ctx) + } + if account.IsOAuth() && s.identityService != nil { + // 1. 获取或创建指纹（包含随机生成的ClientID） + fp, err := s.identityService.GetOrCreateFingerprint(ctx, account.ID, clientHeaders) + if err != nil { + logger.LegacyPrintf("service.gateway", "Warning: failed to get fingerprint for account %d: %v", account.ID, err) + // 失败时降级为透传原始headers + } else { + if enableFP { + fingerprint = fp + } + + // 2. 重写metadata.user_id（需要指纹中的ClientID和账号的account_uuid） + // 如果启用了会话ID伪装，会在重写后替换 session 部分为固定值 + // 当 metadata 透传开启时跳过重写 + if !enableMPT { + accountUUID := account.GetExtraString("account_uuid") + if accountUUID != "" && fp.ClientID != "" { + if newBody, err := s.identityService.RewriteUserIDWithMasking(ctx, body, account, accountUUID, fp.ClientID, fp.UserAgent); err == nil && len(newBody) > 0 { + body = newBody + } + } + } + } + } + + // 同步 billing header cc_version 与实际发送的 User-Agent 版本 + if fingerprint != nil { + body = syncBillingHeaderVersion(body, fingerprint.UserAgent) + } + + // === 计算最终 anthropic-beta header（先于 body sanitize 与 CCH 签名）=== + // + // 顺序约束： + // 1) 算 finalBeta（纯函数，不依赖 req.Header；mimicry 路径会忽略客户端 beta， + // 与原“OAuth + mimicClaudeCode 跳过白名单透传”行为对齐） + // 2) 按 finalBeta 做能力维度 body sanitize（如 context-management beta 缺失 → + // strip body.context_management，与 Bedrock 路径对称） + // 3) CCH 签名（必须使用 strip 后的 body，否则 hash 与最终 body 不一致 → + // 被 Anthropic 判 third-party） + // 4) NewRequest（body 至此最终敲定） + // 5) 透传白名单 / fingerprint / mimic header / 写入 finalBeta + policyFilterSet := s.getBetaPolicyFilterSet(ctx, c, account, modelID) + effectiveDropSet := mergeDropSets(policyFilterSet) + finalBetaHeader, finalBetaShouldSet := s.computeFinalAnthropicBeta( + tokenType, mimicClaudeCode, modelID, clientHeaders, body, effectiveDropSet, + ) + + // 账号覆写了 anthropic-beta 时，覆写值即最终上游值（由下方 ApplyHeaderOverrides 写入）： + // body 能力净化必须以覆写值为准，否则 header/body 不对称会被上游 400。 + if beta, ok := account.HeaderOverrideValue("anthropic-beta"); ok { + finalBetaHeader, finalBetaShouldSet = beta, true + } + + // 能力维度 body sanitize：与最终 anthropic-beta header 对称 + if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(body, finalBetaHeader); changed { + body = sanitized + } + + req, err := http.NewRequestWithContext(ctx, "POST", targetURL, bytes.NewReader(body)) + if err != nil { + return nil, nil, err + } + + // 设置认证头（保持原始大小写） + if tokenType == "oauth" { + setHeaderRaw(req.Header, "authorization", "Bearer "+token) + } else { + setAnthropicAPIKeyAuthHeader(req.Header, account, token) + } + + // 白名单透传 headers + // OAuth mimicry 路径：跳过客户端 header 透传，与 Parrot 对齐。 + // Parrot 的 build_upstream_headers 只发 9 个精确 header，不透传任何客户端 header。 + // 透传客户端 header 会引入不一致的 x-stainless-* / anthropic-beta / user-agent / + // x-claude-code-session-id 等值，和我们注入的伪装 header 冲突，被 Anthropic 判 third-party。 + if tokenType != "oauth" || !mimicClaudeCode { + for key, values := range clientHeaders { + lowerKey := strings.ToLower(key) + if allowedHeaders[lowerKey] { + wireKey := resolveWireCasing(key) + for _, v := range values { + addHeaderRaw(req.Header, wireKey, v) + } + } + } + } + + // OAuth账号：应用缓存的指纹到请求头（覆盖白名单透传的头） + if fingerprint != nil { + s.identityService.ApplyFingerprint(req, fingerprint) + } + + // 确保必要的headers存在（保持原始大小写） + if getHeaderRaw(req.Header, "content-type") == "" { + setHeaderRaw(req.Header, "content-type", "application/json") + } + if getHeaderRaw(req.Header, "anthropic-version") == "" { + setHeaderRaw(req.Header, "anthropic-version", "2023-06-01") + } + if tokenType == "oauth" { + applyClaudeOAuthHeaderDefaults(req) + } + + // OAuth + mimic Claude Code：强制注入 CLI 指纹相关 header + // （user-agent/x-stainless-*/x-app/Accept/x-stainless-helper-method/x-client-request-id） + if tokenType == "oauth" && mimicClaudeCode { + applyClaudeCodeMimicHeaders(req, reqStream) + } + + // 写入最终 anthropic-beta header + // 注：透传分支白名单可能写入了客户端 anthropic-beta，无条件 Del 一次再按 finalBeta + // 决定是否 set，确保 dropSet 过滤后的结果一定覆盖客户端原始值。 + deleteHeaderAllForms(req.Header, "anthropic-beta") + if finalBetaShouldSet { + setHeaderRaw(req.Header, "anthropic-beta", finalBetaHeader) + } + + // 同步 X-Claude-Code-Session-Id 头：取 body 中已处理的 metadata.user_id 的 session_id 覆盖 + if sessionHeader := getHeaderRaw(req.Header, "X-Claude-Code-Session-Id"); sessionHeader != "" { + if uid := gjson.GetBytes(body, "metadata.user_id").String(); uid != "" { + if parsed := ParseMetadataUserID(uid); parsed != nil { + setHeaderRaw(req.Header, "X-Claude-Code-Session-Id", parsed.SessionID) + } + } + } + + // 账号级请求头覆写（仅 anthropic/openai api_key 账号启用时生效；OAuth 路径 no-op）。 + // 放在所有 header 逻辑之后，确保配置值对同名头拥有最终决定权。 + account.ApplyHeaderOverrides(req.Header) + + // === DEBUG: 打印上游转发请求（headers + body 摘要），与 CLIENT_ORIGINAL 对比 === + s.debugLogGatewaySnapshot("UPSTREAM_FORWARD", req.Header, body, map[string]string{ + "url": req.URL.String(), + "token_type": tokenType, + "mimic_claude_code": strconv.FormatBool(mimicClaudeCode), + "fingerprint_applied": strconv.FormatBool(fingerprint != nil), + "enable_fp": strconv.FormatBool(enableFP), + "enable_mpt": strconv.FormatBool(enableMPT), + }) + + // Always capture a compact fingerprint line for later error diagnostics. + // We only print it when needed (or when the explicit debug flag is enabled). + if c != nil && tokenType == "oauth" { + c.Set(claudeMimicDebugInfoKey, buildClaudeMimicDebugLine(req, body, account, tokenType, mimicClaudeCode)) + } + if s.debugClaudeMimicEnabled() { + logClaudeMimicDebug(req, body, account, tokenType, mimicClaudeCode) + } + + return req, body, nil +} + +// vertexSupportedBetaTokens 是 Vertex AI 的 Anthropic 端点接受的 anthropic-beta +// 白名单。Vertex 对任何未知 token 直接 HTTP 400，故采用白名单（与 Bedrock 的 +// bedrockSupportedBetaTokens 同思路）而非黑名单：未来 Claude Code 新增的、Vertex 尚未 +// 支持的 token 天然被剥离。当 Vertex 新增支持某 beta 时在此补充。 +// +// 明确排除（issue #3358 中 Vertex 报 400 的 token）：advisor-tool-2026-03-01、 +// prompt-caching-scope-2026-01-05、redact-thinking-2026-02-12、 +// thinking-token-count-2026-05-13；以及 claude-code-20250219 / oauth-2025-04-20 等 +// 客户端身份 beta——Vertex service_account 走 Bearer 鉴权，不需要它们。 +var vertexSupportedBetaTokens = map[string]bool{ + "context-1m-2025-08-07": true, + "context-management-2025-06-27": true, + "fine-grained-tool-streaming-2025-05-14": true, + "interleaved-thinking-2025-05-14": true, +} + +// filterVertexBetaTokens 解析 client 的 anthropic-beta header，先剔除 drop 集合中的 +// token（BetaPolicy filter + 默认 drop），再只保留 Vertex 支持的 token，去重后逗号拼接。 +// 返回最终 header（可能为空字符串）。 +func filterVertexBetaTokens(header string, drop map[string]struct{}) string { + tokens := parseAnthropicBetaHeader(header) + if len(tokens) == 0 { + return "" + } + out := make([]string, 0, len(tokens)) + seen := make(map[string]bool, len(tokens)) + for _, t := range tokens { + if _, dropped := drop[t]; dropped { + continue + } + if !vertexSupportedBetaTokens[t] { + continue + } + if seen[t] { + continue + } + seen[t] = true + out = append(out, t) + } + return strings.Join(out, ",") +} + +func (s *GatewayService) buildUpstreamRequestAnthropicVertex( + ctx context.Context, + c *gin.Context, + account *Account, + body []byte, + token string, + modelID string, + reqStream bool, +) (*http.Request, error) { + vertexBody, err := buildVertexAnthropicRequestBody(body) + if err != nil { + return nil, err + } + + // 计算最终 outgoing anthropic-beta。Vertex AI 的 Anthropic 端点只接受一小撮 + // beta token，未知 token 会直接 HTTP 400——近期 Claude Code CLI 透传的 + // advisor-tool-2026-03-01 / prompt-caching-scope-2026-01-05 / + // redact-thinking-2026-02-12 / thinking-token-count-2026-05-13 都不被 Vertex 接受 + // （issue #3358）。这里复用 BetaPolicy 的 block 检查（与 Bedrock 的 + // resolveBedrockBetaTokensForRequest 对称），再按 vertexSupportedBetaTokens 白名单 + // 剥离其余 token，使该路径与 Anthropic 直连 / Bedrock 路径行为一致。 + clientBeta := "" + if c != nil && c.Request != nil { + clientBeta = getHeaderRaw(c.Request.Header, "anthropic-beta") + } + policy := s.evaluateBetaPolicy(ctx, clientBeta, account, modelID) + if policy.blockErr != nil { + return nil, policy.blockErr + } + finalBeta := filterVertexBetaTokens(clientBeta, mergeDropSets(policy.filterSet)) + + // 能力维度 sanitize：基于最终 beta（而非原始 client 值）决定是否保留 body 中的 + // context_management，与 Anthropic 直连 / Bedrock 路径对称。 + if sanitized, changed := sanitizeAnthropicBodyForBetaTokens(vertexBody, finalBeta); changed { + vertexBody = sanitized + } + fullURL, err := buildVertexAnthropicURL(account.VertexProjectID(), account.VertexLocation(modelID), modelID, reqStream) + if err != nil { + return nil, err + } + req, err := http.NewRequestWithContext(ctx, http.MethodPost, fullURL, bytes.NewReader(vertexBody)) + if err != nil { + return nil, err + } + + if c != nil && c.Request != nil { + for key, values := range c.Request.Header { + lowerKey := strings.ToLower(strings.TrimSpace(key)) + if !allowedHeaders[lowerKey] || lowerKey == "anthropic-version" { + continue + } + wireKey := resolveWireCasing(key) + for _, v := range values { + addHeaderRaw(req.Header, wireKey, v) + } + } + } + + req.Header.Del("authorization") + req.Header.Del("x-api-key") + req.Header.Del("x-goog-api-key") + req.Header.Del("cookie") + req.Header.Del("anthropic-version") + setHeaderRaw(req.Header, "authorization", "Bearer "+token) + setHeaderRaw(req.Header, "content-type", "application/json") + + // 覆盖上面白名单 loop 写入的原始 client anthropic-beta，使用过滤后的最终值。 + // finalBeta 为空（全部被剥离）时不下发该 header，与 Vertex 无 beta 请求一致。 + deleteHeaderAllForms(req.Header, "anthropic-beta") + if finalBeta != "" { + setHeaderRaw(req.Header, "anthropic-beta", finalBeta) + } + + s.debugLogGatewaySnapshot("UPSTREAM_FORWARD_VERTEX_ANTHROPIC", req.Header, vertexBody, map[string]string{ + "url": req.URL.String(), + "token_type": "service_account", + "model": modelID, + "stream": strconv.FormatBool(reqStream), + }) + + return req, nil +} + +// getBetaHeader 处理anthropic-beta header +// 对于OAuth账号，需要确保包含oauth-2025-04-20 +func (s *GatewayService) getBetaHeader(modelID string, clientBetaHeader string) string { + // 如果客户端传了anthropic-beta + if clientBetaHeader != "" { + // 已包含oauth beta则直接返回 + if strings.Contains(clientBetaHeader, claude.BetaOAuth) { + return clientBetaHeader + } + + // 需要添加oauth beta + parts := strings.Split(clientBetaHeader, ",") + for i, p := range parts { + parts[i] = strings.TrimSpace(p) + } + + // 在claude-code-20250219后面插入oauth beta + claudeCodeIdx := -1 + for i, p := range parts { + if p == claude.BetaClaudeCode { + claudeCodeIdx = i + break + } + } + + if claudeCodeIdx >= 0 { + // 在claude-code后面插入 + newParts := make([]string, 0, len(parts)+1) + newParts = append(newParts, parts[:claudeCodeIdx+1]...) + newParts = append(newParts, claude.BetaOAuth) + newParts = append(newParts, parts[claudeCodeIdx+1:]...) + return strings.Join(newParts, ",") + } + + // 没有claude-code，放在第一位 + return claude.BetaOAuth + "," + clientBetaHeader + } + + // 客户端没传，根据模型生成 + // haiku 模型不需要 claude-code beta + if strings.Contains(strings.ToLower(modelID), "haiku") { + return claude.HaikuBetaHeader + } + + return claude.DefaultBetaHeader +} + +func requestNeedsBetaFeatures(body []byte) bool { + tools := gjson.GetBytes(body, "tools") + if tools.Exists() && tools.IsArray() && len(tools.Array()) > 0 { + return true + } + thinkingType := gjson.GetBytes(body, "thinking.type").String() + if strings.EqualFold(thinkingType, "enabled") || strings.EqualFold(thinkingType, "adaptive") { + return true + } + return false +} + +func defaultAPIKeyBetaHeader(body []byte) string { + modelID := gjson.GetBytes(body, "model").String() + if strings.Contains(strings.ToLower(modelID), "haiku") { + return claude.APIKeyHaikuBetaHeader + } + return claude.APIKeyBetaHeader +} + +func applyClaudeOAuthHeaderDefaults(req *http.Request) { + if req == nil { + return + } + if getHeaderRaw(req.Header, "Accept") == "" { + setHeaderRaw(req.Header, "Accept", "application/json") + } + for key, value := range claude.DefaultHeaders { + if value == "" { + continue + } + if getHeaderRaw(req.Header, key) == "" { + setHeaderRaw(req.Header, resolveWireCasing(key), value) + } + } +} + +func mergeAnthropicBeta(required []string, incoming string) string { + seen := make(map[string]struct{}, len(required)+8) + out := make([]string, 0, len(required)+8) + + add := func(v string) { + v = strings.TrimSpace(v) + if v == "" { + return + } + if _, ok := seen[v]; ok { + return + } + seen[v] = struct{}{} + out = append(out, v) + } + + for _, r := range required { + add(r) + } + for _, p := range strings.Split(incoming, ",") { + add(p) + } + return strings.Join(out, ",") +} + +func mergeAnthropicBetaDropping(required []string, incoming string, drop map[string]struct{}) string { + merged := mergeAnthropicBeta(required, incoming) + if merged == "" || len(drop) == 0 { + return merged + } + out := make([]string, 0, 8) + for _, p := range strings.Split(merged, ",") { + p = strings.TrimSpace(p) + if p == "" { + continue + } + if _, ok := drop[p]; ok { + continue + } + out = append(out, p) + } + return strings.Join(out, ",") +} + +// computeFinalAnthropicBeta 计算发往上游的最终 anthropic-beta header 值。 +// +// 设计动机：将原本在 buildUpstreamRequest 内联在一起、依赖 req.Header 的 +// anthropic-beta 计算逻辑抽成纯函数。这样调用方可以在 NewRequest 之前 +// 就提前拿到最终 beta header，进而能按它对 body 做能力维度 sanitize 后再做 +// CCH 签名——一举修复了以下之前由顺序依赖导致的能力维度 sanitize +// 无法部署的问题（签名与最终 body 不一致可以被判 third-party）。 +// +// 返回 (value, shouldSet)： +// - shouldSet=false 意为“不主动设置 anthropic-beta header”，与原代码“ +// API-key 账号 + 客户端未传 anthropic-beta + InjectBetaForAPIKey 未开启或 +// requestNeedsBetaFeatures=false”的行为对齐。 +// - shouldSet=true 时 value 可能为空字符串（例如客户端透传的 beta 被 dropSet +// 全部过滤掉），这与原代码中 setHeaderRaw 的结果一致。 +// +// clientHeaders 是客户端原始 HTTP header（通常为 c.Request.Header）；nil 时按“客户端 +// 未传”处理。body 是已经 metadata 重写 / billing version sync 之后但未 sanitize 上游 +// 不兼容字段之前的版本。 +func (s *GatewayService) computeFinalAnthropicBeta( + tokenType string, + mimicClaudeCode bool, + modelID string, + clientHeaders http.Header, + body []byte, + effectiveDropSet map[string]struct{}, +) (string, bool) { + clientBeta := "" + if clientHeaders != nil { + clientBeta = getHeaderRaw(clientHeaders, "anthropic-beta") + } + + if tokenType == "oauth" { + if mimicClaudeCode { + // mimic 路径：原代码跳过白名单透传，incomingBeta 总是空字符串。 + // 这里传空 string 以严格对齐原行为。 + requiredBetas := []string{claude.BetaOAuth, claude.BetaInterleavedThinking} + if !strings.Contains(strings.ToLower(modelID), "haiku") { + requiredBetas = claude.FullClaudeCodeMimicryBetas() + } + return mergeAnthropicBetaDropping(requiredBetas, "", effectiveDropSet), true + } + // 真 Claude Code 客户端透传路径 + return stripBetaTokensWithSet(s.getBetaHeader(modelID, clientBeta), effectiveDropSet), true + } + + // API-key accounts + if clientBeta != "" { + return stripBetaTokensWithSet(clientBeta, effectiveDropSet), true + } + if s.cfg != nil && s.cfg.Gateway.InjectBetaForAPIKey { + if requestNeedsBetaFeatures(body) { + if beta := defaultAPIKeyBetaHeader(body); beta != "" { + return beta, true + } + } + } + return "", false +} + +// computeFinalCountTokensAnthropicBeta 是 count_tokens 路径上 anthropic-beta header 的 +// 计算纯函数。语义与 computeFinalAnthropicBeta 对齐，但备份了 count_tokens 独有的 +// 两条特殊规则： +// +// - OAuth mimic：requiredBetas 为 FullClaudeCodeMimicryBetas + BetaTokenCounting +// （与 messages 不同的是：不按 haiku 排除；count_tokens 始终携带 token-counting beta） +// - OAuth 透传 + 客户端未传 anthropic-beta：补齐 CountTokensBetaHeader +// - OAuth 透传 + 客户端传了：补齐 BetaTokenCounting（如果未含） +// +// 返回语义同 computeFinalAnthropicBeta。 +func (s *GatewayService) computeFinalCountTokensAnthropicBeta( + tokenType string, + mimicClaudeCode bool, + modelID string, + clientHeaders http.Header, + body []byte, + effectiveDropSet map[string]struct{}, +) (string, bool) { + clientBeta := "" + if clientHeaders != nil { + clientBeta = getHeaderRaw(clientHeaders, "anthropic-beta") + } + + if tokenType == "oauth" { + if mimicClaudeCode { + // 与原代码严格等价：original buildCountTokensRequest 在 count_tokens mimic + // 分支上**不**会跳过白名单透传（与 messages mimic 路径不同），所以 + // incomingBeta = req.Header[anthropic-beta] = 客户端透传过来的 client beta。 + // 重构后直接从 clientHeaders 拿同一个值，保持行为一致。 + requiredBetas := append(claude.FullClaudeCodeMimicryBetas(), claude.BetaTokenCounting) + return mergeAnthropicBetaDropping(requiredBetas, clientBeta, effectiveDropSet), true + } + if clientBeta == "" { + return claude.CountTokensBetaHeader, true + } + beta := s.getBetaHeader(modelID, clientBeta) + if !strings.Contains(beta, claude.BetaTokenCounting) { + beta = beta + "," + claude.BetaTokenCounting + } + return stripBetaTokensWithSet(beta, effectiveDropSet), true + } + + // API-key accounts + if clientBeta != "" { + return stripBetaTokensWithSet(clientBeta, effectiveDropSet), true + } + if s.cfg != nil && s.cfg.Gateway.InjectBetaForAPIKey { + if requestNeedsBetaFeatures(body) { + if beta := defaultAPIKeyBetaHeader(body); beta != "" { + return beta, true + } + } + } + return "", false +} + +// stripBetaTokens removes the given beta tokens from a comma-separated header value. +func stripBetaTokens(header string, tokens []string) string { + if header == "" || len(tokens) == 0 { + return header + } + return stripBetaTokensWithSet(header, buildBetaTokenSet(tokens)) +} + +func stripBetaTokensWithSet(header string, drop map[string]struct{}) string { + if header == "" || len(drop) == 0 { + return header + } + parts := strings.Split(header, ",") + out := make([]string, 0, len(parts)) + for _, p := range parts { + p = strings.TrimSpace(p) + if p == "" { + continue + } + if _, ok := drop[p]; ok { + continue + } + out = append(out, p) + } + if len(out) == len(parts) { + return header // no change, avoid allocation + } + return strings.Join(out, ",") +} + +// BetaBlockedError indicates a request was blocked by a beta policy rule. +type BetaBlockedError struct { + Message string +} + +func (e *BetaBlockedError) Error() string { return e.Message } + +// betaPolicyResult holds the evaluated result of beta policy rules for a single request. +type betaPolicyResult struct { + blockErr *BetaBlockedError // non-nil if a block rule matched + filterSet map[string]struct{} // tokens to filter (may be nil) +} + +// evaluateBetaPolicy loads settings once and evaluates all rules against the given request. +func (s *GatewayService) evaluateBetaPolicy(ctx context.Context, betaHeader string, account *Account, model string) betaPolicyResult { + if s.settingService == nil { + return betaPolicyResult{} + } + settings, err := s.settingService.GetBetaPolicySettings(ctx) + if err != nil || settings == nil { + return betaPolicyResult{} + } + isOAuth := account.IsOAuth() + isBedrock := account.IsBedrock() + var result betaPolicyResult + for _, rule := range settings.Rules { + if !betaPolicyScopeMatches(rule.Scope, isOAuth, isBedrock) { + continue + } + effectiveAction, effectiveErrMsg := resolveRuleAction(rule, model) + switch effectiveAction { + case BetaPolicyActionBlock: + if result.blockErr == nil && betaHeader != "" && containsBetaToken(betaHeader, rule.BetaToken) { + msg := effectiveErrMsg + if msg == "" { + msg = "beta feature " + rule.BetaToken + " is not allowed" + } + result.blockErr = &BetaBlockedError{Message: msg} + } + case BetaPolicyActionFilter: + if result.filterSet == nil { + result.filterSet = make(map[string]struct{}) + } + result.filterSet[rule.BetaToken] = struct{}{} + } + } + return result +} + +// mergeDropSets merges the static defaultDroppedBetasSet with dynamic policy filter tokens. +// Returns defaultDroppedBetasSet directly when policySet is empty (zero allocation). +func mergeDropSets(policySet map[string]struct{}, extra ...string) map[string]struct{} { + if len(policySet) == 0 && len(extra) == 0 { + return defaultDroppedBetasSet + } + m := make(map[string]struct{}, len(defaultDroppedBetasSet)+len(policySet)+len(extra)) + for t := range defaultDroppedBetasSet { + m[t] = struct{}{} + } + for t := range policySet { + m[t] = struct{}{} + } + for _, t := range extra { + m[t] = struct{}{} + } + return m +} + +// betaPolicyFilterSetKey is the gin.Context key for caching the policy filter set within a request. +const betaPolicyFilterSetKey = "betaPolicyFilterSet" + +// getBetaPolicyFilterSet returns the beta policy filter set, using the gin context cache if available. +// In the /v1/messages path, Forward() evaluates the policy first and caches the result; +// buildUpstreamRequest reuses it (zero extra DB calls). In the count_tokens path, this +// evaluates on demand (one DB call). +func (s *GatewayService) getBetaPolicyFilterSet(ctx context.Context, c *gin.Context, account *Account, model string) map[string]struct{} { + if c != nil { + if v, ok := c.Get(betaPolicyFilterSetKey); ok { + if fs, ok := v.(map[string]struct{}); ok { + return fs + } + } + } + return s.evaluateBetaPolicy(ctx, "", account, model).filterSet +} + +// betaPolicyScopeMatches checks whether a rule's scope matches the current account type. +func betaPolicyScopeMatches(scope string, isOAuth bool, isBedrock bool) bool { + switch scope { + case BetaPolicyScopeAll: + return true + case BetaPolicyScopeOAuth: + return isOAuth + case BetaPolicyScopeAPIKey: + return !isOAuth && !isBedrock + case BetaPolicyScopeBedrock: + return isBedrock + default: + return true // unknown scope → match all (fail-open) + } +} + +// matchModelWhitelist checks if a model matches any pattern in the whitelist. +// Reuses matchModelPattern from group.go which supports exact and wildcard prefix matching. +func matchModelWhitelist(model string, whitelist []string) bool { + for _, pattern := range whitelist { + if matchModelPattern(pattern, model) { + return true + } + } + return false +} + +// resolveRuleAction determines the effective action and error message for a rule given the request model. +// When ModelWhitelist is empty, the rule's primary Action/ErrorMessage applies unconditionally. +// When non-empty, Action applies to matching models; FallbackAction/FallbackErrorMessage applies to others. +func resolveRuleAction(rule BetaPolicyRule, model string) (action, errorMessage string) { + if len(rule.ModelWhitelist) == 0 { + return rule.Action, rule.ErrorMessage + } + if matchModelWhitelist(model, rule.ModelWhitelist) { + return rule.Action, rule.ErrorMessage + } + if rule.FallbackAction != "" { + return rule.FallbackAction, rule.FallbackErrorMessage + } + return BetaPolicyActionPass, "" // default fallback: pass (fail-open) +} + +// droppedBetaSet returns claude.DroppedBetas as a set, with optional extra tokens. +func droppedBetaSet(extra ...string) map[string]struct{} { + m := make(map[string]struct{}, len(defaultDroppedBetasSet)+len(extra)) + for t := range defaultDroppedBetasSet { + m[t] = struct{}{} + } + for _, t := range extra { + m[t] = struct{}{} + } + return m +} + +// containsBetaToken checks if a comma-separated header value contains the given token. +func containsBetaToken(header, token string) bool { + if header == "" || token == "" { + return false + } + for _, p := range strings.Split(header, ",") { + if strings.TrimSpace(p) == token { + return true + } + } + return false +} + +func filterBetaTokens(tokens []string, filterSet map[string]struct{}) []string { + if len(tokens) == 0 || len(filterSet) == 0 { + return tokens + } + kept := make([]string, 0, len(tokens)) + for _, token := range tokens { + if _, filtered := filterSet[token]; !filtered { + kept = append(kept, token) + } + } + return kept +} + +func (s *GatewayService) resolveBedrockBetaTokensForRequest( + ctx context.Context, + account *Account, + betaHeader string, + body []byte, + modelID string, +) ([]string, error) { + // 1. 对原始 header 中的 beta token 做 block 检查（快速失败） + policy := s.evaluateBetaPolicy(ctx, betaHeader, account, modelID) + if policy.blockErr != nil { + return nil, policy.blockErr + } + + // 2. 解析 header + body 自动注入 + Bedrock 转换/过滤 + betaTokens := ResolveBedrockBetaTokens(betaHeader, body, modelID) + + // 3. 对最终 token 列表再做 block 检查，捕获通过 body 自动注入绕过 header block 的情况。 + // 例如：管理员 block 了 interleaved-thinking，客户端不在 header 中带该 token， + // 但请求体中包含 thinking 字段 → autoInjectBedrockBetaTokens 会自动补齐 → + // 如果不做此检查，block 规则会被绕过。 + if blockErr := s.checkBetaPolicyBlockForTokens(ctx, betaTokens, account, modelID); blockErr != nil { + return nil, blockErr + } + + return filterBetaTokens(betaTokens, policy.filterSet), nil +} + +// checkBetaPolicyBlockForTokens 检查 token 列表中是否有被管理员 block 规则命中的 token。 +// 用于补充 evaluateBetaPolicy 对 header 的检查，覆盖 body 自动注入的 token。 +func (s *GatewayService) checkBetaPolicyBlockForTokens(ctx context.Context, tokens []string, account *Account, model string) *BetaBlockedError { + if s.settingService == nil || len(tokens) == 0 { + return nil + } + settings, err := s.settingService.GetBetaPolicySettings(ctx) + if err != nil || settings == nil { + return nil + } + isOAuth := account.IsOAuth() + isBedrock := account.IsBedrock() + tokenSet := buildBetaTokenSet(tokens) + for _, rule := range settings.Rules { + effectiveAction, effectiveErrMsg := resolveRuleAction(rule, model) + if effectiveAction != BetaPolicyActionBlock { + continue + } + if !betaPolicyScopeMatches(rule.Scope, isOAuth, isBedrock) { + continue + } + if _, present := tokenSet[rule.BetaToken]; present { + msg := effectiveErrMsg + if msg == "" { + msg = "beta feature " + rule.BetaToken + " is not allowed" + } + return &BetaBlockedError{Message: msg} + } + } + return nil +} + +func buildBetaTokenSet(tokens []string) map[string]struct{} { + m := make(map[string]struct{}, len(tokens)) + for _, t := range tokens { + if t == "" { + continue + } + m[t] = struct{}{} + } + return m +} + +var defaultDroppedBetasSet = buildBetaTokenSet(claude.DroppedBetas) + +// applyClaudeCodeMimicHeaders forces "Claude Code-like" request headers. +// This mirrors opencode-anthropic-auth behavior: do not trust downstream +// headers when using Claude Code-scoped OAuth credentials. +func applyClaudeCodeMimicHeaders(req *http.Request, isStream bool) { + if req == nil { + return + } + // Start with the standard defaults (fill missing). + applyClaudeOAuthHeaderDefaults(req) + // Then force key headers to match Claude Code fingerprint regardless of what the client sent. + // 使用 resolveWireCasing 确保 key 与真实 wire format 一致（如 "x-app" 而非 "X-App"） + for key, value := range claude.DefaultHeaders { + if value == "" { + continue + } + setHeaderRaw(req.Header, resolveWireCasing(key), value) + } + // Real Claude CLI uses Accept: application/json (even for streaming). + setHeaderRaw(req.Header, "Accept", "application/json") + if isStream { + setHeaderRaw(req.Header, "x-stainless-helper-method", "stream") + } + // Real Claude CLI 每个请求都会生成一个新的 UUID 放在 x-client-request-id。 + // 上游会以此作为会话/请求指纹的一部分，缺失或重复都可能触发第三方判定。 + if getHeaderRaw(req.Header, "x-client-request-id") == "" { + setHeaderRaw(req.Header, "x-client-request-id", uuid.NewString()) + } +} + +func truncateForLog(b []byte, maxBytes int) string { + if maxBytes <= 0 { + maxBytes = 2048 + } + if len(b) > maxBytes { + b = b[:maxBytes] + } + s := string(b) + // 保持一行，避免污染日志格式 + s = strings.ReplaceAll(s, "\n", "\\n") + s = strings.ReplaceAll(s, "\r", "\\r") + return s +} + +// buildCustomRelayURL 构建自定义中继转发 URL +// 在 path 后附加 beta=true 和可选的 proxy 查询参数 +func (s *GatewayService) buildCustomRelayURL(baseURL, path string, account *Account) string { + u := strings.TrimRight(baseURL, "/") + path + "?beta=true" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL := account.Proxy.URL() + if proxyURL != "" { + u += "&proxy=" + url.QueryEscape(proxyURL) + } + } + return u +} + +func (s *GatewayService) validateUpstreamBaseURL(raw string) (string, error) { + if s.cfg != nil && !s.cfg.Security.URLAllowlist.Enabled { + normalized, err := urlvalidator.ValidateURLFormat(raw, s.cfg.Security.URLAllowlist.AllowInsecureHTTP) + if err != nil { + return "", fmt.Errorf("invalid base_url: %w", err) + } + return normalized, nil + } + normalized, err := urlvalidator.ValidateHTTPSURL(raw, urlvalidator.ValidationOptions{ + AllowedHosts: s.cfg.Security.URLAllowlist.UpstreamHosts, + RequireAllowlist: true, + AllowPrivate: s.cfg.Security.URLAllowlist.AllowPrivateHosts, + }) + if err != nil { + return "", fmt.Errorf("invalid base_url: %w", err) + } + return normalized, nil +} diff --git a/backend/internal/service/gateway_upstream_response.go b/backend/internal/service/gateway_upstream_response.go new file mode 100644 index 0000000000..7a66e12a9a --- /dev/null +++ b/backend/internal/service/gateway_upstream_response.go @@ -0,0 +1,1426 @@ +package service + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "strconv" + "strings" + "sync/atomic" + "syscall" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/util/responseheaders" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" + + "github.com/gin-gonic/gin" +) + +// isClaudeCodeClient 判断请求是否来自真正的 Claude Code 客户端。 +// 判定条件： +// 1. User-Agent 匹配 claude-cli/X.Y.Z（大小写不敏感） +// 2. metadata.user_id 符合 Claude Code 格式（legacy 或 JSON 格式） +// +// 只检查 metadata.user_id 非空不够严格：第三方工具（opencode 等）可能伪造 UA +// 并附带任意 metadata.user_id 字符串，从而绕过 mimicry。必须通过 ParseMetadataUserID +// 验证格式才能确认是真正的 Claude Code 客户端。 +func isClaudeCodeClient(userAgent string, metadataUserID string) bool { + if !claudeCliUserAgentRe.MatchString(userAgent) { + return false + } + return ParseMetadataUserID(metadataUserID) != nil +} + +func shouldUseClaudeCodeNoopDeltaKeepalive(userAgent string) bool { + version := ExtractCLIVersion(userAgent) + if version == "" { + return false + } + return CompareVersions(version, claudeCodeNoopDeltaKeepaliveMinVersion) >= 0 +} + +func claudeCodeKeepaliveDeltaTypeForContentBlock(blockType string) string { + switch blockType { + case "text": + return "text_delta" + case "tool_use": + return "input_json_delta" + case "thinking": + return "thinking_delta" + default: + return "" + } +} + +func claudeCodeKeepaliveFieldForDeltaType(deltaType string) string { + switch deltaType { + case "text_delta": + return "text" + case "input_json_delta": + return "partial_json" + case "thinking_delta": + return "thinking" + default: + return "" + } +} + +func buildClaudeCodeNoopDeltaKeepalive(index int, deltaType string) (string, bool) { + fieldName := claudeCodeKeepaliveFieldForDeltaType(deltaType) + if fieldName == "" { + return "", false + } + return fmt.Sprintf("event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"index\":%d,\"delta\":{\"type\":\"%s\",\"%s\":\"\"}}\n\n", index, deltaType, fieldName), true +} + +func sseEventIndex(event map[string]any) (int, bool) { + switch v := event["index"].(type) { + case float64: + return int(v), true + case int: + return v, true + case int64: + return int(v), true + case json.Number: + i, err := v.Int64() + if err != nil { + return 0, false + } + return int(i), true + default: + return 0, false + } +} + +// shouldRectifySignatureError 统一判断是否应触发签名整流（strip thinking blocks 并重试）。 +// 根据账号类型检查对应的开关和匹配模式。 +// +// mappedModel 用于按 thinking 协议族分流：passback-required (DeepSeek/Kimi/GLM 等) 上游 +// 的 400 不是签名缺失问题，retry 任何 thinking 变形都会破坏「原样回传」契约——直接透传 +// 错误给客户端。详见 thinking_protocol.go。 +func (s *GatewayService) shouldRectifySignatureError(ctx context.Context, account *Account, respBody []byte, mappedModel string) bool { + if !ShouldRectifyThinkingSignatureError(mappedModel) { + return false + } + if account.Type == AccountTypeAPIKey { + // API Key 账号：独立开关，一次读取配置 + settings, err := s.settingService.GetRectifierSettings(ctx) + if err != nil || !settings.Enabled || !settings.APIKeySignatureEnabled { + return false + } + // 先检查内置模式（同 OAuth），再检查自定义关键词 + if s.isThinkingBlockSignatureError(respBody) { + return true + } + return matchSignaturePatterns(respBody, settings.APIKeySignaturePatterns) + } + // OAuth/SetupToken/Upstream/Bedrock 等：保持原有行为（内置模式 + 原开关） + return s.isThinkingBlockSignatureError(respBody) && s.settingService.IsSignatureRectifierEnabled(ctx) +} + +// isSignatureErrorPattern 仅做模式匹配，不检查开关。 +// 用于已进入重试流程后的二阶段检测（此时开关已在首次调用时验证过）。 +func (s *GatewayService) isSignatureErrorPattern(ctx context.Context, account *Account, respBody []byte) bool { + if s.isThinkingBlockSignatureError(respBody) { + return true + } + if account.Type == AccountTypeAPIKey { + settings, err := s.settingService.GetRectifierSettings(ctx) + if err != nil { + return false + } + return matchSignaturePatterns(respBody, settings.APIKeySignaturePatterns) + } + return false +} + +// matchSignaturePatterns 检查响应体是否匹配自定义关键词列表（不区分大小写）。 +func matchSignaturePatterns(respBody []byte, patterns []string) bool { + if len(patterns) == 0 { + return false + } + bodyLower := strings.ToLower(string(respBody)) + for _, p := range patterns { + p = strings.TrimSpace(p) + if p == "" { + continue + } + if strings.Contains(bodyLower, strings.ToLower(p)) { + return true + } + } + return false +} + +// isThinkingBlockSignatureError 检测是否是thinking block相关错误 +// 这类错误可以通过过滤thinking blocks并重试来解决 +func (s *GatewayService) isThinkingBlockSignatureError(respBody []byte) bool { + msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) + if msg == "" { + return false + } + + // 检测signature相关的错误（更宽松的匹配） + // 例如: "Invalid `signature` in `thinking` block", "***.signature" 等 + if strings.Contains(msg, "signature") { + return true + } + + // 检测 thinking block 顺序/类型错误 + // 例如: "Expected `thinking` or `redacted_thinking`, but found `text`" + if strings.Contains(msg, "expected") && (strings.Contains(msg, "thinking") || strings.Contains(msg, "redacted_thinking")) { + logger.LegacyPrintf("service.gateway", "[SignatureCheck] Detected thinking block type error") + return true + } + + // 检测 thinking block 被修改的错误 + // 例如: "thinking or redacted_thinking blocks in the latest assistant message cannot be modified" + if strings.Contains(msg, "cannot be modified") && (strings.Contains(msg, "thinking") || strings.Contains(msg, "redacted_thinking")) { + logger.LegacyPrintf("service.gateway", "[SignatureCheck] Detected thinking block modification error") + return true + } + + // 检测空消息内容错误（可能是过滤 thinking blocks 后导致的，或客户端发送了空 text block） + // 例如: "all messages must have non-empty content" + // "messages: text content blocks must be non-empty" + if strings.Contains(msg, "non-empty content") || strings.Contains(msg, "empty content") || + strings.Contains(msg, "content blocks must be non-empty") { + logger.LegacyPrintf("service.gateway", "[SignatureCheck] Detected empty content error") + return true + } + + // 检测 thinking block 缺少 thinking 字段的错误（跨模型切换时常见： + // 其他模型回过的 assistant 历史里有 type=thinking 但没有 thinking 文本， + // 喂给开启 extended thinking 的 claude 时会被拒） + // 例如: "messages.1.content.0.thinking: each thinking block must contain thinking" + if strings.Contains(msg, "thinking block must contain") { + logger.LegacyPrintf("service.gateway", "[SignatureCheck] Detected thinking block missing content error") + return true + } + + return false +} + +func (s *GatewayService) shouldFailoverOn400(respBody []byte) bool { + // 只对"可能是兼容性差异导致"的 400 允许切换，避免无意义重试。 + // 默认保守：无法识别则不切换。 + msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) + if msg == "" { + return false + } + + // 缺少/错误的 beta header：换账号/链路可能成功（尤其是混合调度时）。 + // 更精确匹配 beta 相关的兼容性问题，避免误触发切换。 + if strings.Contains(msg, "anthropic-beta") || + strings.Contains(msg, "beta feature") || + strings.Contains(msg, "requires beta") { + return true + } + + // thinking/tool streaming 等兼容性约束（常见于中间转换链路） + if strings.Contains(msg, "thinking") || strings.Contains(msg, "thought_signature") || strings.Contains(msg, "signature") { + return true + } + if strings.Contains(msg, "tool_use") || strings.Contains(msg, "tool_result") || strings.Contains(msg, "tools") { + return true + } + + return false +} + +// sanitizeStreamError 返回不含网络地址的客户端可见错误描述。 +// 默认 (*net.OpError).Error() 会拼接 Source/Addr 字段，泄露内部 IP/端口与上游 +// 服务器地址（例如 "read tcp 10.0.0.1:54321->52.1.2.3:443: read: connection +// reset by peer"）。该函数只保留可识别的错误类别，原始 err 仍在调用点写入日志。 +func sanitizeStreamError(err error) string { + if err == nil { + return "" + } + switch { + case errors.Is(err, io.ErrUnexpectedEOF): + return "unexpected EOF" + case errors.Is(err, io.EOF): + return "EOF" + case errors.Is(err, context.Canceled): + return "canceled" + case errors.Is(err, context.DeadlineExceeded): + return "deadline exceeded" + case errors.Is(err, syscall.ECONNRESET): + return "connection reset by peer" + case errors.Is(err, syscall.ECONNABORTED): + return "connection aborted" + case errors.Is(err, syscall.ETIMEDOUT): + return "connection timed out" + case errors.Is(err, syscall.EPIPE): + return "broken pipe" + case errors.Is(err, syscall.ECONNREFUSED): + return "connection refused" + } + var netErr *net.OpError + if errors.As(err, &netErr) { + if netErr.Timeout() { + if netErr.Op != "" { + return netErr.Op + " timeout" + } + return "i/o timeout" + } + if netErr.Op != "" { + return netErr.Op + " network error" + } + } + return "upstream connection error" +} + +// ExtractUpstreamErrorMessage 从上游响应体中提取错误消息 +// 支持 Claude 风格的错误格式：{"type":"error","error":{"type":"...","message":"..."}} +func ExtractUpstreamErrorMessage(body []byte) string { + return extractUpstreamErrorMessage(body) +} + +func extractUpstreamErrorMessage(body []byte) string { + // Claude 风格：{"type":"error","error":{"type":"...","message":"..."}} + if m := gjson.GetBytes(body, "error.message").String(); strings.TrimSpace(m) != "" { + inner := strings.TrimSpace(m) + // 有些上游会把完整 JSON 作为字符串塞进 message + if strings.HasPrefix(inner, "{") { + if innerMsg := gjson.Get(inner, "error.message").String(); strings.TrimSpace(innerMsg) != "" { + return innerMsg + } + } + return m + } + + // ChatGPT 内部 API 风格：{"detail":"..."} + if d := gjson.GetBytes(body, "detail").String(); strings.TrimSpace(d) != "" { + return d + } + + // 兜底：尝试顶层 message + return gjson.GetBytes(body, "message").String() +} + +func extractUpstreamErrorCode(body []byte) string { + if code := strings.TrimSpace(gjson.GetBytes(body, "error.code").String()); code != "" { + return code + } + + inner := strings.TrimSpace(gjson.GetBytes(body, "error.message").String()) + if !strings.HasPrefix(inner, "{") { + return "" + } + + if code := strings.TrimSpace(gjson.Get(inner, "error.code").String()); code != "" { + return code + } + + if lastBrace := strings.LastIndex(inner, "}"); lastBrace >= 0 { + if code := strings.TrimSpace(gjson.Get(inner[:lastBrace+1], "error.code").String()); code != "" { + return code + } + } + + return "" +} + +func isCountTokensUnsupported404(statusCode int, body []byte) bool { + if statusCode != http.StatusNotFound { + return false + } + msg := strings.ToLower(strings.TrimSpace(extractUpstreamErrorMessage(body))) + if msg == "" { + return false + } + if strings.Contains(msg, "/v1/messages/count_tokens") { + return true + } + return strings.Contains(msg, "count_tokens") && strings.Contains(msg, "not found") +} + +func (s *GatewayService) readUpstreamErrorBody(resp *http.Response) ([]byte, error) { + if resp == nil || resp.Body == nil { + return nil, nil + } + limit := gatewayUpstreamErrorBodyReadLimit + if s != nil && s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody && s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > int(limit) { + limit = int64(s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes) + } + return io.ReadAll(io.LimitReader(resp.Body, limit)) +} + +func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, requestedModel ...string) (*ForwardResult, error) { + body, _ := s.readUpstreamErrorBody(resp) + + // 调试日志：打印上游错误响应 + logger.LegacyPrintf("service.gateway", "[Forward] Upstream error (non-retryable): Account=%d(%s) Status=%d RequestID=%s Body=%s", + account.ID, account.Name, resp.StatusCode, resp.Header.Get("x-request-id"), truncateString(string(body), 1000)) + + upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(body)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + + // Print a compact upstream request fingerprint when we hit the Claude Code OAuth + // credential scope error. This avoids requiring env-var tweaks in a fixed deploy. + if isClaudeCodeCredentialScopeError(upstreamMsg) && c != nil { + if v, ok := c.Get(claudeMimicDebugInfoKey); ok { + if line, ok := v.(string); ok && strings.TrimSpace(line) != "" { + logger.LegacyPrintf("service.gateway", "[ClaudeMimicDebugOnError] status=%d request_id=%s %s", + resp.StatusCode, + resp.Header.Get("x-request-id"), + line, + ) + } + } + } + + // Enrich Ops error logs with upstream status + message, and optionally a truncated body snippet. + upstreamDetail := "" + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes + if maxBytes <= 0 { + maxBytes = 2048 + } + upstreamDetail = truncateString(string(body), maxBytes) + } + setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "http_error", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + // 处理上游错误，标记账号状态 + shouldDisable := false + if s.rateLimitService != nil { + if len(requestedModel) > 0 { + shouldDisable = s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel[0]) + } else { + shouldDisable = s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body) + } + } + if shouldDisable { + return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: body} + } + + MarkResponseCommitted(c) + + // 记录上游错误响应体摘要便于排障（可选：由配置控制；不回显到客户端） + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + logger.LegacyPrintf("service.gateway", + "Upstream error %d (account=%d platform=%s type=%s): %s", + resp.StatusCode, + account.ID, + account.Platform, + account.Type, + truncateForLog(body, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), + ) + } + + // 非 failover 错误也支持错误透传规则匹配。 + if status, errType, errMsg, matched := applyErrorPassthroughRule( + c, + account.Platform, + resp.StatusCode, + body, + http.StatusBadGateway, + "upstream_error", + "Upstream request failed", + ); matched { + c.JSON(status, gin.H{ + "type": "error", + "error": gin.H{ + "type": errType, + "message": errMsg, + }, + }) + + summary := upstreamMsg + if summary == "" { + summary = errMsg + } + if summary == "" { + return nil, fmt.Errorf("upstream error: %d (passthrough rule matched)", resp.StatusCode) + } + return nil, fmt.Errorf("upstream error: %d (passthrough rule matched) message=%s", resp.StatusCode, summary) + } + + // 根据状态码返回适当的自定义错误响应（不透传上游详细信息） + var errType, errMsg string + var statusCode int + + switch resp.StatusCode { + case 400: + c.Data(http.StatusBadRequest, "application/json", body) + summary := upstreamMsg + if summary == "" { + summary = truncateForLog(body, 512) + } + if summary == "" { + return nil, fmt.Errorf("upstream error: %d", resp.StatusCode) + } + return nil, fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, summary) + case 401: + statusCode = http.StatusBadGateway + errType = "upstream_error" + errMsg = "Upstream authentication failed, please contact administrator" + case 403: + statusCode = http.StatusBadGateway + errType = "upstream_error" + errMsg = "Upstream access forbidden, please contact administrator" + case 429: + statusCode = http.StatusTooManyRequests + errType = "rate_limit_error" + errMsg = "Upstream rate limit exceeded, please retry later" + case 529: + statusCode = http.StatusServiceUnavailable + errType = "overloaded_error" + errMsg = "Upstream service overloaded, please retry later" + case 500, 502, 503, 504: + statusCode = http.StatusBadGateway + errType = "upstream_error" + errMsg = "Upstream service temporarily unavailable" + 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 nil, fmt.Errorf("upstream error: %d", resp.StatusCode) + } + return nil, fmt.Errorf("upstream error: %d message=%s", resp.StatusCode, upstreamMsg) +} + +func (s *GatewayService) handleRetryExhaustedSideEffects(ctx context.Context, resp *http.Response, account *Account) { + body, _ := s.readUpstreamErrorBody(resp) + statusCode := resp.StatusCode + + // OAuth/Setup Token 账号的 403：标记账号异常 + if account.IsOAuth() && statusCode == 403 { + s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, resp.Header, body) + logger.LegacyPrintf("service.gateway", "Account %d: marked as error after %d retries for status %d", account.ID, maxRetryAttempts, statusCode) + } else { + // API Key 未配置错误码：不标记账号状态 + logger.LegacyPrintf("service.gateway", "Account %d: upstream error %d after %d retries (not marking account)", account.ID, statusCode, maxRetryAttempts) + } +} + +func (s *GatewayService) handleFailoverSideEffects(ctx context.Context, resp *http.Response, account *Account, requestedModel ...string) { + body, _ := s.readUpstreamErrorBody(resp) + if len(requestedModel) > 0 { + s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body, requestedModel[0]) + return + } + s.rateLimitService.HandleUpstreamError(ctx, account, resp.StatusCode, resp.Header, body) +} + +// handleRetryExhaustedError 处理重试耗尽后的错误 +// OAuth 403：标记账号异常 +// API Key 未配置错误码：仅返回错误，不标记账号 +func (s *GatewayService) handleRetryExhaustedError(ctx context.Context, resp *http.Response, c *gin.Context, account *Account) (*ForwardResult, error) { + MarkResponseCommitted(c) + // Capture upstream error body before side-effects consume the stream. + respBody, _ := s.readUpstreamErrorBody(resp) + _ = resp.Body.Close() + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + + s.handleRetryExhaustedSideEffects(ctx, resp, account) + + upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(respBody)) + upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg) + + if isClaudeCodeCredentialScopeError(upstreamMsg) && c != nil { + if v, ok := c.Get(claudeMimicDebugInfoKey); ok { + if line, ok := v.(string); ok && strings.TrimSpace(line) != "" { + logger.LegacyPrintf("service.gateway", "[ClaudeMimicDebugOnError] status=%d request_id=%s %s", + resp.StatusCode, + resp.Header.Get("x-request-id"), + line, + ) + } + } + } + + upstreamDetail := "" + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes + if maxBytes <= 0 { + maxBytes = 2048 + } + upstreamDetail = truncateString(string(respBody), maxBytes) + } + setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail) + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: resp.Header.Get("x-request-id"), + Kind: "retry_exhausted", + Message: upstreamMsg, + Detail: upstreamDetail, + }) + + if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody { + logger.LegacyPrintf("service.gateway", + "Upstream error %d retries_exhausted (account=%d platform=%s type=%s): %s", + resp.StatusCode, + account.ID, + account.Platform, + account.Type, + truncateForLog(respBody, s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes), + ) + } + + if status, errType, errMsg, matched := applyErrorPassthroughRule( + c, + account.Platform, + resp.StatusCode, + respBody, + http.StatusBadGateway, + "upstream_error", + "Upstream request failed after retries", + ); matched { + c.JSON(status, gin.H{ + "type": "error", + "error": gin.H{ + "type": errType, + "message": errMsg, + }, + }) + + summary := upstreamMsg + if summary == "" { + summary = errMsg + } + if summary == "" { + return nil, fmt.Errorf("upstream error: %d (retries exhausted, passthrough rule matched)", resp.StatusCode) + } + return nil, fmt.Errorf("upstream error: %d (retries exhausted, passthrough rule matched) message=%s", resp.StatusCode, summary) + } + + // 返回统一的重试耗尽错误响应 + c.JSON(http.StatusBadGateway, gin.H{ + "type": "error", + "error": gin.H{ + "type": "upstream_error", + "message": "Upstream request failed after retries", + }, + }) + + if upstreamMsg == "" { + return nil, fmt.Errorf("upstream error: %d (retries exhausted)", resp.StatusCode) + } + return nil, fmt.Errorf("upstream error: %d (retries exhausted) message=%s", resp.StatusCode, upstreamMsg) +} + +// streamingResult 流式响应结果 +type streamingResult struct { + usage *ClaudeUsage + firstTokenMs *int + clientDisconnect bool // 客户端是否在流式传输过程中断开 +} + +func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel string, mimicClaudeCode bool) (*streamingResult, error) { + // 更新5h窗口状态 + s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header) + + if s.responseHeaderFilter != nil { + responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) + } + + // 设置SSE响应头 + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Header("X-Accel-Buffering", "no") + + // 透传其他响应头 + if v := resp.Header.Get("x-request-id"); v != "" { + c.Header("x-request-id", v) + } + + w := c.Writer + flusher, ok := w.(http.Flusher) + if !ok { + return nil, errors.New("streaming not supported") + } + + usage := &ClaudeUsage{} + var firstTokenMs *int + scanner := bufio.NewScanner(resp.Body) + // 设置更大的buffer以处理长行 + maxLineSize := defaultMaxLineSize + if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 { + maxLineSize = s.cfg.Gateway.MaxLineSize + } + scanBuf := getSSEScannerBuf64K() + scanner.Buffer(scanBuf[:0], maxLineSize) + + type scanEvent struct { + line string + err error + } + // 独立 goroutine 读取上游，避免读取阻塞导致超时/keepalive无法处理 + 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.cfg != nil && s.cfg.Gateway.StreamDataIntervalTimeout > 0 { + streamInterval = time.Duration(s.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.cfg != nil && s.cfg.Gateway.StreamKeepaliveInterval > 0 { + keepaliveInterval = time.Duration(s.cfg.Gateway.StreamKeepaliveInterval) * time.Second + } + var keepaliveTimer *time.Timer + if keepaliveInterval > 0 { + keepaliveTimer = time.NewTimer(keepaliveInterval) + defer keepaliveTimer.Stop() + } + var keepaliveCh <-chan time.Time + if keepaliveTimer != nil { + keepaliveCh = keepaliveTimer.C + } + lastDataAt := time.Now() + resetKeepaliveTimer := func() { + if keepaliveTimer == nil { + return + } + if !keepaliveTimer.Stop() { + select { + case <-keepaliveTimer.C: + default: + } + } + keepaliveTimer.Reset(keepaliveInterval) + } + + // 仅发送一次错误事件，避免多次写入导致协议混乱（写失败时尽力通知客户端）。 + // 事件格式遵循 Anthropic SSE 标准：{"type":"error","error":{"type":,"message":}} + // 这样 Anthropic SDK / Claude Code 等客户端能按标准 error 类型解析，UI 能显示具体错误文案， + // 服务端 ExtractUpstreamErrorMessage 也能从透传的 body 中提取 message。 + errorEventSent := false + sendErrorEvent := func(reason, message string) { + if errorEventSent { + return + } + errorEventSent = true + if message == "" { + message = reason + } + body, err := json.Marshal(map[string]any{ + "type": "error", + "error": map[string]string{ + "type": reason, + "message": message, + }, + }) + if err != nil { + // json.Marshal 不可能在已知 string-only 输入上失败，保守 fallback + body = []byte(fmt.Sprintf(`{"type":"error","error":{"type":%q,"message":%q}}`, reason, message)) + } + _, _ = fmt.Fprintf(w, "event: error\ndata: %s\n\n", body) + flusher.Flush() + } + + needModelReplace := originalModel != mappedModel + clientDisconnected := false // 客户端断开标志，断开后继续读取上游以获取完整usage + sawTerminalEvent := false + useNoopDeltaKeepalive := c != nil && c.Request != nil && shouldUseClaudeCodeNoopDeltaKeepalive(c.GetHeader("User-Agent")) + noopDeltaKeepaliveBlockIndex := -1 + noopDeltaKeepaliveDeltaType := "" + + pendingEventLines := make([]string, 0, 4) + + processSSEEvent := func(lines []string) ([]string, string, *sseUsagePatch, error) { + if len(lines) == 0 { + return nil, "", nil, nil + } + + eventName := "" + dataLine := "" + for _, line := range lines { + trimmed := strings.TrimSpace(line) + if strings.HasPrefix(trimmed, "event:") { + eventName = strings.TrimSpace(strings.TrimPrefix(trimmed, "event:")) + continue + } + if dataLine == "" && sseDataRe.MatchString(trimmed) { + dataLine = sseDataRe.ReplaceAllString(trimmed, "") + } + } + + if eventName == "error" { + return nil, dataLine, nil, &sseStreamErrorEventError{RawData: dataLine} + } + + if dataLine == "" { + return []string{strings.Join(lines, "\n") + "\n\n"}, "", nil, nil + } + + if dataLine == "[DONE]" { + sawTerminalEvent = true + block := "" + if eventName != "" { + block = "event: " + eventName + "\n" + } + block += "data: " + dataLine + "\n\n" + return []string{block}, dataLine, nil, nil + } + + var event map[string]any + if err := json.Unmarshal([]byte(dataLine), &event); err != nil { + // JSON 解析失败，直接透传原始数据 + block := "" + if eventName != "" { + block = "event: " + eventName + "\n" + } + block += "data: " + dataLine + "\n\n" + return []string{block}, dataLine, nil, nil + } + + eventType, _ := event["type"].(string) + if eventName == "" { + eventName = eventType + } + eventChanged := false + + if useNoopDeltaKeepalive { + switch eventType { + case "content_block_start": + if idx, ok := sseEventIndex(event); ok { + noopDeltaKeepaliveBlockIndex = -1 + noopDeltaKeepaliveDeltaType = "" + if contentBlock, ok := event["content_block"].(map[string]any); ok { + blockType, _ := contentBlock["type"].(string) + if deltaType := claudeCodeKeepaliveDeltaTypeForContentBlock(blockType); deltaType != "" { + noopDeltaKeepaliveBlockIndex = idx + noopDeltaKeepaliveDeltaType = deltaType + } + } + } + case "content_block_delta": + if idx, ok := sseEventIndex(event); ok { + if delta, ok := event["delta"].(map[string]any); ok { + deltaType, _ := delta["type"].(string) + if claudeCodeKeepaliveFieldForDeltaType(deltaType) != "" { + noopDeltaKeepaliveBlockIndex = idx + noopDeltaKeepaliveDeltaType = deltaType + } + } + } + case "content_block_stop": + if idx, ok := sseEventIndex(event); ok && idx == noopDeltaKeepaliveBlockIndex { + noopDeltaKeepaliveBlockIndex = -1 + noopDeltaKeepaliveDeltaType = "" + } + case "message_stop": + noopDeltaKeepaliveBlockIndex = -1 + noopDeltaKeepaliveDeltaType = "" + } + } + + // 兼容 Kimi cached_tokens → cache_read_input_tokens + if eventType == "message_start" { + if msg, ok := event["message"].(map[string]any); ok { + if u, ok := msg["usage"].(map[string]any); ok { + eventChanged = reconcileCachedTokens(u) || eventChanged + } + } + } + if eventType == "message_delta" { + if u, ok := event["usage"].(map[string]any); ok { + eventChanged = reconcileCachedTokens(u) || eventChanged + } + } + + // Cache TTL Override: 重写 SSE 事件中的 cache_creation 分类。 + // 账号级设置优先；全局 1h 请求注入开启时，默认把 usage 计费归回 5m。 + if overrideTarget, ok := s.resolveCacheTTLUsageOverrideTarget(ctx, account); ok { + if eventType == "message_start" { + if msg, ok := event["message"].(map[string]any); ok { + if u, ok := msg["usage"].(map[string]any); ok { + eventChanged = rewriteCacheCreationJSON(u, overrideTarget) || eventChanged + } + } + } + if eventType == "message_delta" { + if u, ok := event["usage"].(map[string]any); ok { + eventChanged = rewriteCacheCreationJSON(u, overrideTarget) || eventChanged + } + } + } + + if needModelReplace { + if msg, ok := event["message"].(map[string]any); ok { + if model, ok := msg["model"].(string); ok && model == mappedModel { + msg["model"] = originalModel + eventChanged = true + } + } + } + + usagePatch := s.extractSSEUsagePatch(event) + if anthropicStreamEventIsTerminal(eventName, dataLine) { + sawTerminalEvent = true + } + if !eventChanged { + block := "" + if eventName != "" { + block = "event: " + eventName + "\n" + } + block += "data: " + dataLine + "\n\n" + return []string{block}, dataLine, usagePatch, nil + } + + newData, err := json.Marshal(event) + if err != nil { + // 序列化失败，直接透传原始数据 + block := "" + if eventName != "" { + block = "event: " + eventName + "\n" + } + block += "data: " + dataLine + "\n\n" + return []string{block}, dataLine, usagePatch, nil + } + + block := "" + if eventName != "" { + block = "event: " + eventName + "\n" + } + block += "data: " + string(newData) + "\n\n" + return []string{block}, string(newData), usagePatch, nil + } + + for { + select { + case ev, ok := <-events: + if !ok { + // 上游完成，返回结果 + if !sawTerminalEvent { + return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, fmt.Errorf("stream usage incomplete: missing terminal event") + } + return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, nil + } + if ev.err != nil { + if sawTerminalEvent { + return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: clientDisconnected}, nil + } + // 检测 context 取消（客户端断开会导致 context 取消，进而影响上游读取） + if errors.Is(ev.err, context.Canceled) || errors.Is(ev.err, context.DeadlineExceeded) { + return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, fmt.Errorf("stream usage incomplete: %w", ev.err) + } + // 客户端已通过写入失败检测到断开，上游也出错了，返回已收集的 usage + if clientDisconnected { + return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, fmt.Errorf("stream usage incomplete after disconnect: %w", ev.err) + } + // 客户端未断开，正常的错误处理 + if errors.Is(ev.err, bufio.ErrTooLong) { + logger.LegacyPrintf("service.gateway", "SSE line too long: account=%d max_size=%d error=%v", account.ID, maxLineSize, ev.err) + sendErrorEvent("response_too_large", fmt.Sprintf("upstream SSE line exceeded %d bytes", maxLineSize)) + return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, ev.err + } + // 上游中途读错误（unexpected EOF / connection reset 等，常见于 HTTP/2 GOAWAY）： + // 若尚未向客户端写过任何字节，包成 UpstreamFailoverError 让 handler 层走 failover/重试。 + // 已经开始写流时 SSE 协议无 resume，只能透传错误事件给客户端。 + // 注意:面向客户端的 disconnectMsg 必须用 sanitizeStreamError 剥离地址, + // 默认 *net.OpError 的 Error() 会泄露内部 IP/端口和上游地址。完整 ev.err + // 仅在下方 LegacyPrintf 内部日志中保留供运维诊断。 + disconnectMsg := "upstream stream disconnected: " + sanitizeStreamError(ev.err) + if !c.Writer.Written() { + logger.LegacyPrintf("service.gateway", "Upstream stream read error before any client output (account=%d), failing over: %v", account.ID, ev.err) + body, _ := json.Marshal(map[string]any{ + "type": "error", + "error": map[string]string{ + "type": "upstream_disconnected", + "message": disconnectMsg, + }, + }) + return nil, &UpstreamFailoverError{ + StatusCode: http.StatusBadGateway, + ResponseBody: body, + RetryableOnSameAccount: true, + } + } + sendErrorEvent("stream_read_error", disconnectMsg) + return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream read error: %w", ev.err) + } + line := ev.line + trimmed := strings.TrimSpace(line) + + if trimmed == "" { + if len(pendingEventLines) == 0 { + continue + } + + outputBlocks, data, usagePatch, err := processSSEEvent(pendingEventLines) + pendingEventLines = pendingEventLines[:0] + if err != nil { + if clientDisconnected { + return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, nil + } + return nil, err + } + + for _, block := range outputBlocks { + if !clientDisconnected { + restored := reverseToolNamesIfPresent(c, []byte(block)) + if _, werr := fmt.Fprint(w, string(restored)); werr != nil { + clientDisconnected = true + logger.LegacyPrintf("service.gateway", "Client disconnected during streaming, continuing to drain upstream for billing") + break + } + flusher.Flush() + lastDataAt = time.Now() + resetKeepaliveTimer() + } + if data != "" { + if firstTokenMs == nil && data != "[DONE]" { + ms := int(time.Since(startTime).Milliseconds()) + firstTokenMs = &ms + } + if usagePatch != nil { + mergeSSEUsagePatch(usage, usagePatch) + } + } + } + continue + } + + pendingEventLines = append(pendingEventLines, line) + + case <-intervalCh: + lastRead := time.Unix(0, atomic.LoadInt64(&lastReadAt)) + if time.Since(lastRead) < streamInterval { + continue + } + if clientDisconnected { + return &streamingResult{usage: usage, firstTokenMs: firstTokenMs, clientDisconnect: true}, fmt.Errorf("stream usage incomplete after timeout") + } + logger.LegacyPrintf("service.gateway", "Stream data interval timeout: account=%d model=%s interval=%s", account.ID, originalModel, streamInterval) + // 处理流超时，可能标记账户为临时不可调度或错误状态 + if s.rateLimitService != nil { + s.rateLimitService.HandleStreamTimeout(ctx, account, originalModel) + } + sendErrorEvent("stream_timeout", fmt.Sprintf("upstream stream idle for %s", streamInterval)) + return &streamingResult{usage: usage, firstTokenMs: firstTokenMs}, fmt.Errorf("stream data interval timeout") + + case <-keepaliveCh: + if clientDisconnected { + continue + } + if time.Since(lastDataAt) < keepaliveInterval { + resetKeepaliveTimer() + continue + } + keepaliveBlock := "event: ping\ndata: {\"type\": \"ping\"}\n\n" + if useNoopDeltaKeepalive && noopDeltaKeepaliveBlockIndex >= 0 { + if block, ok := buildClaudeCodeNoopDeltaKeepalive(noopDeltaKeepaliveBlockIndex, noopDeltaKeepaliveDeltaType); ok { + keepaliveBlock = block + } + } + if _, werr := fmt.Fprint(w, keepaliveBlock); werr != nil { + clientDisconnected = true + logger.LegacyPrintf("service.gateway", "Client disconnected during keepalive ping, continuing to drain upstream for billing") + continue + } + flusher.Flush() + lastDataAt = time.Now() + resetKeepaliveTimer() + } + } + +} + +func (s *GatewayService) parseSSEUsage(data string, usage *ClaudeUsage) { + if usage == nil { + return + } + + var event map[string]any + if err := json.Unmarshal([]byte(data), &event); err != nil { + return + } + + if patch := s.extractSSEUsagePatch(event); patch != nil { + mergeSSEUsagePatch(usage, patch) + } +} + +type sseUsagePatch struct { + inputTokens int + hasInputTokens bool + outputTokens int + hasOutputTokens bool + cacheCreationInputTokens int + hasCacheCreationInput bool + cacheReadInputTokens int + hasCacheReadInput bool + cacheCreation5mTokens int + hasCacheCreation5m bool + cacheCreation1hTokens int + hasCacheCreation1h bool +} + +func (s *GatewayService) extractSSEUsagePatch(event map[string]any) *sseUsagePatch { + if len(event) == 0 { + return nil + } + + eventType, _ := event["type"].(string) + switch eventType { + case "message_start": + msg, _ := event["message"].(map[string]any) + usageObj, _ := msg["usage"].(map[string]any) + if len(usageObj) == 0 { + return nil + } + + patch := &sseUsagePatch{} + patch.hasInputTokens = true + if v, ok := parseSSEUsageInt(usageObj["input_tokens"]); ok { + patch.inputTokens = v + } + patch.hasCacheCreationInput = true + if v, ok := parseSSEUsageInt(usageObj["cache_creation_input_tokens"]); ok { + patch.cacheCreationInputTokens = v + } + patch.hasCacheReadInput = true + if v, ok := parseSSEUsageInt(usageObj["cache_read_input_tokens"]); ok { + patch.cacheReadInputTokens = v + } + if cc, ok := usageObj["cache_creation"].(map[string]any); ok { + if v, exists := parseSSEUsageInt(cc["ephemeral_5m_input_tokens"]); exists { + patch.cacheCreation5mTokens = v + patch.hasCacheCreation5m = true + } + if v, exists := parseSSEUsageInt(cc["ephemeral_1h_input_tokens"]); exists { + patch.cacheCreation1hTokens = v + patch.hasCacheCreation1h = true + } + } + return patch + + case "message_delta": + usageObj, _ := event["usage"].(map[string]any) + if len(usageObj) == 0 { + return nil + } + + patch := &sseUsagePatch{} + if v, ok := parseSSEUsageInt(usageObj["input_tokens"]); ok && v > 0 { + patch.inputTokens = v + patch.hasInputTokens = true + } + if v, ok := parseSSEUsageInt(usageObj["output_tokens"]); ok && v > 0 { + patch.outputTokens = v + patch.hasOutputTokens = true + } + if v, ok := parseSSEUsageInt(usageObj["cache_creation_input_tokens"]); ok && v > 0 { + patch.cacheCreationInputTokens = v + patch.hasCacheCreationInput = true + } + if v, ok := parseSSEUsageInt(usageObj["cache_read_input_tokens"]); ok && v > 0 { + patch.cacheReadInputTokens = v + patch.hasCacheReadInput = true + } + if cc, ok := usageObj["cache_creation"].(map[string]any); ok { + if v, exists := parseSSEUsageInt(cc["ephemeral_5m_input_tokens"]); exists && v > 0 { + patch.cacheCreation5mTokens = v + patch.hasCacheCreation5m = true + } + if v, exists := parseSSEUsageInt(cc["ephemeral_1h_input_tokens"]); exists && v > 0 { + patch.cacheCreation1hTokens = v + patch.hasCacheCreation1h = true + } + } + return patch + } + + return nil +} + +func mergeSSEUsagePatch(usage *ClaudeUsage, patch *sseUsagePatch) { + if usage == nil || patch == nil { + return + } + + if patch.hasInputTokens { + usage.InputTokens = patch.inputTokens + } + if patch.hasCacheCreationInput { + usage.CacheCreationInputTokens = patch.cacheCreationInputTokens + } + if patch.hasCacheReadInput { + usage.CacheReadInputTokens = patch.cacheReadInputTokens + } + if patch.hasOutputTokens { + usage.OutputTokens = patch.outputTokens + } + if patch.hasCacheCreation5m { + usage.CacheCreation5mTokens = patch.cacheCreation5mTokens + } + if patch.hasCacheCreation1h { + usage.CacheCreation1hTokens = patch.cacheCreation1hTokens + } +} + +func parseSSEUsageInt(value any) (int, bool) { + switch v := value.(type) { + case float64: + return int(v), true + case float32: + return int(v), true + case int: + return v, true + case int64: + return int(v), true + case int32: + return int(v), true + case json.Number: + if i, err := v.Int64(); err == nil { + return int(i), true + } + if f, err := v.Float64(); err == nil { + return int(f), true + } + case string: + if parsed, err := strconv.Atoi(strings.TrimSpace(v)); err == nil { + return parsed, true + } + } + return 0, false +} + +// applyCacheTTLOverride 将所有 cache creation tokens 归入指定的 TTL 类型。 +// target 为 "5m" 或 "1h"。返回 true 表示发生了变更。 +func applyCacheTTLOverride(usage *ClaudeUsage, target string) bool { + // Fallback: 如果只有聚合字段但无 5m/1h 明细，将聚合字段归入 5m 默认类别 + if usage.CacheCreation5mTokens == 0 && usage.CacheCreation1hTokens == 0 && usage.CacheCreationInputTokens > 0 { + usage.CacheCreation5mTokens = usage.CacheCreationInputTokens + } + + total := usage.CacheCreation5mTokens + usage.CacheCreation1hTokens + if total == 0 { + return false + } + switch target { + case "1h": + if usage.CacheCreation1hTokens == total { + return false // 已经全是 1h + } + usage.CacheCreation1hTokens = total + usage.CacheCreation5mTokens = 0 + default: // "5m" + if usage.CacheCreation5mTokens == total { + return false // 已经全是 5m + } + usage.CacheCreation5mTokens = total + usage.CacheCreation1hTokens = 0 + } + return true +} + +// rewriteCacheCreationJSON 在 JSON usage 对象中重写 cache_creation 嵌套对象的 TTL 分类。 +// usageObj 是 usage JSON 对象（map[string]any）。 +func rewriteCacheCreationJSON(usageObj map[string]any, target string) bool { + ccObj, ok := usageObj["cache_creation"].(map[string]any) + if !ok { + return false + } + v5m, _ := parseSSEUsageInt(ccObj["ephemeral_5m_input_tokens"]) + v1h, _ := parseSSEUsageInt(ccObj["ephemeral_1h_input_tokens"]) + total := v5m + v1h + if total == 0 { + return false + } + switch target { + case "1h": + if v1h == total { + return false + } + ccObj["ephemeral_1h_input_tokens"] = float64(total) + ccObj["ephemeral_5m_input_tokens"] = float64(0) + default: // "5m" + if v5m == total { + return false + } + ccObj["ephemeral_5m_input_tokens"] = float64(total) + ccObj["ephemeral_1h_input_tokens"] = float64(0) + } + return true +} + +func (s *GatewayService) resolveCacheTTLUsageOverrideTarget(ctx context.Context, account *Account) (string, bool) { + if account == nil { + return "", false + } + if account.IsCacheTTLOverrideEnabled() { + return account.GetCacheTTLOverrideTarget(), true + } + if account.IsAnthropicOAuthOrSetupToken() && s != nil && s.settingService != nil && s.settingService.IsAnthropicCacheTTL1hInjectionEnabled(ctx) { + return cacheTTLTarget5m, true + } + return "", false +} + +func (s *GatewayService) handleNonStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, originalModel, mappedModel string) (*ClaudeUsage, error) { + // 更新5h窗口状态 + s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header) + + body, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, anthropicTooLargeError) + if err != nil { + return nil, err + } + + // 解析usage + var response struct { + Usage ClaudeUsage `json:"usage"` + } + if err := json.Unmarshal(body, &response); err != nil { + if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices { + return nil, s.invalidNonStreamingJSONFailoverError(ctx, resp, account, body, err, mappedModel) + } + return nil, fmt.Errorf("parse response: %w", err) + } + + // 解析嵌套的 cache_creation 对象中的 5m/1h 明细 + cc5m := gjson.GetBytes(body, "usage.cache_creation.ephemeral_5m_input_tokens") + cc1h := gjson.GetBytes(body, "usage.cache_creation.ephemeral_1h_input_tokens") + if cc5m.Exists() || cc1h.Exists() { + response.Usage.CacheCreation5mTokens = int(cc5m.Int()) + response.Usage.CacheCreation1hTokens = int(cc1h.Int()) + } + + // 兼容 Kimi cached_tokens → cache_read_input_tokens + if response.Usage.CacheReadInputTokens == 0 { + cachedTokens := gjson.GetBytes(body, "usage.cached_tokens").Int() + if cachedTokens > 0 { + response.Usage.CacheReadInputTokens = int(cachedTokens) + if newBody, err := sjson.SetBytes(body, "usage.cache_read_input_tokens", cachedTokens); err == nil { + body = newBody + } + } + } + + // Cache TTL Override: 重写 non-streaming 响应中的 cache_creation 分类。 + // 账号级设置优先；全局 1h 请求注入开启时，默认把 usage 计费归回 5m。 + if overrideTarget, ok := s.resolveCacheTTLUsageOverrideTarget(ctx, account); ok { + if applyCacheTTLOverride(&response.Usage, overrideTarget) { + // 同步更新 body JSON 中的嵌套 cache_creation 对象 + if newBody, err := sjson.SetBytes(body, "usage.cache_creation.ephemeral_5m_input_tokens", response.Usage.CacheCreation5mTokens); err == nil { + body = newBody + } + if newBody, err := sjson.SetBytes(body, "usage.cache_creation.ephemeral_1h_input_tokens", response.Usage.CacheCreation1hTokens); err == nil { + body = newBody + } + } + } + + // 如果有模型映射，替换响应中的model字段 + if originalModel != mappedModel { + body = s.replaceModelInResponseBody(body, mappedModel, originalModel) + } + + responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) + + contentType := "application/json" + if s.cfg != nil && !s.cfg.Security.ResponseHeaders.Enabled { + if upstreamType := resp.Header.Get("Content-Type"); upstreamType != "" { + contentType = upstreamType + } + } + + body = reverseToolNamesIfPresent(c, body) + + // 写入响应 + c.Data(resp.StatusCode, contentType, body) + + return &response.Usage, nil +} + +// replaceModelInResponseBody 替换响应体中的model字段 +// 使用 gjson/sjson 精确替换，避免全量 JSON 反序列化 +func (s *GatewayService) replaceModelInResponseBody(body []byte, fromModel, toModel string) []byte { + if m := gjson.GetBytes(body, "model"); m.Exists() && m.Str == fromModel { + newBody, err := sjson.SetBytes(body, "model", toModel) + if err != nil { + return body + } + return newBody + } + return body +} + +// reconcileCachedTokens 兼容 Kimi 等上游： +// 将 OpenAI 风格的 cached_tokens 映射到 Claude 标准的 cache_read_input_tokens +func reconcileCachedTokens(usage map[string]any) bool { + if usage == nil { + return false + } + cacheRead, _ := usage["cache_read_input_tokens"].(float64) + if cacheRead > 0 { + return false // 已有标准字段，无需处理 + } + cached, _ := usage["cached_tokens"].(float64) + if cached <= 0 { + return false + } + usage["cache_read_input_tokens"] = cached + return true +} diff --git a/backend/internal/service/gateway_usage_billing.go b/backend/internal/service/gateway_usage_billing.go new file mode 100644 index 0000000000..21685caae8 --- /dev/null +++ b/backend/internal/service/gateway_usage_billing.go @@ -0,0 +1,978 @@ +package service + +import ( + "context" + "log/slog" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + "github.com/Wei-Shaw/sub2api/internal/pkg/timezone" +) + +func (s *GatewayService) getUserGroupRateMultiplier(ctx context.Context, userID, groupID int64, groupDefaultMultiplier float64) float64 { + if s == nil { + return groupDefaultMultiplier + } + resolver := s.userGroupRateResolver + if resolver == nil { + resolver = newUserGroupRateResolver( + s.userGroupRateRepo, + s.userGroupRateCache, + resolveUserGroupRateCacheTTL(s.cfg), + &s.userGroupRateSF, + "service.gateway", + ) + } + return resolver.Resolve(ctx, userID, groupID, groupDefaultMultiplier) +} + +// RecordUsageInput 记录使用量的输入参数。 +// 异步 worker 只接收计费所需快照，不能持有 ParsedRequest/RequestBodyRef 这类大请求体引用。 +type RecordUsageInput struct { + Result *ForwardResult + APIKey *APIKey + User *User + Account *Account + Subscription *UserSubscription // 可选：订阅信息 + InboundEndpoint string // 入站端点（客户端请求路径） + UpstreamEndpoint string // 上游端点（标准化后的上游路径） + UserAgent string // 请求的 User-Agent + IPAddress string // 请求的客户端 IP 地址 + RequestPayloadHash string // 请求体语义哈希，用于降低 request_id 误复用时的静默误去重风险 + ForceCacheBilling bool // 强制缓存计费：将 input_tokens 转为 cache_read 计费（用于粘性会话切换） + APIKeyService APIKeyQuotaUpdater // 可选：用于更新API Key配额 + QuotaPlatform string // user×platform 配额计量平台：handler 在请求 ctx 内经 QuotaPlatform() 算定后传入（后扣运行在 worker 池 background ctx 上，取不到 ForcePlatform） + + ChannelUsageFields // 渠道映射信息（由 handler 在 Forward 前解析） +} + +// APIKeyQuotaUpdater defines the interface for updating API Key quota and rate limit usage +type APIKeyQuotaUpdater interface { + UpdateQuotaUsed(ctx context.Context, apiKeyID int64, cost float64) error + UpdateRateLimitUsage(ctx context.Context, apiKeyID int64, cost float64) error +} + +type apiKeyAuthCacheInvalidator interface { + InvalidateAuthCacheByKey(ctx context.Context, key string) +} + +type usageLogBestEffortWriter interface { + CreateBestEffort(ctx context.Context, log *UsageLog) error +} + +// postUsageBillingParams 统一扣费所需的参数 +type postUsageBillingParams struct { + Cost *CostBreakdown + User *User + APIKey *APIKey + Account *Account + Subscription *UserSubscription + RequestPayloadHash string + IsSubscriptionBill bool + AccountRateMultiplier float64 + APIKeyService APIKeyQuotaUpdater + Platform string // 来自 APIKey 关联 Group 的平台标识 +} + +// PlatformFromAPIKey 从 APIKey 关联的 Group 推导 platform 名称。 +// apiKey 为 nil 或 Group 信息缺失时返回空串（调用方据此 short-circuit quota 累加）。 +// 导出供 handler 层调用。 +func PlatformFromAPIKey(apiKey *APIKey) string { + if apiKey == nil || apiKey.Group == nil { + return "" + } + return apiKey.Group.Platform +} + +// QuotaPlatform 返回 user×platform 配额计量使用的平台标识。 +// 强制平台路由（如 /antigravity）优先按 ctx 中的 ForcePlatform 计量，否则回退到 +// APIKey 关联 Group 的平台。 +// +// 注意：必须用带 ForcePlatform 的请求 context 调用（如 handler 的 c.Request.Context()）。 +// 后扣运行在 worker 池的 background ctx 上没有 ForcePlatform，因此后扣平台由 handler +// 预先算定、经 RecordUsageInput.QuotaPlatform 传入，不要在后扣链路用 worker ctx 调用本函数。 +func QuotaPlatform(ctx context.Context, apiKey *APIKey) string { + if fp, ok := ctx.Value(ctxkey.ForcePlatform).(string); ok && fp != "" { + return fp + } + return PlatformFromAPIKey(apiKey) +} + +func (p *postUsageBillingParams) shouldDeductAPIKeyQuota() bool { + return p.Cost.ActualCost > 0 && p.APIKey.Quota > 0 && p.APIKeyService != nil +} + +func (p *postUsageBillingParams) shouldUpdateRateLimits() bool { + return p.Cost.ActualCost > 0 && p.APIKey.HasRateLimits() && p.APIKeyService != nil +} + +func (p *postUsageBillingParams) shouldUpdateAccountQuota() bool { + return p.Cost.TotalCost > 0 && p.Account.IsAPIKeyOrBedrock() && p.Account.HasAnyQuotaLimit() +} + +// postUsageBilling is the legacy fallback billing path used when the unified +// billing repo is unavailable (nil). Production uses applyUsageBilling → repo.Apply +// for atomic billing. This path only runs in tests or degraded mode. +func postUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *billingDeps) { + billingCtx, cancel := detachedBillingContext(ctx) + defer cancel() + + cost := p.Cost + + if p.IsSubscriptionBill { + // Subscription usage tracked by ActualCost so group rate multiplier + // consumes the quota at the expected speed. + if cost.ActualCost > 0 { + if err := deps.userSubRepo.IncrementUsage(billingCtx, p.Subscription.ID, cost.ActualCost); err != nil { + slog.Error("increment subscription usage failed", "subscription_id", p.Subscription.ID, "error", err) + } + } + } else { + if cost.ActualCost > 0 { + if err := deps.userRepo.DeductBalance(billingCtx, p.User.ID, cost.ActualCost); err != nil { + slog.Error("deduct balance failed", "user_id", p.User.ID, "error", err) + } else if deps.billingCacheService != nil { + if err := deps.billingCacheService.InvalidateUserBalance(billingCtx, p.User.ID); err != nil { + slog.Warn("invalidate balance cache after legacy deduction failed", "user_id", p.User.ID, "error", err) + } + } + } + } + + if p.shouldDeductAPIKeyQuota() { + if err := p.APIKeyService.UpdateQuotaUsed(billingCtx, p.APIKey.ID, cost.ActualCost); err != nil { + slog.Error("update api key quota failed", "api_key_id", p.APIKey.ID, "error", err) + } + } + + if p.shouldUpdateRateLimits() { + if err := p.APIKeyService.UpdateRateLimitUsage(billingCtx, p.APIKey.ID, cost.ActualCost); err != nil { + slog.Error("update api key rate limit usage failed", "api_key_id", p.APIKey.ID, "error", err) + } + } + + if p.shouldUpdateAccountQuota() { + accountCost := cost.TotalCost * p.AccountRateMultiplier + if err := deps.accountRepo.IncrementQuotaUsed(billingCtx, p.Account.ID, accountCost); err != nil { + slog.Error("increment account quota used failed", "account_id", p.Account.ID, "cost", accountCost, "error", err) + } + } + + // Platform quota 累加（legacy 兜底路径）：仅对 standard（余额）模式生效；订阅模式豁免；仅对有 limit 的用户写 + // - HasUserPlatformQuotaLimit 守卫:与正常路径对齐，无 limit 公司跳过 + // - 新增 Redis 同步写:enforcement 走 Redis，legacy 路径也必须同步写，否则 preflight 看不到消费 + // - flusher_enabled=false（降级）:保留原有同步直写 DB + // - flusher_enabled=true:跳过直写 DB，由 flusher 异步批量刷（markDirty 在 IncrementUserPlatformQuotaUsage 内部完成） + // - 失败仅记 ALERT log + counter，不阻断主扣费流程 + if !p.IsSubscriptionBill && p.Platform != "" && cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil { + if deps.billingCacheService.HasUserPlatformQuotaLimit(billingCtx, p.User.ID, p.Platform) { + deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, cost.ActualCost) + if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled { + // 降级路径:flusher 未启用时保留原有同步直写 DB + if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(billingCtx, p.User.ID, p.Platform, cost.ActualCost, time.Now().UTC()); err != nil { + userPlatformQuotaDBIncrLegacyErrorTotal.Add(1) + logger.LegacyPrintf("service.gateway", "ALERT: legacy incr user platform quota DB failed user=%d platform=%s cost=%f: %v", p.User.ID, p.Platform, cost.ActualCost, err) + } + } + // flusher_enabled=true:不直写 DB，flusher 异步批量刷 + } + } + + // NOTE: finalizePostUsageBilling is NOT called here to avoid double-queuing + // cache updates. The legacy path does DB writes directly; the finalize path + // does cache queue + notifications. Notifications are dispatched separately + // by the caller after recording the usage log. +} + +func resolveUsageBillingRequestID(ctx context.Context, upstreamRequestID string) string { + if ctx != nil { + if clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(clientRequestID) != "" { + return "client:" + strings.TrimSpace(clientRequestID) + } + if requestID, _ := ctx.Value(ctxkey.RequestID).(string); strings.TrimSpace(requestID) != "" { + return "local:" + strings.TrimSpace(requestID) + } + } + if requestID := strings.TrimSpace(upstreamRequestID); requestID != "" { + return requestID + } + return "generated:" + generateRequestID() +} + +func resolveUsageBillingPayloadFingerprint(ctx context.Context, requestPayloadHash string) string { + if payloadHash := strings.TrimSpace(requestPayloadHash); payloadHash != "" { + return payloadHash + } + if ctx != nil { + if clientRequestID, _ := ctx.Value(ctxkey.ClientRequestID).(string); strings.TrimSpace(clientRequestID) != "" { + return "client:" + strings.TrimSpace(clientRequestID) + } + if requestID, _ := ctx.Value(ctxkey.RequestID).(string); strings.TrimSpace(requestID) != "" { + return "local:" + strings.TrimSpace(requestID) + } + } + return "" +} + +func buildUsageBillingCommand(requestID string, usageLog *UsageLog, p *postUsageBillingParams) *UsageBillingCommand { + if p == nil || p.Cost == nil || p.APIKey == nil || p.User == nil || p.Account == nil { + return nil + } + + cmd := &UsageBillingCommand{ + RequestID: requestID, + APIKeyID: p.APIKey.ID, + UserID: p.User.ID, + AccountID: p.Account.ID, + AccountType: p.Account.Type, + RequestPayloadHash: strings.TrimSpace(p.RequestPayloadHash), + } + if usageLog != nil { + cmd.Model = usageLog.Model + cmd.BillingType = usageLog.BillingType + cmd.InputTokens = usageLog.InputTokens + cmd.OutputTokens = usageLog.OutputTokens + cmd.CacheCreationTokens = usageLog.CacheCreationTokens + cmd.CacheReadTokens = usageLog.CacheReadTokens + cmd.ImageCount = usageLog.ImageCount + if usageLog.ServiceTier != nil { + cmd.ServiceTier = *usageLog.ServiceTier + } + if usageLog.ReasoningEffort != nil { + cmd.ReasoningEffort = *usageLog.ReasoningEffort + } + if usageLog.SubscriptionID != nil { + cmd.SubscriptionID = usageLog.SubscriptionID + } + } + + // Record subscription / balance cost using ActualCost so the group (and any + // user-specific) rate multiplier consumes subscription quota at the expected + // speed. TotalCost remains the raw (pre-multiplier) value; downstream guards + // on "> 0" still correctly skip free subscriptions (RateMultiplier == 0). + if p.IsSubscriptionBill && p.Subscription != nil && p.Cost.TotalCost > 0 { + cmd.SubscriptionID = &p.Subscription.ID + cmd.SubscriptionCost = p.Cost.ActualCost + } else if p.Cost.ActualCost > 0 { + cmd.BalanceCost = p.Cost.ActualCost + } + + if p.shouldDeductAPIKeyQuota() { + cmd.APIKeyQuotaCost = p.Cost.ActualCost + } + if p.shouldUpdateRateLimits() { + cmd.APIKeyRateLimitCost = p.Cost.ActualCost + } + if p.shouldUpdateAccountQuota() { + cmd.AccountQuotaCost = p.Cost.TotalCost * p.AccountRateMultiplier + } + + cmd.Normalize() + return cmd +} + +func applyUsageBilling(ctx context.Context, requestID string, usageLog *UsageLog, p *postUsageBillingParams, deps *billingDeps, repo UsageBillingRepository) (bool, error) { + if p == nil || deps == nil { + return false, nil + } + + cmd := buildUsageBillingCommand(requestID, usageLog, p) + if cmd == nil || cmd.RequestID == "" || repo == nil { + postUsageBilling(ctx, p, deps) + return true, nil + } + + billingCtx, cancel := detachedBillingContext(ctx) + defer cancel() + + result, err := repo.Apply(billingCtx, cmd) + if err != nil { + return false, err + } + + if result == nil || !result.Applied { + deps.deferredService.ScheduleLastUsedUpdate(p.Account.ID) + return false, nil + } + + if result.APIKeyQuotaExhausted { + if invalidator, ok := p.APIKeyService.(apiKeyAuthCacheInvalidator); ok && p.APIKey != nil && p.APIKey.Key != "" { + invalidator.InvalidateAuthCacheByKey(billingCtx, p.APIKey.Key) + } + } + + finalizePostUsageBilling(billingCtx, p, deps, result) + return true, nil +} + +func finalizePostUsageBilling(ctx context.Context, p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) { + if p == nil || p.Cost == nil || deps == nil { + return + } + + if p.IsSubscriptionBill { + if p.Cost.ActualCost > 0 && p.User != nil && p.APIKey != nil && p.APIKey.GroupID != nil { + deps.billingCacheService.QueueUpdateSubscriptionUsage(p.User.ID, *p.APIKey.GroupID, p.Cost.ActualCost) + } + } else if p.Cost.ActualCost > 0 && p.User != nil { + syncBalanceCacheAfterDeduction(ctx, p, deps, result) + } + + if p.Cost.ActualCost > 0 && p.APIKey != nil && p.APIKey.HasRateLimits() { + deps.billingCacheService.QueueUpdateAPIKeyRateLimitUsage(p.APIKey.ID, p.Cost.ActualCost) + } + + deps.deferredService.ScheduleLastUsedUpdate(p.Account.ID) + + // Platform quota 累加：仅在 standard（余额）模式生效；订阅模式豁免；仅对有 limit 的用户写 + // Redis 同步写 + DB 异步持久化（flag=false 降级）或 flusher 异步刷（flag=true）: + // - HasUserPlatformQuotaLimit 守卫:无 limit 的公司跳过,避免无效写入 + 浪费 Redis 容量 + // - Redis 同步:确保下次 preflight 立即看到最新 usage,把 TOCTOU 超支窗口 + // 限制在并发 in-flight 请求数量内（旧实现的异步入队会让超支无限累积直到 worker 处理） + // - DB 异步(flusher_enabled=false):在独立 goroutine 中走 detached context,失败用 ALERT log 触发 oncall 对账 + // - flusher_enabled=true:不直写 DB,由 flusher 异步批量刷（markDirty 已在 IncrementUserPlatformQuotaUsage 内部完成） + if !p.IsSubscriptionBill && p.Platform != "" && p.Cost.ActualCost > 0 && p.User != nil && deps.userPlatformQuotaRepo != nil { + if deps.billingCacheService.HasUserPlatformQuotaLimit(ctx, p.User.ID, p.Platform) { + deps.billingCacheService.IncrementUserPlatformQuotaUsage(p.User.ID, p.Platform, p.Cost.ActualCost) + if deps.cfg == nil || !deps.cfg.Database.UserPlatformQuotaFlusherEnabled { + // 降级路径:flusher 未启用时保留原有异步直写 DB + dbCtx, dbCancel := detachUpstreamContext(ctx) + userID, platform, cost := p.User.ID, p.Platform, p.Cost.ActualCost + go func() { + defer func() { + if r := recover(); r != nil { + logger.LegacyPrintf("service.gateway", "ALERT: panic in user platform quota incr goroutine user=%d platform=%s: %v", userID, platform, r) + } + }() + defer dbCancel() + if err := deps.userPlatformQuotaRepo.IncrementUsageWithReset(dbCtx, userID, platform, cost, time.Now().UTC()); err != nil { + // 失败计数器:暴露给 GatewayUserPlatformQuotaIncrStats(),由 ops 面板做斜率告警。 + userPlatformQuotaDBIncrErrorTotal.Add(1) + // ALERT 级别:DB 持久化失败意味着 Redis cache 失效后该笔 cost 永久丢失, + // 用户配额视图与实际消费会偏差,oncall 需要据此对账或人工补录。 + logger.LegacyPrintf("service.gateway", "ALERT: incr user platform quota DB failed user=%d platform=%s cost=%f: %v", userID, platform, cost, err) + } + }() + } + // flusher_enabled=true:不直写 DB,flusher 异步批量刷 + } + } + + // Notification checks run async — all parameters are already captured, + // no dependency on the request context or upstream connection. + go notifyBalanceLow(p, deps, result) + go notifyAccountQuota(p, deps, result) +} + +func syncBalanceCacheAfterDeduction(ctx context.Context, p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) { + if p == nil || p.Cost == nil || p.User == nil || deps == nil || deps.billingCacheService == nil { + return + } + if result != nil && result.NewBalance != nil && deps.billingCacheService.balanceBelowEligibilityThreshold(*result.NewBalance) { + if err := deps.billingCacheService.InvalidateUserBalance(ctx, p.User.ID); err != nil { + slog.Warn("invalidate balance cache after exhausted deduction failed", + "user_id", p.User.ID, + "new_balance", *result.NewBalance, + "balance_overdrafted", result.BalanceOverdrafted, + "error", err, + ) + } + return + } + deps.billingCacheService.QueueDeductBalance(p.User.ID, p.Cost.ActualCost) +} + +// notifyBalanceLow sends balance low notification after deduction. +// When result.NewBalance is available (from DB transaction RETURNING), it is used directly +// to reconstruct oldBalance, avoiding stale Redis reads and concurrent-deduction races. +func notifyBalanceLow(p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) { + defer func() { + if r := recover(); r != nil { + slog.Error("panic in notifyBalanceLow", "recover", r) + } + }() + if p.IsSubscriptionBill || p.Cost.ActualCost <= 0 || p.User == nil || deps.balanceNotifyService == nil { + slog.Debug("notifyBalanceLow: skipped", + "is_subscription", p.IsSubscriptionBill, + "actual_cost", p.Cost.ActualCost, + "user_nil", p.User == nil, + "service_nil", deps.balanceNotifyService == nil, + ) + return + } + + oldBalance := resolveOldBalance(p, result) + slog.Debug("notifyBalanceLow: calling CheckBalanceAfterDeduction", + "user_id", p.User.ID, + "old_balance", oldBalance, + "cost", p.Cost.ActualCost, + "notify_enabled", p.User.BalanceNotifyEnabled, + "threshold", p.User.BalanceNotifyThreshold, + "result_has_new_balance", result != nil && result.NewBalance != nil, + ) + deps.balanceNotifyService.CheckBalanceAfterDeduction(context.Background(), p.User, oldBalance, p.Cost.ActualCost) +} + +// resolveOldBalance returns the pre-deduction balance. +// Prefers the DB transaction result (newBalance + cost) over snapshot. +func resolveOldBalance(p *postUsageBillingParams, result *UsageBillingApplyResult) float64 { + if result != nil && result.NewBalance != nil { + return *result.NewBalance + p.Cost.ActualCost + } + // Legacy fallback: snapshot balance from request context + return p.User.Balance +} + +// notifyAccountQuota sends account quota threshold notification after increment. +// When result.QuotaState is available (from DB transaction RETURNING), it is passed directly +// to avoid a separate DB read that may see stale or concurrently-modified data. +func notifyAccountQuota(p *postUsageBillingParams, deps *billingDeps, result *UsageBillingApplyResult) { + defer func() { + if r := recover(); r != nil { + slog.Error("panic in notifyAccountQuota", "recover", r) + } + }() + if p.Cost.TotalCost <= 0 || p.Account == nil || !p.Account.IsAPIKeyOrBedrock() || deps.balanceNotifyService == nil { + slog.Debug("notifyAccountQuota: skipped", + "total_cost", p.Cost.TotalCost, + "account_nil", p.Account == nil, + "is_apikey_or_bedrock", p.Account != nil && p.Account.IsAPIKeyOrBedrock(), + "service_nil", deps.balanceNotifyService == nil, + ) + return + } + accountCost := p.Cost.TotalCost * p.AccountRateMultiplier + var quotaState *AccountQuotaState + if result != nil { + quotaState = result.QuotaState + } + slog.Debug("notifyAccountQuota: calling CheckAccountQuotaAfterIncrement", + "account_id", p.Account.ID, + "account_cost", accountCost, + "has_quota_state", quotaState != nil, + ) + deps.balanceNotifyService.CheckAccountQuotaAfterIncrement(context.Background(), p.Account, accountCost, quotaState) +} + +func detachedBillingContext(ctx context.Context) (context.Context, context.CancelFunc) { + base := context.Background() + if ctx != nil { + base = context.WithoutCancel(ctx) + } + return context.WithTimeout(base, postUsageBillingTimeout) +} + +func detachStreamUpstreamContext(ctx context.Context, stream bool) (context.Context, context.CancelFunc) { + if ctx == nil { + return context.Background(), func() {} + } + if !stream { + return ctx, func() {} + } + return context.WithoutCancel(ctx), func() {} +} + +func detachUpstreamContext(ctx context.Context) (context.Context, context.CancelFunc) { + if ctx == nil { + return context.Background(), func() {} + } + return context.WithoutCancel(ctx), func() {} +} + +// billingDeps 扣费逻辑依赖的服务（由各 gateway service 提供） +type billingDeps struct { + accountRepo AccountRepository + userRepo UserRepository + userSubRepo UserSubscriptionRepository + billingCacheService *BillingCacheService + deferredService *DeferredService + balanceNotifyService *BalanceNotifyService + userPlatformQuotaRepo UserPlatformQuotaRepository + cfg *config.Config +} + +func (s *GatewayService) billingDeps() *billingDeps { + return &billingDeps{ + accountRepo: s.accountRepo, + userRepo: s.userRepo, + userSubRepo: s.userSubRepo, + billingCacheService: s.billingCacheService, + deferredService: s.deferredService, + balanceNotifyService: s.balanceNotifyService, + userPlatformQuotaRepo: s.userPlatformQuotaRepo, + cfg: s.cfg, + } +} + +func writeUsageLogBestEffort(ctx context.Context, repo UsageLogRepository, usageLog *UsageLog, logKey string) { + if repo == nil || usageLog == nil { + return + } + usageCtx, cancel := detachedBillingContext(ctx) + defer cancel() + + if writer, ok := repo.(usageLogBestEffortWriter); ok { + if err := writer.CreateBestEffort(usageCtx, usageLog); err != nil { + logger.LegacyPrintf(logKey, "Create usage log failed: %v", err) + // 计费已在此前完成，日志必须落库：dropped（批处理队列超时）同样走同步兜底， + // 否则会出现“已扣费但无 usage_log”的对账缺口（issue #3656）。 + // 重复写入由 usage_logs 的 ON CONFLICT (request_id, api_key_id) DO NOTHING 防护。 + fallbackCtx := usageCtx + if usageCtx.Err() != nil { + // usageCtx 已耗尽（best-effort 入队阻塞到期限）：换新的 detached 窗口，避免兜底必然失败。 + var fallbackCancel context.CancelFunc + fallbackCtx, fallbackCancel = detachedBillingContext(context.Background()) + defer fallbackCancel() + } + if _, syncErr := repo.Create(fallbackCtx, usageLog); syncErr != nil { + logger.LegacyPrintf(logKey, "Create usage log sync fallback failed: %v", syncErr) + } + } + return + } + + if _, err := repo.Create(usageCtx, usageLog); err != nil { + logger.LegacyPrintf(logKey, "Create usage log failed: %v", err) + } +} + +// recordUsageOpts 内部选项，参数化普通计费与长上下文计费的差异点。 +type recordUsageOpts struct { + // 长上下文计费（仅 Gemini 路径需要） + LongContextThreshold int + LongContextMultiplier float64 +} + +// RecordUsage 记录使用量并扣费（或更新订阅用量） +func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInput) error { + return s.recordUsageCore(ctx, &recordUsageCoreInput{ + Result: input.Result, + APIKey: input.APIKey, + User: input.User, + Account: input.Account, + Subscription: input.Subscription, + InboundEndpoint: input.InboundEndpoint, + UpstreamEndpoint: input.UpstreamEndpoint, + UserAgent: input.UserAgent, + IPAddress: input.IPAddress, + RequestPayloadHash: input.RequestPayloadHash, + ForceCacheBilling: input.ForceCacheBilling, + APIKeyService: input.APIKeyService, + QuotaPlatform: input.QuotaPlatform, + ChannelUsageFields: input.ChannelUsageFields, + }, &recordUsageOpts{}) +} + +// RecordUsageLongContextInput 记录使用量的输入参数（支持长上下文双倍计费） +type RecordUsageLongContextInput struct { + Result *ForwardResult + APIKey *APIKey + User *User + Account *Account + Subscription *UserSubscription // 可选：订阅信息 + InboundEndpoint string // 入站端点（客户端请求路径） + UpstreamEndpoint string // 上游端点（标准化后的上游路径） + UserAgent string // 请求的 User-Agent + IPAddress string // 请求的客户端 IP 地址 + RequestPayloadHash string // 请求体语义哈希，用于降低 request_id 误复用时的静默误去重风险 + LongContextThreshold int // 长上下文阈值（如 200000） + LongContextMultiplier float64 // 超出阈值部分的倍率（如 2.0） + ForceCacheBilling bool // 强制缓存计费：将 input_tokens 转为 cache_read 计费（用于粘性会话切换） + APIKeyService APIKeyQuotaUpdater // API Key 配额服务（可选） + QuotaPlatform string // user×platform 配额计量平台：handler 在请求 ctx 内经 QuotaPlatform() 算定后传入（后扣运行在 worker 池 background ctx 上，取不到 ForcePlatform） + + ChannelUsageFields // 渠道映射信息（由 handler 在 Forward 前解析） +} + +// RecordUsageWithLongContext 记录使用量并扣费，支持长上下文双倍计费（用于 Gemini） +func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input *RecordUsageLongContextInput) error { + return s.recordUsageCore(ctx, &recordUsageCoreInput{ + Result: input.Result, + APIKey: input.APIKey, + User: input.User, + Account: input.Account, + Subscription: input.Subscription, + InboundEndpoint: input.InboundEndpoint, + UpstreamEndpoint: input.UpstreamEndpoint, + UserAgent: input.UserAgent, + IPAddress: input.IPAddress, + RequestPayloadHash: input.RequestPayloadHash, + ForceCacheBilling: input.ForceCacheBilling, + APIKeyService: input.APIKeyService, + QuotaPlatform: input.QuotaPlatform, + ChannelUsageFields: input.ChannelUsageFields, + }, &recordUsageOpts{ + LongContextThreshold: input.LongContextThreshold, + LongContextMultiplier: input.LongContextMultiplier, + }) +} + +// recordUsageCoreInput 是 recordUsageCore 的公共输入字段，从两种输入结构体中提取。 +type recordUsageCoreInput struct { + Result *ForwardResult + APIKey *APIKey + User *User + Account *Account + Subscription *UserSubscription + InboundEndpoint string + UpstreamEndpoint string + UserAgent string + IPAddress string + RequestPayloadHash string + ForceCacheBilling bool + APIKeyService APIKeyQuotaUpdater + QuotaPlatform string + ChannelUsageFields +} + +// recordUsageCore 是 RecordUsage 和 RecordUsageWithLongContext 的统一实现。 +// LongContextThreshold > 0 时 Token 计费回退走 CalculateCostWithLongContext。 +func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsageCoreInput, opts *recordUsageOpts) error { + result := input.Result + apiKey := input.APIKey + user := input.User + account := input.Account + subscription := input.Subscription + ApplyForwardImageBillingResolution(result) + + // 强制缓存计费：将 input_tokens 转为 cache_read_input_tokens + // 用于粘性会话切换时的特殊计费处理 + if input.ForceCacheBilling && result.Usage.InputTokens > 0 { + logger.LegacyPrintf("service.gateway", "force_cache_billing: %d input_tokens → cache_read_input_tokens (account=%d)", + result.Usage.InputTokens, account.ID) + result.Usage.CacheReadInputTokens += result.Usage.InputTokens + result.Usage.InputTokens = 0 + } + + // Cache TTL Override: 确保计费时 token 分类与账号设置一致。 + // 账号级设置优先；全局 1h 请求注入开启时，默认把 usage 计费归回 5m。 + cacheTTLOverridden := false + if overrideTarget, ok := s.resolveCacheTTLUsageOverrideTarget(ctx, account); ok { + applyCacheTTLOverride(&result.Usage, overrideTarget) + cacheTTLOverridden = (result.Usage.CacheCreation5mTokens + result.Usage.CacheCreation1hTokens) > 0 + } + + // 获取费率倍数（优先级：用户专属 > 分组默认 > 系统默认） + multiplier := 1.0 + if s.cfg != nil { + multiplier = s.cfg.Default.RateMultiplier + } + if apiKey.GroupID != nil && apiKey.Group != nil { + groupDefault := apiKey.Group.RateMultiplier + multiplier = s.getUserGroupRateMultiplier(ctx, user.ID, *apiKey.GroupID, groupDefault) + } + // token 倍率叠加高峰因子（token 计费含图片 token，图片按次倍率不受影响）。高峰因子按请求时刻现算， + // 不并入上面的 getUserGroupRateMultiplier，以免污染 user:group 倍率缓存。 + multiplier, imageMultiplier := computePeakAwareMultipliers(apiKey, multiplier, timezone.Now()) + + // 确定计费模型 + billingModel := forwardResultBillingModel(result.Model, result.UpstreamModel) + if input.BillingModelSource == BillingModelSourceChannelMapped && input.ChannelMappedModel != "" { + billingModel = input.ChannelMappedModel + } + if input.BillingModelSource == BillingModelSourceRequested && input.OriginalModel != "" { + billingModel = input.OriginalModel + } + + // 确定 RequestedModel（渠道映射前的原始模型） + requestedModel := result.Model + if input.OriginalModel != "" { + requestedModel = input.OriginalModel + } + + // 计算费用 + cost := s.calculateRecordUsageCost(ctx, result, apiKey, billingModel, multiplier, imageMultiplier, opts) + + // 判断计费方式：订阅模式 vs 余额模式 + isSubscriptionBilling := subscription != nil && apiKey.Group != nil && apiKey.Group.IsSubscriptionType() + billingType := BillingTypeBalance + if isSubscriptionBilling { + billingType = BillingTypeSubscription + } + + // 创建使用日志 + accountRateMultiplier := account.BillingRateMultiplier() + usageLog := s.buildRecordUsageLog(ctx, input, result, apiKey, user, account, subscription, + requestedModel, multiplier, imageMultiplier, accountRateMultiplier, billingType, cacheTTLOverridden, cost, opts) + + // 计算账号统计定价费用（使用最终上游模型匹配自定义规则） + if apiKey.GroupID != nil { + applyAccountStatsCost(ctx, usageLog, s.channelService, s.billingService, + account.ID, *apiKey.GroupID, result.UpstreamModel, result.Model, + // Anthropic's input_tokens excludes cache_read and cache_creation (billed separately); + // OpenAI gateway uses actualInputTokens which also excludes cache_read for the same reason. + UsageTokens{ + InputTokens: result.Usage.InputTokens, + OutputTokens: result.Usage.OutputTokens, + CacheCreationTokens: result.Usage.CacheCreationInputTokens, + CacheReadTokens: result.Usage.CacheReadInputTokens, + ImageOutputTokens: result.Usage.ImageOutputTokens, + }, + cost.TotalCost, + ) + } + + if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple { + writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway") + logger.LegacyPrintf("service.gateway", "[SIMPLE MODE] Usage recorded (not billed): user=%d, tokens=%d", usageLog.UserID, usageLog.TotalTokens()) + s.deferredService.ScheduleLastUsedUpdate(account.ID) + return nil + } + + // 配额平台由 handler 在请求 ctx 内经 QuotaPlatform() 算定并通过 input 传入； + // 后扣运行在 worker 池的 background ctx 上，无法再从 ctx 取 ForcePlatform。 + // 缺省（未设置）时回退到分组平台，保持对其它调用方的兼容。 + quotaPlatform := input.QuotaPlatform + if quotaPlatform == "" { + quotaPlatform = PlatformFromAPIKey(apiKey) + } + requestID := usageLog.RequestID + _, billingErr := applyUsageBilling(ctx, requestID, usageLog, &postUsageBillingParams{ + Cost: cost, + User: user, + APIKey: apiKey, + Account: account, + Subscription: subscription, + RequestPayloadHash: resolveUsageBillingPayloadFingerprint(ctx, input.RequestPayloadHash), + IsSubscriptionBill: isSubscriptionBilling, + AccountRateMultiplier: accountRateMultiplier, + APIKeyService: input.APIKeyService, + Platform: quotaPlatform, + }, s.billingDeps(), s.usageBillingRepo) + + if billingErr != nil { + return billingErr + } + writeUsageLogBestEffort(ctx, s.usageLogRepo, usageLog, "service.gateway") + + return nil +} + +// calculateRecordUsageCost 根据请求类型和选项计算费用。 +func (s *GatewayService) calculateRecordUsageCost( + ctx context.Context, + result *ForwardResult, + apiKey *APIKey, + billingModel string, + multiplier float64, + imageMultiplier float64, + opts *recordUsageOpts, +) *CostBreakdown { + // 图片生成：渠道定价为 token 计费时走 token 路径，否则走图片计费 + if result.ImageCount > 0 { + if resolved := s.resolveChannelPricing(ctx, billingModel, apiKey); resolved != nil && resolved.Mode == BillingModeToken { + return s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts) + } + return s.calculateImageCost(ctx, result, apiKey, billingModel, imageMultiplier) + } + + // Token 计费 + return s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts) +} + +// resolveChannelPricing 检查指定模型是否存在渠道级别定价。 +// 返回非 nil 的 ResolvedPricing 表示有渠道定价，nil 表示走默认定价路径。 +func (s *GatewayService) resolveChannelPricing(ctx context.Context, billingModel string, apiKey *APIKey) *ResolvedPricing { + if s.resolver == nil || apiKey.Group == nil { + return nil + } + gid := apiKey.Group.ID + resolved := s.resolver.Resolve(ctx, PricingInput{Model: billingModel, GroupID: &gid}) + if resolved.Source == PricingSourceChannel { + return resolved + } + return nil +} + +// calculateImageCost 计算图片生成费用：渠道级别定价优先，否则走按次计费。 +func (s *GatewayService) calculateImageCost( + ctx context.Context, + result *ForwardResult, + apiKey *APIKey, + billingModel string, + multiplier float64, +) *CostBreakdown { + sizeTier := NormalizeImageBillingTierOrDefault(result.ImageSize) + if resolved := s.resolveChannelPricing(ctx, billingModel, apiKey); resolved != nil { + tokens := UsageTokens{ + InputTokens: result.Usage.InputTokens, + OutputTokens: result.Usage.OutputTokens, + ImageOutputTokens: result.Usage.ImageOutputTokens, + } + gid := apiKey.Group.ID + cost, err := s.billingService.CalculateCostUnified(CostInput{ + Ctx: ctx, + Model: billingModel, + GroupID: &gid, + Tokens: tokens, + RequestCount: result.ImageCount, + SizeTier: sizeTier, + RateMultiplier: multiplier, + Resolver: s.resolver, + Resolved: resolved, + }) + if err != nil { + logger.LegacyPrintf("service.gateway", "Calculate image token cost failed: %v", err) + return &CostBreakdown{ActualCost: 0} + } + return cost + } + + var groupConfig *ImagePriceConfig + if apiKey.Group != nil { + groupConfig = &ImagePriceConfig{ + Price1K: apiKey.Group.ImagePrice1K, + Price2K: apiKey.Group.ImagePrice2K, + Price4K: apiKey.Group.ImagePrice4K, + } + } + return s.billingService.CalculateImageCost(billingModel, sizeTier, result.ImageCount, groupConfig, multiplier) +} + +// calculateTokenCost 计算 Token 计费：根据 opts 决定走普通/长上下文/渠道统一计费。 +func (s *GatewayService) calculateTokenCost( + ctx context.Context, + result *ForwardResult, + apiKey *APIKey, + billingModel string, + multiplier float64, + opts *recordUsageOpts, +) *CostBreakdown { + tokens := UsageTokens{ + InputTokens: result.Usage.InputTokens, + OutputTokens: result.Usage.OutputTokens, + CacheCreationTokens: result.Usage.CacheCreationInputTokens, + CacheReadTokens: result.Usage.CacheReadInputTokens, + CacheCreation5mTokens: result.Usage.CacheCreation5mTokens, + CacheCreation1hTokens: result.Usage.CacheCreation1hTokens, + ImageOutputTokens: result.Usage.ImageOutputTokens, + } + + var cost *CostBreakdown + var err error + + // 优先尝试渠道定价 → CalculateCostUnified + if resolved := s.resolveChannelPricing(ctx, billingModel, apiKey); resolved != nil { + gid := apiKey.Group.ID + cost, err = s.billingService.CalculateCostUnified(CostInput{ + Ctx: ctx, + Model: billingModel, + GroupID: &gid, + Tokens: tokens, + RequestCount: 1, + RateMultiplier: multiplier, + Resolver: s.resolver, + Resolved: resolved, + }) + } else if opts.LongContextThreshold > 0 { + // 长上下文双倍计费（如 Gemini 200K 阈值） + cost, err = s.billingService.CalculateCostWithLongContext(billingModel, tokens, multiplier, opts.LongContextThreshold, opts.LongContextMultiplier) + } else { + cost, err = s.billingService.CalculateCost(billingModel, tokens, multiplier) + } + if err != nil { + logger.LegacyPrintf("service.gateway", "Calculate cost failed: %v", err) + return &CostBreakdown{ActualCost: 0} + } + return cost +} + +// buildRecordUsageLog 构建使用日志并设置计费模式。 +func (s *GatewayService) buildRecordUsageLog( + ctx context.Context, + input *recordUsageCoreInput, + result *ForwardResult, + apiKey *APIKey, + user *User, + account *Account, + subscription *UserSubscription, + requestedModel string, + multiplier float64, + imageMultiplier float64, + accountRateMultiplier float64, + billingType int8, + cacheTTLOverridden bool, + cost *CostBreakdown, + opts *recordUsageOpts, +) *UsageLog { + durationMs := int(result.Duration.Milliseconds()) + requestID := resolveUsageBillingRequestID(ctx, result.RequestID) + usageLog := &UsageLog{ + UserID: user.ID, + APIKeyID: apiKey.ID, + AccountID: account.ID, + RequestID: requestID, + Model: result.Model, + RequestedModel: requestedModel, + UpstreamModel: optionalNonEqualStringPtr(result.UpstreamModel, result.Model), + ReasoningEffort: result.ReasoningEffort, + InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint), + UpstreamEndpoint: optionalTrimmedStringPtr(input.UpstreamEndpoint), + InputTokens: result.Usage.InputTokens, + OutputTokens: result.Usage.OutputTokens, + CacheCreationTokens: result.Usage.CacheCreationInputTokens, + CacheReadTokens: result.Usage.CacheReadInputTokens, + CacheCreation5mTokens: result.Usage.CacheCreation5mTokens, + CacheCreation1hTokens: result.Usage.CacheCreation1hTokens, + ImageOutputTokens: result.Usage.ImageOutputTokens, + RateMultiplier: multiplier, + AccountRateMultiplier: &accountRateMultiplier, + BillingType: billingType, + BillingMode: resolveBillingMode(result, cost), + Stream: result.Stream, + DurationMs: &durationMs, + FirstTokenMs: result.FirstTokenMs, + ImageCount: result.ImageCount, + ImageSize: optionalTrimmedStringPtr(result.ImageSize), + ImageInputSize: optionalTrimmedStringPtr(result.ImageInputSize), + ImageOutputSize: optionalTrimmedStringPtr(result.ImageOutputSize), + ImageSizeSource: optionalTrimmedStringPtr(result.ImageSizeSource), + ImageSizeBreakdown: result.ImageSizeBreakdown, + CacheTTLOverridden: cacheTTLOverridden, + ChannelID: optionalInt64Ptr(input.ChannelID), + ModelMappingChain: optionalTrimmedStringPtr(input.ModelMappingChain), + UserAgent: optionalTrimmedStringPtr(input.UserAgent), + IPAddress: optionalTrimmedStringPtr(input.IPAddress), + GroupID: apiKey.GroupID, + SubscriptionID: optionalSubscriptionID(subscription), + CreatedAt: time.Now(), + } + if result.ImageCount > 0 && (cost == nil || cost.BillingMode != string(BillingModeToken)) { + usageLog.RateMultiplier = imageMultiplier + } + if cost != nil { + usageLog.InputCost = cost.InputCost + usageLog.OutputCost = cost.OutputCost + usageLog.ImageOutputCost = cost.ImageOutputCost + usageLog.CacheCreationCost = cost.CacheCreationCost + usageLog.CacheReadCost = cost.CacheReadCost + usageLog.TotalCost = cost.TotalCost + usageLog.ActualCost = cost.ActualCost + } + + return usageLog +} + +// resolveBillingMode 根据计费结果和请求类型确定计费模式。 +func resolveBillingMode(result *ForwardResult, cost *CostBreakdown) *string { + var mode string + switch { + case cost != nil && cost.BillingMode != "": + mode = cost.BillingMode + case result.ImageCount > 0: + mode = string(BillingModeImage) + default: + mode = string(BillingModeToken) + } + return &mode +} + +func optionalSubscriptionID(subscription *UserSubscription) *int64 { + if subscription != nil { + return &subscription.ID + } + return nil +}