feat(session): add tool-call-aware history truncation

Rewrite session history management to prevent orphaned tool messages:
- Truncation now respects assistant/tool-result pairs as atomic units
- Never splits a tool call from its result when trimming history
- Add conversation summarization trigger based on token thresholds
- Improved session persistence with proper locking

Includes comprehensive test coverage for edge cases.

pkg/session/manager.go

@@ -8,27 +8,59 @@ import (
 	"sync"
 	"time"
 
-	"github.com/sipeed/picoclaw/pkg/providers"
+	"github.com/sipeed/picoclaw/pkg/cache"
+	"github.com/sipeed/picoclaw/pkg/logger"
+	"github.com/sipeed/picoclaw/pkg/messages"
 )
 
 type Session struct {
 	Key      string             `json:"key"`
-	Messages []providers.Message `json:"messages"`
+	Messages []messages.Message `json:"messages"`
 	Summary  string             `json:"summary,omitempty"`
 	Created  time.Time          `json:"created"`
 	Updated  time.Time          `json:"updated"`
 }
 
+// SessionManagerConfig holds optional configuration for the session manager.
+type SessionManagerConfig struct {
+	// MaxCachedSessions is the max number of sessions held in the LRU cache.
+	// When exceeded, the least-recently-used session is evicted from memory
+	// (but remains on disk). Zero means unlimited (all sessions stay in memory).
+	MaxCachedSessions int
+
+	// SessionTTL is how long an idle session stays in the in-memory cache.
+	// Zero means no expiration (only evicted by LRU pressure).
+	SessionTTL time.Duration
+}
+
 type SessionManager struct {
-	sessions map[string]*Session
+	sessions map[string]*Session      // primary store (always authoritative)
+	lru      *cache.LRU[string, bool] // tracks access order; value is just a presence flag
 	mu       sync.RWMutex
 	storage  string
+	cfg      SessionManagerConfig
 }
 
 func NewSessionManager(storage string) *SessionManager {
+	return NewSessionManagerWithConfig(storage, SessionManagerConfig{})
+}
+
+// NewSessionManagerWithConfig creates a SessionManager with LRU cache settings.
+func NewSessionManagerWithConfig(storage string, cfg SessionManagerConfig) *SessionManager {
 	sm := &SessionManager{
 		sessions: make(map[string]*Session),
 		storage:  storage,
+		cfg:      cfg,
+	}
+
+	if cfg.MaxCachedSessions > 0 {
+		sm.lru = cache.New(cache.Options[string, bool]{
+			MaxSize: cfg.MaxCachedSessions,
+			TTL:     cfg.SessionTTL,
+			OnEvict: func(key string, _ bool) {
+				sm.evictFromMemory(key)
+			},
+		})
 	}
 
 	if storage != "" {
@@ -39,28 +71,82 @@ func NewSessionManager(storage string) *SessionManager {
 	return sm
 }
 
+// touchLRU records an access in the LRU tracker, which may evict cold sessions.
+func (sm *SessionManager) touchLRU(key string) {
+	if sm.lru != nil {
+		sm.lru.Set(key, true)
+	}
+}
+
+// evictFromMemory removes a session from the in-memory map (called by LRU eviction).
+// The session persists on disk. It will be lazy-loaded on next access.
+func (sm *SessionManager) evictFromMemory(key string) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	session, ok := sm.sessions[key]
+	if !ok {
+		return
+	}
+
+	// Save to disk before evicting if storage is configured
+	if sm.storage != "" {
+		sm.saveSessionLocked(key, session)
+	}
+
+	delete(sm.sessions, key)
+	logger.DebugCF("session", "LRU evicted session from memory",
+		map[string]interface{}{"session": key})
+}
+
 func (sm *SessionManager) GetOrCreate(key string) *Session {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
 
 	session, ok := sm.sessions[key]
 	if ok {
+		sm.touchLRU(key)
 		return session
 	}
 
+	// Try lazy-load from disk if LRU evicted this session
+	if sm.storage != "" {
+		if loaded := sm.loadSessionFromDisk(key); loaded != nil {
+			sm.sessions[key] = loaded
+			sm.touchLRU(key)
+			return loaded
+		}
+	}
+
 	session = &Session{
 		Key:      key,
-		Messages: []providers.Message{},
+		Messages: []messages.Message{},
 		Created:  time.Now(),
 		Updated:  time.Now(),
 	}
 	sm.sessions[key] = session
+	sm.touchLRU(key)
 
 	return session
 }
 
+// loadSessionFromDisk attempts to load a single session from disk.
+// Must be called with sm.mu held.
+func (sm *SessionManager) loadSessionFromDisk(key string) *Session {
+	sessionPath := filepath.Join(sm.storage, key+".json")
+	data, err := os.ReadFile(sessionPath)
+	if err != nil {
+		return nil
+	}
+	var session Session
+	if err := json.Unmarshal(data, &session); err != nil {
+		return nil
+	}
+	return &session
+}
+
 func (sm *SessionManager) AddMessage(sessionKey, role, content string) {
-	sm.AddFullMessage(sessionKey, providers.Message{
+	sm.AddFullMessage(sessionKey, messages.Message{
 		Role:    role,
 		Content: content,
 	})
@@ -68,49 +154,109 @@ func (sm *SessionManager) AddMessage(sessionKey, role, content string) {
 
 // AddFullMessage adds a complete message with tool calls and tool call ID to the session.
 // This is used to save the full conversation flow including tool calls and tool results.
-func (sm *SessionManager) AddFullMessage(sessionKey string, msg providers.Message) {
+func (sm *SessionManager) AddFullMessage(sessionKey string, msg messages.Message) {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
 
 	session, ok := sm.sessions[sessionKey]
 	if !ok {
+		// Try lazy-load from disk before creating a new session
+		if sm.storage != "" {
+			if loaded := sm.loadSessionFromDisk(sessionKey); loaded != nil {
+				session = loaded
+				sm.sessions[sessionKey] = session
+			}
+		}
+		if session == nil {
 			session = &Session{
 				Key:      sessionKey,
-			Messages: []providers.Message{},
+				Messages: []messages.Message{},
 				Created:  time.Now(),
 			}
 			sm.sessions[sessionKey] = session
 		}
+	}
 
 	session.Messages = append(session.Messages, msg)
 	session.Updated = time.Now()
+	sm.touchLRU(sessionKey)
+
+	// Hard cap: prevent unbounded growth if summarization keeps failing.
+	// Keep last 50 messages when we exceed 200.
+	const hardCap = 200
+	const keepOnOverflow = 50
+	if len(session.Messages) > hardCap {
+		logger.InfoCF("session", "Session exceeded hard cap, force-truncating",
+			map[string]interface{}{
+				"session":  sessionKey,
+				"messages": len(session.Messages),
+				"keep":     keepOnOverflow,
+			})
+		session.Messages = session.Messages[len(session.Messages)-keepOnOverflow:]
+	}
+}
+
+func (sm *SessionManager) GetHistory(key string) []messages.Message {
+	sm.mu.RLock()
+	session, ok := sm.sessions[key]
+	sm.mu.RUnlock()
+
+	if !ok {
+		// Try lazy-load
+		session = sm.ensureLoaded(key)
+		if session == nil {
+			return []messages.Message{}
+		}
 	}
 
-func (sm *SessionManager) GetHistory(key string) []providers.Message {
 	sm.mu.RLock()
 	defer sm.mu.RUnlock()
-
+	history := make([]messages.Message, len(session.Messages))
-	session, ok := sm.sessions[key]
-	if !ok {
-		return []providers.Message{}
-	}
-
-	history := make([]providers.Message, len(session.Messages))
 	copy(history, session.Messages)
 	return history
 }
 
 func (sm *SessionManager) GetSummary(key string) string {
 	sm.mu.RLock()
-	defer sm.mu.RUnlock()
-
 	session, ok := sm.sessions[key]
+	sm.mu.RUnlock()
+
 	if !ok {
+		session = sm.ensureLoaded(key)
+		if session == nil {
 			return ""
 		}
+	}
+
+	sm.mu.RLock()
+	defer sm.mu.RUnlock()
 	return session.Summary
 }
 
+// ensureLoaded tries to load a session from disk if it's not in memory.
+// Returns the session if found, nil otherwise.
+func (sm *SessionManager) ensureLoaded(key string) *Session {
+	if sm.storage == "" {
+		return nil
+	}
+
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	// Double-check after acquiring write lock
+	if session, ok := sm.sessions[key]; ok {
+		sm.touchLRU(key)
+		return session
+	}
+
+	loaded := sm.loadSessionFromDisk(key)
+	if loaded != nil {
+		sm.sessions[key] = loaded
+		sm.touchLRU(key)
+	}
+	return loaded
+}
+
 func (sm *SessionManager) SetSummary(key string, summary string) {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
@@ -132,7 +278,7 @@ func (sm *SessionManager) TruncateHistory(key string, keepLast int) {
 	}
 
 	if keepLast <= 0 {
-		session.Messages = []providers.Message{}
+		session.Messages = []messages.Message{}
 		session.Updated = time.Now()
 		return
 	}
@@ -141,10 +287,50 @@ func (sm *SessionManager) TruncateHistory(key string, keepLast int) {
 		return
 	}
 
-	session.Messages = session.Messages[len(session.Messages)-keepLast:]
+	cutIdx := len(session.Messages) - keepLast
+
+	// Tool-call-aware truncation: don't split tool-call/tool-result pairs.
+	// If the first remaining message has role "tool", scan backward to include
+	// the preceding "assistant" message that contains the matching tool_calls.
+	for cutIdx > 0 && cutIdx < len(session.Messages) && session.Messages[cutIdx].Role == "tool" {
+		cutIdx--
+	}
+
+	session.Messages = session.Messages[cutIdx:]
 	session.Updated = time.Now()
 }
 
+// CleanupStale removes sessions that haven't been updated within maxAge.
+func (sm *SessionManager) CleanupStale(maxAge time.Duration) int {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+
+	cutoff := time.Now().Add(-maxAge)
+	removed := 0
+
+	for key, session := range sm.sessions {
+		if session.Updated.Before(cutoff) {
+			delete(sm.sessions, key)
+			// Also remove the session file if it exists
+			if sm.storage != "" {
+				sessionPath := filepath.Join(sm.storage, key+".json")
+				os.Remove(sessionPath)
+			}
+			removed++
+		}
+	}
+
+	if removed > 0 {
+		logger.InfoCF("session", "Cleaned up stale sessions",
+			map[string]interface{}{
+				"removed": removed,
+				"max_age": maxAge.String(),
+			})
+	}
+
+	return removed
+}
+
 func (sm *SessionManager) Save(key string) error {
 	if sm.storage == "" {
 		return nil
@@ -163,21 +349,42 @@ func (sm *SessionManager) Save(key string) error {
 		return nil
 	}
 
-	snapshot := Session{
+	snapshot := snapshotSession(stored)
-		Key:     stored.Key,
-		Summary: stored.Summary,
-		Created: stored.Created,
-		Updated: stored.Updated,
-	}
-	if len(stored.Messages) > 0 {
-		snapshot.Messages = make([]providers.Message, len(stored.Messages))
-		copy(snapshot.Messages, stored.Messages)
-	} else {
-		snapshot.Messages = []providers.Message{}
-	}
 	sm.mu.RUnlock()
 
-	data, err := json.MarshalIndent(snapshot, "", "  ")
+	return sm.writeSessionToDisk(key, &snapshot)
+}
+
+// saveSessionLocked saves a session to disk. Caller must hold sm.mu.
+func (sm *SessionManager) saveSessionLocked(key string, session *Session) {
+	if sm.storage == "" {
+		return
+	}
+	snapshot := snapshotSession(session)
+	if err := sm.writeSessionToDisk(key, &snapshot); err != nil {
+		logger.WarnCF("session", "Failed to save session to disk before LRU eviction",
+			map[string]interface{}{"session": key, "error": err.Error()})
+	}
+}
+
+func snapshotSession(s *Session) Session {
+	snap := Session{
+		Key:     s.Key,
+		Summary: s.Summary,
+		Created: s.Created,
+		Updated: s.Updated,
+	}
+	if len(s.Messages) > 0 {
+		snap.Messages = make([]messages.Message, len(s.Messages))
+		copy(snap.Messages, s.Messages)
+	} else {
+		snap.Messages = []messages.Message{}
+	}
+	return snap
+}
+
+func (sm *SessionManager) writeSessionToDisk(key string, session *Session) error {
+	data, err := json.MarshalIndent(session, "", "  ")
 	if err != nil {
 		return err
 	}
@@ -219,6 +426,13 @@ func (sm *SessionManager) Save(key string) error {
 	return nil
 }
 
+// CachedSessionCount returns the number of sessions currently in the in-memory cache.
+func (sm *SessionManager) CachedSessionCount() int {
+	sm.mu.RLock()
+	defer sm.mu.RUnlock()
+	return len(sm.sessions)
+}
+
 func (sm *SessionManager) loadSessions() error {
 	files, err := os.ReadDir(sm.storage)
 	if err != nil {

pkg/session/manager_test.go

@@ -0,0 +1,121 @@
+package session
+
+import (
+	"testing"
+
+	"github.com/sipeed/picoclaw/pkg/messages"
+)
+
+func TestTruncateHistory_ToolCallAware(t *testing.T) {
+	sm := NewSessionManager("") // in-memory only
+
+	key := "test-tool-truncation"
+
+	// Build a history: [user, assistant+tool_calls, tool, tool, assistant]
+	sm.AddFullMessage(key, messages.Message{Role: "user", Content: "hello"})
+	sm.AddFullMessage(key, messages.Message{
+		Role:    "assistant",
+		Content: "",
+		ToolCalls: []messages.ToolCall{
+			{ID: "call_1", Function: &messages.FunctionCall{Name: "exec", Arguments: `{"command":"ls"}`}},
+			{ID: "call_2", Function: &messages.FunctionCall{Name: "read", Arguments: `{"path":"foo"}`}},
+		},
+	})
+	sm.AddFullMessage(key, messages.Message{Role: "tool", Content: "file1\nfile2", ToolCallID: "call_1"})
+	sm.AddFullMessage(key, messages.Message{Role: "tool", Content: "file contents", ToolCallID: "call_2"})
+	sm.AddFullMessage(key, messages.Message{Role: "assistant", Content: "I found the files"})
+
+	// Truncate to keep last 2 messages (assistant response + one tool result).
+	// The tool-call-aware logic should expand to include the full tool call pair.
+	sm.TruncateHistory(key, 2)
+
+	history := sm.GetHistory(key)
+
+	// Verify: first remaining message should NOT be role "tool".
+	// It should be the assistant message with tool_calls.
+	if len(history) == 0 {
+		t.Fatal("expected non-empty history after truncation")
+	}
+
+	if history[0].Role == "tool" {
+		t.Errorf("first message after truncation is 'tool' -- tool-call pair was split")
+	}
+
+	if history[0].Role != "assistant" {
+		t.Errorf("expected first message to be 'assistant', got %q", history[0].Role)
+	}
+
+	// Should have: assistant+tool_calls, tool, tool, assistant = 4 messages
+	if len(history) != 4 {
+		t.Errorf("expected 4 messages (full tool-call group), got %d", len(history))
+	}
+}
+
+func TestTruncateHistory_NoToolCalls(t *testing.T) {
+	sm := NewSessionManager("")
+
+	key := "test-no-tools"
+
+	sm.AddFullMessage(key, messages.Message{Role: "user", Content: "hello"})
+	sm.AddFullMessage(key, messages.Message{Role: "assistant", Content: "hi"})
+	sm.AddFullMessage(key, messages.Message{Role: "user", Content: "how are you"})
+	sm.AddFullMessage(key, messages.Message{Role: "assistant", Content: "fine"})
+
+	sm.TruncateHistory(key, 2)
+
+	history := sm.GetHistory(key)
+	if len(history) != 2 {
+		t.Errorf("expected 2 messages, got %d", len(history))
+	}
+
+	if history[0].Role != "user" {
+		t.Errorf("expected first message 'user', got %q", history[0].Role)
+	}
+}
+
+func TestAddFullMessage_HardCap(t *testing.T) {
+	sm := NewSessionManager("")
+
+	key := "test-hard-cap"
+
+	// Add 201 messages to trigger the hard cap
+	for i := 0; i < 201; i++ {
+		sm.AddFullMessage(key, messages.Message{Role: "user", Content: "msg"})
+	}
+
+	history := sm.GetHistory(key)
+	if len(history) > 200 {
+		t.Errorf("expected <= 200 messages after hard cap, got %d", len(history))
+	}
+	if len(history) != 50 {
+		t.Errorf("expected 50 messages after hard cap truncation, got %d", len(history))
+	}
+}
+
+func TestCleanupStale(t *testing.T) {
+	sm := NewSessionManager("")
+
+	sm.AddFullMessage("active", messages.Message{Role: "user", Content: "hi"})
+	sm.AddFullMessage("stale", messages.Message{Role: "user", Content: "old"})
+
+	// Make "stale" session appear old by directly modifying its Updated time
+	sm.mu.Lock()
+	sm.sessions["stale"].Updated = sm.sessions["stale"].Updated.Add(-8 * 24 * 3600e9) // 8 days ago
+	sm.mu.Unlock()
+
+	removed := sm.CleanupStale(7 * 24 * 3600e9) // 7 day TTL
+
+	if removed != 1 {
+		t.Errorf("expected 1 stale session removed, got %d", removed)
+	}
+
+	history := sm.GetHistory("stale")
+	if len(history) != 0 {
+		t.Errorf("expected stale session to be removed")
+	}
+
+	history = sm.GetHistory("active")
+	if len(history) != 1 {
+		t.Errorf("expected active session to remain")
+	}
+}