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.
This commit is contained in:
ZanzyTHEbar
parent
66b104555b
commit
d982302fa6
2 changed files with 373 additions and 38 deletions
|
|
@ -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) []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 {
|
||||
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 {
|
||||
|
|
|
|||
121
pkg/session/manager_test.go
Normal file
121
pkg/session/manager_test.go
Normal file
|
|
@ -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")
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue