mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
refactor(service): 纯移动拆分 gateway_service.go（7294→1289行）
This commit is contained in:
shaw
File diff suppressed because it is too large
Load Diff
@@ -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,
|
||||
},
|
||||
})
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user