Feature: Enhance Performance, Logging & Stability: Precompile Regex, Command Matching, Golang 1.24, History Cleanup & memLimitMiB Flag. (#182)

* Feat: Add support for logging which "command" was matched for SSH and HTTP strategies.

* Feat: Convert to precompiling regexp at config load time. This allows for errors to be presented to the user during startup, and provides better performance for complex regexp.

* Feat:Bump Golang version to latest stable 1.24

* Feat: Add a cleanup routine for HistoryStore, default TTL for events is 1 hour since last interaction.

* Feat: Add new command line flag "memLimitMiB" with a default value of 100.

---------

Signed-off-by: Bryan Nolen <bryan@arc.net.au>
Signed-off-by: Mario Candela <mario.candela.personal@gmail.com>
Co-authored-by: Mario Candela <mario.candela.personal@gmail.com>

historystore/history_store.go

@@ -2,20 +2,32 @@ package historystore
 
 import (
 	"sync"
+	"time"
 
 	"github.com/mariocandela/beelzebub/v3/plugins"
 )
 
+var (
+	MaxHistoryAge   = 60 * time.Minute
+	CleanerInterval = 1 * time.Minute
+)
+
 // HistoryStore is a thread-safe structure for storing Messages used to build LLM Context.
 type HistoryStore struct {
 	sync.RWMutex
-	sessions map[string][]plugins.Message
+	sessions map[string]HistoryEvent
+}
+
+// HistoryEvent is a container for storing messages
+type HistoryEvent struct {
+	LastSeen time.Time
+	Messages []plugins.Message
 }
 
 // NewHistoryStore returns a prepared HistoryStore
 func NewHistoryStore() *HistoryStore {
 	return &HistoryStore{
-		sessions: make(map[string][]plugins.Message),
+		sessions: make(map[string]HistoryEvent),
 	}
 }
 
@@ -31,7 +43,7 @@ func (hs *HistoryStore) HasKey(key string) bool {
 func (hs *HistoryStore) Query(key string) []plugins.Message {
 	hs.RLock()
 	defer hs.RUnlock()
-	return hs.sessions[key]
+	return hs.sessions[key].Messages
 }
 
 // Append will add the slice of Mesages to the entry for the key.
@@ -41,7 +53,30 @@ func (hs *HistoryStore) Append(key string, message ...plugins.Message) {
 	defer hs.Unlock()
 	// In the unexpected case that the map has not yet been initalised, create it.
 	if hs.sessions == nil {
-		hs.sessions = make(map[string][]plugins.Message)
+		hs.sessions = make(map[string]HistoryEvent)
 	}
-	hs.sessions[key] = append(hs.sessions[key], message...)
+	e, ok := hs.sessions[key]
+	if !ok {
+		e = HistoryEvent{}
+	}
+	e.LastSeen = time.Now()
+	e.Messages = append(e.Messages, message...)
+	hs.sessions[key] = e
+}
+
+// HistoryCleaner is a function that will periodically remove records from the HistoryStore
+// that are older than MaxHistoryAge.
+func (hs *HistoryStore) HistoryCleaner() {
+	cleanerTicker := time.NewTicker(CleanerInterval)
+	go func() {
+		for range cleanerTicker.C {
+			hs.Lock()
+			for k, v := range hs.sessions {
+				if time.Since(v.LastSeen) > MaxHistoryAge {
+					delete(hs.sessions, k)
+				}
+			}
+			hs.Unlock()
+		}
+	}()
 }

historystore/history_store_test.go

@@ -2,6 +2,7 @@ package historystore
 
 import (
 	"testing"
+	"time"
 
 	"github.com/mariocandela/beelzebub/v3/plugins"
 	"github.com/stretchr/testify/assert"
@@ -15,7 +16,7 @@ func TestNewHistoryStore(t *testing.T) {
 
 func TestHasKey(t *testing.T) {
 	hs := NewHistoryStore()
-	hs.sessions["testKey"] = []plugins.Message{}
+	hs.sessions["testKey"] = HistoryEvent{Messages: []plugins.Message{}}
 	assert.True(t, hs.HasKey("testKey"))
 	assert.False(t, hs.HasKey("nonExistentKey"))
 }
@@ -23,7 +24,7 @@ func TestHasKey(t *testing.T) {
 func TestQuery(t *testing.T) {
 	hs := NewHistoryStore()
 	expectedMessages := []plugins.Message{{Role: "user", Content: "Hello"}}
-	hs.sessions["testKey"] = expectedMessages
+	hs.sessions["testKey"] = HistoryEvent{Messages: expectedMessages}
 	actualMessages := hs.Query("testKey")
 	assert.Equal(t, expectedMessages, actualMessages)
 }
@@ -33,9 +34,9 @@ func TestAppend(t *testing.T) {
 	message1 := plugins.Message{Role: "user", Content: "Hello"}
 	message2 := plugins.Message{Role: "assistant", Content: "Hi"}
 	hs.Append("testKey", message1)
-	assert.Equal(t, []plugins.Message{message1}, hs.sessions["testKey"])
+	assert.Equal(t, []plugins.Message{message1}, hs.sessions["testKey"].Messages)
 	hs.Append("testKey", message2)
-	assert.Equal(t, []plugins.Message{message1, message2}, hs.sessions["testKey"])
+	assert.Equal(t, []plugins.Message{message1, message2}, hs.sessions["testKey"].Messages)
 }
 
 func TestAppendNilSessions(t *testing.T) {
@@ -43,5 +44,23 @@ func TestAppendNilSessions(t *testing.T) {
 	message1 := plugins.Message{Role: "user", Content: "Hello"}
 	hs.Append("testKey", message1)
 	assert.NotNil(t, hs.sessions)
-	assert.Equal(t, []plugins.Message{message1}, hs.sessions["testKey"])
+	assert.Equal(t, []plugins.Message{message1}, hs.sessions["testKey"].Messages)
+}
+
+func TestHistoryCleaner(t *testing.T) {
+	hs := NewHistoryStore()
+	hs.Append("testKey", plugins.Message{Role: "user", Content: "Hello"})
+	hs.Append("testKey2", plugins.Message{Role: "user", Content: "Hello"})
+
+	// Make key older than MaxHistoryAge
+	e := hs.sessions["testKey"]
+	e.LastSeen = time.Now().Add(-MaxHistoryAge * 2)
+	hs.sessions["testKey"] = e
+
+	CleanerInterval = 5 * time.Second // Override for the test.
+	hs.HistoryCleaner()
+	time.Sleep(CleanerInterval + (1 * time.Second))
+
+	assert.False(t, hs.HasKey("testKey"))
+	assert.True(t, hs.HasKey("testKey2"))
 }