gochat/backend/internal/autoassignment/rate_limiter.go

package autoassignment

// RateLimiter implements Redis-backed rate limiting for auto-assignment.
//
// Reference: Chatwoot RateLimiter (Redis key-per-assignment)
//   - Tracks count of assignments per agent per inbox per time window
//   - Fair distribution limit (default 5) over window (default 5 min)
//   - Key format: gochat:rate_limit:{inboxID}:{agentID}:{windowKey}
//   - Ensures no single agent gets too many assignments in a short time
//
// In gochat, we use Redis string keys with TTL for rate counting:
//   - Key: gochat:auto_assignment_rate:{inboxID}:{agentID}:{windowStart}
//   - Value: count of assignments in this window
//   - TTL: window duration (so keys expire automatically)

import (
	"context"
	"fmt"
	"strconv"
	"time"

	"github.com/redis/go-redis/v9"
)

const (
	// rateLimitKeyPrefix is the Redis key prefix for rate limit counters.
	rateLimitKeyPrefix = "gochat:auto_assignment_rate:"
)

// RateLimiter tracks assignment counts per agent per inbox per time window.
type RateLimiter struct {
	redis *redis.Client
}

// NewRateLimiter creates a new RateLimiter.
func NewRateLimiter(rdb *redis.Client) *RateLimiter {
	return &RateLimiter{redis: rdb}
}

// key returns the Redis key for rate limiting the given agent in the given inbox.
// The key includes a time window component that changes when the window expires.
func (rl *RateLimiter) key(inboxID uint, agentID uint, windowSeconds int) string {
	windowStart := time.Now().Unix() / int64(windowSeconds)
	return fmt.Sprintf("%s%d:%d:%d", rateLimitKeyPrefix, inboxID, agentID, windowStart)
}

// IsAllowed checks whether the given agent is allowed to receive another
// assignment in the current time window.
//
// Returns true if the agent's assignment count in this window is below the limit.
func (rl *RateLimiter) IsAllowed(ctx context.Context, inboxID uint, agentID uint, limit int, windowSeconds int) bool {
	if limit <= 0 {
		return true // no rate limiting configured
	}

	key := rl.key(inboxID, agentID, windowSeconds)
	count, err := rl.redis.Get(ctx, key).Int()
	if err == redis.Nil {
		return true // no assignments yet in this window
	}
	if err != nil {
		// Redis error — allow assignment to avoid blocking service
		return true
	}

	return count < limit
}

// Increment increments the assignment count for the given agent in the
// current time window.
//
// Sets a TTL on the key equal to the window duration, so the key
// expires automatically when the window ends.
func (rl *RateLimiter) Increment(ctx context.Context, inboxID uint, agentID uint, windowSeconds int) error {
	key := rl.key(inboxID, agentID, windowSeconds)

	// Increment count
	if err := rl.redis.Incr(ctx, key).Err(); err != nil {
		return fmt.Errorf("redis Incr: %w", err)
	}

	// Set TTL (only on first increment in this window)
	ttl, err := rl.redis.TTL(ctx, key).Result()
	if err != nil {
		return fmt.Errorf("redis TTL: %w", err)
	}
	if ttl == -1 { // key has no TTL (first increment)
		if err := rl.redis.Expire(ctx, key, time.Duration(windowSeconds)*time.Second).Err(); err != nil {
			return fmt.Errorf("redis Expire: %w", err)
		}
	}

	return nil
}

// GetCount returns the current assignment count for the given agent
// in the current time window.
func (rl *RateLimiter) GetCount(ctx context.Context, inboxID uint, agentID uint, windowSeconds int) (int, error) {
	key := rl.key(inboxID, agentID, windowSeconds)
	count, err := rl.redis.Get(ctx, key).Int()
	if err == redis.Nil {
		return 0, nil
	}
	if err != nil {
		return 0, fmt.Errorf("redis Get: %w", err)
	}
	return count, nil
}

// Reset resets all rate limit counters for the given inbox.
// This is useful when inbox members change significantly.
func (rl *RateLimiter) Reset(ctx context.Context, inboxID uint) error {
	pattern := fmt.Sprintf("%s%d:*", rateLimitKeyPrefix, inboxID)
	keys, err := rl.redis.Keys(ctx, pattern).Result()
	if err != nil {
		return fmt.Errorf("redis Keys: %w", err)
	}
	if len(keys) > 0 {
		if err := rl.redis.Del(ctx, keys...).Err(); err != nil {
			return fmt.Errorf("redis Del: %w", err)
		}
	}
	return nil
}

// ParseAgentID extracts a uint agent ID from a string representation.
// Utility function for converting Redis string values to uint IDs.
func ParseAgentID(s string) (uint, error) {
	id, err := strconv.ParseUint(s, 10, 32)
	if err != nil {
		return 0, err
	}
	return uint(id), nil
}