ai-agent/internal/ai/runtime/reply_commit_service_test.go

package runtime

import (
	"strings"
	"testing"
	"time"

	"agent-desk/internal/models"
	"agent-desk/internal/pkg/enums"

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

func TestReplyCommitStoresWorkflowRunIDOnAIMessage(t *testing.T) {
	db := setupReplyCommitTestDB(t)
	aiAgent := createReplyCommitTestAIAgent(t, db)
	conversation := createReplyCommitTestConversation(t, db, aiAgent.ID)

	replyMessage, err := newReplyCommitService().CommitAIReply(replyCommitInput{
		Conversation:  *conversation,
		Message:       models.Message{ID: 101, RequestID: "trace-101"},
		AIAgent:       *aiAgent,
		ReplyText:     "AI reply",
		ClientPrefix:  "ai_reply",
		WorkflowRunID: 9988,
	})
	if err != nil {
		t.Fatalf("CommitAIReply() error = %v", err)
	}
	if replyMessage == nil {
		t.Fatalf("expected reply message")
	}
	if replyMessage.WorkflowRunID != 9988 {
		t.Fatalf("replyMessage.WorkflowRunID=%d want 9988", replyMessage.WorkflowRunID)
	}

	var stored models.Message
	if err := db.First(&stored, replyMessage.ID).Error; err != nil {
		t.Fatalf("find reply message: %v", err)
	}
	if stored.WorkflowRunID != 9988 {
		t.Fatalf("stored.WorkflowRunID=%d want 9988", stored.WorkflowRunID)
	}
}

func TestReplyCommitRejectsSensitiveModelOutput(t *testing.T) {
	db := setupReplyCommitTestDB(t)
	aiAgent := createReplyCommitTestAIAgent(t, db)
	conversation := createReplyCommitTestConversation(t, db, aiAgent.ID)
	_, err := newReplyCommitService().CommitAIReply(replyCommitInput{
		Conversation: *conversation, Message: models.Message{ID: 102, RequestID: "trace-102"}, AIAgent: *aiAgent,
		ReplyText: "authorization=Bearer-secret", ClientPrefix: "ai_reply",
	})
	if err == nil {
		t.Fatal("expected sensitive model output to be rejected")
	}
	var count int64
	if err := db.Model(&models.Message{}).Where("conversation_id = ?", conversation.ID).Count(&count).Error; err != nil {
		t.Fatalf("count messages: %v", err)
	}
	if count != 0 {
		t.Fatalf("unexpected message written for rejected output: %d", count)
	}
}

func setupReplyCommitTestDB(t *testing.T) *gorm.DB {
	t.Helper()
	dbName := "reply_commit_test_" + strings.NewReplacer("/", "_").Replace(t.Name())
	db, err := gorm.Open(sqlite.Open("file:"+dbName+"?mode=memory&cache=shared"), &gorm.Config{
		NamingStrategy: schema.NamingStrategy{
			TablePrefix:   "t_",
			SingularTable: true,
		},
	})
	if err != nil {
		t.Fatalf("open sqlite db: %v", err)
	}
	sqlDB, err := db.DB()
	if err != nil {
		t.Fatalf("get sqlite db: %v", err)
	}
	t.Cleanup(func() {
		if err := sqlDB.Close(); err != nil {
			t.Fatalf("close sqlite db: %v", err)
		}
	})
	if err := db.AutoMigrate(
		&models.AIAgent{},
		&models.Channel{},
		&models.ChannelMessageOutbox{},
		&models.Conversation{},
		&models.ConversationReadState{},
		&models.ConversationEventLog{},
		&models.Message{},
	); err != nil {
		t.Fatalf("auto migrate: %v", err)
	}
	sqls.SetDB(db)
	return db
}

func createReplyCommitTestAIAgent(t *testing.T, db *gorm.DB) *models.AIAgent {
	t.Helper()
	now := time.Now()
	item := &models.AIAgent{
		Name:   "reply-agent",
		Status: enums.StatusOk,
		AuditFields: models.AuditFields{
			CreatedAt: now,
			UpdatedAt: now,
		},
	}
	if err := db.Create(item).Error; err != nil {
		t.Fatalf("create ai agent: %v", err)
	}
	return item
}

func createReplyCommitTestConversation(t *testing.T, db *gorm.DB, aiAgentID int64) *models.Conversation {
	t.Helper()
	now := time.Now()
	item := &models.Conversation{
		CustomerID:   1,
		ChannelID:    11,
		AIAgentID:    aiAgentID,
		Status:       enums.IMConversationStatusAIServing,
		LastActiveAt: now,
		AuditFields: models.AuditFields{
			CreatedAt: now,
			UpdatedAt: now,
		},
	}
	if err := db.Create(item).Error; err != nil {
		t.Fatalf("create conversation: %v", err)
	}
	return item
}