first commit

backend/internal/handler/gateway_handler.go

@@ -0,0 +1,874 @@
+package handler
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"strings"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
+	"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
+	pkgerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
+	"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
+	"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
+	middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
+	"github.com/Wei-Shaw/sub2api/internal/service"
+
+	"github.com/gin-gonic/gin"
+)
+
+// GatewayHandler handles API gateway requests
+type GatewayHandler struct {
+	gatewayService            *service.GatewayService
+	geminiCompatService       *service.GeminiMessagesCompatService
+	antigravityGatewayService *service.AntigravityGatewayService
+	userService               *service.UserService
+	billingCacheService       *service.BillingCacheService
+	concurrencyHelper         *ConcurrencyHelper
+}
+
+// NewGatewayHandler creates a new GatewayHandler
+func NewGatewayHandler(
+	gatewayService *service.GatewayService,
+	geminiCompatService *service.GeminiMessagesCompatService,
+	antigravityGatewayService *service.AntigravityGatewayService,
+	userService *service.UserService,
+	concurrencyService *service.ConcurrencyService,
+	billingCacheService *service.BillingCacheService,
+	cfg *config.Config,
+) *GatewayHandler {
+	pingInterval := time.Duration(0)
+	if cfg != nil {
+		pingInterval = time.Duration(cfg.Concurrency.PingInterval) * time.Second
+	}
+	return &GatewayHandler{
+		gatewayService:            gatewayService,
+		geminiCompatService:       geminiCompatService,
+		antigravityGatewayService: antigravityGatewayService,
+		userService:               userService,
+		billingCacheService:       billingCacheService,
+		concurrencyHelper:         NewConcurrencyHelper(concurrencyService, SSEPingFormatClaude, pingInterval),
+	}
+}
+
+// Messages handles Claude API compatible messages endpoint
+// POST /v1/messages
+func (h *GatewayHandler) Messages(c *gin.Context) {
+	// 从context获取apiKey和user（ApiKeyAuth中间件已设置）
+	apiKey, ok := middleware2.GetAPIKeyFromContext(c)
+	if !ok {
+		h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
+		return
+	}
+
+	subject, ok := middleware2.GetAuthSubjectFromContext(c)
+	if !ok {
+		h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
+		return
+	}
+
+	// 读取请求体
+	body, err := io.ReadAll(c.Request.Body)
+	if err != nil {
+		if maxErr, ok := extractMaxBytesError(err); ok {
+			h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
+			return
+		}
+		h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
+		return
+	}
+
+	if len(body) == 0 {
+		h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
+		return
+	}
+
+	// 检查是否为 Claude Code 客户端，设置到 context 中
+	SetClaudeCodeClientContext(c, body)
+
+	setOpsRequestContext(c, "", false, body)
+
+	parsedReq, err := service.ParseGatewayRequest(body)
+	if err != nil {
+		h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
+		return
+	}
+	reqModel := parsedReq.Model
+	reqStream := parsedReq.Stream
+
+	setOpsRequestContext(c, reqModel, reqStream, body)
+
+	// 验证 model 必填
+	if reqModel == "" {
+		h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
+		return
+	}
+
+	// Track if we've started streaming (for error handling)
+	streamStarted := false
+
+	// 获取订阅信息（可能为nil）- 提前获取用于后续检查
+	subscription, _ := middleware2.GetSubscriptionFromContext(c)
+
+	// 0. 检查wait队列是否已满
+	maxWait := service.CalculateMaxWait(subject.Concurrency)
+	canWait, err := h.concurrencyHelper.IncrementWaitCount(c.Request.Context(), subject.UserID, maxWait)
+	waitCounted := false
+	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
+	}
+	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)
+		}
+	}()
+
+	// 1. 首先获取用户并发槽位
+	userReleaseFunc, err := h.concurrencyHelper.AcquireUserSlotWithWait(c, subject.UserID, subject.Concurrency, reqStream, &streamStarted)
+	if err != nil {
+		log.Printf("User concurrency acquire failed: %v", err)
+		h.handleConcurrencyError(c, err, "user", streamStarted)
+		return
+	}
+	// User slot acquired: no longer waiting in the queue.
+	if waitCounted {
+		h.concurrencyHelper.DecrementWaitCount(c.Request.Context(), subject.UserID)
+		waitCounted = false
+	}
+	// 在请求结束或 Context 取消时确保释放槽位，避免客户端断开造成泄漏
+	userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
+	if userReleaseFunc != nil {
+		defer userReleaseFunc()
+	}
+
+	// 2. 【新增】Wait后二次检查余额/订阅
+	if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription); err != nil {
+		log.Printf("Billing eligibility check failed after wait: %v", err)
+		status, code, message := billingErrorDetails(err)
+		h.handleStreamingAwareError(c, status, code, message, streamStarted)
+		return
+	}
+
+	// 计算粘性会话hash
+	sessionHash := h.gatewayService.GenerateSessionHash(parsedReq)
+
+	// 获取平台：优先使用强制平台（/antigravity 路由，中间件已设置 request.Context），否则使用分组平台
+	platform := ""
+	if forcePlatform, ok := middleware2.GetForcePlatformFromContext(c); ok {
+		platform = forcePlatform
+	} else if apiKey.Group != nil {
+		platform = apiKey.Group.Platform
+	}
+	sessionKey := sessionHash
+	if platform == service.PlatformGemini && sessionHash != "" {
+		sessionKey = "gemini:" + sessionHash
+	}
+
+	if platform == service.PlatformGemini {
+		const maxAccountSwitches = 3
+		switchCount := 0
+		failedAccountIDs := make(map[int64]struct{})
+		lastFailoverStatus := 0
+
+		for {
+			selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, failedAccountIDs)
+			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
+			}
+			account := selection.Account
+			setOpsSelectedAccount(c, account.ID)
+
+			// 检查预热请求拦截（在账号选择后、转发前检查）
+			if account.IsInterceptWarmupEnabled() && isWarmupRequest(body) {
+				if selection.Acquired && selection.ReleaseFunc != nil {
+					selection.ReleaseFunc()
+				}
+				if reqStream {
+					sendMockWarmupStream(c, reqModel)
+				} else {
+					sendMockWarmupResponse(c, reqModel)
+				}
+				return
+			}
+
+			// 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 {
+					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
+				}
+				if err == nil && canWait {
+					accountWaitCounted = true
+				}
+				// Ensure the wait counter is decremented if we exit before acquiring the slot.
+				defer func() {
+					if accountWaitCounted {
+						h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
+					}
+				}()
+
+				accountReleaseFunc, err = h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout(
+					c,
+					account.ID,
+					selection.WaitPlan.MaxConcurrency,
+					selection.WaitPlan.Timeout,
+					reqStream,
+					&streamStarted,
+				)
+				if err != nil {
+					log.Printf("Account concurrency acquire failed: %v", err)
+					h.handleConcurrencyError(c, err, "account", streamStarted)
+					return
+				}
+				// Slot acquired: no longer waiting in queue.
+				if accountWaitCounted {
+					h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
+					accountWaitCounted = false
+				}
+				if err := h.gatewayService.BindStickySession(c.Request.Context(), apiKey.GroupID, sessionKey, account.ID); err != nil {
+					log.Printf("Bind sticky session failed: %v", err)
+				}
+			}
+			// 账号槽位/等待计数需要在超时或断开时安全回收
+			accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
+
+			// 转发请求 - 根据账号平台分流
+			var result *service.ForwardResult
+			if account.Platform == service.PlatformAntigravity {
+				result, err = h.antigravityGatewayService.ForwardGemini(c.Request.Context(), c, account, reqModel, "generateContent", reqStream, body)
+			} else {
+				result, err = h.geminiCompatService.Forward(c.Request.Context(), c, account, body)
+			}
+			if accountReleaseFunc != nil {
+				accountReleaseFunc()
+			}
+			if err != nil {
+				var failoverErr *service.UpstreamFailoverError
+				if errors.As(err, &failoverErr) {
+					failedAccountIDs[account.ID] = struct{}{}
+					lastFailoverStatus = failoverErr.StatusCode
+					if switchCount >= maxAccountSwitches {
+						h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
+						return
+					}
+					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
+			}
+
+			// 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context）
+			userAgent := c.GetHeader("User-Agent")
+			clientIP := ip.GetClientIP(c)
+
+			// 异步记录使用量（subscription已在函数开头获取）
+			go func(result *service.ForwardResult, usedAccount *service.Account, ua, clientIP string) {
+				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+				defer cancel()
+				if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
+					Result:       result,
+					APIKey:       apiKey,
+					User:         apiKey.User,
+					Account:      usedAccount,
+					Subscription: subscription,
+					UserAgent:    ua,
+					IPAddress:    clientIP,
+				}); err != nil {
+					log.Printf("Record usage failed: %v", err)
+				}
+			}(result, account, userAgent, clientIP)
+			return
+		}
+	}
+
+	const maxAccountSwitches = 10
+	switchCount := 0
+	failedAccountIDs := make(map[int64]struct{})
+	lastFailoverStatus := 0
+
+	for {
+		// 选择支持该模型的账号
+		selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, failedAccountIDs)
+		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
+		}
+		account := selection.Account
+		setOpsSelectedAccount(c, account.ID)
+
+		// 检查预热请求拦截（在账号选择后、转发前检查）
+		if account.IsInterceptWarmupEnabled() && isWarmupRequest(body) {
+			if selection.Acquired && selection.ReleaseFunc != nil {
+				selection.ReleaseFunc()
+			}
+			if reqStream {
+				sendMockWarmupStream(c, reqModel)
+			} else {
+				sendMockWarmupResponse(c, reqModel)
+			}
+			return
+		}
+
+		// 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 {
+				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
+			}
+			if err == nil && canWait {
+				accountWaitCounted = true
+			}
+			defer func() {
+				if accountWaitCounted {
+					h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
+				}
+			}()
+
+			accountReleaseFunc, err = h.concurrencyHelper.AcquireAccountSlotWithWaitTimeout(
+				c,
+				account.ID,
+				selection.WaitPlan.MaxConcurrency,
+				selection.WaitPlan.Timeout,
+				reqStream,
+				&streamStarted,
+			)
+			if err != nil {
+				log.Printf("Account concurrency acquire failed: %v", err)
+				h.handleConcurrencyError(c, err, "account", streamStarted)
+				return
+			}
+			if accountWaitCounted {
+				h.concurrencyHelper.DecrementAccountWaitCount(c.Request.Context(), account.ID)
+				accountWaitCounted = false
+			}
+			if err := h.gatewayService.BindStickySession(c.Request.Context(), apiKey.GroupID, sessionKey, account.ID); err != nil {
+				log.Printf("Bind sticky session failed: %v", err)
+			}
+		}
+		// 账号槽位/等待计数需要在超时或断开时安全回收
+		accountReleaseFunc = wrapReleaseOnDone(c.Request.Context(), accountReleaseFunc)
+
+		// 转发请求 - 根据账号平台分流
+		var result *service.ForwardResult
+		if account.Platform == service.PlatformAntigravity {
+			result, err = h.antigravityGatewayService.Forward(c.Request.Context(), c, account, body)
+		} else {
+			result, err = h.gatewayService.Forward(c.Request.Context(), c, account, parsedReq)
+		}
+		if accountReleaseFunc != nil {
+			accountReleaseFunc()
+		}
+		if err != nil {
+			var failoverErr *service.UpstreamFailoverError
+			if errors.As(err, &failoverErr) {
+				failedAccountIDs[account.ID] = struct{}{}
+				lastFailoverStatus = failoverErr.StatusCode
+				if switchCount >= maxAccountSwitches {
+					h.handleFailoverExhausted(c, lastFailoverStatus, streamStarted)
+					return
+				}
+				switchCount++
+				log.Printf("Account %d: upstream error %d, switching account %d/%d", account.ID, failoverErr.StatusCode, switchCount, maxAccountSwitches)
+				continue
+			}
+			// 错误响应已在Forward中处理，这里只记录日志
+			log.Printf("Account %d: Forward request failed: %v", account.ID, err)
+			return
+		}
+
+		// 捕获请求信息（用于异步记录，避免在 goroutine 中访问 gin.Context）
+		userAgent := c.GetHeader("User-Agent")
+		clientIP := ip.GetClientIP(c)
+
+		// 异步记录使用量（subscription已在函数开头获取）
+		go func(result *service.ForwardResult, usedAccount *service.Account, ua, clientIP string) {
+			ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+			defer cancel()
+			if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
+				Result:       result,
+				APIKey:       apiKey,
+				User:         apiKey.User,
+				Account:      usedAccount,
+				Subscription: subscription,
+				UserAgent:    ua,
+				IPAddress:    clientIP,
+			}); err != nil {
+				log.Printf("Record usage failed: %v", err)
+			}
+		}(result, account, userAgent, clientIP)
+		return
+	}
+}
+
+// Models handles listing available models
+// GET /v1/models
+// Returns models based on account configurations (model_mapping whitelist)
+// Falls back to default models if no whitelist is configured
+func (h *GatewayHandler) Models(c *gin.Context) {
+	apiKey, _ := middleware2.GetAPIKeyFromContext(c)
+
+	var groupID *int64
+	var platform string
+
+	if apiKey != nil && apiKey.Group != nil {
+		groupID = &apiKey.Group.ID
+		platform = apiKey.Group.Platform
+	}
+
+	// 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" {
+		c.JSON(http.StatusOK, gin.H{
+			"object": "list",
+			"data":   openai.DefaultModels,
+		})
+		return
+	}
+
+	c.JSON(http.StatusOK, gin.H{
+		"object": "list",
+		"data":   claude.DefaultModels,
+	})
+}
+
+// AntigravityModels 返回 Antigravity 支持的全部模型
+// GET /antigravity/models
+func (h *GatewayHandler) AntigravityModels(c *gin.Context) {
+	c.JSON(http.StatusOK, gin.H{
+		"object": "list",
+		"data":   antigravity.DefaultModels(),
+	})
+}
+
+// Usage handles getting account balance for CC Switch integration
+// GET /v1/usage
+func (h *GatewayHandler) Usage(c *gin.Context) {
+	apiKey, ok := middleware2.GetAPIKeyFromContext(c)
+	if !ok {
+		h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
+		return
+	}
+
+	subject, ok := middleware2.GetAuthSubjectFromContext(c)
+	if !ok {
+		h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
+		return
+	}
+
+	// 订阅模式：返回订阅限额信息
+	if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() {
+		subscription, ok := middleware2.GetSubscriptionFromContext(c)
+		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
+	}
+
+	// 余额模式：返回钱包余额
+	latestUser, err := h.userService.GetByID(c.Request.Context(), subject.UserID)
+	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. 否则返回所有已配置周期中剩余额度的最小值
+func (h *GatewayHandler) calculateSubscriptionRemaining(group *service.Group, sub *service.UserSubscription) float64 {
+	var remainingValues []float64
+
+	// 检查日限额
+	if group.HasDailyLimit() {
+		remaining := *group.DailyLimitUSD - sub.DailyUsageUSD
+		if remaining <= 0 {
+			return 0
+		}
+		remainingValues = append(remainingValues, remaining)
+	}
+
+	// 检查周限额
+	if group.HasWeeklyLimit() {
+		remaining := *group.WeeklyLimitUSD - sub.WeeklyUsageUSD
+		if remaining <= 0 {
+			return 0
+		}
+		remainingValues = append(remainingValues, remaining)
+	}
+
+	// 检查月限额
+	if group.HasMonthlyLimit() {
+		remaining := *group.MonthlyLimitUSD - sub.MonthlyUsageUSD
+		if remaining <= 0 {
+			return 0
+		}
+		remainingValues = append(remainingValues, remaining)
+	}
+
+	// 如果没有配置任何限额，返回-1表示无限制
+	if len(remainingValues) == 0 {
+		return -1
+	}
+
+	// 返回最小值
+	min := remainingValues[0]
+	for _, v := range remainingValues[1:] {
+		if v < min {
+			min = v
+		}
+	}
+	return min
+}
+
+// handleConcurrencyError handles concurrency-related 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)
+}
+
+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"
+	}
+}
+
+// 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 {
+			// 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))
+			if _, err := fmt.Fprint(c.Writer, errorEvent); err != nil {
+				_ = c.Error(err)
+			}
+			flusher.Flush()
+		}
+		return
+	}
+
+	// Normal case: return JSON response with proper status code
+	h.errorResponse(c, status, errType, message)
+}
+
+// errorResponse 返回Claude API格式的错误响应
+func (h *GatewayHandler) errorResponse(c *gin.Context, status int, errType, message string) {
+	c.JSON(status, gin.H{
+		"type": "error",
+		"error": gin.H{
+			"type":    errType,
+			"message": message,
+		},
+	})
+}
+
+// CountTokens handles token counting endpoint
+// POST /v1/messages/count_tokens
+// 特点：校验订阅/余额，但不计算并发、不记录使用量
+func (h *GatewayHandler) CountTokens(c *gin.Context) {
+	// 从context获取apiKey和user（ApiKeyAuth中间件已设置）
+	apiKey, ok := middleware2.GetAPIKeyFromContext(c)
+	if !ok {
+		h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
+		return
+	}
+
+	_, ok = middleware2.GetAuthSubjectFromContext(c)
+	if !ok {
+		h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
+		return
+	}
+
+	// 读取请求体
+	body, err := io.ReadAll(c.Request.Body)
+	if err != nil {
+		if maxErr, ok := extractMaxBytesError(err); ok {
+			h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
+			return
+		}
+		h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
+		return
+	}
+
+	if len(body) == 0 {
+		h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
+		return
+	}
+
+	setOpsRequestContext(c, "", false, body)
+
+	parsedReq, err := service.ParseGatewayRequest(body)
+	if err != nil {
+		h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
+		return
+	}
+
+	// 验证 model 必填
+	if parsedReq.Model == "" {
+		h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
+		return
+	}
+
+	setOpsRequestContext(c, parsedReq.Model, parsedReq.Stream, body)
+
+	// 获取订阅信息（可能为nil）
+	subscription, _ := middleware2.GetSubscriptionFromContext(c)
+
+	// 校验 billing eligibility（订阅/余额）
+	// 【注意】不计算并发，但需要校验订阅/余额
+	if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription); err != nil {
+		status, code, message := billingErrorDetails(err)
+		h.errorResponse(c, status, code, message)
+		return
+	}
+
+	// 计算粘性会话 hash
+	sessionHash := h.gatewayService.GenerateSessionHash(parsedReq)
+
+	// 选择支持该模型的账号
+	account, err := h.gatewayService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, sessionHash, parsedReq.Model)
+	if err != nil {
+		h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error())
+		return
+	}
+	setOpsSelectedAccount(c, account.ID)
+
+	// 转发请求（不记录使用量）
+	if err := h.gatewayService.ForwardCountTokens(c.Request.Context(), c, account, parsedReq); err != nil {
+		log.Printf("Forward count_tokens request failed: %v", err)
+		// 错误响应已在 ForwardCountTokens 中处理
+		return
+	}
+}
+
+// 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")
+
+	// 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)
+
+	events := []string{
+		`event: message_start` + "\n" + `data: ` + string(messageStartJSON),
+		`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{
+		"id":          "msg_mock_warmup",
+		"type":        "message",
+		"role":        "assistant",
+		"model":       model,
+		"content":     []gin.H{{"type": "text", "text": "New Conversation"}},
+		"stop_reason": "end_turn",
+		"usage": gin.H{
+			"input_tokens":  10,
+			"output_tokens": 2,
+		},
+	})
+}
+
+func billingErrorDetails(err error) (status int, code, message string) {
+	if errors.Is(err, service.ErrBillingServiceUnavailable) {
+		msg := pkgerrors.Message(err)
+		if msg == "" {
+			msg = "Billing service temporarily unavailable. Please retry later."
+		}
+		return http.StatusServiceUnavailable, "billing_service_error", msg
+	}
+	msg := pkgerrors.Message(err)
+	if msg == "" {
+		msg = err.Error()
+	}
+	return http.StatusForbidden, "billing_error", msg
+}