From 34916332d81e68911a4e811fcbcf6c87836dcc0e Mon Sep 17 00:00:00 2001 From: mingmxren Date: Sun, 1 Mar 2026 11:44:21 +0800 Subject: [PATCH] feat(session): add scope-aware session index and rotation APIs --- pkg/session/manager.go | 391 +++++++++++++++++++++++++++++++++++- pkg/session/manager_test.go | 157 +++++++++++++++ 2 files changed, 545 insertions(+), 3 deletions(-) diff --git a/pkg/session/manager.go b/pkg/session/manager.go index 08f0b0ad2..8497e74fa 100644 --- a/pkg/session/manager.go +++ b/pkg/session/manager.go @@ -2,8 +2,10 @@ package session import ( "encoding/json" + "fmt" "os" "path/filepath" + "strconv" "strings" "sync" "time" @@ -19,26 +21,244 @@ type Session struct { Updated time.Time `json:"updated"` } +const sessionIndexFilename = "index.json" + +type scopeIndex struct { + ActiveSessionKey string `json:"active_session_key"` + OrderedSessions []string `json:"ordered_sessions"` + UpdatedAt time.Time `json:"updated_at"` +} + +type sessionIndex struct { + Version int `json:"version"` + Scopes map[string]*scopeIndex `json:"scopes"` +} + +type SessionMeta struct { + Ordinal int `json:"ordinal"` + SessionKey string `json:"session_key"` + UpdatedAt time.Time `json:"updated_at"` + MessageCnt int `json:"message_cnt"` + Active bool `json:"active"` +} + type SessionManager struct { - sessions map[string]*Session - mu sync.RWMutex - storage string + sessions map[string]*Session + mu sync.RWMutex + storage string + index sessionIndex + indexPath string } func NewSessionManager(storage string) *SessionManager { sm := &SessionManager{ sessions: make(map[string]*Session), storage: storage, + index: sessionIndex{ + Version: 1, + Scopes: make(map[string]*scopeIndex), + }, } if storage != "" { os.MkdirAll(storage, 0o755) + sm.indexPath = filepath.Join(storage, sessionIndexFilename) sm.loadSessions() + sm.loadIndex() } return sm } +func (sm *SessionManager) ResolveActive(scopeKey string) (string, error) { + sm.mu.Lock() + defer sm.mu.Unlock() + + now := time.Now() + scope, changed := sm.ensureScopeLocked(scopeKey, now) + if changed { + if err := sm.saveIndexLocked(); err != nil { + return "", err + } + } + return scope.ActiveSessionKey, nil +} + +func (sm *SessionManager) StartNew(scopeKey string) (string, error) { + sm.mu.Lock() + defer sm.mu.Unlock() + + now := time.Now() + scope, _ := sm.ensureScopeLocked(scopeKey, now) + + newOrdinal := 2 + for _, existing := range scope.OrderedSessions { + ordinal, ok := sessionOrdinal(scopeKey, existing) + if !ok { + continue + } + if ordinal >= newOrdinal { + newOrdinal = ordinal + 1 + } + } + + newSessionKey := scopeKey + "#" + strconv.Itoa(newOrdinal) + scope.ActiveSessionKey = newSessionKey + scope.OrderedSessions = append([]string{newSessionKey}, scope.OrderedSessions...) + scope.UpdatedAt = now + if err := sm.saveIndexLocked(); err != nil { + return "", err + } + return newSessionKey, nil +} + +func (sm *SessionManager) List(scopeKey string) ([]SessionMeta, error) { + sm.mu.Lock() + defer sm.mu.Unlock() + + now := time.Now() + scope, changed := sm.ensureScopeLocked(scopeKey, now) + if changed { + if err := sm.saveIndexLocked(); err != nil { + return nil, err + } + } + + list := make([]SessionMeta, 0, len(scope.OrderedSessions)) + for i, key := range scope.OrderedSessions { + meta := SessionMeta{ + Ordinal: i + 1, + SessionKey: key, + Active: key == scope.ActiveSessionKey, + } + + if session, ok := sm.sessions[key]; ok { + meta.UpdatedAt = session.Updated + meta.MessageCnt = len(session.Messages) + } + list = append(list, meta) + } + return list, nil +} + +func (sm *SessionManager) Resume(scopeKey string, index int) (string, error) { + if index < 1 { + return "", fmt.Errorf("session index must be >= 1") + } + + sm.mu.Lock() + defer sm.mu.Unlock() + + now := time.Now() + scope, changed := sm.ensureScopeLocked(scopeKey, now) + if changed { + if err := sm.saveIndexLocked(); err != nil { + return "", err + } + } + + if index > len(scope.OrderedSessions) { + return "", fmt.Errorf("session index %d out of range", index) + } + + scope.ActiveSessionKey = scope.OrderedSessions[index-1] + scope.UpdatedAt = now + if err := sm.saveIndexLocked(); err != nil { + return "", err + } + return scope.ActiveSessionKey, nil +} + +func (sm *SessionManager) DeleteSession(sessionKey string) error { + sm.mu.Lock() + changed := false + delete(sm.sessions, sessionKey) + + now := time.Now() + for scopeKey, scope := range sm.index.Scopes { + if scope == nil { + delete(sm.index.Scopes, scopeKey) + changed = true + continue + } + + filtered := scope.OrderedSessions[:0] + removed := false + for _, key := range scope.OrderedSessions { + if key == sessionKey { + removed = true + changed = true + continue + } + filtered = append(filtered, key) + } + scope.OrderedSessions = filtered + if !removed { + continue + } + + if scope.ActiveSessionKey == sessionKey { + if len(scope.OrderedSessions) > 0 { + scope.ActiveSessionKey = scope.OrderedSessions[0] + } else { + scope.ActiveSessionKey = "" + } + } + + if len(scope.OrderedSessions) == 0 { + delete(sm.index.Scopes, scopeKey) + continue + } + scope.UpdatedAt = now + } + + if changed { + if err := sm.saveIndexLocked(); err != nil { + sm.mu.Unlock() + return err + } + } + sm.mu.Unlock() + + if err := sm.deleteSessionFile(sessionKey); err != nil { + return err + } + return nil +} + +func (sm *SessionManager) Prune(scopeKey string, limit int) ([]string, error) { + if limit < 1 { + return nil, fmt.Errorf("limit must be >= 1") + } + + sm.mu.Lock() + now := time.Now() + scope, changed := sm.ensureScopeLocked(scopeKey, now) + if changed { + if err := sm.saveIndexLocked(); err != nil { + sm.mu.Unlock() + return nil, err + } + } + + if len(scope.OrderedSessions) <= limit { + sm.mu.Unlock() + return []string{}, nil + } + + candidates := append([]string(nil), scope.OrderedSessions[limit:]...) + sm.mu.Unlock() + + pruned := make([]string, 0, len(candidates)) + for _, sessionKey := range candidates { + if err := sm.DeleteSession(sessionKey); err != nil { + return pruned, err + } + pruned = append(pruned, sessionKey) + } + return pruned, nil +} + func (sm *SessionManager) GetOrCreate(key string) *Session { sm.mu.Lock() defer sm.mu.Unlock() @@ -247,6 +467,9 @@ func (sm *SessionManager) loadSessions() error { if filepath.Ext(file.Name()) != ".json" { continue } + if file.Name() == sessionIndexFilename { + continue + } sessionPath := filepath.Join(sm.storage, file.Name()) data, err := os.ReadFile(sessionPath) @@ -258,6 +481,9 @@ func (sm *SessionManager) loadSessions() error { if err := json.Unmarshal(data, &session); err != nil { continue } + if session.Key == "" { + continue + } sm.sessions[session.Key] = &session } @@ -265,6 +491,165 @@ func (sm *SessionManager) loadSessions() error { return nil } +func (sm *SessionManager) loadIndex() error { + if sm.storage == "" { + return nil + } + + data, err := os.ReadFile(sm.indexPath) + if err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + + var loaded sessionIndex + if err := json.Unmarshal(data, &loaded); err != nil { + return err + } + if loaded.Version == 0 { + loaded.Version = 1 + } + if loaded.Scopes == nil { + loaded.Scopes = make(map[string]*scopeIndex) + } + + for scopeKey, scope := range loaded.Scopes { + if scope == nil { + delete(loaded.Scopes, scopeKey) + continue + } + if len(scope.OrderedSessions) == 0 { + scope.OrderedSessions = []string{scopeKey} + } + if scope.ActiveSessionKey == "" { + scope.ActiveSessionKey = scope.OrderedSessions[0] + } + } + + sm.index = loaded + return nil +} + +func (sm *SessionManager) saveIndexLocked() error { + if sm.storage == "" { + return nil + } + if sm.index.Scopes == nil { + sm.index.Scopes = make(map[string]*scopeIndex) + } + sm.index.Version = 1 + + data, err := json.MarshalIndent(sm.index, "", " ") + if err != nil { + return err + } + + tmpFile, err := os.CreateTemp(sm.storage, "index-*.tmp") + if err != nil { + return err + } + + tmpPath := tmpFile.Name() + cleanup := true + defer func() { + if cleanup { + _ = os.Remove(tmpPath) + } + }() + + if _, err := tmpFile.Write(data); err != nil { + _ = tmpFile.Close() + return err + } + if err := tmpFile.Chmod(0o644); err != nil { + _ = tmpFile.Close() + return err + } + if err := tmpFile.Sync(); err != nil { + _ = tmpFile.Close() + return err + } + if err := tmpFile.Close(); err != nil { + return err + } + if err := os.Rename(tmpPath, sm.indexPath); err != nil { + return err + } + cleanup = false + return nil +} + +func (sm *SessionManager) ensureScopeLocked(scopeKey string, now time.Time) (*scopeIndex, bool) { + if sm.index.Scopes == nil { + sm.index.Scopes = make(map[string]*scopeIndex) + } + + scope, ok := sm.index.Scopes[scopeKey] + if !ok || scope == nil { + scope = &scopeIndex{ + ActiveSessionKey: scopeKey, + OrderedSessions: []string{scopeKey}, + UpdatedAt: now, + } + sm.index.Scopes[scopeKey] = scope + return scope, true + } + + changed := false + if len(scope.OrderedSessions) == 0 { + scope.OrderedSessions = []string{scopeKey} + changed = true + } + if scope.ActiveSessionKey == "" { + scope.ActiveSessionKey = scope.OrderedSessions[0] + changed = true + } + if changed { + scope.UpdatedAt = now + } + return scope, changed +} + +func (sm *SessionManager) deleteSessionFile(sessionKey string) error { + if sm.storage == "" { + return nil + } + + filename := sanitizeFilename(sessionKey) + if filename == "." || !filepath.IsLocal(filename) || strings.ContainsAny(filename, `/\`) { + return os.ErrInvalid + } + + sessionPath := filepath.Join(sm.storage, filename+".json") + if err := os.Remove(sessionPath); err != nil { + if os.IsNotExist(err) { + return nil + } + return err + } + return nil +} + +func sessionOrdinal(scopeKey, sessionKey string) (int, bool) { + if sessionKey == scopeKey { + return 1, true + } + + prefix := scopeKey + "#" + if !strings.HasPrefix(sessionKey, prefix) { + return 0, false + } + + n, err := strconv.Atoi(strings.TrimPrefix(sessionKey, prefix)) + if err != nil || n < 2 { + return 0, false + } + + return n, true +} + // SetHistory updates the messages of a session. func (sm *SessionManager) SetHistory(key string, history []providers.Message) { sm.mu.Lock() diff --git a/pkg/session/manager_test.go b/pkg/session/manager_test.go index 5ef5f4349..ade3b1e25 100644 --- a/pkg/session/manager_test.go +++ b/pkg/session/manager_test.go @@ -1,6 +1,7 @@ package session import ( + "encoding/json" "os" "path/filepath" "testing" @@ -72,3 +73,159 @@ func TestSave_RejectsPathTraversal(t *testing.T) { } } } + +func TestSessionIndex_BootstrapScopeAndPersist(t *testing.T) { + tmp := t.TempDir() + sm := NewSessionManager(tmp) + scope := "agent:main:telegram:direct:user1" + + active, err := sm.ResolveActive(scope) + if err != nil { + t.Fatal(err) + } + if active != scope { + t.Fatalf("active=%q, want %q", active, scope) + } + + indexPath := filepath.Join(tmp, "index.json") + raw, err := os.ReadFile(indexPath) + if err != nil { + t.Fatalf("read index: %v", err) + } + + var decoded map[string]any + if err := json.Unmarshal(raw, &decoded); err != nil { + t.Fatalf("unmarshal index: %v", err) + } + + sm2 := NewSessionManager(tmp) + active2, err := sm2.ResolveActive(scope) + if err != nil { + t.Fatal(err) + } + if active2 != active { + t.Fatalf("active2=%q, want %q", active2, active) + } +} + +func TestStartNew_CreatesMonotonicSessionKeys(t *testing.T) { + sm := NewSessionManager(t.TempDir()) + scope := "agent:main:telegram:direct:user1" + + if _, err := sm.ResolveActive(scope); err != nil { + t.Fatal(err) + } + + s2, err := sm.StartNew(scope) + if err != nil { + t.Fatal(err) + } + if s2 != scope+"#2" { + t.Fatalf("s2=%q, want %q", s2, scope+"#2") + } + + s3, err := sm.StartNew(scope) + if err != nil { + t.Fatal(err) + } + if s3 != scope+"#3" { + t.Fatalf("s3=%q, want %q", s3, scope+"#3") + } +} + +func TestListAndResume_ByScopeOrdinal(t *testing.T) { + sm := NewSessionManager(t.TempDir()) + scope := "agent:main:telegram:direct:user1" + + if _, err := sm.ResolveActive(scope); err != nil { + t.Fatal(err) + } + if _, err := sm.StartNew(scope); err != nil { // #2 + t.Fatal(err) + } + if _, err := sm.StartNew(scope); err != nil { // #3 (active) + t.Fatal(err) + } + + list, err := sm.List(scope) + if err != nil { + t.Fatal(err) + } + if len(list) != 3 { + t.Fatalf("len(list)=%d, want 3", len(list)) + } + if list[0].Ordinal != 1 || list[0].SessionKey != scope+"#3" || !list[0].Active { + t.Fatalf("list[0]=%+v", list[0]) + } + if list[2].Ordinal != 3 || list[2].SessionKey != scope { + t.Fatalf("list[2]=%+v", list[2]) + } + + resumed, err := sm.Resume(scope, 3) + if err != nil { + t.Fatal(err) + } + if resumed != scope { + t.Fatalf("resumed=%q, want %q", resumed, scope) + } + + listAfter, err := sm.List(scope) + if err != nil { + t.Fatal(err) + } + if !listAfter[2].Active { + t.Fatalf("listAfter[2] should be active: %+v", listAfter[2]) + } +} + +func TestPrune_RemovesOldestFromMemoryAndDisk(t *testing.T) { + dir := t.TempDir() + sm := NewSessionManager(dir) + scope := "agent:main:telegram:direct:user1" + + if _, err := sm.ResolveActive(scope); err != nil { + t.Fatal(err) + } + if _, err := sm.StartNew(scope); err != nil { // #2 + t.Fatal(err) + } + if _, err := sm.StartNew(scope); err != nil { // #3 (active) + t.Fatal(err) + } + + keys := []string{scope, scope + "#2", scope + "#3"} + for _, key := range keys { + sm.AddMessage(key, "user", "hello") + if err := sm.Save(key); err != nil { + t.Fatalf("save %q: %v", key, err) + } + } + + pruned, err := sm.Prune(scope, 2) + if err != nil { + t.Fatal(err) + } + if len(pruned) != 1 || pruned[0] != scope { + t.Fatalf("pruned=%v, want [%s]", pruned, scope) + } + + if got := len(sm.GetHistory(scope)); got != 0 { + t.Fatalf("expected deleted session history len=0, got %d", got) + } + + removedFile := filepath.Join(dir, sanitizeFilename(scope)+".json") + if _, err := os.Stat(removedFile); !os.IsNotExist(err) { + t.Fatalf("expected %s deleted, stat err=%v", removedFile, err) + } + + list, err := sm.List(scope) + if err != nil { + t.Fatal(err) + } + if len(list) != 2 { + t.Fatalf("len(list)=%d, want 2", len(list)) + } + if list[0].SessionKey != scope+"#3" || list[1].SessionKey != scope+"#2" { + t.Fatalf("list order after prune = %+v", list) + } +}