package service import ( "context" "database/sql" "encoding/json" "errors" "log" "log/slog" "strings" "sync" "time" ) // SettingKeySpendGuardSettings 烧钱防护配置(JSON,见 SpendGuardSettings)。 const SettingKeySpendGuardSettings = "spend_guard_settings" // SpendGuardSettings 烧钱防护阈值:窗口内某 key 的 token 速率或错误率 // 超阈即冻结该 key(status=disabled),仅管理员手动解冻。 type SpendGuardSettings struct { Enabled bool `json:"enabled"` WindowMinutes int `json:"window_minutes"` TokensPerMinute float64 `json:"tokens_per_minute"` MinRequests int `json:"min_requests"` MaxErrorRate float64 `json:"max_error_rate"` IntervalSeconds int `json:"interval_seconds"` } // DefaultSpendGuardSettings 默认:5 分钟窗口,速率 ≥200 万 token/分钟或 // 错误率 ≥80%(≥5 请求)冻结。默认关闭,由管理员显式启用。 func DefaultSpendGuardSettings() SpendGuardSettings { return SpendGuardSettings{ Enabled: false, WindowMinutes: 5, TokensPerMinute: 2000000, MinRequests: 5, MaxErrorRate: 0.8, IntervalSeconds: 60, } } // SpendGuardOffender 烧钱嫌疑 key 行(管理端展示)。 type SpendGuardOffender struct { APIKeyID int64 `json:"api_key_id"` Name string `json:"name"` UserID int64 `json:"user_id"` Status string `json:"status"` Requests int64 `json:"requests"` Tokens int64 `json:"tokens"` TokensPerMin float64 `json:"tokens_per_min"` Errors int64 `json:"errors"` ErrRate float64 `json:"err_rate"` Frozen bool `json:"frozen"` } // SpendGuardEvent 冻结/解冻事件(内存保留近 100 条)。 type SpendGuardEvent struct { At time.Time `json:"at"` APIKeyID int64 `json:"api_key_id"` Name string `json:"name"` Action string `json:"action"` // frozen / unfrozen Reason string `json:"reason"` } // SpendGuardKeyStore 冻结需要的 key 读写接口(APIKeyService 已实现)。 type SpendGuardKeyStore interface { GetByID(ctx context.Context, id int64) (*APIKey, error) Update(ctx context.Context, id int64, userID int64, req UpdateAPIKeyRequest) (*APIKey, error) } // SpendGuardService 烧钱防护:窗口 token 速率 / 错误率超阈自动冻结 key。 type SpendGuardService struct { db *sql.DB keys SpendGuardKeyStore settings SettingRepository mu sync.RWMutex frozen map[int64]time.Time events []SpendGuardEvent stopCh chan struct{} stopOnce sync.Once wg sync.WaitGroup } // NewSpendGuardService 创建烧钱防护服务(不启动后台循环)。 func NewSpendGuardService(db *sql.DB, keys SpendGuardKeyStore, settingRepo SettingRepository) *SpendGuardService { return &SpendGuardService{ db: db, keys: keys, settings: settingRepo, frozen: make(map[int64]time.Time), stopCh: make(chan struct{}), } } // ProvideSpendGuardService 创建并启动烧钱防护服务(wire 用)。 func ProvideSpendGuardService(db *sql.DB, apiKeyService *APIKeyService, settingRepo SettingRepository) *SpendGuardService { svc := NewSpendGuardService(db, apiKeyService, settingRepo) svc.Start() return svc } // Start 启动周期评估。 func (s *SpendGuardService) Start() { if s == nil || s.db == nil || s.keys == nil { return } s.wg.Add(1) go func() { defer s.wg.Done() ticker := time.NewTicker(time.Duration(s.currentSettings().IntervalSeconds) * time.Second) defer ticker.Stop() s.runOnce() for { select { case <-ticker.C: s.runOnce() case <-s.stopCh: return } } }() } // Stop 停止后台循环。 func (s *SpendGuardService) Stop() { if s == nil { return } s.stopOnce.Do(func() { close(s.stopCh) }) s.wg.Wait() } func (s *SpendGuardService) currentSettings() SpendGuardSettings { cfg := DefaultSpendGuardSettings() if s == nil || s.settings == nil { return cfg } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() raw, err := s.settings.GetValue(ctx, SettingKeySpendGuardSettings) if err != nil || strings.TrimSpace(raw) == "" { return cfg } var parsed SpendGuardSettings if err := json.Unmarshal([]byte(raw), &parsed); err != nil { slog.Warn("spend_guard.bad_settings", "error", err) return cfg } return normalizeSpendGuardSettings(parsed) } func normalizeSpendGuardSettings(in SpendGuardSettings) SpendGuardSettings { def := DefaultSpendGuardSettings() if in.WindowMinutes <= 0 { in.WindowMinutes = def.WindowMinutes } if in.TokensPerMinute < 0 { in.TokensPerMinute = def.TokensPerMinute } if in.MinRequests <= 0 { in.MinRequests = def.MinRequests } if in.MaxErrorRate <= 0 || in.MaxErrorRate > 1 { in.MaxErrorRate = def.MaxErrorRate } if in.IntervalSeconds < 10 { in.IntervalSeconds = def.IntervalSeconds } return in } // GetSettings 读取当前配置。 func (s *SpendGuardService) GetSettings(_ context.Context) SpendGuardSettings { if s == nil { return DefaultSpendGuardSettings() } return s.currentSettings() } // UpdateSettings 保存配置。 func (s *SpendGuardService) UpdateSettings(ctx context.Context, in SpendGuardSettings) (SpendGuardSettings, error) { if s == nil || s.settings == nil { return DefaultSpendGuardSettings(), nil } norm := normalizeSpendGuardSettings(in) raw, err := json.Marshal(norm) if err != nil { return norm, err } if err := s.settings.Set(ctx, SettingKeySpendGuardSettings, string(raw)); err != nil { return norm, err } return norm, nil } // Events 返回近期冻结/解冻事件。 func (s *SpendGuardService) Events() []SpendGuardEvent { if s == nil { return nil } s.mu.RLock() defer s.mu.RUnlock() out := make([]SpendGuardEvent, len(s.events)) copy(out, s.events) return out } func (s *SpendGuardService) pushEvent(ev SpendGuardEvent) { s.mu.Lock() defer s.mu.Unlock() s.events = append([]SpendGuardEvent{ev}, s.events...) if len(s.events) > 100 { s.events = s.events[:100] } } type spendGuardRow struct { keyID int64 name string userID int64 status string reqs int64 tokens int64 errs int64 } // Offenders 查询窗口内各 key 的用量与错误(管理端列表与评估共用)。 func (s *SpendGuardService) Offenders(ctx context.Context, windowMinutes int) ([]SpendGuardOffender, error) { if s == nil || s.db == nil { return nil, nil } if windowMinutes <= 0 { windowMinutes = 5 } // usage 与 errors 双源 UNION:纯报错(零成功行)的 key 也必须可见, // 否则错误率冻结对最坏情况失效。 q := ` SELECT ak.id, COALESCE(ak.name, ''), COALESCE(ak.user_id, 0), COALESCE(ak.status, ''), COALESCE(SUM(t.reqs), 0), COALESCE(SUM(t.toks), 0), COALESCE(SUM(t.errs), 0) FROM ( SELECT api_key_id AS key_id, COUNT(*) AS reqs, SUM(input_tokens + output_tokens + cache_read_tokens + cache_creation_tokens) AS toks, 0 AS errs FROM usage_logs WHERE created_at >= NOW() - ($1 || ' minutes')::interval AND api_key_id IS NOT NULL AND api_key_id > 0 GROUP BY api_key_id UNION ALL SELECT api_key_id AS key_id, 0 AS reqs, 0 AS toks, COUNT(*) AS errs FROM ops_error_logs WHERE created_at >= NOW() - ($1 || ' minutes')::interval AND COALESCE(status_code, 0) >= 400 AND NOT COALESCE(is_business_limited, FALSE) AND api_key_id IS NOT NULL GROUP BY api_key_id ) t JOIN api_keys ak ON ak.id = t.key_id GROUP BY ak.id, ak.name, ak.user_id, ak.status ORDER BY SUM(t.toks) DESC, SUM(t.errs) DESC LIMIT 500` rows, err := s.db.QueryContext(ctx, q, windowMinutes) if err != nil { return nil, err } defer func() { _ = rows.Close() }() var out []SpendGuardOffender for rows.Next() { var r spendGuardRow if err := rows.Scan(&r.keyID, &r.name, &r.userID, &r.status, &r.reqs, &r.tokens, &r.errs); err != nil { return nil, err } o := SpendGuardOffender{ APIKeyID: r.keyID, Name: r.name, UserID: r.userID, Status: r.status, Requests: r.reqs, Tokens: r.tokens, Errors: r.errs, TokensPerMin: float64(r.tokens) / float64(windowMinutes), } if total := r.reqs + r.errs; total > 0 { o.ErrRate = float64(r.errs) / float64(total) } s.mu.RLock() _, o.Frozen = s.frozen[r.keyID] s.mu.RUnlock() out = append(out, o) } return out, rows.Err() } func spendGuardBreach(o SpendGuardOffender, cfg SpendGuardSettings) (bool, string) { if o.Requests+o.Errors < int64(cfg.MinRequests) { return false, "" } if cfg.TokensPerMinute > 0 && o.TokensPerMin >= cfg.TokensPerMinute { return true, "token velocity" } if o.Errors > 0 && o.ErrRate >= cfg.MaxErrorRate { return true, "error rate" } return false, "" } func (s *SpendGuardService) runOnce() { cfg := s.currentSettings() if !cfg.Enabled { return } ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() offenders, err := s.Offenders(ctx, cfg.WindowMinutes) if err != nil { log.Printf("[SpendGuard] offenders query failed: %v", err) return } now := time.Now() for _, o := range offenders { ok, reason := spendGuardBreach(o, cfg) if !ok { continue } s.mu.RLock() _, already := s.frozen[o.APIKeyID] s.mu.RUnlock() if already || o.Status != StatusAPIKeyActive { continue } key, err := s.keys.GetByID(ctx, o.APIKeyID) if err != nil || key == nil || key.Status != StatusAPIKeyActive { continue } disabled := StatusAPIKeyDisabled if _, err := s.keys.Update(ctx, o.APIKeyID, key.UserID, UpdateAPIKeyRequest{Status: &disabled}); err != nil { log.Printf("[SpendGuard] freeze key %d failed: %v", o.APIKeyID, err) continue } s.mu.Lock() s.frozen[o.APIKeyID] = now s.mu.Unlock() slog.Warn("spend_guard.key_frozen", "api_key_id", o.APIKeyID, "reason", reason, "tokens_per_min", o.TokensPerMin, "err_rate", o.ErrRate) s.pushEvent(SpendGuardEvent{At: now, APIKeyID: o.APIKeyID, Name: o.Name, Action: "frozen", Reason: reason}) } } // Unfreeze 手动解冻 key(恢复 active 并清除跟踪)。 // 仅当 key 当前为 disabled 才写回,避免覆盖 quota_exhausted/expired 等中间态。 func (s *SpendGuardService) Unfreeze(ctx context.Context, id int64) error { if s == nil || s.keys == nil { return nil } key, err := s.keys.GetByID(ctx, id) if err != nil { return err } if key == nil { return errors.New("api key not found") } if key.Status != StatusAPIKeyDisabled { s.mu.Lock() delete(s.frozen, id) s.mu.Unlock() return nil } active := StatusAPIKeyActive if _, err := s.keys.Update(ctx, id, key.UserID, UpdateAPIKeyRequest{Status: &active}); err != nil { return err } s.mu.Lock() delete(s.frozen, id) s.mu.Unlock() s.pushEvent(SpendGuardEvent{At: time.Now(), APIKeyID: id, Name: key.Name, Action: "unfrozen", Reason: "manual"}) return nil }