second commit
This commit is contained in:
@@ -0,0 +1,195 @@
|
||||
package autoassignment
|
||||
|
||||
// AutoAssignmentListener implements channel.EventListener for auto-assignment.
|
||||
//
|
||||
// Reference: Chatwoot pattern
|
||||
// - Chatwoot uses an AsyncDispatcher that dispatches to listeners like
|
||||
// AutomationRuleListener, which in turn trigger assignment.
|
||||
// - In gochat, the AutoAssignmentListener subscribes to conversation events
|
||||
// and triggers the AssignmentService when a new conversation is created
|
||||
// or an existing conversation is reopened without an assignee.
|
||||
//
|
||||
// Events this listener subscribes to:
|
||||
// - conversation.created: assign new conversations
|
||||
// - conversation.opened: assign reopened conversations (if unassigned)
|
||||
// - conversation.unassigned: re-assign when agent is removed
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/gochat/gochat/internal/channel"
|
||||
"github.com/gochat/gochat/internal/dispatch"
|
||||
"github.com/gochat/gochat/internal/model"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"gorm.io/gorm"
|
||||
applogger "github.com/gochat/gochat/pkg/logger"
|
||||
)
|
||||
|
||||
// AutoAssignmentListener listens for conversation events and triggers
|
||||
// auto-assignment.
|
||||
type AutoAssignmentListener struct {
|
||||
db *gorm.DB
|
||||
redis *redis.Client
|
||||
service *AssignmentService
|
||||
}
|
||||
|
||||
// NewAutoAssignmentListener creates a new AutoAssignmentListener.
|
||||
func NewAutoAssignmentListener(db *gorm.DB, rdb *redis.Client) *AutoAssignmentListener {
|
||||
return &AutoAssignmentListener{
|
||||
db: db,
|
||||
redis: rdb,
|
||||
service: NewAssignmentService(db, rdb),
|
||||
}
|
||||
}
|
||||
|
||||
// Name returns the unique identifier for this listener.
|
||||
func (l *AutoAssignmentListener) Name() string {
|
||||
return "auto_assignment_listener"
|
||||
}
|
||||
|
||||
// OnEvent processes the given event.
|
||||
// It triggers auto-assignment for conversation events that need it.
|
||||
func (l *AutoAssignmentListener) OnEvent(ctx context.Context, event *channel.ChannelEvent) error {
|
||||
switch event.Type {
|
||||
case channel.EventConversationCreated:
|
||||
return l.onConversationCreated(ctx, event)
|
||||
case channel.EventConversationOpened:
|
||||
return l.onConversationOpened(ctx, event)
|
||||
case channel.EventConversationUnassigned:
|
||||
return l.onConversationUnassigned(ctx, event)
|
||||
default:
|
||||
return nil // not our event
|
||||
}
|
||||
}
|
||||
|
||||
// onConversationCreated handles new conversation events.
|
||||
// It triggers auto-assignment for the conversation's inbox.
|
||||
func (l *AutoAssignmentListener) onConversationCreated(ctx context.Context, event *channel.ChannelEvent) error {
|
||||
conversationID, err := dispatch.ExtractConversationID(event.Data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("extract conversation_id: %w", err)
|
||||
}
|
||||
|
||||
inboxID, err := dispatch.ExtractInboxID(event.Data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("extract inbox_id: %w", err)
|
||||
}
|
||||
|
||||
accountID, err := dispatch.ExtractAccountID(event.Data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("extract account_id: %w", err)
|
||||
}
|
||||
|
||||
applogger.L().Infof("auto-assignment: new conversation %d in inbox %d", conversationID, inboxID)
|
||||
|
||||
assignedAgent, err := l.service.AssignConversation(ctx, conversationID, inboxID, accountID)
|
||||
if err != nil {
|
||||
applogger.L().Errorf("auto-assignment failed for conversation %d: %v", conversationID, err)
|
||||
return err
|
||||
}
|
||||
|
||||
if assignedAgent > 0 {
|
||||
applogger.L().Infof("auto-assignment: conversation %d assigned to agent %d", conversationID, assignedAgent)
|
||||
} else {
|
||||
applogger.L().Infof("auto-assignment: no agent available for conversation %d", conversationID)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// onConversationOpened handles reopened conversation events.
|
||||
// It triggers auto-assignment only if the conversation has no assignee.
|
||||
func (l *AutoAssignmentListener) onConversationOpened(ctx context.Context, event *channel.ChannelEvent) error {
|
||||
conversationID, err := dispatch.ExtractConversationID(event.Data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("extract conversation_id: %w", err)
|
||||
}
|
||||
|
||||
// Check if conversation is unassigned
|
||||
conv, err := l.getConversation(ctx, conversationID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("get conversation: %w", err)
|
||||
}
|
||||
|
||||
// Only auto-assign if unassigned
|
||||
if conv.AssigneeID != nil {
|
||||
return nil // already assigned
|
||||
}
|
||||
|
||||
inboxID, err := dispatch.ExtractInboxID(event.Data)
|
||||
if err != nil {
|
||||
inboxID = conv.InboxID // fallback to conversation's inbox
|
||||
}
|
||||
|
||||
accountID, err := dispatch.ExtractAccountID(event.Data)
|
||||
if err != nil {
|
||||
accountID = conv.AccountID // fallback to conversation's account
|
||||
}
|
||||
|
||||
applogger.L().Infof("auto-assignment: reopened conversation %d in inbox %d", conversationID, inboxID)
|
||||
|
||||
assignedAgent, err := l.service.AssignConversation(ctx, conversationID, inboxID, accountID)
|
||||
if err != nil {
|
||||
applogger.L().Errorf("auto-assignment failed for conversation %d: %v", conversationID, err)
|
||||
return err
|
||||
}
|
||||
|
||||
if assignedAgent > 0 {
|
||||
applogger.L().Infof("auto-assignment: conversation %d assigned to agent %d", conversationID, assignedAgent)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// onConversationUnassigned handles unassigned conversation events.
|
||||
// It triggers re-assignment for the conversation.
|
||||
func (l *AutoAssignmentListener) onConversationUnassigned(ctx context.Context, event *channel.ChannelEvent) error {
|
||||
conversationID, err := dispatch.ExtractConversationID(event.Data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("extract conversation_id: %w", err)
|
||||
}
|
||||
|
||||
inboxID, err := dispatch.ExtractInboxID(event.Data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("extract inbox_id: %w", err)
|
||||
}
|
||||
|
||||
accountID, err := dispatch.ExtractAccountID(event.Data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("extract account_id: %w", err)
|
||||
}
|
||||
|
||||
applogger.L().Infof("auto-assignment: unassigned conversation %d in inbox %d", conversationID, inboxID)
|
||||
|
||||
assignedAgent, err := l.service.AssignConversation(ctx, conversationID, inboxID, accountID)
|
||||
if err != nil {
|
||||
applogger.L().Errorf("auto-assignment re-assign failed for conversation %d: %v", conversationID, err)
|
||||
return err
|
||||
}
|
||||
|
||||
if assignedAgent > 0 {
|
||||
applogger.L().Infof("auto-assignment: conversation %d re-assigned to agent %d", conversationID, assignedAgent)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// getConversation fetches a conversation by ID.
|
||||
func (l *AutoAssignmentListener) getConversation(ctx context.Context, id uint) (*model.Conversation, error) {
|
||||
var conv model.Conversation
|
||||
if err := l.db.WithContext(ctx).First(&conv, id).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &conv, nil
|
||||
}
|
||||
|
||||
// EventNames returns the event names this listener subscribes to.
|
||||
// Used when registering with the EventDispatcher.
|
||||
func EventNames() []string {
|
||||
return []string{
|
||||
string(channel.EventConversationCreated),
|
||||
string(channel.EventConversationOpened),
|
||||
string(channel.EventConversationUnassigned),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,105 @@
|
||||
package autoassignment
|
||||
|
||||
// LowestLoadSelector picks the agent with the fewest open conversations.
|
||||
//
|
||||
// Reference: Chatwoot "least_busy" concept
|
||||
// - Agent with lowest workload (fewest open conversations) gets the next assignment.
|
||||
// - This ensures fair distribution based on actual current load, not just rotation order.
|
||||
//
|
||||
// Implementation:
|
||||
// - Queries the conversations table for open conversations in the given inbox,
|
||||
// grouped by assignee_id to count load per agent.
|
||||
// - Agents not present in the result set have zero load.
|
||||
// - Picks the agent with the lowest load among the candidates.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"gorm.io/gorm"
|
||||
applogger "github.com/gochat/gochat/pkg/logger"
|
||||
)
|
||||
|
||||
// ConversationQueryModel is a minimal GORM model for querying conversation
|
||||
// counts from the conversations table. It only includes the fields needed
|
||||
// for the lowest-load selection query: AssigneeID, InboxID, and Status.
|
||||
type ConversationQueryModel struct {
|
||||
AssigneeID uint `gorm:"column:assignee_id"`
|
||||
InboxID uint `gorm:"column:inbox_id"`
|
||||
Status string `gorm:"column:status"`
|
||||
}
|
||||
|
||||
// TableName returns the GORM table name for ConversationQueryModel.
|
||||
func (ConversationQueryModel) TableName() string {
|
||||
return "conversations"
|
||||
}
|
||||
|
||||
// LowestLoadSelector selects the agent with the fewest open conversations
|
||||
// in a given inbox.
|
||||
type LowestLoadSelector struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewLowestLoadSelector creates a new LowestLoadSelector.
|
||||
func NewLowestLoadSelector(db *gorm.DB) *LowestLoadSelector {
|
||||
return &LowestLoadSelector{db: db}
|
||||
}
|
||||
|
||||
// Select picks the agent with the lowest load (fewest open conversations)
|
||||
// among the given candidate agentIDs for the specified inbox.
|
||||
//
|
||||
// Steps:
|
||||
// 1. Query conversations table for open conversations in the inbox
|
||||
// that are assigned to any of the candidate agents, grouped by
|
||||
// assignee_id, to get the count (load) per agent.
|
||||
// 2. For agents not in the result, assume load = 0.
|
||||
// 3. Return the agentID with the lowest load.
|
||||
// 4. If there are ties, the first agent in the agentIDs list with
|
||||
// that load wins (deterministic tie-break based on caller ordering).
|
||||
func (ll *LowestLoadSelector) Select(ctx context.Context, inboxID uint, agentIDs []uint) (uint, error) {
|
||||
if len(agentIDs) == 0 {
|
||||
return 0, fmt.Errorf("no candidate agents provided")
|
||||
}
|
||||
|
||||
// Build a map of agent load from the database.
|
||||
// Query: SELECT assignee_id, COUNT(*) as load
|
||||
// FROM conversations
|
||||
// WHERE inbox_id = ? AND status = 'open' AND assignee_id IN (?)
|
||||
// GROUP BY assignee_id
|
||||
type loadRow struct {
|
||||
AssigneeID uint
|
||||
Load int
|
||||
}
|
||||
|
||||
var rows []loadRow
|
||||
err := ll.db.WithContext(ctx).
|
||||
Model(ConversationQueryModel{}).
|
||||
Select("assignee_id, COUNT(*) as load").
|
||||
Where("inbox_id = ? AND status = ? AND assignee_id IN ?", inboxID, "open", agentIDs).
|
||||
Group("assignee_id").
|
||||
Find(&rows).Error
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("query agent loads: %w", err)
|
||||
}
|
||||
|
||||
// Build load map: agentID -> load count (default 0 for agents not in result)
|
||||
loadMap := make(map[uint]int, len(agentIDs))
|
||||
for _, row := range rows {
|
||||
loadMap[row.AssigneeID] = row.Load
|
||||
}
|
||||
|
||||
// Find agent with lowest load, using agentIDs order as tie-breaker
|
||||
bestAgentID := agentIDs[0]
|
||||
bestLoad := loadMap[bestAgentID]
|
||||
|
||||
for _, agentID := range agentIDs[1:] {
|
||||
agentLoad := loadMap[agentID]
|
||||
if agentLoad < bestLoad {
|
||||
bestLoad = agentLoad
|
||||
bestAgentID = agentID
|
||||
}
|
||||
}
|
||||
|
||||
applogger.L().Debugf("lowest_load: selected agent %d with load %d for inbox %d", bestAgentID, bestLoad, inboxID)
|
||||
return bestAgentID, nil
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
package autoassignment
|
||||
|
||||
// GORM models for the auto-assignment system.
|
||||
//
|
||||
// Reference: Chatwoot AutoAssignment pattern
|
||||
// - AssignmentPolicy: global policy for how conversations are assigned
|
||||
// (round_robin, longest_waiting). Has fair_distribution_limit and
|
||||
// fair_distribution_window for rate limiting per agent.
|
||||
// - InboxAssignmentPolicy: per-inbox override of the global policy.
|
||||
// Each inbox can have its own assignment policy and rate limits.
|
||||
//
|
||||
// Pattern: gochat models embed model.Base (ID, CreatedAt, UpdatedAt, DeletedAt),
|
||||
// use TableName() method, and use gorm struct tags.
|
||||
|
||||
import (
|
||||
"github.com/gochat/gochat/internal/model"
|
||||
)
|
||||
|
||||
// AssignmentPolicyType defines the type of auto-assignment policy.
|
||||
type AssignmentPolicyType string
|
||||
|
||||
const (
|
||||
// PolicyRoundRobin assigns conversations to agents in round-robin order.
|
||||
// Reference: Chatwoot InboxRoundRobinService
|
||||
PolicyRoundRobin AssignmentPolicyType = "round_robin"
|
||||
|
||||
// PolicyLongestWaiting assigns conversations to the agent with the
|
||||
// longest idle time since their last assignment.
|
||||
// Reference: Chatwoot "longest_waiting" policy (planned feature)
|
||||
PolicyLongestWaiting AssignmentPolicyType = "longest_waiting"
|
||||
|
||||
// PolicyLowestLoad assigns conversations to the agent with the
|
||||
// fewest currently open conversations.
|
||||
// Reference: Chatwoot "least_busy" concept — agent with lowest workload.
|
||||
PolicyLowestLoad AssignmentPolicyType = "lowest_load"
|
||||
)
|
||||
|
||||
// AssignmentPolicy represents the global auto-assignment policy for an account.
|
||||
// Reference: Chatwoot AssignmentPolicy model
|
||||
// - Defines how unassigned conversations are distributed among agents
|
||||
// - fair_distribution_limit: max assignments per agent per window (default 5)
|
||||
// - fair_distribution_window: time window in seconds for rate limiting (default 300 = 5 min)
|
||||
type AssignmentPolicy struct {
|
||||
model.Base
|
||||
AccountID uint `gorm:"index;not null" json:"account_id"`
|
||||
Policy AssignmentPolicyType `gorm:"size:50;default:round_robin" json:"policy"`
|
||||
FairDistributionLimit int `gorm:"default:5" json:"fair_distribution_limit"`
|
||||
FairDistributionWindow int `gorm:"default:300" json:"fair_distribution_window"` // seconds
|
||||
Active bool `gorm:"default:true" json:"active"`
|
||||
}
|
||||
|
||||
func (AssignmentPolicy) TableName() string { return "assignment_policies" }
|
||||
|
||||
// InboxAssignmentPolicy represents a per-inbox override of the global assignment policy.
|
||||
// Reference: Chatwoot InboxAssignmentPolicy model
|
||||
// - Each inbox can override the account-level policy
|
||||
// - If no inbox-specific policy exists, the account policy is used
|
||||
type InboxAssignmentPolicy struct {
|
||||
model.Base
|
||||
AccountID uint `gorm:"index;not null" json:"account_id"`
|
||||
InboxID uint `gorm:"uniqueIndex;not null" json:"inbox_id"`
|
||||
Policy AssignmentPolicyType `gorm:"size:50;default:round_robin" json:"policy"`
|
||||
FairDistributionLimit int `gorm:"default:5" json:"fair_distribution_limit"`
|
||||
FairDistributionWindow int `gorm:"default:300" json:"fair_distribution_window"` // seconds
|
||||
Active bool `gorm:"default:true" json:"active"`
|
||||
}
|
||||
|
||||
func (InboxAssignmentPolicy) TableName() string { return "inbox_assignment_policies" }
|
||||
|
||||
// EffectivePolicy returns the effective assignment policy for a given inbox.
|
||||
// If the inbox has a specific policy, it is used; otherwise the account policy is used.
|
||||
func EffectivePolicy(accountPolicy *AssignmentPolicy, inboxPolicy *InboxAssignmentPolicy) AssignmentPolicyType {
|
||||
if inboxPolicy != nil && inboxPolicy.Active {
|
||||
return inboxPolicy.Policy
|
||||
}
|
||||
if accountPolicy != nil && accountPolicy.Active {
|
||||
return accountPolicy.Policy
|
||||
}
|
||||
return PolicyRoundRobin // default
|
||||
}
|
||||
|
||||
// EffectiveLimit returns the effective fair distribution limit for a given inbox.
|
||||
func EffectiveLimit(accountPolicy *AssignmentPolicy, inboxPolicy *InboxAssignmentPolicy) int {
|
||||
if inboxPolicy != nil && inboxPolicy.Active && inboxPolicy.FairDistributionLimit > 0 {
|
||||
return inboxPolicy.FairDistributionLimit
|
||||
}
|
||||
if accountPolicy != nil && accountPolicy.FairDistributionLimit > 0 {
|
||||
return accountPolicy.FairDistributionLimit
|
||||
}
|
||||
return 5 // default
|
||||
}
|
||||
|
||||
// EffectiveWindow returns the effective fair distribution window (in seconds).
|
||||
func EffectiveWindow(accountPolicy *AssignmentPolicy, inboxPolicy *InboxAssignmentPolicy) int {
|
||||
if inboxPolicy != nil && inboxPolicy.Active && inboxPolicy.FairDistributionWindow > 0 {
|
||||
return inboxPolicy.FairDistributionWindow
|
||||
}
|
||||
if accountPolicy != nil && accountPolicy.FairDistributionWindow > 0 {
|
||||
return accountPolicy.FairDistributionWindow
|
||||
}
|
||||
return 300 // default (5 minutes)
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,454 @@
|
||||
package autoassignment
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// setupTestRedis 创建一个基于 miniredis 的测试 Redis 环境。
|
||||
// 返回 miniredis 服务实例和 redis.Client。
|
||||
// miniredis.RunT 会自动在 t.Cleanup 中关闭服务。
|
||||
func setupTestRedis(t *testing.T) (*miniredis.Miniredis, *redis.Client) {
|
||||
t.Helper()
|
||||
|
||||
mr := miniredis.RunT(t)
|
||||
rdb := redis.NewClient(&redis.Options{
|
||||
Addr: mr.Addr(),
|
||||
})
|
||||
|
||||
t.Cleanup(func() {
|
||||
rdb.Close()
|
||||
})
|
||||
|
||||
return mr, rdb
|
||||
}
|
||||
|
||||
// --- IsAllowed 测试 ---
|
||||
|
||||
// TestIsAllowed_首次请求允许 验证:新窗口中没有任何计数时,IsAllowed 应返回 true。
|
||||
func TestIsAllowed_首次请求允许(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
limit := 3
|
||||
windowSeconds := 300
|
||||
|
||||
// 新窗口中没有任何计数,应该允许
|
||||
assert.True(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
||||
"首次请求应被允许")
|
||||
}
|
||||
|
||||
// TestIsAllowed_达到限制后拒绝 验证:当计数达到 limit 后,IsAllowed 应返回 false。
|
||||
func TestIsAllowed_达到限制后拒绝(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
limit := 3
|
||||
windowSeconds := 300
|
||||
|
||||
// 先 Increment 3 次,达到限制
|
||||
for i := 0; i < limit; i++ {
|
||||
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// 达到限制后,应该拒绝
|
||||
assert.False(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
||||
"达到限制后应被拒绝")
|
||||
}
|
||||
|
||||
// TestIsAllowed_未达到限制时允许 验证:计数小于 limit 时,IsAllowed 应返回 true。
|
||||
func TestIsAllowed_未达到限制时允许(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
limit := 3
|
||||
windowSeconds := 300
|
||||
|
||||
// Increment 2 次,未达到限制
|
||||
for i := 0; i < limit-1; i++ {
|
||||
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// 未达到限制,应该允许
|
||||
assert.True(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
||||
"未达到限制时应被允许")
|
||||
}
|
||||
|
||||
// TestIsAllowed_Limit为零时总是允许 验证:limit <= 0 时,IsAllowed 应返回 true(不限制)。
|
||||
func TestIsAllowed_Limit为零时总是允许(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// limit = 0,应总是允许
|
||||
assert.True(t, rl.IsAllowed(ctx, 1, 10, 0, 300),
|
||||
"limit 为 0 时应总是允许")
|
||||
|
||||
// limit = -1,应总是允许
|
||||
assert.True(t, rl.IsAllowed(ctx, 1, 10, -1, 300),
|
||||
"limit 为负数时应总是允许")
|
||||
}
|
||||
|
||||
// TestIsAllowed_不同Agent独立计数 验证:不同 agent 的限制独立计算。
|
||||
func TestIsAllowed_不同Agent独立计数(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
limit := 2
|
||||
windowSeconds := 300
|
||||
|
||||
// Agent A 达到限制
|
||||
for i := 0; i < limit; i++ {
|
||||
err := rl.Increment(ctx, inboxID, 10, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
assert.False(t, rl.IsAllowed(ctx, inboxID, 10, limit, windowSeconds),
|
||||
"Agent A 达到限制后应被拒绝")
|
||||
|
||||
// Agent B 未达到限制,应允许
|
||||
assert.True(t, rl.IsAllowed(ctx, inboxID, 20, limit, windowSeconds),
|
||||
"Agent B 应独立计算,首次请求应被允许")
|
||||
}
|
||||
|
||||
// TestIsAllowed_不同Inbox独立计数 验证:不同 inbox 的限制独立计算。
|
||||
func TestIsAllowed_不同Inbox独立计数(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
agentID := uint(10)
|
||||
limit := 2
|
||||
windowSeconds := 300
|
||||
|
||||
// Inbox 1 达到限制
|
||||
for i := 0; i < limit; i++ {
|
||||
err := rl.Increment(ctx, 1, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
assert.False(t, rl.IsAllowed(ctx, 1, agentID, limit, windowSeconds),
|
||||
"Inbox 1 达到限制后应被拒绝")
|
||||
|
||||
// Inbox 2 未达到限制,应允许
|
||||
assert.True(t, rl.IsAllowed(ctx, 2, agentID, limit, windowSeconds),
|
||||
"Inbox 2 应独立计算,首次请求应被允许")
|
||||
}
|
||||
|
||||
// --- Increment 测试 ---
|
||||
|
||||
// TestIncrement_正常计数增加 验证:每次 Increment 后计数应递增 1。
|
||||
func TestIncrement_正常计数增加(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
windowSeconds := 300
|
||||
|
||||
// 验证初始计数为 0
|
||||
count, err := rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, count)
|
||||
|
||||
// Increment 1 次
|
||||
err = rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
|
||||
count, err = rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, count)
|
||||
|
||||
// Increment 第 2 次
|
||||
err = rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
|
||||
count, err = rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 2, count)
|
||||
}
|
||||
|
||||
// TestIncrement_设置TTL 验证:Increment 后 key 应有正确的 TTL。
|
||||
func TestIncrement_设置TTL(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
windowSeconds := 300
|
||||
|
||||
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
|
||||
// 验证 key 存在且有 TTL
|
||||
key := rl.key(inboxID, agentID, windowSeconds)
|
||||
ttl, err := rdb.TTL(ctx, key).Result()
|
||||
require.NoError(t, err)
|
||||
|
||||
// TTL 应大于 0,且不超过 windowSeconds
|
||||
assert.Greater(t, ttl, time.Duration(0))
|
||||
assert.LessOrEqual(t, ttl, time.Duration(windowSeconds)*time.Second)
|
||||
}
|
||||
|
||||
// TestIncrement_重复Increment不覆盖TTL 验证:窗口内后续 Increment 不应重新设置 TTL。
|
||||
func TestIncrement_重复Increment不覆盖TTL(t *testing.T) {
|
||||
mr, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
windowSeconds := 60
|
||||
|
||||
// 第一次 Increment,设置 TTL
|
||||
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
|
||||
key := rl.key(inboxID, agentID, windowSeconds)
|
||||
ttl1, err := rdb.TTL(ctx, key).Result()
|
||||
require.NoError(t, err)
|
||||
|
||||
// 快进时间,使 TTL 减少
|
||||
mr.FastForward(30 * time.Second)
|
||||
|
||||
ttlAfterForward, err := rdb.TTL(ctx, key).Result()
|
||||
require.NoError(t, err)
|
||||
assert.Less(t, ttlAfterForward, ttl1, "快进后 TTL 应减少")
|
||||
|
||||
// 第二次 Increment,不应重新设置 TTL(TTL 应继续减少而非恢复)
|
||||
err = rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
|
||||
ttl2, err := rdb.TTL(ctx, key).Result()
|
||||
require.NoError(t, err)
|
||||
assert.Less(t, ttl2, ttl1, "后续 Increment 不应重新设置 TTL")
|
||||
}
|
||||
|
||||
// --- Redis 错误降级测试 ---
|
||||
|
||||
// TestIsAllowed_Redis错误时允许 验证:当 Redis 出错时,IsAllowed 应返回 true(降级允许)。
|
||||
func TestIsAllowed_Redis错误时允许(t *testing.T) {
|
||||
mr, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
limit := 1
|
||||
windowSeconds := 300
|
||||
|
||||
// 先 Increment 1 次,达到限制
|
||||
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
assert.False(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
||||
"正常情况下达到限制应拒绝")
|
||||
|
||||
// 关闭 miniredis,模拟 Redis 连接错误
|
||||
mr.Close()
|
||||
|
||||
// Redis 错误时应降级允许
|
||||
assert.True(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
||||
"Redis 错误时应降级允许请求")
|
||||
}
|
||||
|
||||
// TestIncrement_Redis错误返回错误 验证:当 Redis 出错时,Increment 应返回错误。
|
||||
func TestIncrement_Redis错误返回错误(t *testing.T) {
|
||||
mr, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
windowSeconds := 300
|
||||
|
||||
// 关闭 miniredis,模拟 Redis 连接错误
|
||||
mr.Close()
|
||||
|
||||
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
assert.Error(t, err, "Redis 错误时 Increment 应返回错误")
|
||||
}
|
||||
|
||||
// --- 窗口过期后限制重置测试 ---
|
||||
|
||||
// TestIsAllowed_窗口过期后限制重置 验证:时间窗口过期后,计数应重置,请求应重新被允许。
|
||||
func TestIsAllowed_窗口过期后限制重置(t *testing.T) {
|
||||
mr, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
limit := 2
|
||||
windowSeconds := 60
|
||||
|
||||
// Increment 达到限制
|
||||
for i := 0; i < limit; i++ {
|
||||
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
assert.False(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
||||
"达到限制后应被拒绝")
|
||||
|
||||
// 快进时间超过窗口时长,使 key 过期
|
||||
mr.FastForward(time.Duration(windowSeconds) * time.Second)
|
||||
|
||||
// 窗口过期后,应重新允许
|
||||
assert.True(t, rl.IsAllowed(ctx, inboxID, agentID, limit, windowSeconds),
|
||||
"窗口过期后应重新允许请求")
|
||||
}
|
||||
|
||||
// --- GetCount 测试 ---
|
||||
|
||||
// TestGetCount_无计数时返回零 验证:没有 Increment 时 GetCount 应返回 0。
|
||||
func TestGetCount_无计数时返回零(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
windowSeconds := 300
|
||||
|
||||
count, err := rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, count, "无计数时应返回 0")
|
||||
}
|
||||
|
||||
// TestGetCount_正确返回计数 验证:Increment 后 GetCount 应返回正确的计数。
|
||||
func TestGetCount_正确返回计数(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
agentID := uint(10)
|
||||
windowSeconds := 300
|
||||
|
||||
for i := 1; i <= 5; i++ {
|
||||
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
|
||||
count, err := rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, i, count)
|
||||
}
|
||||
}
|
||||
|
||||
// --- Reset 测试 ---
|
||||
|
||||
// TestReset_清除所有计数 验证:Reset 应清除指定 inbox 下所有 agent 的计数。
|
||||
func TestReset_清除所有计数(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
inboxID := uint(1)
|
||||
windowSeconds := 300
|
||||
|
||||
// 为多个 agent Increment
|
||||
for agentID := uint(10); agentID <= 13; agentID++ {
|
||||
err := rl.Increment(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// 验证计数存在
|
||||
for agentID := uint(10); agentID <= 13; agentID++ {
|
||||
count, err := rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, count)
|
||||
}
|
||||
|
||||
// Reset inbox
|
||||
err := rl.Reset(ctx, inboxID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// 验证计数已被清除
|
||||
for agentID := uint(10); agentID <= 13; agentID++ {
|
||||
count, err := rl.GetCount(ctx, inboxID, agentID, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, count, "Reset 后计数应被清除")
|
||||
}
|
||||
}
|
||||
|
||||
// TestReset_不影响其他Inbox 验证:Reset 一个 inbox 不应影响其他 inbox 的计数。
|
||||
func TestReset_不影响其他Inbox(t *testing.T) {
|
||||
_, rdb := setupTestRedis(t)
|
||||
rl := NewRateLimiter(rdb)
|
||||
|
||||
ctx := context.Background()
|
||||
windowSeconds := 300
|
||||
|
||||
// Inbox 1 和 Inbox 2 都有计数
|
||||
err := rl.Increment(ctx, 1, 10, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
err = rl.Increment(ctx, 2, 10, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Reset Inbox 1
|
||||
err = rl.Reset(ctx, 1)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Inbox 1 计数被清除
|
||||
count, err := rl.GetCount(ctx, 1, 10, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 0, count)
|
||||
|
||||
// Inbox 2 计数不受影响
|
||||
count, err = rl.GetCount(ctx, 2, 10, windowSeconds)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 1, count, "其他 inbox 的计数不应受影响")
|
||||
}
|
||||
|
||||
// --- ParseAgentID 测试 ---
|
||||
|
||||
// TestParseAgentID_正常解析 验证:ParseAgentID 应正确解析合法字符串。
|
||||
func TestParseAgentID_正常解析(t *testing.T) {
|
||||
id, err := ParseAgentID("42")
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, uint(42), id)
|
||||
}
|
||||
|
||||
// TestParseAgentID_非法字符串 验证:ParseAgentID 应对非法字符串返回错误。
|
||||
func TestParseAgentID_非法字符串(t *testing.T) {
|
||||
_, err := ParseAgentID("abc")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
// TestParseAgentID_空字符串 验证:ParseAgentID 应对空字符串返回错误。
|
||||
func TestParseAgentID_空字符串(t *testing.T) {
|
||||
_, err := ParseAgentID("")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
// TestParseAgentID_负数 验证:ParseAgentID 应对负数字符串返回错误。
|
||||
func TestParseAgentID_负数(t *testing.T) {
|
||||
_, err := ParseAgentID("-1")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
// TestParseAgentID_零 验证:ParseAgentID 应正确解析 "0"。
|
||||
func TestParseAgentID_零(t *testing.T) {
|
||||
id, err := ParseAgentID("0")
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, uint(0), id)
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
package autoassignment
|
||||
|
||||
// RoundRobinSelector implements Redis-backed round-robin agent selection.
|
||||
//
|
||||
// Reference: Chatwoot InboxRoundRobinService
|
||||
// - Uses Redis to maintain a per-inbox queue of agent IDs
|
||||
// - Push/pop/rotate agents through the queue
|
||||
// - Reset queue if members changed (SyncQueue method)
|
||||
// - RoundRobinSelector wraps InboxRoundRobinService for single agent selection
|
||||
//
|
||||
// In gochat, we use Redis lists to implement the round-robin queue:
|
||||
// - Key format: gochat:round_robin:{inboxID}
|
||||
// - On Next(): pop from left, push to right (rotate)
|
||||
// - On SyncQueue(): rebuild the list from current inbox members
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
applogger "github.com/gochat/gochat/pkg/logger"
|
||||
)
|
||||
|
||||
const (
|
||||
// roundRobinKeyPrefix is the Redis key prefix for round-robin queues.
|
||||
roundRobinKeyPrefix = "gochat:round_robin:"
|
||||
)
|
||||
|
||||
// RoundRobinSelector implements Redis-backed round-robin agent selection
|
||||
// for a given inbox.
|
||||
type RoundRobinSelector struct {
|
||||
redis *redis.Client
|
||||
}
|
||||
|
||||
// NewRoundRobinSelector creates a new RoundRobinSelector.
|
||||
func NewRoundRobinSelector(rdb *redis.Client) *RoundRobinSelector {
|
||||
return &RoundRobinSelector{redis: rdb}
|
||||
}
|
||||
|
||||
// key returns the Redis key for the given inbox's round-robin queue.
|
||||
func (rr *RoundRobinSelector) key(inboxID uint) string {
|
||||
return fmt.Sprintf("%s%d", roundRobinKeyPrefix, inboxID)
|
||||
}
|
||||
|
||||
// Next returns the next agent ID from the round-robin queue.
|
||||
// It pops from the left of the list and pushes to the right (rotate),
|
||||
// so each agent gets a turn before any agent gets a second turn.
|
||||
//
|
||||
// Returns 0 if the queue is empty.
|
||||
func (rr *RoundRobinSelector) Next(ctx context.Context, inboxID uint) (uint, error) {
|
||||
key := rr.key(inboxID)
|
||||
|
||||
// Pop from left, push to right (rotate)
|
||||
result, err := rr.redis.LPop(ctx, key).Result()
|
||||
if err == redis.Nil {
|
||||
// Queue is empty
|
||||
return 0, nil
|
||||
}
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("redis LPop: %w", err)
|
||||
}
|
||||
|
||||
// Push the same agent to the right (rotate back)
|
||||
if err := rr.redis.RPush(ctx, key, result).Err(); err != nil {
|
||||
return 0, fmt.Errorf("redis RPush: %w", err)
|
||||
}
|
||||
|
||||
agentID, err := strconv.ParseUint(result, 10, 32)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("parse agent ID: %w", err)
|
||||
}
|
||||
|
||||
return uint(agentID), nil
|
||||
}
|
||||
|
||||
// AvailableAgent selects the next agent from the round-robin queue that is also
|
||||
// in the allowedAgentIDs set.
|
||||
// 1:1 Chatwoot: InboxRoundRobinService#available_agent(allowed_agent_ids:)
|
||||
// Logic: queue.intersection(allowed_agent_ids).pop → pop_push_to_queue
|
||||
func (rr *RoundRobinSelector) AvailableAgent(ctx context.Context, inboxID uint, allowedAgentIDs []uint) (uint, error) {
|
||||
if len(allowedAgentIDs) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
key := rr.key(inboxID)
|
||||
|
||||
// Get current queue
|
||||
queue, err := rr.redis.LRange(ctx, key, 0, -1).Result()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("redis LRange: %w", err)
|
||||
}
|
||||
if len(queue) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Build allowed set for intersection
|
||||
allowedSet := make(map[string]bool, len(allowedAgentIDs))
|
||||
for _, id := range allowedAgentIDs {
|
||||
allowedSet[strconv.FormatUint(uint64(id), 10)] = true
|
||||
}
|
||||
|
||||
// Find last agent in queue that is in allowed set (Chatwoot: intersection.pop)
|
||||
var selected string
|
||||
for i := len(queue) - 1; i >= 0; i-- {
|
||||
if allowedSet[queue[i]] {
|
||||
selected = queue[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if selected == "" {
|
||||
return 0, nil // no eligible agent
|
||||
}
|
||||
|
||||
// pop_push_to_queue: remove from current position, push to end
|
||||
rr.redis.LRem(ctx, key, 1, selected)
|
||||
rr.redis.RPush(ctx, key, selected)
|
||||
|
||||
agentID, err := strconv.ParseUint(selected, 10, 32)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("parse agent ID: %w", err)
|
||||
}
|
||||
|
||||
return uint(agentID), nil
|
||||
}
|
||||
|
||||
// SyncQueue rebuilds the round-robin queue for the given inbox.
|
||||
// It compares the current queue with the provided agent list,
|
||||
// and resets the queue if the members have changed.
|
||||
//
|
||||
// Reference: Chatwoot InboxRoundRobinService.reset_round_robin_for_inbox
|
||||
// - Called when inbox members change (agent added/removed)
|
||||
// - Ensures only current members are in the rotation
|
||||
func (rr *RoundRobinSelector) SyncQueue(ctx context.Context, inboxID uint, agents []uint) {
|
||||
key := rr.key(inboxID)
|
||||
|
||||
// Check current queue
|
||||
current, err := rr.redis.LRange(ctx, key, 0, -1).Result()
|
||||
if err != nil && err != redis.Nil {
|
||||
applogger.L().Warnf("redis LRange error for inbox %d: %v", inboxID, err)
|
||||
}
|
||||
|
||||
// Build expected queue
|
||||
expected := make([]string, len(agents))
|
||||
for i, agentID := range agents {
|
||||
expected[i] = strconv.FormatUint(uint64(agentID), 10)
|
||||
}
|
||||
|
||||
// Check if queue needs reset
|
||||
needsReset := len(current) != len(expected)
|
||||
if !needsReset {
|
||||
// Check if all expected agents are in current queue
|
||||
currentSet := make(map[string]bool)
|
||||
for _, id := range current {
|
||||
currentSet[id] = true
|
||||
}
|
||||
for _, id := range expected {
|
||||
if !currentSet[id] {
|
||||
needsReset = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if needsReset {
|
||||
// Delete old queue and create new one
|
||||
rr.redis.Del(ctx, key)
|
||||
if len(expected) > 0 {
|
||||
values := make([]interface{}, len(expected))
|
||||
for i, v := range expected {
|
||||
values[i] = v
|
||||
}
|
||||
if err := rr.redis.RPush(ctx, key, values...).Err(); err != nil {
|
||||
applogger.L().Warnf("redis RPush error for inbox %d: %v", inboxID, err)
|
||||
}
|
||||
}
|
||||
applogger.L().Infof("round-robin queue synced for inbox %d (%d agents)", inboxID, len(agents))
|
||||
}
|
||||
}
|
||||
|
||||
// ClearQueue removes the round-robin queue for the given inbox.
|
||||
func (rr *RoundRobinSelector) ClearQueue(ctx context.Context, inboxID uint) error {
|
||||
return rr.redis.Del(ctx, rr.key(inboxID)).Err()
|
||||
}
|
||||
|
||||
// QueueLength returns the number of agents in the round-robin queue.
|
||||
func (rr *RoundRobinSelector) QueueLength(ctx context.Context, inboxID uint) (int64, error) {
|
||||
return rr.redis.LLen(ctx, rr.key(inboxID)).Result()
|
||||
}
|
||||
@@ -0,0 +1,311 @@
|
||||
package autoassignment
|
||||
|
||||
// AssignmentService provides the core auto-assignment logic.
|
||||
//
|
||||
// Reference: Chatwoot AutoAssignment pattern
|
||||
// - AssignmentService: bulk assigns conversations in an inbox.
|
||||
// Finds unassigned open conversations, filters by team and rate limit,
|
||||
// then uses RoundRobinSelector.
|
||||
// - AgentAssignmentService: assigns a specific conversation.
|
||||
// Intersects online agents with allowed agents, then round-robin.
|
||||
//
|
||||
// In gochat, we combine both into a single AssignmentService that handles:
|
||||
// - Bulk assignment of unassigned conversations in an inbox
|
||||
// - Single conversation assignment for a specific conversation
|
||||
// - Rate limiting per agent via RateLimiter
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/gochat/gochat/internal/model"
|
||||
"github.com/redis/go-redis/v9"
|
||||
"gorm.io/gorm"
|
||||
applogger "github.com/gochat/gochat/pkg/logger"
|
||||
)
|
||||
|
||||
// AssignmentService handles auto-assignment of conversations to agents.
|
||||
type AssignmentService struct {
|
||||
db *gorm.DB
|
||||
redis *redis.Client
|
||||
roundRobin *RoundRobinSelector
|
||||
lowestLoad *LowestLoadSelector
|
||||
rateLimiter *RateLimiter
|
||||
}
|
||||
|
||||
// NewAssignmentService creates a new AssignmentService.
|
||||
func NewAssignmentService(db *gorm.DB, rdb *redis.Client) *AssignmentService {
|
||||
rr := NewRoundRobinSelector(rdb)
|
||||
ll := NewLowestLoadSelector(db)
|
||||
rl := NewRateLimiter(rdb)
|
||||
return &AssignmentService{
|
||||
db: db,
|
||||
redis: rdb,
|
||||
roundRobin: rr,
|
||||
lowestLoad: ll,
|
||||
rateLimiter: rl,
|
||||
}
|
||||
}
|
||||
|
||||
// AssignUnassignedConversations bulk assigns all unassigned open conversations
|
||||
// in the given inbox.
|
||||
// Reference: Chatwoot AssignmentService.assign_unassigned_conversations
|
||||
//
|
||||
// Steps:
|
||||
// 1. Find all unassigned open conversations for the inbox
|
||||
// 2. Get the inbox's assignment policy and rate limits
|
||||
// 3. Get the list of eligible agents (online + members of the inbox)
|
||||
// 4. For each conversation, select an agent via round-robin + rate limit
|
||||
// 5. Assign the conversation to the selected agent
|
||||
func (s *AssignmentService) AssignUnassignedConversations(ctx context.Context, inboxID uint, accountID uint) ([]uint, error) {
|
||||
// Step 1: Check if auto-assignment is enabled for this inbox
|
||||
inbox, err := s.getInbox(ctx, inboxID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get inbox: %w", err)
|
||||
}
|
||||
if !inbox.EnableAutoAssignment {
|
||||
applogger.L().Infof("auto-assignment disabled for inbox %d", inboxID)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Step 2: Find unassigned open conversations
|
||||
conversations, err := s.findUnassignedConversations(ctx, inboxID, accountID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("find unassigned: %w", err)
|
||||
}
|
||||
if len(conversations) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Step 3: Get eligible agents
|
||||
agents, err := s.getEligibleAgents(ctx, inboxID, accountID)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get eligible agents: %w", err)
|
||||
}
|
||||
if len(agents) == 0 {
|
||||
applogger.L().Infof("no eligible agents for inbox %d", inboxID)
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// Step 4: Get assignment policy and limits
|
||||
accountPolicy := s.getAccountPolicy(ctx, accountID)
|
||||
inboxPolicy := s.getInboxPolicy(ctx, inboxID)
|
||||
policy := EffectivePolicy(accountPolicy, inboxPolicy)
|
||||
limit := EffectiveLimit(accountPolicy, inboxPolicy)
|
||||
window := EffectiveWindow(accountPolicy, inboxPolicy)
|
||||
|
||||
// Step 5: Reset round-robin queue if members changed
|
||||
s.roundRobin.SyncQueue(ctx, inboxID, agents)
|
||||
|
||||
// Step 6: Assign each conversation
|
||||
assignedIDs := make([]uint, 0)
|
||||
for _, conv := range conversations {
|
||||
agentID, err := s.selectAgent(ctx, inboxID, agents, policy, limit, window)
|
||||
if err != nil {
|
||||
applogger.L().Warnf("failed to select agent for conversation %d: %v", conv.ID, err)
|
||||
continue
|
||||
}
|
||||
if agentID == 0 {
|
||||
// All agents hit rate limit, skip
|
||||
continue
|
||||
}
|
||||
|
||||
if err := s.assignConversation(ctx, conv.ID, agentID); err != nil {
|
||||
applogger.L().Warnf("failed to assign conversation %d to agent %d: %v", conv.ID, agentID, err)
|
||||
continue
|
||||
}
|
||||
|
||||
// Track rate limit
|
||||
s.rateLimiter.Increment(ctx, inboxID, agentID, window)
|
||||
|
||||
assignedIDs = append(assignedIDs, conv.ID)
|
||||
}
|
||||
|
||||
applogger.L().Infof("auto-assigned %d conversations in inbox %d", len(assignedIDs), inboxID)
|
||||
return assignedIDs, nil
|
||||
}
|
||||
|
||||
// AssignConversation assigns a specific conversation to an agent.
|
||||
// Reference: Chatwoot AgentAssignmentService.assign_conversation
|
||||
//
|
||||
// Steps:
|
||||
// 1. Get eligible agents for the conversation's inbox
|
||||
// 2. Select an agent via round-robin + rate limit
|
||||
// 3. Assign the conversation
|
||||
func (s *AssignmentService) AssignConversation(ctx context.Context, conversationID uint, inboxID uint, accountID uint) (uint, error) {
|
||||
// Check if auto-assignment is enabled
|
||||
inbox, err := s.getInbox(ctx, inboxID)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("get inbox: %w", err)
|
||||
}
|
||||
if !inbox.EnableAutoAssignment {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Get eligible agents
|
||||
agents, err := s.getEligibleAgents(ctx, inboxID, accountID)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("get eligible agents: %w", err)
|
||||
}
|
||||
if len(agents) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// Get policy and limits
|
||||
accountPolicy := s.getAccountPolicy(ctx, accountID)
|
||||
inboxPolicy := s.getInboxPolicy(ctx, inboxID)
|
||||
policy := EffectivePolicy(accountPolicy, inboxPolicy)
|
||||
limit := EffectiveLimit(accountPolicy, inboxPolicy)
|
||||
window := EffectiveWindow(accountPolicy, inboxPolicy)
|
||||
|
||||
// Sync round-robin queue
|
||||
s.roundRobin.SyncQueue(ctx, inboxID, agents)
|
||||
|
||||
// Select agent
|
||||
agentID, err := s.selectAgent(ctx, inboxID, agents, policy, limit, window)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("select agent: %w", err)
|
||||
}
|
||||
if agentID == 0 {
|
||||
return 0, nil // all agents at rate limit
|
||||
}
|
||||
|
||||
// Assign conversation
|
||||
if err := s.assignConversation(ctx, conversationID, agentID); err != nil {
|
||||
return 0, fmt.Errorf("assign conversation: %w", err)
|
||||
}
|
||||
|
||||
// Track rate limit
|
||||
s.rateLimiter.Increment(ctx, inboxID, agentID, window)
|
||||
|
||||
return agentID, nil
|
||||
}
|
||||
|
||||
// selectAgent picks an agent for assignment using the selected policy
|
||||
// and rate limiting.
|
||||
func (s *AssignmentService) selectAgent(ctx context.Context, inboxID uint, agents []uint, policy AssignmentPolicyType, limit int, window int) (uint, error) {
|
||||
switch policy {
|
||||
case PolicyRoundRobin:
|
||||
return s.selectRoundRobin(ctx, inboxID, agents, limit, window)
|
||||
case PolicyLowestLoad:
|
||||
return s.selectLowestLoad(ctx, inboxID, agents, limit, window)
|
||||
case PolicyLongestWaiting:
|
||||
return s.selectLongestWaiting(ctx, inboxID, agents, limit, window)
|
||||
default:
|
||||
return s.selectRoundRobin(ctx, inboxID, agents, limit, window)
|
||||
}
|
||||
}
|
||||
|
||||
// selectRoundRobin selects an agent via round-robin, respecting rate limits.
|
||||
func (s *AssignmentService) selectRoundRobin(ctx context.Context, inboxID uint, agents []uint, limit int, window int) (uint, error) {
|
||||
// Try each agent in round-robin order
|
||||
for i := 0; i < len(agents); i++ {
|
||||
agentID, err := s.roundRobin.Next(ctx, inboxID)
|
||||
if err != nil || agentID == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
// Check rate limit
|
||||
if s.rateLimiter.IsAllowed(ctx, inboxID, agentID, limit, window) {
|
||||
return agentID, nil
|
||||
}
|
||||
// Agent is at rate limit, skip and try next
|
||||
}
|
||||
|
||||
// All agents at rate limit
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
// selectLongestWaiting selects the agent with the longest idle time.
|
||||
// This is a simplified implementation — a full version would track
|
||||
// idle time in Redis. For now, we fall back to round-robin.
|
||||
func (s *AssignmentService) selectLongestWaiting(ctx context.Context, inboxID uint, agents []uint, limit int, window int) (uint, error) {
|
||||
// TODO: Implement longest-waiting tracking in Redis
|
||||
// For now, delegate to round-robin with rate limit checking
|
||||
return s.selectRoundRobin(ctx, inboxID, agents, limit, window)
|
||||
}
|
||||
|
||||
// selectLowestLoad selects the agent with the fewest open conversations,
|
||||
// respecting rate limits.
|
||||
func (s *AssignmentService) selectLowestLoad(ctx context.Context, inboxID uint, agents []uint, limit int, window int) (uint, error) {
|
||||
agentID, err := s.lowestLoad.Select(ctx, inboxID, agents)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("lowest load select: %w", err)
|
||||
}
|
||||
|
||||
// Check rate limit for the selected agent
|
||||
if s.rateLimiter.IsAllowed(ctx, inboxID, agentID, limit, window) {
|
||||
return agentID, nil
|
||||
}
|
||||
|
||||
// Selected agent is at rate limit — fall back to round-robin
|
||||
applogger.L().Debugf("lowest_load: agent %d at rate limit, falling back to round-robin", agentID)
|
||||
return s.selectRoundRobin(ctx, inboxID, agents, limit, window)
|
||||
}
|
||||
|
||||
// findUnassignedConversations returns all open conversations in the inbox
|
||||
// that have no assignee.
|
||||
func (s *AssignmentService) findUnassignedConversations(ctx context.Context, inboxID uint, accountID uint) ([]model.Conversation, error) {
|
||||
var conversations []model.Conversation
|
||||
err := s.db.WithContext(ctx).
|
||||
Where("inbox_id = ? AND account_id = ? AND status = ? AND assignee_id IS NULL",
|
||||
inboxID, accountID, model.ConversationStatusOpen).
|
||||
Find(&conversations).Error
|
||||
return conversations, err
|
||||
}
|
||||
|
||||
// getEligibleAgents returns agents that are:
|
||||
// - Members of the inbox (via inbox_members table)
|
||||
// - Online/available (via user available field)
|
||||
func (s *AssignmentService) getEligibleAgents(ctx context.Context, inboxID uint, accountID uint) ([]uint, error) {
|
||||
var agentIDs []uint
|
||||
err := s.db.WithContext(ctx).
|
||||
Table("inbox_members").
|
||||
Select("inbox_members.user_id").
|
||||
Joins("JOIN users ON users.id = inbox_members.user_id").
|
||||
Where("inbox_members.inbox_id = ? AND users.available = ? AND users.active = ?",
|
||||
inboxID, true, true).
|
||||
Pluck("inbox_members.user_id", &agentIDs).Error
|
||||
return agentIDs, err
|
||||
}
|
||||
|
||||
// getInbox fetches the inbox by ID.
|
||||
func (s *AssignmentService) getInbox(ctx context.Context, inboxID uint) (*model.Inbox, error) {
|
||||
var inbox model.Inbox
|
||||
if err := s.db.WithContext(ctx).First(&inbox, inboxID).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &inbox, nil
|
||||
}
|
||||
|
||||
// getAccountPolicy returns the assignment policy for an account.
|
||||
// Returns nil if no policy is configured.
|
||||
func (s *AssignmentService) getAccountPolicy(ctx context.Context, accountID uint) *AssignmentPolicy {
|
||||
var policy AssignmentPolicy
|
||||
if err := s.db.WithContext(ctx).
|
||||
Where("account_id = ? AND active = ?", accountID, true).
|
||||
First(&policy).Error; err != nil {
|
||||
return nil // no policy configured
|
||||
}
|
||||
return &policy
|
||||
}
|
||||
|
||||
// getInboxPolicy returns the assignment policy override for an inbox.
|
||||
// Returns nil if no inbox-specific policy is configured.
|
||||
func (s *AssignmentService) getInboxPolicy(ctx context.Context, inboxID uint) *InboxAssignmentPolicy {
|
||||
var policy InboxAssignmentPolicy
|
||||
if err := s.db.WithContext(ctx).
|
||||
Where("inbox_id = ? AND active = ?", inboxID, true).
|
||||
First(&policy).Error; err != nil {
|
||||
return nil // no inbox-specific policy
|
||||
}
|
||||
return &policy
|
||||
}
|
||||
|
||||
// assignConversation sets the assignee_id on a conversation.
|
||||
func (s *AssignmentService) assignConversation(ctx context.Context, conversationID uint, agentID uint) error {
|
||||
return s.db.WithContext(ctx).
|
||||
Model(&model.Conversation{}).
|
||||
Where("id = ?", conversationID).
|
||||
Update("assignee_id", agentID).Error
|
||||
}
|
||||
Reference in New Issue
Block a user