ai-agent/internal/services/agent_tool_invocation_service_test.go

package services

import (
	"strings"
	"testing"
	"time"

	"code.tczkiot.com/wlw/ai-agent/internal/models"

	"github.com/glebarez/sqlite"
	"github.com/mlogclub/simple/sqls"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
)

func TestAgentToolInvocationServiceReusesCompletedInvocation(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: "t_", SingularTable: true}})
	if err != nil {
		t.Fatalf("open sqlite: %v", err)
	}
	if err := db.AutoMigrate(&models.AgentToolInvocation{}); err != nil {
		t.Fatalf("auto migrate: %v", err)
	}
	sqls.SetDB(db)

	first, err := AgentToolInvocationService.Claim(10, 20, "graph/handoff_to_human", "message:30:node:handoff")
	if err != nil || first == nil || first.Item == nil || first.Completed {
		t.Fatalf("first claim = %#v, err=%v", first, err)
	}
	if err := AgentToolInvocationService.Complete(first.Item, `{"handoff":true}`); err != nil {
		t.Fatalf("complete invocation: %v", err)
	}
	second, err := AgentToolInvocationService.Claim(10, 20, "graph/handoff_to_human", "message:30:node:handoff")
	if err != nil || second == nil || !second.Completed || second.Item.ResultData != `{"handoff":true}` {
		t.Fatalf("second claim = %#v, err=%v", second, err)
	}
}

func TestAgentToolInvocationServiceAllowsExplicitRetryableFailure(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: "t_", SingularTable: true}})
	if err != nil {
		t.Fatalf("open sqlite: %v", err)
	}
	if err := db.AutoMigrate(&models.AgentToolInvocation{}); err != nil {
		t.Fatalf("auto migrate: %v", err)
	}
	sqls.SetDB(db)

	first, err := AgentToolInvocationService.Claim(11, 21, "graph/handoff_to_human", "message:31:node:handoff")
	if err != nil {
		t.Fatalf("first claim: %v", err)
	}
	if err := AgentToolInvocationService.FailRetryable(first.Item, errTestToolInvocation); err != nil {
		t.Fatalf("fail invocation: %v", err)
	}
	second, err := AgentToolInvocationService.Claim(11, 21, "graph/handoff_to_human", "message:31:node:handoff")
	if err != nil || second == nil || second.Completed || second.Item.Status != agentToolInvocationStatusRunning || second.Item.ErrorMessage != "" {
		t.Fatalf("retry claim = %#v, err=%v", second, err)
	}
}

func TestAgentToolInvocationServiceNeverReclaimsUnknownOutcome(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: "t_", SingularTable: true}})
	if err != nil {
		t.Fatalf("open sqlite: %v", err)
	}
	if err := db.AutoMigrate(&models.AgentToolInvocation{}); err != nil {
		t.Fatalf("auto migrate: %v", err)
	}
	sqls.SetDB(db)

	first, err := AgentToolInvocationService.Claim(12, 22, "business/order", "confirm-unknown")
	if err != nil {
		t.Fatalf("first claim: %v", err)
	}
	if err := AgentToolInvocationService.MarkUnknownOutcome(first.Item, errTestToolInvocation); err != nil {
		t.Fatalf("mark unknown outcome: %v", err)
	}
	second, err := AgentToolInvocationService.Claim(12, 22, "business/order", "confirm-unknown")
	if err != nil || second == nil || second.Acquired || !second.UnknownOutcome || second.Item.Status != agentToolInvocationStatusUnknownOutcome {
		t.Fatalf("unknown outcome claim = %#v, err=%v", second, err)
	}
}

func TestAgentToolInvocationServiceRecoversOnlyStaleOrchestrationRun(t *testing.T) {
	db, err := gorm.Open(sqlite.Open("file:"+strings.ReplaceAll(t.Name(), "/", "_")+"?mode=memory&cache=shared"), &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: "t_", SingularTable: true}})
	if err != nil {
		t.Fatalf("open sqlite: %v", err)
	}
	if err := db.AutoMigrate(&models.AgentToolInvocation{}); err != nil {
		t.Fatalf("auto migrate: %v", err)
	}
	sqls.SetDB(db)

	first, err := AgentToolInvocationService.Claim(13, 23, "runtime/ai_reply", "message:33:revision:2")
	if err != nil {
		t.Fatalf("first claim: %v", err)
	}
	oldUpdatedAt := time.Now().Add(-10 * time.Minute)
	if err := db.Model(&models.AgentToolInvocation{}).Where("id = ?", first.Item.ID).Update("updated_at", oldUpdatedAt).Error; err != nil {
		t.Fatalf("age running claim: %v", err)
	}

	second, err := AgentToolInvocationService.ClaimRecoverable(13, 23, "runtime/ai_reply", "message:33:revision:2", time.Now().Add(-time.Minute))
	if err != nil || second == nil || !second.Acquired || second.Item.Status != agentToolInvocationStatusRunning {
		t.Fatalf("recovered claim = %#v, err=%v", second, err)
	}
	if !second.Recovered || first.Item.ResultData == second.Item.ResultData {
		t.Fatalf("recovered claim did not receive a new lease: first=%q second=%q", first.Item.ResultData, second.Item.ResultData)
	}
	if err := AgentToolInvocationService.Complete(first.Item, `{"owner":"stale"}`); err == nil {
		t.Fatal("stale owner unexpectedly completed the recovered invocation")
	}
	var running models.AgentToolInvocation
	if err := db.First(&running, second.Item.ID).Error; err != nil {
		t.Fatalf("reload active lease: %v", err)
	}
	if running.Status != agentToolInvocationStatusRunning || running.ResultData != second.Item.ResultData {
		t.Fatalf("stale completion changed the active lease: %#v", running)
	}
	if err := AgentToolInvocationService.Complete(second.Item, `{"owner":"current"}`); err != nil {
		t.Fatalf("current owner complete: %v", err)
	}
}

var errTestToolInvocation = &toolInvocationTestError{}

type toolInvocationTestError struct{}

func (e *toolInvocationTestError) Error() string { return "tool failed" }