Merge pull request #266 from kiwamizamurai/issue-199

feat: Add num_history_items for limiting agent conversation memory

Author: rumpl

pkg/agent/agent.go

@@ -40,6 +40,7 @@ type Agent struct {
 	addDate            bool
 	addEnvironmentInfo bool
 	maxIterations      int
+	numHistoryItems    int
 	toolWrapper        toolWrapper
 	memoryManager      memorymanager.Manager
 }
@@ -80,6 +81,10 @@ func (a *Agent) MaxIterations() int {
 	return a.maxIterations
 }
 
+func (a *Agent) NumHistoryItems() int {
+	return a.numHistoryItems
+}
+
 // Description returns the agent's description
 func (a *Agent) Description() string {
 	return a.description

pkg/agent/opts.go

@@ -78,3 +78,9 @@ func WithMaxIterations(maxIterations int) Opt {
 		a.maxIterations = maxIterations
 	}
 }
+
+func WithNumHistoryItems(numHistoryItems int) Opt {
+	return func(a *Agent) {
+		a.numHistoryItems = numHistoryItems
+	}
+}

pkg/config/v0/types.go

@@ -77,16 +77,17 @@ func (t *TodoConfig) UnmarshalYAML(unmarshal func(any) error) error {
 
 // AgentConfig represents a single agent configuration
 type AgentConfig struct {
-	Name         string       `json:"name,omitempty" yaml:"name,omitempty"`
-	Model        string       `json:"model,omitempty" yaml:"model,omitempty"`
-	Description  string       `json:"description,omitempty" yaml:"description,omitempty"`
-	Toolsets     []Toolset    `json:"toolsets,omitempty" yaml:"toolsets,omitempty"`
-	Instruction  string       `json:"instruction,omitempty" yaml:"instruction,omitempty"`
-	SubAgents    []string     `json:"sub_agents,omitempty" yaml:"sub_agents,omitempty"`
-	AddDate      bool         `json:"add_date,omitempty" yaml:"add_date,omitempty"`
-	Think        bool         `json:"think,omitempty" yaml:"think,omitempty"`
-	Todo         TodoConfig   `json:"todo,omitempty" yaml:"todo,omitempty"`
-	MemoryConfig MemoryConfig `json:"memory,omitempty" yaml:"memory,omitempty"`
+	Name            string       `json:"name,omitempty" yaml:"name,omitempty"`
+	Model           string       `json:"model,omitempty" yaml:"model,omitempty"`
+	Description     string       `json:"description,omitempty" yaml:"description,omitempty"`
+	Toolsets        []Toolset    `json:"toolsets,omitempty" yaml:"toolsets,omitempty"`
+	Instruction     string       `json:"instruction,omitempty" yaml:"instruction,omitempty"`
+	SubAgents       []string     `json:"sub_agents,omitempty" yaml:"sub_agents,omitempty"`
+	AddDate         bool         `json:"add_date,omitempty" yaml:"add_date,omitempty"`
+	Think           bool         `json:"think,omitempty" yaml:"think,omitempty"`
+	Todo            TodoConfig   `json:"todo,omitempty" yaml:"todo,omitempty"`
+	MemoryConfig    MemoryConfig `json:"memory,omitempty" yaml:"memory,omitempty"`
+	NumHistoryItems int          `json:"num_history_items,omitempty" yaml:"num_history_items,omitempty"`
 }
 
 type MemoryConfig struct {

pkg/config/v1/types.go

@@ -110,6 +110,7 @@ type AgentConfig struct {
 	SubAgents          []string  `json:"sub_agents,omitempty" yaml:"sub_agents,omitempty"`
 	AddDate            bool      `json:"add_date,omitempty" yaml:"add_date,omitempty"`
 	AddEnvironmentInfo bool      `json:"add_environment_info,omitempty" yaml:"add_environment_info,omitempty"`
+	NumHistoryItems    int       `json:"num_history_items,omitempty" yaml:"num_history_items,omitempty"`
 }
 
 // ModelConfig represents the configuration for a model

pkg/config/v2/types.go

@@ -23,6 +23,7 @@ type AgentConfig struct {
 	AddDate            bool      `json:"add_date,omitempty" yaml:"add_date,omitempty"`
 	AddEnvironmentInfo bool      `json:"add_environment_info,omitempty" yaml:"add_environment_info,omitempty"`
 	MaxIterations      int       `json:"max_iterations,omitempty" yaml:"max_iterations,omitempty"`
+	NumHistoryItems    int       `json:"num_history_items,omitempty" yaml:"num_history_items,omitempty"`
 }
 
 // ModelConfig represents the configuration for a model

pkg/session/session.go

@@ -275,13 +275,31 @@ func (s *Session) GetMessages(a *agent.Agent) []chat.Message {
 		}
 	}
 
-	trimmed := trimMessages(messages)
+	maxItems := a.NumHistoryItems()
+	if maxItems <= 0 {
+		maxItems = maxMessages
+	}
+
+	trimmed := trimMessages(messages, maxItems)
+
+	systemCount := 0
+	conversationCount := 0
+	for i := range trimmed {
+		if trimmed[i].Role == chat.MessageRoleSystem {
+			systemCount++
+		} else {
+			conversationCount++
+		}
+	}
 
 	slog.Debug("Retrieved messages for agent",
 		"agent", a.Name(),
 		"session_id", s.ID,
 		"total_messages", len(messages),
-		"trimmed_messages", len(trimmed))
+		"trimmed_total", len(trimmed),
+		"system_messages", systemCount,
+		"conversation_messages", conversationCount,
+		"max_history_items", maxItems)
 
 	return trimmed
 }
@@ -304,32 +322,51 @@ func (s *Session) GetMostRecentAgentFilename() string {
 }
 
 // trimMessages ensures we don't exceed the maximum number of messages while maintaining
-// consistency between assistant messages and their tool call results
-func trimMessages(messages []chat.Message) []chat.Message {
-	if len(messages) <= maxMessages {
+// consistency between assistant messages and their tool call results.
+// System messages are always preserved and not counted against the limit.
+func trimMessages(messages []chat.Message, maxItems int) []chat.Message {
+	// Separate system messages from conversation messages
+	var systemMessages []chat.Message
+	var conversationMessages []chat.Message
+
+	for i := range messages {
+		if messages[i].Role == chat.MessageRoleSystem {
+			systemMessages = append(systemMessages, messages[i])
+		} else {
+			conversationMessages = append(conversationMessages, messages[i])
+		}
+	}
+
+	// If conversation messages fit within limit, return all messages
+	if len(conversationMessages) <= maxItems {
 		return messages
 	}
 
 	// Keep track of tool call IDs that need to be removed
 	toolCallsToRemove := make(map[string]bool)
 
-	// Calculate how many messages we need to remove
-	toRemove := len(messages) - maxMessages
+	// Calculate how many conversation messages we need to remove
+	toRemove := len(conversationMessages) - maxItems
 
 	// Start from the beginning (oldest messages)
 	for i := range toRemove {
 		// If this is an assistant message with tool calls, mark them for removal
-		if messages[i].Role == chat.MessageRoleAssistant {
-			for _, toolCall := range messages[i].ToolCalls {
+		if conversationMessages[i].Role == chat.MessageRoleAssistant {
+			for _, toolCall := range conversationMessages[i].ToolCalls {
 				toolCallsToRemove[toolCall.ID] = true
 			}
 		}
 	}
 
-	// Filter messages keeping only those we want to keep
-	result := make([]chat.Message, 0, maxMessages)
-	for i := toRemove; i < len(messages); i++ {
-		msg := messages[i]
+	// Combine system messages with trimmed conversation messages
+	result := make([]chat.Message, 0, len(systemMessages)+maxItems)
+
+	// Add all system messages first
+	result = append(result, systemMessages...)
+
+	// Add the most recent conversation messages
+	for i := toRemove; i < len(conversationMessages); i++ {
+		msg := conversationMessages[i]
 
 		// Skip tool messages that correspond to removed assistant messages
 		if msg.Role == chat.MessageRoleTool && toolCallsToRemove[msg.ToolCallID] {