fix(storage): stop session retention from evicting live sessions

Session keys are {StartTime.UnixNano()}_{ID}, and enforceSessionRetention
deleted the lowest keys. That evicts the LONGEST-LIVED session first, which
is exactly the session most likely to still be connected and working. A
client that stayed connected while 100 newer sessions were created vanished
from storage, went missing from the sessions API, and its later close/stat
updates had no record to find.

Choose victims by usefulness instead of key order:

  1. closed sessions, stalest first - finished work, nothing will write to
     these records again;
  2. only if that frees too little, active sessions by last activity,
     stalest first - an abandoned client goes before a working one.

Status is authoritative for liveness and activity only ranks within a tier,
so a connected-but-idle client outranks a closed session that churned more
recently; CloseInactiveSessions is what flips a truly dead session to
closed, and only then does it become evictable.

Tier 2 keeps the cap absolute. Refusing to evict active sessions at all
would let an all-active bucket (abandoned clients, a reconnect loop that
never closes cleanly) grow without bound, which is worse than the bug being
fixed.

Two further defects surfaced in the same function:

- bucket.Stats().KeyN does not count records Put earlier in the same write
  transaction, so the bucket settled permanently at 101 records rather than
  100. The scan now counts keys directly. This one is observable: the new
  cap assertions failed with "expected 100, actual 101" before the change.
- Records were deleted while iterating the bucket's own cursor. bbolt
  documents that as unsafe ("Changing data while traversing with a cursor
  may cause it to be invalidated and return unexpected keys and/or
  values"). A 2000-key probe did not manage to make it skip, so this is a
  documented hazard rather than an observed failure - but sorting requires
  collecting keys first regardless, which is also what
  PruneExcessActivities in this package already does.

Retention also falls back to StartTime when LastActivity is unset, so
records predating that field do not all sort as year 1 and get evicted
ahead of genuinely stale sessions.

The suite is mutation-verified. An early version of the live-session test
passed with the status tier removed - its live session also had the
freshest activity, so the staleness ordering alone kept it - so the
connected-but-idle case was added to pin the tier itself.

Author: Dumbris

internal/storage/manager.go

@@ -1218,6 +1218,11 @@ type SessionRecord struct {
 	WorkSessionID string `json:"work_session_id,omitempty"`
 }
 
+// sessionRetentionLimit is the hard cap on stored session records. The cap is
+// absolute: it holds even when every retained session is "active" (see
+// enforceSessionRetention).
+const sessionRetentionLimit = 100
+
 // CreateSession creates a new session record
 func (m *Manager) CreateSession(session *SessionRecord) error {
 	m.mu.Lock()
@@ -1279,9 +1284,9 @@ func (m *Manager) CreateSession(session *SessionRecord) error {
 			return fmt.Errorf("failed to store session: %w", err)
 		}
 
-		// Enforce retention limit (keep 100 most recent) only when creating new sessions
+		// Enforce retention limit only when creating new sessions
 		if existingKey == nil {
-			return m.enforceSessionRetention(bucket, 100)
+			return m.enforceSessionRetention(bucket, sessionRetentionLimit)
 		}
 		return nil
 	})
@@ -1699,26 +1704,104 @@ func (m *Manager) GetToolCallsBySession(sessionID string, limit, offset int) ([]
 	return toolCalls, total, err
 }
 
-// enforceSessionRetention deletes oldest sessions if count exceeds limit
+// sessionEvictionCandidate is one stored session, reduced to the two properties
+// retention actually cares about.
+type sessionEvictionCandidate struct {
+	key      []byte
+	active   bool
+	activity time.Time // LastActivity, falling back to StartTime when unset
+}
+
+// enforceSessionRetention trims the sessions bucket down to maxSessions records.
+//
+// Victims are chosen by usefulness, NOT by key order. Session keys are
+// {StartTime.UnixNano()}_{ID}, so deleting the lowest keys deletes the
+// longest-lived session first — which is precisely the session most likely to
+// still be connected and working. A client that stayed connected all day used to
+// disappear from storage as soon as 100 newer sessions had been created; after
+// that it was absent from the sessions API and its later close/stat updates had
+// no record to find.
+//
+// Eviction order is therefore two-tiered:
+//
+//  1. closed sessions, stalest first — finished work; nothing will ever write to
+//     these records again;
+//  2. only if tier 1 does not free enough room, active sessions ordered by last
+//     activity, stalest first — a client that died without closing goes before a
+//     client that is genuinely working.
+//
+// Tier 2 is what keeps the cap absolute. Refusing to evict active sessions at
+// all would let an all-active bucket (abandoned clients, a reconnect loop that
+// never closes cleanly) grow without bound, which is a worse bug than the one
+// this fixes.
 func (m *Manager) enforceSessionRetention(bucket *bbolt.Bucket, maxSessions int) error {
-	stats := bucket.Stats()
-	if stats.KeyN <= maxSessions {
+	if maxSessions <= 0 {
 		return nil
 	}
 
-	// Delete oldest sessions (first keys since they have oldest timestamps)
-	toDelete := stats.KeyN - maxSessions
-	deleted := 0
-
+	// Classify every record in one pass. bucket.Stats() is deliberately not used
+	// for the count: inside a write transaction it does not account for records
+	// Put earlier in the same transaction, which let the bucket settle one record
+	// above the cap forever.
+	var candidates []sessionEvictionCandidate
 	c := bucket.Cursor()
-	for k, _ := c.First(); k != nil && deleted < toDelete; k, _ = c.Next() {
-		if err := bucket.Delete(k); err != nil {
+	for k, v := c.First(); k != nil; k, v = c.Next() {
+		cand := sessionEvictionCandidate{key: append([]byte(nil), k...)}
+		var session SessionRecord
+		if err := json.Unmarshal(v, &session); err != nil {
+			// Unreadable record: nothing can use it, so it evicts first.
+			m.logger.Warnw("Unreadable session record during retention", "key", string(k), "error", err)
+		} else {
+			cand.active = session.Status == "active"
+			cand.activity = session.LastActivity
+			if cand.activity.IsZero() {
+				// Records written before LastActivity existed; StartTime is the
+				// best evidence available. Without this they would all sort as
+				// year 1 and be evicted ahead of genuinely stale sessions.
+				cand.activity = session.StartTime
+			}
+		}
+		candidates = append(candidates, cand)
+	}
+
+	toDelete := len(candidates) - maxSessions
+	if toDelete <= 0 {
+		return nil
+	}
+
+	sort.Slice(candidates, func(i, j int) bool {
+		a, b := candidates[i], candidates[j]
+		if a.active != b.active {
+			return !a.active // closed sessions are evicted before active ones
+		}
+		if !a.activity.Equal(b.activity) {
+			return a.activity.Before(b.activity) // stalest first
+		}
+		return bytes.Compare(a.key, b.key) < 0 // deterministic tiebreak
+	})
+
+	// Delete only after the scan has finished. Sorting requires it anyway, and
+	// bbolt documents mutation during traversal as unsafe ("Changing data while
+	// traversing with a cursor may cause it to be invalidated and return
+	// unexpected keys and/or values"), which is what the previous
+	// delete-inside-the-cursor-loop did. PruneExcessActivities already collects
+	// keys first for the same reason.
+	activeEvicted := 0
+	for i := 0; i < toDelete; i++ {
+		if err := bucket.Delete(candidates[i].key); err != nil {
 			return fmt.Errorf("failed to delete old session: %w", err)
 		}
-		deleted++
+		if candidates[i].active {
+			activeEvicted++
+		}
 	}
 
-	m.logger.Debugw("Enforced session retention", "deleted", deleted, "remaining", maxSessions)
+	if activeEvicted > 0 {
+		// Only reachable when there was no closed session left to sacrifice.
+		m.logger.Warnw("Session retention had to evict active sessions",
+			"active_evicted", activeEvicted, "limit", maxSessions)
+	}
+	m.logger.Debugw("Enforced session retention", "deleted", toDelete, "remaining", maxSessions)
 	return nil
 }
 

internal/storage/sessions_retention_test.go

@@ -0,0 +1,205 @@
+package storage
+
+import (
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+)
+
+func newRetentionTestManager(t *testing.T) *Manager {
+	t.Helper()
+	m, err := NewManager(t.TempDir(), zap.NewNop().Sugar())
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = m.Close() })
+	return m
+}
+
+// makeSession builds a session record. startOffset is relative to base, so tests
+// can control the retention key ({StartTime.UnixNano()}_{ID}) precisely.
+func makeSession(id string, base time.Time, startOffset, activityOffset time.Duration, status string) *SessionRecord {
+	start := base.Add(startOffset)
+	rec := &SessionRecord{
+		ID:           id,
+		ClientName:   "test-client",
+		StartTime:    start,
+		LastActivity: base.Add(activityOffset),
+		Status:       status,
+	}
+	if status == "closed" {
+		end := start.Add(time.Second)
+		rec.EndTime = &end
+	}
+	return rec
+}
+
+func sessionIDs(sessions []*SessionRecord) []string {
+	ids := make([]string, 0, len(sessions))
+	for _, s := range sessions {
+		ids = append(ids, s.ID)
+	}
+	return ids
+}
+
+func containsSessionID(sessions []*SessionRecord, id string) bool {
+	for _, s := range sessions {
+		if s.ID == id {
+			return true
+		}
+	}
+	return false
+}
+
+// The bug: retention deletes the oldest KEYS, and keys are ordered by StartTime.
+// A long-lived ACTIVE session has the oldest start time by definition, so it is
+// the first thing evicted once enough newer sessions are created — even though
+// the client is still connected and every one of those newer sessions is closed.
+//
+// Retention must evict finished work before it evicts live work.
+func TestEnforceSessionRetention_KeepsLiveSessionWhenClosedOnesCouldGo(t *testing.T) {
+	m := newRetentionTestManager(t)
+	base := time.Now().Add(-24 * time.Hour)
+
+	// A long-running client connects first and is still working.
+	live := makeSession("live-session", base, 0, 24*time.Hour, "active")
+	require.NoError(t, m.CreateSession(live))
+
+	// Then a chatty client reconnects sessionRetentionLimit+20 times. Every one of
+	// those sessions is finished, so every one of them is a better eviction
+	// candidate than the live session.
+	for i := 0; i < sessionRetentionLimit+20; i++ {
+		s := makeSession(fmt.Sprintf("closed-%04d", i), base,
+			time.Duration(i+1)*time.Minute, time.Duration(i+1)*time.Minute, "closed")
+		require.NoError(t, m.CreateSession(s))
+	}
+
+	got, err := m.GetSessionByID("live-session")
+	require.NoError(t, err, "the live session must still exist in storage — "+
+		"a connected client's record must not be evicted while closed records remain")
+	assert.Equal(t, "active", got.Status)
+
+	sessions, total, err := m.GetRecentSessions(sessionRetentionLimit)
+	require.NoError(t, err)
+	assert.LessOrEqual(t, total, sessionRetentionLimit, "the hard cap must still hold")
+	assert.True(t, containsSessionID(sessions, "live-session"),
+		"the live session must be listed; got %v", sessionIDs(sessions))
+}
+
+// The same bug, with the staleness ordering deliberately pointing the wrong way:
+// a connected-but-idle client (an editor left open in the background) has OLDER
+// last-activity than every one of the closed sessions that churned past it. Only
+// the status tier can save it here — ranking purely by last activity would evict
+// the live session first.
+//
+// Status is authoritative for liveness; activity only ranks within a tier. A
+// session that really is dead gets its status flipped by CloseInactiveSessions,
+// and only then becomes evictable.
+func TestEnforceSessionRetention_KeepsIdleLiveSessionOverFresherClosedOnes(t *testing.T) {
+	m := newRetentionTestManager(t)
+	base := time.Now().Add(-24 * time.Hour)
+
+	// Connected, but last did anything a minute after base.
+	idleLive := makeSession("idle-live-session", base, 0, time.Minute, "active")
+	require.NoError(t, m.CreateSession(idleLive))
+
+	// Closed sessions, every one of them more recently active than the live one.
+	for i := 0; i < sessionRetentionLimit+20; i++ {
+		s := makeSession(fmt.Sprintf("closed-%04d", i), base,
+			time.Duration(i+2)*time.Minute, time.Duration(i+2)*time.Minute, "closed")
+		require.NoError(t, m.CreateSession(s))
+	}
+
+	got, err := m.GetSessionByID("idle-live-session")
+	require.NoError(t, err, "a connected client must outrank every closed session, "+
+		"however recently those closed sessions were active")
+	assert.Equal(t, "active", got.Status)
+
+	sessions, total, err := m.GetRecentSessions(sessionRetentionLimit)
+	require.NoError(t, err)
+	assert.Equal(t, sessionRetentionLimit, total)
+	assert.True(t, containsSessionID(sessions, "idle-live-session"),
+		"the idle live session must be listed; got %v", sessionIDs(sessions))
+}
+
+// The abandoned-client case: if every retained session is "active" (clients that
+// died without closing), preferring active sessions must NOT turn the bucket
+// into an unbounded bucket. The cap still holds, and among active sessions the
+// one with the freshest activity is the one that survives.
+func TestEnforceSessionRetention_CapHoldsWhenEverySessionIsActive(t *testing.T) {
+	m := newRetentionTestManager(t)
+	base := time.Now().Add(-48 * time.Hour)
+
+	// The genuinely live one: oldest start time, but active seconds ago.
+	live := makeSession("live-session", base, 0, 48*time.Hour, "active")
+	require.NoError(t, m.CreateSession(live))
+
+	// Clients that died without closing: newer start times, stale activity.
+	for i := 0; i < sessionRetentionLimit+30; i++ {
+		s := makeSession(fmt.Sprintf("abandoned-%04d", i), base,
+			time.Duration(i+1)*time.Minute, time.Duration(i+1)*time.Minute, "active")
+		require.NoError(t, m.CreateSession(s))
+	}
+
+	sessions, total, err := m.GetRecentSessions(sessionRetentionLimit + 100)
+	require.NoError(t, err)
+	assert.Equal(t, sessionRetentionLimit, total,
+		"an all-active bucket must still be capped — otherwise abandoned sessions grow without bound")
+	assert.Len(t, sessions, sessionRetentionLimit)
+
+	got, err := m.GetSessionByID("live-session")
+	require.NoError(t, err, "among active sessions, the freshest activity must be the last to go")
+	assert.Equal(t, "active", got.Status)
+}
+
+// Closed sessions are still evicted oldest-first among themselves.
+func TestEnforceSessionRetention_EvictsOldestClosedFirst(t *testing.T) {
+	m := newRetentionTestManager(t)
+	base := time.Now().Add(-24 * time.Hour)
+
+	for i := 0; i < sessionRetentionLimit+5; i++ {
+		s := makeSession(fmt.Sprintf("closed-%04d", i), base,
+			time.Duration(i)*time.Minute, time.Duration(i)*time.Minute, "closed")
+		require.NoError(t, m.CreateSession(s))
+	}
+
+	sessions, total, err := m.GetRecentSessions(sessionRetentionLimit + 50)
+	require.NoError(t, err)
+	assert.Equal(t, sessionRetentionLimit, total)
+
+	// The 5 oldest are gone, the newest are kept.
+	for i := 0; i < 5; i++ {
+		id := fmt.Sprintf("closed-%04d", i)
+		assert.False(t, containsSessionID(sessions, id), "%s should have been evicted", id)
+	}
+	assert.True(t, containsSessionID(sessions, fmt.Sprintf("closed-%04d", sessionRetentionLimit+4)),
+		"the newest closed session must be kept")
+}
+
+// A session with no LastActivity (older records predate the field) must fall
+// back to StartTime rather than sorting as "epoch zero" and being evicted first.
+func TestEnforceSessionRetention_ZeroLastActivityFallsBackToStartTime(t *testing.T) {
+	m := newRetentionTestManager(t)
+	base := time.Now().Add(-24 * time.Hour)
+
+	// Legacy active record: recent start, no LastActivity recorded.
+	legacy := &SessionRecord{
+		ID:        "legacy-active",
+		StartTime: base.Add(10 * time.Hour),
+		Status:    "active",
+	}
+	require.NoError(t, m.CreateSession(legacy))
+
+	// Older active records with equally-old activity.
+	for i := 0; i < sessionRetentionLimit+10; i++ {
+		s := makeSession(fmt.Sprintf("old-active-%04d", i), base,
+			time.Duration(i)*time.Minute, time.Duration(i)*time.Minute, "active")
+		require.NoError(t, m.CreateSession(s))
+	}
+
+	got, err := m.GetSessionByID("legacy-active")
+	require.NoError(t, err, "a zero LastActivity must fall back to StartTime, not sort as the year 1")
+	assert.Equal(t, "active", got.Status)
+}