Skip to content

Commit aaf25e9

Browse files
sam-saffron-jarvissam-saffron-jarvis
committed
fix: Memory mining/freshen walks completed sessions with OFFSET pagination, making full scans grow quadratically
1 parent 50e59be commit aaf25e9

4 files changed

Lines changed: 166 additions & 32 deletions

File tree

‎cmd/memory_shared.go‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,14 +21,15 @@ func openReadOnlySessionStore(cfg *config.Config) (session.Store, error) {
2121

2222
func listCompleteSessions(ctx context.Context, store session.Store) ([]session.SessionSummary, error) {
2323
const pageSize = 200
24-
offset := 0
24+
var beforeNumber int64
2525
all := make([]session.SessionSummary, 0, pageSize)
2626

2727
for {
2828
page, err := store.List(ctx, session.ListOptions{
29-
Status: session.StatusComplete,
30-
Limit: pageSize,
31-
Offset: offset,
29+
Status: session.StatusComplete,
30+
Limit: pageSize,
31+
BeforeNumber: beforeNumber,
32+
SortByNumberDesc: true,
3233
})
3334
if err != nil {
3435
return nil, err
@@ -40,7 +41,11 @@ func listCompleteSessions(ctx context.Context, store session.Store) ([]session.S
4041
if len(page) < pageSize {
4142
break
4243
}
43-
offset += len(page)
44+
lastNumber := page[len(page)-1].Number
45+
if lastNumber <= 0 {
46+
return nil, fmt.Errorf("complete session %s is missing a session number", page[len(page)-1].ID)
47+
}
48+
beforeNumber = lastNumber
4449
}
4550

4651
return all, nil

‎internal/session/sqlite.go‎

Lines changed: 31 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1349,12 +1349,19 @@ func (s *SQLiteStore) List(ctx context.Context, opts ListOptions) ([]SessionSumm
13491349
}
13501350
messageCountCol = "(SELECT COUNT(*) FROM messages WHERE session_id = s.id AND role IN ('user', 'assistant')" + compactionTailClause + ")"
13511351
}
1352+
fromClause := "FROM sessions s"
1353+
if opts.SortByNumberDesc {
1354+
// Completed-session walks page by descending session number. Force the
1355+
// number index so SQLite can continue scanning from the last seen number
1356+
// instead of picking a filter-only index and re-sorting each page.
1357+
fromClause = "FROM sessions s INDEXED BY idx_sessions_number"
1358+
}
13521359
query := `
13531360
SELECT s.id, s.number, s.name, s.summary, ` + generatedShortCol + `, ` + generatedLongCol + `, ` + titleSourceCol + `,
13541361
s.provider, COALESCE(s.provider_key, ''), s.model, s.mode, ` + originCol + `, s.archived, ` + pinnedCol + `, s.created_at, s.updated_at, ` + lastMessageAtCol + `,
13551362
` + messageCountCol + ` as message_count,
13561363
s.user_turns, s.llm_turns, s.tool_calls, s.input_tokens, s.cached_input_tokens, ` + cacheWriteCol + `, s.output_tokens, s.status, s.tags
1357-
FROM sessions s
1364+
` + fromClause + `
13581365
WHERE 1=1`
13591366
args := []any{}
13601367

@@ -1418,26 +1425,34 @@ func (s *SQLiteStore) List(ctx context.Context, opts ListOptions) ([]SessionSumm
14181425
query += " AND 1 = 0"
14191426
}
14201427
}
1428+
if opts.BeforeNumber > 0 {
1429+
query += " AND s.number < ?"
1430+
args = append(args, opts.BeforeNumber)
1431+
}
14211432
if !opts.Archived {
14221433
query += " AND s.archived = FALSE"
14231434
}
14241435

1425-
// Sort by last user message time (when the user last interacted), falling back
1426-
// to created_at for sessions with no user messages yet. This prevents background
1427-
// activity (autotitle, mining, status changes) from reordering the sidebar.
1428-
// Web sidebar callers set SortByActivity to use last_message_at instead so
1429-
// assistant-only turns also surface (keeps the top-N window aligned with the
1430-
// client-side "any-message" ordering).
1431-
sortCol := "s.updated_at"
1432-
if opts.SortByActivity && s.hasLastMessageAt {
1433-
sortCol = "COALESCE(s.last_message_at, s.last_user_message_at, s.created_at)"
1434-
} else if s.hasLastUserMessageAt {
1435-
sortCol = "COALESCE(s.last_user_message_at, s.created_at)"
1436-
}
1437-
if s.hasPinned {
1438-
query += " ORDER BY COALESCE(s.pinned, FALSE) DESC, " + sortCol + " DESC"
1436+
if opts.SortByNumberDesc {
1437+
query += " ORDER BY s.number DESC"
14391438
} else {
1440-
query += " ORDER BY " + sortCol + " DESC"
1439+
// Sort by last user message time (when the user last interacted), falling back
1440+
// to created_at for sessions with no user messages yet. This prevents background
1441+
// activity (autotitle, mining, status changes) from reordering the sidebar.
1442+
// Web sidebar callers set SortByActivity to use last_message_at instead so
1443+
// assistant-only turns also surface (keeps the top-N window aligned with the
1444+
// client-side "any-message" ordering).
1445+
sortCol := "s.updated_at"
1446+
if opts.SortByActivity && s.hasLastMessageAt {
1447+
sortCol = "COALESCE(s.last_message_at, s.last_user_message_at, s.created_at)"
1448+
} else if s.hasLastUserMessageAt {
1449+
sortCol = "COALESCE(s.last_user_message_at, s.created_at)"
1450+
}
1451+
if s.hasPinned {
1452+
query += " ORDER BY COALESCE(s.pinned, FALSE) DESC, " + sortCol + " DESC"
1453+
} else {
1454+
query += " ORDER BY " + sortCol + " DESC"
1455+
}
14411456
}
14421457

14431458
limit := opts.Limit

‎internal/session/sqlite_test.go‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import (
66
"fmt"
77
"os"
88
"path/filepath"
9+
"strings"
910
"testing"
1011
"time"
1112

@@ -41,6 +42,117 @@ func TestNewSQLiteStoreMemoryDBUsesSingleConnection(t *testing.T) {
4142
}
4243
}
4344

45+
func TestSQLiteStoreListByNumberCursorReturnsCompleteSessions(t *testing.T) {
46+
t.Setenv("XDG_DATA_HOME", t.TempDir())
47+
48+
store, err := NewSQLiteStore(DefaultConfig())
49+
if err != nil {
50+
t.Fatalf("NewSQLiteStore: %v", err)
51+
}
52+
defer store.Close()
53+
54+
ctx := context.Background()
55+
seed := []struct {
56+
status SessionStatus
57+
archived bool
58+
}{
59+
{status: StatusComplete},
60+
{status: StatusComplete},
61+
{status: StatusActive},
62+
{status: StatusComplete},
63+
{status: StatusComplete},
64+
{status: StatusComplete, archived: true},
65+
}
66+
for i, tc := range seed {
67+
sess := &Session{
68+
ID: NewID(),
69+
Provider: "test",
70+
Model: "test-model",
71+
Mode: ModeChat,
72+
Status: tc.status,
73+
Archived: tc.archived,
74+
Name: fmt.Sprintf("session-%d", i+1),
75+
}
76+
if err := store.Create(ctx, sess); err != nil {
77+
t.Fatalf("Create(%d): %v", i, err)
78+
}
79+
}
80+
81+
page1, err := store.List(ctx, ListOptions{Status: StatusComplete, Limit: 2, SortByNumberDesc: true})
82+
if err != nil {
83+
t.Fatalf("List page1: %v", err)
84+
}
85+
if len(page1) != 2 {
86+
t.Fatalf("len(page1) = %d, want 2", len(page1))
87+
}
88+
if page1[0].Number != 5 || page1[1].Number != 4 {
89+
t.Fatalf("page1 numbers = [%d %d], want [5 4]", page1[0].Number, page1[1].Number)
90+
}
91+
92+
page2, err := store.List(ctx, ListOptions{Status: StatusComplete, Limit: 2, BeforeNumber: page1[len(page1)-1].Number, SortByNumberDesc: true})
93+
if err != nil {
94+
t.Fatalf("List page2: %v", err)
95+
}
96+
if len(page2) != 2 {
97+
t.Fatalf("len(page2) = %d, want 2", len(page2))
98+
}
99+
if page2[0].Number != 2 || page2[1].Number != 1 {
100+
t.Fatalf("page2 numbers = [%d %d], want [2 1]", page2[0].Number, page2[1].Number)
101+
}
102+
103+
page3, err := store.List(ctx, ListOptions{Status: StatusComplete, Limit: 2, BeforeNumber: page2[len(page2)-1].Number, SortByNumberDesc: true})
104+
if err != nil {
105+
t.Fatalf("List page3: %v", err)
106+
}
107+
if len(page3) != 0 {
108+
t.Fatalf("len(page3) = %d, want 0", len(page3))
109+
}
110+
}
111+
112+
func TestSQLiteStoreListByNumberCursorUsesSessionNumberIndex(t *testing.T) {
113+
store, err := NewSQLiteStore(Config{Enabled: true, Path: ":memory:"})
114+
if err != nil {
115+
t.Fatalf("NewSQLiteStore: %v", err)
116+
}
117+
defer store.Close()
118+
119+
plan := sqliteExplainPlan(t, store.db, `EXPLAIN QUERY PLAN
120+
SELECT s.id, s.number
121+
FROM sessions s INDEXED BY idx_sessions_number
122+
WHERE s.status = ? AND s.archived = FALSE AND s.number < ?
123+
ORDER BY s.number DESC
124+
LIMIT ?`, string(StatusComplete), 1000, 200)
125+
if !strings.Contains(plan, "idx_sessions_number") {
126+
t.Fatalf("query plan = %q, want session number index", plan)
127+
}
128+
if strings.Contains(plan, "USE TEMP B-TREE") {
129+
t.Fatalf("query plan = %q, want no temp sort", plan)
130+
}
131+
}
132+
133+
func sqliteExplainPlan(t *testing.T, db *sql.DB, query string, args ...any) string {
134+
t.Helper()
135+
rows, err := db.Query(query, args...)
136+
if err != nil {
137+
t.Fatalf("explain query: %v", err)
138+
}
139+
defer rows.Close()
140+
141+
var details []string
142+
for rows.Next() {
143+
var id, parent, notused int
144+
var detail string
145+
if err := rows.Scan(&id, &parent, &notused, &detail); err != nil {
146+
t.Fatalf("scan explain row: %v", err)
147+
}
148+
details = append(details, detail)
149+
}
150+
if err := rows.Err(); err != nil {
151+
t.Fatalf("explain rows: %v", err)
152+
}
153+
return strings.Join(details, "\n")
154+
}
155+
44156
func TestSQLiteStoreGetMessagesFromHonorsLimit(t *testing.T) {
45157
t.Setenv("XDG_DATA_HOME", t.TempDir())
46158

‎internal/session/types.go‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -143,17 +143,19 @@ type SessionSummary struct {
143143

144144
// ListOptions configures session listing.
145145
type ListOptions struct {
146-
Name string // Filter by name
147-
Provider string // Filter by provider
148-
Model string // Filter by model
149-
Mode SessionMode // Filter by mode (chat, ask, plan, exec)
150-
Status SessionStatus // Filter by status
151-
Tag string // Filter by tag (substring match)
152-
Categories []string // Sidebar/web categories (all, chat, web, ask, plan, exec)
153-
Limit int // Max results (0 = use default)
154-
Offset int // Pagination offset
155-
Archived bool // Include archived sessions
156-
SortByActivity bool // Sort by last_message_at (web sidebar); defaults to last_user_message_at
146+
Name string // Filter by name
147+
Provider string // Filter by provider
148+
Model string // Filter by model
149+
Mode SessionMode // Filter by mode (chat, ask, plan, exec)
150+
Status SessionStatus // Filter by status
151+
Tag string // Filter by tag (substring match)
152+
Categories []string // Sidebar/web categories (all, chat, web, ask, plan, exec)
153+
Limit int // Max results (0 = use default)
154+
Offset int // Pagination offset
155+
BeforeNumber int64 // Keyset cursor: only sessions with number < this value
156+
SortByNumberDesc bool // Order by session number descending instead of activity sort
157+
Archived bool // Include archived sessions
158+
SortByActivity bool // Sort by last_message_at (web sidebar); defaults to last_user_message_at
157159
}
158160

159161
// SearchResult represents a search match.

0 commit comments

Comments
 (0)