diff --git a/pkg/session/manager.go b/pkg/session/manager.go index 193ad2b30..d1bcb4545 100644 --- a/pkg/session/manager.go +++ b/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"` - Summary string `json:"summary,omitempty"` - Created time.Time `json:"created"` - Updated time.Time `json:"updated"` + Key string `json:"key"` + 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 { - session = &Session{ - Key: sessionKey, - Messages: []providers.Message{}, - Created: time.Now(), + // 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: []messages.Message{}, + Created: time.Now(), + } + sm.sessions[sessionKey] = session } - 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) []providers.Message { +func (sm *SessionManager) GetHistory(key string) []messages.Message { sm.mu.RLock() - defer sm.mu.RUnlock() - session, ok := sm.sessions[key] + sm.mu.RUnlock() + if !ok { - return []providers.Message{} + // Try lazy-load + session = sm.ensureLoaded(key) + if session == nil { + return []messages.Message{} + } } - history := make([]providers.Message, len(session.Messages)) + sm.mu.RLock() + defer sm.mu.RUnlock() + history := make([]messages.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 { - return "" + 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{ - 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{} - } + snapshot := snapshotSession(stored) 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 { diff --git a/pkg/session/manager_test.go b/pkg/session/manager_test.go new file mode 100644 index 000000000..81a3f8656 --- /dev/null +++ b/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") + } +}