From d8cd90ae3f08455d9448e3488e784bb1e3710231 Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Wed, 29 Jul 2026 20:12:03 +0300 Subject: [PATCH 1/7] 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. --- internal/storage/manager.go | 109 +++++++++-- internal/storage/sessions_retention_test.go | 205 ++++++++++++++++++++ 2 files changed, 301 insertions(+), 13 deletions(-) create mode 100644 internal/storage/sessions_retention_test.go diff --git a/internal/storage/manager.go b/internal/storage/manager.go index 301010de..bad5c1ae 100644 --- a/internal/storage/manager.go +++ b/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 } diff --git a/internal/storage/sessions_retention_test.go b/internal/storage/sessions_retention_test.go new file mode 100644 index 00000000..99caffd4 --- /dev/null +++ b/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) +} From cd4c9974565a5d80cabd4ae6adbf43cb66dbbcea Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Wed, 29 Jul 2026 21:45:17 +0300 Subject: [PATCH 2/7] fix(storage): make the session cap an invariant of every write path Cross-model review round 1 on PR #928. All four findings were verified against the code before being taken. Migration bypassed the cap (P2). CreateSession was the only caller of retention, but migrateLegacySessions is a second write path that can move an arbitrary number of records into the sessions bucket. A process that migrated and then only updated or closed existing sessions - never creating a new ID - would sit above the cap indefinitely. Retention now runs at database-open time, right after the migration, which also repairs a bucket left oversized by any older version for any reason. Selection is now bounded (P2). Instead of collecting and sorting every record, keep a heap of the best maxSessions seen so far and stream the rest into the victim list: O(n log maxSessions) time and O(maxSessions) ranking state, rather than O(n log n) and O(n). This matters on the one pass that can see a bucket far above the cap - the open-time trim above. Unreadable records now have their own tier (P2). The policy said they evict first; the code could not deliver it. An unmarshal failure left the zero value, which is indistinguishable from a VALID record whose status is not "active" and whose timestamps are unset, so the two shared a tier and a timestamp and key order decided between them - the usable record could be deleted while the corrupt one survived. Ranking is now an explicit tier (unreadable < closed < active) rather than a pair of derived booleans. The zero-activity test proved too little (P3). It only asserted that a zero-activity record survives, which also passes if the fallback is replaced by time.Now() - a substitution that would make every legacy record eternally the freshest thing in the bucket. It now pins both directions in one bucket: a zero-activity record with an ancient StartTime must be evicted, one with a recent StartTime must survive. New tests: migration leaves the bucket within the cap and spares the user login; an unreadable record is evicted ahead of a valid zero-valued record that sorts before it by key; a 500-record bucket trims to exactly the newest 100 in one call. All mutation-verified, including the time.Now() substitution named in the review. --- internal/storage/bbolt.go | 8 + internal/storage/manager.go | 197 +++++++++++++----- internal/storage/sessions_retention_test.go | 220 +++++++++++++++++++- 3 files changed, 369 insertions(+), 56 deletions(-) diff --git a/internal/storage/bbolt.go b/internal/storage/bbolt.go index 784dd32a..96afbaa1 100644 --- a/internal/storage/bbolt.go +++ b/internal/storage/bbolt.go @@ -114,6 +114,14 @@ func (b *BoltDB) initBuckets() error { return fmt.Errorf("failed to migrate legacy sessions bucket: %w", err) } + // Migration is a write path into the sessions bucket like any other, and + // it can move in an arbitrary number of records. Enforce the retention + // cap here so it is an invariant of an open database rather than + // something only CreateSession happens to maintain. + if err := enforceSessionRetentionOnOpen(tx, b.logger); err != nil { + return fmt.Errorf("failed to enforce session retention: %w", err) + } + // Backfill the scan-job index for databases created before MCP-2205. // Idempotent: only runs when the index is empty but jobs exist. if err := backfillScanJobIndex(tx); err != nil { diff --git a/internal/storage/manager.go b/internal/storage/manager.go index bad5c1ae..2daa22c7 100644 --- a/internal/storage/manager.go +++ b/internal/storage/manager.go @@ -2,6 +2,7 @@ package storage import ( "bytes" + "container/heap" "encoding/hex" "encoding/json" "fmt" @@ -1704,14 +1705,95 @@ func (m *Manager) GetToolCallsBySession(sessionID string, limit, offset int) ([] return toolCalls, total, err } -// sessionEvictionCandidate is one stored session, reduced to the two properties -// retention actually cares about. +// Eviction tiers, worst first. The tier is the primary ranking key; activity +// only breaks ties inside a tier. +const ( + // sessionTierUnreadable is a record that cannot be unmarshalled. Nothing can + // display, close, or update it, so it is worth strictly less than any record + // that parses. This has to be its own tier rather than a derived property: + // an unmarshal failure leaves the zero value, which is indistinguishable + // from a VALID record whose status is not "active" and whose timestamps are + // unset. Sharing a tier with those meant key order decided between them, and + // the usable record could lose. + sessionTierUnreadable = iota + // sessionTierClosed is finished work — nothing will ever write to it again. + sessionTierClosed + // sessionTierActive is a session the proxy still believes is live. + sessionTierActive +) + +// sessionEvictionCandidate is one stored session, reduced to the properties +// retention ranks by. type sessionEvictionCandidate struct { key []byte - active bool + tier int activity time.Time // LastActivity, falling back to StartTime when unset } +// worseThan reports whether a should be evicted before b. +func (a sessionEvictionCandidate) worseThan(b sessionEvictionCandidate) bool { + if a.tier != b.tier { + return a.tier < b.tier + } + if !a.activity.Equal(b.activity) { + return a.activity.Before(b.activity) // stalest first + } + return bytes.Compare(a.key, b.key) < 0 // deterministic tiebreak +} + +// survivorHeap is a min-heap under worseThan, so its root is the weakest record +// currently being kept — i.e. the next one to be displaced. +type survivorHeap []sessionEvictionCandidate + +func (h survivorHeap) Len() int { return len(h) } +func (h survivorHeap) Less(i, j int) bool { return h[i].worseThan(h[j]) } +func (h survivorHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } + +func (h *survivorHeap) Push(x any) { + c, ok := x.(sessionEvictionCandidate) + if !ok { + return + } + *h = append(*h, c) +} + +func (h *survivorHeap) Pop() any { + old := *h + n := len(old) + item := old[n-1] + *h = old[:n-1] + return item +} + +// classifySessionForRetention reduces a stored record to its ranking properties. +func classifySessionForRetention(key, value []byte, logger *zap.SugaredLogger) sessionEvictionCandidate { + cand := sessionEvictionCandidate{ + key: append([]byte(nil), key...), + tier: sessionTierUnreadable, + } + + var session SessionRecord + if err := json.Unmarshal(value, &session); err != nil { + logger.Warnw("Unreadable session record during retention", "key", string(key), "error", err) + return cand + } + + if session.Status == "active" { + cand.tier = sessionTierActive + } else { + cand.tier = sessionTierClosed + } + + cand.activity = session.LastActivity + if cand.activity.IsZero() { + // Records written before LastActivity existed. StartTime is the only + // evidence available, and it cuts both ways: an ancient legacy record + // must rank as stale, not as freshly active. + cand.activity = session.StartTime + } + return cand +} + // enforceSessionRetention trims the sessions bucket down to maxSessions records. // // Victims are chosen by usefulness, NOT by key order. Session keys are @@ -1735,76 +1817,97 @@ type sessionEvictionCandidate struct { // 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 { + return trimSessionsToLimit(bucket, maxSessions, m.logger) +} + +// trimSessionsToLimit is the retention pass itself, callable from any write +// path that can put the bucket over the cap — CreateSession is not the only +// one (see enforceSessionRetentionOnOpen). +// +// Selection is bounded: rather than collecting and sorting every record, it +// keeps a heap of the best maxSessions seen so far and streams everything else +// straight into the victim list. That makes ranking O(n log maxSessions) with +// O(maxSessions) ranking state, instead of O(n log n) with O(n) state, which +// matters on the one pass that can see a bucket far above the cap: the +// migration/open-time trim. +func trimSessionsToLimit(bucket *bbolt.Bucket, maxSessions int, logger *zap.SugaredLogger) error { if maxSessions <= 0 { return nil } - // 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 + // bucket.Stats() is deliberately not used to size this: 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. + survivors := make(survivorHeap, 0, maxSessions) + var victims []sessionEvictionCandidate + c := bucket.Cursor() 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 + cand := classifySessionForRetention(k, v, logger) + switch { + case len(survivors) < maxSessions: + heap.Push(&survivors, cand) + case cand.worseThan(survivors[0]): + // Weaker than every record currently being kept. + victims = append(victims, cand) + default: + displaced, ok := heap.Pop(&survivors).(sessionEvictionCandidate) + if ok { + victims = append(victims, displaced) } + heap.Push(&survivors, cand) } - candidates = append(candidates, cand) } - toDelete := len(candidates) - maxSessions - if toDelete <= 0 { + if len(victims) == 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. + // Delete only after the scan has finished. 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 { + for _, victim := range victims { + if err := bucket.Delete(victim.key); err != nil { return fmt.Errorf("failed to delete old session: %w", err) } - if candidates[i].active { + if victim.tier == sessionTierActive { activeEvicted++ } } if activeEvicted > 0 { - // Only reachable when there was no closed session left to sacrifice. - m.logger.Warnw("Session retention had to evict active sessions", + // Only reachable when there was no closed or unreadable record left to + // sacrifice. + logger.Warnw("Session retention had to evict active sessions", "active_evicted", activeEvicted, "limit", maxSessions) } - m.logger.Debugw("Enforced session retention", "deleted", toDelete, "remaining", maxSessions) + logger.Debugw("Enforced session retention", "deleted", len(victims), "remaining", maxSessions) return nil } +// enforceSessionRetentionOnOpen brings the sessions bucket within the cap at +// database-open time. +// +// The cap is only an invariant if EVERY write path enforces it, and +// CreateSession is not the only one. The legacy-bucket migration moves an +// arbitrary number of records into this bucket, and a process that then only +// updates or closes existing sessions — never creating a new ID — would never +// call retention again and would sit above the cap indefinitely. Running here +// also repairs a bucket left oversized by any older version, whatever the +// reason. +// +// Cost is one scan of an already-capped bucket in the common case. +func enforceSessionRetentionOnOpen(tx *bbolt.Tx, logger *zap.SugaredLogger) error { + bucket := tx.Bucket([]byte(SessionsBucket)) + if bucket == nil { + return nil + } + return trimSessionsToLimit(bucket, sessionRetentionLimit, logger) +} + // GetOAuthToken retrieves an OAuth token for a server from storage func (m *Manager) GetOAuthToken(serverName string) (*OAuthTokenRecord, error) { m.mu.RLock() diff --git a/internal/storage/sessions_retention_test.go b/internal/storage/sessions_retention_test.go index 99caffd4..430fa175 100644 --- a/internal/storage/sessions_retention_test.go +++ b/internal/storage/sessions_retention_test.go @@ -1,12 +1,15 @@ package storage import ( + "encoding/json" "fmt" + "path/filepath" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.etcd.io/bbolt" "go.uber.org/zap" ) @@ -53,6 +56,108 @@ func containsSessionID(sessions []*SessionRecord, id string) bool { return false } +// putRawSession writes a session value verbatim, so tests can plant records +// CreateSession would never produce (corrupt bytes, zero-valued records). +func putRawSession(m *Manager, key string, value []byte) error { + return m.db.db.Update(func(tx *bbolt.Tx) error { + b, err := tx.CreateBucketIfNotExists([]byte(SessionsBucket)) + if err != nil { + return err + } + return b.Put([]byte(key), value) + }) +} + +func countSessions(t *testing.T, m *Manager) int { + t.Helper() + var n int + require.NoError(t, m.db.db.View(func(tx *bbolt.Tx) error { + b := tx.Bucket([]byte(SessionsBucket)) + if b == nil { + return nil + } + return b.ForEach(func(_, _ []byte) error { n++; return nil }) + })) + return n +} + +func rawSessionExists(t *testing.T, m *Manager, key string) bool { + t.Helper() + var found bool + require.NoError(t, m.db.db.View(func(tx *bbolt.Tx) error { + b := tx.Bucket([]byte(SessionsBucket)) + if b == nil { + return nil + } + found = b.Get([]byte(key)) != nil + return nil + })) + return found +} + +// The cap is only an invariant if EVERY write path enforces it. CreateSession is +// not the only one: the legacy-bucket migration moves an arbitrary number of +// records into the sessions bucket at open time. A database that migrates 500 +// records and then only ever updates or closes existing sessions — never +// creating a new ID — would sit above the cap forever. +func TestMigrateLegacySessions_LeavesBucketWithinRetentionLimit(t *testing.T) { + dir := t.TempDir() + + // Build a pre-migration database by hand: MCP records in the shared bucket, + // plus a user login that must survive untouched. + db, err := bbolt.Open(filepath.Join(dir, "config.db"), 0o600, nil) + require.NoError(t, err) + + base := time.Now().Add(-96 * time.Hour) + const legacyCount = 500 + require.NoError(t, db.Update(func(tx *bbolt.Tx) error { + b, err := tx.CreateBucketIfNotExists([]byte(LegacySessionsBucket)) + if err != nil { + return err + } + for i := 0; i < legacyCount; i++ { + rec := makeSession(fmt.Sprintf("legacy-%04d", i), base, + time.Duration(i)*time.Minute, time.Duration(i)*time.Minute, "closed") + data, err := json.Marshal(rec) + if err != nil { + return err + } + if err := b.Put([]byte(fmt.Sprintf("%d_%s", rec.StartTime.UnixNano(), rec.ID)), data); err != nil { + return err + } + } + login, err := json.Marshal(authSession{ + ID: "login-1", UserID: "u1", BearerToken: "jwt", ExpiresAt: time.Now().Add(24 * time.Hour), + }) + if err != nil { + return err + } + return b.Put([]byte("login-1"), login) + })) + require.NoError(t, db.Close()) + + // Opening the database runs the migration. + m, err := NewManager(dir, zap.NewNop().Sugar()) + require.NoError(t, err) + t.Cleanup(func() { _ = m.Close() }) + + assert.Equal(t, sessionRetentionLimit, countSessions(t, m), + "migration must leave the sessions bucket within the cap — it is a write path like any other") + + // The newest records are the ones kept, and the user login is untouched. + _, err = m.GetSessionByID(fmt.Sprintf("legacy-%04d", legacyCount-1)) + assert.NoError(t, err, "the newest migrated record must survive") + _, err = m.GetSessionByID("legacy-0000") + assert.Error(t, err, "the oldest migrated record must be trimmed") + + require.NoError(t, m.db.db.View(func(tx *bbolt.Tx) error { + b := tx.Bucket([]byte(LegacySessionsBucket)) + require.NotNil(t, b, "the legacy bucket must survive — it still holds the user login") + assert.NotNil(t, b.Get([]byte("login-1")), "a user login must never be trimmed by session retention") + return nil + })) +} + // 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 @@ -178,28 +283,125 @@ func TestEnforceSessionRetention_EvictsOldestClosedFirst(t *testing.T) { "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) { +// A session with no LastActivity (older records predate the field) must be +// ranked by StartTime — its actual evidence — rather than sorting as "epoch +// zero" or as "right now". +// +// Both directions are asserted on purpose. Only proving that a zero-activity +// record SURVIVES is a test that passes for the wrong reason: substituting +// time.Now() for the fallback, which would make every legacy record eternally +// the freshest thing in the bucket, passes that assertion too. So a +// zero-activity record with an OLD StartTime must be evicted, and one with a +// NEW StartTime must survive, in the same bucket. +func TestEnforceSessionRetention_ZeroLastActivityIsRankedByStartTime(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", + // Legacy record with no LastActivity and a start time older than everything + // else in the bucket. StartTime is the only evidence, and it says: evict. + stale := &SessionRecord{ + ID: "legacy-stale", + StartTime: base.Add(-10 * time.Hour), + Status: "active", + } + require.NoError(t, m.CreateSession(stale)) + + // Legacy record with no LastActivity and a start time newer than everything + // else. Same missing field, opposite verdict: keep. + fresh := &SessionRecord{ + ID: "legacy-fresh", StartTime: base.Add(10 * time.Hour), Status: "active", } - require.NoError(t, m.CreateSession(legacy)) + require.NoError(t, m.CreateSession(fresh)) - // Older active records with equally-old activity. + // Enough middle-aged active records to force both legacy records to be judged. 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") + _, err := m.GetSessionByID("legacy-stale") + require.Error(t, err, + "a zero LastActivity must be ranked by StartTime — an ancient legacy record "+ + "must not be treated as freshly active and outrank everything") + + got, err := m.GetSessionByID("legacy-fresh") require.NoError(t, err, "a zero LastActivity must fall back to StartTime, not sort as the year 1") assert.Equal(t, "active", got.Status) } + +// An unreadable record cannot be shown, closed, or updated by anything — it is +// the one record with no value at all, so it must be evicted before any usable +// record. +// +// Deriving that from field values alone does not work: an unmarshal failure +// leaves the zero value, which is indistinguishable from a VALID record whose +// status is not "active" and whose timestamps are zero. Both then land in the +// same tier with the same timestamp and the key-order tiebreak decides, which +// can delete the usable record and keep the corrupt one. Unreadability is +// therefore its own explicit tier. +func TestEnforceSessionRetention_EvictsUnreadableRecordsBeforeUsableOnes(t *testing.T) { + m := newRetentionTestManager(t) + + // A valid record that looks exactly like the zero value: closed, no + // timestamps at all. Its key sorts BEFORE the corrupt one, so a key-order + // tiebreak would choose this one as the victim. + require.NoError(t, putRawSession(m, "00000000000000000001_valid-zero", []byte( + `{"id":"valid-zero","status":"closed"}`))) + + // A record that cannot be parsed. Higher key, so only an explicit tier can + // make it lose. + require.NoError(t, putRawSession(m, "00000000000000000002_corrupt", []byte(`{"id":`))) + + // Exactly one record over the limit, so exactly one must go. + base := time.Now().Add(-time.Hour) + for i := 0; i < sessionRetentionLimit-1; 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)) + } + + require.Equal(t, sessionRetentionLimit, countSessions(t, m)) + assert.False(t, rawSessionExists(t, m, "00000000000000000002_corrupt"), + "the unreadable record must be the one evicted") + assert.True(t, rawSessionExists(t, m, "00000000000000000001_valid-zero"), + "a usable record must not be deleted while an unreadable one survives") +} + +// The bounded selection must pick the same victims a full sort would, including +// when the bucket starts far above the limit (the migration case). +func TestEnforceSessionRetention_TrimsFromFarAboveTheLimitInOneGo(t *testing.T) { + m := newRetentionTestManager(t) + base := time.Now().Add(-72 * time.Hour) + + const total = 500 + // Seed far above the cap without going through CreateSession, so retention + // has to do the whole trim in a single call. + for i := 0; i < total; i++ { + rec := makeSession(fmt.Sprintf("bulk-%04d", i), base, + time.Duration(i)*time.Minute, time.Duration(i)*time.Minute, "closed") + data, err := json.Marshal(rec) + require.NoError(t, err) + require.NoError(t, putRawSession(m, + fmt.Sprintf("%d_%s", rec.StartTime.UnixNano(), rec.ID), data)) + } + require.Equal(t, total, countSessions(t, m)) + + require.NoError(t, m.db.db.Update(func(tx *bbolt.Tx) error { + return m.enforceSessionRetention(tx.Bucket([]byte(SessionsBucket)), sessionRetentionLimit) + })) + + assert.Equal(t, sessionRetentionLimit, countSessions(t, m)) + + // The survivors must be the newest 100, exactly as a full sort would choose. + for i := total - sessionRetentionLimit; i < total; i++ { + _, err := m.GetSessionByID(fmt.Sprintf("bulk-%04d", i)) + assert.NoError(t, err, "bulk-%04d is among the newest 100 and must survive", i) + } + for i := 0; i < 5; i++ { + _, err := m.GetSessionByID(fmt.Sprintf("bulk-%04d", i)) + assert.Error(t, err, "bulk-%04d is among the oldest and must be evicted", i) + } +} From 7ede87395a33c5de79c364094847bc7757bde406 Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Wed, 29 Jul 2026 21:48:49 +0300 Subject: [PATCH 3/7] test(storage): cover the case where key order and rank disagree Mutation testing found the bounded selection was under-tested: replacing the "candidate is worse than every survivor" short-circuit with a constant false passed the whole suite. Every existing test feeds records whose rank happens to ascend with their key, so the short-circuit branch was never taken - the displacement branch was always the correct one. Add the case that needs it: a fleet of long-lived active sessions at low keys, then a burst of closed sessions with newer start times at high keys. Every closed record ranks below every active one despite sorting after it, so the selection has to reject them outright rather than displace records it already decided to keep. --- internal/storage/sessions_retention_test.go | 41 +++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/internal/storage/sessions_retention_test.go b/internal/storage/sessions_retention_test.go index 430fa175..6a011f26 100644 --- a/internal/storage/sessions_retention_test.go +++ b/internal/storage/sessions_retention_test.go @@ -370,6 +370,47 @@ func TestEnforceSessionRetention_EvictsUnreadableRecordsBeforeUsableOnes(t *test "a usable record must not be deleted while an unreadable one survives") } +// Records arrive in key order, which is start-time order — but rank is not +// start-time order. When a record that sorts LATE by key is WORSE by rank, the +// bounded selection has to recognise it as a victim outright rather than +// displacing a better record that it already decided to keep. +// +// This is the case where key order and rank point in opposite directions: a +// fleet of long-lived active sessions, then a burst of closed sessions with +// newer start times. Every closed record is worse than every active one despite +// having a higher key, so all of them must lose. +func TestEnforceSessionRetention_EvictsLateKeyedRecordsThatRankWorse(t *testing.T) { + m := newRetentionTestManager(t) + base := time.Now().Add(-24 * time.Hour) + + // Long-lived active sessions occupying the whole budget, at LOW keys. + for i := 0; i < sessionRetentionLimit; i++ { + s := makeSession(fmt.Sprintf("active-%04d", i), base, + time.Duration(i)*time.Minute, time.Duration(i)*time.Minute, "active") + require.NoError(t, m.CreateSession(s)) + } + + // A burst of finished sessions with NEWER start times, at HIGH keys. + const closedBurst = 50 + for i := 0; i < closedBurst; i++ { + s := makeSession(fmt.Sprintf("closed-%04d", i), base, + time.Duration(200+i)*time.Minute, time.Duration(200+i)*time.Minute, "closed") + require.NoError(t, m.CreateSession(s)) + } + + assert.Equal(t, sessionRetentionLimit, countSessions(t, m)) + + // Every active session survives; every closed one is gone, newer keys and all. + for i := 0; i < sessionRetentionLimit; i++ { + _, err := m.GetSessionByID(fmt.Sprintf("active-%04d", i)) + assert.NoError(t, err, "active-%04d must survive a burst of newer closed sessions", i) + } + for i := 0; i < closedBurst; i++ { + _, err := m.GetSessionByID(fmt.Sprintf("closed-%04d", i)) + assert.Error(t, err, "closed-%04d ranks below every active session and must be evicted", i) + } +} + // The bounded selection must pick the same victims a full sort would, including // when the bucket starts far above the limit (the migration case). func TestEnforceSessionRetention_TrimsFromFarAboveTheLimitInOneGo(t *testing.T) { From 50a6c341218db9e895168e4881bbcdc0aed1f55e Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Wed, 29 Jul 2026 21:59:55 +0300 Subject: [PATCH 4/7] fix(storage): bound retention memory, not just its sorting work Cross-model review round 2 on PR #928. The heap bounded the SORT but not the memory. Every rejected or displaced candidate was appended to a victims slice, key and all, so an oversized bucket still cost O(n) state - on precisely the path the bounding was introduced for, the open-time repair of an old database. The comment claimed O(maxSessions) state; the code did not deliver it. Retention is now two passes: Pass 1 ranks and keeps only the heap of the best maxSessions records. Victims are counted, not collected - remembering them is what made state O(n). Pass 2 deletes everything that is not a survivor, in bounded batches, reading keys only. Deleting still cannot happen during traversal, so each batch is collected, the traversal abandoned, the batch deleted, and the scan resumed from the first key the batch did not reach. Peak state is O(maxSessions + retentionDeleteBatch) whatever the bucket size, and the expensive JSON decode still happens exactly once per record across both passes. Also pin what "runs on every open" actually means. The existing regression test only fed records in through the legacy migration, so moving the retention call back inside migrateLegacySessions would still have passed while silently dropping the repair of an already-oversized namespaced bucket. The new test seeds the namespaced bucket directly, closes, reopens, and asserts the trim - no legacy bucket anywhere in the picture. The far-above-the-limit test now uses 1200 records so the delete pass is split across several batches, which is the only way the resume path gets exercised, and asserts the survivors are exactly the newest 100 with no gaps - the shape a mis-resumed scan would break. It seeds in a single transaction, which also takes it from 4.5s to 0.06s. --- internal/storage/manager.go | 106 +++++++++++++++----- internal/storage/sessions_retention_test.go | 94 +++++++++++++---- 2 files changed, 154 insertions(+), 46 deletions(-) diff --git a/internal/storage/manager.go b/internal/storage/manager.go index 2daa22c7..7e011c33 100644 --- a/internal/storage/manager.go +++ b/internal/storage/manager.go @@ -1224,6 +1224,12 @@ type SessionRecord struct { // enforceSessionRetention). const sessionRetentionLimit = 100 +// retentionDeleteBatch bounds how many victim keys retention holds at once. +// Deleting cannot happen while a cursor is traversing the bucket, so victims +// are collected in batches; this is what keeps peak memory independent of how +// oversized the bucket is. +const retentionDeleteBatch = 512 + // CreateSession creates a new session record func (m *Manager) CreateSession(session *SessionRecord) error { m.mu.Lock() @@ -1824,57 +1830,103 @@ func (m *Manager) enforceSessionRetention(bucket *bbolt.Bucket, maxSessions int) // path that can put the bucket over the cap — CreateSession is not the only // one (see enforceSessionRetentionOnOpen). // -// Selection is bounded: rather than collecting and sorting every record, it -// keeps a heap of the best maxSessions seen so far and streams everything else -// straight into the victim list. That makes ranking O(n log maxSessions) with -// O(maxSessions) ranking state, instead of O(n log n) with O(n) state, which -// matters on the one pass that can see a bucket far above the cap: the -// migration/open-time trim. +// Memory is bounded in BOTH directions, which takes two passes: +// +// Pass 1 ranks, keeping only a heap of the best maxSessions records seen so +// far. Victims are counted but not collected — remembering them is what would +// make state O(n). +// Pass 2 deletes everything that is not a survivor, in bounded batches, and +// reads keys only. The expensive JSON decode therefore still happens exactly +// once per record across both passes. +// +// Peak state is O(maxSessions + retentionDeleteBatch) regardless of how +// oversized the bucket is, which is the point: the pass that can see a bucket +// far above the cap is the open-time repair of an old database, and that is +// exactly when a spike proportional to the bucket would hurt. func trimSessionsToLimit(bucket *bbolt.Bucket, maxSessions int, logger *zap.SugaredLogger) error { if maxSessions <= 0 { return nil } - // bucket.Stats() is deliberately not used to size this: inside a write - // transaction it does not account for records Put earlier in the same + // Pass 1 — rank. bucket.Stats() is deliberately not used to 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. survivors := make(survivorHeap, 0, maxSessions) - var victims []sessionEvictionCandidate + total := 0 + activeEvicted := 0 c := bucket.Cursor() for k, v := c.First(); k != nil; k, v = c.Next() { + total++ cand := classifySessionForRetention(k, v, logger) switch { case len(survivors) < maxSessions: heap.Push(&survivors, cand) case cand.worseThan(survivors[0]): - // Weaker than every record currently being kept. - victims = append(victims, cand) + // Weaker than every record currently being kept, so it is a victim. + // Its key is not retained; pass 2 rediscovers it as "not a survivor". + if cand.tier == sessionTierActive { + activeEvicted++ + } default: - displaced, ok := heap.Pop(&survivors).(sessionEvictionCandidate) - if ok { - victims = append(victims, displaced) + if displaced, ok := heap.Pop(&survivors).(sessionEvictionCandidate); ok && + displaced.tier == sessionTierActive { + activeEvicted++ } heap.Push(&survivors, cand) } } - if len(victims) == 0 { + if total <= maxSessions { return nil } - // Delete only after the scan has finished. 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 _, victim := range victims { - if err := bucket.Delete(victim.key); err != nil { - return fmt.Errorf("failed to delete old session: %w", err) + // The survivor set, by key: exactly maxSessions entries. + keep := make(map[string]struct{}, len(survivors)) + for _, s := range survivors { + keep[string(s.key)] = struct{}{} + } + + // Pass 2 — delete. Deleting cannot happen while a cursor is traversing the + // bucket: 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"). So each batch is collected, the + // traversal abandoned, the batch deleted, and the scan resumed from the + // first key the batch did not reach. + deleted := 0 + var resume []byte + for { + batch := make([][]byte, 0, retentionDeleteBatch) + + cur := bucket.Cursor() + var k []byte + if resume == nil { + k, _ = cur.First() + } else { + k, _ = cur.Seek(resume) } - if victim.tier == sessionTierActive { - activeEvicted++ + for ; k != nil && len(batch) < retentionDeleteBatch; k, _ = cur.Next() { + if _, survives := keep[string(k)]; survives { + continue + } + batch = append(batch, append([]byte(nil), k...)) + } + + // k is the first key this batch did not examine. It is never one of the + // keys about to be deleted, so it is still there to seek back to. + if k != nil { + resume = append([]byte(nil), k...) + } + + for _, key := range batch { + if err := bucket.Delete(key); err != nil { + return fmt.Errorf("failed to delete old session: %w", err) + } + deleted++ + } + + if k == nil { + break } } @@ -1884,7 +1936,7 @@ func trimSessionsToLimit(bucket *bbolt.Bucket, maxSessions int, logger *zap.Suga logger.Warnw("Session retention had to evict active sessions", "active_evicted", activeEvicted, "limit", maxSessions) } - logger.Debugw("Enforced session retention", "deleted", len(victims), "remaining", maxSessions) + logger.Debugw("Enforced session retention", "deleted", deleted, "remaining", maxSessions) return nil } diff --git a/internal/storage/sessions_retention_test.go b/internal/storage/sessions_retention_test.go index 6a011f26..6fd42201 100644 --- a/internal/storage/sessions_retention_test.go +++ b/internal/storage/sessions_retention_test.go @@ -68,6 +68,33 @@ func putRawSession(m *Manager, key string, value []byte) error { }) } +// seedClosedSessions writes count closed session records in a SINGLE +// transaction, bypassing CreateSession (and therefore retention). Used to build +// an oversized bucket cheaply — one Update per record makes 1000+ record +// fixtures unbearably slow. +func seedClosedSessions(t *testing.T, m *Manager, count int, base time.Time) { + t.Helper() + require.NoError(t, m.db.db.Update(func(tx *bbolt.Tx) error { + b, err := tx.CreateBucketIfNotExists([]byte(SessionsBucket)) + if err != nil { + return err + } + for i := 0; i < count; i++ { + rec := makeSession(fmt.Sprintf("bulk-%05d", i), base, + time.Duration(i)*time.Second, time.Duration(i)*time.Second, "closed") + data, err := json.Marshal(rec) + if err != nil { + return err + } + key := fmt.Sprintf("%d_%s", rec.StartTime.UnixNano(), rec.ID) + if err := b.Put([]byte(key), data); err != nil { + return err + } + } + return nil + })) +} + func countSessions(t *testing.T, m *Manager) int { t.Helper() var n int @@ -412,22 +439,17 @@ func TestEnforceSessionRetention_EvictsLateKeyedRecordsThatRankWorse(t *testing. } // The bounded selection must pick the same victims a full sort would, including -// when the bucket starts far above the limit (the migration case). -func TestEnforceSessionRetention_TrimsFromFarAboveTheLimitInOneGo(t *testing.T) { +// when the bucket starts far above the limit — and the delete pass must survive +// being split across several batches, resuming the scan in the right place each +// time. 1200 records is comfortably more than one batch. +func TestEnforceSessionRetention_TrimsFromFarAboveTheLimitAcrossBatches(t *testing.T) { m := newRetentionTestManager(t) base := time.Now().Add(-72 * time.Hour) - const total = 500 - // Seed far above the cap without going through CreateSession, so retention - // has to do the whole trim in a single call. - for i := 0; i < total; i++ { - rec := makeSession(fmt.Sprintf("bulk-%04d", i), base, - time.Duration(i)*time.Minute, time.Duration(i)*time.Minute, "closed") - data, err := json.Marshal(rec) - require.NoError(t, err) - require.NoError(t, putRawSession(m, - fmt.Sprintf("%d_%s", rec.StartTime.UnixNano(), rec.ID), data)) - } + const total = 1200 + require.Greater(t, total-sessionRetentionLimit, retentionDeleteBatch, + "the fixture must force more than one delete batch, or the resume path is untested") + seedClosedSessions(t, m, total, base) require.Equal(t, total, countSessions(t, m)) require.NoError(t, m.db.db.Update(func(tx *bbolt.Tx) error { @@ -436,13 +458,47 @@ func TestEnforceSessionRetention_TrimsFromFarAboveTheLimitInOneGo(t *testing.T) assert.Equal(t, sessionRetentionLimit, countSessions(t, m)) - // The survivors must be the newest 100, exactly as a full sort would choose. + // The survivors must be exactly the newest 100 — every one of them, with no + // gaps, which is what a mis-resumed batch scan would produce. for i := total - sessionRetentionLimit; i < total; i++ { - _, err := m.GetSessionByID(fmt.Sprintf("bulk-%04d", i)) - assert.NoError(t, err, "bulk-%04d is among the newest 100 and must survive", i) + _, err := m.GetSessionByID(fmt.Sprintf("bulk-%05d", i)) + assert.NoError(t, err, "bulk-%05d is among the newest 100 and must survive", i) } - for i := 0; i < 5; i++ { - _, err := m.GetSessionByID(fmt.Sprintf("bulk-%04d", i)) - assert.Error(t, err, "bulk-%04d is among the oldest and must be evicted", i) + for i := 0; i < total-sessionRetentionLimit; i++ { + _, err := m.GetSessionByID(fmt.Sprintf("bulk-%05d", i)) + assert.Error(t, err, "bulk-%05d is outside the newest 100 and must be evicted", i) } } + +// "Retention runs on every open" is the claim; "retention runs as part of the +// legacy migration" is a strictly weaker behaviour that satisfies the migration +// test alone. This pins the difference: an already-oversized NAMESPACED bucket, +// no legacy bucket anywhere in the picture, must be trimmed when the database is +// reopened. +func TestEnforceSessionRetention_TrimsAnOversizedBucketOnReopen(t *testing.T) { + dir := t.TempDir() + base := time.Now().Add(-48 * time.Hour) + + m, err := NewManager(dir, zap.NewNop().Sugar()) + require.NoError(t, err) + + const total = 400 + seedClosedSessions(t, m, total, base) + require.Equal(t, total, countSessions(t, m)) + require.NoError(t, m.Close()) + + // Reopen. Nothing here involves the legacy bucket — the records are already + // in the namespaced one. + reopened, err := NewManager(dir, zap.NewNop().Sugar()) + require.NoError(t, err) + t.Cleanup(func() { _ = reopened.Close() }) + + assert.Equal(t, sessionRetentionLimit, countSessions(t, reopened), + "opening a database must bring an oversized sessions bucket within the cap, "+ + "however the records got there") + + _, err = reopened.GetSessionByID(fmt.Sprintf("bulk-%05d", total-1)) + assert.NoError(t, err, "the newest record must survive the open-time trim") + _, err = reopened.GetSessionByID("bulk-00000") + assert.Error(t, err, "the oldest record must be trimmed on open") +} From 34c2a4c08c648572dacfd149d7dc63c4ce2d79d9 Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Wed, 29 Jul 2026 22:11:50 +0300 Subject: [PATCH 5/7] docs(storage): state what retention actually bounds, and what it does not Cross-model review round 3 on PR #928. The finding is correct and the claim was wrong for the third time on this branch. The 512-key batches bound the explicit victim-key slice. They do not bound peak memory for the trim. Every Bucket.Delete seeks to the key and calls Cursor.node(), which materialises the containing leaf page as an in-memory node cached in the bucket and retained until the transaction spills at commit. A trim touching P pages holds O(P) - in practice O(n) - however the victim keys are batched. Measured: a 100k-record trim holds ~90 MB inside the transaction (10k holds ~9 MB, linear as expected), of which the batched victim slice is ~27 KB. Batching bounds ~0.03% of the peak. Restating rather than re-engineering, deliberately. Delivering the bound means committing between batches, which trades away the property that migration and retention apply atomically under bbolt's exclusive file lock. That trade is not worth making for a bucket that has no known way to get large: enforceSessionRetention has been called from CreateSession since session tracking was introduced (it is present in the commit that added the feature), and now runs on every open too, so the bucket has never been uncapped. The realistic oversized case is the off-by-one that let it settle at 101 records, not 10^5. No behaviour change - comments only. The batching stays: it is correct, tested and does bound one real term, just not the dominant one. --- internal/storage/manager.go | 44 +++++++++++++++++++++++++++---------- 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/internal/storage/manager.go b/internal/storage/manager.go index 7e011c33..98ffb3e9 100644 --- a/internal/storage/manager.go +++ b/internal/storage/manager.go @@ -1226,8 +1226,10 @@ const sessionRetentionLimit = 100 // retentionDeleteBatch bounds how many victim keys retention holds at once. // Deleting cannot happen while a cursor is traversing the bucket, so victims -// are collected in batches; this is what keeps peak memory independent of how -// oversized the bucket is. +// are collected in batches. +// +// This bounds the victim slice only. It does NOT bound peak memory for the +// trim — bbolt's own transaction state dominates that. See trimSessionsToLimit. const retentionDeleteBatch = 512 // CreateSession creates a new session record @@ -1830,19 +1832,37 @@ func (m *Manager) enforceSessionRetention(bucket *bbolt.Bucket, maxSessions int) // path that can put the bucket over the cap — CreateSession is not the only // one (see enforceSessionRetentionOnOpen). // -// Memory is bounded in BOTH directions, which takes two passes: +// It runs in two passes: // // Pass 1 ranks, keeping only a heap of the best maxSessions records seen so -// far. Victims are counted but not collected — remembering them is what would -// make state O(n). -// Pass 2 deletes everything that is not a survivor, in bounded batches, and -// reads keys only. The expensive JSON decode therefore still happens exactly -// once per record across both passes. +// far. Victims are counted but not collected. +// Pass 2 deletes everything that is not a survivor, in batches of +// retentionDeleteBatch, and reads keys only. The expensive JSON decode +// therefore still happens exactly once per record across both passes. +// +// What that bounds, and what it does NOT: +// +// Bounded — the ranking state (a heap of maxSessions, not the whole bucket) +// and the explicit victim-key slice (one batch at a time, not every victim). +// +// NOT bounded — bbolt's transaction state. Every Bucket.Delete seeks to the +// key and calls Cursor.node(), which materialises the containing leaf page as +// an in-memory node cached in the bucket and retained until the transaction +// spills at commit. A trim touching P pages therefore holds O(P) — in practice +// O(n) — however the victim keys are batched. Measured on a 100k-record bucket: +// ~90 MB held inside the transaction, of which the batched victim slice is +// ~27 KB. Batching bounds a real term, but not the dominant one. // -// Peak state is O(maxSessions + retentionDeleteBatch) regardless of how -// oversized the bucket is, which is the point: the pass that can see a bucket -// far above the cap is the open-time repair of an old database, and that is -// exactly when a spike proportional to the bucket would hurt. +// This is accepted rather than fixed. Bounding it means committing between +// batches, which would trade away the property that migration and retention +// apply atomically under bbolt's exclusive file lock. The trade is not worth +// making, because the oversized bucket it protects against has no known way to +// occur: retention has run on every CreateSession since session tracking was +// introduced and now runs on every open, so the bucket has never been +// uncapped. The realistic "oversized" case is the off-by-one that let it settle +// at 101 — not 10^5. If a path is ever found that can put a genuinely large +// number of records in this bucket, this is the comment that has to change with +// it. func trimSessionsToLimit(bucket *bbolt.Bucket, maxSessions int, logger *zap.SugaredLogger) error { if maxSessions <= 0 { return nil From cd10da0ccabfa3806aa351ee6fa97a29006827fe Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Wed, 29 Jul 2026 22:21:12 +0300 Subject: [PATCH 6/7] docs(storage): scope the retention memory claim to in-process writers Cross-model review round 3 confirmation pass, plus two corrections of my own. Comments only, no behaviour change. "No known way to occur" was too broad, and contradicted the comment two lines up: enforceSessionRetentionOnOpen promises to repair an oversized bucket whatever the reason, so claiming nothing can produce one cannot also be true. Direct bbolt manipulation, a database import or merge, or any other process writing the file can produce an arbitrarily large bucket. The claim is now scoped to supported in-process write paths, which is what the evidence actually establishes, and out-of-band writes are named as the excluded case that keeps the open-time repair justified. The in-process claim is now stated as the audit that supports it rather than as a bare assertion: only two writers add keys - CreateSession, with retention immediately after every insert, and the legacy migration, with retention immediately after at open. The other six writers Put on a key they already located by scanning and cannot grow the bucket. Also separate measured from computed. The ~90 MB was measured; the ~27 KB victim slice is arithmetic on retentionDeleteBatch. Both sat under one "Measured" lead-in, which is how a derived number ends up being trusted as an observed one. --- internal/storage/manager.go | 28 +++++++++++++++++++--------- 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/internal/storage/manager.go b/internal/storage/manager.go index 98ffb3e9..abe54e65 100644 --- a/internal/storage/manager.go +++ b/internal/storage/manager.go @@ -1850,19 +1850,29 @@ func (m *Manager) enforceSessionRetention(bucket *bbolt.Bucket, maxSessions int) // an in-memory node cached in the bucket and retained until the transaction // spills at commit. A trim touching P pages therefore holds O(P) — in practice // O(n) — however the victim keys are batched. Measured on a 100k-record bucket: -// ~90 MB held inside the transaction, of which the batched victim slice is -// ~27 KB. Batching bounds a real term, but not the dominant one. +// ~90 MB held inside the transaction. The victim slice is ~27 KB by +// construction (retentionDeleteBatch keys) — computed, not measured. Batching +// bounds a real term, but not the dominant one. // // This is accepted rather than fixed. Bounding it means committing between // batches, which would trade away the property that migration and retention // apply atomically under bbolt's exclusive file lock. The trade is not worth -// making, because the oversized bucket it protects against has no known way to -// occur: retention has run on every CreateSession since session tracking was -// introduced and now runs on every open, so the bucket has never been -// uncapped. The realistic "oversized" case is the off-by-one that let it settle -// at 101 — not 10^5. If a path is ever found that can put a genuinely large -// number of records in this bucket, this is the comment that has to change with -// it. +// making, because no supported, IN-PROCESS write path produces the oversized +// bucket it would defend. Only two writers add keys to this bucket: +// CreateSession, where retention runs immediately after every insert, and the +// legacy migration, where it runs immediately after at open. Every other writer +// Puts on a key it already located by scanning, so it cannot grow the bucket. +// The realistic "oversized" case is the off-by-one that let it settle at 101 — +// not 10^5. +// +// Out-of-band writes are deliberately EXCLUDED from that claim. Direct bbolt +// manipulation, a database import or merge, or any other process writing this +// file can produce an arbitrarily large bucket — which is precisely why +// enforceSessionRetentionOnOpen repairs an oversized bucket whatever the +// reason, and why this function must stay correct (if not thrifty) at any size. +// If such a database turns up in practice, or a supported in-process path is +// ever added that inserts many records at once, this is the comment that has to +// change with it. func trimSessionsToLimit(bucket *bbolt.Bucket, maxSessions int, logger *zap.SugaredLogger) error { if maxSessions <= 0 { return nil From c98c5379a8601016196d8727cdd69fb129aeeb11 Mon Sep 17 00:00:00 2001 From: Algis Dumbris Date: Thu, 30 Jul 2026 10:32:39 +0300 Subject: [PATCH 7/7] docs(storage): correct a retention comment that describes deleted behaviour MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TestEnforceSessionRetention_OnlyTouchesItsOwnBucket still claimed retention 'keeps the newest 100 by raw key order with no type check'. This PR removed exactly that: records are now decoded and ranked by usefulness. The test's real subject is containment — retention touches only the sessions bucket — so say that, and point at sessions_retention_test.go for eviction order. --- internal/storage/sessions_migration_test.go | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/internal/storage/sessions_migration_test.go b/internal/storage/sessions_migration_test.go index eaef4c44..6831f61c 100644 --- a/internal/storage/sessions_migration_test.go +++ b/internal/storage/sessions_migration_test.go @@ -173,9 +173,11 @@ func TestIsMCPSessionRecord(t *testing.T) { assert.False(t, isMCPSessionRecord(both)) } -// Retention keeps the newest 100 by raw key order with no type check. That is -// safe now only because the bucket is exclusively ours — this test pins the -// behaviour the namespacing depends on. +// Retention ranks every record it can decode and evicts the least useful, so +// it no longer depends on raw key order — but it still only ever touches the +// sessions bucket. That containment is what the namespacing relies on, and it +// is what this test pins. Retention's actual eviction ORDER is covered by +// sessions_retention_test.go, which drives the real CreateSession path. func TestEnforceSessionRetention_OnlyTouchesItsOwnBucket(t *testing.T) { db := openTempDB(t)