internal/task/task_wait_token_repository.go

Rama main · réplica de lectura del repositorio original.

package task

import (
	"sort"
	"strings"
	"sync"
	"time"
)

type InMemoryTaskWaitTokenRepository struct {
	mutex          sync.RWMutex
	taskWaitTokens map[string]TaskWaitToken
}

func NewInMemoryTaskWaitTokenRepository() *InMemoryTaskWaitTokenRepository {
	return &InMemoryTaskWaitTokenRepository{
		taskWaitTokens: map[string]TaskWaitToken{},
	}
}

func (repository *InMemoryTaskWaitTokenRepository) InsertTaskWaitToken(taskWaitToken TaskWaitToken) error {
	taskWaitToken = normalizeTaskWaitToken(taskWaitToken)
	if taskWaitToken.WaitID == "" {
		taskWaitToken.WaitID = NewIdentifier()
	}
	repository.mutex.Lock()
	repository.taskWaitTokens[taskWaitToken.WaitID] = taskWaitToken
	repository.mutex.Unlock()
	return nil
}

func (repository *InMemoryTaskWaitTokenRepository) FindOpenByWaitID(waitID string) (TaskWaitToken, bool, error) {
	return repository.findOpen(func(taskWaitToken TaskWaitToken) bool {
		return taskWaitToken.WaitID == strings.TrimSpace(waitID)
	})
}

func (repository *InMemoryTaskWaitTokenRepository) FindOpenByPersonConversationAndReplyTarget(personID string, platform string, conversationID string, replyTargetID string) (TaskWaitToken, bool, error) {
	return repository.findOpen(func(taskWaitToken TaskWaitToken) bool {
		return taskWaitTokenMatchesConversation(taskWaitToken, personID, platform, conversationID) && taskWaitToken.ReplyTargetID == strings.TrimSpace(replyTargetID)
	})
}

func (repository *InMemoryTaskWaitTokenRepository) FindOpenByPersonConversationAndThreadRoot(personID string, platform string, conversationID string, threadRootID string) (TaskWaitToken, bool, error) {
	return repository.findOpen(func(taskWaitToken TaskWaitToken) bool {
		return taskWaitTokenMatchesConversation(taskWaitToken, personID, platform, conversationID) && taskWaitToken.ThreadRootID == strings.TrimSpace(threadRootID)
	})
}

func (repository *InMemoryTaskWaitTokenRepository) FindOpenByPersonConversationAndDispatchID(personID string, platform string, conversationID string, dispatchID string) (TaskWaitToken, bool, error) {
	return repository.findOpen(func(taskWaitToken TaskWaitToken) bool {
		return taskWaitTokenMatchesConversation(taskWaitToken, personID, platform, conversationID) && taskWaitToken.DispatchID == strings.TrimSpace(dispatchID)
	})
}

func (repository *InMemoryTaskWaitTokenRepository) FindOpenByPersonTaskRunAndInteraction(personID string, taskRunID string, interactionID string) (TaskWaitToken, bool, error) {
	return repository.findOpen(func(taskWaitToken TaskWaitToken) bool {
		return taskWaitToken.PersonID == strings.TrimSpace(personID) && taskWaitToken.TaskRunID == strings.TrimSpace(taskRunID) && taskWaitToken.InteractionID == strings.TrimSpace(interactionID)
	})
}

func (repository *InMemoryTaskWaitTokenRepository) FindOpenByPersonAndConversation(personID string, platform string, conversationID string) ([]TaskWaitToken, error) {
	taskWaitTokens := []TaskWaitToken{}
	repository.mutex.RLock()
	for _, taskWaitToken := range repository.taskWaitTokens {
		if !taskWaitTokenIsOpen(taskWaitToken, time.Now().UTC()) {
			continue
		}
		if taskWaitTokenMatchesConversation(taskWaitToken, personID, platform, conversationID) {
			taskWaitTokens = append(taskWaitTokens, taskWaitToken)
		}
	}
	repository.mutex.RUnlock()
	sortTaskWaitTokens(taskWaitTokens)
	return taskWaitTokens, nil
}

func (repository *InMemoryTaskWaitTokenRepository) ResolveTaskWait(waitID string, resolvedAt time.Time) error {
	repository.mutex.Lock()
	defer repository.mutex.Unlock()
	taskWaitToken, isFound := repository.taskWaitTokens[strings.TrimSpace(waitID)]
	if !isFound {
		return nil
	}
	if resolvedAt.IsZero() {
		resolvedAt = time.Now().UTC()
	}
	taskWaitToken.State = "resolved"
	taskWaitToken.ResolvedAt = &resolvedAt
	repository.taskWaitTokens[taskWaitToken.WaitID] = taskWaitToken
	return nil
}

func (repository *InMemoryTaskWaitTokenRepository) ExpireOldTaskWaits(before time.Time) ([]string, error) {
	if before.IsZero() {
		before = time.Now().UTC()
	}
	return repository.expireMatchingTaskWaits(before, func(taskWaitToken TaskWaitToken) bool {
		return taskWaitToken.ExpiresAt.Before(before) || taskWaitToken.ExpiresAt.Equal(before)
	}), nil
}

func (repository *InMemoryTaskWaitTokenRepository) ExpireTaskWaitTokensForPerson(personID string, expiredAt time.Time) ([]string, error) {
	if expiredAt.IsZero() {
		expiredAt = time.Now().UTC()
	}
	trimmedPersonID := strings.TrimSpace(personID)
	return repository.expireMatchingTaskWaits(expiredAt, func(taskWaitToken TaskWaitToken) bool {
		return taskWaitToken.PersonID == trimmedPersonID
	}), nil
}

func (repository *InMemoryTaskWaitTokenRepository) findOpen(matches func(TaskWaitToken) bool) (TaskWaitToken, bool, error) {
	taskWaitTokens := []TaskWaitToken{}
	repository.mutex.RLock()
	for _, taskWaitToken := range repository.taskWaitTokens {
		if taskWaitTokenIsOpen(taskWaitToken, time.Now().UTC()) && matches(taskWaitToken) {
			taskWaitTokens = append(taskWaitTokens, taskWaitToken)
		}
	}
	repository.mutex.RUnlock()
	sortTaskWaitTokens(taskWaitTokens)
	if len(taskWaitTokens) == 0 {
		return TaskWaitToken{}, false, nil
	}
	return taskWaitTokens[0], true, nil
}

func (repository *InMemoryTaskWaitTokenRepository) expireMatchingTaskWaits(expiredAt time.Time, matches func(TaskWaitToken) bool) []string {
	taskRunIDs := []string{}
	repository.mutex.Lock()
	defer repository.mutex.Unlock()
	for waitID, taskWaitToken := range repository.taskWaitTokens {
		if taskWaitToken.State != "open" || !matches(taskWaitToken) {
			continue
		}
		taskWaitToken.State = "expired"
		taskWaitToken.ResolvedAt = &expiredAt
		repository.taskWaitTokens[waitID] = taskWaitToken
		taskRunIDs = append(taskRunIDs, taskWaitToken.TaskRunID)
	}
	sort.Strings(taskRunIDs)
	return taskRunIDs
}

func normalizeTaskWaitToken(taskWaitToken TaskWaitToken) TaskWaitToken {
	taskWaitToken.WaitID = strings.TrimSpace(taskWaitToken.WaitID)
	taskWaitToken.TaskRunID = strings.TrimSpace(taskWaitToken.TaskRunID)
	taskWaitToken.PersonID = strings.TrimSpace(taskWaitToken.PersonID)
	taskWaitToken.Platform = strings.TrimSpace(taskWaitToken.Platform)
	taskWaitToken.ConversationID = strings.TrimSpace(taskWaitToken.ConversationID)
	taskWaitToken.ReplyTargetID = strings.TrimSpace(taskWaitToken.ReplyTargetID)
	taskWaitToken.ThreadRootID = strings.TrimSpace(taskWaitToken.ThreadRootID)
	taskWaitToken.DispatchID = strings.TrimSpace(taskWaitToken.DispatchID)
	taskWaitToken.InteractionID = strings.TrimSpace(taskWaitToken.InteractionID)
	taskWaitToken.Kind = strings.TrimSpace(taskWaitToken.Kind)
	taskWaitToken.State = strings.TrimSpace(taskWaitToken.State)
	if taskWaitToken.State == "" {
		taskWaitToken.State = "open"
	}
	if taskWaitToken.CreatedAt.IsZero() {
		taskWaitToken.CreatedAt = time.Now().UTC()
	}
	return taskWaitToken
}

func taskWaitTokenMatchesConversation(taskWaitToken TaskWaitToken, personID string, platform string, conversationID string) bool {
	return taskWaitToken.PersonID == strings.TrimSpace(personID) &&
		taskWaitToken.Platform == strings.TrimSpace(platform) &&
		taskWaitToken.ConversationID == strings.TrimSpace(conversationID)
}

func taskWaitTokenIsOpen(taskWaitToken TaskWaitToken, now time.Time) bool {
	if taskWaitToken.State != "open" {
		return false
	}
	if taskWaitToken.ExpiresAt.IsZero() {
		return true
	}
	return taskWaitToken.ExpiresAt.After(now)
}

func sortTaskWaitTokens(taskWaitTokens []TaskWaitToken) {
	sort.SliceStable(taskWaitTokens, func(leftIndex int, rightIndex int) bool {
		left := taskWaitTokens[leftIndex]
		right := taskWaitTokens[rightIndex]
		if !left.CreatedAt.Equal(right.CreatedAt) {
			return left.CreatedAt.Before(right.CreatedAt)
		}
		return left.WaitID < right.WaitID
	})
}

var _ TaskWaitTokenRepository = (*InMemoryTaskWaitTokenRepository)(nil)