LocalAI/tests/e2e/distributed/skills_distributed_test.go

package distributed_test

import (
	"sync/atomic"

	"github.com/mudler/LocalAI/core/config"
	"github.com/mudler/LocalAI/core/services/distributed"
	"github.com/mudler/LocalAI/core/services/messaging"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"

	pgdriver "gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

var _ = Describe("Skills Distributed", Label("Distributed"), func() {
	var (
		infra      *TestInfra
		db         *gorm.DB
		skillStore *distributed.SkillStore
	)

	BeforeEach(func() {
		infra = SetupInfra("localai_skills_dist_test")

		var err error
		db, err = gorm.Open(pgdriver.Open(infra.PGURL), &gorm.Config{
			Logger: logger.Default.LogMode(logger.Silent),
		})
		Expect(err).ToNot(HaveOccurred())

		skillStore, err = distributed.NewSkillStore(db)
		Expect(err).ToNot(HaveOccurred())
	})

	Context("PostgreSQL metadata storage", func() {
		It("should store skill metadata in PostgreSQL", func() {
			rec := &distributed.SkillMetadataRecord{
				UserID:     "u1",
				Name:       "web-search",
				Definition: "# Web Search\nSearches the web for information.",
				SourceType: "inline",
				Enabled:    true,
			}
			Expect(skillStore.Save(rec)).To(Succeed())
			Expect(rec.ID).ToNot(BeEmpty())

			retrieved, err := skillStore.Get("u1", "web-search")
			Expect(err).ToNot(HaveOccurred())
			Expect(retrieved.Name).To(Equal("web-search"))
			Expect(retrieved.Definition).To(ContainSubstring("Web Search"))
			Expect(retrieved.SourceType).To(Equal("inline"))
			Expect(retrieved.Enabled).To(BeTrue())

			// Update skill
			rec.Definition = "# Web Search v2\nImproved search."
			Expect(skillStore.Save(rec)).To(Succeed())

			updated, _ := skillStore.Get("u1", "web-search")
			Expect(updated.Definition).To(ContainSubstring("v2"))

			// List skills
			skillStore.Save(&distributed.SkillMetadataRecord{
				UserID: "u1", Name: "code-gen", SourceType: "inline",
			})
			skillStore.Save(&distributed.SkillMetadataRecord{
				UserID: "u2", Name: "translate", SourceType: "git",
				SourceURL: "https://github.com/example/translate-skill",
			})

			u1Skills, _ := skillStore.List("u1")
			Expect(u1Skills).To(HaveLen(2))

			allSkills, _ := skillStore.List("")
			Expect(allSkills).To(HaveLen(3))

			// Git skills for sync
			gitSkills, err := skillStore.ListGitSkills()
			Expect(err).ToNot(HaveOccurred())
			Expect(gitSkills).To(HaveLen(1))
			Expect(gitSkills[0].Name).To(Equal("translate"))

			// Delete
			Expect(skillStore.Delete("u1", "web-search")).To(Succeed())
			_, err = skillStore.Get("u1", "web-search")
			Expect(err).To(HaveOccurred())
		})
	})

	Context("NATS cache invalidation", func() {
		It("should publish cache invalidation via NATS on skill change", func() {
			// Subscribe to skills cache invalidation
			var received atomic.Int32
			sub, err := infra.NC.Subscribe(messaging.SubjectCacheInvalidateSkills, func(data []byte) {
				received.Add(1)
			})
			Expect(err).ToNot(HaveOccurred())
			defer sub.Unsubscribe()

			FlushNATS(infra.NC)

			// Save a skill and publish cache invalidation
			rec := &distributed.SkillMetadataRecord{
				UserID: "u1", Name: "new-skill", SourceType: "inline",
			}
			Expect(skillStore.Save(rec)).To(Succeed())

			// Publish invalidation (in production this is done by the service layer)
			Expect(infra.NC.Publish(messaging.SubjectCacheInvalidateSkills, map[string]string{
				"user_id": "u1",
				"name":    "new-skill",
				"action":  "save",
			})).To(Succeed())

			Eventually(func() int32 { return received.Load() }, "5s").Should(Equal(int32(1)))

			// Delete and publish another invalidation
			Expect(skillStore.Delete("u1", "new-skill")).To(Succeed())
			Expect(infra.NC.Publish(messaging.SubjectCacheInvalidateSkills, map[string]string{
				"user_id": "u1",
				"name":    "new-skill",
				"action":  "delete",
			})).To(Succeed())

			Eventually(func() int32 { return received.Load() }, "5s").Should(Equal(int32(2)))
		})

		It("should broadcast collection cache invalidation", func() {
			var received atomic.Int32
			sub, err := infra.NC.Subscribe(messaging.SubjectCacheInvalidateCollection("my-collection"), func(data []byte) {
				received.Add(1)
			})
			Expect(err).ToNot(HaveOccurred())
			defer sub.Unsubscribe()

			FlushNATS(infra.NC)

			Expect(infra.NC.Publish(messaging.SubjectCacheInvalidateCollection("my-collection"), map[string]string{
				"reason": "skill_updated",
			})).To(Succeed())

			Eventually(func() int32 { return received.Load() }, "5s").Should(Equal(int32(1)))
		})
	})

	Context("Without --distributed", func() {
		It("should use filesystem without --distributed", func() {
			appCfg := config.NewApplicationConfig()
			Expect(appCfg.Distributed.Enabled).To(BeFalse())

			// Without distributed mode, skills are stored on the local
			// filesystem. No PostgreSQL metadata or NATS cache invalidation.
			Expect(appCfg.Distributed.NatsURL).To(BeEmpty())
		})
	})
})