From 79109058d0e1a4fc3609fcec70ab0a6ac7c80ee4 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 3 Nov 2026 13:43:08 +0000 Subject: [PATCH] feat(studio): keep a list of folders on a saved agent An agent may now have up to eight folders, each with its own mode. The list is stored in a new table for GORM and in a folders field for the JSON store. The old single-folder pair stays and mirrors the first entry for one release, so a rollback keeps working. At start the GORM store copies a single folder into the list; the copy is safe to repeat. Assisted-by: Claude:claude-sonnet-5-5 go gofmt --- .../studio/agents/agentstest/suite.go | 39 +++++ core/services/studio/agents/gorm.go | 142 ++++++++++++++---- .../studio/agents/gorm_sqlite_test.go | 42 ++++++ core/services/studio/agents/json.go | 2 +- core/services/studio/agents/types.go | 101 +++++++++++-- core/services/studio/agents/types_test.go | 18 +++ 6 files changed, 299 insertions(+), 45 deletions(-) diff --git a/core/services/studio/agents/agentstest/suite.go b/core/services/studio/agents/agentstest/suite.go index 21e4c6c35c..6e7095f66b 100644 --- a/core/services/studio/agents/agentstest/suite.go +++ b/core/services/studio/agents/agentstest/suite.go @@ -138,6 +138,45 @@ func Run(newStore NewStore) { Expect(all[0].FolderID).To(BeZero()) Expect(all[0].FolderMode).To(BeEmpty(), "after clearing") }) + It("round-trips a list of folders in order and mirrors the first into the old pair", func() { + s := newStore() + a := sample("a1", "u1", "Multi") + a.Folders = []agents.AgentFolder{{Node: 9, Mode: agents.FolderWrite}, {Node: 3, Mode: agents.FolderRead}, {Node: 5, Mode: agents.FolderRead}} + Expect(s.Put(ctx, a)).To(Succeed()) + got, err := s.Get(ctx, "u1", "a1") + Expect(err).NotTo(HaveOccurred()) + Expect(got.Folders).To(Equal(a.Folders)) + Expect(got.FolderID).To(Equal(int64(9))) + Expect(got.FolderMode).To(Equal(agents.FolderWrite)) + + // A shorter list replaces the longer one. + a.Folders = a.Folders[1:2] + Expect(s.Put(ctx, a)).To(Succeed()) + l, _ := s.List(ctx, "u1") + Expect(l[0].Folders).To(Equal([]agents.AgentFolder{{Node: 3, Mode: agents.FolderRead}})) + + // An agent that has only the old pair reads as a one-entry list. + old := sample("a2", "u1", "Old") + old.FolderID, old.FolderMode = 4, agents.FolderRead + Expect(s.Put(ctx, old)).To(Succeed()) + got, _ = s.Get(ctx, "u1", "a2") + Expect(got.Folders).To(Equal([]agents.AgentFolder{{Node: 4, Mode: agents.FolderRead}})) + + // No folders at all clears the list. + a.Folders, a.FolderID, a.FolderMode = nil, 0, "" + Expect(s.Put(ctx, a)).To(Succeed()) + got, _ = s.Get(ctx, "u1", "a1") + Expect(got.Folders).To(BeEmpty()) + Expect(got.HasFolder()).To(BeFalse()) + all, _ := s.ListAll(ctx) + Expect(all).To(HaveLen(2)) + + Expect(s.Delete(ctx, "u1", "a2")).To(Succeed()) + again := sample("a2", "u1", "Old") + Expect(s.Put(ctx, again)).To(Succeed()) + got, _ = s.Get(ctx, "u1", "a2") + Expect(got.Folders).To(BeEmpty(), "the folders of a deleted agent do not come back") + }) It("ReturnedListsAreCopies", func() { s := newStore() _ = s.Put(ctx, sample("a1", "u1", "x")) diff --git a/core/services/studio/agents/gorm.go b/core/services/studio/agents/gorm.go index b448b6a7f2..312debdcb1 100644 --- a/core/services/studio/agents/gorm.go +++ b/core/services/studio/agents/gorm.go @@ -10,6 +10,7 @@ import ( "time" "gorm.io/gorm" + "gorm.io/gorm/clause" "github.com/mudler/LocalAI/core/services/studio/internal/dberr" ) @@ -41,6 +42,18 @@ type agentRow struct { func (agentRow) TableName() string { return "studio_agents" } +// folderRow is one folder of an agent, at position Pos in its list. The pair +// (AgentID, Pos) is the key, so two frontends that run the startup migration +// at once cannot make a folder twice. +type folderRow struct { + AgentID string `gorm:"primaryKey;size:64"` + Pos int `gorm:"primaryKey;autoIncrement:false"` + Node int64 `gorm:"not null;default:0"` + Mode string `gorm:"size:8;not null;default:''"` +} + +func (folderRow) TableName() string { return "studio_agent_folders" } + type gormStore struct{ db *gorm.DB } // NewGORM returns a Store on db (SQLite or PostgreSQL) and creates its table. @@ -48,12 +61,37 @@ func NewGORM(db *gorm.DB) (Store, error) { if db == nil { return nil, errors.New("agents: nil *gorm.DB") } - if err := db.AutoMigrate(&agentRow{}); err != nil { + if err := db.AutoMigrate(&agentRow{}, &folderRow{}); err != nil { return nil, fmt.Errorf("agents: migrate: %w", err) } + if err := migrateScalarFolders(db); err != nil { + return nil, fmt.Errorf("agents: migrate folders: %w", err) + } return &gormStore{db: db}, nil } +// migrateScalarFolders copies the one folder that an agent kept in the columns +// folder_id and folder_mode into the folder table. The columns stay as they +// are for one release, so a rollback works. It is safe to run on every start +// and from several frontends at once: it only adds the first row of an agent +// that has none, and a clash on the key is ignored. +func migrateScalarFolders(db *gorm.DB) error { + var rows []agentRow + if err := db.Where("folder_mode <> ''"). + Where("id NOT IN (?)", db.Model(&folderRow{}).Select("agent_id")). + Find(&rows).Error; err != nil { + return err + } + for _, r := range rows { + err := db.Clauses(clause.OnConflict{DoNothing: true}). + Create(&folderRow{AgentID: r.ID, Pos: 0, Node: r.FolderID, Mode: r.FolderMode}).Error + if err != nil { + return err + } + } + return nil +} + func listJSON(l []string) string { if len(l) == 0 { return "" @@ -71,9 +109,12 @@ func listFromJSON(s string) ([]string, error) { return l, err } -func (r agentRow) agent() (Agent, error) { +func (r agentRow) agent(folders []folderRow) (Agent, error) { a := Agent{ID: r.ID, Principal: r.Principal, Name: r.Name, Description: r.Description, Instructions: r.Instructions, Model: r.Model, Network: r.Network, FolderID: r.FolderID, FolderMode: r.FolderMode, Created: r.Created, Updated: r.Updated} + for _, f := range folders { + a.Folders = append(a.Folders, AgentFolder{Node: f.Node, Mode: f.Mode}) + } var err error for _, f := range []struct { dst *[]string @@ -83,7 +124,44 @@ func (r agentRow) agent() (Agent, error) { return a, err } } - return a, nil + return a.normalizeFolders(), nil +} + +// foldersOf loads the folder rows of the agents, by agent ID, in list order. +func foldersOf(db *gorm.DB, ids []string) (map[string][]folderRow, error) { + out := map[string][]folderRow{} + for len(ids) > 0 { + n := min(len(ids), 500) + var rows []folderRow + if err := db.Where("agent_id IN ?", ids[:n]).Order("agent_id, pos").Find(&rows).Error; err != nil { + return nil, err + } + for _, r := range rows { + out[r.AgentID] = append(out[r.AgentID], r) + } + ids = ids[n:] + } + return out, nil +} + +func (s *gormStore) agentsOf(rows []agentRow) ([]Agent, error) { + ids := make([]string, len(rows)) + for i, r := range rows { + ids[i] = r.ID + } + fs, err := foldersOf(s.db, ids) + if err != nil { + return nil, err + } + out := make([]Agent, 0, len(rows)) + for _, r := range rows { + a, err := r.agent(fs[r.ID]) + if err != nil { + return nil, err + } + out = append(out, a) + } + return out, nil } // isNameTaken reports a unique violation of the name index. A collision on the @@ -98,6 +176,7 @@ func isNameTaken(err error) bool { } func (s *gormStore) Put(ctx context.Context, a Agent) error { + a = a.normalizeFolders() return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { var old agentRow err := tx.Where("id = ?", a.ID).Take(&old).Error @@ -122,7 +201,18 @@ func (s *gormStore) Put(ctx context.Context, a Agent) error { if err != nil && isNameTaken(err) { return ErrNameTaken } - return err + if err != nil { + return err + } + if err := tx.Where("agent_id = ?", a.ID).Delete(&folderRow{}).Error; err != nil { + return err + } + for i, f := range a.Folders { + if err := tx.Create(&folderRow{AgentID: a.ID, Pos: i, Node: f.Node, Mode: f.Mode}).Error; err != nil { + return err + } + } + return nil }) } @@ -135,7 +225,11 @@ func (s *gormStore) Get(ctx context.Context, principal, id string) (Agent, error if err != nil { return Agent{}, err } - return r.agent() + fs, err := foldersOf(s.db.WithContext(ctx), []string{r.ID}) + if err != nil { + return Agent{}, err + } + return r.agent(fs[r.ID]) } func (s *gormStore) List(ctx context.Context, principal string) ([]Agent, error) { @@ -143,13 +237,9 @@ func (s *gormStore) List(ctx context.Context, principal string) ([]Agent, error) if err := s.db.WithContext(ctx).Where("principal = ?", principal).Order("name_key, id").Find(&rows).Error; err != nil { return nil, err } - out := make([]Agent, 0, len(rows)) - for _, r := range rows { - a, err := r.agent() - if err != nil { - return nil, err - } - out = append(out, a) + out, err := s.agentsOf(rows) + if err != nil { + return nil, err } sort.SliceStable(out, func(i, j int) bool { return NameKey(out[i].Name) < NameKey(out[j].Name) }) return out, nil @@ -160,24 +250,18 @@ func (s *gormStore) ListAll(ctx context.Context) ([]Agent, error) { if err := s.db.WithContext(ctx).Order("principal, name_key, id").Find(&rows).Error; err != nil { return nil, err } - out := make([]Agent, 0, len(rows)) - for _, r := range rows { - a, err := r.agent() - if err != nil { - return nil, err - } - out = append(out, a) - } - return out, nil + return s.agentsOf(rows) } func (s *gormStore) Delete(ctx context.Context, principal, id string) error { - res := s.db.WithContext(ctx).Where("id = ? AND principal = ?", id, principal).Delete(&agentRow{}) - if res.Error != nil { - return res.Error - } - if res.RowsAffected == 0 { - return ErrNotFound - } - return nil + return s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + res := tx.Where("id = ? AND principal = ?", id, principal).Delete(&agentRow{}) + if res.Error != nil { + return res.Error + } + if res.RowsAffected == 0 { + return ErrNotFound + } + return tx.Where("agent_id = ?", id).Delete(&folderRow{}).Error + }) } diff --git a/core/services/studio/agents/gorm_sqlite_test.go b/core/services/studio/agents/gorm_sqlite_test.go index e7df190d69..c1b5956211 100644 --- a/core/services/studio/agents/gorm_sqlite_test.go +++ b/core/services/studio/agents/gorm_sqlite_test.go @@ -45,6 +45,26 @@ var _ = Describe("the GORM agent store on SQLite", func() { Expect(got.FolderMode).To(BeEmpty()) }) + It("moves the one-folder columns into the folder list at start, once", func() { + db := openSQLite(filepath.Join(GinkgoT().TempDir(), "studio.db")) + Expect(db.AutoMigrate(&oneFolderAgentRow{})).To(Succeed()) + now := time.Date(2026, 10, 1, 9, 0, 0, 0, time.UTC) + Expect(db.Create(&oneFolderAgentRow{ID: "one", Principal: "u1", NameKey: "one", Name: "One", FolderID: 12, FolderMode: "write", Created: now, Updated: now}).Error).To(Succeed()) + Expect(db.Create(&oneFolderAgentRow{ID: "none", Principal: "u1", NameKey: "none", Name: "None", Created: now, Updated: now}).Error).To(Succeed()) + for range 2 { // the second start finds nothing to do + s, err := agents.NewGORM(db) + Expect(err).NotTo(HaveOccurred()) + got, err := s.Get(context.Background(), "u1", "one") + Expect(err).NotTo(HaveOccurred()) + Expect(got.Folders).To(Equal([]agents.AgentFolder{{Node: 12, Mode: agents.FolderWrite}})) + none, _ := s.Get(context.Background(), "u1", "none") + Expect(none.HasFolder()).To(BeFalse()) + } + var n int64 + Expect(db.Table("studio_agent_folders").Count(&n).Error).To(Succeed()) + Expect(n).To(Equal(int64(1))) + }) + It("survives a reopen", func() { path := filepath.Join(GinkgoT().TempDir(), "studio.db") first, err := agents.NewGORM(openSQLite(path)) @@ -75,3 +95,25 @@ type legacyAgentRow struct { } func (legacyAgentRow) TableName() string { return "studio_agents" } + +// oneFolderAgentRow is studio_agents as the one-folder release left it. +type oneFolderAgentRow struct { + ID string `gorm:"primaryKey;size:64"` + Principal string `gorm:"index:idx_studio_agents_name,unique;size:255"` + NameKey string `gorm:"index:idx_studio_agents_name,unique;size:255"` + Name string `gorm:"size:255"` + Description string + Instructions string + Model string `gorm:"size:255"` + MCPServers string + Addons string + Collections string + Skills string + Network string `gorm:"size:16"` + FolderID int64 `gorm:"not null;default:0"` + FolderMode string `gorm:"size:8;not null;default:''"` + Created time.Time + Updated time.Time +} + +func (oneFolderAgentRow) TableName() string { return "studio_agents" } diff --git a/core/services/studio/agents/json.go b/core/services/studio/agents/json.go index eae3b1494b..4d81a78f44 100644 --- a/core/services/studio/agents/json.go +++ b/core/services/studio/agents/json.go @@ -45,7 +45,7 @@ func NewJSON(dir string) (Store, error) { return nil, fmt.Errorf("agents: read %s: %w", s.path, err) } for _, a := range all { - s.byID[a.ID] = a + s.byID[a.ID] = a.clone() } return s, nil } diff --git a/core/services/studio/agents/types.go b/core/services/studio/agents/types.go index e741ee667d..068cfb0f76 100644 --- a/core/services/studio/agents/types.go +++ b/core/services/studio/agents/types.go @@ -56,26 +56,68 @@ type Agent struct { // by the administrator's ceiling. Empty is off. A run never asks a // person: a host outside the preset is blocked. Network string - // FolderID is the workspace node of the agent's folder, kept by node so a - // rename or a move of the folder is followed. FolderMode is FolderRead or - // FolderWrite for an agent with a folder and "" for an agent that works - // on the whole home. FolderMode set with FolderID 0 is an agent whose - // folder is not chosen yet (an import whose folder was not found): it - // must not start. + // Folders are the folders the agent works in, at most MaxFolders, in the + // order the person gave them. Each is kept by workspace node so a rename + // or a move of the folder is followed. No folder means the agent works on + // the whole home. An entry with node 0 is a folder that is not chosen yet + // (an import whose folder was not found): the agent must not start. + Folders []AgentFolder + // FolderID and FolderMode are the first entry of Folders, kept for one + // release so a rollback to a server that knows one folder still works. + // + // Deprecated: use Folders. A store fills Folders from these two fields + // when it reads an agent that has only them, and fills these two from the + // first of Folders when it writes. FolderID int64 FolderMode string Created time.Time Updated time.Time } +// MaxFolders is the most folders one agent may have. +const MaxFolders = 8 + +// AgentFolder is one folder of an agent: a workspace node and the most the +// agent may do there. Where one folder is inside another, the outer mode caps +// the inner one. +type AgentFolder struct { + Node int64 + Mode string // FolderRead or FolderWrite +} + // The access modes of an agent's folder. const ( FolderRead = "read" FolderWrite = "write" ) -// HasFolder reports whether the agent is limited to a folder, chosen or not. -func (a Agent) HasFolder() bool { return a.FolderMode != "" } +// HasFolder reports whether the agent is limited to folders, chosen or not. +func (a Agent) HasFolder() bool { return len(a.FolderList()) > 0 } + +// FolderList returns the folders of the agent. An agent read from a store that +// kept one folder in FolderID and FolderMode has them as a one-entry list. +func (a Agent) FolderList() []AgentFolder { + if len(a.Folders) > 0 { + return a.Folders + } + if a.FolderMode != "" { + return []AgentFolder{{Node: a.FolderID, Mode: a.FolderMode}} + } + return nil +} + +// normalizeFolders makes Folders and the deprecated pair agree: Folders is the +// truth, the pair mirrors its first entry, and the slice is a copy. +func (a Agent) normalizeFolders() Agent { + list := a.FolderList() + a.Folders = append([]AgentFolder(nil), list...) + if len(list) == 0 { + a.Folders, a.FolderID, a.FolderMode = nil, 0, "" + return a + } + a.FolderID, a.FolderMode = list[0].Node, list[0].Mode + return a +} // Store is the storage of saved agents. type Store interface { @@ -119,14 +161,12 @@ func Validate(a Agent) error { default: return bad("the network must be off, registries or all") } - switch a.FolderMode { - case "", FolderRead, FolderWrite: - default: - return bad("the folder mode must be read or write") - } - if a.FolderID < 0 || (a.FolderID != 0 && a.FolderMode == "") { + if a.FolderMode == "" && a.FolderID != 0 && len(a.Folders) == 0 { return bad("a folder needs a mode, read or write") } + if err := ValidateFolders(a.FolderList()); err != nil { + return err + } for _, l := range []struct { name string entries []string @@ -149,11 +189,42 @@ func Validate(a Agent) error { return nil } +// ValidateFolders checks a folder list: at most MaxFolders entries, a mode +// on each, no node twice. Node 0 (not chosen yet) may repeat. The error wraps +// ErrInvalid. A caller that knows which nodes are shared folders must also +// refuse write on those; the store does not know. +func ValidateFolders(list []AgentFolder) error { + bad := func(format string, args ...any) error { + return fmt.Errorf("%w: %s", ErrInvalid, fmt.Sprintf(format, args...)) + } + if len(list) > MaxFolders { + return bad("an agent may have at most %d folders", MaxFolders) + } + seen := map[int64]bool{} + for _, f := range list { + switch f.Mode { + case FolderRead, FolderWrite: + case "": + return bad("a folder needs a mode, read or write") + default: + return bad("the folder mode must be read or write") + } + if f.Node < 0 { + return bad("a folder node cannot be negative") + } + if f.Node != 0 && seen[f.Node] { + return bad("the same folder is listed twice") + } + seen[f.Node] = true + } + return nil +} + // clone copies the lists so a caller cannot change a stored agent. func (a Agent) clone() Agent { a.MCPServers = append([]string(nil), a.MCPServers...) a.Addons = append([]string(nil), a.Addons...) a.Collections = append([]string(nil), a.Collections...) a.Skills = append([]string(nil), a.Skills...) - return a + return a.normalizeFolders() } diff --git a/core/services/studio/agents/types_test.go b/core/services/studio/agents/types_test.go index bf6de8fc12..9ab926b039 100644 --- a/core/services/studio/agents/types_test.go +++ b/core/services/studio/agents/types_test.go @@ -45,4 +45,22 @@ var _ = Describe("agent validation", func() { Expect(agents.Validate(a)).To(MatchError(agents.ErrInvalid), name) } }) + + It("limits the folder list", func() { + list := func(n int) []agents.AgentFolder { + var l []agents.AgentFolder + for i := 1; i <= n; i++ { + l = append(l, agents.AgentFolder{Node: int64(i), Mode: agents.FolderRead}) + } + return l + } + Expect(agents.Validate(agents.Agent{Name: "a", Folders: list(agents.MaxFolders)})).To(Succeed()) + Expect(agents.Validate(agents.Agent{Name: "a", Folders: list(agents.MaxFolders + 1)})).To(MatchError(agents.ErrInvalid), "nine folders") + dup := []agents.AgentFolder{{Node: 3, Mode: agents.FolderRead}, {Node: 3, Mode: agents.FolderWrite}} + Expect(agents.Validate(agents.Agent{Name: "a", Folders: dup})).To(MatchError(agents.ErrInvalid), "same node twice") + Expect(agents.Validate(agents.Agent{Name: "a", Folders: []agents.AgentFolder{{Node: 3}}})).To(MatchError(agents.ErrInvalid), "no mode") + Expect(agents.Validate(agents.Agent{Name: "a", Folders: []agents.AgentFolder{{Node: 3, Mode: "all"}}})).To(MatchError(agents.ErrInvalid), "bad mode") + unset := []agents.AgentFolder{{Mode: agents.FolderRead}, {Mode: agents.FolderRead}} + Expect(agents.Validate(agents.Agent{Name: "a", Folders: unset})).To(Succeed(), "several unchosen rows") + }) })