package application import ( "context" "encoding/json" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "github.com/mudler/LocalAI/core/http/auth" "github.com/mudler/LocalAI/core/services/addons/apps" "github.com/mudler/LocalAI/core/services/addons/gateway" "github.com/mudler/LocalAI/core/services/addons/hostops" addontools "github.com/mudler/LocalAI/core/services/addons/tools" "github.com/mudler/LocalAI/core/services/studio" "github.com/mudler/LocalAI/pkg/addons/contract" ) type changeOwner struct{ app apps.App } func (o changeOwner) Apps(context.Context, contract.Actor) (apps.List, error) { return apps.List{Apps: []apps.App{o.app}}, nil } func (changeOwner) Dispatch(context.Context, contract.Actor, contract.Request) (contract.Result, error) { return contract.Result{}, nil } func (changeOwner) Propose(context.Context, contract.Actor, contract.Request) (gateway.Proposal, error) { return gateway.Proposal{}, nil } type forbiddenOwner struct{ changeOwner } func (forbiddenOwner) Apps(context.Context, contract.Actor) (apps.List, error) { return apps.List{}, gateway.ErrForbidden } var _ = Describe("the studio add-ons", func() { It("lets a chat act as its user with the chat credential", func() { s := &studioAddons{mode: auth.AddonModeNamedUser} a, err := s.actor("alice") Expect(err).ToNot(HaveOccurred()) Expect(a.UserID).To(Equal("alice")) Expect(a.CredentialID).To(Equal(auth.AddonChatsCredential)) Expect(a.Mode).To(Equal(auth.AddonModeNamedUser)) for _, bad := range []string{"", studio.AnonymousPrincipal, auth.AddonLocalOwnerID, "legacy-api-key"} { _, err := s.actor(bad) Expect(err).To(HaveOccurred(), "actor(%q) was accepted in named-user mode", bad) } owner := &studioAddons{mode: auth.AddonModeLocalOwner} oa, err := owner.actor(studio.AnonymousPrincipal) Expect(err).ToNot(HaveOccurred(), "local-owner actor") Expect(oa.UserID).To(Equal(auth.AddonLocalOwnerID)) _, err = (&studioAddons{}).actor("alice") Expect(err).To(HaveOccurred(), "an unknown mode gave an actor") _, err = s.actor("") Expect(err).To(HaveOccurred()) Expect(err.Error()).To(ContainSubstring(gateway.ErrForbidden.Error()), "a refused actor is not a denial") }) It("shows a proposal to the model as a text that says nothing changed", func() { res := addontools.ToolResult{Kind: addontools.ResultProposal, Tool: "t", Proposal: &addontools.ProposalRef{ID: "p1", AddonID: "tables", Action: "rows-delete"}} got := studioAddonText(res) for _, want := range []string{"tables", "rows-delete", "p1", "Nothing has been changed", "Add-ons page"} { Expect(got).To(ContainSubstring(want), "proposal text") } data := addontools.ToolResult{Kind: addontools.ResultData, Tool: "t", Result: []byte(`{"n":1}`)} Expect(studioAddonText(data)).To(ContainSubstring(`"n":1`), "a read result lost its content") }) It("lets a chat credential confirm while session and key rules stay the same", func() { named := confirmCredential(auth.AddonModeNamedUser) chat := contract.Actor{UserID: "alice", CredentialID: auth.AddonChatsCredential, Mode: auth.AddonModeNamedUser} Expect(named(chat)).To(BeTrue(), "the chat credential cannot confirm what the chat stored") Expect(named(contract.Actor{UserID: "alice", CredentialID: "session:s1", Mode: auth.AddonModeNamedUser})).To(BeTrue(), "a browser session lost the right to confirm") for _, id := range []string{"key:k1", "session:", "", "mcp:1", "Chats"} { Expect(named(contract.Actor{UserID: "alice", CredentialID: id, Mode: auth.AddonModeNamedUser})).To(BeFalse(), "credential %q may confirm", id) } owner := confirmCredential(auth.AddonModeLocalOwner) Expect(owner(contract.Actor{UserID: auth.AddonLocalOwnerID, CredentialID: auth.AddonLocalOwnerID, Mode: auth.AddonModeLocalOwner})).To(BeTrue(), "the local owner cannot confirm") Expect(owner(chat)).To(BeFalse(), "the chat credential confirms in local-owner mode, where chats act as the local owner") }) It("marks a change tool as an action, so an unattended run without an agent does not get it", func() { schema := json.RawMessage(`{"type":"object","properties":{"resourceId":{"type":"string"},"table":{"type":"string"},"rows":{"type":"array","items":{"type":"object"}}},"required":["resourceId","table","rows"],"additionalProperties":false}`) read := json.RawMessage(`{"type":"object","properties":{"resourceId":{"type":"string"}},"required":["resourceId"],"additionalProperties":false}`) app := apps.App{ID: "tables", Generation: 1, Enabled: true, State: "running", Tools: []string{"rows-insert", "rows-get"}, ToolDefinitions: []contract.ToolDefinition{ {Name: "rows-insert", Kind: contract.ToolKindWrite, Description: "Insert rows.", InputSchema: schema}, {Name: "rows-get", Kind: contract.ToolKindRead, Description: "Read rows.", InputSchema: read}, }} s := &studioAddons{adapter: addontools.NewChatToolAdapter(changeOwner{app}), mode: auth.AddonModeNamedUser} got, err := s.Tools(context.Background(), "alice") Expect(err).ToNot(HaveOccurred()) action := map[string]bool{} for _, d := range got { action[d.Name] = d.Action } Expect(action["addon__tables__rows-insert"]).To(BeTrue(), "Action flags = %v, want the write tool marked", action) Expect(action["addon__tables__rows-get"]).To(BeFalse(), "Action flags = %v, want the read tool not marked", action) }) It("carries what the Sandbox's bash shadows: the tool that runs commands is the one that shadows bash", func() { obj := json.RawMessage(`{"type":"object","properties":{"command":{"type":"string"}},"required":["command"],"additionalProperties":false}`) none := json.RawMessage(`{"type":"object","properties":{},"additionalProperties":false}`) app := apps.App{ID: "sandbox", Generation: 1, Enabled: true, State: "running", Tools: []string{"bash", "status"}, ToolDefinitions: []contract.ToolDefinition{ {Name: "bash", Kind: contract.ToolKindAction, Description: "Run a shell script.", InputSchema: obj, Shadows: "bash"}, {Name: "status", Kind: contract.ToolKindRead, Description: "Say whether it works.", InputSchema: none}, }} s := &studioAddons{adapter: addontools.NewChatToolAdapter(changeOwner{app}), mode: auth.AddonModeNamedUser} got, err := s.Tools(context.Background(), "alice") Expect(err).ToNot(HaveOccurred()) shadows := map[string]string{} for _, d := range got { shadows[d.Name] = d.Shadows } Expect(shadows).To(Equal(map[string]string{"addon__sandbox__bash": "bash", "addon__sandbox__status": ""})) }) It("shows an applied change to the model as a document that says it cannot be undone", func() { res := addontools.ToolResult{Kind: addontools.ResultApplied, Tool: "t", Applied: &addontools.AppliedRef{ID: "o1", AddonID: "tables", Action: "write", Count: 2}, Note: addontools.AppliedNote} got := studioAddonText(res) for _, want := range []string{`"kind":"applied"`, `"undoable":false`, "cannot be undone", "tables"} { Expect(got).To(ContainSubstring(want), "applied text") } }) It("makes an add-on that needs the workspace unavailable while files are off", func() { s := &studioAddons{mode: auth.AddonModeNamedUser, hostOps: map[string][]string{"sandbox": {"workspace."}, "office": {"workspace."}, "fine-tuning": {"finetune."}}} ctx := context.Background() for _, id := range []string{"sandbox", "office"} { Expect(s.Unavailable(ctx, "alice", id, false)).To(ContainSubstring("workspace"), "%s with files off", id) } Expect(s.Unavailable(ctx, "alice", "fine-tuning", false)).To(BeEmpty(), "an add-on that does not use the workspace was gated") Expect(s.Unavailable(ctx, "alice", "office", true)).To(BeEmpty(), "office with files on") }) It("gives a user without the add-ons feature no add-on tools", func() { s := &studioAddons{adapter: addontools.NewChatToolAdapter(forbiddenOwner{}), mode: auth.AddonModeNamedUser} got, err := s.Tools(context.Background(), "alice") Expect(err).ToNot(HaveOccurred(), "a refused user is not an error for the chats") Expect(got).To(BeEmpty(), "tools") }) }) // bindingOwner notes the chat binding of each read that reaches it. type bindingOwner struct { changeOwner seen *[]hostops.Binding } func (o bindingOwner) Dispatch(ctx context.Context, _ contract.Actor, _ contract.Request) (contract.Result, error) { b, _ := hostops.BindingOf(ctx) *o.seen = append(*o.seen, b) return contract.Result{Payload: json.RawMessage(`{"ok":true}`)}, nil } var _ = Describe("a chat's add-on call", func() { It("binds a read tool with no resource to the chat it runs in", func() { var seen []hostops.Binding app := apps.App{ID: "office", Name: "office", Generation: 1, Enabled: true, State: "running", Tools: []string{"read-docx"}, ToolDefinitions: []contract.ToolDefinition{{Name: "read-docx", Kind: contract.ToolKindRead, Description: "Read a document.", InputSchema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}},"required":["path"],"additionalProperties":false}`)}}} s := &studioAddons{adapter: addontools.NewChatToolAdapter(bindingOwner{changeOwner: changeOwner{app: app}, seen: &seen}), mode: auth.AddonModeNamedUser} _, err := s.Call(context.Background(), "alice", "chat1", addontools.ToolName("office", "read-docx"), json.RawMessage(`{"path":"/a.docx"}`)) Expect(err).ToNot(HaveOccurred()) Expect(seen).To(Equal([]hostops.Binding{{Principal: "alice", ChatID: "chat1"}})) }) })