Merge pull request #6 from KongZ/optmize-mem

Optimize memory uage

Author: KongZ

pkg/agent/conversation.go

@@ -121,6 +121,9 @@ type Agent struct {
 	// SessionBackend is the configured backend for session persistence (e.g., memory, filesystem).
 	SessionBackend string
 
+	// sessionManager is a cached SessionManager instance, initialized once in Init.
+	sessionManager *sessions.SessionManager
+
 	// lastErr is the most recent error run into, for use across the stack
 	lastErr error
 
@@ -134,6 +137,18 @@ type Agent struct {
 // Assert InMemoryChatStore implements ChatMessageStore
 var _ api.ChatMessageStore = &sessions.InMemoryChatStore{}
 
+// getSessionManager returns the cached session manager, lazily initializing it if needed.
+func (c *Agent) getSessionManager() (*sessions.SessionManager, error) {
+	if c.sessionManager == nil {
+		sm, err := sessions.NewSessionManager(c.SessionBackend)
+		if err != nil {
+			return nil, fmt.Errorf("failed to create session manager: %w", err)
+		}
+		c.sessionManager = sm
+	}
+	return c.sessionManager, nil
+}
+
 func (s *Agent) GetSession() *api.Session {
 	s.sessionMu.Lock()
 	defer s.sessionMu.Unlock()
@@ -156,7 +171,6 @@ func (c *Agent) addMessage(source api.MessageSource, messageType api.MessageType
 		Type:      messageType,
 		Payload:   payload,
 		Timestamp: time.Now(),
-		Metadata:  make(map[string]string),
 	}
 
 	// session should always have a ChatMessageStore at this point
@@ -225,6 +239,13 @@ func (s *Agent) Init(ctx context.Context) error {
 		return fmt.Errorf("agent requires a session to be provided")
 	}
 
+	// Initialize session manager for reuse across agent methods
+	sessionMgr, err := sessions.NewSessionManager(s.SessionBackend)
+	if err != nil {
+		return fmt.Errorf("failed to create session manager: %w", err)
+	}
+	s.sessionManager = sessionMgr
+
 	// Create a session working directory in the user's home directory
 	// to avoid read-only filesystem issues in containers
 	homeDir, err := os.UserHomeDir()
@@ -595,7 +616,7 @@ func (c *Agent) Run(ctx context.Context, initialQuery string) error {
 					msg := c.addMessage(api.MessageSourceModel, api.MessageTypeText, streamedText)
 					// If no function calls to be made, this is the final message of the turn
 					if len(functionCalls) == 0 {
-						msg.Metadata["is_final"] = "true"
+						msg.SetMetadata("is_final", "true")
 					}
 				}
 				// If no function calls to be made, we're done
@@ -612,7 +633,7 @@ func (c *Agent) Run(ctx context.Context, initialQuery string) error {
 						// IMPORTANT: This also prevents UIs from getting blocked on reading from the output channel.
 						log.V(2).Info("Empty response with no tool calls from LLM.")
 						msg := c.addMessage(api.MessageSourceAgent, api.MessageTypeText, "Empty response from LLM")
-						msg.Metadata["is_final"] = "true"
+						msg.SetMetadata("is_final", "true")
 					}
 					continue
 				}
@@ -788,12 +809,11 @@ func (c *Agent) handleMetaQuery(ctx context.Context, query string) (answer strin
 		return "Saved session as " + savedSessionID, true, nil
 
 	case "sessions":
-		manager, err := sessions.NewSessionManager(c.SessionBackend)
+		mgr, err := c.getSessionManager()
 		if err != nil {
-			return "", false, fmt.Errorf("failed to create session manager: %w", err)
+			return "", false, err
 		}
-
-		sessionList, err := manager.ListSessions()
+		sessionList, err := mgr.ListSessions()
 		if err != nil {
 			return "", false, fmt.Errorf("failed to list sessions: %w", err)
 		}
@@ -840,11 +860,6 @@ func (c *Agent) NewSession() (string, error) {
 		return "", fmt.Errorf("failed to save current session: %w", err)
 	}
 
-	manager, err := sessions.NewSessionManager(c.SessionBackend)
-	if err != nil {
-		return "", fmt.Errorf("failed to create session manager: %w", err)
-	}
-
 	metadata := sessions.Metadata{
 		ModelID:    c.Model,
 		ProviderID: c.Provider,
@@ -853,7 +868,11 @@ func (c *Agent) NewSession() (string, error) {
 		metadata.SlackUserID = c.Session.SlackUserID
 	}
 
-	newSession, err := manager.NewSession(metadata)
+	mgr, err := c.getSessionManager()
+	if err != nil {
+		return "", err
+	}
+	newSession, err := mgr.NewSession(metadata)
 	if err != nil {
 		return "", fmt.Errorf("failed to create new session: %w", err)
 	}
@@ -910,13 +929,13 @@ func (c *Agent) SaveSession() (string, error) {
 	c.sessionMu.Lock()
 	defer c.sessionMu.Unlock()
 
-	manager, err := sessions.NewSessionManager(c.SessionBackend)
+	mgr, err := c.getSessionManager()
 	if err != nil {
-		return "", fmt.Errorf("failed to create session manager: %w", err)
+		return "", err
 	}
 
 	if c.Session != nil {
-		foundSession, _ := manager.FindSessionByID(c.Session.ID)
+		foundSession, _ := mgr.FindSessionByID(c.Session.ID)
 		if foundSession != nil {
 			return foundSession.ID, nil
 		}
@@ -930,7 +949,7 @@ func (c *Agent) SaveSession() (string, error) {
 		SlackUserID:  c.Session.SlackUserID,
 	}
 
-	newSession, err := manager.NewSession(metadata)
+	newSession, err := mgr.NewSession(metadata)
 	if err != nil {
 		return "", fmt.Errorf("failed to create new session: %w", err)
 	}
@@ -953,14 +972,14 @@ func (c *Agent) SaveSession() (string, error) {
 
 // LoadSession loads a session by ID (or latest), updates the agent's state, and re-initializes the chat.
 func (c *Agent) LoadSession(sessionID string) error {
-	manager, err := sessions.NewSessionManager(c.SessionBackend)
+	mgr, err := c.getSessionManager()
 	if err != nil {
-		return fmt.Errorf("failed to create session manager: %w", err)
+		return err
 	}
 
 	var session *api.Session
 	if sessionID == "" || sessionID == "latest" {
-		s, err := manager.GetLatestSession()
+		s, err := mgr.GetLatestSession()
 		if err != nil {
 			return fmt.Errorf("failed to get latest session: %w", err)
 		}
@@ -969,7 +988,7 @@ func (c *Agent) LoadSession(sessionID string) error {
 		}
 		session = s
 	} else {
-		s, err := manager.FindSessionByID(sessionID)
+		s, err := mgr.FindSessionByID(sessionID)
 		if err != nil {
 			return fmt.Errorf("failed to get session %q: %w", sessionID, err)
 		}
@@ -993,7 +1012,7 @@ func (c *Agent) LoadSession(sessionID string) error {
 		c.Session.AgentState = api.AgentStateIdle
 	}
 
-	if err := manager.UpdateLastAccessed(session); err != nil {
+	if err := mgr.UpdateLastAccessed(session); err != nil {
 		return fmt.Errorf("failed to update session metadata: %w", err)
 	}
 
@@ -1262,7 +1281,7 @@ func toMap(v any) (map[string]any, error) {
 
 func candidateToShimCandidate(iterator gollm.ChatResponseIterator) (gollm.ChatResponseIterator, error) {
 	return func(yield func(gollm.ChatResponse, error) bool) {
-		buffer := ""
+		var buf strings.Builder
 		for response, err := range iterator {
 			if err != nil {
 				yield(nil, err)
@@ -1278,7 +1297,7 @@ func candidateToShimCandidate(iterator gollm.ChatResponseIterator) (gollm.ChatRe
 
 			for _, part := range candidate.Parts() {
 				if text, ok := part.AsText(); ok {
-					buffer += text
+					buf.WriteString(text)
 					klog.Infof("text is %q", text)
 				} else {
 					yield(nil, fmt.Errorf("no text part found in candidate"))
@@ -1287,6 +1306,7 @@ func candidateToShimCandidate(iterator gollm.ChatResponseIterator) (gollm.ChatRe
 			}
 		}
 
+		buffer := buf.String()
 		if buffer == "" {
 			yield(nil, nil)
 			return

pkg/agent/manager.go

@@ -19,6 +19,7 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"time"
 
 	"github.com/KongZ/kubeai-chatbot/pkg/api"
 	"github.com/KongZ/kubeai-chatbot/pkg/sessions"
@@ -28,22 +29,30 @@ import (
 // Factory is a function that creates a new Agent instance.
 type Factory func(context.Context) (*Agent, error)
 
+// defaultIdleTimeout is how long an agent can sit idle before being evicted.
+const defaultIdleTimeout = 30 * time.Minute
+
 // AgentManager manages the lifecycle of agents and their sessions.
 type AgentManager struct {
 	factory        Factory
 	sessionManager *sessions.SessionManager
 	agents         map[string]*Agent // sessionID -> agent
 	mu             sync.RWMutex
 	onAgentCreated func(*Agent)
+	stopCh         chan struct{}
+	closeOnce      sync.Once
 }
 
 // NewAgentManager creates a new Manager.
 func NewAgentManager(factory Factory, sessionManager *sessions.SessionManager) *AgentManager {
-	return &AgentManager{
+	m := &AgentManager{
 		factory:        factory,
 		sessionManager: sessionManager,
 		agents:         make(map[string]*Agent),
+		stopCh:         make(chan struct{}),
 	}
+	go m.evictionLoop(defaultIdleTimeout)
+	return m
 }
 
 // SetAgentCreatedCallback sets the callback to be called when a new agent is created.
@@ -80,8 +89,12 @@ func (sm *AgentManager) GetAgent(ctx context.Context, sessionID string) (*Agent,
 	return sm.startAgent(ctx, session, newAgent)
 }
 
-// Close closes all active agents.
+// Close closes all active agents and stops the eviction loop.
 func (sm *AgentManager) Close() error {
+	sm.closeOnce.Do(func() {
+		close(sm.stopCh)
+	})
+
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
 
@@ -96,6 +109,44 @@ func (sm *AgentManager) Close() error {
 	return nil
 }
 
+// evictionLoop periodically checks for and evicts idle agents.
+func (sm *AgentManager) evictionLoop(idleTimeout time.Duration) {
+	ticker := time.NewTicker(idleTimeout / 2)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			sm.evictIdleAgents(idleTimeout)
+		case <-sm.stopCh:
+			return
+		}
+	}
+}
+
+// evictIdleAgents removes agents that have been idle longer than the timeout.
+func (sm *AgentManager) evictIdleAgents(idleTimeout time.Duration) {
+	now := time.Now()
+
+	sm.mu.Lock()
+	var toClose []*Agent
+	for id, a := range sm.agents {
+		session := a.GetSession()
+		if now.Sub(session.LastModified) > idleTimeout {
+			klog.Infof("Evicting idle agent for session %s (idle for %v)", id, now.Sub(session.LastModified))
+			toClose = append(toClose, a)
+			delete(sm.agents, id)
+		}
+	}
+	sm.mu.Unlock()
+
+	// Close evicted agents outside the lock to avoid blocking
+	for _, a := range toClose {
+		if err := a.Close(); err != nil {
+			klog.Errorf("Error closing evicted agent: %v", err)
+		}
+	}
+}
+
 // ListSessions delegates to the underlying store.
 func (sm *AgentManager) ListSessions() ([]*api.Session, error) {
 	return sm.sessionManager.ListSessions()

pkg/api/models.go

@@ -89,6 +89,14 @@ func (m *Message) Validate() error {
 	return validate.Struct(m)
 }
 
+// SetMetadata sets a metadata key-value pair, lazily initializing the map if needed.
+func (m *Message) SetMetadata(key, value string) {
+	if m.Metadata == nil {
+		m.Metadata = make(map[string]string)
+	}
+	m.Metadata[key] = value
+}
+
 type MessageSource string
 
 const (

pkg/sessions/memory.go

@@ -117,6 +117,10 @@ func NewInMemoryChatStore() *InMemoryChatStore {
 	}
 }
 
+// maxChatMessages is the upper bound on in-memory messages per session.
+// When exceeded, the oldest messages are trimmed to prevent unbounded growth.
+const maxChatMessages = 500
+
 // AddChatMessage adds a message to the store.
 func (s *InMemoryChatStore) AddChatMessage(record *api.Message) error {
 	if err := record.Validate(); err != nil {
@@ -125,6 +129,10 @@ func (s *InMemoryChatStore) AddChatMessage(record *api.Message) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 	s.messages = append(s.messages, record)
+	// Trim oldest messages to prevent unbounded memory growth
+	if len(s.messages) > maxChatMessages {
+		s.messages = s.messages[len(s.messages)-maxChatMessages:]
+	}
 	return nil
 }
 

pkg/tools/kubectl_tool.go

@@ -21,13 +21,30 @@ import (
 	"os"
 	"os/exec"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/KongZ/kubeai-chatbot/gollm"
 	"github.com/KongZ/kubeai-chatbot/pkg/api"
 	"mvdan.cc/sh/v3/shell"
 )
 
+var (
+	cachedBaseEnv     []string
+	cachedBaseEnvOnce sync.Once
+)
+
+// baseEnviron returns a copy of the cached process environment.
+// The base environment is captured once to avoid repeated os.Environ() syscalls.
+func baseEnviron() []string {
+	cachedBaseEnvOnce.Do(func() {
+		cachedBaseEnv = os.Environ()
+	})
+	env := make([]string, len(cachedBaseEnv))
+	copy(env, cachedBaseEnv)
+	return env
+}
+
 type ExecResult struct {
 	Command    string `json:"command"`
 	Stdout     string `json:"stdout,omitempty"`
@@ -114,8 +131,8 @@ func (t *Kubectl) Run(ctx context.Context, args map[string]any) (any, error) {
 	// Reconstruct command for reporting
 	fullCommand := strings.Join(cmdArgs, " ")
 
-	// Prepare environment
-	env := os.Environ()
+	// Prepare environment (use cached base to avoid repeated syscalls)
+	env := baseEnviron()
 	if kubeconfig != "" {
 		kubeconfig, err := ExpandShellVar(kubeconfig)
 		if err != nil {