greeneclipse
8580c996c3
Drop one-shot tmp/axe scripts and agent i18n scratch so the Gitea tree is deployable.
137 lines
3.1 KiB
Go
137 lines
3.1 KiB
Go
package support
|
||
|
||
import (
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/google/uuid"
|
||
)
|
||
|
||
const (
|
||
defaultAICompanyPerHour = 10
|
||
defaultAIPlatformPerMinute = 30
|
||
)
|
||
|
||
// AIRateLimiter bounds AI auto-reply jobs (not FAQ matches).
|
||
// In-process only — effective limit ≈ N × replicas (same pattern as processing.StartLimiter).
|
||
// RATE_LIMIT_REPLICAS does not divide this limiter; multi-replica hard caps need edge/WAF.
|
||
type AIRateLimiter struct {
|
||
mu sync.Mutex
|
||
|
||
companyLimit int
|
||
companyWindow time.Duration
|
||
companyHits map[uuid.UUID][]time.Time
|
||
|
||
platformLimit int
|
||
platformWindow time.Duration
|
||
platformHits []time.Time
|
||
|
||
lastGC time.Time
|
||
}
|
||
|
||
// NewAIRateLimiter builds a limiter with contract defaults (10/company/hour, 30/platform/min).
|
||
func NewAIRateLimiter(companyPerHour, platformPerMinute int) *AIRateLimiter {
|
||
if companyPerHour <= 0 {
|
||
companyPerHour = defaultAICompanyPerHour
|
||
}
|
||
if platformPerMinute <= 0 {
|
||
platformPerMinute = defaultAIPlatformPerMinute
|
||
}
|
||
return &AIRateLimiter{
|
||
companyLimit: companyPerHour,
|
||
companyWindow: time.Hour,
|
||
companyHits: make(map[uuid.UUID][]time.Time),
|
||
platformLimit: platformPerMinute,
|
||
platformWindow: time.Minute,
|
||
lastGC: time.Now(),
|
||
}
|
||
}
|
||
|
||
// Allow reports whether an AI auto-reply job may proceed for companyID.
|
||
// On deny, no counters are incremented (caller may retry later).
|
||
func (l *AIRateLimiter) Allow(companyID uuid.UUID) bool {
|
||
if l == nil {
|
||
return true
|
||
}
|
||
now := time.Now()
|
||
l.mu.Lock()
|
||
defer l.mu.Unlock()
|
||
|
||
l.gcLocked(now)
|
||
|
||
companyCut := now.Add(-l.companyWindow)
|
||
ch := l.companyHits[companyID]
|
||
keptC := ch[:0]
|
||
for _, t := range ch {
|
||
if t.After(companyCut) {
|
||
keptC = append(keptC, t)
|
||
}
|
||
}
|
||
if len(keptC) >= l.companyLimit {
|
||
l.companyHits[companyID] = keptC
|
||
return false
|
||
}
|
||
|
||
platformCut := now.Add(-l.platformWindow)
|
||
keptP := l.platformHits[:0]
|
||
for _, t := range l.platformHits {
|
||
if t.After(platformCut) {
|
||
keptP = append(keptP, t)
|
||
}
|
||
}
|
||
if len(keptP) >= l.platformLimit {
|
||
l.platformHits = keptP
|
||
l.companyHits[companyID] = keptC
|
||
return false
|
||
}
|
||
|
||
l.companyHits[companyID] = append(keptC, now)
|
||
l.platformHits = append(keptP, now)
|
||
return true
|
||
}
|
||
|
||
func (l *AIRateLimiter) gcLocked(now time.Time) {
|
||
if now.Sub(l.lastGC) < l.companyWindow {
|
||
return
|
||
}
|
||
companyCut := now.Add(-l.companyWindow)
|
||
for id, ts := range l.companyHits {
|
||
kept := ts[:0]
|
||
for _, t := range ts {
|
||
if t.After(companyCut) {
|
||
kept = append(kept, t)
|
||
}
|
||
}
|
||
if len(kept) == 0 {
|
||
delete(l.companyHits, id)
|
||
} else {
|
||
l.companyHits[id] = kept
|
||
}
|
||
}
|
||
platformCut := now.Add(-l.platformWindow)
|
||
keptP := l.platformHits[:0]
|
||
for _, t := range l.platformHits {
|
||
if t.After(platformCut) {
|
||
keptP = append(keptP, t)
|
||
}
|
||
}
|
||
l.platformHits = keptP
|
||
l.lastGC = now
|
||
}
|
||
|
||
// package-level limiter used by TryAutoReplyLLM until Service gains an injected field.
|
||
var defaultAIRateLimiter = NewAIRateLimiter(0, 0)
|
||
|
||
// SetAIRateLimiter replaces the package default (tests / wiring).
|
||
func SetAIRateLimiter(l *AIRateLimiter) {
|
||
if l == nil {
|
||
l = NewAIRateLimiter(0, 0)
|
||
}
|
||
defaultAIRateLimiter = l
|
||
}
|
||
|
||
// AIRateLimiterDefault returns the package limiter.
|
||
func AIRateLimiterDefault() *AIRateLimiter {
|
||
return defaultAIRateLimiter
|
||
}
|