LocalAI/core/application/credential_source_test.go

// SPDX-License-Identifier: MIT
package application

import (
	"context"
	"errors"

	"github.com/mudler/LocalAI/core/services/addons/gateway"
	"github.com/mudler/LocalAI/core/services/addons/hostops"
	"github.com/mudler/LocalAI/core/services/studio"
	"github.com/mudler/LocalAI/core/services/studio/agents"
	"github.com/mudler/LocalAI/core/services/studio/connectors"
	"github.com/mudler/LocalAI/core/services/studio/connsets"
	"github.com/mudler/LocalAI/core/services/studio/provider"
	"github.com/mudler/LocalAI/core/services/studio/session"
	"github.com/mudler/LocalAI/core/services/studio/store"
	"github.com/mudler/LocalAI/core/services/studio/studiotest"
	"github.com/mudler/LocalAI/pkg/addons/contract"
	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

// tokenTransport serves one credential-only kind and answers nothing else.
type tokenTransport struct{}

func (tokenTransport) Kinds(context.Context) ([]studio.ConnectorKind, error) {
	return []studio.ConnectorKind{{AddonID: "tracker", ConnectorKind: contract.ConnectorKind{ID: "tracker", Title: "Tracker", Delivery: contract.DeliveryPull, AgentOptional: true}}}, nil
}

func (tokenTransport) Call(context.Context, studio.ConnectorCallScope, contract.ConnectorRequest) (contract.ConnectorResponse, error) {
	return contract.ConnectorResponse{}, nil
}

// The connection a tool call of an add-on with a credential uses: one of the
// set the chat enabled, never a first match. The add-on here is a stand-in; the
// rules do not depend on which product it is.
var _ = Describe("the credential source of the add-on gateway", func() {
	ctx := context.Background()
	const user = "alice"
	var (
		app   *Application
		chats store.Store
		src   credentialSource
		work  string
		home  string
	)
	refusal := func(err error) string {
		var r *gateway.CredentialRefusal
		Expect(errors.As(err, &r)).To(BeTrue(), "error = %v", err)
		return r.Code
	}
	actor := contract.Actor{UserID: user, CredentialID: "chats", Mode: "named-user"}
	inChat := func(chat string) context.Context {
		return hostops.WithBinding(ctx, hostops.Binding{Principal: user, ChatID: chat})
	}
	newChat := func(rec session.ChatRecord) string {
		rec.ID, rec.Principal, rec.Kind, rec.Model = "chat-"+rec.Title, user, session.KindChat, "m"
		Expect(chats.CreateChat(ctx, rec)).To(Succeed())
		return rec.ID
	}
	BeforeEach(func() {
		var err error
		chats, err = store.NewJSON("")
		Expect(err).NotTo(HaveOccurred())
		agentStore, err := agents.NewJSON("")
		Expect(err).NotTo(HaveOccurred())
		svc, err := studio.New(studio.Deps{Store: chats, Agents: agentStore, Provider: provider.New(), Harness: studiotest.NewGateHarness(), StateDir: GinkgoT().TempDir()})
		Expect(err).NotTo(HaveOccurred())
		DeferCleanup(func() { Expect(svc.Close(ctx)).To(Succeed()) })
		bindings, err := connectors.NewJSON("")
		Expect(err).NotTo(HaveOccurred())
		conns, err := studio.NewConnectors(svc, studio.ConnectorOptions{Store: bindings, Transport: tokenTransport{}})
		Expect(err).NotTo(HaveOccurred())
		DeferCleanup(conns.Close)
		svc.SetConnectors(conns)
		app = &Application{studioStore: chats, studioConnectors: conns}
		src = credentialSource{app}
		mk := func(name, token string) string {
			b, err := conns.CreateConnector(ctx, user, studio.ConnectorInput{Kind: "tracker", Name: name, Token: token, Resources: []string{"acme/widgets"}})
			Expect(err).NotTo(HaveOccurred())
			return b.ID
		}
		work, home = mk("Work", "token-work"), mk("Home", "token-home")
	})

	It("refuses a call whose chat selected nothing, with several connections", func() {
		_, err := src.Credential(inChat(newChat(session.ChatRecord{Title: "none"})), actor, "tracker", "acme/widgets")
		Expect(refusal(err)).To(Equal(gateway.CodeSelectConnection))
	})
	It("uses the connection the chat selected, and only its resource list", func() {
		id := newChat(session.ChatRecord{Title: "sel", Connections: connsets.Sets{"tracker": {home}}})
		grant, err := src.Credential(inChat(id), actor, "tracker", "acme/widgets")
		Expect(err).NotTo(HaveOccurred())
		Expect(grant.Credential).To(Equal("token-home"))
		Expect(grant.BindingID).To(Equal(home))
		Expect(grant.BindingName).To(Equal("Home"))
		_, err = src.Credential(inChat(id), actor, "tracker", "other/repo")
		Expect(refusal(err)).To(Equal(gateway.CodeResourceNotAllowed))
	})
	It("uses the one connection a chat enables, with no name in the call", func() {
		id := newChat(session.ChatRecord{Title: "one", Connections: connsets.Sets{"tracker": {work}}})
		grant, err := src.Credential(inChat(id), actor, "tracker", "acme/widgets")
		Expect(err).NotTo(HaveOccurred())
		Expect(grant.BindingID).To(Equal(work))
	})
	It("refuses a call that names no connection when the chat enables two", func() {
		id := newChat(session.ChatRecord{Title: "two", Connections: connsets.Sets{"tracker": {work, home}}})
		_, err := src.Credential(inChat(id), actor, "tracker", "acme/widgets")
		Expect(refusal(err)).To(Equal(gateway.CodeSelectConnection), "the host picks neither")
	})
	It("takes the connection the studio named, when it is in the chat's set", func() {
		id := newChat(session.ChatRecord{Title: "two2", Connections: connsets.Sets{"tracker": {work, home}}})
		for _, want := range []string{work, home} {
			grant, err := src.Credential(gateway.WithConnection(inChat(id), want), actor, "tracker", "acme/widgets")
			Expect(err).NotTo(HaveOccurred())
			Expect(grant.BindingID).To(Equal(want))
		}
	})
	It("refuses a named connection that is the person's but not in the chat's set", func() {
		id := newChat(session.ChatRecord{Title: "sel2", Connections: connsets.Sets{"tracker": {home}}})
		_, err := src.Credential(gateway.WithConnection(inChat(id), work), actor, "tracker", "acme/widgets")
		Expect(refusal(err)).To(Equal(gateway.CodeSelectConnection))
	})
	It("lets the broker apply a proposal with the connection it was previewed with, outside a chat", func() {
		grant, err := src.Credential(gateway.WithConnection(ctx, work), actor, "tracker", "acme/widgets")
		Expect(err).NotTo(HaveOccurred())
		Expect(grant.BindingID).To(Equal(work))
		_, err = src.Credential(gateway.WithConnection(ctx, "no-such-binding"), actor, "tracker", "acme/widgets")
		Expect(refusal(err)).To(Equal(gateway.CodeNoCredential))
	})
	It("refuses a connection that belongs to someone else", func() {
		_, err := src.Credential(gateway.WithConnection(ctx, work), contract.Actor{UserID: "mallory", CredentialID: "chats", Mode: "named-user"}, "tracker", "acme/widgets")
		Expect(refusal(err)).To(Equal(gateway.CodeNoCredential))
	})
	It("refuses a call from no chat when there are several connections, and names no first", func() {
		_, err := src.Credential(ctx, actor, "tracker", "acme/widgets")
		Expect(refusal(err)).To(Equal(gateway.CodeSelectConnection))
	})
	It("refuses a chat of another person", func() {
		id := newChat(session.ChatRecord{Title: "mine", Connections: connsets.Sets{"tracker": {work}}})
		bad := hostops.WithBinding(ctx, hostops.Binding{Principal: "mallory", ChatID: id})
		_, err := src.Credential(bad, contract.Actor{UserID: "mallory", CredentialID: "chats", Mode: "named-user"}, "tracker", "acme/widgets")
		Expect(refusal(err)).To(Equal(gateway.CodeSelectConnection))
	})
})

// A model chooses among the connections of a chat by name, so a person cannot
// have two credential-only connections of one add-on that fold to one name.
var _ = Describe("the names of connections", func() {
	ctx := context.Background()
	const user = "alice"
	var conns *studio.Connectors
	BeforeEach(func() {
		chats, err := store.NewJSON("")
		Expect(err).NotTo(HaveOccurred())
		svc, err := studio.New(studio.Deps{Store: chats, Provider: provider.New(), Harness: studiotest.NewGateHarness(), StateDir: GinkgoT().TempDir()})
		Expect(err).NotTo(HaveOccurred())
		DeferCleanup(func() { Expect(svc.Close(ctx)).To(Succeed()) })
		bindings, err := connectors.NewJSON("")
		Expect(err).NotTo(HaveOccurred())
		conns, err = studio.NewConnectors(svc, studio.ConnectorOptions{Store: bindings, Transport: tokenTransport{}})
		Expect(err).NotTo(HaveOccurred())
		DeferCleanup(conns.Close)
		svc.SetConnectors(conns)
	})
	mk := func(owner, name string) (studio.ConnectorInfo, error) {
		return conns.CreateConnector(ctx, owner, studio.ConnectorInput{Kind: "tracker", Name: name, Token: "t", Resources: []string{"a/b"}})
	}
	It("refuses a name that another connection of the person already has, and says why", func() {
		_, err := mk(user, "Work")
		Expect(err).NotTo(HaveOccurred())
		_, err = mk(user, "Work")
		Expect(err).To(MatchError(ContainSubstring(`already have a connection named "Work"`)))
		Expect(errors.Is(err, connectors.ErrInvalid)).To(BeTrue(), "a 400, not an outage")
	})
	It("compares by handle: case and punctuation do not make a new name", func() {
		_, err := mk(user, "Work Acme")
		Expect(err).NotTo(HaveOccurred())
		_, err = mk(user, "work-acme!")
		Expect(err).To(MatchError(ContainSubstring("already have a connection")))
	})
	It("refuses a name with no letter or digit", func() {
		_, err := mk(user, "---")
		Expect(err).To(MatchError(ContainSubstring("letter or digit")))
	})
	It("lets two people use the same name", func() {
		_, err := mk(user, "Work")
		Expect(err).NotTo(HaveOccurred())
		_, err = mk("bob", "Work")
		Expect(err).NotTo(HaveOccurred())
	})
	It("refuses a rename onto another name, and allows keeping its own", func() {
		a, err := mk(user, "Work")
		Expect(err).NotTo(HaveOccurred())
		b, err := mk(user, "Home")
		Expect(err).NotTo(HaveOccurred())
		_, err = conns.UpdateConnector(ctx, user, b.ID, studio.ConnectorInput{Name: "WORK", Resources: []string{"a/b"}})
		Expect(err).To(MatchError(ContainSubstring("already have a connection")))
		_, err = conns.UpdateConnector(ctx, user, a.ID, studio.ConnectorInput{Name: "work", Resources: []string{"a/b"}})
		Expect(err).NotTo(HaveOccurred(), "a connection may change its own capitalisation")
	})
})