ai-agent/internal/services/agent_tool_invocation_service.go

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
}