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 }