mirror of
https://github.com/mudler/LocalAI.git
synced 2026-11-06 12:33:16 -05:00
feat(connectors): store a file limit per binding
A binding keeps an optional limit for the files it reads and the files it sends. Zero means the owner sets none. Both stores keep the values and the validation bounds them. Assisted-by: Claude Code:claude-sonnet-5-5 Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
4 files changed
+41
-14
No files matched your search
@@ -94,6 +94,15 @@ func Run(newStore NewStore) {
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(got.ReadsThread()).To(BeFalse())
|
||||
})
|
||||
It("keeps the file limits an owner set for a binding", func() {
|
||||
s := newStore()
|
||||
b := binding("limits", "u1", "Limits")
|
||||
b.FileReceiveMax, b.FileSendMax = 512<<10, 2<<20
|
||||
Expect(s.PutBinding(ctx, b)).To(Succeed())
|
||||
got, err := s.GetBinding(ctx, "limits")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect([]int64{got.FileReceiveMax, got.FileSendMax}).To(Equal([]int64{512 << 10, 2 << 20}))
|
||||
})
|
||||
It("gives an old binding the add-on of its kind", func() {
|
||||
// A binding stored before connectors became add-ons has no add-on ID.
|
||||
s := newStore()
|
||||
|
||||
@@ -33,9 +33,13 @@ type bindingRow struct {
|
||||
ThreadContextOff bool
|
||||
// Options is JSON, added by AutoMigrate.
|
||||
Options string
|
||||
Token []byte
|
||||
Created time.Time
|
||||
Updated time.Time
|
||||
// The columns of the file limits are added by AutoMigrate; a row stored
|
||||
// before them has no limit of the owner.
|
||||
FileReceiveMax int64
|
||||
FileSendMax int64
|
||||
Token []byte
|
||||
Created time.Time
|
||||
Updated time.Time
|
||||
}
|
||||
|
||||
func (bindingRow) TableName() string { return "studio_connectors" }
|
||||
@@ -90,6 +94,7 @@ func NewGORM(db *gorm.DB) (Store, error) {
|
||||
func (r bindingRow) binding() (Binding, error) {
|
||||
b := Binding{ID: r.ID, Owner: r.Owner, Kind: r.Kind, AddonID: r.AddonID, Name: r.Name, AgentID: r.AgentID, Enabled: r.Enabled, AcceptAny: r.AcceptAny,
|
||||
NotifyChat: r.NotifyChat, Reply: r.Reply, Conversation: r.Conversation, WindowMinutes: r.WindowMinutes, FilesOff: r.FilesOff, ThreadContextOff: r.ThreadContextOff,
|
||||
FileReceiveMax: r.FileReceiveMax, FileSendMax: r.FileSendMax,
|
||||
Token: append([]byte(nil), r.Token...), Created: r.Created.UTC(), Updated: r.Updated.UTC()}
|
||||
if r.Allow != "" {
|
||||
if err := json.Unmarshal([]byte(r.Allow), &b.Allow); err != nil {
|
||||
@@ -117,7 +122,8 @@ func (s *gormStore) PutBinding(ctx context.Context, b Binding) error {
|
||||
}
|
||||
return s.db.WithContext(ctx).Save(&bindingRow{ID: b.ID, Owner: b.Owner, Kind: b.Kind, AddonID: b.AddonID, Name: b.Name, AgentID: b.AgentID, Enabled: b.Enabled,
|
||||
AcceptAny: b.AcceptAny, Allow: allow, NotifyChat: b.NotifyChat, Reply: b.Reply, Conversation: b.Conversation, WindowMinutes: b.WindowMinutes,
|
||||
FilesOff: b.FilesOff, ThreadContextOff: b.ThreadContextOff, Options: options, Token: b.Token, Created: b.Created, Updated: b.Updated}).Error
|
||||
FilesOff: b.FilesOff, ThreadContextOff: b.ThreadContextOff, Options: options, FileReceiveMax: b.FileReceiveMax, FileSendMax: b.FileSendMax,
|
||||
Token: b.Token, Created: b.Created, Updated: b.Updated}).Error
|
||||
}
|
||||
|
||||
func (s *gormStore) GetBinding(ctx context.Context, id string) (Binding, error) {
|
||||
|
||||
@@ -46,16 +46,18 @@ var _ = Describe("sealing and validating bindings", func() {
|
||||
many[i] = "1"
|
||||
}
|
||||
for name, b := range map[string]connectors.Binding{
|
||||
"no name": {Kind: connectors.KindTelegram, AgentID: "a"},
|
||||
"long name": {Name: strings.Repeat("n", 81), Kind: connectors.KindTelegram, AgentID: "a"},
|
||||
"bad kind": {Name: "b", Kind: "Not A Kind", AgentID: "a"},
|
||||
"no agent": {Name: "b", Kind: connectors.KindTelegram},
|
||||
"bad allow": {Name: "b", Kind: connectors.KindTelegram, AgentID: "a", Allow: []string{"alice"}},
|
||||
"short name": {Name: "b", Kind: connectors.KindTelegram, AgentID: "a", Allow: []string{"@ab"}},
|
||||
"many": {Name: "b", Kind: connectors.KindTelegram, AgentID: "a", Allow: many},
|
||||
"bad notify": {Name: "b", Kind: connectors.KindTelegram, AgentID: "a", NotifyChat: "me"},
|
||||
"lower id": {Name: "b", Kind: "slack", AgentID: "a", Allow: []string{"u012ab"}},
|
||||
"id notify": {Name: "b", Kind: "slack", AgentID: "a", NotifyChat: "@ada"},
|
||||
"no name": {Kind: connectors.KindTelegram, AgentID: "a"},
|
||||
"long name": {Name: strings.Repeat("n", 81), Kind: connectors.KindTelegram, AgentID: "a"},
|
||||
"bad kind": {Name: "b", Kind: "Not A Kind", AgentID: "a"},
|
||||
"no agent": {Name: "b", Kind: connectors.KindTelegram},
|
||||
"bad allow": {Name: "b", Kind: connectors.KindTelegram, AgentID: "a", Allow: []string{"alice"}},
|
||||
"short name": {Name: "b", Kind: connectors.KindTelegram, AgentID: "a", Allow: []string{"@ab"}},
|
||||
"many": {Name: "b", Kind: connectors.KindTelegram, AgentID: "a", Allow: many},
|
||||
"bad notify": {Name: "b", Kind: connectors.KindTelegram, AgentID: "a", NotifyChat: "me"},
|
||||
"lower id": {Name: "b", Kind: "slack", AgentID: "a", Allow: []string{"u012ab"}},
|
||||
"id notify": {Name: "b", Kind: "slack", AgentID: "a", NotifyChat: "@ada"},
|
||||
"negative receive limit": {Name: "b", Kind: "slack", AgentID: "a", FileReceiveMax: -1},
|
||||
"huge send limit": {Name: "b", Kind: "slack", AgentID: "a", FileSendMax: connectors.MaxFileLimit + 1},
|
||||
} {
|
||||
Expect(connectors.Validate(b)).To(MatchError(connectors.ErrInvalid), name)
|
||||
}
|
||||
|
||||
@@ -68,6 +68,8 @@ const (
|
||||
MaxSendersKept = 200
|
||||
MaxOptions = 8
|
||||
MaxOptionRunes = 32
|
||||
// MaxFileLimit is the largest file limit a binding can hold: 1 TiB.
|
||||
MaxFileLimit int64 = 1 << 40
|
||||
)
|
||||
|
||||
// Errors.
|
||||
@@ -116,6 +118,11 @@ type Binding struct {
|
||||
// the binding declares (contract.ConnectorKind.Options), by option ID. An
|
||||
// option that is not here has its default.
|
||||
Options map[string]string
|
||||
// FileReceiveMax and FileSendMax are the largest file, in bytes, that the
|
||||
// owner lets the bot read and send. Zero means no limit of the owner: the
|
||||
// limit of the kind applies. A value can only lower what the kind allows;
|
||||
// the host takes the smaller one.
|
||||
FileReceiveMax, FileSendMax int64
|
||||
// Token is the bot token. In a Store it is sealed bytes; the Store returns
|
||||
// what it was given.
|
||||
Token []byte
|
||||
@@ -193,6 +200,9 @@ func Validate(b Binding) error {
|
||||
return bad("%q is not a chat ID, user ID or @username", a)
|
||||
}
|
||||
}
|
||||
if b.FileReceiveMax < 0 || b.FileReceiveMax > MaxFileLimit || b.FileSendMax < 0 || b.FileSendMax > MaxFileLimit {
|
||||
return bad("a file limit must be 0 (the limit of the kind) or 1 to %d bytes", MaxFileLimit)
|
||||
}
|
||||
if b.NotifyChat != "" && !chatShape.MatchString(b.NotifyChat) {
|
||||
return bad("the notification chat must be a chat ID")
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user