descrybe/apps/api/internal/support/auto_ratelimit.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
}