package services import ( "errors" "fmt" "strings" "time" "code.tczkiot.com/wlw/ai-agent/contract" "code.tczkiot.com/wlw/ai-agent/internal/models" "code.tczkiot.com/wlw/ai-agent/internal/repositories" "github.com/google/uuid" "github.com/mlogclub/simple/sqls" ) const ( agentToolInvocationStatusRunning = "running" agentToolInvocationStatusCompleted = "completed" agentToolInvocationStatusRetryableFailed = "retryable_failed" agentToolInvocationStatusUnknownOutcome = "unknown_outcome" agentToolInvocationStatusLegacyFailed = "failed" ) var AgentToolInvocationService = newAgentToolInvocationService() type AgentToolInvocationClaim struct { Item *models.AgentToolInvocation Completed bool Acquired bool UnknownOutcome bool Recovered bool } type agentToolInvocationService struct{} func newAgentToolInvocationService() *agentToolInvocationService { return &agentToolInvocationService{} } // Claim obtains the persistent idempotency boundary. A completed invocation // can be returned to callers; an in-flight invocation is never executed again. func (s *agentToolInvocationService) Claim(conversationID, aiAgentID int64, toolCode, idempotencyKey string) (*AgentToolInvocationClaim, error) { return s.claim(conversationID, aiAgentID, toolCode, idempotencyKey, time.Time{}) } // ClaimRecoverable is reserved for side-effect-free orchestration runs. A // stale running claim may be recovered after the caller's lease expires. It // must never be used for an external write operation whose outcome is unknown. func (s *agentToolInvocationService) ClaimRecoverable(conversationID, aiAgentID int64, toolCode, idempotencyKey string, staleBefore time.Time) (*AgentToolInvocationClaim, error) { return s.claim(conversationID, aiAgentID, toolCode, idempotencyKey, staleBefore) } func (s *agentToolInvocationService) claim(conversationID, aiAgentID int64, toolCode, idempotencyKey string, staleBefore time.Time) (*AgentToolInvocationClaim, error) { toolCode = strings.TrimSpace(toolCode) idempotencyKey = strings.TrimSpace(idempotencyKey) if conversationID <= 0 || toolCode == "" || idempotencyKey == "" { return nil, nil } if item := repositories.AgentToolInvocationRepository.GetByIdempotencyKey(sqls.DB(), conversationID, toolCode, idempotencyKey); item != nil { if item.Status == agentToolInvocationStatusCompleted { return &AgentToolInvocationClaim{Item: item, Completed: true}, nil } if item.Status == agentToolInvocationStatusUnknownOutcome { return &AgentToolInvocationClaim{Item: item, UnknownOutcome: true}, nil } if item.Status == agentToolInvocationStatusRunning { if !staleBefore.IsZero() && item.UpdatedAt.Before(staleBefore) { leaseToken := uuid.NewString() values := map[string]any{"error_message": "", "result_data": leaseToken, "updated_at": time.Now()} acquired, err := repositories.AgentToolInvocationRepository.RecoverStaleRunning(sqls.DB(), item.ID, staleBefore, values) if err != nil { return nil, err } if acquired { item.ErrorMessage, item.ResultData, item.UpdatedAt = "", leaseToken, values["updated_at"].(time.Time) return &AgentToolInvocationClaim{Item: item, Acquired: true, Recovered: true}, nil } } return &AgentToolInvocationClaim{Item: item}, nil } if item.Status != agentToolInvocationStatusRetryableFailed && item.Status != agentToolInvocationStatusLegacyFailed { return &AgentToolInvocationClaim{Item: item, UnknownOutcome: true}, nil } leaseToken := uuid.NewString() values := map[string]any{"status": agentToolInvocationStatusRunning, "error_message": "", "result_data": leaseToken, "updated_at": time.Now()} acquired, err := repositories.AgentToolInvocationRepository.TransitionStatus(sqls.DB(), item.ID, item.Status, values) if err != nil { return nil, err } if !acquired { current := repositories.AgentToolInvocationRepository.GetByIdempotencyKey(sqls.DB(), conversationID, toolCode, idempotencyKey) return &AgentToolInvocationClaim{Item: current, Completed: current != nil && current.Status == agentToolInvocationStatusCompleted, UnknownOutcome: current != nil && current.Status == agentToolInvocationStatusUnknownOutcome}, nil } item.Status, item.ErrorMessage, item.ResultData = agentToolInvocationStatusRunning, "", leaseToken return &AgentToolInvocationClaim{Item: item, Acquired: true}, nil } item := &models.AgentToolInvocation{ConversationID: conversationID, AIAgentID: aiAgentID, ToolCode: toolCode, IdempotencyKey: idempotencyKey, Status: agentToolInvocationStatusRunning, ResultData: uuid.NewString()} if err := repositories.AgentToolInvocationRepository.Create(sqls.DB(), item); err != nil { // A concurrent caller may have created the unique invocation first. if existing := repositories.AgentToolInvocationRepository.GetByIdempotencyKey(sqls.DB(), conversationID, toolCode, idempotencyKey); existing != nil { return &AgentToolInvocationClaim{ Item: existing, Completed: existing.Status == agentToolInvocationStatusCompleted, UnknownOutcome: existing.Status == agentToolInvocationStatusUnknownOutcome, }, nil } return nil, err } return &AgentToolInvocationClaim{Item: item, Acquired: true}, nil } func (s *agentToolInvocationService) Complete(item *models.AgentToolInvocation, resultData string) error { if item == nil || item.ID <= 0 { return nil } updated, err := repositories.AgentToolInvocationRepository.TransitionLease(sqls.DB(), item.ID, item.ResultData, map[string]any{"status": agentToolInvocationStatusCompleted, "result_data": resultData, "error_message": "", "updated_at": time.Now()}) if err != nil { return err } if !updated { return fmt.Errorf("agent tool invocation lease lost: %d", item.ID) } return nil } func (s *agentToolInvocationService) FailRetryable(item *models.AgentToolInvocation, cause error) error { return s.failWithStatus(item, cause, agentToolInvocationStatusRetryableFailed) } func (s *agentToolInvocationService) MarkUnknownOutcome(item *models.AgentToolInvocation, cause error) error { return s.failWithStatus(item, cause, agentToolInvocationStatusUnknownOutcome) } // Fail is kept as a compatibility alias for failures known to have occurred // before side effects. New write paths should call the explicit method. func (s *agentToolInvocationService) Fail(item *models.AgentToolInvocation, cause error) error { return s.FailRetryable(item, cause) } func (s *agentToolInvocationService) failWithStatus(item *models.AgentToolInvocation, cause error, status string) error { if item == nil || item.ID <= 0 { return nil } message := "" if cause != nil { message = cause.Error() var actionErr *contract.BusinessActionError if errors.As(cause, &actionErr) && actionErr.Cause != nil { message = actionErr.Cause.Error() } } updated, err := repositories.AgentToolInvocationRepository.TransitionLease(sqls.DB(), item.ID, item.ResultData, map[string]any{"status": status, "error_message": message, "updated_at": time.Now()}) if err != nil { return err } if !updated { return fmt.Errorf("agent tool invocation lease lost: %d", item.ID) } return nil }