李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
2468 lines
92 KiB
Go
2468 lines
92 KiB
Go
package handler
|
||
|
||
import (
|
||
"context"
|
||
"crypto/rand"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"math"
|
||
"net/http"
|
||
"strconv"
|
||
"strings"
|
||
"sync/atomic"
|
||
"time"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||
pkgerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/geminicli"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"go.uber.org/zap"
|
||
)
|
||
|
||
const gatewayCompatibilityMetricsLogInterval = 1024
|
||
|
||
var gatewayCompatibilityMetricsLogCounter atomic.Uint64
|
||
|
||
// GatewayHandler handles API gateway requests
|
||
type GatewayHandler struct {
|
||
gatewayService *service.GatewayService
|
||
openAIGatewayService *service.OpenAIGatewayService
|
||
geminiCompatService *service.GeminiMessagesCompatService
|
||
antigravityGatewayService *service.AntigravityGatewayService
|
||
userService *service.UserService
|
||
billingCacheService *service.BillingCacheService
|
||
usageService *service.UsageService
|
||
apiKeyService *service.APIKeyService
|
||
usageRecordWorkerPool *service.UsageRecordWorkerPool
|
||
errorPassthroughService *service.ErrorPassthroughService
|
||
contentModerationService *service.ContentModerationService
|
||
securityAuditCoordinator *securityaudit.Coordinator
|
||
concurrencyHelper *ConcurrencyHelper
|
||
userMsgQueueHelper *UserMsgQueueHelper
|
||
maxAccountSwitches int
|
||
maxAccountSwitchesGemini int
|
||
cfg *config.Config
|
||
settingService *service.SettingService
|
||
}
|
||
|
||
// NewGatewayHandler creates a new GatewayHandler
|
||
func NewGatewayHandler(
|
||
gatewayService *service.GatewayService,
|
||
openAIGatewayService *service.OpenAIGatewayService,
|
||
geminiCompatService *service.GeminiMessagesCompatService,
|
||
antigravityGatewayService *service.AntigravityGatewayService,
|
||
userService *service.UserService,
|
||
concurrencyService *service.ConcurrencyService,
|
||
billingCacheService *service.BillingCacheService,
|
||
usageService *service.UsageService,
|
||
apiKeyService *service.APIKeyService,
|
||
usageRecordWorkerPool *service.UsageRecordWorkerPool,
|
||
errorPassthroughService *service.ErrorPassthroughService,
|
||
contentModerationService *service.ContentModerationService,
|
||
userMsgQueueService *service.UserMessageQueueService,
|
||
cfg *config.Config,
|
||
settingService *service.SettingService,
|
||
) *GatewayHandler {
|
||
pingInterval := time.Duration(0)
|
||
maxAccountSwitches := 10
|
||
maxAccountSwitchesGemini := 3
|
||
if cfg != nil {
|
||
pingInterval = time.Duration(cfg.Concurrency.PingInterval) * time.Second
|
||
if cfg.Gateway.MaxAccountSwitches > 0 {
|
||
maxAccountSwitches = cfg.Gateway.MaxAccountSwitches
|
||
}
|
||
if cfg.Gateway.MaxAccountSwitchesGemini > 0 {
|
||
maxAccountSwitchesGemini = cfg.Gateway.MaxAccountSwitchesGemini
|
||
}
|
||
}
|
||
|
||
// 初始化用户消息串行队列 helper
|
||
var umqHelper *UserMsgQueueHelper
|
||
if userMsgQueueService != nil && cfg != nil {
|
||
umqHelper = NewUserMsgQueueHelper(userMsgQueueService, SSEPingFormatClaude, pingInterval)
|
||
}
|
||
|
||
return &GatewayHandler{
|
||
gatewayService: gatewayService,
|
||
openAIGatewayService: openAIGatewayService,
|
||
geminiCompatService: geminiCompatService,
|
||
antigravityGatewayService: antigravityGatewayService,
|
||
userService: userService,
|
||
billingCacheService: billingCacheService,
|
||
usageService: usageService,
|
||
apiKeyService: apiKeyService,
|
||
usageRecordWorkerPool: usageRecordWorkerPool,
|
||
errorPassthroughService: errorPassthroughService,
|
||
contentModerationService: contentModerationService,
|
||
concurrencyHelper: NewConcurrencyHelper(concurrencyService, SSEPingFormatClaude, pingInterval),
|
||
userMsgQueueHelper: umqHelper,
|
||
maxAccountSwitches: maxAccountSwitches,
|
||
maxAccountSwitchesGemini: maxAccountSwitchesGemini,
|
||
cfg: cfg,
|
||
settingService: settingService,
|
||
}
|
||
}
|
||
|
||
// Messages handles Claude API compatible messages endpoint
|
||
// POST /v1/messages
|
||
func (h *GatewayHandler) Messages(c *gin.Context) {
|
||
// 从context获取apiKey和user（ApiKeyAuth中间件已设置）
|
||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||
return
|
||
}
|
||
|
||
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
|
||
return
|
||
}
|
||
reqLog := requestLogger(
|
||
c,
|
||
"handler.gateway.messages",
|
||
zap.Int64("user_id", subject.UserID),
|
||
zap.Int64("api_key_id", apiKey.ID),
|
||
zap.Any("group_id", apiKey.GroupID),
|
||
)
|
||
defer h.maybeLogCompatibilityFallbackMetrics(reqLog)
|
||
|
||
// 读取请求体
|
||
body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg)
|
||
if err != nil {
|
||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||
h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
|
||
return
|
||
}
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
|
||
return
|
||
}
|
||
|
||
if len(body) == 0 {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
|
||
return
|
||
}
|
||
|
||
setOpsRequestContext(c, "", false)
|
||
|
||
bodyRef := service.NewRequestBodyRef(body)
|
||
parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic)
|
||
if err != nil {
|
||
logRequestBodyParseFailure(reqLog, body, err)
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||
return
|
||
}
|
||
body = parsedReq.Body.Bytes()
|
||
reqModel := parsedReq.Model
|
||
reqStream := parsedReq.Stream
|
||
ensureCompositeTargetPlatform(c, apiKey, reqModel)
|
||
reqLog = reqLog.With(zap.String("model", reqModel), zap.Bool("stream", reqStream))
|
||
|
||
// 解析渠道级模型映射
|
||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel)
|
||
|
||
// 设置 max_tokens=1 + haiku 探测请求标识到 context 中
|
||
// 必须在 SetClaudeCodeClientContext 之前设置，因为 ClaudeCodeValidator 需要读取此标识进行绕过判断
|
||
if isMaxTokensOneHaikuRequest(reqModel, parsedReq.MaxTokens) {
|
||
ctx := service.WithIsMaxTokensOneHaikuRequest(c.Request.Context(), true, h.metadataBridgeEnabled())
|
||
c.Request = c.Request.WithContext(ctx)
|
||
}
|
||
|
||
// 检查是否为 Claude Code 客户端，设置到 context 中（复用已解析请求，避免二次反序列化）。
|
||
SetClaudeCodeClientContext(c, body, parsedReq)
|
||
isClaudeCodeClient := service.IsClaudeCodeClient(c.Request.Context())
|
||
|
||
// 版本检查：仅对 Claude Code 客户端，拒绝低于最低版本的请求
|
||
if !h.checkClaudeCodeVersion(c) {
|
||
return
|
||
}
|
||
|
||
// 在请求上下文中记录 thinking 状态，供 Antigravity 最终模型 key 推导/模型维度限流使用
|
||
c.Request = c.Request.WithContext(service.WithThinkingEnabled(c.Request.Context(), parsedReq.ThinkingEnabled, h.metadataBridgeEnabled()))
|
||
|
||
setOpsRequestContext(c, reqModel, reqStream)
|
||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(reqStream, false)))
|
||
pricingCtx, pricingAt := service.WithGatewayTokenRequestPricing(c.Request.Context())
|
||
c.Request = c.Request.WithContext(pricingCtx)
|
||
|
||
// 验证 model 必填
|
||
if reqModel == "" {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||
return
|
||
}
|
||
if !compositeTargetPlatformResolved(c, apiKey, reqModel) {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by composite groups")
|
||
return
|
||
}
|
||
|
||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolAnthropicMessages, reqModel, body); decision != nil && !decision.AllowNextStage {
|
||
h.anthropicSecurityAuditError(c, decision)
|
||
return
|
||
}
|
||
|
||
// Track if we've started streaming (for error handling)
|
||
streamStarted := false
|
||
|
||
// 绑定错误透传服务，允许 service 层在非 failover 错误场景复用规则。
|
||
if h.errorPassthroughService != nil {
|
||
service.BindErrorPassthroughService(c, h.errorPassthroughService)
|
||
}
|
||
|
||
// 获取订阅信息（可能为nil）- 提前获取用于后续检查
|
||
subscription, _ := middleware2.GetSubscriptionFromContext(c)
|
||
|
||
// 1. 首先获取用户并发槽位
|
||
userReleaseFunc, err := h.concurrencyHelper.AcquireUserSlotWithWait(c, subject.UserID, subject.Concurrency, reqStream, &streamStarted)
|
||
if err != nil {
|
||
reqLog.Warn("gateway.user_slot_acquire_failed", zap.Error(err))
|
||
h.handleConcurrencyError(c, err, "user", streamStarted)
|
||
return
|
||
}
|
||
// 在请求结束或 Context 取消时确保释放槽位，避免客户端断开造成泄漏
|
||
userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
|
||
if userReleaseFunc != nil {
|
||
defer userReleaseFunc()
|
||
}
|
||
|
||
// 2. 【新增】Wait后二次检查余额/订阅
|
||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil {
|
||
reqLog.Info("gateway.billing_eligibility_check_failed", zap.Error(err))
|
||
status, code, message, retryAfter := billingErrorDetails(err)
|
||
if retryAfter > 0 {
|
||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||
}
|
||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||
return
|
||
}
|
||
|
||
// 设置请求所属分组 ID（用于渠道级功能判断，如 WebSearch 模拟）
|
||
parsedReq.GroupID = apiKey.GroupID
|
||
|
||
// 计算粘性会话hash
|
||
parsedReq.SessionContext = &service.SessionContext{
|
||
ClientIP: ip.GetClientIP(c),
|
||
UserAgent: c.GetHeader("User-Agent"),
|
||
APIKeyID: apiKey.ID,
|
||
}
|
||
sessionHash := h.gatewayService.GenerateSessionHash(parsedReq)
|
||
|
||
// [DEBUG-STICKY] 打印会话 hash 生成结果
|
||
reqLog.Info("sticky.session_hash_generated",
|
||
zap.String("session_hash", sessionHash),
|
||
zap.String("metadata_user_id_raw", parsedReq.MetadataUserID),
|
||
)
|
||
|
||
// 获取平台：优先使用强制平台（/antigravity 路由），其次使用 composite 解析出的目标平台，否则使用分组平台
|
||
platform := ""
|
||
if forcePlatform, ok := middleware2.GetForcePlatformFromContext(c); ok {
|
||
platform = forcePlatform
|
||
} else if resolvedPlatform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context()); ok {
|
||
platform = resolvedPlatform
|
||
} else if apiKey.Group != nil {
|
||
platform = apiKey.Group.Platform
|
||
}
|
||
sessionKey := sessionHash
|
||
if platform == service.PlatformGemini && sessionHash != "" {
|
||
sessionKey = "gemini:" + sessionHash
|
||
}
|
||
|
||
// 查询粘性会话绑定的账号 ID
|
||
var sessionBoundAccountID int64
|
||
if sessionKey != "" {
|
||
sessionBoundAccountID, _ = h.gatewayService.GetCachedSessionAccountID(c.Request.Context(), apiKey.GroupID, sessionKey)
|
||
// [DEBUG-STICKY] 打印粘性会话查询结果
|
||
reqLog.Info("sticky.cache_lookup",
|
||
zap.String("session_key", sessionKey),
|
||
zap.Int64("bound_account_id", sessionBoundAccountID),
|
||
)
|
||
if sessionBoundAccountID > 0 {
|
||
prefetchedGroupID := int64(0)
|
||
if apiKey.GroupID != nil {
|
||
prefetchedGroupID = *apiKey.GroupID
|
||
}
|
||
ctx := service.WithPrefetchedStickySession(c.Request.Context(), sessionBoundAccountID, prefetchedGroupID, h.metadataBridgeEnabled())
|
||
c.Request = c.Request.WithContext(ctx)
|
||
}
|
||
} else {
|
||
reqLog.Info("sticky.no_session_key", zap.String("session_hash", sessionHash))
|
||
}
|
||
// 判断是否真的绑定了粘性会话：有 sessionKey 且已经绑定到某个账号
|
||
hasBoundSession := sessionKey != "" && sessionBoundAccountID > 0
|
||
|
||
if platform == service.PlatformGemini {
|
||
fs := NewFailoverState(h.maxAccountSwitchesGemini, hasBoundSession)
|
||
|
||
// 单账号分组提前设置 SingleAccountRetry 标记，让 Service 层首次 503 就不设模型限流标记。
|
||
// 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流，导致后续请求连续快速失败。
|
||
if h.gatewayService.IsSingleAntigravityAccountGroup(c.Request.Context(), apiKey.GroupID) {
|
||
ctx := service.WithSingleAccountRetry(c.Request.Context(), true, h.metadataBridgeEnabled())
|
||
c.Request = c.Request.WithContext(ctx)
|
||
}
|
||
|
||
for {
|
||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, "", int64(0)) // Gemini 不使用会话限制
|
||
if err != nil {
|
||
if len(fs.FailedAccountIDs) == 0 {
|
||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, service.PlatformGemini)
|
||
if !cls.ModelNotFound {
|
||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||
}
|
||
reqLog.Warn("gateway.select_account_no_available",
|
||
zap.String("model", reqModel),
|
||
zap.Int64p("group_id", apiKey.GroupID),
|
||
zap.String("platform", platform),
|
||
zap.Bool("model_not_found", cls.ModelNotFound),
|
||
zap.Error(err),
|
||
)
|
||
message := cls.Message
|
||
if !cls.ModelNotFound {
|
||
message = "No available accounts: " + err.Error()
|
||
}
|
||
h.handleStreamingAwareError(c, cls.Status, cls.ErrType, message, streamStarted)
|
||
return
|
||
}
|
||
action := fs.HandleSelectionExhausted(c.Request.Context())
|
||
switch action {
|
||
case FailoverContinue:
|
||
ctx := service.WithSingleAccountRetry(c.Request.Context(), true, h.metadataBridgeEnabled())
|
||
c.Request = c.Request.WithContext(ctx)
|
||
continue
|
||
case FailoverCanceled:
|
||
failoverClientGone(c)
|
||
return
|
||
default: // FailoverExhausted
|
||
if fs.LastFailoverErr != nil {
|
||
h.handleFailoverExhausted(c, fs.LastFailoverErr, service.PlatformGemini, streamStarted)
|
||
} else {
|
||
h.handleFailoverExhaustedSimple(c, 502, streamStarted)
|
||
}
|
||
return
|
||
}
|
||
}
|
||
account := selection.Account
|
||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||
|
||
// 检查请求拦截（预热请求、SUGGESTION MODE等）
|
||
if account.IsInterceptWarmupEnabled() {
|
||
interceptType := detectInterceptType(body, reqModel, parsedReq.MaxTokens, isClaudeCodeClient)
|
||
if interceptType != InterceptTypeNone {
|
||
if selection.Acquired && selection.ReleaseFunc != nil {
|
||
selection.ReleaseFunc()
|
||
}
|
||
if reqStream {
|
||
sendMockInterceptStream(c, reqModel, interceptType)
|
||
} else {
|
||
sendMockInterceptResponse(c, reqModel, interceptType)
|
||
}
|
||
return
|
||
}
|
||
}
|
||
|
||
// 3. 获取账号并发槽位
|
||
accountReleaseFunc := selection.ReleaseFunc
|
||
if !selection.Acquired {
|
||
if selection.WaitPlan == nil {
|
||
markOpsRoutingCapacityLimited(c)
|
||
reqLog.Warn("gateway.select_account_no_slot_no_wait_plan",
|
||
zap.Int64("account_id", account.ID),
|
||
zap.String("model", reqModel),
|
||
zap.String("platform", platform),
|
||
)
|
||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts", streamStarted)
|
||
return
|
||
}
|
||
accountWaitCounted := false
|
||
canWait, err := h.concurrencyHelper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
|
||
if err != nil {
|
||
reqLog.Warn("gateway.account_wait_counter_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||
} else if !canWait {
|
||
reqLog.Info("gateway.account_wait_queue_full",
|
||
zap.Int64("account_id", account.ID),
|
||
zap.Int("max_waiting", selection.WaitPlan.MaxWaiting),
|
||
)
|
||
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later", streamStarted)
|
||
return
|
||
}
|
||
if err == nil && canWait {
|
||
accountWaitCounted = true
|
||
}
|
||
releaseWait := func() {
|
||
if accountWaitCounted {
|
||
h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
||
accountWaitCounted = false
|
||
}
|
||
}
|
||
|
||
accountReleaseFunc, err = h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout(
|
||
c,
|
||
account.ID,
|
||
selection.WaitPlan.MaxConcurrency,
|
||
selection.WaitPlan.Timeout,
|
||
reqStream,
|
||
&streamStarted,
|
||
)
|
||
if err != nil {
|
||
reqLog.Warn("gateway.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||
releaseWait()
|
||
h.handleConcurrencyError(c, err, "account", streamStarted)
|
||
return
|
||
}
|
||
// Slot acquired: no longer waiting in queue.
|
||
releaseWait()
|
||
}
|
||
// 终检与准入后绑定使用选号结果携带的门（见 responses 同名注释）。
|
||
admissionCtx := service.ContextWithSelectionProfitGate(c.Request.Context(), selection)
|
||
latest, vetoed, reason := h.gatewayService.GatewayProfitControlVetoLatest(admissionCtx, account)
|
||
if vetoed {
|
||
if accountReleaseFunc != nil {
|
||
accountReleaseFunc()
|
||
}
|
||
reqLog.Debug("gateway.account_slot_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason))
|
||
if fs.RecordProfitVeto(account.ID) == FailoverExhausted {
|
||
reqLog.Warn("gateway.profit_veto_attempts_exhausted", zap.Int("profit_veto_count", fs.ProfitVetoCount()))
|
||
markOpsRoutingCapacityLimited(c)
|
||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", profitVetoExhaustedMessage, streamStarted)
|
||
return
|
||
}
|
||
continue
|
||
}
|
||
account = latest
|
||
selection.Account = latest
|
||
// 等待路径保持既有 eager 绑定（无门时 helper 直接绑定）；调度器已
|
||
// 抢槽的直达路径无门时由选号内部绑定，这里只在门下补准入后绑定。
|
||
if selection.ProfitGateActive() || !selection.Acquired {
|
||
if err := h.gatewayService.BindStickySessionAfterProfitAdmission(admissionCtx, apiKey.GroupID, sessionKey, account.ID); err != nil {
|
||
reqLog.Warn("gateway.bind_sticky_session_after_profit_admission_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||
}
|
||
}
|
||
// 账号槽位/等待计数需要在超时或断开时安全回收
|
||
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
|
||
|
||
// 转发请求 - 根据账号平台分流
|
||
var result *service.ForwardResult
|
||
requestCtx := c.Request.Context()
|
||
if fs.SwitchCount > 0 {
|
||
requestCtx = service.WithAccountSwitchCount(requestCtx, fs.SwitchCount, h.metadataBridgeEnabled())
|
||
}
|
||
// 记录 Forward 前已写入字节数，Forward 后若增加则说明 SSE 内容已发，禁止 failover
|
||
writerSizeBeforeForward := c.Writer.Size()
|
||
if account.Platform == service.PlatformAntigravity {
|
||
result, err = h.antigravityGatewayService.ForwardGemini(
|
||
requestCtx,
|
||
c,
|
||
account,
|
||
reqModel,
|
||
"generateContent",
|
||
reqStream,
|
||
body,
|
||
hasBoundSession,
|
||
service.WithForwardGeminiSession(derefGroupID(apiKey.GroupID), sessionKey),
|
||
)
|
||
} else {
|
||
result, err = h.geminiCompatService.Forward(requestCtx, c, account, body)
|
||
}
|
||
if accountReleaseFunc != nil {
|
||
accountReleaseFunc()
|
||
}
|
||
if err != nil {
|
||
var failoverErr *service.UpstreamFailoverError
|
||
if errors.As(err, &failoverErr) {
|
||
// 流式内容已写入客户端，无法撤销，禁止 failover 以防止流拼接腐化
|
||
if c.Writer.Size() != writerSizeBeforeForward {
|
||
h.handleFailoverExhausted(c, failoverErr, service.PlatformGemini, true)
|
||
return
|
||
}
|
||
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
|
||
switch action {
|
||
case FailoverContinue:
|
||
continue
|
||
case FailoverExhausted:
|
||
h.handleFailoverExhausted(c, fs.LastFailoverErr, service.PlatformGemini, streamStarted)
|
||
return
|
||
case FailoverCanceled:
|
||
failoverClientGone(c)
|
||
return
|
||
}
|
||
}
|
||
upstreamErrorAlreadyCommunicated := gatewayForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err)
|
||
wroteFallback := false
|
||
if !upstreamErrorAlreadyCommunicated {
|
||
wroteFallback = h.ensureForwardErrorResponse(c, streamStarted)
|
||
}
|
||
forwardFailedFields := []zap.Field{
|
||
zap.Int64("account_id", account.ID),
|
||
zap.String("account_name", account.Name),
|
||
zap.String("account_platform", account.Platform),
|
||
zap.Bool("fallback_error_response_written", wroteFallback),
|
||
zap.Bool("upstream_error_response_already_written", upstreamErrorAlreadyCommunicated),
|
||
zap.Error(err),
|
||
}
|
||
if account.Proxy != nil {
|
||
forwardFailedFields = append(forwardFailedFields,
|
||
zap.Int64("proxy_id", account.Proxy.ID),
|
||
zap.String("proxy_name", account.Proxy.Name),
|
||
zap.String("proxy_host", account.Proxy.Host),
|
||
zap.Int("proxy_port", account.Proxy.Port),
|
||
)
|
||
} else if account.ProxyID != nil {
|
||
forwardFailedFields = append(forwardFailedFields, zap.Int64p("proxy_id", account.ProxyID))
|
||
}
|
||
reqLog.Error("gateway.forward_failed", forwardFailedFields...)
|
||
return
|
||
}
|
||
|
||
// RPM 计数递增（Forward 成功后）
|
||
// 注意：TOCTOU 竞态是已知且可接受的设计权衡，与 WindowCost 一致的 soft-limit 模式。
|
||
// 在高并发下可能短暂超出 RPM 限制，但不会导致请求失败。
|
||
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
|
||
if err := h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID); err != nil {
|
||
reqLog.Warn("gateway.rpm_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||
}
|
||
}
|
||
|
||
// 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context）
|
||
userAgent := c.GetHeader("User-Agent")
|
||
clientIP := ip.GetClientIP(c)
|
||
requestPayloadHash := service.HashUsageRequestPayload(body)
|
||
inboundEndpoint := GetInboundEndpoint(c)
|
||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||
|
||
if result.ReasoningEffort == nil {
|
||
result.ReasoningEffort = service.NormalizeClaudeOutputEffort(parsedReq.OutputEffort)
|
||
}
|
||
// 国产模型 thinking-enabled 默认 effort 填充：Kimi/GLM/MiniMax 这些不支持 effort 档位的
|
||
// passback-required 上游，仅要 thinking 启用且 OutputEffort 未明确传递时，在 usage_log 写 "high"
|
||
// 避免该字段长期为 NULL（详见 DefaultEffortForThinkingEnabled 文档）。
|
||
if result.ReasoningEffort == nil && parsedReq.ThinkingEnabled {
|
||
protocolModel := result.UpstreamModel
|
||
if protocolModel == "" {
|
||
protocolModel = result.Model
|
||
}
|
||
result.ReasoningEffort = service.DefaultEffortForThinkingEnabled(protocolModel)
|
||
}
|
||
|
||
// 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。
|
||
// ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。
|
||
forceCacheBilling := fs.ForceCacheBilling
|
||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||
sessionID := service.ExtractClientSessionID(c)
|
||
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||
Result: result,
|
||
QuotaPlatform: quotaPlatform,
|
||
APIKey: apiKey,
|
||
User: apiKey.User,
|
||
Account: account,
|
||
Subscription: subscription,
|
||
PricingAt: pricingAt,
|
||
InboundEndpoint: inboundEndpoint,
|
||
UpstreamEndpoint: upstreamEndpoint,
|
||
UserAgent: userAgent,
|
||
IPAddress: clientIP,
|
||
SessionID: sessionID,
|
||
RequestPayloadHash: requestPayloadHash,
|
||
ForceCacheBilling: forceCacheBilling,
|
||
APIKeyService: h.apiKeyService,
|
||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
|
||
}); err != nil {
|
||
logger.L().With(
|
||
zap.String("component", "handler.gateway.messages"),
|
||
zap.Int64("user_id", subject.UserID),
|
||
zap.Int64("api_key_id", apiKey.ID),
|
||
zap.Any("group_id", apiKey.GroupID),
|
||
zap.String("model", reqModel),
|
||
zap.Int64("account_id", account.ID),
|
||
).Error("gateway.record_usage_failed", zap.Error(err))
|
||
}
|
||
})
|
||
return
|
||
}
|
||
}
|
||
|
||
currentAPIKey := apiKey
|
||
currentSubscription := subscription
|
||
var fallbackGroupID *int64
|
||
if apiKey.Group != nil {
|
||
fallbackGroupID = apiKey.Group.FallbackGroupIDOnInvalidRequest
|
||
}
|
||
fallbackUsed := false
|
||
|
||
// 单账号分组提前设置 SingleAccountRetry 标记，让 Service 层首次 503 就不设模型限流标记。
|
||
// 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流，导致后续请求连续快速失败。
|
||
if h.gatewayService.IsSingleAntigravityAccountGroup(c.Request.Context(), currentAPIKey.GroupID) {
|
||
ctx := service.WithSingleAccountRetry(c.Request.Context(), true, h.metadataBridgeEnabled())
|
||
c.Request = c.Request.WithContext(ctx)
|
||
}
|
||
|
||
for {
|
||
fs := NewFailoverState(h.maxAccountSwitches, hasBoundSession)
|
||
retryWithFallback := false
|
||
|
||
for {
|
||
attemptParsedReq, err := parsedReq.CloneForBody(body)
|
||
if err != nil {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||
return
|
||
}
|
||
|
||
// 选择支持该模型的账号
|
||
reqLog.Info("sticky.selecting_account",
|
||
zap.String("session_key", sessionKey),
|
||
zap.Int64("sticky_bound_account_id", sessionBoundAccountID),
|
||
zap.Bool("has_bound_session", hasBoundSession),
|
||
zap.Int("failed_account_count", len(fs.FailedAccountIDs)),
|
||
)
|
||
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), currentAPIKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, parsedReq.MetadataUserID, subject.UserID)
|
||
if err != nil {
|
||
if len(fs.FailedAccountIDs) == 0 {
|
||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, currentAPIKey, reqModel, reqModel, platform)
|
||
if !cls.ModelNotFound {
|
||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||
}
|
||
reqLog.Warn("gateway.select_account_no_available",
|
||
zap.String("model", reqModel),
|
||
zap.Int64p("group_id", currentAPIKey.GroupID),
|
||
zap.String("platform", platform),
|
||
zap.Bool("fallback_used", fallbackUsed),
|
||
zap.Bool("model_not_found", cls.ModelNotFound),
|
||
zap.Error(err),
|
||
)
|
||
message := cls.Message
|
||
if !cls.ModelNotFound {
|
||
message = "No available accounts: " + err.Error()
|
||
}
|
||
h.handleStreamingAwareError(c, cls.Status, cls.ErrType, message, streamStarted)
|
||
return
|
||
}
|
||
action := fs.HandleSelectionExhausted(c.Request.Context())
|
||
switch action {
|
||
case FailoverContinue:
|
||
ctx := service.WithSingleAccountRetry(c.Request.Context(), true, h.metadataBridgeEnabled())
|
||
c.Request = c.Request.WithContext(ctx)
|
||
continue
|
||
case FailoverCanceled:
|
||
failoverClientGone(c)
|
||
return
|
||
default: // FailoverExhausted
|
||
if fs.LastFailoverErr != nil {
|
||
h.handleFailoverExhausted(c, fs.LastFailoverErr, platform, streamStarted)
|
||
} else {
|
||
h.handleFailoverExhaustedSimple(c, 502, streamStarted)
|
||
}
|
||
return
|
||
}
|
||
}
|
||
account := selection.Account
|
||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||
|
||
// [DEBUG-STICKY] 打印账号选择结果
|
||
reqLog.Info("sticky.account_selected",
|
||
zap.Int64("selected_account_id", account.ID),
|
||
zap.String("account_name", account.Name),
|
||
zap.Bool("slot_acquired", selection.Acquired),
|
||
zap.Bool("has_wait_plan", selection.WaitPlan != nil),
|
||
zap.Int64("sticky_bound_account_id", sessionBoundAccountID),
|
||
zap.Bool("sticky_honored", sessionBoundAccountID > 0 && sessionBoundAccountID == account.ID),
|
||
)
|
||
|
||
// 检查请求拦截（预热请求、SUGGESTION MODE等）
|
||
if account.IsInterceptWarmupEnabled() {
|
||
interceptType := detectInterceptType(body, reqModel, parsedReq.MaxTokens, isClaudeCodeClient)
|
||
if interceptType != InterceptTypeNone {
|
||
if selection.Acquired && selection.ReleaseFunc != nil {
|
||
selection.ReleaseFunc()
|
||
}
|
||
if reqStream {
|
||
sendMockInterceptStream(c, reqModel, interceptType)
|
||
} else {
|
||
sendMockInterceptResponse(c, reqModel, interceptType)
|
||
}
|
||
return
|
||
}
|
||
}
|
||
|
||
// 3. 获取账号并发槽位
|
||
accountReleaseFunc := selection.ReleaseFunc
|
||
if !selection.Acquired {
|
||
if selection.WaitPlan == nil {
|
||
markOpsRoutingCapacityLimited(c)
|
||
reqLog.Warn("gateway.select_account_no_slot_no_wait_plan",
|
||
zap.Int64("account_id", account.ID),
|
||
zap.String("model", reqModel),
|
||
zap.String("platform", platform),
|
||
)
|
||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts", streamStarted)
|
||
return
|
||
}
|
||
accountWaitCounted := false
|
||
canWait, err := h.concurrencyHelper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
|
||
if err != nil {
|
||
reqLog.Warn("gateway.account_wait_counter_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||
} else if !canWait {
|
||
reqLog.Info("gateway.account_wait_queue_full",
|
||
zap.Int64("account_id", account.ID),
|
||
zap.Int("max_waiting", selection.WaitPlan.MaxWaiting),
|
||
)
|
||
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later", streamStarted)
|
||
return
|
||
}
|
||
if err == nil && canWait {
|
||
accountWaitCounted = true
|
||
}
|
||
releaseWait := func() {
|
||
if accountWaitCounted {
|
||
h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
||
accountWaitCounted = false
|
||
}
|
||
}
|
||
|
||
accountReleaseFunc, err = h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout(
|
||
c,
|
||
account.ID,
|
||
selection.WaitPlan.MaxConcurrency,
|
||
selection.WaitPlan.Timeout,
|
||
reqStream,
|
||
&streamStarted,
|
||
)
|
||
if err != nil {
|
||
reqLog.Warn("gateway.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||
releaseWait()
|
||
h.handleConcurrencyError(c, err, "account", streamStarted)
|
||
return
|
||
}
|
||
// Slot acquired: no longer waiting in queue.
|
||
releaseWait()
|
||
}
|
||
// 终检与准入后绑定使用选号结果携带的门（见 responses 同名注释）。
|
||
admissionCtx := service.ContextWithSelectionProfitGate(c.Request.Context(), selection)
|
||
latest, vetoed, reason := h.gatewayService.GatewayProfitControlVetoLatest(admissionCtx, account)
|
||
if vetoed {
|
||
if accountReleaseFunc != nil {
|
||
accountReleaseFunc()
|
||
}
|
||
reqLog.Debug("gateway.account_slot_profit_vetoed", zap.Int64("account_id", account.ID), zap.String("reason", reason))
|
||
if fs.RecordProfitVeto(account.ID) == FailoverExhausted {
|
||
reqLog.Warn("gateway.profit_veto_attempts_exhausted", zap.Int("profit_veto_count", fs.ProfitVetoCount()))
|
||
markOpsRoutingCapacityLimited(c)
|
||
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", profitVetoExhaustedMessage, streamStarted)
|
||
return
|
||
}
|
||
continue
|
||
}
|
||
account = latest
|
||
selection.Account = latest
|
||
// 等待路径保持既有 eager 绑定（无门时 helper 直接绑定）；调度器已
|
||
// 抢槽的直达路径无门时由选号内部绑定，这里只在门下补准入后绑定。
|
||
if selection.ProfitGateActive() || !selection.Acquired {
|
||
if err := h.gatewayService.BindStickySessionAfterProfitAdmission(admissionCtx, currentAPIKey.GroupID, sessionKey, account.ID); err != nil {
|
||
reqLog.Warn("gateway.bind_sticky_session_after_profit_admission_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||
}
|
||
}
|
||
// 账号槽位/等待计数需要在超时或断开时安全回收
|
||
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
|
||
|
||
// ===== 用户消息串行队列 START =====
|
||
var queueRelease func()
|
||
umqMode := h.getUserMsgQueueMode(account, attemptParsedReq)
|
||
|
||
switch umqMode {
|
||
case config.UMQModeSerialize:
|
||
// 串行模式：获取锁 + RPM 延迟 + 释放（当前行为不变）
|
||
baseRPM := account.GetBaseRPM()
|
||
release, qErr := h.userMsgQueueHelper.AcquireWithWait(
|
||
c, account.ID, baseRPM, reqStream, &streamStarted,
|
||
h.cfg.Gateway.UserMessageQueue.WaitTimeout(),
|
||
reqLog,
|
||
)
|
||
if qErr != nil {
|
||
// fail-open: 记录 warn，不阻止请求
|
||
reqLog.Warn("gateway.umq_acquire_failed",
|
||
zap.Int64("account_id", account.ID),
|
||
zap.Error(qErr),
|
||
)
|
||
} else {
|
||
queueRelease = release
|
||
}
|
||
|
||
case config.UMQModeThrottle:
|
||
// 软性限速：仅施加 RPM 自适应延迟，不阻塞并发
|
||
baseRPM := account.GetBaseRPM()
|
||
if tErr := h.userMsgQueueHelper.ThrottleWithPing(
|
||
c, account.ID, baseRPM, reqStream, &streamStarted,
|
||
h.cfg.Gateway.UserMessageQueue.WaitTimeout(),
|
||
reqLog,
|
||
); tErr != nil {
|
||
reqLog.Warn("gateway.umq_throttle_failed",
|
||
zap.Int64("account_id", account.ID),
|
||
zap.Error(tErr),
|
||
)
|
||
}
|
||
|
||
default:
|
||
if umqMode != "" {
|
||
reqLog.Warn("gateway.umq_unknown_mode",
|
||
zap.String("mode", umqMode),
|
||
zap.Int64("account_id", account.ID),
|
||
)
|
||
}
|
||
}
|
||
|
||
// 用 wrapReleaseOnDone 确保 context 取消时自动释放（仅 serialize 模式有 queueRelease）
|
||
queueRelease = wrapReleaseOnDone(c.Request.Context(), queueRelease)
|
||
// 注入回调到 ParsedRequest：使用外层 wrapper 以便提前清理 AfterFunc
|
||
attemptParsedReq.OnUpstreamAccepted = queueRelease
|
||
// ===== 用户消息串行队列 END =====
|
||
|
||
// 渠道模型映射只作用于本次账号尝试，避免 failover 后污染原始 ParsedRequest。
|
||
if channelMapping.Mapped {
|
||
attemptParsedReq.Model = channelMapping.MappedModel
|
||
if err := attemptParsedReq.ReplaceBody(h.gatewayService.ReplaceModelInBody(attemptParsedReq.Body.Bytes(), channelMapping.MappedModel)); err != nil {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||
return
|
||
}
|
||
}
|
||
// Bedrock CC 兼容：清理 body 专有字段 + 过滤 anthropic-beta header，适用于所有转发路径
|
||
if err := attemptParsedReq.ReplaceBody(h.gatewayService.ApplyBedrockCCCompat(c, attemptParsedReq.Body.Bytes(), attemptParsedReq.Model, account, apiKey.GroupID)); err != nil {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||
return
|
||
}
|
||
attemptBody := attemptParsedReq.Body.Bytes()
|
||
|
||
// 转发请求 - 根据账号平台分流
|
||
c.Set("parsed_request", attemptParsedReq)
|
||
var result *service.ForwardResult
|
||
requestCtx := c.Request.Context()
|
||
if fs.SwitchCount > 0 {
|
||
requestCtx = service.WithAccountSwitchCount(requestCtx, fs.SwitchCount, h.metadataBridgeEnabled())
|
||
}
|
||
if fs.ForceCacheBilling {
|
||
requestCtx = service.WithForceCacheBilling(requestCtx)
|
||
}
|
||
// 记录 Forward 前已写入字节数，Forward 后若增加则说明 SSE 内容已发，禁止 failover
|
||
writerSizeBeforeForward := c.Writer.Size()
|
||
if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey {
|
||
result, err = h.antigravityGatewayService.Forward(requestCtx, c, account, attemptBody, hasBoundSession)
|
||
} else {
|
||
result, err = h.gatewayService.Forward(requestCtx, c, account, attemptParsedReq)
|
||
}
|
||
|
||
// 兜底释放串行锁（正常情况已通过回调提前释放）
|
||
if queueRelease != nil {
|
||
queueRelease()
|
||
}
|
||
// 清理回调引用，防止 failover 重试时旧回调被错误调用
|
||
attemptParsedReq.OnUpstreamAccepted = nil
|
||
|
||
if accountReleaseFunc != nil {
|
||
accountReleaseFunc()
|
||
}
|
||
|
||
// 提交 usage 记录。成功路径与"流中断但 Forward 已观测到 usage 的部分结果"
|
||
// 错误路径共用：后者若不入账，上游已计量的请求会完全漏记漏计费（#5148）。
|
||
submitForwardUsage := func(result *service.ForwardResult) {
|
||
// 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context）
|
||
userAgent := c.GetHeader("User-Agent")
|
||
clientIP := ip.GetClientIP(c)
|
||
// Forward 内部可能继续改写 body，usage 去重指纹必须使用最终上游接受的当前 body。
|
||
requestPayloadHash := service.HashUsageRequestPayload(attemptParsedReq.Body.Bytes())
|
||
inboundEndpoint := GetInboundEndpoint(c)
|
||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||
|
||
if result.ReasoningEffort == nil {
|
||
result.ReasoningEffort = service.NormalizeClaudeOutputEffort(attemptParsedReq.OutputEffort)
|
||
}
|
||
// 同上（重试路径中的对称填充）。详见非重试路径同名注释。
|
||
if result.ReasoningEffort == nil && attemptParsedReq.ThinkingEnabled {
|
||
protocolModel := result.UpstreamModel
|
||
if protocolModel == "" {
|
||
protocolModel = result.Model
|
||
}
|
||
result.ReasoningEffort = service.DefaultEffortForThinkingEnabled(protocolModel)
|
||
}
|
||
|
||
// 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。
|
||
// ForceCacheBilling 提前拍成标量，避免 worker 闭包保活 failover 状态里的响应体。
|
||
forceCacheBilling := fs.ForceCacheBilling
|
||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), currentAPIKey)
|
||
sessionID := service.ExtractClientSessionID(c)
|
||
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||
Result: result,
|
||
QuotaPlatform: quotaPlatform,
|
||
APIKey: currentAPIKey,
|
||
User: currentAPIKey.User,
|
||
Account: account,
|
||
Subscription: currentSubscription,
|
||
PricingAt: pricingAt,
|
||
InboundEndpoint: inboundEndpoint,
|
||
UpstreamEndpoint: upstreamEndpoint,
|
||
UserAgent: userAgent,
|
||
IPAddress: clientIP,
|
||
SessionID: sessionID,
|
||
RequestPayloadHash: requestPayloadHash,
|
||
ForceCacheBilling: forceCacheBilling,
|
||
APIKeyService: h.apiKeyService,
|
||
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
|
||
}); err != nil {
|
||
logger.L().With(
|
||
zap.String("component", "handler.gateway.messages"),
|
||
zap.Int64("user_id", subject.UserID),
|
||
zap.Int64("api_key_id", currentAPIKey.ID),
|
||
zap.Any("group_id", currentAPIKey.GroupID),
|
||
zap.String("model", reqModel),
|
||
zap.Int64("account_id", account.ID),
|
||
).Error("gateway.record_usage_failed", zap.Error(err))
|
||
}
|
||
})
|
||
}
|
||
|
||
if err != nil {
|
||
// Beta policy block: return 400 immediately, no failover
|
||
var betaBlockedErr *service.BetaBlockedError
|
||
if errors.As(err, &betaBlockedErr) {
|
||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalPolicyDenied)
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", betaBlockedErr.Message)
|
||
return
|
||
}
|
||
|
||
var promptTooLongErr *service.PromptTooLongError
|
||
if errors.As(err, &promptTooLongErr) {
|
||
reqLog.Warn("gateway.prompt_too_long_from_antigravity",
|
||
zap.Any("current_group_id", currentAPIKey.GroupID),
|
||
zap.Any("fallback_group_id", fallbackGroupID),
|
||
zap.Bool("fallback_used", fallbackUsed),
|
||
)
|
||
if !fallbackUsed && fallbackGroupID != nil && *fallbackGroupID > 0 {
|
||
fallbackGroup, err := h.gatewayService.ResolveGroupByID(c.Request.Context(), *fallbackGroupID)
|
||
if err != nil {
|
||
reqLog.Warn("gateway.resolve_fallback_group_failed", zap.Int64("fallback_group_id", *fallbackGroupID), zap.Error(err))
|
||
_ = h.antigravityGatewayService.WriteMappedClaudeError(c, account, promptTooLongErr.StatusCode, promptTooLongErr.RequestID, promptTooLongErr.Body)
|
||
return
|
||
}
|
||
if fallbackGroup.Platform != service.PlatformAnthropic ||
|
||
fallbackGroup.SubscriptionType == service.SubscriptionTypeSubscription ||
|
||
fallbackGroup.FallbackGroupIDOnInvalidRequest != nil {
|
||
reqLog.Warn("gateway.fallback_group_invalid",
|
||
zap.Int64("fallback_group_id", fallbackGroup.ID),
|
||
zap.String("fallback_platform", fallbackGroup.Platform),
|
||
zap.String("fallback_subscription_type", fallbackGroup.SubscriptionType),
|
||
)
|
||
_ = h.antigravityGatewayService.WriteMappedClaudeError(c, account, promptTooLongErr.StatusCode, promptTooLongErr.RequestID, promptTooLongErr.Body)
|
||
return
|
||
}
|
||
fallbackAPIKey := cloneAPIKeyWithGroup(apiKey, fallbackGroup)
|
||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), fallbackAPIKey.User, fallbackAPIKey, fallbackGroup, nil, service.PlatformFromAPIKey(fallbackAPIKey)); err != nil {
|
||
status, code, message, retryAfter := billingErrorDetails(err)
|
||
if retryAfter > 0 {
|
||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||
}
|
||
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
||
return
|
||
}
|
||
// 兜底重试按"直接请求兜底分组"处理：清除强制平台，允许按分组平台调度
|
||
ctx := context.WithValue(c.Request.Context(), ctxkey.ForcePlatform, "")
|
||
c.Request = c.Request.WithContext(ctx)
|
||
currentAPIKey = fallbackAPIKey
|
||
currentSubscription = nil
|
||
fallbackUsed = true
|
||
retryWithFallback = true
|
||
break
|
||
}
|
||
_ = h.antigravityGatewayService.WriteMappedClaudeError(c, account, promptTooLongErr.StatusCode, promptTooLongErr.RequestID, promptTooLongErr.Body)
|
||
return
|
||
}
|
||
var failoverErr *service.UpstreamFailoverError
|
||
if errors.As(err, &failoverErr) {
|
||
// 流式内容已写入客户端，无法撤销，禁止 failover 以防止流拼接腐化
|
||
if c.Writer.Size() != writerSizeBeforeForward {
|
||
h.handleFailoverExhausted(c, failoverErr, account.Platform, true)
|
||
return
|
||
}
|
||
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
|
||
switch action {
|
||
case FailoverContinue:
|
||
continue
|
||
case FailoverExhausted:
|
||
h.handleFailoverExhausted(c, fs.LastFailoverErr, account.Platform, streamStarted)
|
||
return
|
||
case FailoverCanceled:
|
||
failoverClientGone(c)
|
||
return
|
||
}
|
||
}
|
||
upstreamErrorAlreadyCommunicated := gatewayForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err)
|
||
wroteFallback := false
|
||
if !upstreamErrorAlreadyCommunicated {
|
||
wroteFallback = h.ensureForwardErrorResponse(c, streamStarted)
|
||
}
|
||
forwardFailedFields := []zap.Field{
|
||
zap.Int64("account_id", account.ID),
|
||
zap.String("account_name", account.Name),
|
||
zap.String("account_platform", account.Platform),
|
||
zap.Bool("fallback_error_response_written", wroteFallback),
|
||
zap.Bool("upstream_error_response_already_written", upstreamErrorAlreadyCommunicated),
|
||
zap.Error(err),
|
||
}
|
||
if account.Proxy != nil {
|
||
forwardFailedFields = append(forwardFailedFields,
|
||
zap.Int64("proxy_id", account.Proxy.ID),
|
||
zap.String("proxy_name", account.Proxy.Name),
|
||
zap.String("proxy_host", account.Proxy.Host),
|
||
zap.Int("proxy_port", account.Proxy.Port),
|
||
)
|
||
} else if account.ProxyID != nil {
|
||
forwardFailedFields = append(forwardFailedFields, zap.Int64p("proxy_id", account.ProxyID))
|
||
}
|
||
reqLog.Error("gateway.forward_failed", forwardFailedFields...)
|
||
// Forward 与错误一起返回的部分结果：流中断前上游已计量的 usage 照常入账，
|
||
// 避免上游已产生消耗的请求完全漏记（#5148）。failover 错误恒定 result=nil，
|
||
// 不会走到这里重复计费。
|
||
if result != nil {
|
||
submitForwardUsage(result)
|
||
}
|
||
return
|
||
}
|
||
|
||
// RPM 计数递增（Forward 成功后）
|
||
// 注意：TOCTOU 竞态是已知且可接受的设计权衡，与 WindowCost 一致的 soft-limit 模式。
|
||
// 在高并发下可能短暂超出 RPM 限制，但不会导致请求失败。
|
||
if account.IsAnthropicOAuthOrSetupToken() && account.GetBaseRPM() > 0 {
|
||
if err := h.gatewayService.IncrementAccountRPM(c.Request.Context(), account.ID); err != nil {
|
||
reqLog.Warn("gateway.rpm_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||
}
|
||
}
|
||
|
||
// 绑定粘性会话（成功转发后绑定/刷新）
|
||
// - 无现有绑定（首次请求）：创建绑定
|
||
// - 选中账号与粘性账号一致：刷新 TTL
|
||
// - 粘性账号因负载/RPM 被跳过、选中了其他账号：不覆盖原绑定，
|
||
// 下次请求粘性账号恢复后仍可命中
|
||
if sessionKey != "" && (sessionBoundAccountID == 0 || sessionBoundAccountID == account.ID) {
|
||
if err := h.gatewayService.BindStickySession(c.Request.Context(), currentAPIKey.GroupID, sessionKey, account.ID); err != nil {
|
||
reqLog.Warn("gateway.bind_sticky_session_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||
}
|
||
}
|
||
|
||
submitForwardUsage(result)
|
||
return
|
||
}
|
||
if !retryWithFallback {
|
||
return
|
||
}
|
||
}
|
||
}
|
||
|
||
// Models handles listing available models
|
||
// GET /v1/models
|
||
// Returns models based on account configurations (model_mapping whitelist)
|
||
// Falls back to default models if no whitelist is configured
|
||
func (h *GatewayHandler) Models(c *gin.Context) {
|
||
apiKey, _ := middleware2.GetAPIKeyFromContext(c)
|
||
|
||
var groupID *int64
|
||
var platform string
|
||
|
||
if apiKey != nil && apiKey.Group != nil {
|
||
groupID = &apiKey.Group.ID
|
||
platform = apiKey.Group.Platform
|
||
}
|
||
if forcedPlatform, ok := middleware2.GetForcePlatformFromContext(c); ok && strings.TrimSpace(forcedPlatform) != "" {
|
||
platform = forcedPlatform
|
||
}
|
||
|
||
if platform == service.PlatformComposite {
|
||
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID)
|
||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.CustomModelsListEnabled() {
|
||
availableModels = filterModelsByCustomList(availableModels, defaultModelIDsForPlatform(service.PlatformComposite), apiKey.Group.ModelsListConfig.Models)
|
||
writeCustomModelsList(c, service.PlatformComposite, availableModels)
|
||
return
|
||
}
|
||
if len(availableModels) > 0 {
|
||
writeModelsList(c, service.PlatformComposite, availableModels)
|
||
return
|
||
}
|
||
writeModelsList(c, service.PlatformComposite, defaultModelIDsForPlatform(service.PlatformComposite))
|
||
return
|
||
}
|
||
|
||
// Get available models from account configurations for the selected group platform.
|
||
availableModels := h.gatewayService.GetAvailableModels(c.Request.Context(), groupID, platform)
|
||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.CustomModelsListEnabled() {
|
||
fallbackModels := defaultModelIDsForPlatform(platform)
|
||
availableModels = filterModelsByCustomList(customModelsListSource(platform, availableModels, fallbackModels), fallbackModels, apiKey.Group.ModelsListConfig.Models)
|
||
writeCustomModelsList(c, platform, availableModels)
|
||
return
|
||
}
|
||
|
||
if len(availableModels) > 0 {
|
||
writeModelsList(c, platform, availableModels)
|
||
return
|
||
}
|
||
|
||
// Fallback to default models
|
||
if platform == service.PlatformOpenAI {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "list",
|
||
"data": openai.DefaultModels,
|
||
})
|
||
return
|
||
}
|
||
|
||
if platform == service.PlatformGemini {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "list",
|
||
"data": geminicli.DefaultModels,
|
||
})
|
||
return
|
||
}
|
||
if platform == service.PlatformGrok {
|
||
writeGrokModelsList(c, xai.DefaultModelIDs())
|
||
return
|
||
}
|
||
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "list",
|
||
"data": claude.DefaultModels,
|
||
})
|
||
}
|
||
|
||
func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64) []string {
|
||
if h == nil || h.gatewayService == nil {
|
||
return nil
|
||
}
|
||
seen := make(map[string]struct{})
|
||
models := make([]string, 0)
|
||
schedulablePlatforms := h.gatewayService.GetSchedulablePlatforms(ctx, groupID)
|
||
for _, platform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek} {
|
||
platformModels := h.gatewayService.GetAvailableModels(ctx, groupID, platform)
|
||
if len(platformModels) == 0 {
|
||
// CN 供应商没有静态默认模型列表（defaultModelIDsForPlatform 的
|
||
// default 分支是 Claude 列表），composite 下只暴露账号映射键。
|
||
if _, ok := schedulablePlatforms[platform]; ok && !service.IsCNProvider(platform) {
|
||
platformModels = defaultModelIDsForPlatform(platform)
|
||
}
|
||
}
|
||
for _, model := range platformModels {
|
||
model = strings.TrimSpace(model)
|
||
if model == "" {
|
||
continue
|
||
}
|
||
if _, ok := seen[model]; ok {
|
||
continue
|
||
}
|
||
seen[model] = struct{}{}
|
||
models = append(models, model)
|
||
}
|
||
}
|
||
return models
|
||
}
|
||
|
||
func writeModelsList(c *gin.Context, platform string, modelIDs []string) {
|
||
if platform == service.PlatformGrok {
|
||
writeGrokModelsList(c, modelIDs)
|
||
return
|
||
}
|
||
models := make([]claude.Model, 0, len(modelIDs))
|
||
for _, modelID := range modelIDs {
|
||
models = append(models, claude.Model{
|
||
ID: modelID,
|
||
Type: "model",
|
||
DisplayName: modelID,
|
||
CreatedAt: "2024-01-01T00:00:00Z",
|
||
})
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "list",
|
||
"data": models,
|
||
})
|
||
}
|
||
|
||
func writeCustomModelsList(c *gin.Context, platform string, modelIDs []string) {
|
||
if platform == service.PlatformOpenAI {
|
||
writeOpenAIModelsList(c, modelIDs)
|
||
return
|
||
}
|
||
writeModelsList(c, platform, modelIDs)
|
||
}
|
||
|
||
type grokReasoningEffortOption struct {
|
||
Value string `json:"value"`
|
||
Label string `json:"label"`
|
||
Default bool `json:"default,omitempty"`
|
||
}
|
||
|
||
type grokModelListItem struct {
|
||
xai.Model
|
||
SupportsReasoningEffort bool `json:"supportsReasoningEffort,omitempty"`
|
||
ReasoningEffort string `json:"reasoningEffort,omitempty"`
|
||
ReasoningEfforts []grokReasoningEffortOption `json:"reasoningEfforts,omitempty"`
|
||
}
|
||
|
||
func writeGrokModelsList(c *gin.Context, modelIDs []string) {
|
||
defaults := xai.DefaultModels()
|
||
defaultsByID := make(map[string]xai.Model, len(defaults))
|
||
for _, model := range defaults {
|
||
defaultsByID[model.ID] = model
|
||
}
|
||
|
||
models := make([]grokModelListItem, 0, len(modelIDs))
|
||
for _, modelID := range modelIDs {
|
||
model, ok := defaultsByID[modelID]
|
||
if !ok {
|
||
model = xai.Model{
|
||
ID: modelID,
|
||
Object: "model",
|
||
OwnedBy: "xai",
|
||
DisplayName: modelID,
|
||
}
|
||
}
|
||
item := grokModelListItem{Model: model}
|
||
if grokModelSupportsConfigurableReasoning(modelID) {
|
||
item.SupportsReasoningEffort = true
|
||
item.ReasoningEffort = "high"
|
||
item.ReasoningEfforts = []grokReasoningEffortOption{
|
||
{Value: "low", Label: "Low"},
|
||
{Value: "medium", Label: "Medium"},
|
||
{Value: "high", Label: "High", Default: true},
|
||
}
|
||
}
|
||
models = append(models, item)
|
||
}
|
||
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "list",
|
||
"data": models,
|
||
})
|
||
}
|
||
|
||
func grokModelSupportsConfigurableReasoning(modelID string) bool {
|
||
switch strings.ToLower(strings.TrimSpace(modelID)) {
|
||
case "grok-4.6", "grok-4.6-latest", "grok-4.5", "grok-4.5-latest", "grok", "grok-latest", "grok-build", "grok-build-latest", "grok-build-0.1":
|
||
return true
|
||
default:
|
||
return false
|
||
}
|
||
}
|
||
|
||
func writeOpenAIModelsList(c *gin.Context, modelIDs []string) {
|
||
defaultsByID := make(map[string]openai.Model, len(openai.DefaultModels))
|
||
for _, model := range openai.DefaultModels {
|
||
defaultsByID[model.ID] = model
|
||
}
|
||
|
||
models := make([]openai.Model, 0, len(modelIDs))
|
||
for _, modelID := range modelIDs {
|
||
if model, ok := defaultsByID[modelID]; ok {
|
||
models = append(models, model)
|
||
continue
|
||
}
|
||
models = append(models, openai.Model{
|
||
ID: modelID,
|
||
Object: "model",
|
||
Created: 1704067200,
|
||
OwnedBy: "openai",
|
||
Type: "model",
|
||
DisplayName: modelID,
|
||
})
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "list",
|
||
"data": models,
|
||
})
|
||
}
|
||
|
||
func customModelsListSource(platform string, availableModels, fallbackModels []string) []string {
|
||
if platform == service.PlatformAnthropic && len(availableModels) > 0 {
|
||
return mergeModelIDs(availableModels, fallbackModels)
|
||
}
|
||
return availableModels
|
||
}
|
||
|
||
func filterModelsByCustomList(availableModels, fallbackModels, selectedModels []string) []string {
|
||
if len(selectedModels) == 0 {
|
||
return availableModels
|
||
}
|
||
source := availableModels
|
||
if len(source) == 0 {
|
||
source = fallbackModels
|
||
}
|
||
if len(source) == 0 {
|
||
return nil
|
||
}
|
||
|
||
allowed := make([]string, 0, len(source))
|
||
for _, model := range source {
|
||
model = strings.TrimSpace(model)
|
||
if model != "" {
|
||
allowed = append(allowed, model)
|
||
}
|
||
}
|
||
|
||
seen := make(map[string]struct{}, len(selectedModels))
|
||
filtered := make([]string, 0, len(selectedModels))
|
||
for _, model := range selectedModels {
|
||
model = strings.TrimSpace(model)
|
||
if model == "" {
|
||
continue
|
||
}
|
||
if !customModelsListAllowsModel(allowed, model) {
|
||
continue
|
||
}
|
||
if _, ok := seen[model]; ok {
|
||
continue
|
||
}
|
||
seen[model] = struct{}{}
|
||
filtered = append(filtered, model)
|
||
}
|
||
return filtered
|
||
}
|
||
|
||
func customModelsListAllowsModel(availablePatterns []string, model string) bool {
|
||
for _, pattern := range availablePatterns {
|
||
if pattern == model {
|
||
return true
|
||
}
|
||
if strings.HasSuffix(pattern, "*") && strings.HasPrefix(model, strings.TrimSuffix(pattern, "*")) {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
func defaultModelIDsForPlatform(platform string) []string {
|
||
switch platform {
|
||
case service.PlatformOpenAI:
|
||
return openai.DefaultModelIDs()
|
||
case service.PlatformGemini:
|
||
ids := make([]string, 0, len(geminicli.DefaultModels))
|
||
for _, model := range geminicli.DefaultModels {
|
||
ids = append(ids, model.ID)
|
||
}
|
||
return ids
|
||
case service.PlatformAntigravity:
|
||
models := antigravity.DefaultModels()
|
||
ids := make([]string, 0, len(models))
|
||
for _, model := range models {
|
||
ids = append(ids, model.ID)
|
||
}
|
||
return ids
|
||
case service.PlatformAnthropic:
|
||
ids := make([]string, 0, len(claude.DefaultModels)+len(antigravity.DefaultModels()))
|
||
for _, model := range claude.DefaultModels {
|
||
ids = append(ids, model.ID)
|
||
}
|
||
for _, model := range antigravity.DefaultModels() {
|
||
ids = append(ids, model.ID)
|
||
}
|
||
return mergeModelIDs(ids, nil)
|
||
case service.PlatformGrok:
|
||
return xai.DefaultModelIDs()
|
||
case service.PlatformComposite:
|
||
ids := make([]string, 0)
|
||
seen := make(map[string]struct{})
|
||
for _, concretePlatform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek} {
|
||
for _, id := range defaultModelIDsForPlatform(concretePlatform) {
|
||
if _, ok := seen[id]; ok {
|
||
continue
|
||
}
|
||
seen[id] = struct{}{}
|
||
ids = append(ids, id)
|
||
}
|
||
}
|
||
return ids
|
||
default:
|
||
ids := make([]string, 0, len(claude.DefaultModels))
|
||
for _, model := range claude.DefaultModels {
|
||
ids = append(ids, model.ID)
|
||
}
|
||
return ids
|
||
}
|
||
}
|
||
|
||
func mergeModelIDs(primary, secondary []string) []string {
|
||
seen := make(map[string]struct{}, len(primary)+len(secondary))
|
||
merged := make([]string, 0, len(primary)+len(secondary))
|
||
for _, models := range [][]string{primary, secondary} {
|
||
for _, model := range models {
|
||
model = strings.TrimSpace(model)
|
||
if model == "" {
|
||
continue
|
||
}
|
||
if _, ok := seen[model]; ok {
|
||
continue
|
||
}
|
||
seen[model] = struct{}{}
|
||
merged = append(merged, model)
|
||
}
|
||
}
|
||
return merged
|
||
}
|
||
|
||
// AntigravityModels 返回 Antigravity 支持的全部模型
|
||
// GET /antigravity/models
|
||
func (h *GatewayHandler) AntigravityModels(c *gin.Context) {
|
||
c.JSON(http.StatusOK, gin.H{
|
||
"object": "list",
|
||
"data": antigravity.DefaultModels(),
|
||
})
|
||
}
|
||
|
||
func cloneAPIKeyWithGroup(apiKey *service.APIKey, group *service.Group) *service.APIKey {
|
||
if apiKey == nil || group == nil {
|
||
return apiKey
|
||
}
|
||
cloned := *apiKey
|
||
groupID := group.ID
|
||
cloned.GroupID = &groupID
|
||
cloned.Group = group
|
||
return &cloned
|
||
}
|
||
|
||
// Usage handles getting account balance and usage statistics for CC Switch integration
|
||
// GET /v1/usage
|
||
//
|
||
// Two modes:
|
||
// - quota_limited: API Key has quota or rate limits configured. Returns key-level limits/usage.
|
||
// - unrestricted: No key-level limits. Returns subscription or wallet balance info.
|
||
func (h *GatewayHandler) Usage(c *gin.Context) {
|
||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||
return
|
||
}
|
||
|
||
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||
return
|
||
}
|
||
|
||
ctx := c.Request.Context()
|
||
|
||
// 解析可选的日期范围参数（用于 model_stats 查询）
|
||
startTime, endTime := h.parseUsageDateRange(c)
|
||
days, ok := parseAPIKeyDailyUsageDays(c.DefaultQuery("days", ""))
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Invalid days, allowed range is 1-90")
|
||
return
|
||
}
|
||
|
||
// Best-effort: 获取用量统计（按当前 API Key 过滤），失败不影响基础响应
|
||
usageData := h.buildUsageData(ctx, apiKey.ID)
|
||
dailyUsage := h.buildAPIKeyDailyUsage(c, subject.UserID, apiKey.ID, days)
|
||
|
||
// Best-effort: 获取模型统计
|
||
var modelStats any
|
||
if h.usageService != nil {
|
||
if stats, err := h.usageService.GetAPIKeyModelStats(ctx, apiKey.ID, startTime, endTime); err == nil && len(stats) > 0 {
|
||
modelStats = stats
|
||
}
|
||
}
|
||
|
||
// 判断模式: key 有总额度或速率限制 → quota_limited，否则 → unrestricted
|
||
isQuotaLimited := apiKey.Quota > 0 || apiKey.HasRateLimits()
|
||
|
||
if isQuotaLimited {
|
||
h.usageQuotaLimited(c, ctx, apiKey, usageData, dailyUsage, modelStats)
|
||
return
|
||
}
|
||
|
||
h.usageUnrestricted(c, ctx, apiKey, subject, usageData, dailyUsage, modelStats)
|
||
}
|
||
|
||
// parseUsageDateRange 解析 start_date / end_date query params，默认返回近 30 天范围
|
||
func (h *GatewayHandler) parseUsageDateRange(c *gin.Context) (time.Time, time.Time) {
|
||
now := timezone.Now()
|
||
endTime := now
|
||
startTime := now.AddDate(0, 0, -30)
|
||
|
||
if s := c.Query("start_date"); s != "" {
|
||
if t, err := timezone.ParseInLocation("2006-01-02", s); err == nil {
|
||
startTime = t
|
||
}
|
||
}
|
||
if s := c.Query("end_date"); s != "" {
|
||
if t, err := timezone.ParseInLocation("2006-01-02", s); err == nil {
|
||
endTime = t.AddDate(0, 0, 1) // half-open range upper bound
|
||
}
|
||
}
|
||
return startTime, endTime
|
||
}
|
||
|
||
// buildUsageData 构建 today/total 用量摘要
|
||
func (h *GatewayHandler) buildUsageData(ctx context.Context, apiKeyID int64) gin.H {
|
||
if h.usageService == nil {
|
||
return nil
|
||
}
|
||
dashStats, err := h.usageService.GetAPIKeyDashboardStats(ctx, apiKeyID)
|
||
if err != nil || dashStats == nil {
|
||
return nil
|
||
}
|
||
return gin.H{
|
||
"today": gin.H{
|
||
"requests": dashStats.TodayRequests,
|
||
"input_tokens": dashStats.TodayInputTokens,
|
||
"output_tokens": dashStats.TodayOutputTokens,
|
||
"cache_creation_tokens": dashStats.TodayCacheCreationTokens,
|
||
"cache_read_tokens": dashStats.TodayCacheReadTokens,
|
||
"total_tokens": dashStats.TodayTokens,
|
||
"cost": dashStats.TodayCost,
|
||
"actual_cost": dashStats.TodayActualCost,
|
||
},
|
||
"total": gin.H{
|
||
"requests": dashStats.TotalRequests,
|
||
"input_tokens": dashStats.TotalInputTokens,
|
||
"output_tokens": dashStats.TotalOutputTokens,
|
||
"cache_creation_tokens": dashStats.TotalCacheCreationTokens,
|
||
"cache_read_tokens": dashStats.TotalCacheReadTokens,
|
||
"total_tokens": dashStats.TotalTokens,
|
||
"cost": dashStats.TotalCost,
|
||
"actual_cost": dashStats.TotalActualCost,
|
||
},
|
||
"average_duration_ms": dashStats.AverageDurationMs,
|
||
"rpm": dashStats.Rpm,
|
||
"tpm": dashStats.Tpm,
|
||
}
|
||
}
|
||
|
||
func (h *GatewayHandler) buildAPIKeyDailyUsage(c *gin.Context, userID, apiKeyID int64, days int) any {
|
||
if h.usageService == nil {
|
||
return nil
|
||
}
|
||
startTime, endTime := apiKeyDailyUsageRange(days, c.Query("timezone"))
|
||
stats, err := h.usageService.GetAPIKeyDailyUsage(c.Request.Context(), userID, apiKeyID, startTime, endTime)
|
||
if err != nil {
|
||
return nil
|
||
}
|
||
return stats
|
||
}
|
||
|
||
// usageQuotaLimited 处理 quota_limited 模式的响应
|
||
func (h *GatewayHandler) usageQuotaLimited(c *gin.Context, ctx context.Context, apiKey *service.APIKey, usageData gin.H, dailyUsage any, modelStats any) {
|
||
resp := gin.H{
|
||
"mode": "quota_limited",
|
||
"isValid": apiKey.Status == service.StatusAPIKeyActive || apiKey.Status == service.StatusAPIKeyQuotaExhausted || apiKey.Status == service.StatusAPIKeyExpired,
|
||
"status": apiKey.Status,
|
||
}
|
||
|
||
// 总额度信息
|
||
if apiKey.Quota > 0 {
|
||
remaining := apiKey.GetQuotaRemaining()
|
||
resp["quota"] = gin.H{
|
||
"limit": apiKey.Quota,
|
||
"used": apiKey.QuotaUsed,
|
||
"remaining": remaining,
|
||
"unit": "USD",
|
||
}
|
||
resp["remaining"] = remaining
|
||
resp["unit"] = "USD"
|
||
}
|
||
|
||
// 速率限制信息（从 DB 获取实时用量）
|
||
if apiKey.HasRateLimits() && h.apiKeyService != nil {
|
||
rateLimitData, err := h.apiKeyService.GetRateLimitData(ctx, apiKey.ID)
|
||
if err == nil && rateLimitData != nil {
|
||
var rateLimits []gin.H
|
||
if apiKey.RateLimit5h > 0 {
|
||
used := rateLimitData.EffectiveUsage5h()
|
||
entry := gin.H{
|
||
"window": "5h",
|
||
"limit": apiKey.RateLimit5h,
|
||
"used": used,
|
||
"remaining": max(0, apiKey.RateLimit5h-used),
|
||
"window_start": rateLimitData.Window5hStart,
|
||
}
|
||
if rateLimitData.Window5hStart != nil && !service.IsWindowExpired(rateLimitData.Window5hStart, service.RateLimitWindow5h) {
|
||
entry["reset_at"] = rateLimitData.Window5hStart.Add(service.RateLimitWindow5h)
|
||
}
|
||
rateLimits = append(rateLimits, entry)
|
||
}
|
||
if apiKey.RateLimit1d > 0 {
|
||
used := rateLimitData.EffectiveUsage1d()
|
||
entry := gin.H{
|
||
"window": "1d",
|
||
"limit": apiKey.RateLimit1d,
|
||
"used": used,
|
||
"remaining": max(0, apiKey.RateLimit1d-used),
|
||
"window_start": rateLimitData.Window1dStart,
|
||
}
|
||
if rateLimitData.Window1dStart != nil && !service.IsWindowExpired(rateLimitData.Window1dStart, service.RateLimitWindow1d) {
|
||
entry["reset_at"] = rateLimitData.Window1dStart.Add(service.RateLimitWindow1d)
|
||
}
|
||
rateLimits = append(rateLimits, entry)
|
||
}
|
||
if apiKey.RateLimit7d > 0 {
|
||
used := rateLimitData.EffectiveUsage7d()
|
||
entry := gin.H{
|
||
"window": "7d",
|
||
"limit": apiKey.RateLimit7d,
|
||
"used": used,
|
||
"remaining": max(0, apiKey.RateLimit7d-used),
|
||
"window_start": rateLimitData.Window7dStart,
|
||
}
|
||
if rateLimitData.Window7dStart != nil && !service.IsWindowExpired(rateLimitData.Window7dStart, service.RateLimitWindow7d) {
|
||
entry["reset_at"] = rateLimitData.Window7dStart.Add(service.RateLimitWindow7d)
|
||
}
|
||
rateLimits = append(rateLimits, entry)
|
||
}
|
||
if len(rateLimits) > 0 {
|
||
resp["rate_limits"] = rateLimits
|
||
}
|
||
}
|
||
}
|
||
|
||
// 过期时间
|
||
if apiKey.ExpiresAt != nil {
|
||
resp["expires_at"] = apiKey.ExpiresAt
|
||
resp["days_until_expiry"] = apiKey.GetDaysUntilExpiry()
|
||
}
|
||
|
||
if usageData != nil {
|
||
resp["usage"] = usageData
|
||
}
|
||
if dailyUsage != nil {
|
||
resp["daily_usage"] = dailyUsage
|
||
}
|
||
if modelStats != nil {
|
||
resp["model_stats"] = modelStats
|
||
}
|
||
|
||
c.JSON(http.StatusOK, resp)
|
||
}
|
||
|
||
// usageUnrestricted 处理 unrestricted 模式的响应（向后兼容）
|
||
func (h *GatewayHandler) usageUnrestricted(c *gin.Context, ctx context.Context, apiKey *service.APIKey, subject middleware2.AuthSubject, usageData gin.H, dailyUsage any, modelStats any) {
|
||
// 订阅模式
|
||
if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() {
|
||
resp := gin.H{
|
||
"mode": "unrestricted",
|
||
"isValid": true,
|
||
"planName": apiKey.Group.Name,
|
||
"unit": "USD",
|
||
}
|
||
|
||
// 订阅信息可能不在 context 中（/v1/usage 路径跳过了中间件的计费检查）
|
||
subscription, ok := middleware2.GetSubscriptionFromContext(c)
|
||
if ok {
|
||
remaining := h.calculateSubscriptionRemaining(apiKey.Group, subscription)
|
||
resp["remaining"] = remaining
|
||
resp["subscription"] = gin.H{
|
||
"daily_usage_usd": subscription.DailyUsageUSD,
|
||
"weekly_usage_usd": subscription.WeeklyUsageUSD,
|
||
"monthly_usage_usd": subscription.MonthlyUsageUSD,
|
||
"daily_limit_usd": apiKey.Group.DailyLimitUSD,
|
||
"weekly_limit_usd": apiKey.Group.WeeklyLimitUSD,
|
||
"monthly_limit_usd": apiKey.Group.MonthlyLimitUSD,
|
||
"weekly_window_start": subscription.WeeklyWindowStart,
|
||
"expires_at": subscription.ExpiresAt,
|
||
}
|
||
}
|
||
|
||
if usageData != nil {
|
||
resp["usage"] = usageData
|
||
}
|
||
if dailyUsage != nil {
|
||
resp["daily_usage"] = dailyUsage
|
||
}
|
||
if modelStats != nil {
|
||
resp["model_stats"] = modelStats
|
||
}
|
||
c.JSON(http.StatusOK, resp)
|
||
return
|
||
}
|
||
|
||
// 余额模式
|
||
latestUser, err := h.userService.GetByID(ctx, subject.UserID)
|
||
if err != nil {
|
||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Failed to get user info")
|
||
return
|
||
}
|
||
|
||
resp := gin.H{
|
||
"mode": "unrestricted",
|
||
"isValid": true,
|
||
"planName": "钱包余额",
|
||
"remaining": latestUser.Balance,
|
||
"unit": "USD",
|
||
"balance": latestUser.Balance,
|
||
}
|
||
if usageData != nil {
|
||
resp["usage"] = usageData
|
||
}
|
||
if dailyUsage != nil {
|
||
resp["daily_usage"] = dailyUsage
|
||
}
|
||
if modelStats != nil {
|
||
resp["model_stats"] = modelStats
|
||
}
|
||
c.JSON(http.StatusOK, resp)
|
||
}
|
||
|
||
// calculateSubscriptionRemaining 计算订阅剩余可用额度
|
||
// 逻辑：
|
||
// 1. 如果日/周/月任一限额达到100%，返回0
|
||
// 2. 否则返回所有已配置周期中剩余额度的最小值
|
||
func (h *GatewayHandler) calculateSubscriptionRemaining(group *service.Group, sub *service.UserSubscription) float64 {
|
||
var remainingValues []float64
|
||
|
||
// 检查日限额
|
||
if group.HasDailyLimit() {
|
||
remaining := *group.DailyLimitUSD - sub.DailyUsageUSD
|
||
if remaining <= 0 {
|
||
return 0
|
||
}
|
||
remainingValues = append(remainingValues, remaining)
|
||
}
|
||
|
||
// 检查周限额
|
||
if group.HasWeeklyLimit() {
|
||
remaining := *group.WeeklyLimitUSD - sub.WeeklyUsageUSD
|
||
if remaining <= 0 {
|
||
return 0
|
||
}
|
||
remainingValues = append(remainingValues, remaining)
|
||
}
|
||
|
||
// 检查月限额
|
||
if group.HasMonthlyLimit() {
|
||
remaining := *group.MonthlyLimitUSD - sub.MonthlyUsageUSD
|
||
if remaining <= 0 {
|
||
return 0
|
||
}
|
||
remainingValues = append(remainingValues, remaining)
|
||
}
|
||
|
||
// 如果没有配置任何限额，返回-1表示无限制
|
||
if len(remainingValues) == 0 {
|
||
return -1
|
||
}
|
||
|
||
// 返回最小值
|
||
min := remainingValues[0]
|
||
for _, v := range remainingValues[1:] {
|
||
if v < min {
|
||
min = v
|
||
}
|
||
}
|
||
return min
|
||
}
|
||
|
||
// handleConcurrencyError handles concurrency-related acquire errors.
|
||
func (h *GatewayHandler) handleConcurrencyError(c *gin.Context, err error, slotType string, streamStarted bool) {
|
||
status, errType, message := concurrencyErrorResponse(err, slotType)
|
||
h.handleStreamingAwareError(c, status, errType, message, streamStarted)
|
||
}
|
||
|
||
func (h *GatewayHandler) handleFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError, platform string, streamStarted bool) {
|
||
statusCode := failoverErr.StatusCode
|
||
responseBody := failoverErr.ResponseBody
|
||
if service.IsOpenAISilentRefusalErrorBody(responseBody) {
|
||
service.SetOpsUpstreamError(c, statusCode, service.OpenAISilentRefusalClientMessage(), "")
|
||
h.handleStreamingAwareError(c, http.StatusBadGateway, "upstream_error", service.OpenAISilentRefusalClientMessage(), streamStarted)
|
||
return
|
||
}
|
||
|
||
// 先检查透传规则
|
||
if h.errorPassthroughService != nil && len(responseBody) > 0 {
|
||
if rule := h.errorPassthroughService.MatchRule(platform, statusCode, responseBody); rule != nil {
|
||
// 确定响应状态码
|
||
respCode := statusCode
|
||
if !rule.PassthroughCode && rule.ResponseCode != nil {
|
||
respCode = *rule.ResponseCode
|
||
}
|
||
|
||
// 确定响应消息
|
||
msg := service.ExtractUpstreamErrorMessage(responseBody)
|
||
if !rule.PassthroughBody && rule.CustomMessage != nil {
|
||
msg = *rule.CustomMessage
|
||
}
|
||
|
||
if rule.SkipMonitoring {
|
||
c.Set(service.OpsSkipPassthroughKey, true)
|
||
}
|
||
|
||
h.handleStreamingAwareError(c, respCode, "upstream_error", msg, streamStarted)
|
||
return
|
||
}
|
||
}
|
||
|
||
// 记录原始上游状态码，以便 ops 错误日志捕获真实的上游错误
|
||
upstreamMsg := service.ExtractUpstreamErrorMessage(responseBody)
|
||
service.SetOpsUpstreamError(c, statusCode, upstreamMsg, "")
|
||
|
||
// 使用默认的错误映射
|
||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||
}
|
||
|
||
// handleFailoverExhaustedSimple 简化版本，用于没有响应体的情况
|
||
func (h *GatewayHandler) handleFailoverExhaustedSimple(c *gin.Context, statusCode int, streamStarted bool) {
|
||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||
service.SetOpsUpstreamError(c, statusCode, errMsg, "")
|
||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||
}
|
||
|
||
func (h *GatewayHandler) mapUpstreamError(statusCode int) (int, string, string) {
|
||
switch statusCode {
|
||
case 401:
|
||
return http.StatusBadGateway, "upstream_error", "Upstream authentication failed, please contact administrator"
|
||
case 403:
|
||
return http.StatusBadGateway, "upstream_error", "Upstream access forbidden, please contact administrator"
|
||
case 429:
|
||
return http.StatusTooManyRequests, "rate_limit_error", "Upstream rate limit exceeded, please retry later"
|
||
case 529:
|
||
return http.StatusServiceUnavailable, "overloaded_error", "Upstream service overloaded, please retry later"
|
||
case 500, 502, 503, 504:
|
||
return http.StatusBadGateway, "upstream_error", "Upstream service temporarily unavailable"
|
||
default:
|
||
return http.StatusBadGateway, "upstream_error", "Upstream request failed"
|
||
}
|
||
}
|
||
|
||
// handleStreamingAwareError handles errors that may occur after streaming has started
|
||
func (h *GatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) {
|
||
if streamStarted {
|
||
// 响应状态码已固化为 200（ping/部分数据已 flush），错误只能就地以 SSE 帧回传。
|
||
// 标记本次流内错误，供 ops_error_logger 补记——否则该中间件按 status>=400 采集，
|
||
// 这类挂在 200 流上的失败（如并发限流回退）不会进错误看板。
|
||
service.MarkOpsStreamError(c, errType, message, status)
|
||
|
||
// /v1/responses 的严格 SDK（Codex CLI）要求终止事件必须属于
|
||
// response.completed/failed/incomplete/cancelled 集合。
|
||
// Anthropic-backed Responses 路径同样会因为通用 error 帧被拒。
|
||
if inboundIsResponses(c) {
|
||
if writeResponsesFailedSSE(c, errType, message) {
|
||
return
|
||
}
|
||
}
|
||
// Stream already started, send error as SSE event then close
|
||
flusher, ok := c.Writer.(http.Flusher)
|
||
if ok {
|
||
// SSE 错误事件固定 schema，使用 Quote 直拼可避免额外 Marshal 分配。
|
||
errorEvent := `data: {"type":"error","error":{"type":` + strconv.Quote(errType) + `,"message":` + strconv.Quote(message) + `}}` + "\n\n"
|
||
if _, err := fmt.Fprint(c.Writer, errorEvent); err != nil {
|
||
_ = c.Error(err)
|
||
}
|
||
flusher.Flush()
|
||
}
|
||
return
|
||
}
|
||
|
||
// Normal case: return JSON response with proper status code
|
||
h.errorResponse(c, status, errType, message)
|
||
}
|
||
|
||
// ensureForwardErrorResponse 在 Forward 返回错误但尚未写响应时补写统一错误响应。
|
||
// Writer 已被写过时（ping 已 flush）走 streamStarted 分支，
|
||
// 让 handleStreamingAwareError 通过 SSE 发协议合规的终止事件，
|
||
// 否则下游收到的就是 silent EOF。
|
||
func (h *GatewayHandler) ensureForwardErrorResponse(c *gin.Context, streamStarted bool) bool {
|
||
if c == nil || c.Writer == nil {
|
||
return false
|
||
}
|
||
if service.IsResponseCommitted(c) {
|
||
return false
|
||
}
|
||
if c.Writer.Written() {
|
||
streamStarted = true
|
||
}
|
||
h.handleStreamingAwareError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed", streamStarted)
|
||
return true
|
||
}
|
||
|
||
// gatewayForwardErrorAlreadyCommunicated reports whether a Forward implementation
|
||
// has already written a complete error response to the client before returning
|
||
// an error to the handler.
|
||
//
|
||
// This is intentionally narrower than "writer size changed": a stream may have
|
||
// only emitted keepalive pings or partial data, in which case the handler still
|
||
// needs to append a protocol-level terminal error. Non-SSE output from Forward
|
||
// is different: service-level helpers such as handleErrorResponse/writeClaudeError
|
||
// already wrote the client-visible JSON body, so adding the generic streaming
|
||
// fallback would corrupt the response by appending a second `data: ...` frame.
|
||
func gatewayForwardErrorAlreadyCommunicated(c *gin.Context, writerSizeBeforeForward int, err error) bool {
|
||
if err == nil || c == nil || c.Writer == nil {
|
||
return false
|
||
}
|
||
if c.Writer.Size() == writerSizeBeforeForward {
|
||
return false
|
||
}
|
||
|
||
contentType := strings.ToLower(strings.TrimSpace(c.Writer.Header().Get("Content-Type")))
|
||
if contentType == "" {
|
||
return false
|
||
}
|
||
return !strings.Contains(contentType, "text/event-stream")
|
||
}
|
||
|
||
// checkClaudeCodeVersion 检查 Claude Code 客户端版本是否满足版本要求
|
||
// 仅对已识别的 Claude Code 客户端执行，count_tokens 路径除外
|
||
func (h *GatewayHandler) checkClaudeCodeVersion(c *gin.Context) bool {
|
||
ctx := c.Request.Context()
|
||
if !service.IsClaudeCodeClient(ctx) {
|
||
return true
|
||
}
|
||
|
||
// 排除 count_tokens 子路径
|
||
if strings.HasSuffix(c.Request.URL.Path, "/count_tokens") {
|
||
return true
|
||
}
|
||
|
||
minVersion, maxVersion := h.settingService.GetClaudeCodeVersionBounds(ctx)
|
||
if minVersion == "" && maxVersion == "" {
|
||
return true // 未设置，不检查
|
||
}
|
||
|
||
clientVersion := service.GetClaudeCodeVersion(ctx)
|
||
if clientVersion == "" {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
||
"Unable to determine Claude Code version. Please update Claude Code: npm update -g @anthropic-ai/claude-code")
|
||
return false
|
||
}
|
||
|
||
if minVersion != "" && service.CompareVersions(clientVersion, minVersion) < 0 {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
||
fmt.Sprintf("Your Claude Code version (%s) is below the minimum required version (%s). Please update: npm update -g @anthropic-ai/claude-code",
|
||
clientVersion, minVersion))
|
||
return false
|
||
}
|
||
|
||
if maxVersion != "" && service.CompareVersions(clientVersion, maxVersion) > 0 {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
||
fmt.Sprintf("Your Claude Code version (%s) exceeds the maximum allowed version (%s). "+
|
||
"Please downgrade: npm install -g @anthropic-ai/claude-code@%s && "+
|
||
"set CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 to prevent auto-upgrade",
|
||
clientVersion, maxVersion, maxVersion))
|
||
return false
|
||
}
|
||
|
||
return true
|
||
}
|
||
|
||
// errorResponse 返回Claude API格式的错误响应
|
||
func (h *GatewayHandler) errorResponse(c *gin.Context, status int, errType, message string) {
|
||
c.JSON(status, gin.H{
|
||
"type": "error",
|
||
"error": gin.H{
|
||
"type": errType,
|
||
"message": message,
|
||
},
|
||
})
|
||
}
|
||
|
||
// CountTokens handles token counting endpoint
|
||
// POST /v1/messages/count_tokens
|
||
// 特点：校验订阅/余额，但不计算并发、不记录使用量
|
||
func (h *GatewayHandler) CountTokens(c *gin.Context) {
|
||
// 从context获取apiKey和user（ApiKeyAuth中间件已设置）
|
||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
||
return
|
||
}
|
||
|
||
_, ok = middleware2.GetAuthSubjectFromContext(c)
|
||
if !ok {
|
||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
|
||
return
|
||
}
|
||
reqLog := requestLogger(
|
||
c,
|
||
"handler.gateway.count_tokens",
|
||
zap.Int64("api_key_id", apiKey.ID),
|
||
zap.Any("group_id", apiKey.GroupID),
|
||
)
|
||
defer h.maybeLogCompatibilityFallbackMetrics(reqLog)
|
||
|
||
// 读取请求体
|
||
body, err := readLenientJSONRequestBodyWithPrealloc(c.Request, h.cfg)
|
||
if err != nil {
|
||
if maxErr, ok := extractMaxBytesError(err); ok {
|
||
h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
|
||
return
|
||
}
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
|
||
return
|
||
}
|
||
|
||
if len(body) == 0 {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
|
||
return
|
||
}
|
||
|
||
setOpsRequestContext(c, "", false)
|
||
|
||
bodyRef := service.NewRequestBodyRef(body)
|
||
parsedReq, err := service.ParseGatewayRequest(bodyRef, domain.PlatformAnthropic)
|
||
if err != nil {
|
||
logRequestBodyParseFailure(reqLog, body, err)
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||
return
|
||
}
|
||
body = parsedReq.Body.Bytes()
|
||
// count_tokens 走 messages 严格校验时，复用已解析请求，避免二次反序列化。
|
||
SetClaudeCodeClientContext(c, body, parsedReq)
|
||
ensureCompositeTargetPlatform(c, apiKey, parsedReq.Model)
|
||
reqLog = reqLog.With(zap.String("model", parsedReq.Model), zap.Bool("stream", parsedReq.Stream))
|
||
// 在请求上下文中记录 thinking 状态，供 Antigravity 最终模型 key 推导/模型维度限流使用
|
||
c.Request = c.Request.WithContext(service.WithThinkingEnabled(c.Request.Context(), parsedReq.ThinkingEnabled, h.metadataBridgeEnabled()))
|
||
|
||
// 验证 model 必填
|
||
if parsedReq.Model == "" {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||
return
|
||
}
|
||
if !compositeTargetPlatformResolved(c, apiKey, parsedReq.Model) {
|
||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Model is not supported by composite groups")
|
||
return
|
||
}
|
||
|
||
setOpsRequestContext(c, parsedReq.Model, parsedReq.Stream)
|
||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(parsedReq.Stream, false)))
|
||
|
||
// 获取订阅信息（可能为nil）
|
||
subscription, _ := middleware2.GetSubscriptionFromContext(c)
|
||
|
||
// 校验 billing eligibility（订阅/余额）
|
||
// 【注意】不计算并发，但需要校验订阅/余额
|
||
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil {
|
||
status, code, message, retryAfter := billingErrorDetails(err)
|
||
if retryAfter > 0 {
|
||
c.Header("Retry-After", strconv.Itoa(retryAfter))
|
||
}
|
||
h.errorResponse(c, status, code, message)
|
||
return
|
||
}
|
||
|
||
// 计算粘性会话 hash
|
||
parsedReq.SessionContext = &service.SessionContext{
|
||
ClientIP: ip.GetClientIP(c),
|
||
UserAgent: c.GetHeader("User-Agent"),
|
||
APIKeyID: apiKey.ID,
|
||
}
|
||
sessionHash := h.gatewayService.GenerateSessionHash(parsedReq)
|
||
|
||
// 选择支持该模型的账号
|
||
account, err := h.gatewayService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, sessionHash, parsedReq.Model)
|
||
if err != nil {
|
||
reqLog.Warn("gateway.count_tokens_select_account_failed", zap.Error(err))
|
||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, parsedReq.Model, parsedReq.Model, service.PlatformAnthropic)
|
||
if !cls.ModelNotFound {
|
||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||
}
|
||
h.errorResponse(c, cls.Status, cls.ErrType, cls.Message)
|
||
return
|
||
}
|
||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||
|
||
// 转发请求（不记录使用量）
|
||
if err := h.gatewayService.ForwardCountTokens(c.Request.Context(), c, account, parsedReq); err != nil {
|
||
reqLog.Error("gateway.count_tokens_forward_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
||
// 错误响应已在 ForwardCountTokens 中处理
|
||
return
|
||
}
|
||
}
|
||
|
||
// InterceptType 表示请求拦截类型
|
||
type InterceptType int
|
||
|
||
const (
|
||
InterceptTypeNone InterceptType = iota
|
||
InterceptTypeWarmup // 预热请求（返回 "New Conversation"）
|
||
InterceptTypeSuggestionMode // SUGGESTION MODE（返回空字符串）
|
||
InterceptTypeMaxTokensOneHaiku // max_tokens=1 + haiku 探测请求（返回 "#"）
|
||
)
|
||
|
||
// isHaikuModel 检查模型名称是否包含 "haiku"（大小写不敏感）
|
||
func isHaikuModel(model string) bool {
|
||
return strings.Contains(strings.ToLower(model), "haiku")
|
||
}
|
||
|
||
// isMaxTokensOneHaikuRequest 检查是否为 max_tokens=1 + haiku 模型的探测请求
|
||
// 这类请求用于 Claude Code 验证 API 连通性（流式/非流式均会出现，如 cc-switch v3.9.0 起的健康检查探测为流式）
|
||
// 条件：max_tokens == 1 且 model 包含 "haiku"
|
||
func isMaxTokensOneHaikuRequest(model string, maxTokens int) bool {
|
||
return maxTokens == 1 && isHaikuModel(model)
|
||
}
|
||
|
||
// detectInterceptType 检测请求是否需要拦截，返回拦截类型
|
||
// 参数说明：
|
||
// - body: 请求体字节
|
||
// - model: 请求的模型名称
|
||
// - maxTokens: max_tokens 值
|
||
// - isClaudeCodeClient: 是否已通过 Claude Code 客户端校验
|
||
func detectInterceptType(body []byte, model string, maxTokens int, isClaudeCodeClient bool) InterceptType {
|
||
// 优先检查 max_tokens=1 + haiku 探测请求（流式/非流式均适用）
|
||
if isClaudeCodeClient && isMaxTokensOneHaikuRequest(model, maxTokens) {
|
||
return InterceptTypeMaxTokensOneHaiku
|
||
}
|
||
|
||
// 快速检查：如果不包含任何关键字，直接返回
|
||
bodyStr := string(body)
|
||
hasSuggestionMode := strings.Contains(bodyStr, "[SUGGESTION MODE:")
|
||
hasWarmupKeyword := strings.Contains(bodyStr, "title") || strings.Contains(bodyStr, "Warmup")
|
||
|
||
if !hasSuggestionMode && !hasWarmupKeyword {
|
||
return InterceptTypeNone
|
||
}
|
||
|
||
// 解析请求（只解析一次）
|
||
var req struct {
|
||
Messages []struct {
|
||
Role string `json:"role"`
|
||
Content []struct {
|
||
Type string `json:"type"`
|
||
Text string `json:"text"`
|
||
} `json:"content"`
|
||
} `json:"messages"`
|
||
System []struct {
|
||
Text string `json:"text"`
|
||
} `json:"system"`
|
||
}
|
||
if err := json.Unmarshal(body, &req); err != nil {
|
||
return InterceptTypeNone
|
||
}
|
||
|
||
// 检查 SUGGESTION MODE（最后一条 user 消息）
|
||
if hasSuggestionMode && len(req.Messages) > 0 {
|
||
lastMsg := req.Messages[len(req.Messages)-1]
|
||
if lastMsg.Role == "user" && len(lastMsg.Content) > 0 &&
|
||
lastMsg.Content[0].Type == "text" &&
|
||
strings.HasPrefix(lastMsg.Content[0].Text, "[SUGGESTION MODE:") {
|
||
return InterceptTypeSuggestionMode
|
||
}
|
||
}
|
||
|
||
// 检查 Warmup 请求
|
||
if hasWarmupKeyword {
|
||
// 检查 messages 中的标题提示模式
|
||
for _, msg := range req.Messages {
|
||
for _, content := range msg.Content {
|
||
if content.Type == "text" {
|
||
if strings.Contains(content.Text, "Please write a 5-10 word title for the following conversation:") ||
|
||
content.Text == "Warmup" {
|
||
return InterceptTypeWarmup
|
||
}
|
||
}
|
||
}
|
||
}
|
||
// 检查 system 中的标题提取模式
|
||
for _, sys := range req.System {
|
||
if strings.Contains(sys.Text, "nalyze if this message indicates a new conversation topic. If it does, extract a 2-3 word title") {
|
||
return InterceptTypeWarmup
|
||
}
|
||
}
|
||
}
|
||
|
||
return InterceptTypeNone
|
||
}
|
||
|
||
// sendMockInterceptStream 发送流式 mock 响应（用于请求拦截）
|
||
func sendMockInterceptStream(c *gin.Context, model string, interceptType InterceptType) {
|
||
c.Header("Content-Type", "text/event-stream")
|
||
c.Header("Cache-Control", "no-cache")
|
||
c.Header("Connection", "keep-alive")
|
||
c.Header("X-Accel-Buffering", "no")
|
||
|
||
// 根据拦截类型决定响应内容
|
||
var msgID string
|
||
var outputTokens int
|
||
var textDeltas []string
|
||
|
||
switch interceptType {
|
||
case InterceptTypeSuggestionMode:
|
||
msgID = generateRealisticMsgID()
|
||
outputTokens = 1
|
||
textDeltas = []string{""} // 空内容
|
||
default: // InterceptTypeWarmup
|
||
msgID = generateRealisticMsgID()
|
||
outputTokens = 2
|
||
textDeltas = []string{"New", " Conversation"}
|
||
}
|
||
|
||
// Build message_start event — field order matches real Anthropic API response.
|
||
messageStartJSON := `{"type":"message_start","message":{"model":` + strconv.Quote(model) + `,"id":` + strconv.Quote(msgID) + `,"type":"message","role":"assistant","content":[],"stop_reason":null,"stop_sequence":null,"stop_details":null,"usage":{"input_tokens":10,"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"output_tokens":0}}}`
|
||
|
||
// Build events
|
||
events := []string{
|
||
`event: message_start` + "\n" + `data: ` + string(messageStartJSON),
|
||
`event: content_block_start` + "\n" + `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
|
||
}
|
||
|
||
// Add text deltas
|
||
for _, text := range textDeltas {
|
||
deltaJSON := `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":` + strconv.Quote(text) + `}}`
|
||
events = append(events, `event: content_block_delta`+"\n"+`data: `+string(deltaJSON))
|
||
}
|
||
|
||
// Add final events
|
||
messageDeltaJSON := `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null,"stop_details":null},"usage":{"output_tokens":` + strconv.Itoa(outputTokens) + `}}`
|
||
|
||
events = append(events,
|
||
`event: content_block_stop`+"\n"+`data: {"index":0,"type":"content_block_stop"}`,
|
||
`event: message_delta`+"\n"+`data: `+string(messageDeltaJSON),
|
||
`event: message_stop`+"\n"+`data: {"type":"message_stop"}`,
|
||
)
|
||
|
||
for _, event := range events {
|
||
_, _ = c.Writer.WriteString(event + "\n\n")
|
||
c.Writer.Flush()
|
||
time.Sleep(20 * time.Millisecond)
|
||
}
|
||
}
|
||
|
||
// generateRealisticMsgID 生成仿真的消息 ID（msg_01XXXXXXX 格式）
|
||
// 格式与 Anthropic API 官方响应一致：msg_01 + 22 位 Base62 随机字符
|
||
func generateRealisticMsgID() string {
|
||
const charset = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||
const idLen = 22
|
||
randomBytes := make([]byte, idLen)
|
||
if _, err := rand.Read(randomBytes); err != nil {
|
||
return fmt.Sprintf("msg_01%d", time.Now().UnixNano())
|
||
}
|
||
b := make([]byte, idLen)
|
||
for i := range b {
|
||
b[i] = charset[int(randomBytes[i])%len(charset)]
|
||
}
|
||
return "msg_01" + string(b)
|
||
}
|
||
|
||
// sendMockInterceptResponse 发送非流式 mock 响应（用于请求拦截）
|
||
func sendMockInterceptResponse(c *gin.Context, model string, interceptType InterceptType) {
|
||
var msgID, text, stopReason string
|
||
var outputTokens int
|
||
|
||
switch interceptType {
|
||
case InterceptTypeSuggestionMode:
|
||
msgID = generateRealisticMsgID()
|
||
text = ""
|
||
outputTokens = 1
|
||
stopReason = "end_turn"
|
||
case InterceptTypeMaxTokensOneHaiku:
|
||
msgID = generateRealisticMsgID()
|
||
text = "#"
|
||
outputTokens = 1
|
||
stopReason = "max_tokens" // max_tokens=1 探测请求的 stop_reason 应为 max_tokens
|
||
default: // InterceptTypeWarmup
|
||
msgID = generateRealisticMsgID()
|
||
text = "New Conversation"
|
||
outputTokens = 2
|
||
stopReason = "end_turn"
|
||
}
|
||
|
||
// 构建完整的响应格式（与 Anthropic API 官方响应格式一致）
|
||
response := gin.H{
|
||
"model": model,
|
||
"id": msgID,
|
||
"type": "message",
|
||
"role": "assistant",
|
||
"content": []gin.H{{"type": "text", "text": text}},
|
||
"stop_reason": stopReason,
|
||
"stop_sequence": nil,
|
||
"stop_details": nil,
|
||
"usage": gin.H{
|
||
"input_tokens": 10,
|
||
"cache_creation_input_tokens": 0,
|
||
"cache_read_input_tokens": 0,
|
||
"cache_creation": gin.H{
|
||
"ephemeral_5m_input_tokens": 0,
|
||
"ephemeral_1h_input_tokens": 0,
|
||
},
|
||
"output_tokens": outputTokens,
|
||
},
|
||
}
|
||
|
||
c.JSON(http.StatusOK, response)
|
||
}
|
||
|
||
// extractQuotaResetSeconds 从 quota 错误的 metadata 中提取 window_resets_at 并计算
|
||
// 距重置剩余秒数。fallback 路径必须返回 ≥1 秒，避免客户端立即重试无限循环。
|
||
func extractQuotaResetSeconds(err error) int {
|
||
const fallback = 60
|
||
appErr := pkgerrors.FromError(err)
|
||
if appErr == nil {
|
||
return fallback
|
||
}
|
||
raw, ok := appErr.Metadata["window_resets_at"]
|
||
if !ok || raw == "" {
|
||
return fallback
|
||
}
|
||
resetAt, parseErr := time.Parse(time.RFC3339, raw)
|
||
if parseErr != nil {
|
||
logger.L().With(
|
||
zap.String("component", "handler.gateway.billing"),
|
||
zap.String("raw", raw),
|
||
zap.Error(parseErr),
|
||
).Warn("quota.invalid_window_resets_at_format")
|
||
return fallback
|
||
}
|
||
secs := time.Until(resetAt).Seconds()
|
||
if secs <= 0 {
|
||
// reset 时间已过：cache 与 DB 应该正在自愈，返回 fallback 让客户端按常规节奏退避，
|
||
// 避免返回 1 秒导致客户端立即重试仍触发限额的退避循环。
|
||
return fallback
|
||
}
|
||
return int(math.Ceil(secs))
|
||
}
|
||
|
||
func billingErrorDetails(err error) (status int, code, message string, retryAfter int) {
|
||
if errors.Is(err, service.ErrBillingServiceUnavailable) {
|
||
msg := pkgerrors.Message(err)
|
||
if msg == "" {
|
||
msg = "Billing service temporarily unavailable. Please retry later."
|
||
}
|
||
return http.StatusServiceUnavailable, "billing_service_error", msg, 0
|
||
}
|
||
if errors.Is(err, service.ErrAPIKeyRateLimit5hExceeded) {
|
||
msg := pkgerrors.Message(err)
|
||
return http.StatusTooManyRequests, "rate_limit_exceeded", msg, 0
|
||
}
|
||
if errors.Is(err, service.ErrAPIKeyRateLimit1dExceeded) {
|
||
msg := pkgerrors.Message(err)
|
||
return http.StatusTooManyRequests, "rate_limit_exceeded", msg, 0
|
||
}
|
||
if errors.Is(err, service.ErrAPIKeyRateLimit7dExceeded) {
|
||
msg := pkgerrors.Message(err)
|
||
return http.StatusTooManyRequests, "rate_limit_exceeded", msg, 0
|
||
}
|
||
// 用户/分组 RPM 超限统一映射为 HTTP 429；保留与其它 rate_limit 一致的错误码便于客户端分类。
|
||
// 返回 Retry-After 秒数（当前分钟剩余秒数），让 SDK 自动退避。
|
||
if errors.Is(err, service.ErrGroupRPMExceeded) || errors.Is(err, service.ErrUserRPMExceeded) {
|
||
msg := pkgerrors.Message(err)
|
||
retrySeconds := 60 - int(time.Now().Unix()%60)
|
||
return http.StatusTooManyRequests, "rate_limit_exceeded", msg, retrySeconds
|
||
}
|
||
if errors.Is(err, service.ErrUserPlatformDailyQuotaExhausted) ||
|
||
errors.Is(err, service.ErrUserPlatformWeeklyQuotaExhausted) ||
|
||
errors.Is(err, service.ErrUserPlatformMonthlyQuotaExhausted) {
|
||
// 与 RPM 超限一致映射 429 + Retry-After，让 SDK 自动退避（而非 403 直接失败）。
|
||
// 错误码用 rate_limit_exceeded 与 OpenAI 兼容客户端一致；细分类型由 ErrCode + window_resets_at metadata 区分。
|
||
msg := pkgerrors.Message(err)
|
||
return http.StatusTooManyRequests, "rate_limit_exceeded", msg, extractQuotaResetSeconds(err)
|
||
}
|
||
msg := pkgerrors.Message(err)
|
||
if msg == "" {
|
||
logger.L().With(
|
||
zap.String("component", "handler.gateway.billing"),
|
||
zap.Error(err),
|
||
).Warn("gateway.billing_error_missing_message")
|
||
msg = "Billing error"
|
||
}
|
||
return http.StatusForbidden, "billing_error", msg, 0
|
||
}
|
||
|
||
func (h *GatewayHandler) metadataBridgeEnabled() bool {
|
||
if h == nil || h.cfg == nil {
|
||
return true
|
||
}
|
||
return h.cfg.Gateway.OpenAIWS.MetadataBridgeEnabled
|
||
}
|
||
|
||
func (h *GatewayHandler) maybeLogCompatibilityFallbackMetrics(reqLog *zap.Logger) {
|
||
if reqLog == nil {
|
||
return
|
||
}
|
||
if gatewayCompatibilityMetricsLogCounter.Add(1)%gatewayCompatibilityMetricsLogInterval != 0 {
|
||
return
|
||
}
|
||
metrics := service.SnapshotOpenAICompatibilityFallbackMetrics()
|
||
reqLog.Info("gateway.compatibility_fallback_metrics",
|
||
zap.Int64("session_hash_legacy_read_fallback_total", metrics.SessionHashLegacyReadFallbackTotal),
|
||
zap.Int64("session_hash_legacy_read_fallback_hit", metrics.SessionHashLegacyReadFallbackHit),
|
||
zap.Int64("session_hash_legacy_dual_write_total", metrics.SessionHashLegacyDualWriteTotal),
|
||
zap.Float64("session_hash_legacy_read_hit_rate", metrics.SessionHashLegacyReadHitRate),
|
||
zap.Int64("metadata_legacy_fallback_total", metrics.MetadataLegacyFallbackTotal),
|
||
)
|
||
}
|
||
|
||
func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task service.UsageRecordTask) {
|
||
if task == nil {
|
||
return
|
||
}
|
||
task = wrapUsageRecordTaskContext(parent, task)
|
||
if h.usageRecordWorkerPool != nil {
|
||
if mode := h.usageRecordWorkerPool.Submit(task); mode != service.UsageRecordSubmitModeDroppedStopped {
|
||
return
|
||
}
|
||
// 池已停止（进程关停窗口）：计费任务不能静默丢失，降级为内联同步执行。
|
||
// 显式配置的 drop/sample 溢出丢弃仍按配置语义保留。
|
||
logger.L().With(
|
||
zap.String("component", "handler.gateway.messages"),
|
||
).Warn("gateway.usage_record_task_stopped_sync_fallback")
|
||
}
|
||
// 回退路径：worker 池未注入或已停止时同步执行，避免退回到无界 goroutine 模式。
|
||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
defer cancel()
|
||
defer func() {
|
||
if recovered := recover(); recovered != nil {
|
||
logger.L().With(
|
||
zap.String("component", "handler.gateway.messages"),
|
||
zap.Any("panic", recovered),
|
||
).Error("gateway.usage_record_task_panic_recovered")
|
||
}
|
||
}()
|
||
task(ctx)
|
||
}
|
||
|
||
// submitMandatoryUsageRecordTask never silently drops billing work on pool overflow.
|
||
func (h *GatewayHandler) submitMandatoryUsageRecordTask(parent context.Context, task service.UsageRecordTask) {
|
||
if task == nil {
|
||
return
|
||
}
|
||
task = wrapUsageRecordTaskContext(parent, task)
|
||
if h.usageRecordWorkerPool != nil {
|
||
if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() {
|
||
return
|
||
}
|
||
logger.L().With(
|
||
zap.String("component", "handler.gateway.usage"),
|
||
).Warn("gateway.usage_record_task_mandatory_sync_fallback")
|
||
}
|
||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||
defer cancel()
|
||
defer func() {
|
||
if recovered := recover(); recovered != nil {
|
||
logger.L().With(
|
||
zap.String("component", "handler.gateway.usage"),
|
||
zap.Any("panic", recovered),
|
||
).Error("gateway.usage_record_task_panic_recovered")
|
||
}
|
||
}()
|
||
task(ctx)
|
||
}
|
||
|
||
// getUserMsgQueueMode 获取当前请求的 UMQ 模式
|
||
// 返回 "serialize" | "throttle" | ""
|
||
func (h *GatewayHandler) getUserMsgQueueMode(account *service.Account, parsed *service.ParsedRequest) string {
|
||
if h.userMsgQueueHelper == nil {
|
||
return ""
|
||
}
|
||
// 仅适用于 Anthropic OAuth/SetupToken 账号
|
||
if !account.IsAnthropicOAuthOrSetupToken() {
|
||
return ""
|
||
}
|
||
if !service.IsRealUserMessage(parsed) {
|
||
return ""
|
||
}
|
||
// 账号级模式优先，fallback 到全局配置
|
||
mode := account.GetUserMsgQueueMode()
|
||
if mode == "" {
|
||
mode = h.cfg.Gateway.UserMessageQueue.GetEffectiveMode()
|
||
}
|
||
return mode
|
||
}
|