2025-12-18 13:50:39 +08:00
|
|
|
|
package service
|
|
|
|
|
|
|
|
|
|
|
|
import (
|
|
|
|
|
|
"context"
|
2026-01-03 06:34:00 -08:00
|
|
|
|
"encoding/json"
|
2026-01-14 15:55:44 +08:00
|
|
|
|
"log/slog"
|
2025-12-18 13:50:39 +08:00
|
|
|
|
"net/http"
|
|
|
|
|
|
"strconv"
|
2026-01-01 04:22:39 +08:00
|
|
|
|
"strings"
|
2026-01-01 08:29:57 +08:00
|
|
|
|
"sync"
|
2025-12-18 13:50:39 +08:00
|
|
|
|
"time"
|
|
|
|
|
|
|
2025-12-24 21:07:21 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
2026-02-28 15:01:20 +08:00
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
2025-12-18 13:50:39 +08:00
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
// RateLimitService 处理限流和过载状态管理
|
|
|
|
|
|
type RateLimitService struct {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
accountRepo AccountRepository
|
|
|
|
|
|
usageRepo UsageLogRepository
|
|
|
|
|
|
cfg *config.Config
|
|
|
|
|
|
geminiQuotaService *GeminiQuotaService
|
|
|
|
|
|
tempUnschedCache TempUnschedCache
|
|
|
|
|
|
timeoutCounterCache TimeoutCounterCache
|
|
|
|
|
|
settingService *SettingService
|
|
|
|
|
|
tokenCacheInvalidator TokenCacheInvalidator
|
|
|
|
|
|
usageCacheMu sync.RWMutex
|
|
|
|
|
|
usageCache map[int64]*geminiUsageCacheEntry
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-03-08 06:59:53 +08:00
|
|
|
|
// SuccessfulTestRecoveryResult 表示测试成功后恢复了哪些运行时状态。
|
|
|
|
|
|
type SuccessfulTestRecoveryResult struct {
|
|
|
|
|
|
ClearedError bool
|
|
|
|
|
|
ClearedRateLimit bool
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// AccountRecoveryOptions 控制账号恢复时的附加行为。
|
|
|
|
|
|
type AccountRecoveryOptions struct {
|
|
|
|
|
|
InvalidateToken bool
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-01 08:29:57 +08:00
|
|
|
|
type geminiUsageCacheEntry struct {
|
|
|
|
|
|
windowStart time.Time
|
|
|
|
|
|
cachedAt time.Time
|
|
|
|
|
|
totals GeminiUsageTotals
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-28 15:01:20 +08:00
|
|
|
|
type geminiUsageTotalsBatchProvider interface {
|
|
|
|
|
|
GetGeminiUsageTotalsBatch(ctx context.Context, accountIDs []int64, startTime, endTime time.Time) (map[int64]GeminiUsageTotals, error)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-01 08:29:57 +08:00
|
|
|
|
const geminiPrecheckCacheTTL = time.Minute
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// NewRateLimitService 创建RateLimitService实例
|
2026-01-03 06:34:00 -08:00
|
|
|
|
func NewRateLimitService(accountRepo AccountRepository, usageRepo UsageLogRepository, cfg *config.Config, geminiQuotaService *GeminiQuotaService, tempUnschedCache TempUnschedCache) *RateLimitService {
|
2025-12-18 13:50:39 +08:00
|
|
|
|
return &RateLimitService{
|
2026-01-01 04:22:39 +08:00
|
|
|
|
accountRepo: accountRepo,
|
|
|
|
|
|
usageRepo: usageRepo,
|
|
|
|
|
|
cfg: cfg,
|
|
|
|
|
|
geminiQuotaService: geminiQuotaService,
|
2026-01-03 06:34:00 -08:00
|
|
|
|
tempUnschedCache: tempUnschedCache,
|
2026-01-01 08:29:57 +08:00
|
|
|
|
usageCache: make(map[int64]*geminiUsageCacheEntry),
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-11 21:54:52 -08:00
|
|
|
|
// SetTimeoutCounterCache 设置超时计数器缓存(可选依赖)
|
|
|
|
|
|
func (s *RateLimitService) SetTimeoutCounterCache(cache TimeoutCounterCache) {
|
|
|
|
|
|
s.timeoutCounterCache = cache
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// SetSettingService 设置系统设置服务(可选依赖)
|
|
|
|
|
|
func (s *RateLimitService) SetSettingService(settingService *SettingService) {
|
|
|
|
|
|
s.settingService = settingService
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-14 15:55:44 +08:00
|
|
|
|
// SetTokenCacheInvalidator 设置 token 缓存清理器(可选依赖)
|
|
|
|
|
|
func (s *RateLimitService) SetTokenCacheInvalidator(invalidator TokenCacheInvalidator) {
|
|
|
|
|
|
s.tokenCacheInvalidator = invalidator
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-09 06:54:42 +08:00
|
|
|
|
// ErrorPolicyResult 表示错误策略检查的结果
|
|
|
|
|
|
type ErrorPolicyResult int
|
|
|
|
|
|
|
|
|
|
|
|
const (
|
|
|
|
|
|
ErrorPolicyNone ErrorPolicyResult = iota // 未命中任何策略,继续默认逻辑
|
|
|
|
|
|
ErrorPolicySkipped // 自定义错误码开启但未命中,跳过处理
|
|
|
|
|
|
ErrorPolicyMatched // 自定义错误码命中,应停止调度
|
|
|
|
|
|
ErrorPolicyTempUnscheduled // 临时不可调度规则命中
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
// CheckErrorPolicy 检查自定义错误码和临时不可调度规则。
|
|
|
|
|
|
// 自定义错误码开启时覆盖后续所有逻辑(包括临时不可调度)。
|
|
|
|
|
|
func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Account, statusCode int, responseBody []byte) ErrorPolicyResult {
|
|
|
|
|
|
if account.IsCustomErrorCodesEnabled() {
|
|
|
|
|
|
if account.ShouldHandleErrorCode(statusCode) {
|
|
|
|
|
|
return ErrorPolicyMatched
|
|
|
|
|
|
}
|
|
|
|
|
|
slog.Info("account_error_code_skipped", "account_id", account.ID, "status_code", statusCode)
|
|
|
|
|
|
return ErrorPolicySkipped
|
|
|
|
|
|
}
|
2026-03-08 13:57:23 +08:00
|
|
|
|
if account.IsPoolMode() {
|
|
|
|
|
|
return ErrorPolicySkipped
|
|
|
|
|
|
}
|
2026-02-09 06:54:42 +08:00
|
|
|
|
if s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
|
|
|
|
|
|
return ErrorPolicyTempUnscheduled
|
|
|
|
|
|
}
|
|
|
|
|
|
return ErrorPolicyNone
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// HandleUpstreamError 处理上游错误响应,标记账号状态
|
|
|
|
|
|
// 返回是否应该停止该账号的调度
|
2025-12-26 15:40:24 +08:00
|
|
|
|
func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) (shouldDisable bool) {
|
2026-03-08 13:57:23 +08:00
|
|
|
|
customErrorCodesEnabled := account.IsCustomErrorCodesEnabled()
|
|
|
|
|
|
|
|
|
|
|
|
// 池模式默认不标记本地账号状态;仅当用户显式配置自定义错误码时按本地策略处理。
|
|
|
|
|
|
if account.IsPoolMode() && !customErrorCodesEnabled {
|
|
|
|
|
|
slog.Info("pool_mode_error_skipped", "account_id", account.ID, "status_code", statusCode)
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// apikey 类型账号:检查自定义错误码配置
|
|
|
|
|
|
// 如果启用且错误码不在列表中,则不处理(不停止调度、不标记限流/过载)
|
|
|
|
|
|
if !account.ShouldHandleErrorCode(statusCode) {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Info("account_error_code_skipped", "account_id", account.ID, "status_code", statusCode)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-19 10:51:51 +08:00
|
|
|
|
// 先尝试临时不可调度规则(401除外)
|
|
|
|
|
|
// 如果匹配成功,直接返回,不执行后续禁用逻辑
|
2026-01-15 15:06:34 +08:00
|
|
|
|
if statusCode != 401 {
|
2026-01-19 10:51:51 +08:00
|
|
|
|
if s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
2026-01-14 15:55:44 +08:00
|
|
|
|
}
|
2026-01-19 10:51:51 +08:00
|
|
|
|
|
2026-01-11 15:30:27 +08:00
|
|
|
|
upstreamMsg := strings.TrimSpace(extractUpstreamErrorMessage(responseBody))
|
|
|
|
|
|
upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
|
|
|
|
|
|
if upstreamMsg != "" {
|
|
|
|
|
|
upstreamMsg = truncateForLog([]byte(upstreamMsg), 512)
|
|
|
|
|
|
}
|
2026-01-03 06:34:00 -08:00
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
switch statusCode {
|
2026-01-19 10:51:51 +08:00
|
|
|
|
case 400:
|
|
|
|
|
|
// 只有当错误信息包含 "organization has been disabled" 时才禁用
|
|
|
|
|
|
if strings.Contains(strings.ToLower(upstreamMsg), "organization has been disabled") {
|
|
|
|
|
|
msg := "Organization disabled (400): " + upstreamMsg
|
|
|
|
|
|
s.handleAuthError(ctx, account, msg)
|
|
|
|
|
|
shouldDisable = true
|
|
|
|
|
|
}
|
|
|
|
|
|
// 其他 400 错误(如参数问题)不处理,不禁用账号
|
2025-12-18 13:50:39 +08:00
|
|
|
|
case 401:
|
2026-01-15 19:08:07 +08:00
|
|
|
|
// 对所有 OAuth 账号在 401 错误时调用缓存失效并强制下次刷新
|
2026-01-15 18:27:06 +08:00
|
|
|
|
if account.Type == AccountTypeOAuth {
|
2026-01-15 19:08:07 +08:00
|
|
|
|
// 1. 失效缓存
|
2026-01-15 15:06:34 +08:00
|
|
|
|
if s.tokenCacheInvalidator != nil {
|
|
|
|
|
|
if err := s.tokenCacheInvalidator.InvalidateToken(ctx, account); err != nil {
|
|
|
|
|
|
slog.Warn("oauth_401_invalidate_cache_failed", "account_id", account.ID, "error", err)
|
2026-01-14 15:55:44 +08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
2026-01-15 19:08:07 +08:00
|
|
|
|
// 2. 设置 expires_at 为当前时间,强制下次请求刷新 token
|
|
|
|
|
|
if account.Credentials == nil {
|
|
|
|
|
|
account.Credentials = make(map[string]any)
|
|
|
|
|
|
}
|
|
|
|
|
|
account.Credentials["expires_at"] = time.Now().Format(time.RFC3339)
|
|
|
|
|
|
if err := s.accountRepo.Update(ctx, account); err != nil {
|
|
|
|
|
|
slog.Warn("oauth_401_force_refresh_update_failed", "account_id", account.ID, "error", err)
|
|
|
|
|
|
} else {
|
|
|
|
|
|
slog.Info("oauth_401_force_refresh_set", "account_id", account.ID, "platform", account.Platform)
|
|
|
|
|
|
}
|
2026-03-02 22:54:38 +08:00
|
|
|
|
// 3. 临时不可调度,替代 SetError(保持 status=active 让刷新服务能拾取)
|
|
|
|
|
|
msg := "Authentication failed (401): invalid or expired credentials"
|
|
|
|
|
|
if upstreamMsg != "" {
|
|
|
|
|
|
msg = "OAuth 401: " + upstreamMsg
|
|
|
|
|
|
}
|
|
|
|
|
|
cooldownMinutes := s.cfg.RateLimit.OAuth401CooldownMinutes
|
|
|
|
|
|
if cooldownMinutes <= 0 {
|
|
|
|
|
|
cooldownMinutes = 10
|
|
|
|
|
|
}
|
|
|
|
|
|
until := time.Now().Add(time.Duration(cooldownMinutes) * time.Minute)
|
|
|
|
|
|
if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, msg); err != nil {
|
|
|
|
|
|
slog.Warn("oauth_401_set_temp_unschedulable_failed", "account_id", account.ID, "error", err)
|
|
|
|
|
|
}
|
|
|
|
|
|
shouldDisable = true
|
|
|
|
|
|
} else {
|
|
|
|
|
|
// 非 OAuth 账号(APIKey):保持原有 SetError 行为
|
|
|
|
|
|
msg := "Authentication failed (401): invalid or expired credentials"
|
|
|
|
|
|
if upstreamMsg != "" {
|
|
|
|
|
|
msg = "Authentication failed (401): " + upstreamMsg
|
|
|
|
|
|
}
|
|
|
|
|
|
s.handleAuthError(ctx, account, msg)
|
|
|
|
|
|
shouldDisable = true
|
2026-01-14 15:55:44 +08:00
|
|
|
|
}
|
2025-12-31 11:46:53 +08:00
|
|
|
|
case 402:
|
|
|
|
|
|
// 支付要求:余额不足或计费问题,停止调度
|
2026-01-11 15:30:27 +08:00
|
|
|
|
msg := "Payment required (402): insufficient balance or billing issue"
|
|
|
|
|
|
if upstreamMsg != "" {
|
|
|
|
|
|
msg = "Payment required (402): " + upstreamMsg
|
|
|
|
|
|
}
|
|
|
|
|
|
s.handleAuthError(ctx, account, msg)
|
2026-01-03 06:34:00 -08:00
|
|
|
|
shouldDisable = true
|
2025-12-18 13:50:39 +08:00
|
|
|
|
case 403:
|
|
|
|
|
|
// 禁止访问:停止调度,记录错误
|
2026-01-11 15:30:27 +08:00
|
|
|
|
msg := "Access forbidden (403): account may be suspended or lack permissions"
|
|
|
|
|
|
if upstreamMsg != "" {
|
|
|
|
|
|
msg = "Access forbidden (403): " + upstreamMsg
|
|
|
|
|
|
}
|
2026-02-28 15:01:20 +08:00
|
|
|
|
logger.LegacyPrintf(
|
|
|
|
|
|
"service.ratelimit",
|
|
|
|
|
|
"[HandleUpstreamErrorRaw] account_id=%d platform=%s type=%s status=403 request_id=%s cf_ray=%s upstream_msg=%s raw_body=%s",
|
|
|
|
|
|
account.ID,
|
|
|
|
|
|
account.Platform,
|
|
|
|
|
|
account.Type,
|
|
|
|
|
|
strings.TrimSpace(headers.Get("x-request-id")),
|
|
|
|
|
|
strings.TrimSpace(headers.Get("cf-ray")),
|
|
|
|
|
|
upstreamMsg,
|
|
|
|
|
|
truncateForLog(responseBody, 1024),
|
|
|
|
|
|
)
|
2026-01-11 15:30:27 +08:00
|
|
|
|
s.handleAuthError(ctx, account, msg)
|
2026-01-03 06:34:00 -08:00
|
|
|
|
shouldDisable = true
|
2025-12-18 13:50:39 +08:00
|
|
|
|
case 429:
|
2026-01-16 13:03:04 +08:00
|
|
|
|
s.handle429(ctx, account, headers, responseBody)
|
2026-01-03 06:34:00 -08:00
|
|
|
|
shouldDisable = false
|
2025-12-18 13:50:39 +08:00
|
|
|
|
case 529:
|
|
|
|
|
|
s.handle529(ctx, account)
|
2026-01-03 06:34:00 -08:00
|
|
|
|
shouldDisable = false
|
2025-12-18 13:50:39 +08:00
|
|
|
|
default:
|
2026-01-11 22:20:02 -08:00
|
|
|
|
// 自定义错误码启用时:在列表中的错误码都应该停止调度
|
|
|
|
|
|
if customErrorCodesEnabled {
|
|
|
|
|
|
msg := "Custom error code triggered"
|
|
|
|
|
|
if upstreamMsg != "" {
|
|
|
|
|
|
msg = upstreamMsg
|
|
|
|
|
|
}
|
|
|
|
|
|
s.handleCustomErrorCode(ctx, account, statusCode, msg)
|
|
|
|
|
|
shouldDisable = true
|
|
|
|
|
|
} else if statusCode >= 500 {
|
|
|
|
|
|
// 未启用自定义错误码时:仅记录5xx错误
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("account_upstream_error", "account_id", account.ID, "status_code", statusCode)
|
2026-01-11 22:20:02 -08:00
|
|
|
|
shouldDisable = false
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
2026-01-03 06:34:00 -08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
return shouldDisable
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-01 04:22:39 +08:00
|
|
|
|
// PreCheckUsage proactively checks local quota before dispatching a request.
|
|
|
|
|
|
// Returns false when the account should be skipped.
|
|
|
|
|
|
func (s *RateLimitService) PreCheckUsage(ctx context.Context, account *Account, requestedModel string) (bool, error) {
|
2026-01-04 15:36:00 +08:00
|
|
|
|
if account == nil || account.Platform != PlatformGemini {
|
2026-01-01 04:22:39 +08:00
|
|
|
|
return true, nil
|
|
|
|
|
|
}
|
2026-01-01 08:29:57 +08:00
|
|
|
|
if s.usageRepo == nil || s.geminiQuotaService == nil {
|
2026-01-01 04:22:39 +08:00
|
|
|
|
return true, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
quota, ok := s.geminiQuotaService.QuotaForAccount(ctx, account)
|
|
|
|
|
|
if !ok {
|
|
|
|
|
|
return true, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
now := time.Now()
|
2026-01-04 15:36:00 +08:00
|
|
|
|
modelClass := geminiModelClassFromName(requestedModel)
|
|
|
|
|
|
|
|
|
|
|
|
// 1) Daily quota precheck (RPD; resets at PST midnight)
|
|
|
|
|
|
{
|
|
|
|
|
|
var limit int64
|
|
|
|
|
|
if quota.SharedRPD > 0 {
|
|
|
|
|
|
limit = quota.SharedRPD
|
|
|
|
|
|
} else {
|
|
|
|
|
|
switch modelClass {
|
|
|
|
|
|
case geminiModelFlash:
|
|
|
|
|
|
limit = quota.FlashRPD
|
|
|
|
|
|
default:
|
|
|
|
|
|
limit = quota.ProRPD
|
|
|
|
|
|
}
|
2026-01-01 08:29:57 +08:00
|
|
|
|
}
|
2026-01-01 04:22:39 +08:00
|
|
|
|
|
2026-01-04 15:36:00 +08:00
|
|
|
|
if limit > 0 {
|
|
|
|
|
|
start := geminiDailyWindowStart(now)
|
|
|
|
|
|
totals, ok := s.getGeminiUsageTotals(account.ID, start, now)
|
|
|
|
|
|
if !ok {
|
2026-02-28 15:01:20 +08:00
|
|
|
|
stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, now, 0, 0, account.ID, 0, nil, nil, nil)
|
2026-01-04 15:36:00 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
|
return true, err
|
|
|
|
|
|
}
|
|
|
|
|
|
totals = geminiAggregateUsage(stats)
|
|
|
|
|
|
s.setGeminiUsageTotals(account.ID, start, now, totals)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
var used int64
|
|
|
|
|
|
if quota.SharedRPD > 0 {
|
|
|
|
|
|
used = totals.ProRequests + totals.FlashRequests
|
|
|
|
|
|
} else {
|
|
|
|
|
|
switch modelClass {
|
|
|
|
|
|
case geminiModelFlash:
|
|
|
|
|
|
used = totals.FlashRequests
|
|
|
|
|
|
default:
|
|
|
|
|
|
used = totals.ProRequests
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if used >= limit {
|
|
|
|
|
|
resetAt := geminiDailyResetTime(now)
|
|
|
|
|
|
// NOTE:
|
|
|
|
|
|
// - This is a local precheck to reduce upstream 429s.
|
|
|
|
|
|
// - Do NOT mark the account as rate-limited here; rate_limit_reset_at should reflect real upstream 429s.
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Info("gemini_precheck_daily_quota_reached", "account_id", account.ID, "used", used, "limit", limit, "reset_at", resetAt)
|
2026-01-04 15:36:00 +08:00
|
|
|
|
return false, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
2026-01-01 04:22:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-04 15:36:00 +08:00
|
|
|
|
// 2) Minute quota precheck (RPM; fixed window current minute)
|
|
|
|
|
|
{
|
|
|
|
|
|
var limit int64
|
|
|
|
|
|
if quota.SharedRPM > 0 {
|
|
|
|
|
|
limit = quota.SharedRPM
|
|
|
|
|
|
} else {
|
|
|
|
|
|
switch modelClass {
|
|
|
|
|
|
case geminiModelFlash:
|
|
|
|
|
|
limit = quota.FlashRPM
|
|
|
|
|
|
default:
|
|
|
|
|
|
limit = quota.ProRPM
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if limit > 0 {
|
|
|
|
|
|
start := now.Truncate(time.Minute)
|
2026-02-28 15:01:20 +08:00
|
|
|
|
stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, now, 0, 0, account.ID, 0, nil, nil, nil)
|
2026-01-04 15:36:00 +08:00
|
|
|
|
if err != nil {
|
|
|
|
|
|
return true, err
|
|
|
|
|
|
}
|
|
|
|
|
|
totals := geminiAggregateUsage(stats)
|
|
|
|
|
|
|
|
|
|
|
|
var used int64
|
|
|
|
|
|
if quota.SharedRPM > 0 {
|
|
|
|
|
|
used = totals.ProRequests + totals.FlashRequests
|
|
|
|
|
|
} else {
|
|
|
|
|
|
switch modelClass {
|
|
|
|
|
|
case geminiModelFlash:
|
|
|
|
|
|
used = totals.FlashRequests
|
|
|
|
|
|
default:
|
|
|
|
|
|
used = totals.ProRequests
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if used >= limit {
|
|
|
|
|
|
resetAt := start.Add(time.Minute)
|
|
|
|
|
|
// Do not persist "rate limited" status from local precheck. See note above.
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Info("gemini_precheck_minute_quota_reached", "account_id", account.ID, "used", used, "limit", limit, "reset_at", resetAt)
|
2026-01-04 15:36:00 +08:00
|
|
|
|
return false, nil
|
|
|
|
|
|
}
|
2026-01-01 04:22:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
return true, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-28 15:01:20 +08:00
|
|
|
|
// PreCheckUsageBatch performs quota precheck for multiple accounts in one request.
|
|
|
|
|
|
// Returned map value=false means the account should be skipped.
|
|
|
|
|
|
func (s *RateLimitService) PreCheckUsageBatch(ctx context.Context, accounts []*Account, requestedModel string) (map[int64]bool, error) {
|
|
|
|
|
|
result := make(map[int64]bool, len(accounts))
|
|
|
|
|
|
for _, account := range accounts {
|
|
|
|
|
|
if account == nil {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
result[account.ID] = true
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if len(accounts) == 0 || requestedModel == "" {
|
|
|
|
|
|
return result, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
if s.usageRepo == nil || s.geminiQuotaService == nil {
|
|
|
|
|
|
return result, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
modelClass := geminiModelClassFromName(requestedModel)
|
|
|
|
|
|
now := time.Now()
|
|
|
|
|
|
dailyStart := geminiDailyWindowStart(now)
|
|
|
|
|
|
minuteStart := now.Truncate(time.Minute)
|
|
|
|
|
|
|
|
|
|
|
|
type quotaAccount struct {
|
|
|
|
|
|
account *Account
|
|
|
|
|
|
quota GeminiQuota
|
|
|
|
|
|
}
|
|
|
|
|
|
quotaAccounts := make([]quotaAccount, 0, len(accounts))
|
|
|
|
|
|
for _, account := range accounts {
|
|
|
|
|
|
if account == nil || account.Platform != PlatformGemini {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
quota, ok := s.geminiQuotaService.QuotaForAccount(ctx, account)
|
|
|
|
|
|
if !ok {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
quotaAccounts = append(quotaAccounts, quotaAccount{
|
|
|
|
|
|
account: account,
|
|
|
|
|
|
quota: quota,
|
|
|
|
|
|
})
|
|
|
|
|
|
}
|
|
|
|
|
|
if len(quotaAccounts) == 0 {
|
|
|
|
|
|
return result, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 1) Daily precheck (cached + batch DB fallback)
|
|
|
|
|
|
dailyTotalsByID := make(map[int64]GeminiUsageTotals, len(quotaAccounts))
|
|
|
|
|
|
dailyMissIDs := make([]int64, 0, len(quotaAccounts))
|
|
|
|
|
|
for _, item := range quotaAccounts {
|
|
|
|
|
|
limit := geminiDailyLimit(item.quota, modelClass)
|
|
|
|
|
|
if limit <= 0 {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
accountID := item.account.ID
|
|
|
|
|
|
if totals, ok := s.getGeminiUsageTotals(accountID, dailyStart, now); ok {
|
|
|
|
|
|
dailyTotalsByID[accountID] = totals
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
dailyMissIDs = append(dailyMissIDs, accountID)
|
|
|
|
|
|
}
|
|
|
|
|
|
if len(dailyMissIDs) > 0 {
|
|
|
|
|
|
totalsBatch, err := s.getGeminiUsageTotalsBatch(ctx, dailyMissIDs, dailyStart, now)
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
return result, err
|
|
|
|
|
|
}
|
|
|
|
|
|
for _, accountID := range dailyMissIDs {
|
|
|
|
|
|
totals := totalsBatch[accountID]
|
|
|
|
|
|
dailyTotalsByID[accountID] = totals
|
|
|
|
|
|
s.setGeminiUsageTotals(accountID, dailyStart, now, totals)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
for _, item := range quotaAccounts {
|
|
|
|
|
|
limit := geminiDailyLimit(item.quota, modelClass)
|
|
|
|
|
|
if limit <= 0 {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
accountID := item.account.ID
|
|
|
|
|
|
used := geminiUsedRequests(item.quota, modelClass, dailyTotalsByID[accountID], true)
|
|
|
|
|
|
if used >= limit {
|
|
|
|
|
|
resetAt := geminiDailyResetTime(now)
|
|
|
|
|
|
slog.Info("gemini_precheck_daily_quota_reached_batch", "account_id", accountID, "used", used, "limit", limit, "reset_at", resetAt)
|
|
|
|
|
|
result[accountID] = false
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 2) Minute precheck (batch DB)
|
|
|
|
|
|
minuteIDs := make([]int64, 0, len(quotaAccounts))
|
|
|
|
|
|
for _, item := range quotaAccounts {
|
|
|
|
|
|
accountID := item.account.ID
|
|
|
|
|
|
if !result[accountID] {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
if geminiMinuteLimit(item.quota, modelClass) <= 0 {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
minuteIDs = append(minuteIDs, accountID)
|
|
|
|
|
|
}
|
|
|
|
|
|
if len(minuteIDs) == 0 {
|
|
|
|
|
|
return result, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
minuteTotalsByID, err := s.getGeminiUsageTotalsBatch(ctx, minuteIDs, minuteStart, now)
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
return result, err
|
|
|
|
|
|
}
|
|
|
|
|
|
for _, item := range quotaAccounts {
|
|
|
|
|
|
accountID := item.account.ID
|
|
|
|
|
|
if !result[accountID] {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
limit := geminiMinuteLimit(item.quota, modelClass)
|
|
|
|
|
|
if limit <= 0 {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
used := geminiUsedRequests(item.quota, modelClass, minuteTotalsByID[accountID], false)
|
|
|
|
|
|
if used >= limit {
|
|
|
|
|
|
resetAt := minuteStart.Add(time.Minute)
|
|
|
|
|
|
slog.Info("gemini_precheck_minute_quota_reached_batch", "account_id", accountID, "used", used, "limit", limit, "reset_at", resetAt)
|
|
|
|
|
|
result[accountID] = false
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
return result, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func (s *RateLimitService) getGeminiUsageTotalsBatch(ctx context.Context, accountIDs []int64, start, end time.Time) (map[int64]GeminiUsageTotals, error) {
|
|
|
|
|
|
result := make(map[int64]GeminiUsageTotals, len(accountIDs))
|
|
|
|
|
|
if len(accountIDs) == 0 {
|
|
|
|
|
|
return result, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
ids := make([]int64, 0, len(accountIDs))
|
|
|
|
|
|
seen := make(map[int64]struct{}, len(accountIDs))
|
|
|
|
|
|
for _, accountID := range accountIDs {
|
|
|
|
|
|
if accountID <= 0 {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
if _, ok := seen[accountID]; ok {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
seen[accountID] = struct{}{}
|
|
|
|
|
|
ids = append(ids, accountID)
|
|
|
|
|
|
}
|
|
|
|
|
|
if len(ids) == 0 {
|
|
|
|
|
|
return result, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if batchReader, ok := s.usageRepo.(geminiUsageTotalsBatchProvider); ok {
|
|
|
|
|
|
stats, err := batchReader.GetGeminiUsageTotalsBatch(ctx, ids, start, end)
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
return nil, err
|
|
|
|
|
|
}
|
|
|
|
|
|
for _, accountID := range ids {
|
|
|
|
|
|
result[accountID] = stats[accountID]
|
|
|
|
|
|
}
|
|
|
|
|
|
return result, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
for _, accountID := range ids {
|
|
|
|
|
|
stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, start, end, 0, 0, accountID, 0, nil, nil, nil)
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
return nil, err
|
|
|
|
|
|
}
|
|
|
|
|
|
result[accountID] = geminiAggregateUsage(stats)
|
|
|
|
|
|
}
|
|
|
|
|
|
return result, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func geminiDailyLimit(quota GeminiQuota, modelClass geminiModelClass) int64 {
|
|
|
|
|
|
if quota.SharedRPD > 0 {
|
|
|
|
|
|
return quota.SharedRPD
|
|
|
|
|
|
}
|
|
|
|
|
|
switch modelClass {
|
|
|
|
|
|
case geminiModelFlash:
|
|
|
|
|
|
return quota.FlashRPD
|
|
|
|
|
|
default:
|
|
|
|
|
|
return quota.ProRPD
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func geminiMinuteLimit(quota GeminiQuota, modelClass geminiModelClass) int64 {
|
|
|
|
|
|
if quota.SharedRPM > 0 {
|
|
|
|
|
|
return quota.SharedRPM
|
|
|
|
|
|
}
|
|
|
|
|
|
switch modelClass {
|
|
|
|
|
|
case geminiModelFlash:
|
|
|
|
|
|
return quota.FlashRPM
|
|
|
|
|
|
default:
|
|
|
|
|
|
return quota.ProRPM
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func geminiUsedRequests(quota GeminiQuota, modelClass geminiModelClass, totals GeminiUsageTotals, daily bool) int64 {
|
|
|
|
|
|
if daily {
|
|
|
|
|
|
if quota.SharedRPD > 0 {
|
|
|
|
|
|
return totals.ProRequests + totals.FlashRequests
|
|
|
|
|
|
}
|
|
|
|
|
|
} else {
|
|
|
|
|
|
if quota.SharedRPM > 0 {
|
|
|
|
|
|
return totals.ProRequests + totals.FlashRequests
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
switch modelClass {
|
|
|
|
|
|
case geminiModelFlash:
|
|
|
|
|
|
return totals.FlashRequests
|
|
|
|
|
|
default:
|
|
|
|
|
|
return totals.ProRequests
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-01 08:29:57 +08:00
|
|
|
|
func (s *RateLimitService) getGeminiUsageTotals(accountID int64, windowStart, now time.Time) (GeminiUsageTotals, bool) {
|
2026-01-01 15:35:08 +08:00
|
|
|
|
s.usageCacheMu.RLock()
|
|
|
|
|
|
defer s.usageCacheMu.RUnlock()
|
2026-01-01 08:29:57 +08:00
|
|
|
|
|
|
|
|
|
|
if s.usageCache == nil {
|
|
|
|
|
|
return GeminiUsageTotals{}, false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
entry, ok := s.usageCache[accountID]
|
|
|
|
|
|
if !ok || entry == nil {
|
|
|
|
|
|
return GeminiUsageTotals{}, false
|
|
|
|
|
|
}
|
|
|
|
|
|
if !entry.windowStart.Equal(windowStart) {
|
|
|
|
|
|
return GeminiUsageTotals{}, false
|
|
|
|
|
|
}
|
|
|
|
|
|
if now.Sub(entry.cachedAt) >= geminiPrecheckCacheTTL {
|
|
|
|
|
|
return GeminiUsageTotals{}, false
|
|
|
|
|
|
}
|
|
|
|
|
|
return entry.totals, true
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func (s *RateLimitService) setGeminiUsageTotals(accountID int64, windowStart, now time.Time, totals GeminiUsageTotals) {
|
|
|
|
|
|
s.usageCacheMu.Lock()
|
|
|
|
|
|
defer s.usageCacheMu.Unlock()
|
|
|
|
|
|
if s.usageCache == nil {
|
|
|
|
|
|
s.usageCache = make(map[int64]*geminiUsageCacheEntry)
|
|
|
|
|
|
}
|
|
|
|
|
|
s.usageCache[accountID] = &geminiUsageCacheEntry{
|
|
|
|
|
|
windowStart: windowStart,
|
|
|
|
|
|
cachedAt: now,
|
|
|
|
|
|
totals: totals,
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-01 04:22:39 +08:00
|
|
|
|
// GeminiCooldown returns the fallback cooldown duration for Gemini 429s based on tier.
|
|
|
|
|
|
func (s *RateLimitService) GeminiCooldown(ctx context.Context, account *Account) time.Duration {
|
|
|
|
|
|
if account == nil {
|
|
|
|
|
|
return 5 * time.Minute
|
|
|
|
|
|
}
|
2026-01-04 15:36:00 +08:00
|
|
|
|
if s.geminiQuotaService == nil {
|
|
|
|
|
|
return 5 * time.Minute
|
|
|
|
|
|
}
|
|
|
|
|
|
return s.geminiQuotaService.CooldownForAccount(ctx, account)
|
2026-01-01 04:22:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// handleAuthError 处理认证类错误(401/403),停止账号调度
|
2025-12-26 15:40:24 +08:00
|
|
|
|
func (s *RateLimitService) handleAuthError(ctx context.Context, account *Account, errorMsg string) {
|
2025-12-19 21:26:19 +08:00
|
|
|
|
if err := s.accountRepo.SetError(ctx, account.ID, errorMsg); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("account_set_error_failed", "account_id", account.ID, "error", err)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("account_disabled_auth_error", "account_id", account.ID, "error", errorMsg)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-11 22:20:02 -08:00
|
|
|
|
// handleCustomErrorCode 处理自定义错误码,停止账号调度
|
|
|
|
|
|
func (s *RateLimitService) handleCustomErrorCode(ctx context.Context, account *Account, statusCode int, errorMsg string) {
|
|
|
|
|
|
msg := "Custom error code " + strconv.Itoa(statusCode) + ": " + errorMsg
|
|
|
|
|
|
if err := s.accountRepo.SetError(ctx, account.ID, msg); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("account_set_error_failed", "account_id", account.ID, "status_code", statusCode, "error", err)
|
2026-01-11 22:20:02 -08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("account_disabled_custom_error", "account_id", account.ID, "status_code", statusCode, "error", errorMsg)
|
2026-01-11 22:20:02 -08:00
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// handle429 处理429限流错误
|
|
|
|
|
|
// 解析响应头获取重置时间,标记账号为限流状态
|
2026-01-16 13:03:04 +08:00
|
|
|
|
func (s *RateLimitService) handle429(ctx context.Context, account *Account, headers http.Header, responseBody []byte) {
|
2026-01-25 00:50:42 +08:00
|
|
|
|
// 1. OpenAI 平台:优先尝试解析 x-codex-* 响应头(用于 rate_limit_exceeded)
|
2026-01-25 13:32:08 +08:00
|
|
|
|
if account.Platform == PlatformOpenAI {
|
2026-03-07 20:02:58 +08:00
|
|
|
|
s.persistOpenAICodexSnapshot(ctx, account, headers)
|
2026-01-25 13:32:08 +08:00
|
|
|
|
if resetAt := s.calculateOpenAI429ResetTime(headers); resetAt != nil {
|
|
|
|
|
|
if err := s.accountRepo.SetRateLimited(ctx, account.ID, *resetAt); err != nil {
|
|
|
|
|
|
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
slog.Info("openai_account_rate_limited", "account_id", account.ID, "reset_at", *resetAt)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-14 00:21:56 +08:00
|
|
|
|
// 2. Anthropic 平台:尝试解析 per-window 头(5h / 7d),选择实际触发的窗口
|
|
|
|
|
|
if result := calculateAnthropic429ResetTime(headers); result != nil {
|
|
|
|
|
|
if err := s.accountRepo.SetRateLimited(ctx, account.ID, result.resetAt); err != nil {
|
|
|
|
|
|
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 更新 session window:优先使用 5h-reset 头精确计算,否则从 resetAt 反推
|
|
|
|
|
|
windowEnd := result.resetAt
|
|
|
|
|
|
if result.fiveHourReset != nil {
|
|
|
|
|
|
windowEnd = *result.fiveHourReset
|
|
|
|
|
|
}
|
|
|
|
|
|
windowStart := windowEnd.Add(-5 * time.Hour)
|
|
|
|
|
|
if err := s.accountRepo.UpdateSessionWindow(ctx, account.ID, &windowStart, &windowEnd, "rejected"); err != nil {
|
|
|
|
|
|
slog.Warn("rate_limit_update_session_window_failed", "account_id", account.ID, "error", err)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
slog.Info("anthropic_account_rate_limited", "account_id", account.ID, "reset_at", result.resetAt, "reset_in", time.Until(result.resetAt).Truncate(time.Second))
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 3. 尝试从响应头解析重置时间(Anthropic 聚合头,向后兼容)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
resetTimestamp := headers.Get("anthropic-ratelimit-unified-reset")
|
2026-01-25 00:50:42 +08:00
|
|
|
|
|
2026-02-14 00:21:56 +08:00
|
|
|
|
// 4. 如果响应头没有,尝试从响应体解析(OpenAI usage_limit_reached, Gemini)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
if resetTimestamp == "" {
|
2026-01-25 00:50:42 +08:00
|
|
|
|
switch account.Platform {
|
|
|
|
|
|
case PlatformOpenAI:
|
|
|
|
|
|
// 尝试解析 OpenAI 的 usage_limit_reached 错误
|
|
|
|
|
|
if resetAt := parseOpenAIRateLimitResetTime(responseBody); resetAt != nil {
|
|
|
|
|
|
resetTime := time.Unix(*resetAt, 0)
|
|
|
|
|
|
if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetTime); err != nil {
|
|
|
|
|
|
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
slog.Info("account_rate_limited", "account_id", account.ID, "platform", account.Platform, "reset_at", resetTime, "reset_in", time.Until(resetTime).Truncate(time.Second))
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
case PlatformGemini, PlatformAntigravity:
|
|
|
|
|
|
// 尝试解析 Gemini 格式(用于其他平台)
|
|
|
|
|
|
if resetAt := ParseGeminiRateLimitResetTime(responseBody); resetAt != nil {
|
|
|
|
|
|
resetTime := time.Unix(*resetAt, 0)
|
|
|
|
|
|
if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetTime); err != nil {
|
|
|
|
|
|
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
slog.Info("account_rate_limited", "account_id", account.ID, "platform", account.Platform, "reset_at", resetTime, "reset_in", time.Until(resetTime).Truncate(time.Second))
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-03-04 09:36:24 +08:00
|
|
|
|
// Anthropic 平台:没有限流重置时间的 429 可能是非真实限流(如 Extra usage required),
|
|
|
|
|
|
// 不标记账号限流状态,直接透传错误给客户端
|
|
|
|
|
|
if account.Platform == PlatformAnthropic {
|
|
|
|
|
|
slog.Warn("rate_limit_429_no_reset_time_skipped",
|
|
|
|
|
|
"account_id", account.ID,
|
|
|
|
|
|
"platform", account.Platform,
|
|
|
|
|
|
"reason", "no rate limit reset time in headers, likely not a real rate limit")
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 其他平台:没有重置时间,使用默认5分钟
|
2025-12-18 13:50:39 +08:00
|
|
|
|
resetAt := time.Now().Add(5 * time.Minute)
|
2026-01-25 00:50:42 +08:00
|
|
|
|
slog.Warn("rate_limit_no_reset_time", "account_id", account.ID, "platform", account.Platform, "using_default", "5m")
|
2025-12-19 21:26:19 +08:00
|
|
|
|
if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetAt); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 解析Unix时间戳
|
|
|
|
|
|
ts, err := strconv.ParseInt(resetTimestamp, 10, 64)
|
|
|
|
|
|
if err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("rate_limit_reset_parse_failed", "reset_timestamp", resetTimestamp, "error", err)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
resetAt := time.Now().Add(5 * time.Minute)
|
2025-12-19 21:26:19 +08:00
|
|
|
|
if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetAt); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
resetAt := time.Unix(ts, 0)
|
|
|
|
|
|
|
|
|
|
|
|
// 标记限流状态
|
2025-12-19 21:26:19 +08:00
|
|
|
|
if err := s.accountRepo.SetRateLimited(ctx, account.ID, resetAt); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("rate_limit_set_failed", "account_id", account.ID, "error", err)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 根据重置时间反推5h窗口
|
|
|
|
|
|
windowEnd := resetAt
|
|
|
|
|
|
windowStart := resetAt.Add(-5 * time.Hour)
|
2025-12-19 21:26:19 +08:00
|
|
|
|
if err := s.accountRepo.UpdateSessionWindow(ctx, account.ID, &windowStart, &windowEnd, "rejected"); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("rate_limit_update_session_window_failed", "account_id", account.ID, "error", err)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Info("account_rate_limited", "account_id", account.ID, "reset_at", resetAt)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-25 13:32:08 +08:00
|
|
|
|
// calculateOpenAI429ResetTime 从 OpenAI 429 响应头计算正确的重置时间
|
|
|
|
|
|
// 返回 nil 表示无法从响应头中确定重置时间
|
|
|
|
|
|
func (s *RateLimitService) calculateOpenAI429ResetTime(headers http.Header) *time.Time {
|
|
|
|
|
|
snapshot := ParseCodexRateLimitHeaders(headers)
|
|
|
|
|
|
if snapshot == nil {
|
|
|
|
|
|
return nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
normalized := snapshot.Normalize()
|
|
|
|
|
|
if normalized == nil {
|
|
|
|
|
|
return nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
now := time.Now()
|
|
|
|
|
|
|
|
|
|
|
|
// 判断哪个限制被触发(used_percent >= 100)
|
|
|
|
|
|
is7dExhausted := normalized.Used7dPercent != nil && *normalized.Used7dPercent >= 100
|
|
|
|
|
|
is5hExhausted := normalized.Used5hPercent != nil && *normalized.Used5hPercent >= 100
|
|
|
|
|
|
|
|
|
|
|
|
// 优先使用被触发限制的重置时间
|
|
|
|
|
|
if is7dExhausted && normalized.Reset7dSeconds != nil {
|
|
|
|
|
|
resetAt := now.Add(time.Duration(*normalized.Reset7dSeconds) * time.Second)
|
|
|
|
|
|
slog.Info("openai_429_7d_limit_exhausted", "reset_after_seconds", *normalized.Reset7dSeconds, "reset_at", resetAt)
|
|
|
|
|
|
return &resetAt
|
|
|
|
|
|
}
|
|
|
|
|
|
if is5hExhausted && normalized.Reset5hSeconds != nil {
|
|
|
|
|
|
resetAt := now.Add(time.Duration(*normalized.Reset5hSeconds) * time.Second)
|
|
|
|
|
|
slog.Info("openai_429_5h_limit_exhausted", "reset_after_seconds", *normalized.Reset5hSeconds, "reset_at", resetAt)
|
|
|
|
|
|
return &resetAt
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 都未达到100%但收到429,使用较长的重置时间
|
|
|
|
|
|
var maxResetSecs int
|
|
|
|
|
|
if normalized.Reset7dSeconds != nil && *normalized.Reset7dSeconds > maxResetSecs {
|
|
|
|
|
|
maxResetSecs = *normalized.Reset7dSeconds
|
|
|
|
|
|
}
|
|
|
|
|
|
if normalized.Reset5hSeconds != nil && *normalized.Reset5hSeconds > maxResetSecs {
|
|
|
|
|
|
maxResetSecs = *normalized.Reset5hSeconds
|
|
|
|
|
|
}
|
|
|
|
|
|
if maxResetSecs > 0 {
|
|
|
|
|
|
resetAt := now.Add(time.Duration(maxResetSecs) * time.Second)
|
|
|
|
|
|
slog.Info("openai_429_using_max_reset", "max_reset_seconds", maxResetSecs, "reset_at", resetAt)
|
|
|
|
|
|
return &resetAt
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
return nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-02-14 00:21:56 +08:00
|
|
|
|
// anthropic429Result holds the parsed Anthropic 429 rate-limit information.
|
|
|
|
|
|
type anthropic429Result struct {
|
|
|
|
|
|
resetAt time.Time // The correct reset time to use for SetRateLimited
|
|
|
|
|
|
fiveHourReset *time.Time // 5h window reset timestamp (for session window calculation), nil if not available
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// calculateAnthropic429ResetTime parses Anthropic's per-window rate-limit headers
|
|
|
|
|
|
// to determine which window (5h or 7d) actually triggered the 429.
|
|
|
|
|
|
//
|
|
|
|
|
|
// Headers used:
|
|
|
|
|
|
// - anthropic-ratelimit-unified-5h-utilization / anthropic-ratelimit-unified-5h-surpassed-threshold
|
|
|
|
|
|
// - anthropic-ratelimit-unified-5h-reset
|
|
|
|
|
|
// - anthropic-ratelimit-unified-7d-utilization / anthropic-ratelimit-unified-7d-surpassed-threshold
|
|
|
|
|
|
// - anthropic-ratelimit-unified-7d-reset
|
|
|
|
|
|
//
|
|
|
|
|
|
// Returns nil when the per-window headers are absent (caller should fall back to
|
|
|
|
|
|
// the aggregated anthropic-ratelimit-unified-reset header).
|
|
|
|
|
|
func calculateAnthropic429ResetTime(headers http.Header) *anthropic429Result {
|
|
|
|
|
|
reset5hStr := headers.Get("anthropic-ratelimit-unified-5h-reset")
|
|
|
|
|
|
reset7dStr := headers.Get("anthropic-ratelimit-unified-7d-reset")
|
|
|
|
|
|
|
|
|
|
|
|
if reset5hStr == "" && reset7dStr == "" {
|
|
|
|
|
|
return nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
var reset5h, reset7d *time.Time
|
|
|
|
|
|
if ts, err := strconv.ParseInt(reset5hStr, 10, 64); err == nil {
|
|
|
|
|
|
t := time.Unix(ts, 0)
|
|
|
|
|
|
reset5h = &t
|
|
|
|
|
|
}
|
|
|
|
|
|
if ts, err := strconv.ParseInt(reset7dStr, 10, 64); err == nil {
|
|
|
|
|
|
t := time.Unix(ts, 0)
|
|
|
|
|
|
reset7d = &t
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
is5hExceeded := isAnthropicWindowExceeded(headers, "5h")
|
|
|
|
|
|
is7dExceeded := isAnthropicWindowExceeded(headers, "7d")
|
|
|
|
|
|
|
|
|
|
|
|
slog.Info("anthropic_429_window_analysis",
|
|
|
|
|
|
"is_5h_exceeded", is5hExceeded,
|
|
|
|
|
|
"is_7d_exceeded", is7dExceeded,
|
|
|
|
|
|
"reset_5h", reset5hStr,
|
|
|
|
|
|
"reset_7d", reset7dStr,
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
// Select the correct reset time based on which window(s) are exceeded.
|
|
|
|
|
|
var chosen *time.Time
|
|
|
|
|
|
switch {
|
|
|
|
|
|
case is5hExceeded && is7dExceeded:
|
|
|
|
|
|
// Both exceeded → prefer 7d (longer cooldown), fall back to 5h
|
|
|
|
|
|
chosen = reset7d
|
|
|
|
|
|
if chosen == nil {
|
|
|
|
|
|
chosen = reset5h
|
|
|
|
|
|
}
|
|
|
|
|
|
case is5hExceeded:
|
|
|
|
|
|
chosen = reset5h
|
|
|
|
|
|
case is7dExceeded:
|
|
|
|
|
|
chosen = reset7d
|
|
|
|
|
|
default:
|
|
|
|
|
|
// Neither flag clearly exceeded — pick the sooner reset as best guess
|
|
|
|
|
|
chosen = pickSooner(reset5h, reset7d)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if chosen == nil {
|
|
|
|
|
|
return nil
|
|
|
|
|
|
}
|
|
|
|
|
|
return &anthropic429Result{resetAt: *chosen, fiveHourReset: reset5h}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// isAnthropicWindowExceeded checks whether a given Anthropic rate-limit window
|
|
|
|
|
|
// (e.g. "5h" or "7d") has been exceeded, using utilization and surpassed-threshold headers.
|
|
|
|
|
|
func isAnthropicWindowExceeded(headers http.Header, window string) bool {
|
|
|
|
|
|
prefix := "anthropic-ratelimit-unified-" + window + "-"
|
|
|
|
|
|
|
|
|
|
|
|
// Check surpassed-threshold first (most explicit signal)
|
|
|
|
|
|
if st := headers.Get(prefix + "surpassed-threshold"); strings.EqualFold(st, "true") {
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// Fall back to utilization >= 1.0
|
|
|
|
|
|
if utilStr := headers.Get(prefix + "utilization"); utilStr != "" {
|
|
|
|
|
|
if util, err := strconv.ParseFloat(utilStr, 64); err == nil && util >= 1.0-1e-9 {
|
|
|
|
|
|
// Use a small epsilon to handle floating point: treat 0.9999999... as >= 1.0
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// pickSooner returns whichever of the two time pointers is earlier.
|
|
|
|
|
|
// If only one is non-nil, it is returned. If both are nil, returns nil.
|
|
|
|
|
|
func pickSooner(a, b *time.Time) *time.Time {
|
|
|
|
|
|
switch {
|
|
|
|
|
|
case a != nil && b != nil:
|
|
|
|
|
|
if a.Before(*b) {
|
|
|
|
|
|
return a
|
|
|
|
|
|
}
|
|
|
|
|
|
return b
|
|
|
|
|
|
case a != nil:
|
|
|
|
|
|
return a
|
|
|
|
|
|
default:
|
|
|
|
|
|
return b
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-03-07 20:02:58 +08:00
|
|
|
|
func (s *RateLimitService) persistOpenAICodexSnapshot(ctx context.Context, account *Account, headers http.Header) {
|
|
|
|
|
|
if s == nil || s.accountRepo == nil || account == nil || headers == nil {
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
snapshot := ParseCodexRateLimitHeaders(headers)
|
|
|
|
|
|
if snapshot == nil {
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
updates := buildCodexUsageExtraUpdates(snapshot, time.Now())
|
|
|
|
|
|
if len(updates) == 0 {
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
if err := s.accountRepo.UpdateExtra(ctx, account.ID, updates); err != nil {
|
|
|
|
|
|
slog.Warn("openai_codex_snapshot_persist_failed", "account_id", account.ID, "error", err)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-25 00:50:42 +08:00
|
|
|
|
// parseOpenAIRateLimitResetTime 解析 OpenAI 格式的 429 响应,返回重置时间的 Unix 时间戳
|
|
|
|
|
|
// OpenAI 的 usage_limit_reached 错误格式:
|
|
|
|
|
|
//
|
|
|
|
|
|
// {
|
|
|
|
|
|
// "error": {
|
|
|
|
|
|
// "message": "The usage limit has been reached",
|
|
|
|
|
|
// "type": "usage_limit_reached",
|
|
|
|
|
|
// "resets_at": 1769404154,
|
|
|
|
|
|
// "resets_in_seconds": 133107
|
|
|
|
|
|
// }
|
|
|
|
|
|
// }
|
|
|
|
|
|
func parseOpenAIRateLimitResetTime(body []byte) *int64 {
|
|
|
|
|
|
var parsed map[string]any
|
|
|
|
|
|
if err := json.Unmarshal(body, &parsed); err != nil {
|
|
|
|
|
|
return nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
errObj, ok := parsed["error"].(map[string]any)
|
|
|
|
|
|
if !ok {
|
|
|
|
|
|
return nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 检查是否为 usage_limit_reached 或 rate_limit_exceeded 类型
|
|
|
|
|
|
errType, _ := errObj["type"].(string)
|
|
|
|
|
|
if errType != "usage_limit_reached" && errType != "rate_limit_exceeded" {
|
|
|
|
|
|
return nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 优先使用 resets_at(Unix 时间戳)
|
|
|
|
|
|
if resetsAt, ok := errObj["resets_at"].(float64); ok {
|
|
|
|
|
|
ts := int64(resetsAt)
|
|
|
|
|
|
return &ts
|
|
|
|
|
|
}
|
|
|
|
|
|
if resetsAt, ok := errObj["resets_at"].(string); ok {
|
|
|
|
|
|
if ts, err := strconv.ParseInt(resetsAt, 10, 64); err == nil {
|
|
|
|
|
|
return &ts
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 如果没有 resets_at,尝试使用 resets_in_seconds
|
|
|
|
|
|
if resetsInSeconds, ok := errObj["resets_in_seconds"].(float64); ok {
|
|
|
|
|
|
ts := time.Now().Unix() + int64(resetsInSeconds)
|
|
|
|
|
|
return &ts
|
|
|
|
|
|
}
|
|
|
|
|
|
if resetsInSeconds, ok := errObj["resets_in_seconds"].(string); ok {
|
|
|
|
|
|
if sec, err := strconv.ParseInt(resetsInSeconds, 10, 64); err == nil {
|
|
|
|
|
|
ts := time.Now().Unix() + sec
|
|
|
|
|
|
return &ts
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
return nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// handle529 处理529过载错误
|
|
|
|
|
|
// 根据配置设置过载冷却时间
|
2025-12-26 15:40:24 +08:00
|
|
|
|
func (s *RateLimitService) handle529(ctx context.Context, account *Account) {
|
2025-12-18 13:50:39 +08:00
|
|
|
|
cooldownMinutes := s.cfg.RateLimit.OverloadCooldownMinutes
|
|
|
|
|
|
if cooldownMinutes <= 0 {
|
|
|
|
|
|
cooldownMinutes = 10 // 默认10分钟
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
until := time.Now().Add(time.Duration(cooldownMinutes) * time.Minute)
|
2025-12-19 21:26:19 +08:00
|
|
|
|
if err := s.accountRepo.SetOverloaded(ctx, account.ID, until); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("overload_set_failed", "account_id", account.ID, "error", err)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Info("account_overloaded", "account_id", account.ID, "until", until)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// UpdateSessionWindow 从成功响应更新5h窗口状态
|
2025-12-26 15:40:24 +08:00
|
|
|
|
func (s *RateLimitService) UpdateSessionWindow(ctx context.Context, account *Account, headers http.Header) {
|
2025-12-18 13:50:39 +08:00
|
|
|
|
status := headers.Get("anthropic-ratelimit-unified-5h-status")
|
|
|
|
|
|
if status == "" {
|
|
|
|
|
|
return
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 检查是否需要初始化时间窗口
|
|
|
|
|
|
// 对于 Setup Token 账号,首次成功请求时需要预测时间窗口
|
|
|
|
|
|
var windowStart, windowEnd *time.Time
|
|
|
|
|
|
needInitWindow := account.SessionWindowEnd == nil || time.Now().After(*account.SessionWindowEnd)
|
|
|
|
|
|
|
|
|
|
|
|
if needInitWindow && (status == "allowed" || status == "allowed_warning") {
|
|
|
|
|
|
// 预测时间窗口:从当前时间的整点开始,+5小时为结束
|
|
|
|
|
|
// 例如:现在是 14:30,窗口为 14:00 ~ 19:00
|
|
|
|
|
|
now := time.Now()
|
|
|
|
|
|
start := time.Date(now.Year(), now.Month(), now.Day(), now.Hour(), 0, 0, 0, now.Location())
|
|
|
|
|
|
end := start.Add(5 * time.Hour)
|
|
|
|
|
|
windowStart = &start
|
|
|
|
|
|
windowEnd = &end
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Info("account_session_window_initialized", "account_id", account.ID, "window_start", start, "window_end", end, "status", status)
|
2026-03-06 17:53:04 +08:00
|
|
|
|
// 窗口重置时清除旧的 utilization,避免残留上个窗口的数据
|
|
|
|
|
|
_ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{
|
|
|
|
|
|
"session_window_utilization": nil,
|
|
|
|
|
|
})
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-19 21:26:19 +08:00
|
|
|
|
if err := s.accountRepo.UpdateSessionWindow(ctx, account.ID, windowStart, windowEnd, status); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("session_window_update_failed", "account_id", account.ID, "error", err)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
|
2026-03-06 17:53:04 +08:00
|
|
|
|
// 存储真实的 utilization 值(0-1 小数),供 estimateSetupTokenUsage 使用
|
|
|
|
|
|
if utilStr := headers.Get("anthropic-ratelimit-unified-5h-utilization"); utilStr != "" {
|
|
|
|
|
|
if util, err := strconv.ParseFloat(utilStr, 64); err == nil {
|
|
|
|
|
|
if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{
|
|
|
|
|
|
"session_window_utilization": util,
|
|
|
|
|
|
}); err != nil {
|
|
|
|
|
|
slog.Warn("session_window_utilization_update_failed", "account_id", account.ID, "error", err)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2025-12-18 13:50:39 +08:00
|
|
|
|
// 如果状态为allowed且之前有限流,说明窗口已重置,清除限流状态
|
|
|
|
|
|
if status == "allowed" && account.IsRateLimited() {
|
2026-01-09 22:00:14 +08:00
|
|
|
|
if err := s.ClearRateLimit(ctx, account.ID); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("rate_limit_clear_failed", "account_id", account.ID, "error", err)
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// ClearRateLimit 清除账号的限流状态
|
|
|
|
|
|
func (s *RateLimitService) ClearRateLimit(ctx context.Context, accountID int64) error {
|
2026-01-09 22:00:14 +08:00
|
|
|
|
if err := s.accountRepo.ClearRateLimit(ctx, accountID); err != nil {
|
|
|
|
|
|
return err
|
|
|
|
|
|
}
|
2026-01-16 13:03:04 +08:00
|
|
|
|
if err := s.accountRepo.ClearAntigravityQuotaScopes(ctx, accountID); err != nil {
|
|
|
|
|
|
return err
|
|
|
|
|
|
}
|
2026-02-22 17:00:29 +08:00
|
|
|
|
if err := s.accountRepo.ClearModelRateLimits(ctx, accountID); err != nil {
|
|
|
|
|
|
return err
|
|
|
|
|
|
}
|
|
|
|
|
|
// 清除限流时一并清理临时不可调度状态,避免周限/窗口重置后仍被本地临时状态阻断。
|
|
|
|
|
|
if err := s.accountRepo.ClearTempUnschedulable(ctx, accountID); err != nil {
|
|
|
|
|
|
return err
|
|
|
|
|
|
}
|
|
|
|
|
|
if s.tempUnschedCache != nil {
|
|
|
|
|
|
if err := s.tempUnschedCache.DeleteTempUnsched(ctx, accountID); err != nil {
|
|
|
|
|
|
slog.Warn("temp_unsched_cache_delete_failed", "account_id", accountID, "error", err)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
return nil
|
2025-12-18 13:50:39 +08:00
|
|
|
|
}
|
2026-01-03 06:34:00 -08:00
|
|
|
|
|
2026-03-08 06:59:53 +08:00
|
|
|
|
// RecoverAccountState 按需恢复账号的可恢复运行时状态。
|
|
|
|
|
|
func (s *RateLimitService) RecoverAccountState(ctx context.Context, accountID int64, options AccountRecoveryOptions) (*SuccessfulTestRecoveryResult, error) {
|
|
|
|
|
|
account, err := s.accountRepo.GetByID(ctx, accountID)
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
return nil, err
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
result := &SuccessfulTestRecoveryResult{}
|
|
|
|
|
|
if account.Status == StatusError {
|
|
|
|
|
|
if err := s.accountRepo.ClearError(ctx, accountID); err != nil {
|
|
|
|
|
|
return nil, err
|
|
|
|
|
|
}
|
|
|
|
|
|
result.ClearedError = true
|
|
|
|
|
|
if options.InvalidateToken && s.tokenCacheInvalidator != nil && account.IsOAuth() {
|
|
|
|
|
|
if invalidateErr := s.tokenCacheInvalidator.InvalidateToken(ctx, account); invalidateErr != nil {
|
|
|
|
|
|
slog.Warn("recover_account_state_invalidate_token_failed", "account_id", accountID, "error", invalidateErr)
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if hasRecoverableRuntimeState(account) {
|
|
|
|
|
|
if err := s.ClearRateLimit(ctx, accountID); err != nil {
|
|
|
|
|
|
return nil, err
|
|
|
|
|
|
}
|
|
|
|
|
|
result.ClearedRateLimit = true
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
return result, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// RecoverAccountAfterSuccessfulTest 将一次成功测试视为正常请求,
|
|
|
|
|
|
// 按需恢复 error / rate-limit / overload / temp-unsched / model-rate-limit 等运行时状态。
|
|
|
|
|
|
func (s *RateLimitService) RecoverAccountAfterSuccessfulTest(ctx context.Context, accountID int64) (*SuccessfulTestRecoveryResult, error) {
|
|
|
|
|
|
return s.RecoverAccountState(ctx, accountID, AccountRecoveryOptions{})
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-03 06:34:00 -08:00
|
|
|
|
func (s *RateLimitService) ClearTempUnschedulable(ctx context.Context, accountID int64) error {
|
|
|
|
|
|
if err := s.accountRepo.ClearTempUnschedulable(ctx, accountID); err != nil {
|
|
|
|
|
|
return err
|
|
|
|
|
|
}
|
|
|
|
|
|
if s.tempUnschedCache != nil {
|
|
|
|
|
|
if err := s.tempUnschedCache.DeleteTempUnsched(ctx, accountID); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("temp_unsched_cache_delete_failed", "account_id", accountID, "error", err)
|
2026-01-03 06:34:00 -08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
2026-02-09 20:08:00 +08:00
|
|
|
|
// 同时清除模型级别限流
|
|
|
|
|
|
if err := s.accountRepo.ClearModelRateLimits(ctx, accountID); err != nil {
|
|
|
|
|
|
slog.Warn("clear_model_rate_limits_on_temp_unsched_reset_failed", "account_id", accountID, "error", err)
|
|
|
|
|
|
}
|
2026-01-03 06:34:00 -08:00
|
|
|
|
return nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-03-08 06:59:53 +08:00
|
|
|
|
func hasRecoverableRuntimeState(account *Account) bool {
|
|
|
|
|
|
if account == nil {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
if account.RateLimitedAt != nil || account.RateLimitResetAt != nil || account.OverloadUntil != nil || account.TempUnschedulableUntil != nil {
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
if len(account.Extra) == 0 {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
return hasNonEmptyMapValue(account.Extra, "model_rate_limits") || hasNonEmptyMapValue(account.Extra, "antigravity_quota_scopes")
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func hasNonEmptyMapValue(extra map[string]any, key string) bool {
|
|
|
|
|
|
raw, ok := extra[key]
|
|
|
|
|
|
if !ok || raw == nil {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
switch typed := raw.(type) {
|
|
|
|
|
|
case map[string]any:
|
|
|
|
|
|
return len(typed) > 0
|
|
|
|
|
|
case map[string]string:
|
|
|
|
|
|
return len(typed) > 0
|
|
|
|
|
|
case []any:
|
|
|
|
|
|
return len(typed) > 0
|
|
|
|
|
|
default:
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-03 06:34:00 -08:00
|
|
|
|
func (s *RateLimitService) GetTempUnschedStatus(ctx context.Context, accountID int64) (*TempUnschedState, error) {
|
|
|
|
|
|
now := time.Now().Unix()
|
|
|
|
|
|
if s.tempUnschedCache != nil {
|
|
|
|
|
|
state, err := s.tempUnschedCache.GetTempUnsched(ctx, accountID)
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
return nil, err
|
|
|
|
|
|
}
|
|
|
|
|
|
if state != nil && state.UntilUnix > now {
|
|
|
|
|
|
return state, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
account, err := s.accountRepo.GetByID(ctx, accountID)
|
|
|
|
|
|
if err != nil {
|
|
|
|
|
|
return nil, err
|
|
|
|
|
|
}
|
|
|
|
|
|
if account.TempUnschedulableUntil == nil {
|
|
|
|
|
|
return nil, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
if account.TempUnschedulableUntil.Unix() <= now {
|
|
|
|
|
|
return nil, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
state := &TempUnschedState{
|
|
|
|
|
|
UntilUnix: account.TempUnschedulableUntil.Unix(),
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if account.TempUnschedulableReason != "" {
|
|
|
|
|
|
var parsed TempUnschedState
|
|
|
|
|
|
if err := json.Unmarshal([]byte(account.TempUnschedulableReason), &parsed); err == nil {
|
|
|
|
|
|
if parsed.UntilUnix == 0 {
|
|
|
|
|
|
parsed.UntilUnix = state.UntilUnix
|
|
|
|
|
|
}
|
|
|
|
|
|
state = &parsed
|
|
|
|
|
|
} else {
|
|
|
|
|
|
state.ErrorMessage = account.TempUnschedulableReason
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if s.tempUnschedCache != nil {
|
|
|
|
|
|
if err := s.tempUnschedCache.SetTempUnsched(ctx, accountID, state); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("temp_unsched_cache_set_failed", "account_id", accountID, "error", err)
|
2026-01-03 06:34:00 -08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
return state, nil
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func (s *RateLimitService) HandleTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte) bool {
|
|
|
|
|
|
if account == nil {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
if !account.ShouldHandleErrorCode(statusCode) {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
return s.tryTempUnschedulable(ctx, account, statusCode, responseBody)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
const tempUnschedBodyMaxBytes = 64 << 10
|
|
|
|
|
|
const tempUnschedMessageMaxBytes = 2048
|
|
|
|
|
|
|
|
|
|
|
|
func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Account, statusCode int, responseBody []byte) bool {
|
|
|
|
|
|
if account == nil {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
if !account.IsTempUnschedulableEnabled() {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
2026-03-04 20:25:15 +08:00
|
|
|
|
// 401 首次命中可临时不可调度(给 token 刷新窗口);
|
|
|
|
|
|
// 若历史上已因 401 进入过临时不可调度,则本次应升级为 error(返回 false 交由默认错误逻辑处理)。
|
|
|
|
|
|
if statusCode == http.StatusUnauthorized {
|
|
|
|
|
|
reason := account.TempUnschedulableReason
|
|
|
|
|
|
// 缓存可能没有 reason,从 DB 回退读取
|
|
|
|
|
|
if reason == "" {
|
|
|
|
|
|
if dbAcc, err := s.accountRepo.GetByID(ctx, account.ID); err == nil && dbAcc != nil {
|
|
|
|
|
|
reason = dbAcc.TempUnschedulableReason
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
if wasTempUnschedByStatusCode(reason, statusCode) {
|
|
|
|
|
|
slog.Info("401_escalated_to_error", "account_id", account.ID,
|
|
|
|
|
|
"reason", "previous temp-unschedulable was also 401")
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
2026-01-03 06:34:00 -08:00
|
|
|
|
rules := account.GetTempUnschedulableRules()
|
|
|
|
|
|
if len(rules) == 0 {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
if statusCode <= 0 || len(responseBody) == 0 {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
body := responseBody
|
|
|
|
|
|
if len(body) > tempUnschedBodyMaxBytes {
|
|
|
|
|
|
body = body[:tempUnschedBodyMaxBytes]
|
|
|
|
|
|
}
|
|
|
|
|
|
bodyLower := strings.ToLower(string(body))
|
|
|
|
|
|
|
|
|
|
|
|
for idx, rule := range rules {
|
|
|
|
|
|
if rule.ErrorCode != statusCode || len(rule.Keywords) == 0 {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
matchedKeyword := matchTempUnschedKeyword(bodyLower, rule.Keywords)
|
|
|
|
|
|
if matchedKeyword == "" {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if s.triggerTempUnschedulable(ctx, account, rule, idx, statusCode, matchedKeyword, responseBody) {
|
|
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-03-04 20:25:15 +08:00
|
|
|
|
func wasTempUnschedByStatusCode(reason string, statusCode int) bool {
|
|
|
|
|
|
if statusCode <= 0 {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
reason = strings.TrimSpace(reason)
|
|
|
|
|
|
if reason == "" {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
var state TempUnschedState
|
|
|
|
|
|
if err := json.Unmarshal([]byte(reason), &state); err != nil {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
return state.StatusCode == statusCode
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-03 06:34:00 -08:00
|
|
|
|
func matchTempUnschedKeyword(bodyLower string, keywords []string) string {
|
|
|
|
|
|
if bodyLower == "" {
|
|
|
|
|
|
return ""
|
|
|
|
|
|
}
|
|
|
|
|
|
for _, keyword := range keywords {
|
|
|
|
|
|
k := strings.TrimSpace(keyword)
|
|
|
|
|
|
if k == "" {
|
|
|
|
|
|
continue
|
|
|
|
|
|
}
|
|
|
|
|
|
if strings.Contains(bodyLower, strings.ToLower(k)) {
|
|
|
|
|
|
return k
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
return ""
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func (s *RateLimitService) triggerTempUnschedulable(ctx context.Context, account *Account, rule TempUnschedulableRule, ruleIndex int, statusCode int, matchedKeyword string, responseBody []byte) bool {
|
|
|
|
|
|
if account == nil {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
if rule.DurationMinutes <= 0 {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
now := time.Now()
|
|
|
|
|
|
until := now.Add(time.Duration(rule.DurationMinutes) * time.Minute)
|
|
|
|
|
|
|
|
|
|
|
|
state := &TempUnschedState{
|
|
|
|
|
|
UntilUnix: until.Unix(),
|
|
|
|
|
|
TriggeredAtUnix: now.Unix(),
|
|
|
|
|
|
StatusCode: statusCode,
|
|
|
|
|
|
MatchedKeyword: matchedKeyword,
|
|
|
|
|
|
RuleIndex: ruleIndex,
|
|
|
|
|
|
ErrorMessage: truncateTempUnschedMessage(responseBody, tempUnschedMessageMaxBytes),
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
reason := ""
|
|
|
|
|
|
if raw, err := json.Marshal(state); err == nil {
|
|
|
|
|
|
reason = string(raw)
|
|
|
|
|
|
}
|
|
|
|
|
|
if reason == "" {
|
|
|
|
|
|
reason = strings.TrimSpace(state.ErrorMessage)
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("temp_unsched_set_failed", "account_id", account.ID, "error", err)
|
2026-01-03 06:34:00 -08:00
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if s.tempUnschedCache != nil {
|
|
|
|
|
|
if err := s.tempUnschedCache.SetTempUnsched(ctx, account.ID, state); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("temp_unsched_cache_set_failed", "account_id", account.ID, "error", err)
|
2026-01-03 06:34:00 -08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Info("account_temp_unschedulable", "account_id", account.ID, "until", until, "rule_index", ruleIndex, "status_code", statusCode)
|
2026-01-03 06:34:00 -08:00
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
func truncateTempUnschedMessage(body []byte, maxBytes int) string {
|
|
|
|
|
|
if maxBytes <= 0 || len(body) == 0 {
|
|
|
|
|
|
return ""
|
|
|
|
|
|
}
|
|
|
|
|
|
if len(body) > maxBytes {
|
|
|
|
|
|
body = body[:maxBytes]
|
|
|
|
|
|
}
|
|
|
|
|
|
return strings.TrimSpace(string(body))
|
|
|
|
|
|
}
|
2026-01-11 21:54:52 -08:00
|
|
|
|
|
|
|
|
|
|
// HandleStreamTimeout 处理流数据超时
|
|
|
|
|
|
// 根据系统设置决定是否标记账户为临时不可调度或错误状态
|
|
|
|
|
|
// 返回是否应该停止该账号的调度
|
|
|
|
|
|
func (s *RateLimitService) HandleStreamTimeout(ctx context.Context, account *Account, model string) bool {
|
|
|
|
|
|
if account == nil {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 获取系统设置
|
|
|
|
|
|
if s.settingService == nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("stream_timeout_setting_service_missing", "account_id", account.ID)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
settings, err := s.settingService.GetStreamTimeoutSettings(ctx)
|
|
|
|
|
|
if err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("stream_timeout_get_settings_failed", "account_id", account.ID, "error", err)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if !settings.Enabled {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if settings.Action == StreamTimeoutActionNone {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 增加超时计数
|
|
|
|
|
|
var count int64 = 1
|
|
|
|
|
|
if s.timeoutCounterCache != nil {
|
|
|
|
|
|
count, err = s.timeoutCounterCache.IncrementTimeoutCount(ctx, account.ID, settings.ThresholdWindowMinutes)
|
|
|
|
|
|
if err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("stream_timeout_increment_count_failed", "account_id", account.ID, "error", err)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
// 继续处理,使用 count=1
|
|
|
|
|
|
count = 1
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Info("stream_timeout_count", "account_id", account.ID, "count", count, "threshold", settings.ThresholdCount, "window_minutes", settings.ThresholdWindowMinutes, "model", model)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
|
|
|
|
|
|
// 检查是否达到阈值
|
|
|
|
|
|
if count < int64(settings.ThresholdCount) {
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 达到阈值,执行相应操作
|
|
|
|
|
|
switch settings.Action {
|
|
|
|
|
|
case StreamTimeoutActionTempUnsched:
|
|
|
|
|
|
return s.triggerStreamTimeoutTempUnsched(ctx, account, settings, model)
|
|
|
|
|
|
case StreamTimeoutActionError:
|
|
|
|
|
|
return s.triggerStreamTimeoutError(ctx, account, model)
|
|
|
|
|
|
default:
|
|
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// triggerStreamTimeoutTempUnsched 触发流超时临时不可调度
|
|
|
|
|
|
func (s *RateLimitService) triggerStreamTimeoutTempUnsched(ctx context.Context, account *Account, settings *StreamTimeoutSettings, model string) bool {
|
|
|
|
|
|
now := time.Now()
|
|
|
|
|
|
until := now.Add(time.Duration(settings.TempUnschedMinutes) * time.Minute)
|
|
|
|
|
|
|
|
|
|
|
|
state := &TempUnschedState{
|
|
|
|
|
|
UntilUnix: until.Unix(),
|
|
|
|
|
|
TriggeredAtUnix: now.Unix(),
|
|
|
|
|
|
StatusCode: 0, // 超时没有状态码
|
|
|
|
|
|
MatchedKeyword: "stream_timeout",
|
|
|
|
|
|
RuleIndex: -1, // 表示系统级规则
|
|
|
|
|
|
ErrorMessage: "Stream data interval timeout for model: " + model,
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
reason := ""
|
|
|
|
|
|
if raw, err := json.Marshal(state); err == nil {
|
|
|
|
|
|
reason = string(raw)
|
|
|
|
|
|
}
|
|
|
|
|
|
if reason == "" {
|
|
|
|
|
|
reason = state.ErrorMessage
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, until, reason); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("stream_timeout_set_temp_unsched_failed", "account_id", account.ID, "error", err)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
if s.tempUnschedCache != nil {
|
|
|
|
|
|
if err := s.tempUnschedCache.SetTempUnsched(ctx, account.ID, state); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("stream_timeout_set_temp_unsched_cache_failed", "account_id", account.ID, "error", err)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 重置超时计数
|
|
|
|
|
|
if s.timeoutCounterCache != nil {
|
|
|
|
|
|
if err := s.timeoutCounterCache.ResetTimeoutCount(ctx, account.ID); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("stream_timeout_reset_count_failed", "account_id", account.ID, "error", err)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Info("stream_timeout_temp_unschedulable", "account_id", account.ID, "until", until, "model", model)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
return true
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// triggerStreamTimeoutError 触发流超时错误状态
|
|
|
|
|
|
func (s *RateLimitService) triggerStreamTimeoutError(ctx context.Context, account *Account, model string) bool {
|
|
|
|
|
|
errorMsg := "Stream data interval timeout (repeated failures) for model: " + model
|
|
|
|
|
|
|
|
|
|
|
|
if err := s.accountRepo.SetError(ctx, account.ID, errorMsg); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("stream_timeout_set_error_failed", "account_id", account.ID, "error", err)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
return false
|
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
// 重置超时计数
|
|
|
|
|
|
if s.timeoutCounterCache != nil {
|
|
|
|
|
|
if err := s.timeoutCounterCache.ResetTimeoutCount(ctx, account.ID); err != nil {
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("stream_timeout_reset_count_failed", "account_id", account.ID, "error", err)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
}
|
|
|
|
|
|
}
|
|
|
|
|
|
|
2026-01-14 15:55:44 +08:00
|
|
|
|
slog.Warn("stream_timeout_account_error", "account_id", account.ID, "model", model)
|
2026-01-11 21:54:52 -08:00
|
|
|
|
return true
|
|
|
|
|
|
}
|