diff --git a/core/application/studio_connectors.go b/core/application/studio_connectors.go index 8382d5f55e..87f998b91a 100644 --- a/core/application/studio_connectors.go +++ b/core/application/studio_connectors.go @@ -136,6 +136,11 @@ func (a *Application) startStudioConnectors(svc *studio.Service) { c.AddCapabilitySwitch(studio.TalkbackSwitch("turned off on this server (LOCALAI_DISABLE_CONNECTOR_TALKBACK)")) } a.studioConnectorsHook.Store(c) + if a.applicationConfig.SkipBindingCapsMigration { + xlog.Info("Connectors: the grant of the usual permissions to old bindings is skipped (LOCALAI_SKIP_BINDING_CAPS_MIGRATION)") + } else { + go a.migrateBindingCaps(a.backgroundContext(), tr) + } if a.ownsAddonsOnly() { // The manager stays: frontends get the stored bot credentials from it // and add-ons push messages to it. Only the polling is left to the @@ -162,3 +167,44 @@ func (a *Application) transfers() *hostops.Transfers { }) return a.addonTransfers } + +// bindingCapsRetry bounds how long a start waits for the connector kinds, which +// come from add-ons that may still be starting. +const ( + bindingCapsRetryEvery = 5 * time.Second + bindingCapsRetryFor = 10 * time.Minute +) + +// migrateBindingCaps gives the usual permissions, once, to the connector +// bindings saved before permissions were enforced. The store decides who runs +// it, so any number of processes on one store may call this at start. It needs +// the kinds that add-ons serve, to leave out a binding whose kind declares no +// permissions, so it waits for them. With none (add-ons off, or none installed) +// it does nothing and tries again at the next start. +func (a *Application) migrateBindingCaps(ctx context.Context, tr studio.ConnectorTransport) { + deadline := time.Now().Add(bindingCapsRetryFor) + for { + kinds, err := tr.Kinds(ctx) + if err == nil && len(kinds) > 0 { + rep, err := studio.MigrateBindingCaps(ctx, a.studioConnStore, kinds, time.Now(), nil) + if err == nil { + if !rep.Ran { + xlog.Debug("Connectors: the permissions migration already ran") + } + return + } + xlog.Warn("Connectors: could not give the usual permissions to old bindings; it is tried again at the next start", "error", err) + return + } + // No kinds yet, or the add-on gateway is not up: wait, as an add-on that + // serves a kind may be about to start. + if time.Now().After(deadline) { + return + } + select { + case <-ctx.Done(): + return + case <-time.After(bindingCapsRetryEvery): + } + } +} diff --git a/core/cli/run.go b/core/cli/run.go index a0c703a46f..d4eb3f597f 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -175,6 +175,7 @@ type RunCMD struct { DisableChats bool `env:"LOCALAI_DISABLE_CHATS" default:"false" help:"Turn off the chats: the /api/chats API, saved agents, triggers, connectors and the agent tasks and jobs. Nothing of them starts" group:"chats"` DisableTriggers bool `env:"LOCALAI_DISABLE_TRIGGERS" default:"false" help:"Turn off the trigger runner and the triggers API (the rest of the chats stays on)" group:"chats"` DisableConnectors bool `env:"LOCALAI_DISABLE_CONNECTORS" default:"false" help:"Turn off the connector bots (Telegram, Slack) and the connectors API (the rest of the chats stays on)" group:"chats"` + SkipBindingCapsMigration bool `env:"LOCALAI_SKIP_BINDING_CAPS_MIGRATION" default:"false" help:"Do not give the usual permissions, once, to connector bindings saved before permissions were enforced (they stay with none until the owner chooses)" group:"chats"` DisableConnectorTalkback bool `env:"LOCALAI_DISABLE_CONNECTOR_TALKBACK" default:"false" help:"Turn off the tool with which an agent writes to a chat of its connection, and the reports of triggers to a channel (the rest of the connectors stays on)" group:"chats"` AgentJobRetentionDays int `env:"LOCALAI_AGENT_JOB_RETENTION_DAYS,AGENT_JOB_RETENTION_DAYS" default:"30" help:"Days to keep finished agent jobs" group:"chats"` TriggerRunKeep int `env:"LOCALAI_TRIGGER_RUN_KEEP" default:"200" help:"Finished runs kept for each trigger, beyond those of the last two days (-1 keeps all)" group:"chats"` @@ -687,6 +688,9 @@ func (r *RunCMD) Run(ctx *cliContext.Context) (result error) { if r.DisableConnectorTalkback { opts = append(opts, config.DisableConnectorTalkback) } + if r.SkipBindingCapsMigration { + opts = append(opts, config.SkipBindingCapsMigration) + } opts = append(opts, config.WithAgentJobRetentionDays(r.AgentJobRetentionDays), config.WithMCPCIJobTimeout(r.MCPCIJobTimeout)) opts = append(opts, config.WithTriggerRunRetention(r.TriggerRunKeep, r.TriggerRunRetention)) diff --git a/core/config/application_config.go b/core/config/application_config.go index 9be9414c6a..e9c4d13781 100644 --- a/core/config/application_config.go +++ b/core/config/application_config.go @@ -353,6 +353,9 @@ type ApplicationConfig struct { // chat of its connection and stops the reports of triggers to a channel. // Owner notices keep working. DisableConnectorTalkback bool + // SkipBindingCapsMigration keeps the one-time grant of the usual + // permissions to old connector bindings from running. + SkipBindingCapsMigration bool // AgentJobRetentionDays is how long finished agent jobs are kept. Zero // means DefaultAgentJobRetentionDays. AgentJobRetentionDays int @@ -1407,6 +1410,8 @@ var ( DisableConnectors = func(o *ApplicationConfig) { o.DisableConnectors = true } // DisableConnectorTalkback is the option of LOCALAI_DISABLE_CONNECTOR_TALKBACK. DisableConnectorTalkback = func(o *ApplicationConfig) { o.DisableConnectorTalkback = true } + // SkipBindingCapsMigration is the option of LOCALAI_SKIP_BINDING_CAPS_MIGRATION. + SkipBindingCapsMigration = func(o *ApplicationConfig) { o.SkipBindingCapsMigration = true } ) // WithAgentJobRetentionDays sets how long finished agent jobs are kept. diff --git a/core/services/studio/connectorcapsmigrate.go b/core/services/studio/connectorcapsmigrate.go new file mode 100644 index 0000000000..32a8e53680 --- /dev/null +++ b/core/services/studio/connectorcapsmigrate.go @@ -0,0 +1,118 @@ +package studio + +import ( + "context" + "errors" + "fmt" + "log/slog" + "time" + + "github.com/mudler/LocalAI/core/services/studio/connectors" + "github.com/mudler/LocalAI/pkg/addons/contract" +) + +// BindingCapsMigration is the name of the one-time migration that gives the +// usual permissions to the bindings saved before connector permissions were +// enforced. The record in the store is the only marker: a binding keeps no flag +// of its own, because an empty Caps list is stored the same whether the owner +// never chose or cleared every switch. The migration therefore runs once for +// the whole store, and after it an empty list always means the owner's choice. +const BindingCapsMigration = "binding-caps-usual-v1" + +// bindingCapsLease is how long a claim that was never finished blocks another +// frontend. A run takes seconds, so a longer claim means the process died. +const bindingCapsLease = 10 * time.Minute + +// BindingCapsReport is the outcome of MigrateBindingCaps. +type BindingCapsReport struct { + // Ran is false when another process already ran or is running it. + Ran bool + // Changed is how many bindings got the usual permissions. + Changed int + // Skipped is how many bindings were left as they were: they had a choice + // already, or no kind, or a kind that declares no permissions. + Skipped int +} + +// MigrateBindingCaps gives the usual permissions (what "Allow the usual" turns +// on, never a sensitive one) to each binding that has no permission at all, once +// per store. kinds are the connector kinds that run now; a binding is changed +// only when its kind is among them and declares permissions. It must be called +// with a non-empty kinds list: with no kind known it would skip every binding +// and still mark the migration done. +// +// It is safe to call from several frontends at once: the store lets one claim +// it, and each binding is written only while it still has no permissions, so an +// owner who saves at the same moment is never overwritten. When it fails it +// releases the claim and the next start tries again. It logs one line for each +// binding changed (ID, kind and permissions, nothing else) and one summary. +func MigrateBindingCaps(ctx context.Context, st connectors.Store, kinds []ConnectorKind, now time.Time, log *slog.Logger) (BindingCapsReport, error) { + var rep BindingCapsReport + if st == nil { + return rep, errors.New("studio: the caps migration needs a store") + } + if log == nil { + log = slog.Default() + } + if len(kinds) == 0 { + return rep, errors.New("studio: the caps migration needs the connector kinds") + } + won, err := st.ClaimMigration(ctx, BindingCapsMigration, now, bindingCapsLease) + if err != nil || !won { + return rep, err + } + rep.Ran = true + done := false + defer func() { + if !done { + _ = st.ReleaseMigration(context.WithoutCancel(ctx), BindingCapsMigration) + } + }() + usual := contract.UsualConnectorCapabilities() + all, err := st.AllBindings(ctx) + if err != nil { + return rep, fmt.Errorf("list the connector bindings: %w", err) + } + for _, b := range all { + kind, ok := kindForMigration(kinds, b) + if len(b.Caps) > 0 || b.Kind == "" || !ok || len(kind.Permissions) == 0 { + rep.Skipped++ + continue + } + wrote, err := st.SetCapsIfEmpty(ctx, b.ID, usual) + if err != nil { + return rep, fmt.Errorf("give the usual permissions to binding %s: %w", b.ID, err) + } + if !wrote { + rep.Skipped++ + continue + } + rep.Changed++ + log.Info("studio: connector binding got the usual permissions", "binding", b.ID, "kind", b.Kind, "caps", usual) + } + if err := st.FinishMigration(ctx, BindingCapsMigration, now); err != nil { + return rep, err + } + done = true + log.Info("studio: connector permissions migration finished", "changed", rep.Changed, "skipped", rep.Skipped) + return rep, nil +} + +// kindForMigration finds the kind of a binding. A binding saved before +// connectors were add-ons names no add-on; it matches a kind ID that only one +// add-on serves. +func kindForMigration(kinds []ConnectorKind, b connectors.Binding) (ConnectorKind, bool) { + if b.AddonID != "" { + return kindFor(kinds, b) + } + var found []ConnectorKind + for _, k := range kinds { + if k.ID == b.Kind { + found = append(found, k) + } + } + if len(found) != 1 { + return ConnectorKind{}, false + } + return found[0], true +} diff --git a/core/services/studio/connectorcapsmigrate_helpers_test.go b/core/services/studio/connectorcapsmigrate_helpers_test.go new file mode 100644 index 0000000000..da2ae2f86c --- /dev/null +++ b/core/services/studio/connectorcapsmigrate_helpers_test.go @@ -0,0 +1,29 @@ +package studio_test + +import ( + "io" + "log/slog" + "strings" + "sync" +) + +type lockedBuffer struct { + mu sync.Mutex + b strings.Builder +} + +func (l *lockedBuffer) Write(p []byte) (int, error) { + l.mu.Lock() + defer l.mu.Unlock() + return l.b.Write(p) +} + +func (l *lockedBuffer) String() string { + l.mu.Lock() + defer l.mu.Unlock() + return l.b.String() +} + +func newTextLogger(w io.Writer) *slog.Logger { return slog.New(slog.NewTextHandler(w, nil)) } + +func hasText(s, sub string) bool { return strings.Contains(s, sub) } diff --git a/core/services/studio/connectorcapsmigrate_test.go b/core/services/studio/connectorcapsmigrate_test.go new file mode 100644 index 0000000000..f49b224e24 --- /dev/null +++ b/core/services/studio/connectorcapsmigrate_test.go @@ -0,0 +1,243 @@ +package studio_test + +import ( + "context" + "errors" + "slices" + "sync" + "testing" + "time" + + "github.com/mudler/LocalAI/core/services/studio" + "github.com/mudler/LocalAI/core/services/studio/connectors" + "github.com/mudler/LocalAI/pkg/addons/contract" +) + +var usualCaps = []string{"files.upload", "memory.write", "proposals.confirm", "talkback.channels", "triggers.propose", "triggers.run"} + +func migrationKinds() []studio.ConnectorKind { + perms := []string{"memory.write", "files.upload", "triggers.run"} + return []studio.ConnectorKind{ + {AddonID: "a-chat", ConnectorKind: contract.ConnectorKind{ID: "chat", Permissions: perms}}, + {AddonID: "a-bare", ConnectorKind: contract.ConnectorKind{ID: "bare"}}, + } +} + +func putBinding(t *testing.T, st connectors.Store, b connectors.Binding) { + t.Helper() + if b.Name == "" { + b.Name = "Bot " + b.ID + } + if b.AgentID == "" { + b.AgentID = "ag1" + } + if err := st.PutBinding(context.Background(), b); err != nil { + t.Fatal(err) + } +} + +func capsOf(t *testing.T, st connectors.Store, id string) []string { + t.Helper() + b, err := st.GetBinding(context.Background(), id) + if err != nil { + t.Fatal(err) + } + return b.Caps +} + +func TestUsualSetIsWhatThePresetGrants(t *testing.T) { + got := contract.UsualConnectorCapabilities() + if !slices.Equal(got, usualCaps) { + t.Fatalf("usual set = %v, want %v", got, usualCaps) + } + for _, c := range got { + if contract.SensitiveCapability(c) { + t.Fatalf("sensitive capability %s is in the usual set", c) + } + } + // The preset button turns on the capabilities that the list marks usual. + var ui []string + for _, c := range contract.ConnectorCapabilities() { + if c.Usual && !c.Sensitive { + ui = append(ui, c.Name) + } + } + if !slices.Equal(ui, got) { + t.Fatalf("the settings list marks %v usual, the migration grants %v", ui, got) + } +} + +func TestMigrateBindingCapsTable(t *testing.T) { + now := time.Date(2026, 11, 3, 10, 0, 0, 0, time.UTC) + cases := []struct { + name string + b connectors.Binding + want []string + }{ + {"never set, kind with permissions", connectors.Binding{ID: "old", Kind: "chat", AddonID: "a-chat"}, usualCaps}, + {"never set, no add-on recorded", connectors.Binding{ID: "noaddon", Kind: "chat"}, usualCaps}, + {"already chose some", connectors.Binding{ID: "some", Kind: "chat", AddonID: "a-chat", Caps: []string{"memory.write"}}, []string{"memory.write"}}, + {"already chose a sensitive one", connectors.Binding{ID: "sens", Kind: "chat", AddonID: "a-chat", Caps: []string{"profile.write"}}, []string{"profile.write"}}, + {"kind declares no permissions", connectors.Binding{ID: "bare", Kind: "bare", AddonID: "a-bare"}, nil}, + {"kind not running", connectors.Binding{ID: "gone", Kind: "other", AddonID: "a-other"}, nil}, + {"other add-on serves that kind id", connectors.Binding{ID: "hijack", Kind: "chat", AddonID: "a-other"}, nil}, + {"no kind", connectors.Binding{ID: "nokind"}, nil}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + st, _ := connectors.NewJSON("") + tc.b.Kind = cmpOr(tc.b.Kind, "chat") + if tc.b.ID == "nokind" { + // Validate refuses a binding without a kind, but an old row can have none. + tc.b.Kind = "" + } + putBinding(t, st, tc.b) + rep, err := studio.MigrateBindingCaps(context.Background(), st, migrationKinds(), now, nil) + if err != nil || !rep.Ran { + t.Fatalf("rep=%+v err=%v", rep, err) + } + if got := capsOf(t, st, tc.b.ID); !slices.Equal(got, tc.want) { + t.Fatalf("caps = %v, want %v", got, tc.want) + } + }) + } +} + +func cmpOr(a, b string) string { + if a != "" { + return a + } + return b +} + +func TestMigrateBindingCapsRunsOnce(t *testing.T) { + ctx := context.Background() + now := time.Date(2026, 11, 3, 10, 0, 0, 0, time.UTC) + st, _ := connectors.NewJSON("") + putBinding(t, st, connectors.Binding{ID: "b1", Kind: "chat", AddonID: "a-chat"}) + rep, err := studio.MigrateBindingCaps(ctx, st, migrationKinds(), now, nil) + if err != nil || !rep.Ran || rep.Changed != 1 { + t.Fatalf("first run: %+v %v", rep, err) + } + // The owner clears every switch on purpose, and a new empty binding is made. + b, _ := st.GetBinding(ctx, "b1") + b.Caps = nil + putBinding(t, st, b) + putBinding(t, st, connectors.Binding{ID: "b2", Kind: "chat", AddonID: "a-chat"}) + for i := range 3 { + rep, err = studio.MigrateBindingCaps(ctx, st, migrationKinds(), now.Add(time.Duration(i+1)*24*time.Hour), nil) + if err != nil || rep.Ran || rep.Changed != 0 { + t.Fatalf("run %d: %+v %v", i+2, rep, err) + } + } + if len(capsOf(t, st, "b1")) != 0 || len(capsOf(t, st, "b2")) != 0 { + t.Fatal("a binding that was cleared or made after the migration got permissions") + } +} + +func TestMigrateBindingCapsConcurrent(t *testing.T) { + ctx := context.Background() + now := time.Date(2026, 11, 3, 10, 0, 0, 0, time.UTC) + st, _ := connectors.NewJSON("") + for _, id := range []string{"a", "b", "c"} { + putBinding(t, st, connectors.Binding{ID: id, Kind: "chat", AddonID: "a-chat"}) + } + var wg sync.WaitGroup + var mu sync.Mutex + ran, changed := 0, 0 + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + rep, err := studio.MigrateBindingCaps(ctx, st, migrationKinds(), now, nil) + if err != nil { + t.Error(err) + return + } + mu.Lock() + defer mu.Unlock() + if rep.Ran { + ran++ + } + changed += rep.Changed + }() + } + wg.Wait() + if ran != 1 || changed != 3 { + t.Fatalf("ran=%d changed=%d, want 1 and 3", ran, changed) + } + for _, id := range []string{"a", "b", "c"} { + if !slices.Equal(capsOf(t, st, id), usualCaps) { + t.Fatalf("binding %s has %v", id, capsOf(t, st, id)) + } + } +} + +// failingStore makes the write of one binding fail, to show that a failed run +// releases its claim and a later run finishes the job. +type failingStore struct { + connectors.Store + failOn string + fail bool +} + +func (f *failingStore) SetCapsIfEmpty(ctx context.Context, id string, caps []string) (bool, error) { + if f.fail && id == f.failOn { + return false, errors.New("disk full") + } + return f.Store.SetCapsIfEmpty(ctx, id, caps) +} + +func TestMigrateBindingCapsRetriesAfterAFailure(t *testing.T) { + ctx := context.Background() + now := time.Date(2026, 11, 3, 10, 0, 0, 0, time.UTC) + inner, _ := connectors.NewJSON("") + st := &failingStore{Store: inner, failOn: "b", fail: true} + for _, id := range []string{"a", "b"} { + putBinding(t, st, connectors.Binding{ID: id, Kind: "chat", AddonID: "a-chat"}) + } + if _, err := studio.MigrateBindingCaps(ctx, st, migrationKinds(), now, nil); err == nil { + t.Fatal("expected the failure") + } + st.fail = false + rep, err := studio.MigrateBindingCaps(ctx, st, migrationKinds(), now.Add(time.Second), nil) + if err != nil || !rep.Ran || rep.Changed != 1 { + t.Fatalf("retry: %+v %v (the first binding was done already)", rep, err) + } + if !slices.Equal(capsOf(t, st, "b"), usualCaps) { + t.Fatal("the failed binding was not migrated on retry") + } +} + +func TestMigrateBindingCapsNeedsKinds(t *testing.T) { + st, _ := connectors.NewJSON("") + putBinding(t, st, connectors.Binding{ID: "b1", Kind: "chat", AddonID: "a-chat"}) + if _, err := studio.MigrateBindingCaps(context.Background(), st, nil, time.Now(), nil); err == nil { + t.Fatal("with no kinds the migration must not run, and must not be marked done") + } + rep, err := studio.MigrateBindingCaps(context.Background(), st, migrationKinds(), time.Now(), nil) + if err != nil || !rep.Ran || rep.Changed != 1 { + t.Fatalf("later run: %+v %v", rep, err) + } +} + +func TestMigrateBindingCapsLogsWithoutSecrets(t *testing.T) { + var buf lockedBuffer + log := newTextLogger(&buf) + st, _ := connectors.NewJSON("") + putBinding(t, st, connectors.Binding{ID: "b1", Kind: "chat", AddonID: "a-chat", Token: []byte("sealed-secret-token"), Allow: []string{"123"}}) + if _, err := studio.MigrateBindingCaps(context.Background(), st, migrationKinds(), time.Now(), log); err != nil { + t.Fatal(err) + } + out := buf.String() + for _, want := range []string{"binding=b1", "kind=chat", "memory.write", "changed=1"} { + if !hasText(out, want) { + t.Fatalf("log lacks %q:\n%s", want, out) + } + } + for _, bad := range []string{"sealed-secret-token", "123", "profile.write"} { + if hasText(out, bad) { + t.Fatalf("log has %q:\n%s", bad, out) + } + } +} diff --git a/core/services/studio/connectors/connectorstest/suite.go b/core/services/studio/connectors/connectorstest/suite.go index d9be2f6063..a79761fb1f 100644 --- a/core/services/studio/connectors/connectorstest/suite.go +++ b/core/services/studio/connectors/connectorstest/suite.go @@ -267,6 +267,71 @@ func Run(newStore NewStore) { id, _, _ = s.ChatOf(ctx, "b1", "2") Expect(id).To(Equal("c2")) }) + It("sets the capabilities only while a binding has none", func() { + s := newStore() + Expect(s.PutBinding(ctx, binding("none", "u1", "None"))).To(Succeed()) + withCaps := binding("has", "u1", "Has") + withCaps.Caps = []string{"memory.write"} + Expect(s.PutBinding(ctx, withCaps)).To(Succeed()) + ok, err := s.SetCapsIfEmpty(ctx, "none", []string{"files.upload", "memory.write"}) + Expect(err).NotTo(HaveOccurred()) + Expect(ok).To(BeTrue()) + ok, _ = s.SetCapsIfEmpty(ctx, "none", []string{"triggers.run"}) + Expect(ok).To(BeFalse(), "a second write must not replace the first") + ok, _ = s.SetCapsIfEmpty(ctx, "has", []string{"triggers.run"}) + Expect(ok).To(BeFalse()) + ok, err = s.SetCapsIfEmpty(ctx, "missing", []string{"triggers.run"}) + Expect(err).NotTo(HaveOccurred()) + Expect(ok).To(BeFalse()) + got, _ := s.GetBinding(ctx, "none") + Expect(got.Caps).To(Equal([]string{"files.upload", "memory.write"})) + got, _ = s.GetBinding(ctx, "has") + Expect(got.Caps).To(Equal([]string{"memory.write"})) + }) + It("lets one caller claim a migration, once, until a stale claim is taken over", func() { + s := newStore() + lease := 10 * time.Minute + ok, err := s.ClaimMigration(ctx, "m1", t0, lease) + Expect(err).NotTo(HaveOccurred()) + Expect(ok).To(BeTrue()) + ok, _ = s.ClaimMigration(ctx, "m1", t0.Add(time.Minute), lease) + Expect(ok).To(BeFalse(), "a live claim blocks another caller") + ok, _ = s.ClaimMigration(ctx, "m2", t0, lease) + Expect(ok).To(BeTrue(), "another name is another migration") + ok, _ = s.ClaimMigration(ctx, "m1", t0.Add(lease+time.Second), lease) + Expect(ok).To(BeTrue(), "an unfinished claim past its lease is taken over") + Expect(s.FinishMigration(ctx, "m1", t0.Add(lease+time.Minute))).To(Succeed()) + ok, _ = s.ClaimMigration(ctx, "m1", t0.Add(100*time.Hour), lease) + Expect(ok).To(BeFalse(), "a finished migration is never claimed again") + Expect(s.ReleaseMigration(ctx, "m1")).To(Succeed()) + ok, _ = s.ClaimMigration(ctx, "m1", t0.Add(200*time.Hour), lease) + Expect(ok).To(BeFalse(), "release does not undo a finished migration") + Expect(s.ReleaseMigration(ctx, "m2")).To(Succeed()) + ok, _ = s.ClaimMigration(ctx, "m2", t0.Add(time.Second), lease) + Expect(ok).To(BeTrue(), "a released claim can be made again") + }) + It("lets only one of many racing callers claim a migration", func() { + s := newStore() + var wg sync.WaitGroup + var mu sync.Mutex + wins := 0 + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + defer GinkgoRecover() + ok, err := s.ClaimMigration(ctx, "race", t0, time.Hour) + Expect(err).NotTo(HaveOccurred()) + if ok { + mu.Lock() + wins++ + mu.Unlock() + } + }() + } + wg.Wait() + Expect(wins).To(Equal(1)) + }) It("makes the key once even when calls race", func() { s := newStore() var mu sync.Mutex diff --git a/core/services/studio/connectors/gorm.go b/core/services/studio/connectors/gorm.go index 987a98cfc2..46c0c01ca3 100644 --- a/core/services/studio/connectors/gorm.go +++ b/core/services/studio/connectors/gorm.go @@ -99,6 +99,17 @@ type keyRow struct { func (keyRow) TableName() string { return "studio_connector_keys" } +// migrationRow records a one-time migration. Done is set when it finished; +// until then Started is the lease of the process that runs it. +type migrationRow struct { + Name string `gorm:"primaryKey;size:64"` + Started time.Time + Finished time.Time + Done bool +} + +func (migrationRow) TableName() string { return "studio_connector_migrations" } + type gormStore struct{ db *gorm.DB } // NewGORM returns a Store on db (SQLite or PostgreSQL) and creates its tables. @@ -106,7 +117,7 @@ func NewGORM(db *gorm.DB) (Store, error) { if db == nil { return nil, errors.New("connectors: nil *gorm.DB") } - if err := db.AutoMigrate(&bindingRow{}, &stateRow{}, &senderRow{}, &chatRow{}, &keyRow{}, &pairingRow{}); err != nil { + if err := db.AutoMigrate(&bindingRow{}, &stateRow{}, &senderRow{}, &chatRow{}, &keyRow{}, &pairingRow{}, &migrationRow{}); err != nil { return nil, fmt.Errorf("connectors: migrate: %w", err) } return &gormStore{db: db}, nil @@ -324,3 +335,42 @@ func (s *gormStore) ConsumePairing(ctx context.Context, binding, hash string, no res := s.db.WithContext(ctx).Where("binding = ? AND hash = ? AND expires > ?", binding, hash, now.UTC()).Delete(&pairingRow{}) return res.RowsAffected == 1, res.Error } + +func (s *gormStore) SetCapsIfEmpty(ctx context.Context, id string, caps []string) (bool, error) { + if len(caps) == 0 { + return false, nil + } + raw, _ := json.Marshal(caps) + // Caps is "" when a row has none (PutBinding writes it so), NULL for a row + // from before the column existed, and "[]" is never written but is + // accepted for safety. + res := s.db.WithContext(ctx).Model(&bindingRow{}). + Where("id = ? AND (caps IS NULL OR caps = '' OR caps = '[]' OR caps = 'null')", id). + Update("caps", string(raw)) + return res.RowsAffected == 1, res.Error +} + +func (s *gormStore) ClaimMigration(ctx context.Context, name string, now time.Time, lease time.Duration) (bool, error) { + db := s.db.WithContext(ctx) + res := db.Clauses(clause.OnConflict{DoNothing: true}).Create(&migrationRow{Name: name, Started: now.UTC()}) + if res.Error != nil { + return false, res.Error + } + if res.RowsAffected == 1 { + return true, nil + } + // A claim that was never finished and has outlived its lease is taken + // over with a conditional update, so two takers cannot both win. + res = db.Model(&migrationRow{}).Where("name = ? AND done = ? AND started <= ?", name, false, now.Add(-lease).UTC()). + Update("started", now.UTC()) + return res.RowsAffected == 1, res.Error +} + +func (s *gormStore) FinishMigration(ctx context.Context, name string, now time.Time) error { + return s.db.WithContext(ctx).Model(&migrationRow{}).Where("name = ?", name). + Updates(map[string]any{"done": true, "finished": now.UTC()}).Error +} + +func (s *gormStore) ReleaseMigration(ctx context.Context, name string) error { + return s.db.WithContext(ctx).Where("name = ? AND done = ?", name, false).Delete(&migrationRow{}).Error +} diff --git a/core/services/studio/connectors/json.go b/core/services/studio/connectors/json.go index f4b9360e9c..34521e25b0 100644 --- a/core/services/studio/connectors/json.go +++ b/core/services/studio/connectors/json.go @@ -32,10 +32,18 @@ type jsonData struct { Senders map[string][]Sender `json:"senders"` Chats map[string]string `json:"chats"` // binding + "\x00" + sender Pairings map[string]Pairing `json:"pairings,omitempty"` + // Migrations are the one-time migrations that ran or are running. + Migrations map[string]migrationRecord `json:"migrations,omitempty"` +} + +type migrationRecord struct { + Started time.Time `json:"started"` + Finished time.Time `json:"finished,omitzero"` + Done bool `json:"done"` } func newData() jsonData { - return jsonData{Bindings: map[string]Binding{}, Offsets: map[string]int64{}, Senders: map[string][]Sender{}, Chats: map[string]string{}, Pairings: map[string]Pairing{}} + return jsonData{Bindings: map[string]Binding{}, Offsets: map[string]int64{}, Senders: map[string][]Sender{}, Chats: map[string]string{}, Pairings: map[string]Pairing{}, Migrations: map[string]migrationRecord{}} } // NewJSON returns a Store. With an empty dir it keeps everything in memory. @@ -64,6 +72,9 @@ func NewJSON(dir string) (Store, error) { if s.d.Pairings == nil { s.d.Pairings = map[string]Pairing{} } + if s.d.Migrations == nil { + s.d.Migrations = map[string]migrationRecord{} + } return s, nil } @@ -266,3 +277,59 @@ func (s *jsonStore) ConsumePairing(_ context.Context, binding, hash string, now } return true, nil } + +func (s *jsonStore) SetCapsIfEmpty(_ context.Context, id string, caps []string) (bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + b, ok := s.d.Bindings[id] + if !ok || len(b.Caps) > 0 || len(caps) == 0 { + return false, nil + } + old := b + b.Caps = append([]string(nil), caps...) + s.d.Bindings[id] = b + if err := s.save(); err != nil { + s.d.Bindings[id] = old + return false, err + } + return true, nil +} + +func (s *jsonStore) ClaimMigration(_ context.Context, name string, now time.Time, lease time.Duration) (bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + old, had := s.d.Migrations[name] + if had && (old.Done || now.Sub(old.Started) < lease) { + return false, nil + } + s.d.Migrations[name] = migrationRecord{Started: now.UTC()} + if err := s.save(); err != nil { + if had { + s.d.Migrations[name] = old + } else { + delete(s.d.Migrations, name) + } + return false, err + } + return true, nil +} + +func (s *jsonStore) FinishMigration(_ context.Context, name string, now time.Time) error { + s.mu.Lock() + defer s.mu.Unlock() + r := s.d.Migrations[name] + r.Done, r.Finished = true, now.UTC() + s.d.Migrations[name] = r + return s.save() +} + +func (s *jsonStore) ReleaseMigration(_ context.Context, name string) error { + s.mu.Lock() + defer s.mu.Unlock() + r, ok := s.d.Migrations[name] + if !ok || r.Done { + return nil + } + delete(s.d.Migrations, name) + return s.save() +} diff --git a/core/services/studio/connectors/types.go b/core/services/studio/connectors/types.go index 9abca06639..74250663c0 100644 --- a/core/services/studio/connectors/types.go +++ b/core/services/studio/connectors/types.go @@ -291,6 +291,22 @@ type Store interface { // leaves the code in place. The check and the removal are one step, so a // code works once on any frontend. ConsumePairing(ctx context.Context, binding, hash string, now time.Time) (bool, error) + + // SetCapsIfEmpty sets the capabilities of a binding only while it has none, + // as one step, so it never overwrites a choice that an owner saved at the + // same time. It reports whether it wrote. A missing binding is false. + SetCapsIfEmpty(ctx context.Context, id string, caps []string) (bool, error) + + // ClaimMigration lets one caller, on any frontend, run the named one-time + // migration. It is true for the first caller, and for a later one only when + // an earlier claim was never finished and is older than lease. A finished + // migration is never claimed again. + ClaimMigration(ctx context.Context, name string, now time.Time, lease time.Duration) (bool, error) + // FinishMigration records that the claimed migration is done for good. + FinishMigration(ctx context.Context, name string, now time.Time) error + // ReleaseMigration drops an unfinished claim so that the next start can + // claim it again. It does nothing for a finished migration. + ReleaseMigration(ctx context.Context, name string) error } // An entry of the allowlist is a numeric ID (Telegram), an @username, or an diff --git a/docs/content/features/chats-connectors.md b/docs/content/features/chats-connectors.md index 5543fea21c..46de5a91c6 100644 --- a/docs/content/features/chats-connectors.md +++ b/docs/content/features/chats-connectors.md @@ -504,7 +504,16 @@ add-on. `triggers.run`, `talkback.channels`, `tools.write`, and the sensitive `talkback.new_channel`, `profile.write`, `skills.write`, `pages.write`, `dashboard.write`). The default is no: a binding with none enabled lets a sender talk to the agent, but the agent can - not use any of these on the sender's behalf. + not use any of these on the sender's behalf. A binding saved before permissions + were enforced has none, so the first start after the upgrade gives it the usual set + once, the same six that **Allow the usual** turns on (`proposals.confirm`, + `memory.write`, `files.upload`, `triggers.propose`, `triggers.run`, + `talkback.channels`; never a sensitive one). It only touches a binding with no + capability whose kind lists permissions, and it logs one line for each binding it + changes. It runs once for the whole store, also with several frontends, and never + again, so a binding you clear later stays cleared. A binding you cleared before that + start cannot be told from one never set and gets the set too. To keep the strict + state, set `LOCALAI_SKIP_BINDING_CAPS_MIGRATION=true` before the first start. - **Roles and grants** (`grants`). An add-on can declare roles, such as the admins of a workspace, and report the roles of each sender. You map a role to capabilities with `grants: [{"role", "caps"}]`. A grant is within `caps`, within what the kind diff --git a/docs/content/features/chats.md b/docs/content/features/chats.md index 19e3336975..7d898c3ca7 100644 --- a/docs/content/features/chats.md +++ b/docs/content/features/chats.md @@ -88,6 +88,7 @@ what runs (`chats`, `agents`, `triggers`, `connectors`). | `LOCALAI_DISABLE_TRIGGERS` | `false` | Turns off the trigger runner. Trigger routes answer `409 triggers_off`. | | `LOCALAI_DISABLE_CONNECTORS` | `false` | Turns off the connector bots. Connector routes answer `409 connectors_off`. | | `LOCALAI_DISABLE_CONNECTOR_TALKBACK` | `false` | Turns off `send_message` and the reports of triggers to a channel for every connector binding. | +| `LOCALAI_SKIP_BINDING_CAPS_MIGRATION` | `false` | Skips the one-time grant of the usual permissions to connector bindings that have none. | | `LOCALAI_DISABLE_AGENTS` | `false` | Deprecated. It kept the meaning it had in the old agent pool, which also held the knowledge base: it sets `LOCALAI_DISABLE_CHATS` and `LOCALAI_DISABLE_COLLECTIONS`, and logs a warning. | `LOCALAI_DISABLE_COLLECTIONS` stays separate and turns off only the diff --git a/docs/content/reference/cli-reference.md b/docs/content/reference/cli-reference.md index 404eb865ec..0bde389e8e 100644 --- a/docs/content/reference/cli-reference.md +++ b/docs/content/reference/cli-reference.md @@ -227,6 +227,7 @@ See [Authentication & Authorization]({{%relref "features/authentication" %}}) fo | `--disable-triggers` | `false` | Turn off the trigger runner and the triggers API; the rest of the chats stays on | `$LOCALAI_DISABLE_TRIGGERS` | | `--disable-connectors` | `false` | Turn off the connector bots (Telegram, Slack) and the connectors API; the rest of the chats stays on | `$LOCALAI_DISABLE_CONNECTORS` | | `--disable-connector-talkback` | `false` | Turn off the tool with which an agent writes to a chat of its connection, and the reports of triggers to a channel (the rest of the connectors stays on) | `$LOCALAI_DISABLE_CONNECTOR_TALKBACK` | +| `--skip-binding-caps-migration` | `false` | Do not give the usual permissions, once, to connector bindings saved before permissions were enforced (they stay with none until the owner chooses) | `$LOCALAI_SKIP_BINDING_CAPS_MIGRATION` | | `--agent-job-retention-days` | `30` | Days to keep ended agent jobs | `$LOCALAI_AGENT_JOB_RETENTION_DAYS`, `$AGENT_JOB_RETENTION_DAYS` | | `--trigger-run-keep` | `200` | Finished runs kept for each trigger, beyond the runs of the last two days. `-1` keeps all | `$LOCALAI_TRIGGER_RUN_KEEP` | | `--trigger-run-retention-days` | `30` | Days to keep finished trigger runs. `-1` keeps them forever | `$LOCALAI_TRIGGER_RUN_RETENTION_DAYS` | diff --git a/pkg/addons/contract/permissions.go b/pkg/addons/contract/permissions.go index 5cf0a77b57..1973b07d50 100644 --- a/pkg/addons/contract/permissions.go +++ b/pkg/addons/contract/permissions.go @@ -104,6 +104,20 @@ var usualCapabilities = map[string]bool{ CapTriggersPropose: true, CapTriggersRun: true, CapTalkbackChannels: true, } +// UsualConnectorCapabilities returns what the "Allow the usual" preset of a +// settings page turns on, sorted: the capabilities marked usual, never a +// sensitive one. A migration that applies the preset uses this list so it +// grants exactly what the button grants. +func UsualConnectorCapabilities() []string { + var out []string + for _, n := range ConnectorCapabilityNames() { + if usualCapabilities[n] && !capabilities[n] { + out = append(out, n) + } + } + return out +} + // ConnectorCapabilities returns the closed list with its texts, in the order of // ConnectorCapabilityNames. func ConnectorCapabilities() []CapabilityInfo {