2025-12-18 13:50:39 +08:00
|
|
|
|
package handler
|
|
|
|
|
|
|
|
|
|
|
|
import (
|
|
|
|
|
|
"context"
|
|
|
|
|
|
"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-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-04 22:10:32 +08:00
|
|
|
|
pkgerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
2026-01-09 21:59:32 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
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"
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
// 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
|
|
|
|
|
|
concurrencyHelper *ConcurrencyHelper
|
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-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)
|
|
|
|
|
|
if cfg != nil {
|
|
|
|
|
|
pingInterval = time.Duration(cfg.Concurrency.PingInterval) * time.Second
|
|
|
|
|
|
}
|
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-01-04 19:49:59 +08:00
|
|
|
|
concurrencyHelper: NewConcurrencyHelper(concurrencyService, SSEPingFormatClaude, pingInterval),
|
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
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 读取请求体
|
|
|
|
|
|
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
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-31 08:50:12 +08:00
|
|
|
|
parsedReq, err := service.ParseGatewayRequest(body)
|
|
|
|
|
|
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
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
2026-01-08 23:07:00 +08:00
|
|
|
|
// 设置 Claude Code 客户端标识到 context(用于分组限制检查)
|
|
|
|
|
|
SetClaudeCodeClientContext(c, body)
|
|
|
|
|
|
|
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
|
|
|
|
|
|
|
|
|
|
|
|
// 获取订阅信息(可能为nil)- 提前获取用于后续检查
|
2025-12-26 10:42:08 +08:00
|
|
|
|
subscription, _ := middleware2.GetSubscriptionFromContext(c)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
2026-01-06 16:23:56 +08:00
|
|
|
|
// 获取 User-Agent
|
|
|
|
|
|
userAgent := c.Request.UserAgent()
|
|
|
|
|
|
|
2026-01-09 21:59:32 +08:00
|
|
|
|
// 获取客户端 IP
|
|
|
|
|
|
clientIP := ip.GetClientIP(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)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
|
log.Printf("Increment wait count failed: %v", err)
|
|
|
|
|
|
// On error, allow request to proceed
|
|
|
|
|
|
} else if !canWait {
|
|
|
|
|
|
h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
// 确保在函数退出时减少wait计数
|
2025-12-26 15:40:24 +08:00
|
|
|
|
defer 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 {
|
|
|
|
|
|
log.Printf("User concurrency acquire failed: %v", err)
|
|
|
|
|
|
h.handleConcurrencyError(c, err, "user", streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
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 {
|
2025-12-18 13:50:39 +08:00
|
|
|
|
log.Printf("Billing eligibility check failed after wait: %v", 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
|
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
|
|
|
|
|
2025-12-26 22:07:55 +08:00
|
|
|
|
if platform == service.PlatformGemini {
|
2025-12-27 12:27:47 +08:00
|
|
|
|
const maxAccountSwitches = 3
|
|
|
|
|
|
switchCount := 0
|
|
|
|
|
|
failedAccountIDs := make(map[int64]struct{})
|
|
|
|
|
|
lastFailoverStatus := 0
|
|
|
|
|
|
|
|
|
|
|
|
for {
|
2026-01-01 04:01:51 +08:00
|
|
|
|
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, failedAccountIDs)
|
2025-12-27 12:27:47 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
|
if len(failedAccountIDs) == 0 {
|
|
|
|
|
|
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-01 04:01:51 +08:00
|
|
|
|
account := selection.Account
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
2025-12-27 12:27:47 +08:00
|
|
|
|
// 检查预热请求拦截(在账号选择后、转发前检查)
|
|
|
|
|
|
if account.IsInterceptWarmupEnabled() && isWarmupRequest(body) {
|
2026-01-01 04:01:51 +08:00
|
|
|
|
if selection.Acquired && selection.ReleaseFunc != nil {
|
|
|
|
|
|
selection.ReleaseFunc()
|
|
|
|
|
|
}
|
2025-12-31 08:50:12 +08:00
|
|
|
|
if reqStream {
|
|
|
|
|
|
sendMockWarmupStream(c, reqModel)
|
2025-12-27 12:27:47 +08:00
|
|
|
|
} else {
|
2025-12-31 08:50:12 +08:00
|
|
|
|
sendMockWarmupResponse(c, reqModel)
|
2025-12-27 12:27:47 +08:00
|
|
|
|
}
|
|
|
|
|
|
return
|
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
|
|
|
|
|
|
var accountWaitRelease func()
|
|
|
|
|
|
if !selection.Acquired {
|
|
|
|
|
|
if selection.WaitPlan == nil {
|
|
|
|
|
|
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts", streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
canWait, err := h.concurrencyHelper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
log.Printf("Increment account wait count failed: %v", err)
|
|
|
|
|
|
} else if !canWait {
|
|
|
|
|
|
log.Printf("Account wait queue full: account=%d", account.ID)
|
|
|
|
|
|
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later", streamStarted)
|
|
|
|
|
|
return
|
2026-01-01 04:26:01 +08:00
|
|
|
|
} else {
|
|
|
|
|
|
// Only set release function if increment succeeded
|
|
|
|
|
|
accountWaitRelease = func() {
|
|
|
|
|
|
h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
|
|
|
|
|
}
|
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 {
|
|
|
|
|
|
if accountWaitRelease != nil {
|
|
|
|
|
|
accountWaitRelease()
|
|
|
|
|
|
}
|
|
|
|
|
|
log.Printf("Account concurrency acquire failed: %v", err)
|
|
|
|
|
|
h.handleConcurrencyError(c, err, "account", streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-08 23:07:00 +08:00
|
|
|
|
if err := h.gatewayService.BindStickySession(c.Request.Context(), apiKey.GroupID, sessionKey, account.ID); err != nil {
|
2026-01-01 04:01:51 +08:00
|
|
|
|
log.Printf("Bind sticky session failed: %v", err)
|
|
|
|
|
|
}
|
2025-12-27 12:27:47 +08:00
|
|
|
|
}
|
2026-01-04 19:49:59 +08:00
|
|
|
|
// 账号槽位/等待计数需要在超时或断开时安全回收
|
|
|
|
|
|
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
|
|
|
|
|
|
accountWaitRelease = wrapReleaseOnDone(c.Request.Context(), accountWaitRelease)
|
2025-12-27 12:27:47 +08:00
|
|
|
|
|
2025-12-28 17:48:52 +08:00
|
|
|
|
// 转发请求 - 根据账号平台分流
|
|
|
|
|
|
var result *service.ForwardResult
|
|
|
|
|
|
if account.Platform == service.PlatformAntigravity {
|
2025-12-31 08:50:12 +08:00
|
|
|
|
result, err = h.antigravityGatewayService.ForwardGemini(c.Request.Context(), c, account, reqModel, "generateContent", reqStream, body)
|
2025-12-28 17:48:52 +08:00
|
|
|
|
} else {
|
|
|
|
|
|
result, err = h.geminiCompatService.Forward(c.Request.Context(), c, account, body)
|
|
|
|
|
|
}
|
2025-12-27 12:27:47 +08:00
|
|
|
|
if accountReleaseFunc != nil {
|
|
|
|
|
|
accountReleaseFunc()
|
|
|
|
|
|
}
|
2026-01-01 04:01:51 +08:00
|
|
|
|
if accountWaitRelease != nil {
|
|
|
|
|
|
accountWaitRelease()
|
|
|
|
|
|
}
|
2025-12-27 12:27:47 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
|
var failoverErr *service.UpstreamFailoverError
|
|
|
|
|
|
if errors.As(err, &failoverErr) {
|
|
|
|
|
|
failedAccountIDs[account.ID] = struct{}{}
|
|
|
|
|
|
if switchCount >= maxAccountSwitches {
|
|
|
|
|
|
lastFailoverStatus = failoverErr.StatusCode
|
|
|
|
|
|
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
lastFailoverStatus = failoverErr.StatusCode
|
|
|
|
|
|
switchCount++
|
|
|
|
|
|
log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
// 错误响应已在Forward中处理,这里只记录日志
|
|
|
|
|
|
log.Printf("Forward request failed: %v", err)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2025-12-18 13:50:39 +08:00
|
|
|
|
|
2025-12-27 12:27:47 +08:00
|
|
|
|
// 异步记录使用量(subscription已在函数开头获取)
|
2026-01-09 21:59:32 +08:00
|
|
|
|
go func(result *service.ForwardResult, usedAccount *service.Account, ua string, cip string) {
|
2025-12-27 12:27:47 +08:00
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
|
|
|
|
defer cancel()
|
|
|
|
|
|
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
|
|
|
|
|
Result: result,
|
2026-01-04 19:27:53 +08:00
|
|
|
|
APIKey: apiKey,
|
2025-12-27 12:27:47 +08:00
|
|
|
|
User: apiKey.User,
|
|
|
|
|
|
Account: usedAccount,
|
|
|
|
|
|
Subscription: subscription,
|
2026-01-06 16:23:56 +08:00
|
|
|
|
UserAgent: ua,
|
2026-01-09 21:59:32 +08:00
|
|
|
|
IPAddress: cip,
|
2025-12-27 12:27:47 +08:00
|
|
|
|
}); err != nil {
|
|
|
|
|
|
log.Printf("Record usage failed: %v", err)
|
|
|
|
|
|
}
|
2026-01-09 21:59:32 +08:00
|
|
|
|
}(result, account, userAgent, clientIP)
|
2025-12-27 11:44:00 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-30 10:48:55 +08:00
|
|
|
|
const maxAccountSwitches = 10
|
2025-12-27 11:44:00 +08:00
|
|
|
|
switchCount := 0
|
|
|
|
|
|
failedAccountIDs := make(map[int64]struct{})
|
|
|
|
|
|
lastFailoverStatus := 0
|
|
|
|
|
|
|
|
|
|
|
|
for {
|
|
|
|
|
|
// 选择支持该模型的账号
|
2026-01-01 04:01:51 +08:00
|
|
|
|
selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, failedAccountIDs)
|
2025-12-27 11:44:00 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
|
if len(failedAccountIDs) == 0 {
|
|
|
|
|
|
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-01 04:01:51 +08:00
|
|
|
|
account := selection.Account
|
2025-12-27 11:44:00 +08:00
|
|
|
|
|
|
|
|
|
|
// 检查预热请求拦截(在账号选择后、转发前检查)
|
|
|
|
|
|
if account.IsInterceptWarmupEnabled() && isWarmupRequest(body) {
|
2026-01-01 04:01:51 +08:00
|
|
|
|
if selection.Acquired && selection.ReleaseFunc != nil {
|
|
|
|
|
|
selection.ReleaseFunc()
|
|
|
|
|
|
}
|
2025-12-31 08:50:12 +08:00
|
|
|
|
if reqStream {
|
|
|
|
|
|
sendMockWarmupStream(c, reqModel)
|
2025-12-27 11:44:00 +08:00
|
|
|
|
} else {
|
2025-12-31 08:50:12 +08:00
|
|
|
|
sendMockWarmupResponse(c, reqModel)
|
2025-12-27 11:44:00 +08:00
|
|
|
|
}
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 3. 获取账号并发槽位
|
2026-01-01 04:01:51 +08:00
|
|
|
|
accountReleaseFunc := selection.ReleaseFunc
|
|
|
|
|
|
var accountWaitRelease func()
|
|
|
|
|
|
if !selection.Acquired {
|
|
|
|
|
|
if selection.WaitPlan == nil {
|
|
|
|
|
|
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts", streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
canWait, err := h.concurrencyHelper.IncrementAccountWaitCount(c.Request.Context(), account.ID, selection.WaitPlan.MaxWaiting)
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
log.Printf("Increment account wait count failed: %v", err)
|
|
|
|
|
|
} else if !canWait {
|
|
|
|
|
|
log.Printf("Account wait queue full: account=%d", account.ID)
|
|
|
|
|
|
h.handleStreamingAwareError(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later", streamStarted)
|
|
|
|
|
|
return
|
2026-01-01 04:26:01 +08:00
|
|
|
|
} else {
|
|
|
|
|
|
// Only set release function if increment succeeded
|
|
|
|
|
|
accountWaitRelease = func() {
|
|
|
|
|
|
h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
|
|
|
|
|
|
}
|
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 {
|
|
|
|
|
|
if accountWaitRelease != nil {
|
|
|
|
|
|
accountWaitRelease()
|
|
|
|
|
|
}
|
|
|
|
|
|
log.Printf("Account concurrency acquire failed: %v", err)
|
|
|
|
|
|
h.handleConcurrencyError(c, err, "account", streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-08 23:07:00 +08:00
|
|
|
|
if err := h.gatewayService.BindStickySession(c.Request.Context(), apiKey.GroupID, sessionKey, account.ID); err != nil {
|
2026-01-01 04:01:51 +08:00
|
|
|
|
log.Printf("Bind sticky session failed: %v", err)
|
|
|
|
|
|
}
|
2025-12-27 11:44:00 +08:00
|
|
|
|
}
|
2026-01-04 19:49:59 +08:00
|
|
|
|
// 账号槽位/等待计数需要在超时或断开时安全回收
|
|
|
|
|
|
accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
|
|
|
|
|
|
accountWaitRelease = wrapReleaseOnDone(c.Request.Context(), accountWaitRelease)
|
2025-12-27 11:44:00 +08:00
|
|
|
|
|
2025-12-28 17:48:52 +08:00
|
|
|
|
// 转发请求 - 根据账号平台分流
|
|
|
|
|
|
var result *service.ForwardResult
|
|
|
|
|
|
if account.Platform == service.PlatformAntigravity {
|
|
|
|
|
|
result, err = h.antigravityGatewayService.Forward(c.Request.Context(), c, account, body)
|
|
|
|
|
|
} else {
|
2025-12-31 08:50:12 +08:00
|
|
|
|
result, err = h.gatewayService.Forward(c.Request.Context(), c, account, parsedReq)
|
2025-12-28 17:48:52 +08:00
|
|
|
|
}
|
2025-12-27 11:44:00 +08:00
|
|
|
|
if accountReleaseFunc != nil {
|
|
|
|
|
|
accountReleaseFunc()
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
2026-01-01 04:01:51 +08:00
|
|
|
|
if accountWaitRelease != nil {
|
|
|
|
|
|
accountWaitRelease()
|
|
|
|
|
|
}
|
2025-12-27 11:44:00 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
|
var failoverErr *service.UpstreamFailoverError
|
|
|
|
|
|
if errors.As(err, &failoverErr) {
|
|
|
|
|
|
failedAccountIDs[account.ID] = struct{}{}
|
|
|
|
|
|
if switchCount >= maxAccountSwitches {
|
|
|
|
|
|
lastFailoverStatus = failoverErr.StatusCode
|
|
|
|
|
|
h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
lastFailoverStatus = failoverErr.StatusCode
|
|
|
|
|
|
switchCount++
|
|
|
|
|
|
log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
// 错误响应已在Forward中处理,这里只记录日志
|
2026-01-04 16:45:11 +08:00
|
|
|
|
log.Printf("Account %d: Forward request failed: %v", account.ID, err)
|
2025-12-27 11:44:00 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 异步记录使用量(subscription已在函数开头获取)
|
2026-01-09 21:59:32 +08:00
|
|
|
|
go func(result *service.ForwardResult, usedAccount *service.Account, ua string, cip string) {
|
2025-12-27 11:44:00 +08:00
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
|
|
|
|
defer cancel()
|
|
|
|
|
|
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
|
|
|
|
|
Result: result,
|
2026-01-04 19:27:53 +08:00
|
|
|
|
APIKey: apiKey,
|
2025-12-27 11:44:00 +08:00
|
|
|
|
User: apiKey.User,
|
|
|
|
|
|
Account: usedAccount,
|
|
|
|
|
|
Subscription: subscription,
|
2026-01-06 16:23:56 +08:00
|
|
|
|
UserAgent: ua,
|
2026-01-09 21:59:32 +08:00
|
|
|
|
IPAddress: cip,
|
2025-12-27 11:44:00 +08:00
|
|
|
|
}); err != nil {
|
|
|
|
|
|
log.Printf("Record usage failed: %v", err)
|
|
|
|
|
|
}
|
2026-01-09 21:59:32 +08:00
|
|
|
|
}(result, account, userAgent, clientIP)
|
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
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 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(),
|
|
|
|
|
|
})
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// Usage handles getting account balance for CC Switch integration
|
|
|
|
|
|
// 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
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 订阅模式:返回订阅限额信息
|
|
|
|
|
|
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)
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
|
|
|
|
"isValid": true,
|
|
|
|
|
|
"planName": apiKey.Group.Name,
|
|
|
|
|
|
"remaining": remaining,
|
|
|
|
|
|
"unit": "USD",
|
|
|
|
|
|
})
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 余额模式:返回钱包余额
|
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
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
|
|
|
|
|
"isValid": true,
|
|
|
|
|
|
"planName": "钱包余额",
|
|
|
|
|
|
"remaining": latestUser.Balance,
|
|
|
|
|
|
"unit": "USD",
|
|
|
|
|
|
})
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 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)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-27 11:44:00 +08:00
|
|
|
|
func (h *GatewayHandler) handleFailoverExhausted(c *gin.Context, statusCode int, streamStarted bool) {
|
|
|
|
|
|
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)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 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
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 读取请求体
|
|
|
|
|
|
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
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-31 08:50:12 +08:00
|
|
|
|
parsedReq, err := service.ParseGatewayRequest(body)
|
|
|
|
|
|
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-01-08 23:07:00 +08:00
|
|
|
|
// 设置 Claude Code 客户端标识到 context(用于分组限制检查)
|
|
|
|
|
|
SetClaudeCodeClientContext(c, body)
|
|
|
|
|
|
|
2025-12-31 16:17:45 +08:00
|
|
|
|
// 验证 model 必填
|
|
|
|
|
|
if parsedReq.Model == "" {
|
|
|
|
|
|
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
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
|
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 {
|
|
|
|
|
|
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error())
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 转发请求(不记录使用量)
|
2025-12-31 08:50:12 +08:00
|
|
|
|
if err := h.gatewayService.ForwardCountTokens(c.Request.Context(), c, account, parsedReq); err != nil {
|
2025-12-19 11:12:41 +08:00
|
|
|
|
log.Printf("Forward count_tokens request failed: %v", err)
|
|
|
|
|
|
// 错误响应已在 ForwardCountTokens 中处理
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
2025-12-19 16:39:25 +08:00
|
|
|
|
|
|
|
|
|
|
// isWarmupRequest 检测是否为预热请求(标题生成、Warmup等)
|
|
|
|
|
|
func isWarmupRequest(body []byte) bool {
|
|
|
|
|
|
// 快速检查:如果body不包含关键字,直接返回false
|
|
|
|
|
|
bodyStr := string(body)
|
|
|
|
|
|
if !strings.Contains(bodyStr, "title") && !strings.Contains(bodyStr, "Warmup") {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 解析完整请求
|
|
|
|
|
|
var req struct {
|
|
|
|
|
|
Messages []struct {
|
|
|
|
|
|
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 false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 检查 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 true
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 检查 system 中的标题提取模式
|
|
|
|
|
|
for _, system := range req.System {
|
|
|
|
|
|
if strings.Contains(system.Text, "nalyze if this message indicates a new conversation topic. If it does, extract a 2-3 word title") {
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// sendMockWarmupStream 发送流式 mock 响应(用于预热请求拦截)
|
|
|
|
|
|
func sendMockWarmupStream(c *gin.Context, model string) {
|
|
|
|
|
|
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-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{
|
|
|
|
|
|
"id": "msg_mock_warmup",
|
|
|
|
|
|
"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)
|
|
|
|
|
|
|
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"}`,
|
|
|
|
|
|
`event: content_block_delta` + "\n" + `data: {"delta":{"text":"New","type":"text_delta"},"index":0,"type":"content_block_delta"}`,
|
|
|
|
|
|
`event: content_block_delta` + "\n" + `data: {"delta":{"text":" Conversation","type":"text_delta"},"index":0,"type":"content_block_delta"}`,
|
|
|
|
|
|
`event: content_block_stop` + "\n" + `data: {"index":0,"type":"content_block_stop"}`,
|
|
|
|
|
|
`event: message_delta` + "\n" + `data: {"delta":{"stop_reason":"end_turn","stop_sequence":null},"type":"message_delta","usage":{"input_tokens":10,"output_tokens":2}}`,
|
|
|
|
|
|
`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)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// sendMockWarmupResponse 发送非流式 mock 响应(用于预热请求拦截)
|
|
|
|
|
|
func sendMockWarmupResponse(c *gin.Context, model string) {
|
|
|
|
|
|
c.JSON(http.StatusOK, gin.H{
|
2025-12-20 15:29:52 +08:00
|
|
|
|
"id": "msg_mock_warmup",
|
|
|
|
|
|
"type": "message",
|
|
|
|
|
|
"role": "assistant",
|
|
|
|
|
|
"model": model,
|
|
|
|
|
|
"content": []gin.H{{"type": "text", "text": "New Conversation"}},
|
2025-12-19 16:39:25 +08:00
|
|
|
|
"stop_reason": "end_turn",
|
|
|
|
|
|
"usage": gin.H{
|
|
|
|
|
|
"input_tokens": 10,
|
|
|
|
|
|
"output_tokens": 2,
|
|
|
|
|
|
},
|
|
|
|
|
|
})
|
|
|
|
|
|
}
|
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 == "" {
|
|
|
|
|
|
msg = err.Error()
|
|
|
|
|
|
}
|
|
|
|
|
|
return http.StatusForbidden, "billing_error", msg
|
|
|
|
|
|
}
|