refactor split gateway handler helpers
This commit is contained in:
@@ -2,12 +2,8 @@ package handler
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"crypto/rand"
|
|
||||||
"encoding/json"
|
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
|
||||||
"net/http"
|
"net/http"
|
||||||
"strconv"
|
|
||||||
"strings"
|
"strings"
|
||||||
"sync/atomic"
|
"sync/atomic"
|
||||||
"time"
|
"time"
|
||||||
@@ -17,12 +13,10 @@ import (
|
|||||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||||
pkgerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
|
||||||
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
|
||||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||||
|
|
||||||
@@ -904,17 +898,6 @@ func (h *GatewayHandler) AntigravityModels(c *gin.Context) {
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func cloneAPIKeyWithGroup(apiKey *service.APIKey, group *service.Group) *service.APIKey {
|
|
||||||
if apiKey == nil || group == nil {
|
|
||||||
return apiKey
|
|
||||||
}
|
|
||||||
cloned := *apiKey
|
|
||||||
groupID := group.ID
|
|
||||||
cloned.GroupID = &groupID
|
|
||||||
cloned.Group = group
|
|
||||||
return &cloned
|
|
||||||
}
|
|
||||||
|
|
||||||
// Usage handles getting account balance and usage statistics for CC Switch integration
|
// Usage handles getting account balance and usage statistics for CC Switch integration
|
||||||
// GET /v1/usage
|
// GET /v1/usage
|
||||||
//
|
//
|
||||||
@@ -961,412 +944,6 @@ func (h *GatewayHandler) Usage(c *gin.Context) {
|
|||||||
h.usageUnrestricted(c, ctx, apiKey, subject, usageData, modelStats)
|
h.usageUnrestricted(c, ctx, apiKey, subject, usageData, modelStats)
|
||||||
}
|
}
|
||||||
|
|
||||||
// parseUsageDateRange 解析 start_date / end_date query params，默认返回近 30 天范围
|
|
||||||
func (h *GatewayHandler) parseUsageDateRange(c *gin.Context) (time.Time, time.Time) {
|
|
||||||
now := timezone.Now()
|
|
||||||
endTime := now
|
|
||||||
startTime := now.AddDate(0, 0, -30)
|
|
||||||
|
|
||||||
if s := c.Query("start_date"); s != "" {
|
|
||||||
if t, err := timezone.ParseInLocation("2006-01-02", s); err == nil {
|
|
||||||
startTime = t
|
|
||||||
}
|
|
||||||
}
|
|
||||||
if s := c.Query("end_date"); s != "" {
|
|
||||||
if t, err := timezone.ParseInLocation("2006-01-02", s); err == nil {
|
|
||||||
endTime = t.AddDate(0, 0, 1) // half-open range upper bound
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return startTime, endTime
|
|
||||||
}
|
|
||||||
|
|
||||||
// buildUsageData 构建 today/total 用量摘要
|
|
||||||
func (h *GatewayHandler) buildUsageData(ctx context.Context, apiKeyID int64) gin.H {
|
|
||||||
if h.usageService == nil {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
dashStats, err := h.usageService.GetAPIKeyDashboardStats(ctx, apiKeyID)
|
|
||||||
if err != nil || dashStats == nil {
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
return gin.H{
|
|
||||||
"today": gin.H{
|
|
||||||
"requests": dashStats.TodayRequests,
|
|
||||||
"input_tokens": dashStats.TodayInputTokens,
|
|
||||||
"output_tokens": dashStats.TodayOutputTokens,
|
|
||||||
"cache_creation_tokens": dashStats.TodayCacheCreationTokens,
|
|
||||||
"cache_read_tokens": dashStats.TodayCacheReadTokens,
|
|
||||||
"total_tokens": dashStats.TodayTokens,
|
|
||||||
"cost": dashStats.TodayCost,
|
|
||||||
"actual_cost": dashStats.TodayActualCost,
|
|
||||||
},
|
|
||||||
"total": gin.H{
|
|
||||||
"requests": dashStats.TotalRequests,
|
|
||||||
"input_tokens": dashStats.TotalInputTokens,
|
|
||||||
"output_tokens": dashStats.TotalOutputTokens,
|
|
||||||
"cache_creation_tokens": dashStats.TotalCacheCreationTokens,
|
|
||||||
"cache_read_tokens": dashStats.TotalCacheReadTokens,
|
|
||||||
"total_tokens": dashStats.TotalTokens,
|
|
||||||
"cost": dashStats.TotalCost,
|
|
||||||
"actual_cost": dashStats.TotalActualCost,
|
|
||||||
},
|
|
||||||
"average_duration_ms": dashStats.AverageDurationMs,
|
|
||||||
"rpm": dashStats.Rpm,
|
|
||||||
"tpm": dashStats.Tpm,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// usageQuotaLimited 处理 quota_limited 模式的响应
|
|
||||||
func (h *GatewayHandler) usageQuotaLimited(c *gin.Context, ctx context.Context, apiKey *service.APIKey, usageData gin.H, modelStats any) {
|
|
||||||
resp := gin.H{
|
|
||||||
"mode": "quota_limited",
|
|
||||||
"isValid": apiKey.Status == service.StatusAPIKeyActive || apiKey.Status == service.StatusAPIKeyQuotaExhausted || apiKey.Status == service.StatusAPIKeyExpired,
|
|
||||||
"status": apiKey.Status,
|
|
||||||
}
|
|
||||||
|
|
||||||
// 总额度信息
|
|
||||||
if apiKey.Quota > 0 {
|
|
||||||
remaining := apiKey.GetQuotaRemaining()
|
|
||||||
resp["quota"] = gin.H{
|
|
||||||
"limit": apiKey.Quota,
|
|
||||||
"used": apiKey.QuotaUsed,
|
|
||||||
"remaining": remaining,
|
|
||||||
"unit": "USD",
|
|
||||||
}
|
|
||||||
resp["remaining"] = remaining
|
|
||||||
resp["unit"] = "USD"
|
|
||||||
}
|
|
||||||
|
|
||||||
// 速率限制信息（从 DB 获取实时用量）
|
|
||||||
if apiKey.HasRateLimits() && h.apiKeyService != nil {
|
|
||||||
rateLimitData, err := h.apiKeyService.GetRateLimitData(ctx, apiKey.ID)
|
|
||||||
if err == nil && rateLimitData != nil {
|
|
||||||
var rateLimits []gin.H
|
|
||||||
if apiKey.RateLimit5h > 0 {
|
|
||||||
used := rateLimitData.EffectiveUsage5h()
|
|
||||||
entry := gin.H{
|
|
||||||
"window": "5h",
|
|
||||||
"limit": apiKey.RateLimit5h,
|
|
||||||
"used": used,
|
|
||||||
"remaining": max(0, apiKey.RateLimit5h-used),
|
|
||||||
"window_start": rateLimitData.Window5hStart,
|
|
||||||
}
|
|
||||||
if rateLimitData.Window5hStart != nil && !service.IsWindowExpired(rateLimitData.Window5hStart, service.RateLimitWindow5h) {
|
|
||||||
entry["reset_at"] = rateLimitData.Window5hStart.Add(service.RateLimitWindow5h)
|
|
||||||
}
|
|
||||||
rateLimits = append(rateLimits, entry)
|
|
||||||
}
|
|
||||||
if apiKey.RateLimit1d > 0 {
|
|
||||||
used := rateLimitData.EffectiveUsage1d()
|
|
||||||
entry := gin.H{
|
|
||||||
"window": "1d",
|
|
||||||
"limit": apiKey.RateLimit1d,
|
|
||||||
"used": used,
|
|
||||||
"remaining": max(0, apiKey.RateLimit1d-used),
|
|
||||||
"window_start": rateLimitData.Window1dStart,
|
|
||||||
}
|
|
||||||
if rateLimitData.Window1dStart != nil && !service.IsWindowExpired(rateLimitData.Window1dStart, service.RateLimitWindow1d) {
|
|
||||||
entry["reset_at"] = rateLimitData.Window1dStart.Add(service.RateLimitWindow1d)
|
|
||||||
}
|
|
||||||
rateLimits = append(rateLimits, entry)
|
|
||||||
}
|
|
||||||
if apiKey.RateLimit7d > 0 {
|
|
||||||
used := rateLimitData.EffectiveUsage7d()
|
|
||||||
entry := gin.H{
|
|
||||||
"window": "7d",
|
|
||||||
"limit": apiKey.RateLimit7d,
|
|
||||||
"used": used,
|
|
||||||
"remaining": max(0, apiKey.RateLimit7d-used),
|
|
||||||
"window_start": rateLimitData.Window7dStart,
|
|
||||||
}
|
|
||||||
if rateLimitData.Window7dStart != nil && !service.IsWindowExpired(rateLimitData.Window7dStart, service.RateLimitWindow7d) {
|
|
||||||
entry["reset_at"] = rateLimitData.Window7dStart.Add(service.RateLimitWindow7d)
|
|
||||||
}
|
|
||||||
rateLimits = append(rateLimits, entry)
|
|
||||||
}
|
|
||||||
if len(rateLimits) > 0 {
|
|
||||||
resp["rate_limits"] = rateLimits
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// 过期时间
|
|
||||||
if apiKey.ExpiresAt != nil {
|
|
||||||
resp["expires_at"] = apiKey.ExpiresAt
|
|
||||||
resp["days_until_expiry"] = apiKey.GetDaysUntilExpiry()
|
|
||||||
}
|
|
||||||
|
|
||||||
if usageData != nil {
|
|
||||||
resp["usage"] = usageData
|
|
||||||
}
|
|
||||||
if modelStats != nil {
|
|
||||||
resp["model_stats"] = modelStats
|
|
||||||
}
|
|
||||||
|
|
||||||
c.JSON(http.StatusOK, resp)
|
|
||||||
}
|
|
||||||
|
|
||||||
// usageUnrestricted 处理 unrestricted 模式的响应（向后兼容）
|
|
||||||
func (h *GatewayHandler) usageUnrestricted(c *gin.Context, ctx context.Context, apiKey *service.APIKey, subject middleware2.AuthSubject, usageData gin.H, modelStats any) {
|
|
||||||
// 订阅模式
|
|
||||||
if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() {
|
|
||||||
resp := gin.H{
|
|
||||||
"mode": "unrestricted",
|
|
||||||
"isValid": true,
|
|
||||||
"planName": apiKey.Group.Name,
|
|
||||||
"unit": "USD",
|
|
||||||
}
|
|
||||||
|
|
||||||
// 订阅信息可能不在 context 中（/v1/usage 路径跳过了中间件的计费检查）
|
|
||||||
subscription, ok := middleware2.GetSubscriptionFromContext(c)
|
|
||||||
if ok {
|
|
||||||
remaining := h.calculateSubscriptionRemaining(apiKey.Group, subscription)
|
|
||||||
resp["remaining"] = remaining
|
|
||||||
resp["subscription"] = gin.H{
|
|
||||||
"daily_usage_usd": subscription.DailyUsageUSD,
|
|
||||||
"weekly_usage_usd": subscription.WeeklyUsageUSD,
|
|
||||||
"monthly_usage_usd": subscription.MonthlyUsageUSD,
|
|
||||||
"daily_limit_usd": apiKey.Group.DailyLimitUSD,
|
|
||||||
"weekly_limit_usd": apiKey.Group.WeeklyLimitUSD,
|
|
||||||
"monthly_limit_usd": apiKey.Group.MonthlyLimitUSD,
|
|
||||||
"expires_at": subscription.ExpiresAt,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if usageData != nil {
|
|
||||||
resp["usage"] = usageData
|
|
||||||
}
|
|
||||||
if modelStats != nil {
|
|
||||||
resp["model_stats"] = modelStats
|
|
||||||
}
|
|
||||||
c.JSON(http.StatusOK, resp)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
// 余额模式
|
|
||||||
latestUser, err := h.userService.GetByID(ctx, subject.UserID)
|
|
||||||
if err != nil {
|
|
||||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Failed to get user info")
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
resp := gin.H{
|
|
||||||
"mode": "unrestricted",
|
|
||||||
"isValid": true,
|
|
||||||
"planName": "钱包余额",
|
|
||||||
"remaining": latestUser.Balance,
|
|
||||||
"unit": "USD",
|
|
||||||
"balance": latestUser.Balance,
|
|
||||||
}
|
|
||||||
if usageData != nil {
|
|
||||||
resp["usage"] = usageData
|
|
||||||
}
|
|
||||||
if modelStats != nil {
|
|
||||||
resp["model_stats"] = modelStats
|
|
||||||
}
|
|
||||||
c.JSON(http.StatusOK, resp)
|
|
||||||
}
|
|
||||||
|
|
||||||
// calculateSubscriptionRemaining 计算订阅剩余可用额度
|
|
||||||
// 逻辑：
|
|
||||||
// 1. 如果日/周/月任一限额达到100%，返回0
|
|
||||||
// 2. 否则返回所有已配置周期中剩余额度的最小值
|
|
||||||
func (h *GatewayHandler) calculateSubscriptionRemaining(group *service.Group, sub *service.UserSubscription) float64 {
|
|
||||||
var remainingValues []float64
|
|
||||||
|
|
||||||
// 检查日限额
|
|
||||||
if group.HasDailyLimit() {
|
|
||||||
remaining := *group.DailyLimitUSD - sub.DailyUsageUSD
|
|
||||||
if remaining <= 0 {
|
|
||||||
return 0
|
|
||||||
}
|
|
||||||
remainingValues = append(remainingValues, remaining)
|
|
||||||
}
|
|
||||||
|
|
||||||
// 检查周限额
|
|
||||||
if group.HasWeeklyLimit() {
|
|
||||||
remaining := *group.WeeklyLimitUSD - sub.WeeklyUsageUSD
|
|
||||||
if remaining <= 0 {
|
|
||||||
return 0
|
|
||||||
}
|
|
||||||
remainingValues = append(remainingValues, remaining)
|
|
||||||
}
|
|
||||||
|
|
||||||
// 检查月限额
|
|
||||||
if group.HasMonthlyLimit() {
|
|
||||||
remaining := *group.MonthlyLimitUSD - sub.MonthlyUsageUSD
|
|
||||||
if remaining <= 0 {
|
|
||||||
return 0
|
|
||||||
}
|
|
||||||
remainingValues = append(remainingValues, remaining)
|
|
||||||
}
|
|
||||||
|
|
||||||
// 如果没有配置任何限额，返回-1表示无限制
|
|
||||||
if len(remainingValues) == 0 {
|
|
||||||
return -1
|
|
||||||
}
|
|
||||||
|
|
||||||
// 返回最小值
|
|
||||||
min := remainingValues[0]
|
|
||||||
for _, v := range remainingValues[1:] {
|
|
||||||
if v < min {
|
|
||||||
min = v
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return min
|
|
||||||
}
|
|
||||||
|
|
||||||
// handleConcurrencyError handles concurrency-related 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, 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
|
|
||||||
}
|
|
||||||
|
|
||||||
if rule.SkipMonitoring {
|
|
||||||
c.Set(service.OpsSkipPassthroughKey, true)
|
|
||||||
}
|
|
||||||
|
|
||||||
h.handleStreamingAwareError(c, respCode, "upstream_error", msg, streamStarted)
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// 记录原始上游状态码，以便 ops 错误日志捕获真实的上游错误
|
|
||||||
upstreamMsg := service.ExtractUpstreamErrorMessage(responseBody)
|
|
||||||
service.SetOpsUpstreamError(c, statusCode, upstreamMsg, "")
|
|
||||||
|
|
||||||
// 使用默认的错误映射
|
|
||||||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
|
||||||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
|
||||||
}
|
|
||||||
|
|
||||||
// handleFailoverExhaustedSimple 简化版本，用于没有响应体的情况
|
|
||||||
func (h *GatewayHandler) handleFailoverExhaustedSimple(c *gin.Context, statusCode int, streamStarted bool) {
|
|
||||||
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
|
||||||
service.SetOpsUpstreamError(c, statusCode, errMsg, "")
|
|
||||||
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *GatewayHandler) mapUpstreamError(statusCode int) (int, string, string) {
|
|
||||||
switch statusCode {
|
|
||||||
case 401:
|
|
||||||
return http.StatusBadGateway, "upstream_error", "Upstream authentication failed, please contact administrator"
|
|
||||||
case 403:
|
|
||||||
return http.StatusBadGateway, "upstream_error", "Upstream access forbidden, please contact administrator"
|
|
||||||
case 429:
|
|
||||||
return http.StatusTooManyRequests, "rate_limit_error", "Upstream rate limit exceeded, please retry later"
|
|
||||||
case 529:
|
|
||||||
return http.StatusServiceUnavailable, "overloaded_error", "Upstream service overloaded, please retry later"
|
|
||||||
case 500, 502, 503, 504:
|
|
||||||
return http.StatusBadGateway, "upstream_error", "Upstream service temporarily unavailable"
|
|
||||||
default:
|
|
||||||
return http.StatusBadGateway, "upstream_error", "Upstream request failed"
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// handleStreamingAwareError handles errors that may occur after streaming has started
|
|
||||||
func (h *GatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) {
|
|
||||||
if streamStarted {
|
|
||||||
// Stream already started, send error as SSE event then close
|
|
||||||
flusher, ok := c.Writer.(http.Flusher)
|
|
||||||
if ok {
|
|
||||||
// SSE 错误事件固定 schema，使用 Quote 直拼可避免额外 Marshal 分配。
|
|
||||||
errorEvent := `data: {"type":"error","error":{"type":` + strconv.Quote(errType) + `,"message":` + strconv.Quote(message) + `}}` + "\n\n"
|
|
||||||
if _, err := fmt.Fprint(c.Writer, errorEvent); err != nil {
|
|
||||||
_ = c.Error(err)
|
|
||||||
}
|
|
||||||
flusher.Flush()
|
|
||||||
}
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
// Normal case: return JSON response with proper status code
|
|
||||||
h.errorResponse(c, status, errType, message)
|
|
||||||
}
|
|
||||||
|
|
||||||
// ensureForwardErrorResponse 在 Forward 返回错误但尚未写响应时补写统一错误响应。
|
|
||||||
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
|
|
||||||
}
|
|
||||||
|
|
||||||
// checkClaudeCodeVersion 检查 Claude Code 客户端版本是否满足版本要求
|
|
||||||
// 仅对已识别的 Claude Code 客户端执行，count_tokens 路径除外
|
|
||||||
func (h *GatewayHandler) checkClaudeCodeVersion(c *gin.Context) bool {
|
|
||||||
ctx := c.Request.Context()
|
|
||||||
if !service.IsClaudeCodeClient(ctx) {
|
|
||||||
return true
|
|
||||||
}
|
|
||||||
|
|
||||||
// 排除 count_tokens 子路径
|
|
||||||
if strings.HasSuffix(c.Request.URL.Path, "/count_tokens") {
|
|
||||||
return true
|
|
||||||
}
|
|
||||||
|
|
||||||
minVersion, maxVersion := h.settingService.GetClaudeCodeVersionBounds(ctx)
|
|
||||||
if minVersion == "" && maxVersion == "" {
|
|
||||||
return true // 未设置，不检查
|
|
||||||
}
|
|
||||||
|
|
||||||
clientVersion := service.GetClaudeCodeVersion(ctx)
|
|
||||||
if clientVersion == "" {
|
|
||||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
|
||||||
"Unable to determine Claude Code version. Please update Claude Code: npm update -g @anthropic-ai/claude-code")
|
|
||||||
return false
|
|
||||||
}
|
|
||||||
|
|
||||||
if minVersion != "" && service.CompareVersions(clientVersion, minVersion) < 0 {
|
|
||||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
|
||||||
fmt.Sprintf("Your Claude Code version (%s) is below the minimum required version (%s). Please update: npm update -g @anthropic-ai/claude-code",
|
|
||||||
clientVersion, minVersion))
|
|
||||||
return false
|
|
||||||
}
|
|
||||||
|
|
||||||
if maxVersion != "" && service.CompareVersions(clientVersion, maxVersion) > 0 {
|
|
||||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
|
||||||
fmt.Sprintf("Your Claude Code version (%s) exceeds the maximum allowed version (%s). "+
|
|
||||||
"Please downgrade: npm install -g @anthropic-ai/claude-code@%s && "+
|
|
||||||
"set CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 to prevent auto-upgrade",
|
|
||||||
clientVersion, maxVersion, maxVersion))
|
|
||||||
return false
|
|
||||||
}
|
|
||||||
|
|
||||||
return true
|
|
||||||
}
|
|
||||||
|
|
||||||
// errorResponse 返回Claude API格式的错误响应
|
|
||||||
func (h *GatewayHandler) errorResponse(c *gin.Context, status int, errType, message string) {
|
|
||||||
c.JSON(status, gin.H{
|
|
||||||
"type": "error",
|
|
||||||
"error": gin.H{
|
|
||||||
"type": errType,
|
|
||||||
"message": message,
|
|
||||||
},
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
// CountTokens handles token counting endpoint
|
// CountTokens handles token counting endpoint
|
||||||
// POST /v1/messages/count_tokens
|
// POST /v1/messages/count_tokens
|
||||||
// 特点：校验订阅/余额，但不计算并发、不记录使用量
|
// 特点：校验订阅/余额，但不计算并发、不记录使用量
|
||||||
@@ -1464,314 +1041,3 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// InterceptType 表示请求拦截类型
|
|
||||||
type InterceptType int
|
|
||||||
|
|
||||||
const (
|
|
||||||
InterceptTypeNone InterceptType = iota
|
|
||||||
InterceptTypeWarmup // 预热请求（返回 "New Conversation"）
|
|
||||||
InterceptTypeSuggestionMode // SUGGESTION MODE（返回空字符串）
|
|
||||||
InterceptTypeMaxTokensOneHaiku // max_tokens=1 + haiku 探测请求（返回 "#"）
|
|
||||||
)
|
|
||||||
|
|
||||||
// isHaikuModel 检查模型名称是否包含 "haiku"（大小写不敏感）
|
|
||||||
func isHaikuModel(model string) bool {
|
|
||||||
return strings.Contains(strings.ToLower(model), "haiku")
|
|
||||||
}
|
|
||||||
|
|
||||||
// isMaxTokensOneHaikuRequest 检查是否为 max_tokens=1 + haiku 模型的探测请求
|
|
||||||
// 这类请求用于 Claude Code 验证 API 连通性
|
|
||||||
// 条件：max_tokens == 1 且 model 包含 "haiku" 且非流式请求
|
|
||||||
func isMaxTokensOneHaikuRequest(model string, maxTokens int, isStream bool) bool {
|
|
||||||
return maxTokens == 1 && isHaikuModel(model) && !isStream
|
|
||||||
}
|
|
||||||
|
|
||||||
// detectInterceptType 检测请求是否需要拦截，返回拦截类型
|
|
||||||
// 参数说明：
|
|
||||||
// - 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
|
|
||||||
}
|
|
||||||
|
|
||||||
// 快速检查：如果不包含任何关键字，直接返回
|
|
||||||
bodyStr := string(body)
|
|
||||||
hasSuggestionMode := strings.Contains(bodyStr, "[SUGGESTION MODE:")
|
|
||||||
hasWarmupKeyword := strings.Contains(bodyStr, "title") || strings.Contains(bodyStr, "Warmup")
|
|
||||||
|
|
||||||
if !hasSuggestionMode && !hasWarmupKeyword {
|
|
||||||
return InterceptTypeNone
|
|
||||||
}
|
|
||||||
|
|
||||||
// 解析请求（只解析一次）
|
|
||||||
var req struct {
|
|
||||||
Messages []struct {
|
|
||||||
Role string `json:"role"`
|
|
||||||
Content []struct {
|
|
||||||
Type string `json:"type"`
|
|
||||||
Text string `json:"text"`
|
|
||||||
} `json:"content"`
|
|
||||||
} `json:"messages"`
|
|
||||||
System []struct {
|
|
||||||
Text string `json:"text"`
|
|
||||||
} `json:"system"`
|
|
||||||
}
|
|
||||||
if err := json.Unmarshal(body, &req); err != nil {
|
|
||||||
return InterceptTypeNone
|
|
||||||
}
|
|
||||||
|
|
||||||
// 检查 SUGGESTION MODE（最后一条 user 消息）
|
|
||||||
if hasSuggestionMode && len(req.Messages) > 0 {
|
|
||||||
lastMsg := req.Messages[len(req.Messages)-1]
|
|
||||||
if lastMsg.Role == "user" && len(lastMsg.Content) > 0 &&
|
|
||||||
lastMsg.Content[0].Type == "text" &&
|
|
||||||
strings.HasPrefix(lastMsg.Content[0].Text, "[SUGGESTION MODE:") {
|
|
||||||
return InterceptTypeSuggestionMode
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// 检查 Warmup 请求
|
|
||||||
if hasWarmupKeyword {
|
|
||||||
// 检查 messages 中的标题提示模式
|
|
||||||
for _, msg := range req.Messages {
|
|
||||||
for _, content := range msg.Content {
|
|
||||||
if content.Type == "text" {
|
|
||||||
if strings.Contains(content.Text, "Please write a 5-10 word title for the following conversation:") ||
|
|
||||||
content.Text == "Warmup" {
|
|
||||||
return InterceptTypeWarmup
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
// 检查 system 中的标题提取模式
|
|
||||||
for _, sys := range req.System {
|
|
||||||
if strings.Contains(sys.Text, "nalyze if this message indicates a new conversation topic. If it does, extract a 2-3 word title") {
|
|
||||||
return InterceptTypeWarmup
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return InterceptTypeNone
|
|
||||||
}
|
|
||||||
|
|
||||||
// sendMockInterceptStream 发送流式 mock 响应（用于请求拦截）
|
|
||||||
func sendMockInterceptStream(c *gin.Context, model string, interceptType InterceptType) {
|
|
||||||
c.Header("Content-Type", "text/event-stream")
|
|
||||||
c.Header("Cache-Control", "no-cache")
|
|
||||||
c.Header("Connection", "keep-alive")
|
|
||||||
c.Header("X-Accel-Buffering", "no")
|
|
||||||
|
|
||||||
// 根据拦截类型决定响应内容
|
|
||||||
var msgID string
|
|
||||||
var outputTokens int
|
|
||||||
var textDeltas []string
|
|
||||||
|
|
||||||
switch interceptType {
|
|
||||||
case InterceptTypeSuggestionMode:
|
|
||||||
msgID = "msg_mock_suggestion"
|
|
||||||
outputTokens = 1
|
|
||||||
textDeltas = []string{""} // 空内容
|
|
||||||
default: // InterceptTypeWarmup
|
|
||||||
msgID = "msg_mock_warmup"
|
|
||||||
outputTokens = 2
|
|
||||||
textDeltas = []string{"New", " Conversation"}
|
|
||||||
}
|
|
||||||
|
|
||||||
// Build message_start event with fixed schema.
|
|
||||||
messageStartJSON := `{"type":"message_start","message":{"id":` + strconv.Quote(msgID) + `,"type":"message","role":"assistant","model":` + strconv.Quote(model) + `,"content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}`
|
|
||||||
|
|
||||||
// Build events
|
|
||||||
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"}`,
|
|
||||||
}
|
|
||||||
|
|
||||||
// Add text deltas
|
|
||||||
for _, text := range textDeltas {
|
|
||||||
deltaJSON := `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":` + strconv.Quote(text) + `}}`
|
|
||||||
events = append(events, `event: content_block_delta`+"\n"+`data: `+string(deltaJSON))
|
|
||||||
}
|
|
||||||
|
|
||||||
// Add final events
|
|
||||||
messageDeltaJSON := `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":10,"output_tokens":` + strconv.Itoa(outputTokens) + `}}`
|
|
||||||
|
|
||||||
events = append(events,
|
|
||||||
`event: content_block_stop`+"\n"+`data: {"index":0,"type":"content_block_stop"}`,
|
|
||||||
`event: message_delta`+"\n"+`data: `+string(messageDeltaJSON),
|
|
||||||
`event: message_stop`+"\n"+`data: {"type":"message_stop"}`,
|
|
||||||
)
|
|
||||||
|
|
||||||
for _, event := range events {
|
|
||||||
_, _ = c.Writer.WriteString(event + "\n\n")
|
|
||||||
c.Writer.Flush()
|
|
||||||
time.Sleep(20 * time.Millisecond)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// generateRealisticMsgID 生成仿真的消息 ID（msg_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)
|
|
||||||
}
|
|
||||||
|
|
||||||
// sendMockInterceptResponse 发送非流式 mock 响应（用于请求拦截）
|
|
||||||
func sendMockInterceptResponse(c *gin.Context, model string, interceptType InterceptType) {
|
|
||||||
var msgID, text, stopReason string
|
|
||||||
var outputTokens int
|
|
||||||
|
|
||||||
switch interceptType {
|
|
||||||
case InterceptTypeSuggestionMode:
|
|
||||||
msgID = "msg_mock_suggestion"
|
|
||||||
text = ""
|
|
||||||
outputTokens = 1
|
|
||||||
stopReason = "end_turn"
|
|
||||||
case InterceptTypeMaxTokensOneHaiku:
|
|
||||||
msgID = generateRealisticMsgID()
|
|
||||||
text = "#"
|
|
||||||
outputTokens = 1
|
|
||||||
stopReason = "max_tokens" // max_tokens=1 探测请求的 stop_reason 应为 max_tokens
|
|
||||||
default: // InterceptTypeWarmup
|
|
||||||
msgID = "msg_mock_warmup"
|
|
||||||
text = "New Conversation"
|
|
||||||
outputTokens = 2
|
|
||||||
stopReason = "end_turn"
|
|
||||||
}
|
|
||||||
|
|
||||||
// 构建完整的响应格式（与 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,
|
|
||||||
"usage": gin.H{
|
|
||||||
"input_tokens": 10,
|
|
||||||
"cache_creation_input_tokens": 0,
|
|
||||||
"cache_read_input_tokens": 0,
|
|
||||||
"cache_creation": gin.H{
|
|
||||||
"ephemeral_5m_input_tokens": 0,
|
|
||||||
"ephemeral_1h_input_tokens": 0,
|
|
||||||
},
|
|
||||||
"output_tokens": outputTokens,
|
|
||||||
"total_tokens": 10 + outputTokens,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
c.JSON(http.StatusOK, response)
|
|
||||||
}
|
|
||||||
|
|
||||||
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
|
|
||||||
}
|
|
||||||
if errors.Is(err, service.ErrAPIKeyRateLimit5hExceeded) {
|
|
||||||
msg := pkgerrors.Message(err)
|
|
||||||
return http.StatusTooManyRequests, "rate_limit_exceeded", msg
|
|
||||||
}
|
|
||||||
if errors.Is(err, service.ErrAPIKeyRateLimit1dExceeded) {
|
|
||||||
msg := pkgerrors.Message(err)
|
|
||||||
return http.StatusTooManyRequests, "rate_limit_exceeded", msg
|
|
||||||
}
|
|
||||||
if errors.Is(err, service.ErrAPIKeyRateLimit7dExceeded) {
|
|
||||||
msg := pkgerrors.Message(err)
|
|
||||||
return http.StatusTooManyRequests, "rate_limit_exceeded", msg
|
|
||||||
}
|
|
||||||
msg := pkgerrors.Message(err)
|
|
||||||
if msg == "" {
|
|
||||||
logger.L().With(
|
|
||||||
zap.String("component", "handler.gateway.billing"),
|
|
||||||
zap.Error(err),
|
|
||||||
).Warn("gateway.billing_error_missing_message")
|
|
||||||
msg = "Billing error"
|
|
||||||
}
|
|
||||||
return http.StatusForbidden, "billing_error", msg
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *GatewayHandler) metadataBridgeEnabled() bool {
|
|
||||||
if h == nil || h.cfg == nil {
|
|
||||||
return true
|
|
||||||
}
|
|
||||||
return h.cfg.Gateway.OpenAIWS.MetadataBridgeEnabled
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *GatewayHandler) maybeLogCompatibilityFallbackMetrics(reqLog *zap.Logger) {
|
|
||||||
if reqLog == nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
if gatewayCompatibilityMetricsLogCounter.Add(1)%gatewayCompatibilityMetricsLogInterval != 0 {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
metrics := service.SnapshotOpenAICompatibilityFallbackMetrics()
|
|
||||||
reqLog.Info("gateway.compatibility_fallback_metrics",
|
|
||||||
zap.Int64("session_hash_legacy_read_fallback_total", metrics.SessionHashLegacyReadFallbackTotal),
|
|
||||||
zap.Int64("session_hash_legacy_read_fallback_hit", metrics.SessionHashLegacyReadFallbackHit),
|
|
||||||
zap.Int64("session_hash_legacy_dual_write_total", metrics.SessionHashLegacyDualWriteTotal),
|
|
||||||
zap.Float64("session_hash_legacy_read_hit_rate", metrics.SessionHashLegacyReadHitRate),
|
|
||||||
zap.Int64("metadata_legacy_fallback_total", metrics.MetadataLegacyFallbackTotal),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (h *GatewayHandler) submitUsageRecordTask(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()
|
|
||||||
defer func() {
|
|
||||||
if recovered := recover(); recovered != nil {
|
|
||||||
logger.L().With(
|
|
||||||
zap.String("component", "handler.gateway.messages"),
|
|
||||||
zap.Any("panic", recovered),
|
|
||||||
).Error("gateway.usage_record_task_panic_recovered")
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
task(ctx)
|
|
||||||
}
|
|
||||||
|
|
||||||
// getUserMsgQueueMode 获取当前请求的 UMQ 模式
|
|
||||||
// 返回 "serialize" | "throttle" | ""
|
|
||||||
func (h *GatewayHandler) getUserMsgQueueMode(account *service.Account, parsed *service.ParsedRequest) string {
|
|
||||||
if h.userMsgQueueHelper == nil {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
// 仅适用于 Anthropic OAuth/SetupToken 账号
|
|
||||||
if !account.IsAnthropicOAuthOrSetupToken() {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
if !service.IsRealUserMessage(parsed) {
|
|
||||||
return ""
|
|
||||||
}
|
|
||||||
// 账号级模式优先，fallback 到全局配置
|
|
||||||
mode := account.GetUserMsgQueueMode()
|
|
||||||
if mode == "" {
|
|
||||||
mode = h.cfg.Gateway.UserMessageQueue.GetEffectiveMode()
|
|
||||||
}
|
|
||||||
return mode
|
|
||||||
}
|
|
||||||
|
|||||||
197
backend/internal/handler/gateway_handler_intercept.go
Normal file
197
backend/internal/handler/gateway_handler_intercept.go
Normal file
@@ -0,0 +1,197 @@
|
|||||||
|
package handler
|
||||||
|
|
||||||
|
import (
|
||||||
|
"crypto/rand"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
)
|
||||||
|
|
||||||
|
type InterceptType int
|
||||||
|
|
||||||
|
const (
|
||||||
|
InterceptTypeNone InterceptType = iota
|
||||||
|
InterceptTypeWarmup
|
||||||
|
InterceptTypeSuggestionMode
|
||||||
|
InterceptTypeMaxTokensOneHaiku
|
||||||
|
)
|
||||||
|
|
||||||
|
func isHaikuModel(model string) bool {
|
||||||
|
return strings.Contains(strings.ToLower(model), "haiku")
|
||||||
|
}
|
||||||
|
|
||||||
|
func isMaxTokensOneHaikuRequest(model string, maxTokens int, isStream bool) bool {
|
||||||
|
return maxTokens == 1 && isHaikuModel(model) && !isStream
|
||||||
|
}
|
||||||
|
|
||||||
|
func detectInterceptType(body []byte, model string, maxTokens int, isStream bool, isClaudeCodeClient bool) InterceptType {
|
||||||
|
if isClaudeCodeClient && isMaxTokensOneHaikuRequest(model, maxTokens, isStream) {
|
||||||
|
return InterceptTypeMaxTokensOneHaiku
|
||||||
|
}
|
||||||
|
|
||||||
|
bodyStr := string(body)
|
||||||
|
hasSuggestionMode := strings.Contains(bodyStr, "[SUGGESTION MODE:")
|
||||||
|
hasWarmupKeyword := strings.Contains(bodyStr, "title") || strings.Contains(bodyStr, "Warmup")
|
||||||
|
|
||||||
|
if !hasSuggestionMode && !hasWarmupKeyword {
|
||||||
|
return InterceptTypeNone
|
||||||
|
}
|
||||||
|
|
||||||
|
var req struct {
|
||||||
|
Messages []struct {
|
||||||
|
Role string `json:"role"`
|
||||||
|
Content []struct {
|
||||||
|
Type string `json:"type"`
|
||||||
|
Text string `json:"text"`
|
||||||
|
} `json:"content"`
|
||||||
|
} `json:"messages"`
|
||||||
|
System []struct {
|
||||||
|
Text string `json:"text"`
|
||||||
|
} `json:"system"`
|
||||||
|
}
|
||||||
|
if err := json.Unmarshal(body, &req); err != nil {
|
||||||
|
return InterceptTypeNone
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if hasWarmupKeyword {
|
||||||
|
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
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, sys := range req.System {
|
||||||
|
if strings.Contains(sys.Text, "nalyze if this message indicates a new conversation topic. If it does, extract a 2-3 word title") {
|
||||||
|
return InterceptTypeWarmup
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return InterceptTypeNone
|
||||||
|
}
|
||||||
|
|
||||||
|
func sendMockInterceptStream(c *gin.Context, model string, interceptType InterceptType) {
|
||||||
|
c.Header("Content-Type", "text/event-stream")
|
||||||
|
c.Header("Cache-Control", "no-cache")
|
||||||
|
c.Header("Connection", "keep-alive")
|
||||||
|
c.Header("X-Accel-Buffering", "no")
|
||||||
|
|
||||||
|
var msgID string
|
||||||
|
var outputTokens int
|
||||||
|
var textDeltas []string
|
||||||
|
|
||||||
|
switch interceptType {
|
||||||
|
case InterceptTypeSuggestionMode:
|
||||||
|
msgID = "msg_mock_suggestion"
|
||||||
|
outputTokens = 1
|
||||||
|
textDeltas = []string{""}
|
||||||
|
default:
|
||||||
|
msgID = "msg_mock_warmup"
|
||||||
|
outputTokens = 2
|
||||||
|
textDeltas = []string{"New", " Conversation"}
|
||||||
|
}
|
||||||
|
|
||||||
|
messageStartJSON := `{"type":"message_start","message":{"id":` + strconv.Quote(msgID) + `,"type":"message","role":"assistant","model":` + strconv.Quote(model) + `,"content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}`
|
||||||
|
|
||||||
|
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"}`,
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, text := range textDeltas {
|
||||||
|
deltaJSON := `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":` + strconv.Quote(text) + `}}`
|
||||||
|
events = append(events, `event: content_block_delta`+"\n"+`data: `+string(deltaJSON))
|
||||||
|
}
|
||||||
|
|
||||||
|
messageDeltaJSON := `{"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"input_tokens":10,"output_tokens":` + strconv.Itoa(outputTokens) + `}}`
|
||||||
|
|
||||||
|
events = append(events,
|
||||||
|
`event: content_block_stop`+"\n"+`data: {"index":0,"type":"content_block_stop"}`,
|
||||||
|
`event: message_delta`+"\n"+`data: `+string(messageDeltaJSON),
|
||||||
|
`event: message_stop`+"\n"+`data: {"type":"message_stop"}`,
|
||||||
|
)
|
||||||
|
|
||||||
|
for _, event := range events {
|
||||||
|
_, _ = c.Writer.WriteString(event + "\n\n")
|
||||||
|
c.Writer.Flush()
|
||||||
|
time.Sleep(20 * time.Millisecond)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
|
||||||
|
func sendMockInterceptResponse(c *gin.Context, model string, interceptType InterceptType) {
|
||||||
|
var msgID, text, stopReason string
|
||||||
|
var outputTokens int
|
||||||
|
|
||||||
|
switch interceptType {
|
||||||
|
case InterceptTypeSuggestionMode:
|
||||||
|
msgID = "msg_mock_suggestion"
|
||||||
|
text = ""
|
||||||
|
outputTokens = 1
|
||||||
|
stopReason = "end_turn"
|
||||||
|
case InterceptTypeMaxTokensOneHaiku:
|
||||||
|
msgID = generateRealisticMsgID()
|
||||||
|
text = "#"
|
||||||
|
outputTokens = 1
|
||||||
|
stopReason = "max_tokens"
|
||||||
|
default:
|
||||||
|
msgID = "msg_mock_warmup"
|
||||||
|
text = "New Conversation"
|
||||||
|
outputTokens = 2
|
||||||
|
stopReason = "end_turn"
|
||||||
|
}
|
||||||
|
|
||||||
|
response := gin.H{
|
||||||
|
"model": model,
|
||||||
|
"id": msgID,
|
||||||
|
"type": "message",
|
||||||
|
"role": "assistant",
|
||||||
|
"content": []gin.H{{"type": "text", "text": text}},
|
||||||
|
"stop_reason": stopReason,
|
||||||
|
"stop_sequence": nil,
|
||||||
|
"usage": gin.H{
|
||||||
|
"input_tokens": 10,
|
||||||
|
"cache_creation_input_tokens": 0,
|
||||||
|
"cache_read_input_tokens": 0,
|
||||||
|
"cache_creation": gin.H{
|
||||||
|
"ephemeral_5m_input_tokens": 0,
|
||||||
|
"ephemeral_1h_input_tokens": 0,
|
||||||
|
},
|
||||||
|
"output_tokens": outputTokens,
|
||||||
|
"total_tokens": 10 + outputTokens,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
c.JSON(http.StatusOK, response)
|
||||||
|
}
|
||||||
246
backend/internal/handler/gateway_handler_support.go
Normal file
246
backend/internal/handler/gateway_handler_support.go
Normal file
@@ -0,0 +1,246 @@
|
|||||||
|
package handler
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"strconv"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
pkgerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||||
|
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||||
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||||
|
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
"go.uber.org/zap"
|
||||||
|
)
|
||||||
|
|
||||||
|
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, 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
|
||||||
|
}
|
||||||
|
|
||||||
|
if rule.SkipMonitoring {
|
||||||
|
c.Set(service.OpsSkipPassthroughKey, true)
|
||||||
|
}
|
||||||
|
|
||||||
|
h.handleStreamingAwareError(c, respCode, "upstream_error", msg, streamStarted)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
upstreamMsg := service.ExtractUpstreamErrorMessage(responseBody)
|
||||||
|
service.SetOpsUpstreamError(c, statusCode, upstreamMsg, "")
|
||||||
|
|
||||||
|
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||||||
|
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) handleFailoverExhaustedSimple(c *gin.Context, statusCode int, streamStarted bool) {
|
||||||
|
status, errType, errMsg := h.mapUpstreamError(statusCode)
|
||||||
|
service.SetOpsUpstreamError(c, statusCode, errMsg, "")
|
||||||
|
h.handleStreamingAwareError(c, status, errType, errMsg, streamStarted)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) mapUpstreamError(statusCode int) (int, string, string) {
|
||||||
|
switch statusCode {
|
||||||
|
case 401:
|
||||||
|
return http.StatusBadGateway, "upstream_error", "Upstream authentication failed, please contact administrator"
|
||||||
|
case 403:
|
||||||
|
return http.StatusBadGateway, "upstream_error", "Upstream access forbidden, please contact administrator"
|
||||||
|
case 429:
|
||||||
|
return http.StatusTooManyRequests, "rate_limit_error", "Upstream rate limit exceeded, please retry later"
|
||||||
|
case 529:
|
||||||
|
return http.StatusServiceUnavailable, "overloaded_error", "Upstream service overloaded, please retry later"
|
||||||
|
case 500, 502, 503, 504:
|
||||||
|
return http.StatusBadGateway, "upstream_error", "Upstream service temporarily unavailable"
|
||||||
|
default:
|
||||||
|
return http.StatusBadGateway, "upstream_error", "Upstream request failed"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) {
|
||||||
|
if streamStarted {
|
||||||
|
flusher, ok := c.Writer.(http.Flusher)
|
||||||
|
if ok {
|
||||||
|
errorEvent := `data: {"type":"error","error":{"type":` + strconv.Quote(errType) + `,"message":` + strconv.Quote(message) + `}}` + "\n\n"
|
||||||
|
if _, err := fmt.Fprint(c.Writer, errorEvent); err != nil {
|
||||||
|
_ = c.Error(err)
|
||||||
|
}
|
||||||
|
flusher.Flush()
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
h.errorResponse(c, status, errType, message)
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) checkClaudeCodeVersion(c *gin.Context) bool {
|
||||||
|
ctx := c.Request.Context()
|
||||||
|
if !service.IsClaudeCodeClient(ctx) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
if strings.HasSuffix(c.Request.URL.Path, "/count_tokens") {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
minVersion, maxVersion := h.settingService.GetClaudeCodeVersionBounds(ctx)
|
||||||
|
if minVersion == "" && maxVersion == "" {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
clientVersion := service.GetClaudeCodeVersion(ctx)
|
||||||
|
if clientVersion == "" {
|
||||||
|
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
||||||
|
"Unable to determine Claude Code version. Please update Claude Code: npm update -g @anthropic-ai/claude-code")
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
if minVersion != "" && service.CompareVersions(clientVersion, minVersion) < 0 {
|
||||||
|
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
||||||
|
fmt.Sprintf("Your Claude Code version (%s) is below the minimum required version (%s). Please update: npm update -g @anthropic-ai/claude-code",
|
||||||
|
clientVersion, minVersion))
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
if maxVersion != "" && service.CompareVersions(clientVersion, maxVersion) > 0 {
|
||||||
|
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error",
|
||||||
|
fmt.Sprintf("Your Claude Code version (%s) exceeds the maximum allowed version (%s). "+
|
||||||
|
"Please downgrade: npm install -g @anthropic-ai/claude-code@%s && "+
|
||||||
|
"set CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC=1 to prevent auto-upgrade",
|
||||||
|
clientVersion, maxVersion, maxVersion))
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
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,
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
}
|
||||||
|
if errors.Is(err, service.ErrAPIKeyRateLimit5hExceeded) {
|
||||||
|
msg := pkgerrors.Message(err)
|
||||||
|
return http.StatusTooManyRequests, "rate_limit_exceeded", msg
|
||||||
|
}
|
||||||
|
if errors.Is(err, service.ErrAPIKeyRateLimit1dExceeded) {
|
||||||
|
msg := pkgerrors.Message(err)
|
||||||
|
return http.StatusTooManyRequests, "rate_limit_exceeded", msg
|
||||||
|
}
|
||||||
|
if errors.Is(err, service.ErrAPIKeyRateLimit7dExceeded) {
|
||||||
|
msg := pkgerrors.Message(err)
|
||||||
|
return http.StatusTooManyRequests, "rate_limit_exceeded", msg
|
||||||
|
}
|
||||||
|
msg := pkgerrors.Message(err)
|
||||||
|
if msg == "" {
|
||||||
|
logger.L().With(
|
||||||
|
zap.String("component", "handler.gateway.billing"),
|
||||||
|
zap.Error(err),
|
||||||
|
).Warn("gateway.billing_error_missing_message")
|
||||||
|
msg = "Billing error"
|
||||||
|
}
|
||||||
|
return http.StatusForbidden, "billing_error", msg
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) metadataBridgeEnabled() bool {
|
||||||
|
if h == nil || h.cfg == nil {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return h.cfg.Gateway.OpenAIWS.MetadataBridgeEnabled
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) maybeLogCompatibilityFallbackMetrics(reqLog *zap.Logger) {
|
||||||
|
if reqLog == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if gatewayCompatibilityMetricsLogCounter.Add(1)%gatewayCompatibilityMetricsLogInterval != 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
metrics := service.SnapshotOpenAICompatibilityFallbackMetrics()
|
||||||
|
reqLog.Info("gateway.compatibility_fallback_metrics",
|
||||||
|
zap.Int64("session_hash_legacy_read_fallback_total", metrics.SessionHashLegacyReadFallbackTotal),
|
||||||
|
zap.Int64("session_hash_legacy_read_fallback_hit", metrics.SessionHashLegacyReadFallbackHit),
|
||||||
|
zap.Int64("session_hash_legacy_dual_write_total", metrics.SessionHashLegacyDualWriteTotal),
|
||||||
|
zap.Float64("session_hash_legacy_read_hit_rate", metrics.SessionHashLegacyReadHitRate),
|
||||||
|
zap.Int64("metadata_legacy_fallback_total", metrics.MetadataLegacyFallbackTotal),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) submitUsageRecordTask(task service.UsageRecordTask) {
|
||||||
|
if task == nil {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
if h.usageRecordWorkerPool != nil {
|
||||||
|
h.usageRecordWorkerPool.Submit(task)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||||
|
defer cancel()
|
||||||
|
defer func() {
|
||||||
|
if recovered := recover(); recovered != nil {
|
||||||
|
logger.L().With(
|
||||||
|
zap.String("component", "handler.gateway.messages"),
|
||||||
|
zap.Any("panic", recovered),
|
||||||
|
).Error("gateway.usage_record_task_panic_recovered")
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
task(ctx)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) getUserMsgQueueMode(account *service.Account, parsed *service.ParsedRequest) string {
|
||||||
|
if h.userMsgQueueHelper == nil {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
if !account.IsAnthropicOAuthOrSetupToken() {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
if !service.IsRealUserMessage(parsed) {
|
||||||
|
return ""
|
||||||
|
}
|
||||||
|
mode := account.GetUserMsgQueueMode()
|
||||||
|
if mode == "" {
|
||||||
|
mode = h.cfg.Gateway.UserMessageQueue.GetEffectiveMode()
|
||||||
|
}
|
||||||
|
return mode
|
||||||
|
}
|
||||||
261
backend/internal/handler/gateway_handler_usage.go
Normal file
261
backend/internal/handler/gateway_handler_usage.go
Normal file
@@ -0,0 +1,261 @@
|
|||||||
|
package handler
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"net/http"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||||
|
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||||
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||||
|
|
||||||
|
"github.com/gin-gonic/gin"
|
||||||
|
)
|
||||||
|
|
||||||
|
func cloneAPIKeyWithGroup(apiKey *service.APIKey, group *service.Group) *service.APIKey {
|
||||||
|
if apiKey == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
cloned := *apiKey
|
||||||
|
cloned.Group = group
|
||||||
|
if group != nil {
|
||||||
|
cloned.GroupID = &group.ID
|
||||||
|
}
|
||||||
|
return &cloned
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) parseUsageDateRange(c *gin.Context) (time.Time, time.Time) {
|
||||||
|
now := timezone.Now()
|
||||||
|
endTime := now
|
||||||
|
startTime := now.AddDate(0, 0, -30)
|
||||||
|
|
||||||
|
if s := c.Query("start_date"); s != "" {
|
||||||
|
if t, err := timezone.ParseInLocation("2006-01-02", s); err == nil {
|
||||||
|
startTime = t
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if s := c.Query("end_date"); s != "" {
|
||||||
|
if t, err := timezone.ParseInLocation("2006-01-02", s); err == nil {
|
||||||
|
endTime = t.AddDate(0, 0, 1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return startTime, endTime
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) buildUsageData(ctx context.Context, apiKeyID int64) gin.H {
|
||||||
|
if h.usageService == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
dashStats, err := h.usageService.GetAPIKeyDashboardStats(ctx, apiKeyID)
|
||||||
|
if err != nil || dashStats == nil {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
return gin.H{
|
||||||
|
"today": gin.H{
|
||||||
|
"requests": dashStats.TodayRequests,
|
||||||
|
"input_tokens": dashStats.TodayInputTokens,
|
||||||
|
"output_tokens": dashStats.TodayOutputTokens,
|
||||||
|
"cache_creation_tokens": dashStats.TodayCacheCreationTokens,
|
||||||
|
"cache_read_tokens": dashStats.TodayCacheReadTokens,
|
||||||
|
"total_tokens": dashStats.TodayTokens,
|
||||||
|
"cost": dashStats.TodayCost,
|
||||||
|
"actual_cost": dashStats.TodayActualCost,
|
||||||
|
},
|
||||||
|
"total": gin.H{
|
||||||
|
"requests": dashStats.TotalRequests,
|
||||||
|
"input_tokens": dashStats.TotalInputTokens,
|
||||||
|
"output_tokens": dashStats.TotalOutputTokens,
|
||||||
|
"cache_creation_tokens": dashStats.TotalCacheCreationTokens,
|
||||||
|
"cache_read_tokens": dashStats.TotalCacheReadTokens,
|
||||||
|
"total_tokens": dashStats.TotalTokens,
|
||||||
|
"cost": dashStats.TotalCost,
|
||||||
|
"actual_cost": dashStats.TotalActualCost,
|
||||||
|
},
|
||||||
|
"average_duration_ms": dashStats.AverageDurationMs,
|
||||||
|
"rpm": dashStats.Rpm,
|
||||||
|
"tpm": dashStats.Tpm,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) usageQuotaLimited(c *gin.Context, ctx context.Context, apiKey *service.APIKey, usageData gin.H, modelStats any) {
|
||||||
|
resp := gin.H{
|
||||||
|
"mode": "quota_limited",
|
||||||
|
"isValid": apiKey.Status == service.StatusAPIKeyActive || apiKey.Status == service.StatusAPIKeyQuotaExhausted || apiKey.Status == service.StatusAPIKeyExpired,
|
||||||
|
"status": apiKey.Status,
|
||||||
|
}
|
||||||
|
|
||||||
|
if apiKey.Quota > 0 {
|
||||||
|
remaining := apiKey.GetQuotaRemaining()
|
||||||
|
resp["quota"] = gin.H{
|
||||||
|
"limit": apiKey.Quota,
|
||||||
|
"used": apiKey.QuotaUsed,
|
||||||
|
"remaining": remaining,
|
||||||
|
"unit": "USD",
|
||||||
|
}
|
||||||
|
resp["remaining"] = remaining
|
||||||
|
resp["unit"] = "USD"
|
||||||
|
}
|
||||||
|
|
||||||
|
if apiKey.HasRateLimits() && h.apiKeyService != nil {
|
||||||
|
rateLimitData, err := h.apiKeyService.GetRateLimitData(ctx, apiKey.ID)
|
||||||
|
if err == nil && rateLimitData != nil {
|
||||||
|
var rateLimits []gin.H
|
||||||
|
if apiKey.RateLimit5h > 0 {
|
||||||
|
used := rateLimitData.EffectiveUsage5h()
|
||||||
|
entry := gin.H{
|
||||||
|
"window": "5h",
|
||||||
|
"limit": apiKey.RateLimit5h,
|
||||||
|
"used": used,
|
||||||
|
"remaining": max(0, apiKey.RateLimit5h-used),
|
||||||
|
"window_start": rateLimitData.Window5hStart,
|
||||||
|
}
|
||||||
|
if rateLimitData.Window5hStart != nil && !service.IsWindowExpired(rateLimitData.Window5hStart, service.RateLimitWindow5h) {
|
||||||
|
entry["reset_at"] = rateLimitData.Window5hStart.Add(service.RateLimitWindow5h)
|
||||||
|
}
|
||||||
|
rateLimits = append(rateLimits, entry)
|
||||||
|
}
|
||||||
|
if apiKey.RateLimit1d > 0 {
|
||||||
|
used := rateLimitData.EffectiveUsage1d()
|
||||||
|
entry := gin.H{
|
||||||
|
"window": "1d",
|
||||||
|
"limit": apiKey.RateLimit1d,
|
||||||
|
"used": used,
|
||||||
|
"remaining": max(0, apiKey.RateLimit1d-used),
|
||||||
|
"window_start": rateLimitData.Window1dStart,
|
||||||
|
}
|
||||||
|
if rateLimitData.Window1dStart != nil && !service.IsWindowExpired(rateLimitData.Window1dStart, service.RateLimitWindow1d) {
|
||||||
|
entry["reset_at"] = rateLimitData.Window1dStart.Add(service.RateLimitWindow1d)
|
||||||
|
}
|
||||||
|
rateLimits = append(rateLimits, entry)
|
||||||
|
}
|
||||||
|
if apiKey.RateLimit7d > 0 {
|
||||||
|
used := rateLimitData.EffectiveUsage7d()
|
||||||
|
entry := gin.H{
|
||||||
|
"window": "7d",
|
||||||
|
"limit": apiKey.RateLimit7d,
|
||||||
|
"used": used,
|
||||||
|
"remaining": max(0, apiKey.RateLimit7d-used),
|
||||||
|
"window_start": rateLimitData.Window7dStart,
|
||||||
|
}
|
||||||
|
if rateLimitData.Window7dStart != nil && !service.IsWindowExpired(rateLimitData.Window7dStart, service.RateLimitWindow7d) {
|
||||||
|
entry["reset_at"] = rateLimitData.Window7dStart.Add(service.RateLimitWindow7d)
|
||||||
|
}
|
||||||
|
rateLimits = append(rateLimits, entry)
|
||||||
|
}
|
||||||
|
if len(rateLimits) > 0 {
|
||||||
|
resp["rate_limits"] = rateLimits
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if apiKey.ExpiresAt != nil {
|
||||||
|
resp["expires_at"] = apiKey.ExpiresAt
|
||||||
|
resp["days_until_expiry"] = apiKey.GetDaysUntilExpiry()
|
||||||
|
}
|
||||||
|
|
||||||
|
if usageData != nil {
|
||||||
|
resp["usage"] = usageData
|
||||||
|
}
|
||||||
|
if modelStats != nil {
|
||||||
|
resp["model_stats"] = modelStats
|
||||||
|
}
|
||||||
|
|
||||||
|
c.JSON(http.StatusOK, resp)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (h *GatewayHandler) usageUnrestricted(c *gin.Context, ctx context.Context, apiKey *service.APIKey, subject middleware2.AuthSubject, usageData gin.H, modelStats any) {
|
||||||
|
if apiKey.Group != nil && apiKey.Group.IsSubscriptionType() {
|
||||||
|
resp := gin.H{
|
||||||
|
"mode": "unrestricted",
|
||||||
|
"isValid": true,
|
||||||
|
"planName": apiKey.Group.Name,
|
||||||
|
"unit": "USD",
|
||||||
|
}
|
||||||
|
|
||||||
|
subscription, ok := middleware2.GetSubscriptionFromContext(c)
|
||||||
|
if ok {
|
||||||
|
remaining := h.calculateSubscriptionRemaining(apiKey.Group, subscription)
|
||||||
|
resp["remaining"] = remaining
|
||||||
|
resp["subscription"] = gin.H{
|
||||||
|
"daily_usage_usd": subscription.DailyUsageUSD,
|
||||||
|
"weekly_usage_usd": subscription.WeeklyUsageUSD,
|
||||||
|
"monthly_usage_usd": subscription.MonthlyUsageUSD,
|
||||||
|
"daily_limit_usd": apiKey.Group.DailyLimitUSD,
|
||||||
|
"weekly_limit_usd": apiKey.Group.WeeklyLimitUSD,
|
||||||
|
"monthly_limit_usd": apiKey.Group.MonthlyLimitUSD,
|
||||||
|
"expires_at": subscription.ExpiresAt,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if usageData != nil {
|
||||||
|
resp["usage"] = usageData
|
||||||
|
}
|
||||||
|
if modelStats != nil {
|
||||||
|
resp["model_stats"] = modelStats
|
||||||
|
}
|
||||||
|
c.JSON(http.StatusOK, resp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
latestUser, err := h.userService.GetByID(ctx, subject.UserID)
|
||||||
|
if err != nil {
|
||||||
|
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Failed to get user info")
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
resp := gin.H{
|
||||||
|
"mode": "unrestricted",
|
||||||
|
"isValid": true,
|
||||||
|
"planName": "钱包余额",
|
||||||
|
"remaining": latestUser.Balance,
|
||||||
|
"unit": "USD",
|
||||||
|
"balance": latestUser.Balance,
|
||||||
|
}
|
||||||
|
if usageData != nil {
|
||||||
|
resp["usage"] = usageData
|
||||||
|
}
|
||||||
|
if modelStats != nil {
|
||||||
|
resp["model_stats"] = modelStats
|
||||||
|
}
|
||||||
|
c.JSON(http.StatusOK, resp)
|
||||||
|
}
|
||||||
|
|
||||||
|
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)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(remainingValues) == 0 {
|
||||||
|
return -1
|
||||||
|
}
|
||||||
|
|
||||||
|
minValue := remainingValues[0]
|
||||||
|
for _, v := range remainingValues[1:] {
|
||||||
|
if v < minValue {
|
||||||
|
minValue = v
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return minValue
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user