2025-12-18 13:50:39 +08:00
|
|
|
|
package handler
|
|
|
|
|
|
|
|
|
|
|
|
import (
|
|
|
|
|
|
"context"
|
2026-02-07 19:04:08 +08:00
|
|
|
|
"crypto/rand"
|
2025-12-18 13:50:39 +08:00
|
|
|
|
"encoding/json"
|
2025-12-27 11:44:00 +08:00
|
|
|
|
"errors"
|
2025-12-18 13:50:39 +08:00
|
|
|
|
"fmt"
|
|
|
|
|
|
"io"
|
|
|
|
|
|
"log"
|
|
|
|
|
|
"net/http"
|
2025-12-19 16:39:25 +08:00
|
|
|
|
"strings"
|
2025-12-18 13:50:39 +08:00
|
|
|
|
"time"
|
|
|
|
|
|
|
2026-01-04 19:49:59 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
2026-02-09 06:47:22 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/domain"
|
2026-01-02 10:21:05 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
2025-12-24 21:07:21 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
2026-01-23 22:24:46 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
2026-01-04 22:10:32 +08:00
|
|
|
|
pkgerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
2026-01-12 15:35:54 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
2026-02-12 19:01:09 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
2025-12-24 21:07:21 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
2025-12-26 10:42:08 +08:00
|
|
|
|
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
2025-12-24 21:07:21 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
|
|
|
|
|
"github.com/gin-gonic/gin"
|
2026-02-12 19:01:09 +08:00
|
|
|
|
"go.uber.org/zap"
|
2025-12-18 13:50:39 +08:00
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
// GatewayHandler handles API gateway requests
|
|
|
|
|
|
type GatewayHandler struct {
|
2025-12-28 17:48:52 +08:00
|
|
|
|
gatewayService *service.GatewayService
|
|
|
|
|
|
geminiCompatService *service.GeminiMessagesCompatService
|
|
|
|
|
|
antigravityGatewayService *service.AntigravityGatewayService
|
|
|
|
|
|
userService *service.UserService
|
|
|
|
|
|
billingCacheService *service.BillingCacheService
|
2026-02-02 18:30:06 +08:00
|
|
|
|
usageService *service.UsageService
|
2026-02-03 19:01:49 +08:00
|
|
|
|
apiKeyService *service.APIKeyService
|
2026-02-22 12:56:57 +08:00
|
|
|
|
usageRecordWorkerPool *service.UsageRecordWorkerPool
|
2026-02-05 21:52:54 +08:00
|
|
|
|
errorPassthroughService *service.ErrorPassthroughService
|
2025-12-28 17:48:52 +08:00
|
|
|
|
concurrencyHelper *ConcurrencyHelper
|
2026-01-16 20:18:30 +08:00
|
|
|
|
maxAccountSwitches int
|
|
|
|
|
|
maxAccountSwitchesGemini int
|
2026-02-01 21:37:10 +08:00
|
|
|
|
cfg *config.Config
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// NewGatewayHandler creates a new GatewayHandler
|
2025-12-25 06:45:03 -08:00
|
|
|
|
func NewGatewayHandler(
|
|
|
|
|
|
gatewayService *service.GatewayService,
|
|
|
|
|
|
geminiCompatService *service.GeminiMessagesCompatService,
|
2025-12-28 17:48:52 +08:00
|
|
|
|
antigravityGatewayService *service.AntigravityGatewayService,
|
2025-12-25 06:45:03 -08:00
|
|
|
|
userService *service.UserService,
|
|
|
|
|
|
concurrencyService *service.ConcurrencyService,
|
|
|
|
|
|
billingCacheService *service.BillingCacheService,
|
2026-02-02 18:30:06 +08:00
|
|
|
|
usageService *service.UsageService,
|
2026-02-03 19:01:49 +08:00
|
|
|
|
apiKeyService *service.APIKeyService,
|
2026-02-22 12:56:57 +08:00
|
|
|
|
usageRecordWorkerPool *service.UsageRecordWorkerPool,
|
2026-02-05 21:52:54 +08:00
|
|
|
|
errorPassthroughService *service.ErrorPassthroughService,
|
2026-01-04 19:49:59 +08:00
|
|
|
|
cfg *config.Config,
|
2025-12-25 06:45:03 -08:00
|
|
|
|
) *GatewayHandler {
|
2026-01-04 19:49:59 +08:00
|
|
|
|
pingInterval := time.Duration(0)
|
2026-01-16 20:18:30 +08:00
|
|
|
|
maxAccountSwitches := 10
|
|
|
|
|
|
maxAccountSwitchesGemini := 3
|
2026-01-04 19:49:59 +08:00
|
|
|
|
if cfg != nil {
|
|
|
|
|
|
pingInterval = time.Duration(cfg.Concurrency.PingInterval) * time.Second
|
2026-01-16 20:18:30 +08:00
|
|
|
|
if cfg.Gateway.MaxAccountSwitches > 0 {
|
|
|
|
|
|
maxAccountSwitches = cfg.Gateway.MaxAccountSwitches
|
|
|
|
|
|
}
|
|
|
|
|
|
if cfg.Gateway.MaxAccountSwitchesGemini > 0 {
|
|
|
|
|
|
maxAccountSwitchesGemini = cfg.Gateway.MaxAccountSwitchesGemini
|
|
|
|
|
|
}
|
2026-01-04 19:49:59 +08:00
|
|
|
|
}
|
2025-12-18 13:50:39 +08:00
|
|
|
|
return &GatewayHandler{
|
2025-12-28 17:48:52 +08:00
|
|
|
|
gatewayService: gatewayService,
|
|
|
|
|
|
geminiCompatService: geminiCompatService,
|
|
|
|
|
|
antigravityGatewayService: antigravityGatewayService,
|
|
|
|
|
|
userService: userService,
|
|
|
|
|
|
billingCacheService: billingCacheService,
|
2026-02-02 18:30:06 +08:00
|
|
|
|
usageService: usageService,
|
2026-02-03 19:01:49 +08:00
|
|
|
|
apiKeyService: apiKeyService,
|
2026-02-22 12:56:57 +08:00
|
|
|
|
usageRecordWorkerPool: usageRecordWorkerPool,
|
2026-02-05 21:52:54 +08:00
|
|
|
|
errorPassthroughService: errorPassthroughService,
|
2026-01-04 19:49:59 +08:00
|
|
|
|
concurrencyHelper: NewConcurrencyHelper(concurrencyService, SSEPingFormatClaude, pingInterval),
|
2026-01-16 20:18:30 +08:00
|
|
|
|
maxAccountSwitches: maxAccountSwitches,
|
|
|
|
|
|
maxAccountSwitchesGemini: maxAccountSwitchesGemini,
|
2026-02-01 21:37:10 +08:00
|
|
|
|
cfg: cfg,
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// Messages handles Claude API compatible messages endpoint
|
|
|
|
|
|
// POST /v1/messages
|
|
|
|
|
|
func (h *GatewayHandler) Messages(c *gin.Context) {
|
|
|
|
|
|
// 从context获取apiKey和user(ApiKeyAuth中间件已设置)
|
2026-01-04 19:27:53 +08:00
|
|
|
|
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if !ok {
|
|
|
|
|
|
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-26 15:40:24 +08:00
|
|
|
|
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if !ok {
|
|
|
|
|
|
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-02-12 19:01:09 +08:00
|
|
|
|
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),
|
|
|
|
|
|
)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
|
|
|
|
|
// 读取请求体
|
|
|
|
|
|
body, err := io.ReadAll(c.Request.Body)
|
|
|
|
|
|
if err != nil {
|
2025-12-31 08:50:12 +08:00
|
|
|
|
if maxErr, ok := extractMaxBytesError(err); ok {
|
|
|
|
|
|
h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2025-12-18 13:50:39 +08:00
|
|
|
|
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
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-09 20:56:37 +08:00
|
|
|
|
setOpsRequestContext(c, "", false, body)
|
|
|
|
|
|
|
2026-02-09 06:47:22 +08:00
|
|
|
|
parsedReq, err := service.ParseGatewayRequest(body, domain.PlatformAnthropic)
|
2025-12-31 08:50:12 +08:00
|
|
|
|
if err != nil {
|
2025-12-18 13:50:39 +08:00
|
|
|
|
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2025-12-31 08:50:12 +08:00
|
|
|
|
reqModel := parsedReq.Model
|
|
|
|
|
|
reqStream := parsedReq.Stream
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog = reqLog.With(zap.String("model", reqModel), zap.Bool("stream", reqStream))
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
2026-02-07 19:04:08 +08:00
|
|
|
|
// 设置 max_tokens=1 + haiku 探测请求标识到 context 中
|
|
|
|
|
|
// 必须在 SetClaudeCodeClientContext 之前设置,因为 ClaudeCodeValidator 需要读取此标识进行绕过判断
|
|
|
|
|
|
if isMaxTokensOneHaikuRequest(reqModel, parsedReq.MaxTokens, reqStream) {
|
|
|
|
|
|
ctx := context.WithValue(c.Request.Context(), ctxkey.IsMaxTokensOneHaikuRequest, true)
|
|
|
|
|
|
c.Request = c.Request.WithContext(ctx)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 检查是否为 Claude Code 客户端,设置到 context 中
|
|
|
|
|
|
SetClaudeCodeClientContext(c, body)
|
|
|
|
|
|
isClaudeCodeClient := service.IsClaudeCodeClient(c.Request.Context())
|
|
|
|
|
|
|
|
|
|
|
|
// 在请求上下文中记录 thinking 状态,供 Antigravity 最终模型 key 推导/模型维度限流使用
|
|
|
|
|
|
c.Request = c.Request.WithContext(context.WithValue(c.Request.Context(), ctxkey.ThinkingEnabled, parsedReq.ThinkingEnabled))
|
|
|
|
|
|
|
2026-01-09 20:56:37 +08:00
|
|
|
|
setOpsRequestContext(c, reqModel, reqStream, body)
|
2026-01-08 23:07:00 +08:00
|
|
|
|
|
2025-12-31 16:17:45 +08:00
|
|
|
|
// 验证 model 必填
|
|
|
|
|
|
if reqModel == "" {
|
|
|
|
|
|
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// Track if we've started streaming (for error handling)
|
|
|
|
|
|
streamStarted := false
|
|
|
|
|
|
|
2026-02-07 14:24:55 +08:00
|
|
|
|
// 绑定错误透传服务,允许 service 层在非 failover 错误场景复用规则。
|
|
|
|
|
|
if h.errorPassthroughService != nil {
|
|
|
|
|
|
service.BindErrorPassthroughService(c, h.errorPassthroughService)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// 获取订阅信息(可能为nil)- 提前获取用于后续检查
|
2025-12-26 10:42:08 +08:00
|
|
|
|
subscription, _ := middleware2.GetSubscriptionFromContext(c)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
|
|
|
|
|
// 0. 检查wait队列是否已满
|
2025-12-26 15:40:24 +08:00
|
|
|
|
maxWait := service.CalculateMaxWait(subject.Concurrency)
|
|
|
|
|
|
canWait, err := h.concurrencyHelper.IncrementWaitCount(c.Request.Context(), subject.UserID, maxWait)
|
2026-01-09 20:56:37 +08:00
|
|
|
|
waitCounted := false
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.user_wait_counter_increment_failed", zap.Error(err))
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// On error, allow request to proceed
|
|
|
|
|
|
} else if !canWait {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Info("gateway.user_wait_queue_full", zap.Int("max_wait", maxWait))
|
2025-12-18 13:50:39 +08:00
|
|
|
|
h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-09 20:56:37 +08:00
|
|
|
|
if err == nil && canWait {
|
|
|
|
|
|
waitCounted = true
|
|
|
|
|
|
}
|
|
|
|
|
|
// Ensure we decrement if we exit before acquiring the user slot.
|
|
|
|
|
|
defer func() {
|
|
|
|
|
|
if waitCounted {
|
|
|
|
|
|
h.concurrencyHelper.DecrementWaitCount(c.Request.Context(), subject.UserID)
|
|
|
|
|
|
}
|
|
|
|
|
|
}()
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
|
|
|
|
|
// 1. 首先获取用户并发槽位
|
2025-12-31 08:50:12 +08:00
|
|
|
|
userReleaseFunc, err := h.concurrencyHelper.AcquireUserSlotWithWait(c, subject.UserID, subject.Concurrency, reqStream, &streamStarted)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.user_slot_acquire_failed", zap.Error(err))
|
2025-12-18 13:50:39 +08:00
|
|
|
|
h.handleConcurrencyError(c, err, "user", streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-09 20:56:37 +08:00
|
|
|
|
// User slot acquired: no longer waiting in the queue.
|
|
|
|
|
|
if waitCounted {
|
|
|
|
|
|
h.concurrencyHelper.DecrementWaitCount(c.Request.Context(), subject.UserID)
|
|
|
|
|
|
waitCounted = false
|
|
|
|
|
|
}
|
2026-01-04 19:49:59 +08:00
|
|
|
|
// 在请求结束或 Context 取消时确保释放槽位,避免客户端断开造成泄漏
|
|
|
|
|
|
userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if userReleaseFunc != nil {
|
|
|
|
|
|
defer userReleaseFunc()
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 2. 【新增】Wait后二次检查余额/订阅
|
2025-12-26 15:40:24 +08:00
|
|
|
|
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription); err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Info("gateway.billing_eligibility_check_failed", zap.Error(err))
|
2026-01-02 17:40:57 +08:00
|
|
|
|
status, code, message := billingErrorDetails(err)
|
|
|
|
|
|
h.handleStreamingAwareError(c, status, code, message, streamStarted)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 计算粘性会话hash
|
2026-02-09 06:46:32 +08:00
|
|
|
|
parsedReq.SessionContext = &service.SessionContext{
|
|
|
|
|
|
ClientIP: ip.GetClientIP(c),
|
|
|
|
|
|
UserAgent: c.GetHeader("User-Agent"),
|
|
|
|
|
|
APIKeyID: apiKey.ID,
|
|
|
|
|
|
}
|
2025-12-31 08:50:12 +08:00
|
|
|
|
sessionHash := h.gatewayService.GenerateSessionHash(parsedReq)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
2025-12-29 16:52:55 +08:00
|
|
|
|
// 获取平台:优先使用强制平台(/antigravity 路由,中间件已设置 request.Context),否则使用分组平台
|
2025-12-25 06:45:03 -08:00
|
|
|
|
platform := ""
|
2025-12-29 16:52:55 +08:00
|
|
|
|
if forcePlatform, ok := middleware2.GetForcePlatformFromContext(c); ok {
|
|
|
|
|
|
platform = forcePlatform
|
|
|
|
|
|
} else if apiKey.Group != nil {
|
2025-12-25 06:45:03 -08:00
|
|
|
|
platform = apiKey.Group.Platform
|
|
|
|
|
|
}
|
2026-01-01 04:01:51 +08:00
|
|
|
|
sessionKey := sessionHash
|
|
|
|
|
|
if platform == service.PlatformGemini && sessionHash != "" {
|
|
|
|
|
|
sessionKey = "gemini:" + sessionHash
|
|
|
|
|
|
}
|
2025-12-25 06:45:03 -08:00
|
|
|
|
|
2026-02-07 12:31:10 +08:00
|
|
|
|
// 查询粘性会话绑定的账号 ID
|
|
|
|
|
|
var sessionBoundAccountID int64
|
|
|
|
|
|
if sessionKey != "" {
|
|
|
|
|
|
sessionBoundAccountID, _ = h.gatewayService.GetCachedSessionAccountID(c.Request.Context(), apiKey.GroupID, sessionKey)
|
2026-02-22 13:31:30 +08:00
|
|
|
|
if sessionBoundAccountID > 0 {
|
|
|
|
|
|
ctx := context.WithValue(c.Request.Context(), ctxkey.PrefetchedStickyAccountID, sessionBoundAccountID)
|
|
|
|
|
|
c.Request = c.Request.WithContext(ctx)
|
|
|
|
|
|
}
|
2026-02-07 12:31:10 +08:00
|
|
|
|
}
|
|
|
|
|
|
// 判断是否真的绑定了粘性会话:有 sessionKey 且已经绑定到某个账号
|
|
|
|
|
|
hasBoundSession := sessionKey != "" && sessionBoundAccountID > 0
|
|
|
|
|
|
|
2025-12-26 22:07:55 +08:00
|
|
|
|
if platform == service.PlatformGemini {
|
2026-01-16 20:18:30 +08:00
|
|
|
|
maxAccountSwitches := h.maxAccountSwitchesGemini
|
2025-12-27 12:27:47 +08:00
|
|
|
|
switchCount := 0
|
|
|
|
|
|
failedAccountIDs := make(map[int64]struct{})
|
2026-02-10 00:53:54 +08:00
|
|
|
|
sameAccountRetryCount := make(map[int64]int) // 同账号重试计数
|
2026-02-05 21:52:54 +08:00
|
|
|
|
var lastFailoverErr *service.UpstreamFailoverError
|
2026-02-07 12:31:10 +08:00
|
|
|
|
var forceCacheBilling bool // 粘性会话切换时的缓存计费标记
|
2025-12-27 12:27:47 +08:00
|
|
|
|
|
2026-02-09 17:25:36 +08:00
|
|
|
|
// 单账号分组提前设置 SingleAccountRetry 标记,让 Service 层首次 503 就不设模型限流标记。
|
|
|
|
|
|
// 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流,导致后续请求连续快速失败。
|
|
|
|
|
|
if h.gatewayService.IsSingleAntigravityAccountGroup(c.Request.Context(), apiKey.GroupID) {
|
|
|
|
|
|
ctx := context.WithValue(c.Request.Context(), ctxkey.SingleAccountRetry, true)
|
|
|
|
|
|
c.Request = c.Request.WithContext(ctx)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-27 12:27:47 +08:00
|
|
|
|
for {
|
2026-01-16 23:36:52 +08:00
|
|
|
|
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, failedAccountIDs, "") // Gemini 不使用会话限制
|
2025-12-27 12:27:47 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
|
if len(failedAccountIDs) == 0 {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.account_select_failed", zap.Error(err), zap.Int("excluded_account_count", len(failedAccountIDs)))
|
2026-02-07 19:46:42 +08:00
|
|
|
|
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "Service temporarily unavailable", streamStarted)
|
2025-12-27 12:27:47 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-02-09 14:26:01 +08:00
|
|
|
|
// Antigravity 单账号退避重试:分组内没有其他可用账号时,
|
|
|
|
|
|
// 对 503 错误不直接返回,而是清除排除列表、等待退避后重试同一个账号。
|
|
|
|
|
|
// 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的,等几秒就能恢复。
|
|
|
|
|
|
if lastFailoverErr != nil && lastFailoverErr.StatusCode == http.StatusServiceUnavailable && switchCount <= maxAccountSwitches {
|
|
|
|
|
|
if sleepAntigravitySingleAccountBackoff(c.Request.Context(), switchCount) {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.single_account_retrying",
|
|
|
|
|
|
zap.Int("retry_count", switchCount),
|
|
|
|
|
|
zap.Int("max_retries", maxAccountSwitches),
|
|
|
|
|
|
)
|
2026-02-09 14:26:01 +08:00
|
|
|
|
failedAccountIDs = make(map[int64]struct{})
|
|
|
|
|
|
// 设置 context 标记,让 Service 层预检查等待限流过期而非直接切换
|
|
|
|
|
|
ctx := context.WithValue(c.Request.Context(), ctxkey.SingleAccountRetry, true)
|
|
|
|
|
|
c.Request = c.Request.WithContext(ctx)
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
2026-02-05 21:52:54 +08:00
|
|
|
|
if lastFailoverErr != nil {
|
|
|
|
|
|
h.handleFailoverExhausted(c, lastFailoverErr, service.PlatformGemini, streamStarted)
|
|
|
|
|
|
} else {
|
|
|
|
|
|
h.handleFailoverExhaustedSimple(c, 502, streamStarted)
|
|
|
|
|
|
}
|
2025-12-27 12:27:47 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-01 04:01:51 +08:00
|
|
|
|
account := selection.Account
|
2026-02-12 16:27:29 +08:00
|
|
|
|
setOpsSelectedAccount(c, account.ID, account.Platform)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// 检查请求拦截(预热请求、SUGGESTION MODE等)
|
|
|
|
|
|
if account.IsInterceptWarmupEnabled() {
|
2026-02-07 19:04:08 +08:00
|
|
|
|
interceptType := detectInterceptType(body, reqModel, parsedReq.MaxTokens, reqStream, isClaudeCodeClient)
|
2026-01-23 16:57:25 +08:00
|
|
|
|
if interceptType != InterceptTypeNone {
|
|
|
|
|
|
if selection.Acquired && selection.ReleaseFunc != nil {
|
|
|
|
|
|
selection.ReleaseFunc()
|
|
|
|
|
|
}
|
|
|
|
|
|
if reqStream {
|
|
|
|
|
|
sendMockInterceptStream(c, reqModel, interceptType)
|
|
|
|
|
|
} else {
|
|
|
|
|
|
sendMockInterceptResponse(c, reqModel, interceptType)
|
|
|
|
|
|
}
|
|
|
|
|
|
return
|
2025-12-27 12:27:47 +08:00
|
|
|
|
}
|
2025-12-27 11:44:00 +08:00
|
|
|
|
}
|
2025-12-19 16:39:25 +08:00
|
|
|
|
|
2025-12-27 12:27:47 +08:00
|
|
|
|
// 3. 获取账号并发槽位
|
2026-01-01 04:01:51 +08:00
|
|
|
|
accountReleaseFunc := selection.ReleaseFunc
|
|
|
|
|
|
if !selection.Acquired {
|
|
|
|
|
|
if selection.WaitPlan == nil {
|
|
|
|
|
|
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts", streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-09 20:56:37 +08:00
|
|
|
|
accountWaitCounted := false
|
2026-01-01 04:01:51 +08:00
|
|
|
|
canWait, err := h.concurrencyHelper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
|
|
|
|
|
|
if err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.account_wait_counter_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
2026-01-01 04:01:51 +08:00
|
|
|
|
} else if !canWait {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Info("gateway.account_wait_queue_full",
|
|
|
|
|
|
zap.Int64("account_id", account.ID),
|
|
|
|
|
|
zap.Int("max_waiting", selection.WaitPlan.MaxWaiting),
|
|
|
|
|
|
)
|
2026-01-01 04:01:51 +08:00
|
|
|
|
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later", streamStarted)
|
|
|
|
|
|
return
|
2026-01-09 20:56:37 +08:00
|
|
|
|
}
|
|
|
|
|
|
if err == nil && canWait {
|
|
|
|
|
|
accountWaitCounted = true
|
|
|
|
|
|
}
|
2026-02-07 19:46:42 +08:00
|
|
|
|
releaseWait := func() {
|
2026-01-09 20:56:37 +08:00
|
|
|
|
if accountWaitCounted {
|
2026-01-01 04:26:01 +08:00
|
|
|
|
h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
2026-02-07 19:46:42 +08:00
|
|
|
|
accountWaitCounted = false
|
2026-01-01 04:26:01 +08:00
|
|
|
|
}
|
2026-02-07 19:46:42 +08:00
|
|
|
|
}
|
2026-01-01 04:01:51 +08:00
|
|
|
|
|
|
|
|
|
|
accountReleaseFunc, err = h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout(
|
|
|
|
|
|
c,
|
|
|
|
|
|
account.ID,
|
|
|
|
|
|
selection.WaitPlan.MaxConcurrency,
|
|
|
|
|
|
selection.WaitPlan.Timeout,
|
|
|
|
|
|
reqStream,
|
|
|
|
|
|
&streamStarted,
|
|
|
|
|
|
)
|
|
|
|
|
|
if err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
2026-02-07 19:46:42 +08:00
|
|
|
|
releaseWait()
|
2026-01-01 04:01:51 +08:00
|
|
|
|
h.handleConcurrencyError(c, err, "account", streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-09 20:56:37 +08:00
|
|
|
|
// Slot acquired: no longer waiting in queue.
|
2026-02-07 19:46:42 +08:00
|
|
|
|
releaseWait()
|
2026-01-08 23:07:00 +08:00
|
|
|
|
if err := h.gatewayService.BindStickySession(c.Request.Context(), apiKey.GroupID, sessionKey, account.ID); err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.bind_sticky_session_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
2026-01-01 04:01:51 +08:00
|
|
|
|
}
|
2025-12-27 12:27:47 +08:00
|
|
|
|
}
|
2026-01-04 19:49:59 +08:00
|
|
|
|
// 账号槽位/等待计数需要在超时或断开时安全回收
|
|
|
|
|
|
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
|
2025-12-27 12:27:47 +08:00
|
|
|
|
|
2025-12-28 17:48:52 +08:00
|
|
|
|
// 转发请求 - 根据账号平台分流
|
|
|
|
|
|
var result *service.ForwardResult
|
2026-01-28 00:01:03 +08:00
|
|
|
|
requestCtx := c.Request.Context()
|
|
|
|
|
|
if switchCount > 0 {
|
|
|
|
|
|
requestCtx = context.WithValue(requestCtx, ctxkey.AccountSwitchCount, switchCount)
|
|
|
|
|
|
}
|
2025-12-28 17:48:52 +08:00
|
|
|
|
if account.Platform == service.PlatformAntigravity {
|
2026-02-07 12:31:10 +08:00
|
|
|
|
result, err = h.antigravityGatewayService.ForwardGemini(requestCtx, c, account, reqModel, "generateContent", reqStream, body, hasBoundSession)
|
2025-12-28 17:48:52 +08:00
|
|
|
|
} else {
|
2026-01-28 00:01:03 +08:00
|
|
|
|
result, err = h.geminiCompatService.Forward(requestCtx, c, account, body)
|
2025-12-28 17:48:52 +08:00
|
|
|
|
}
|
2025-12-27 12:27:47 +08:00
|
|
|
|
if accountReleaseFunc != nil {
|
|
|
|
|
|
accountReleaseFunc()
|
|
|
|
|
|
}
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
var failoverErr *service.UpstreamFailoverError
|
|
|
|
|
|
if errors.As(err, &failoverErr) {
|
2026-02-05 21:52:54 +08:00
|
|
|
|
lastFailoverErr = failoverErr
|
2026-02-09 06:57:07 +08:00
|
|
|
|
if needForceCacheBilling(hasBoundSession, failoverErr) {
|
2026-02-07 12:31:10 +08:00
|
|
|
|
forceCacheBilling = true
|
|
|
|
|
|
}
|
2026-02-10 00:53:54 +08:00
|
|
|
|
|
|
|
|
|
|
// 同账号重试:对 RetryableOnSameAccount 的临时性错误,先在同一账号上重试
|
|
|
|
|
|
if failoverErr.RetryableOnSameAccount && sameAccountRetryCount[account.ID] < maxSameAccountRetries {
|
|
|
|
|
|
sameAccountRetryCount[account.ID]++
|
|
|
|
|
|
log.Printf("Account %d: retryable error %d, same-account retry %d/%d",
|
|
|
|
|
|
account.ID, failoverErr.StatusCode, sameAccountRetryCount[account.ID], maxSameAccountRetries)
|
|
|
|
|
|
if !sleepSameAccountRetryDelay(c.Request.Context()) {
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 同账号重试用尽,执行临时封禁并切换账号
|
|
|
|
|
|
if failoverErr.RetryableOnSameAccount {
|
|
|
|
|
|
h.gatewayService.TempUnscheduleRetryableError(c.Request.Context(), account.ID, failoverErr)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
failedAccountIDs[account.ID] = struct{}{}
|
2025-12-27 12:27:47 +08:00
|
|
|
|
if switchCount >= maxAccountSwitches {
|
2026-02-05 21:52:54 +08:00
|
|
|
|
h.handleFailoverExhausted(c, failoverErr, service.PlatformGemini, streamStarted)
|
2025-12-27 12:27:47 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
switchCount++
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.upstream_failover_switching",
|
|
|
|
|
|
zap.Int64("account_id", account.ID),
|
|
|
|
|
|
zap.Int("upstream_status", failoverErr.StatusCode),
|
|
|
|
|
|
zap.Int("switch_count", switchCount),
|
|
|
|
|
|
zap.Int("max_switches", maxAccountSwitches),
|
|
|
|
|
|
)
|
2026-02-09 06:56:29 +08:00
|
|
|
|
if account.Platform == service.PlatformAntigravity {
|
|
|
|
|
|
if !sleepFailoverDelay(c.Request.Context(), switchCount) {
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
2025-12-27 12:27:47 +08:00
|
|
|
|
continue
|
|
|
|
|
|
}
|
2026-02-14 11:23:10 +08:00
|
|
|
|
wroteFallback := h.ensureForwardErrorResponse(c, streamStarted)
|
|
|
|
|
|
reqLog.Error("gateway.forward_failed",
|
|
|
|
|
|
zap.Int64("account_id", account.ID),
|
|
|
|
|
|
zap.Bool("fallback_error_response_written", wroteFallback),
|
|
|
|
|
|
zap.Error(err),
|
|
|
|
|
|
)
|
2025-12-27 12:27:47 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
2026-01-12 15:19:40 +08:00
|
|
|
|
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context)
|
|
|
|
|
|
userAgent := c.GetHeader("User-Agent")
|
2026-01-12 15:35:54 +08:00
|
|
|
|
clientIP := ip.GetClientIP(c)
|
2026-01-12 15:19:40 +08:00
|
|
|
|
|
2026-02-22 12:56:57 +08:00
|
|
|
|
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
|
|
|
|
|
|
h.submitUsageRecordTask(func(ctx context.Context) {
|
2025-12-27 12:27:47 +08:00
|
|
|
|
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
2026-02-07 12:31:10 +08:00
|
|
|
|
Result: result,
|
|
|
|
|
|
APIKey: apiKey,
|
|
|
|
|
|
User: apiKey.User,
|
2026-02-22 12:56:57 +08:00
|
|
|
|
Account: account,
|
2026-02-07 12:31:10 +08:00
|
|
|
|
Subscription: subscription,
|
2026-02-22 12:56:57 +08:00
|
|
|
|
UserAgent: userAgent,
|
2026-02-07 12:31:10 +08:00
|
|
|
|
IPAddress: clientIP,
|
2026-02-22 12:56:57 +08:00
|
|
|
|
ForceCacheBilling: forceCacheBilling,
|
2026-02-07 12:31:10 +08:00
|
|
|
|
APIKeyService: h.apiKeyService,
|
2025-12-27 12:27:47 +08:00
|
|
|
|
}); err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
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),
|
2026-02-22 12:56:57 +08:00
|
|
|
|
zap.Int64("account_id", account.ID),
|
2026-02-12 19:01:09 +08:00
|
|
|
|
).Error("gateway.record_usage_failed", zap.Error(err))
|
2025-12-27 12:27:47 +08:00
|
|
|
|
}
|
2026-02-22 12:56:57 +08:00
|
|
|
|
})
|
2025-12-27 11:44:00 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 22:24:46 +08:00
|
|
|
|
currentAPIKey := apiKey
|
|
|
|
|
|
currentSubscription := subscription
|
|
|
|
|
|
var fallbackGroupID *int64
|
|
|
|
|
|
if apiKey.Group != nil {
|
|
|
|
|
|
fallbackGroupID = apiKey.Group.FallbackGroupIDOnInvalidRequest
|
|
|
|
|
|
}
|
|
|
|
|
|
fallbackUsed := false
|
2025-12-27 11:44:00 +08:00
|
|
|
|
|
2026-02-09 17:25:36 +08:00
|
|
|
|
// 单账号分组提前设置 SingleAccountRetry 标记,让 Service 层首次 503 就不设模型限流标记。
|
|
|
|
|
|
// 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流,导致后续请求连续快速失败。
|
|
|
|
|
|
if h.gatewayService.IsSingleAntigravityAccountGroup(c.Request.Context(), currentAPIKey.GroupID) {
|
|
|
|
|
|
ctx := context.WithValue(c.Request.Context(), ctxkey.SingleAccountRetry, true)
|
|
|
|
|
|
c.Request = c.Request.WithContext(ctx)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-27 11:44:00 +08:00
|
|
|
|
for {
|
2026-01-23 22:24:46 +08:00
|
|
|
|
maxAccountSwitches := h.maxAccountSwitches
|
|
|
|
|
|
switchCount := 0
|
|
|
|
|
|
failedAccountIDs := make(map[int64]struct{})
|
2026-02-10 00:53:54 +08:00
|
|
|
|
sameAccountRetryCount := make(map[int64]int) // 同账号重试计数
|
2026-02-05 21:52:54 +08:00
|
|
|
|
var lastFailoverErr *service.UpstreamFailoverError
|
2026-01-23 22:24:46 +08:00
|
|
|
|
retryWithFallback := false
|
2026-02-07 12:31:10 +08:00
|
|
|
|
var forceCacheBilling bool // 粘性会话切换时的缓存计费标记
|
2025-12-27 11:44:00 +08:00
|
|
|
|
|
2026-01-23 22:24:46 +08:00
|
|
|
|
for {
|
|
|
|
|
|
// 选择支持该模型的账号
|
|
|
|
|
|
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), currentAPIKey.GroupID, sessionKey, reqModel, failedAccountIDs, parsedReq.MetadataUserID)
|
2026-01-01 04:01:51 +08:00
|
|
|
|
if err != nil {
|
2026-01-23 22:24:46 +08:00
|
|
|
|
if len(failedAccountIDs) == 0 {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.account_select_failed", zap.Error(err), zap.Int("excluded_account_count", len(failedAccountIDs)))
|
2026-02-07 19:46:42 +08:00
|
|
|
|
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "Service temporarily unavailable", streamStarted)
|
2026-01-23 22:24:46 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-02-09 14:26:01 +08:00
|
|
|
|
// Antigravity 单账号退避重试:分组内没有其他可用账号时,
|
|
|
|
|
|
// 对 503 错误不直接返回,而是清除排除列表、等待退避后重试同一个账号。
|
|
|
|
|
|
// 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的,等几秒就能恢复。
|
|
|
|
|
|
if lastFailoverErr != nil && lastFailoverErr.StatusCode == http.StatusServiceUnavailable && switchCount <= maxAccountSwitches {
|
|
|
|
|
|
if sleepAntigravitySingleAccountBackoff(c.Request.Context(), switchCount) {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.single_account_retrying",
|
|
|
|
|
|
zap.Int("retry_count", switchCount),
|
|
|
|
|
|
zap.Int("max_retries", maxAccountSwitches),
|
|
|
|
|
|
)
|
2026-02-09 14:26:01 +08:00
|
|
|
|
failedAccountIDs = make(map[int64]struct{})
|
|
|
|
|
|
// 设置 context 标记,让 Service 层预检查等待限流过期而非直接切换
|
|
|
|
|
|
ctx := context.WithValue(c.Request.Context(), ctxkey.SingleAccountRetry, true)
|
|
|
|
|
|
c.Request = c.Request.WithContext(ctx)
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
2026-02-05 21:52:54 +08:00
|
|
|
|
if lastFailoverErr != nil {
|
|
|
|
|
|
h.handleFailoverExhausted(c, lastFailoverErr, platform, streamStarted)
|
|
|
|
|
|
} else {
|
|
|
|
|
|
h.handleFailoverExhaustedSimple(c, 502, streamStarted)
|
|
|
|
|
|
}
|
2026-01-01 04:01:51 +08:00
|
|
|
|
return
|
2026-01-09 20:56:37 +08:00
|
|
|
|
}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
account := selection.Account
|
2026-02-12 16:27:29 +08:00
|
|
|
|
setOpsSelectedAccount(c, account.ID, account.Platform)
|
2026-01-23 22:24:46 +08:00
|
|
|
|
|
2026-02-02 22:13:50 +08:00
|
|
|
|
// 检查请求拦截(预热请求、SUGGESTION MODE等)
|
|
|
|
|
|
if account.IsInterceptWarmupEnabled() {
|
2026-02-07 19:04:08 +08:00
|
|
|
|
interceptType := detectInterceptType(body, reqModel, parsedReq.MaxTokens, reqStream, isClaudeCodeClient)
|
2026-02-02 22:13:50 +08:00
|
|
|
|
if interceptType != InterceptTypeNone {
|
|
|
|
|
|
if selection.Acquired && selection.ReleaseFunc != nil {
|
|
|
|
|
|
selection.ReleaseFunc()
|
|
|
|
|
|
}
|
|
|
|
|
|
if reqStream {
|
|
|
|
|
|
sendMockInterceptStream(c, reqModel, interceptType)
|
|
|
|
|
|
} else {
|
|
|
|
|
|
sendMockInterceptResponse(c, reqModel, interceptType)
|
|
|
|
|
|
}
|
|
|
|
|
|
return
|
2026-01-23 22:24:46 +08:00
|
|
|
|
}
|
2026-01-09 20:56:37 +08:00
|
|
|
|
}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
|
|
|
|
|
|
// 3. 获取账号并发槽位
|
|
|
|
|
|
accountReleaseFunc := selection.ReleaseFunc
|
|
|
|
|
|
if !selection.Acquired {
|
|
|
|
|
|
if selection.WaitPlan == nil {
|
|
|
|
|
|
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 {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.account_wait_counter_increment_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
2026-01-23 22:24:46 +08:00
|
|
|
|
} else if !canWait {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Info("gateway.account_wait_queue_full",
|
|
|
|
|
|
zap.Int64("account_id", account.ID),
|
|
|
|
|
|
zap.Int("max_waiting", selection.WaitPlan.MaxWaiting),
|
|
|
|
|
|
)
|
2026-01-23 22:24:46 +08:00
|
|
|
|
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later", streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
if err == nil && canWait {
|
|
|
|
|
|
accountWaitCounted = true
|
|
|
|
|
|
}
|
2026-02-07 19:46:42 +08:00
|
|
|
|
releaseWait := func() {
|
2026-01-23 22:24:46 +08:00
|
|
|
|
if accountWaitCounted {
|
|
|
|
|
|
h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
2026-02-07 19:46:42 +08:00
|
|
|
|
accountWaitCounted = false
|
2026-01-23 22:24:46 +08:00
|
|
|
|
}
|
2026-02-07 19:46:42 +08:00
|
|
|
|
}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
|
|
|
|
|
|
accountReleaseFunc, err = h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout(
|
|
|
|
|
|
c,
|
|
|
|
|
|
account.ID,
|
|
|
|
|
|
selection.WaitPlan.MaxConcurrency,
|
|
|
|
|
|
selection.WaitPlan.Timeout,
|
|
|
|
|
|
reqStream,
|
|
|
|
|
|
&streamStarted,
|
|
|
|
|
|
)
|
|
|
|
|
|
if err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.account_slot_acquire_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
2026-02-07 19:46:42 +08:00
|
|
|
|
releaseWait()
|
2026-01-23 22:24:46 +08:00
|
|
|
|
h.handleConcurrencyError(c, err, "account", streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-02-07 19:46:42 +08:00
|
|
|
|
// Slot acquired: no longer waiting in queue.
|
|
|
|
|
|
releaseWait()
|
2026-01-23 22:24:46 +08:00
|
|
|
|
if err := h.gatewayService.BindStickySession(c.Request.Context(), currentAPIKey.GroupID, sessionKey, account.ID); err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.bind_sticky_session_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
2026-01-01 04:26:01 +08:00
|
|
|
|
}
|
2026-01-01 04:01:51 +08:00
|
|
|
|
}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
// 账号槽位/等待计数需要在超时或断开时安全回收
|
|
|
|
|
|
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
|
|
|
|
|
|
|
|
|
|
|
|
// 转发请求 - 根据账号平台分流
|
|
|
|
|
|
var result *service.ForwardResult
|
2026-01-28 00:01:03 +08:00
|
|
|
|
requestCtx := c.Request.Context()
|
|
|
|
|
|
if switchCount > 0 {
|
|
|
|
|
|
requestCtx = context.WithValue(requestCtx, ctxkey.AccountSwitchCount, switchCount)
|
2026-01-01 04:01:51 +08:00
|
|
|
|
}
|
2026-02-08 13:06:25 +08:00
|
|
|
|
if account.Platform == service.PlatformAntigravity && account.Type != service.AccountTypeAPIKey {
|
2026-02-07 12:31:10 +08:00
|
|
|
|
result, err = h.antigravityGatewayService.Forward(requestCtx, c, account, body, hasBoundSession)
|
2026-01-23 22:24:46 +08:00
|
|
|
|
} else {
|
2026-01-28 00:01:03 +08:00
|
|
|
|
result, err = h.gatewayService.Forward(requestCtx, c, account, parsedReq)
|
2026-01-09 20:56:37 +08:00
|
|
|
|
}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
if accountReleaseFunc != nil {
|
|
|
|
|
|
accountReleaseFunc()
|
2026-01-01 04:01:51 +08:00
|
|
|
|
}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
|
var promptTooLongErr *service.PromptTooLongError
|
|
|
|
|
|
if errors.As(err, &promptTooLongErr) {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
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),
|
|
|
|
|
|
)
|
2026-01-23 22:24:46 +08:00
|
|
|
|
if !fallbackUsed && fallbackGroupID != nil && *fallbackGroupID > 0 {
|
|
|
|
|
|
fallbackGroup, err := h.gatewayService.ResolveGroupByID(c.Request.Context(), *fallbackGroupID)
|
|
|
|
|
|
if err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.resolve_fallback_group_failed", zap.Int64("fallback_group_id", *fallbackGroupID), zap.Error(err))
|
2026-01-23 22:24:46 +08:00
|
|
|
|
_ = h.antigravityGatewayService.WriteMappedClaudeError(c, account, promptTooLongErr.StatusCode, promptTooLongErr.RequestID, promptTooLongErr.Body)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
if fallbackGroup.Platform != service.PlatformAnthropic ||
|
|
|
|
|
|
fallbackGroup.SubscriptionType == service.SubscriptionTypeSubscription ||
|
|
|
|
|
|
fallbackGroup.FallbackGroupIDOnInvalidRequest != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
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),
|
|
|
|
|
|
)
|
2026-01-23 22:24:46 +08:00
|
|
|
|
_ = 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); err != nil {
|
|
|
|
|
|
status, code, message := billingErrorDetails(err)
|
|
|
|
|
|
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)
|
2025-12-27 11:44:00 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
var failoverErr *service.UpstreamFailoverError
|
|
|
|
|
|
if errors.As(err, &failoverErr) {
|
2026-02-05 21:52:54 +08:00
|
|
|
|
lastFailoverErr = failoverErr
|
2026-02-09 06:57:07 +08:00
|
|
|
|
if needForceCacheBilling(hasBoundSession, failoverErr) {
|
2026-02-07 12:31:10 +08:00
|
|
|
|
forceCacheBilling = true
|
|
|
|
|
|
}
|
2026-02-10 00:53:54 +08:00
|
|
|
|
|
|
|
|
|
|
// 同账号重试:对 RetryableOnSameAccount 的临时性错误,先在同一账号上重试
|
|
|
|
|
|
if failoverErr.RetryableOnSameAccount && sameAccountRetryCount[account.ID] < maxSameAccountRetries {
|
|
|
|
|
|
sameAccountRetryCount[account.ID]++
|
|
|
|
|
|
log.Printf("Account %d: retryable error %d, same-account retry %d/%d",
|
|
|
|
|
|
account.ID, failoverErr.StatusCode, sameAccountRetryCount[account.ID], maxSameAccountRetries)
|
|
|
|
|
|
if !sleepSameAccountRetryDelay(c.Request.Context()) {
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 同账号重试用尽,执行临时封禁并切换账号
|
|
|
|
|
|
if failoverErr.RetryableOnSameAccount {
|
|
|
|
|
|
h.gatewayService.TempUnscheduleRetryableError(c.Request.Context(), account.ID, failoverErr)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
failedAccountIDs[account.ID] = struct{}{}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
if switchCount >= maxAccountSwitches {
|
2026-02-05 21:52:54 +08:00
|
|
|
|
h.handleFailoverExhausted(c, failoverErr, account.Platform, streamStarted)
|
2026-01-23 22:24:46 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
switchCount++
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.upstream_failover_switching",
|
|
|
|
|
|
zap.Int64("account_id", account.ID),
|
|
|
|
|
|
zap.Int("upstream_status", failoverErr.StatusCode),
|
|
|
|
|
|
zap.Int("switch_count", switchCount),
|
|
|
|
|
|
zap.Int("max_switches", maxAccountSwitches),
|
|
|
|
|
|
)
|
2026-02-09 06:56:29 +08:00
|
|
|
|
if account.Platform == service.PlatformAntigravity {
|
|
|
|
|
|
if !sleepFailoverDelay(c.Request.Context(), switchCount) {
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
continue
|
|
|
|
|
|
}
|
2026-02-14 11:23:10 +08:00
|
|
|
|
wroteFallback := h.ensureForwardErrorResponse(c, streamStarted)
|
|
|
|
|
|
reqLog.Error("gateway.forward_failed",
|
|
|
|
|
|
zap.Int64("account_id", account.ID),
|
|
|
|
|
|
zap.Bool("fallback_error_response_written", wroteFallback),
|
|
|
|
|
|
zap.Error(err),
|
|
|
|
|
|
)
|
2026-01-23 22:24:46 +08:00
|
|
|
|
return
|
2025-12-27 11:44:00 +08:00
|
|
|
|
}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
|
|
|
|
|
|
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context)
|
|
|
|
|
|
userAgent := c.GetHeader("User-Agent")
|
|
|
|
|
|
clientIP := ip.GetClientIP(c)
|
|
|
|
|
|
|
2026-02-22 12:56:57 +08:00
|
|
|
|
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
|
|
|
|
|
|
h.submitUsageRecordTask(func(ctx context.Context) {
|
2026-01-23 22:24:46 +08:00
|
|
|
|
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
2026-02-07 12:31:10 +08:00
|
|
|
|
Result: result,
|
|
|
|
|
|
APIKey: currentAPIKey,
|
|
|
|
|
|
User: currentAPIKey.User,
|
2026-02-22 12:56:57 +08:00
|
|
|
|
Account: account,
|
2026-02-07 12:31:10 +08:00
|
|
|
|
Subscription: currentSubscription,
|
2026-02-22 12:56:57 +08:00
|
|
|
|
UserAgent: userAgent,
|
2026-02-07 12:31:10 +08:00
|
|
|
|
IPAddress: clientIP,
|
2026-02-22 12:56:57 +08:00
|
|
|
|
ForceCacheBilling: forceCacheBilling,
|
2026-02-07 12:31:10 +08:00
|
|
|
|
APIKeyService: h.apiKeyService,
|
2026-01-23 22:24:46 +08:00
|
|
|
|
}); err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
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),
|
2026-02-22 12:56:57 +08:00
|
|
|
|
zap.Int64("account_id", account.ID),
|
2026-02-12 19:01:09 +08:00
|
|
|
|
).Error("gateway.record_usage_failed", zap.Error(err))
|
2026-01-23 22:24:46 +08:00
|
|
|
|
}
|
2026-02-22 12:56:57 +08:00
|
|
|
|
})
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Debug("gateway.request_completed",
|
|
|
|
|
|
zap.Int64("account_id", account.ID),
|
|
|
|
|
|
zap.Int("switch_count", switchCount),
|
|
|
|
|
|
zap.Bool("fallback_used", fallbackUsed),
|
|
|
|
|
|
)
|
2025-12-27 11:44:00 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-23 22:24:46 +08:00
|
|
|
|
if !retryWithFallback {
|
2025-12-27 11:44:00 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// Models handles listing available models
|
|
|
|
|
|
// GET /v1/models
|
2026-01-01 16:03:48 +08:00
|
|
|
|
// Returns models based on account configurations (model_mapping whitelist)
|
|
|
|
|
|
// Falls back to default models if no whitelist is configured
|
2025-12-18 13:50:39 +08:00
|
|
|
|
func (h *GatewayHandler) Models(c *gin.Context) {
|
2026-01-04 19:27:53 +08:00
|
|
|
|
apiKey, _ := middleware2.GetAPIKeyFromContext(c)
|
2025-12-22 22:58:31 +08:00
|
|
|
|
|
2026-01-01 16:03:48 +08:00
|
|
|
|
var groupID *int64
|
|
|
|
|
|
var platform string
|
|
|
|
|
|
|
|
|
|
|
|
if apiKey != nil && apiKey.Group != nil {
|
|
|
|
|
|
groupID = &apiKey.Group.ID
|
|
|
|
|
|
platform = apiKey.Group.Platform
|
|
|
|
|
|
}
|
2026-01-31 20:22:22 +08:00
|
|
|
|
if forcedPlatform, ok := middleware2.GetForcePlatformFromContext(c); ok && strings.TrimSpace(forcedPlatform) != "" {
|
|
|
|
|
|
platform = forcedPlatform
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if platform == service.PlatformSora {
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
|
|
|
|
"object": "list",
|
2026-02-01 21:37:10 +08:00
|
|
|
|
"data": service.DefaultSoraModels(h.cfg),
|
2026-01-31 20:22:22 +08:00
|
|
|
|
})
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-01 16:03:48 +08:00
|
|
|
|
|
|
|
|
|
|
// Get available models from account configurations (without platform filter)
|
|
|
|
|
|
availableModels := h.gatewayService.GetAvailableModels(c.Request.Context(), groupID, "")
|
|
|
|
|
|
|
|
|
|
|
|
if len(availableModels) > 0 {
|
|
|
|
|
|
// Build model list from whitelist
|
|
|
|
|
|
models := make([]claude.Model, 0, len(availableModels))
|
|
|
|
|
|
for _, modelID := range availableModels {
|
|
|
|
|
|
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,
|
|
|
|
|
|
})
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// Fallback to default models
|
|
|
|
|
|
if platform == "openai" {
|
2025-12-22 22:58:31 +08:00
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
|
|
|
|
"object": "list",
|
|
|
|
|
|
"data": openai.DefaultModels,
|
|
|
|
|
|
})
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
|
|
|
|
"object": "list",
|
2025-12-22 22:58:31 +08:00
|
|
|
|
"data": claude.DefaultModels,
|
2025-12-18 13:50:39 +08:00
|
|
|
|
})
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-02 10:21:05 +08:00
|
|
|
|
// AntigravityModels 返回 Antigravity 支持的全部模型
|
|
|
|
|
|
// GET /antigravity/models
|
|
|
|
|
|
func (h *GatewayHandler) AntigravityModels(c *gin.Context) {
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
|
|
|
|
"object": "list",
|
|
|
|
|
|
"data": antigravity.DefaultModels(),
|
|
|
|
|
|
})
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 22:24:46 +08:00
|
|
|
|
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
|
|
|
|
|
|
}
|
2026-02-03 16:55:13 +08:00
|
|
|
|
|
2026-02-02 18:30:06 +08:00
|
|
|
|
// Usage handles getting account balance and usage statistics for CC Switch integration
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// GET /v1/usage
|
|
|
|
|
|
func (h *GatewayHandler) Usage(c *gin.Context) {
|
2026-01-04 19:27:53 +08:00
|
|
|
|
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if !ok {
|
|
|
|
|
|
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-26 15:40:24 +08:00
|
|
|
|
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if !ok {
|
|
|
|
|
|
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-05 11:41:25 +08:00
|
|
|
|
// Best-effort: 获取用量统计(按当前 API Key 过滤),失败不影响基础响应
|
2026-02-02 18:30:06 +08:00
|
|
|
|
var usageData gin.H
|
|
|
|
|
|
if h.usageService != nil {
|
2026-02-05 11:41:25 +08:00
|
|
|
|
dashStats, err := h.usageService.GetAPIKeyDashboardStats(c.Request.Context(), apiKey.ID)
|
2026-02-02 18:30:06 +08:00
|
|
|
|
if err == nil && dashStats != nil {
|
|
|
|
|
|
usageData = 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,
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 订阅模式:返回订阅限额信息 + 用量统计
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() {
|
2025-12-26 10:42:08 +08:00
|
|
|
|
subscription, ok := middleware2.GetSubscriptionFromContext(c)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if !ok {
|
|
|
|
|
|
h.errorResponse(c, http.StatusForbidden, "subscription_error", "No active subscription")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
remaining := h.calculateSubscriptionRemaining(apiKey.Group, subscription)
|
2026-02-02 18:30:06 +08:00
|
|
|
|
resp := gin.H{
|
2025-12-18 13:50:39 +08:00
|
|
|
|
"isValid": true,
|
|
|
|
|
|
"planName": apiKey.Group.Name,
|
|
|
|
|
|
"remaining": remaining,
|
|
|
|
|
|
"unit": "USD",
|
2026-02-02 18:30:06 +08:00
|
|
|
|
"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,
|
|
|
|
|
|
"expires_at": subscription.ExpiresAt,
|
|
|
|
|
|
},
|
|
|
|
|
|
}
|
|
|
|
|
|
if usageData != nil {
|
|
|
|
|
|
resp["usage"] = usageData
|
|
|
|
|
|
}
|
|
|
|
|
|
c.JSON(http.StatusOK, resp)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-02 18:30:06 +08:00
|
|
|
|
// 余额模式:返回钱包余额 + 用量统计
|
2025-12-26 15:40:24 +08:00
|
|
|
|
latestUser, err := h.userService.GetByID(c.Request.Context(), subject.UserID)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
|
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Failed to get user info")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-02 18:30:06 +08:00
|
|
|
|
resp := gin.H{
|
2025-12-18 13:50:39 +08:00
|
|
|
|
"isValid": true,
|
|
|
|
|
|
"planName": "钱包余额",
|
|
|
|
|
|
"remaining": latestUser.Balance,
|
|
|
|
|
|
"unit": "USD",
|
2026-02-02 18:30:06 +08:00
|
|
|
|
"balance": latestUser.Balance,
|
|
|
|
|
|
}
|
|
|
|
|
|
if usageData != nil {
|
|
|
|
|
|
resp["usage"] = usageData
|
|
|
|
|
|
}
|
|
|
|
|
|
c.JSON(http.StatusOK, resp)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// calculateSubscriptionRemaining 计算订阅剩余可用额度
|
|
|
|
|
|
// 逻辑:
|
|
|
|
|
|
// 1. 如果日/周/月任一限额达到100%,返回0
|
|
|
|
|
|
// 2. 否则返回所有已配置周期中剩余额度的最小值
|
2025-12-26 15:40:24 +08:00
|
|
|
|
func (h *GatewayHandler) calculateSubscriptionRemaining(group *service.Group, sub *service.UserSubscription) float64 {
|
2025-12-18 13:50:39 +08:00
|
|
|
|
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 errors with proper 429 response
|
|
|
|
|
|
func (h *GatewayHandler) handleConcurrencyError(c *gin.Context, err error, slotType string, streamStarted bool) {
|
|
|
|
|
|
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error",
|
|
|
|
|
|
fmt.Sprintf("Concurrency limit exceeded for %s, please retry later", slotType), streamStarted)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-09 06:57:07 +08:00
|
|
|
|
// needForceCacheBilling 判断 failover 时是否需要强制缓存计费
|
|
|
|
|
|
// 粘性会话切换账号、或上游明确标记时,将 input_tokens 转为 cache_read 计费
|
|
|
|
|
|
func needForceCacheBilling(hasBoundSession bool, failoverErr *service.UpstreamFailoverError) bool {
|
|
|
|
|
|
return hasBoundSession || (failoverErr != nil && failoverErr.ForceCacheBilling)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-10 00:53:54 +08:00
|
|
|
|
const (
|
|
|
|
|
|
// maxSameAccountRetries 同账号重试次数上限(针对 RetryableOnSameAccount 错误)
|
|
|
|
|
|
maxSameAccountRetries = 2
|
|
|
|
|
|
// sameAccountRetryDelay 同账号重试间隔
|
|
|
|
|
|
sameAccountRetryDelay = 500 * time.Millisecond
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
// sleepSameAccountRetryDelay 同账号重试固定延时,返回 false 表示 context 已取消。
|
|
|
|
|
|
func sleepSameAccountRetryDelay(ctx context.Context) bool {
|
|
|
|
|
|
select {
|
|
|
|
|
|
case <-ctx.Done():
|
|
|
|
|
|
return false
|
|
|
|
|
|
case <-time.After(sameAccountRetryDelay):
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-09 06:56:29 +08:00
|
|
|
|
// sleepFailoverDelay 账号切换线性递增延时:第1次0s、第2次1s、第3次2s…
|
|
|
|
|
|
// 返回 false 表示 context 已取消。
|
|
|
|
|
|
func sleepFailoverDelay(ctx context.Context, switchCount int) bool {
|
|
|
|
|
|
delay := time.Duration(switchCount-1) * time.Second
|
|
|
|
|
|
if delay <= 0 {
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
select {
|
|
|
|
|
|
case <-ctx.Done():
|
|
|
|
|
|
return false
|
|
|
|
|
|
case <-time.After(delay):
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-09 14:26:01 +08:00
|
|
|
|
// sleepAntigravitySingleAccountBackoff Antigravity 平台单账号分组的 503 退避重试延时。
|
|
|
|
|
|
// 当分组内只有一个可用账号且上游返回 503(MODEL_CAPACITY_EXHAUSTED)时使用,
|
|
|
|
|
|
// 采用短固定延时策略。Service 层在 SingleAccountRetry 模式下已经做了充分的原地重试
|
|
|
|
|
|
// (最多 3 次、总等待 30s),所以 Handler 层的退避只需短暂等待即可。
|
|
|
|
|
|
// 返回 false 表示 context 已取消。
|
|
|
|
|
|
func sleepAntigravitySingleAccountBackoff(ctx context.Context, retryCount int) bool {
|
|
|
|
|
|
// 固定短延时:2s
|
|
|
|
|
|
// Service 层已经在原地等待了足够长的时间(retryDelay × 重试次数),
|
|
|
|
|
|
// Handler 层只需短暂间隔后重新进入 Service 层即可。
|
|
|
|
|
|
const delay = 2 * time.Second
|
|
|
|
|
|
|
2026-02-12 19:01:09 +08:00
|
|
|
|
logger.L().With(
|
|
|
|
|
|
zap.String("component", "handler.gateway.failover"),
|
|
|
|
|
|
zap.Duration("delay", delay),
|
|
|
|
|
|
zap.Int("retry_count", retryCount),
|
|
|
|
|
|
).Info("gateway.single_account_backoff_waiting")
|
2026-02-09 14:26:01 +08:00
|
|
|
|
|
|
|
|
|
|
select {
|
|
|
|
|
|
case <-ctx.Done():
|
|
|
|
|
|
return false
|
|
|
|
|
|
case <-time.After(delay):
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-05 21:52:54 +08:00
|
|
|
|
func (h *GatewayHandler) handleFailoverExhausted(c *gin.Context, failoverErr *service.UpstreamFailoverError, platform string, streamStarted bool) {
|
|
|
|
|
|
statusCode := failoverErr.StatusCode
|
|
|
|
|
|
responseBody := failoverErr.ResponseBody
|
|
|
|
|
|
|
|
|
|
|
|
// 先检查透传规则
|
|
|
|
|
|
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
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-10 20:56:01 +08:00
|
|
|
|
if rule.SkipMonitoring {
|
|
|
|
|
|
c.Set(service.OpsSkipPassthroughKey, true)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-05 21:52:54 +08:00
|
|
|
|
h.handleStreamingAwareError(c, respCode, "upstream_error", msg, streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 使用默认的错误映射
|
|
|
|
|
|
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) {
|
2025-12-27 11:44:00 +08:00
|
|
|
|
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
|
|
|
|
|
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"
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// 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 {
|
|
|
|
|
|
// Stream already started, send error as SSE event then close
|
|
|
|
|
|
flusher, ok := c.Writer.(http.Flusher)
|
|
|
|
|
|
if ok {
|
2026-01-01 19:47:26 -08:00
|
|
|
|
// Send error event in SSE format with proper JSON marshaling
|
|
|
|
|
|
errorData := map[string]any{
|
|
|
|
|
|
"type": "error",
|
|
|
|
|
|
"error": map[string]string{
|
|
|
|
|
|
"type": errType,
|
|
|
|
|
|
"message": message,
|
|
|
|
|
|
},
|
|
|
|
|
|
}
|
|
|
|
|
|
jsonBytes, err := json.Marshal(errorData)
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
_ = c.Error(err)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
errorEvent := fmt.Sprintf("data: %s\n\n", string(jsonBytes))
|
2025-12-20 15:29:52 +08:00
|
|
|
|
if _, err := fmt.Fprint(c.Writer, errorEvent); err != nil {
|
|
|
|
|
|
_ = c.Error(err)
|
|
|
|
|
|
}
|
2025-12-18 13:50:39 +08:00
|
|
|
|
flusher.Flush()
|
|
|
|
|
|
}
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// Normal case: return JSON response with proper status code
|
|
|
|
|
|
h.errorResponse(c, status, errType, message)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-14 11:23:10 +08:00
|
|
|
|
// ensureForwardErrorResponse 在 Forward 返回错误但尚未写响应时补写统一错误响应。
|
|
|
|
|
|
func (h *GatewayHandler) ensureForwardErrorResponse(c *gin.Context, streamStarted bool) bool {
|
|
|
|
|
|
if c == nil || c.Writer == nil || c.Writer.Written() {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
h.handleStreamingAwareError(c, http.StatusBadGateway, "upstream_error", "Upstream request failed", streamStarted)
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// 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,
|
|
|
|
|
|
},
|
|
|
|
|
|
})
|
|
|
|
|
|
}
|
2025-12-19 11:12:41 +08:00
|
|
|
|
|
|
|
|
|
|
// CountTokens handles token counting endpoint
|
|
|
|
|
|
// POST /v1/messages/count_tokens
|
|
|
|
|
|
// 特点:校验订阅/余额,但不计算并发、不记录使用量
|
|
|
|
|
|
func (h *GatewayHandler) CountTokens(c *gin.Context) {
|
|
|
|
|
|
// 从context获取apiKey和user(ApiKeyAuth中间件已设置)
|
2026-01-04 19:27:53 +08:00
|
|
|
|
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
2025-12-19 11:12:41 +08:00
|
|
|
|
if !ok {
|
|
|
|
|
|
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-26 15:40:24 +08:00
|
|
|
|
_, ok = middleware2.GetAuthSubjectFromContext(c)
|
2025-12-19 11:12:41 +08:00
|
|
|
|
if !ok {
|
|
|
|
|
|
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog := requestLogger(
|
|
|
|
|
|
c,
|
|
|
|
|
|
"handler.gateway.count_tokens",
|
|
|
|
|
|
zap.Int64("api_key_id", apiKey.ID),
|
|
|
|
|
|
zap.Any("group_id", apiKey.GroupID),
|
|
|
|
|
|
)
|
2025-12-19 11:12:41 +08:00
|
|
|
|
|
|
|
|
|
|
// 读取请求体
|
|
|
|
|
|
body, err := io.ReadAll(c.Request.Body)
|
|
|
|
|
|
if err != nil {
|
2025-12-31 08:50:12 +08:00
|
|
|
|
if maxErr, ok := extractMaxBytesError(err); ok {
|
|
|
|
|
|
h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2025-12-19 11:12:41 +08:00
|
|
|
|
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
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-15 19:17:07 +08:00
|
|
|
|
// 检查是否为 Claude Code 客户端,设置到 context 中
|
|
|
|
|
|
SetClaudeCodeClientContext(c, body)
|
|
|
|
|
|
|
2026-01-09 20:56:37 +08:00
|
|
|
|
setOpsRequestContext(c, "", false, body)
|
|
|
|
|
|
|
2026-02-09 06:47:22 +08:00
|
|
|
|
parsedReq, err := service.ParseGatewayRequest(body, domain.PlatformAnthropic)
|
2025-12-31 08:50:12 +08:00
|
|
|
|
if err != nil {
|
2025-12-19 11:12:41 +08:00
|
|
|
|
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog = reqLog.With(zap.String("model", parsedReq.Model), zap.Bool("stream", parsedReq.Stream))
|
2026-02-07 12:31:10 +08:00
|
|
|
|
// 在请求上下文中记录 thinking 状态,供 Antigravity 最终模型 key 推导/模型维度限流使用
|
|
|
|
|
|
c.Request = c.Request.WithContext(context.WithValue(c.Request.Context(), ctxkey.ThinkingEnabled, parsedReq.ThinkingEnabled))
|
2025-12-19 11:12:41 +08:00
|
|
|
|
|
2025-12-31 16:17:45 +08:00
|
|
|
|
// 验证 model 必填
|
|
|
|
|
|
if parsedReq.Model == "" {
|
|
|
|
|
|
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-09 20:56:37 +08:00
|
|
|
|
setOpsRequestContext(c, parsedReq.Model, parsedReq.Stream, body)
|
|
|
|
|
|
|
2025-12-19 11:12:41 +08:00
|
|
|
|
// 获取订阅信息(可能为nil)
|
2025-12-26 10:42:08 +08:00
|
|
|
|
subscription, _ := middleware2.GetSubscriptionFromContext(c)
|
2025-12-19 11:12:41 +08:00
|
|
|
|
|
|
|
|
|
|
// 校验 billing eligibility(订阅/余额)
|
|
|
|
|
|
// 【注意】不计算并发,但需要校验订阅/余额
|
2025-12-26 15:40:24 +08:00
|
|
|
|
if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription); err != nil {
|
2026-01-02 17:40:57 +08:00
|
|
|
|
status, code, message := billingErrorDetails(err)
|
|
|
|
|
|
h.errorResponse(c, status, code, message)
|
2025-12-19 11:12:41 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 计算粘性会话 hash
|
2026-02-09 06:46:32 +08:00
|
|
|
|
parsedReq.SessionContext = &service.SessionContext{
|
|
|
|
|
|
ClientIP: ip.GetClientIP(c),
|
|
|
|
|
|
UserAgent: c.GetHeader("User-Agent"),
|
|
|
|
|
|
APIKeyID: apiKey.ID,
|
|
|
|
|
|
}
|
2025-12-31 08:50:12 +08:00
|
|
|
|
sessionHash := h.gatewayService.GenerateSessionHash(parsedReq)
|
2025-12-19 11:12:41 +08:00
|
|
|
|
|
|
|
|
|
|
// 选择支持该模型的账号
|
2025-12-31 08:50:12 +08:00
|
|
|
|
account, err := h.gatewayService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, sessionHash, parsedReq.Model)
|
2025-12-19 11:12:41 +08:00
|
|
|
|
if err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Warn("gateway.count_tokens_select_account_failed", zap.Error(err))
|
2026-02-07 19:46:42 +08:00
|
|
|
|
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "Service temporarily unavailable")
|
2025-12-19 11:12:41 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-02-12 16:27:29 +08:00
|
|
|
|
setOpsSelectedAccount(c, account.ID, account.Platform)
|
2025-12-19 11:12:41 +08:00
|
|
|
|
|
|
|
|
|
|
// 转发请求(不记录使用量)
|
2025-12-31 08:50:12 +08:00
|
|
|
|
if err := h.gatewayService.ForwardCountTokens(c.Request.Context(), c, account, parsedReq); err != nil {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
reqLog.Error("gateway.count_tokens_forward_failed", zap.Int64("account_id", account.ID), zap.Error(err))
|
2025-12-19 11:12:41 +08:00
|
|
|
|
// 错误响应已在 ForwardCountTokens 中处理
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
2025-12-19 16:39:25 +08:00
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// InterceptType 表示请求拦截类型
|
|
|
|
|
|
type InterceptType int
|
|
|
|
|
|
|
|
|
|
|
|
const (
|
2026-02-07 19:04:08 +08:00
|
|
|
|
InterceptTypeNone InterceptType = iota
|
|
|
|
|
|
InterceptTypeWarmup // 预热请求(返回 "New Conversation")
|
|
|
|
|
|
InterceptTypeSuggestionMode // SUGGESTION MODE(返回空字符串)
|
|
|
|
|
|
InterceptTypeMaxTokensOneHaiku // max_tokens=1 + haiku 探测请求(返回 "#")
|
2026-01-23 16:57:25 +08:00
|
|
|
|
)
|
|
|
|
|
|
|
2026-02-07 19:04:08 +08:00
|
|
|
|
// isHaikuModel 检查模型名称是否包含 "haiku"(大小写不敏感)
|
|
|
|
|
|
func isHaikuModel(model string) bool {
|
|
|
|
|
|
return strings.Contains(strings.ToLower(model), "haiku")
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// isMaxTokensOneHaikuRequest 检查是否为 max_tokens=1 + haiku 模型的探测请求
|
|
|
|
|
|
// 这类请求用于 Claude Code 验证 API 连通性
|
|
|
|
|
|
// 条件:max_tokens == 1 且 model 包含 "haiku" 且非流式请求
|
|
|
|
|
|
func isMaxTokensOneHaikuRequest(model string, maxTokens int, isStream bool) bool {
|
|
|
|
|
|
return maxTokens == 1 && isHaikuModel(model) && !isStream
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// detectInterceptType 检测请求是否需要拦截,返回拦截类型
|
2026-02-07 19:04:08 +08:00
|
|
|
|
// 参数说明:
|
|
|
|
|
|
// - body: 请求体字节
|
|
|
|
|
|
// - model: 请求的模型名称
|
|
|
|
|
|
// - maxTokens: max_tokens 值
|
|
|
|
|
|
// - isStream: 是否为流式请求
|
|
|
|
|
|
// - isClaudeCodeClient: 是否已通过 Claude Code 客户端校验
|
|
|
|
|
|
func detectInterceptType(body []byte, model string, maxTokens int, isStream bool, isClaudeCodeClient bool) InterceptType {
|
|
|
|
|
|
// 优先检查 max_tokens=1 + haiku 探测请求(仅非流式)
|
|
|
|
|
|
if isClaudeCodeClient && isMaxTokensOneHaikuRequest(model, maxTokens, isStream) {
|
|
|
|
|
|
return InterceptTypeMaxTokensOneHaiku
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// 快速检查:如果不包含任何关键字,直接返回
|
2025-12-19 16:39:25 +08:00
|
|
|
|
bodyStr := string(body)
|
2026-01-23 16:57:25 +08:00
|
|
|
|
hasSuggestionMode := strings.Contains(bodyStr, "[SUGGESTION MODE:")
|
|
|
|
|
|
hasWarmupKeyword := strings.Contains(bodyStr, "title") || strings.Contains(bodyStr, "Warmup")
|
|
|
|
|
|
|
|
|
|
|
|
if !hasSuggestionMode && !hasWarmupKeyword {
|
|
|
|
|
|
return InterceptTypeNone
|
2025-12-19 16:39:25 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// 解析请求(只解析一次)
|
2025-12-19 16:39:25 +08:00
|
|
|
|
var req struct {
|
|
|
|
|
|
Messages []struct {
|
2026-01-23 16:57:25 +08:00
|
|
|
|
Role string `json:"role"`
|
2025-12-19 16:39:25 +08:00
|
|
|
|
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 {
|
2026-01-23 16:57:25 +08:00
|
|
|
|
return InterceptTypeNone
|
2025-12-19 16:39:25 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// 检查 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
|
2025-12-19 16:39:25 +08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// 检查 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
|
|
|
|
|
|
}
|
2025-12-19 16:39:25 +08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
return InterceptTypeNone
|
2025-12-19 16:39:25 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// sendMockInterceptStream 发送流式 mock 响应(用于请求拦截)
|
|
|
|
|
|
func sendMockInterceptStream(c *gin.Context, model string, interceptType InterceptType) {
|
2025-12-19 16:39:25 +08:00
|
|
|
|
c.Header("Content-Type", "text/event-stream")
|
|
|
|
|
|
c.Header("Cache-Control", "no-cache")
|
|
|
|
|
|
c.Header("Connection", "keep-alive")
|
|
|
|
|
|
c.Header("X-Accel-Buffering", "no")
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// 根据拦截类型决定响应内容
|
|
|
|
|
|
var msgID string
|
|
|
|
|
|
var outputTokens int
|
|
|
|
|
|
var textDeltas []string
|
|
|
|
|
|
|
|
|
|
|
|
switch interceptType {
|
|
|
|
|
|
case InterceptTypeSuggestionMode:
|
|
|
|
|
|
msgID = "msg_mock_suggestion"
|
|
|
|
|
|
outputTokens = 1
|
|
|
|
|
|
textDeltas = []string{""} // 空内容
|
|
|
|
|
|
default: // InterceptTypeWarmup
|
|
|
|
|
|
msgID = "msg_mock_warmup"
|
|
|
|
|
|
outputTokens = 2
|
|
|
|
|
|
textDeltas = []string{"New", " Conversation"}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-01 19:47:26 -08:00
|
|
|
|
// Build message_start event with proper JSON marshaling
|
|
|
|
|
|
messageStart := map[string]any{
|
|
|
|
|
|
"type": "message_start",
|
|
|
|
|
|
"message": map[string]any{
|
2026-01-23 16:57:25 +08:00
|
|
|
|
"id": msgID,
|
2026-01-01 19:47:26 -08:00
|
|
|
|
"type": "message",
|
|
|
|
|
|
"role": "assistant",
|
|
|
|
|
|
"model": model,
|
|
|
|
|
|
"content": []any{},
|
|
|
|
|
|
"stop_reason": nil,
|
|
|
|
|
|
"stop_sequence": nil,
|
|
|
|
|
|
"usage": map[string]int{
|
|
|
|
|
|
"input_tokens": 10,
|
|
|
|
|
|
"output_tokens": 0,
|
|
|
|
|
|
},
|
|
|
|
|
|
},
|
|
|
|
|
|
}
|
|
|
|
|
|
messageStartJSON, _ := json.Marshal(messageStart)
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// Build events
|
2025-12-19 16:39:25 +08:00
|
|
|
|
events := []string{
|
2026-01-01 19:47:26 -08:00
|
|
|
|
`event: message_start` + "\n" + `data: ` + string(messageStartJSON),
|
2025-12-19 16:39:25 +08:00
|
|
|
|
`event: content_block_start` + "\n" + `data: {"content_block":{"text":"","type":"text"},"index":0,"type":"content_block_start"}`,
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// Add text deltas
|
|
|
|
|
|
for _, text := range textDeltas {
|
|
|
|
|
|
delta := map[string]any{
|
|
|
|
|
|
"type": "content_block_delta",
|
|
|
|
|
|
"index": 0,
|
|
|
|
|
|
"delta": map[string]string{
|
|
|
|
|
|
"type": "text_delta",
|
|
|
|
|
|
"text": text,
|
|
|
|
|
|
},
|
|
|
|
|
|
}
|
|
|
|
|
|
deltaJSON, _ := json.Marshal(delta)
|
|
|
|
|
|
events = append(events, `event: content_block_delta`+"\n"+`data: `+string(deltaJSON))
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// Add final events
|
|
|
|
|
|
messageDelta := map[string]any{
|
|
|
|
|
|
"type": "message_delta",
|
|
|
|
|
|
"delta": map[string]any{
|
|
|
|
|
|
"stop_reason": "end_turn",
|
|
|
|
|
|
"stop_sequence": nil,
|
|
|
|
|
|
},
|
|
|
|
|
|
"usage": map[string]int{
|
|
|
|
|
|
"input_tokens": 10,
|
|
|
|
|
|
"output_tokens": outputTokens,
|
|
|
|
|
|
},
|
|
|
|
|
|
}
|
|
|
|
|
|
messageDeltaJSON, _ := json.Marshal(messageDelta)
|
|
|
|
|
|
|
|
|
|
|
|
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"}`,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
2025-12-19 16:39:25 +08:00
|
|
|
|
for _, event := range events {
|
|
|
|
|
|
_, _ = c.Writer.WriteString(event + "\n\n")
|
|
|
|
|
|
c.Writer.Flush()
|
|
|
|
|
|
time.Sleep(20 * time.Millisecond)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-07 19:04:08 +08:00
|
|
|
|
// generateRealisticMsgID 生成仿真的消息 ID(msg_bdrk_XXXXXXX 格式)
|
|
|
|
|
|
// 格式与 Claude API 真实响应一致,24 位随机字母数字
|
|
|
|
|
|
func generateRealisticMsgID() string {
|
|
|
|
|
|
const charset = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
|
|
|
|
|
const idLen = 24
|
|
|
|
|
|
randomBytes := make([]byte, idLen)
|
|
|
|
|
|
if _, err := rand.Read(randomBytes); err != nil {
|
|
|
|
|
|
return fmt.Sprintf("msg_bdrk_%d", time.Now().UnixNano())
|
|
|
|
|
|
}
|
|
|
|
|
|
b := make([]byte, idLen)
|
|
|
|
|
|
for i := range b {
|
|
|
|
|
|
b[i] = charset[int(randomBytes[i])%len(charset)]
|
|
|
|
|
|
}
|
|
|
|
|
|
return "msg_bdrk_" + string(b)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-23 16:57:25 +08:00
|
|
|
|
// sendMockInterceptResponse 发送非流式 mock 响应(用于请求拦截)
|
|
|
|
|
|
func sendMockInterceptResponse(c *gin.Context, model string, interceptType InterceptType) {
|
2026-02-07 19:04:08 +08:00
|
|
|
|
var msgID, text, stopReason string
|
2026-01-23 16:57:25 +08:00
|
|
|
|
var outputTokens int
|
|
|
|
|
|
|
|
|
|
|
|
switch interceptType {
|
|
|
|
|
|
case InterceptTypeSuggestionMode:
|
|
|
|
|
|
msgID = "msg_mock_suggestion"
|
|
|
|
|
|
text = ""
|
|
|
|
|
|
outputTokens = 1
|
2026-02-07 19:04:08 +08:00
|
|
|
|
stopReason = "end_turn"
|
|
|
|
|
|
case InterceptTypeMaxTokensOneHaiku:
|
|
|
|
|
|
msgID = generateRealisticMsgID()
|
|
|
|
|
|
text = "#"
|
|
|
|
|
|
outputTokens = 1
|
|
|
|
|
|
stopReason = "max_tokens" // max_tokens=1 探测请求的 stop_reason 应为 max_tokens
|
2026-01-23 16:57:25 +08:00
|
|
|
|
default: // InterceptTypeWarmup
|
|
|
|
|
|
msgID = "msg_mock_warmup"
|
|
|
|
|
|
text = "New Conversation"
|
|
|
|
|
|
outputTokens = 2
|
2026-02-07 19:04:08 +08:00
|
|
|
|
stopReason = "end_turn"
|
2026-01-23 16:57:25 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-07 19:04:08 +08:00
|
|
|
|
// 构建完整的响应格式(与 Claude 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,
|
2025-12-19 16:39:25 +08:00
|
|
|
|
"usage": gin.H{
|
2026-02-07 19:04:08 +08:00
|
|
|
|
"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,
|
|
|
|
|
|
},
|
2026-01-23 16:57:25 +08:00
|
|
|
|
"output_tokens": outputTokens,
|
2026-02-07 19:04:08 +08:00
|
|
|
|
"total_tokens": 10 + outputTokens,
|
2025-12-19 16:39:25 +08:00
|
|
|
|
},
|
2026-02-07 19:04:08 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
c.JSON(http.StatusOK, response)
|
2025-12-19 16:39:25 +08:00
|
|
|
|
}
|
2026-01-02 17:40:57 +08:00
|
|
|
|
|
|
|
|
|
|
func billingErrorDetails(err error) (status int, code, message string) {
|
|
|
|
|
|
if errors.Is(err, service.ErrBillingServiceUnavailable) {
|
2026-01-03 11:36:31 +08:00
|
|
|
|
msg := pkgerrors.Message(err)
|
2026-01-02 17:40:57 +08:00
|
|
|
|
if msg == "" {
|
|
|
|
|
|
msg = "Billing service temporarily unavailable. Please retry later."
|
|
|
|
|
|
}
|
|
|
|
|
|
return http.StatusServiceUnavailable, "billing_service_error", msg
|
|
|
|
|
|
}
|
2026-01-03 11:36:31 +08:00
|
|
|
|
msg := pkgerrors.Message(err)
|
2026-01-02 17:40:57 +08:00
|
|
|
|
if msg == "" {
|
2026-02-12 19:01:09 +08:00
|
|
|
|
logger.L().With(
|
|
|
|
|
|
zap.String("component", "handler.gateway.billing"),
|
|
|
|
|
|
zap.Error(err),
|
|
|
|
|
|
).Warn("gateway.billing_error_missing_message")
|
2026-02-07 19:46:42 +08:00
|
|
|
|
msg = "Billing error"
|
2026-01-02 17:40:57 +08:00
|
|
|
|
}
|
|
|
|
|
|
return http.StatusForbidden, "billing_error", msg
|
|
|
|
|
|
}
|
2026-02-22 12:56:57 +08:00
|
|
|
|
|
|
|
|
|
|
func (h *GatewayHandler) submitUsageRecordTask(task service.UsageRecordTask) {
|
|
|
|
|
|
if task == nil {
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
if h.usageRecordWorkerPool != nil {
|
|
|
|
|
|
h.usageRecordWorkerPool.Submit(task)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
// 回退路径:worker 池未注入时同步执行,避免退回到无界 goroutine 模式。
|
|
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
|
|
|
|
defer cancel()
|
|
|
|
|
|
task(ctx)
|
|
|
|
|
|
}
|