fix(search): bound semantic pagination

src/db/search.rs

@@ -9,6 +9,8 @@ use crate::project_scope::ProjectScope;
 use crate::types::{MatchSource, SearchResult, Session};
 use crate::utils::f32_slice_to_bytes;
 
+const SQLITE_VEC_MAX_K: usize = 4096;
+
 pub(crate) struct SearchEngine<'a> {
     conn: &'a Connection,
 }
@@ -109,38 +111,65 @@ impl<'a> SearchEngine<'a> {
         limit: usize,
         fetch_multiplier: usize,
     ) -> anyhow::Result<Vec<SearchResult>> {
-        let fetch_size = limit * fetch_multiplier;
-        let fts_hits = self.fts_search(query, filters, fetch_size)?;
+        let fetch_size = limit.saturating_mul(fetch_multiplier).max(1);
+        let fts_hits = self.fts_search(query, filters, Some(fetch_size))?;
         let vec_hits = match embedding {
-            Some(e) => self.vec_search(e, filters, fetch_size)?,
+            Some(embedding) => self.vec_search(embedding, filters, fetch_size.saturating_mul(5))?,
             None => vec![],
         };
-        let merged = rrf_merge(&fts_hits, &vec_hits, 10);
+        self.search_results(fts_hits, vec_hits, 0, Some(limit))
+    }
 
-        let session_ids: Vec<&str> =
-            merged.iter().take(limit).map(|(id, _, _)| id.as_str()).collect();
+    pub(crate) fn hybrid_search_page(
+        &self,
+        query: &str,
+        embedding: Option<&[f32]>,
+        filters: &SearchFilters,
+        limit: Option<usize>,
+        offset: usize,
+    ) -> anyhow::Result<Vec<SearchResult>> {
+        if limit == Some(0) {
+            return Ok(vec![]);
+        }
+
+        let fts_hits = self.fts_search(query, filters, None)?;
+        let vec_hits = match embedding {
+            Some(embedding) => self.vec_search(embedding, filters, SQLITE_VEC_MAX_K)?,
+            None => vec![],
+        };
+        self.search_results(fts_hits, vec_hits, offset, limit)
+    }
 
+    fn search_results(
+        &self,
+        fts_hits: Vec<Hit>,
+        vec_hits: Vec<Hit>,
+        offset: usize,
+        limit: Option<usize>,
+    ) -> anyhow::Result<Vec<SearchResult>> {
+        let merged = rrf_merge(&fts_hits, &vec_hits, 10);
+        let snippets: HashMap<_, _> =
+            fts_hits.into_iter().map(|hit| (hit.session_id, hit.snippet)).collect();
+        let limit = limit.unwrap_or(usize::MAX);
+        let session_ids: Vec<&str> =
+            merged.iter().skip(offset).take(limit).map(|(id, _, _)| id.as_str()).collect();
         let sessions = self.load_sessions(&session_ids)?;
 
         let mut results = Vec::new();
-        for (session_id, _score, match_source) in merged.into_iter().take(limit) {
+        for (session_id, _score, match_source) in merged.into_iter().skip(offset).take(limit) {
             if let Some(session) = sessions.get(&session_id) {
-                let snippet = fts_hits
-                    .iter()
-                    .find(|h| h.session_id == session_id)
-                    .and_then(|h| h.snippet.clone());
+                let snippet = snippets.get(&session_id).cloned().flatten();
                 results.push(SearchResult { session: session.clone(), match_source, snippet });
             }
         }
-
         Ok(results)
     }
 
     fn fts_search(
         &self,
         query: &str,
         filters: &SearchFilters,
-        limit: usize,
+        limit: Option<usize>,
     ) -> anyhow::Result<Vec<Hit>> {
         let escaped = fts5_escape(query);
         if escaped.is_empty() {
@@ -160,10 +189,14 @@ impl<'a> SearchEngine<'a> {
         let mut param_idx = 2;
         apply_filters(&mut sql, &mut params, &mut param_idx, filters);
 
-        sql.push_str(&format!(" GROUP BY m.session_id ORDER BY best_rank LIMIT {limit}"));
+        sql.push_str(" GROUP BY m.session_id ORDER BY best_rank, m.session_id");
+        if let Some(limit) = limit {
+            let limit = i64::try_from(limit).unwrap_or(i64::MAX);
+            sql.push_str(&format!(" LIMIT {limit}"));
+        }
 
         let param_refs: Vec<&dyn rusqlite::types::ToSql> =
-            params.iter().map(|p| p.as_ref()).collect();
+            params.iter().map(|param| param.as_ref()).collect();
 
         let mut stmt = self.conn.prepare(&sql)?;
         let rows = stmt.query_map(param_refs.as_slice(), |row| {
@@ -181,10 +214,10 @@ impl<'a> SearchEngine<'a> {
         &self,
         embedding: &[f32],
         filters: &SearchFilters,
-        limit: usize,
+        requested_k: usize,
     ) -> anyhow::Result<Vec<Hit>> {
         let blob = f32_slice_to_bytes(embedding);
-        let fetch_k = (limit * 5) as i64;
+        let fetch_k = requested_k.clamp(1, SQLITE_VEC_MAX_K) as i64;
 
         let mut sql = String::from(
             "SELECT m.session_id, MIN(mv.distance) AS best_distance
@@ -200,10 +233,10 @@ impl<'a> SearchEngine<'a> {
         let mut param_idx = 3;
         apply_filters(&mut sql, &mut params, &mut param_idx, filters);
 
-        sql.push_str(" GROUP BY m.session_id ORDER BY best_distance");
+        sql.push_str(" GROUP BY m.session_id ORDER BY best_distance, m.session_id");
 
         let param_refs: Vec<&dyn rusqlite::types::ToSql> =
-            params.iter().map(|p| p.as_ref()).collect();
+            params.iter().map(|param| param.as_ref()).collect();
 
         let mut stmt = self.conn.prepare(&sql)?;
         let rows = stmt.query_map(param_refs.as_slice(), |row| {
@@ -218,27 +251,25 @@ impl<'a> SearchEngine<'a> {
     }
 
     fn load_sessions(&self, ids: &[&str]) -> anyhow::Result<HashMap<String, Session>> {
-        let mut map = HashMap::new();
-        if ids.is_empty() {
-            return Ok(map);
-        }
+        const SESSION_LOAD_CHUNK_SIZE: usize = 900;
 
-        let placeholders: Vec<String> = (1..=ids.len()).map(|i| format!("?{i}")).collect();
-        let sql = format!(
-            "SELECT {SESSION_COLUMNS}
-             FROM sessions WHERE id IN ({})",
-            placeholders.join(", ")
-        );
-
-        let params: Vec<&dyn rusqlite::types::ToSql> =
-            ids.iter().map(|id| id as &dyn rusqlite::types::ToSql).collect();
-
-        let mut stmt = self.conn.prepare(&sql)?;
-        let rows = stmt.query_map(params.as_slice(), session_from_row)?;
-
-        for row in rows {
-            let session = row?;
-            map.insert(session.id.clone(), session);
+        let mut map = HashMap::new();
+        for ids in ids.chunks(SESSION_LOAD_CHUNK_SIZE) {
+            let placeholders: Vec<String> = (1..=ids.len()).map(|i| format!("?{i}")).collect();
+            let sql = format!(
+                "SELECT {SESSION_COLUMNS}
+                 FROM sessions WHERE id IN ({})",
+                placeholders.join(", ")
+            );
+            let params: Vec<&dyn rusqlite::types::ToSql> =
+                ids.iter().map(|id| id as &dyn rusqlite::types::ToSql).collect();
+            let mut stmt = self.conn.prepare(&sql)?;
+            let rows = stmt.query_map(params.as_slice(), session_from_row)?;
+
+            for row in rows {
+                let session = row?;
+                map.insert(session.id.clone(), session);
+            }
         }
         Ok(map)
     }
@@ -300,7 +331,7 @@ fn rrf_merge(fts_hits: &[Hit], vec_hits: &[Hit], k: u32) -> Vec<(String, f64, Ma
         })
         .collect();
 
-    results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
+    results.sort_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
     results
 }
 

src/integration/regression.rs

@@ -680,6 +680,178 @@ fn hybrid_search_fts_only_without_embedding() {
     assert_eq!(results.len(), 1);
 }
 
+fn seed_semantic_boundary_sessions(store: &Store, count: usize) {
+    store
+        .conn
+        .execute_batch(&format!(
+            "WITH RECURSIVE seq(n) AS (
+                 SELECT 0
+                 UNION ALL
+                 SELECT n + 1 FROM seq WHERE n + 1 < {count}
+             )
+             INSERT INTO sessions (id, source, source_id, title, started_at, message_count)
+             SELECT printf('semantic-%05d', n), 'test', printf('raw-%05d', n),
+                    printf('Semantic session %05d', n), n, 1
+             FROM seq;
+             INSERT INTO messages (session_id, role, content, timestamp, seq)
+             SELECT id, 'user', 'semanticboundary ' || id, started_at, 0
+             FROM sessions
+             WHERE source = 'test';"
+        ))
+        .unwrap();
+}
+
+fn add_semantic_boundary_embedding(store: &Store) -> Vec<f32> {
+    let message_id: i64 = store
+        .conn
+        .query_row("SELECT id FROM messages ORDER BY id LIMIT 1", [], |row| row.get(0))
+        .unwrap();
+    let embedding = vec![0.1f32; 384];
+    store.upsert_embeddings(&[(message_id, &embedding)]).unwrap();
+    embedding
+}
+
+fn seed_semantic_page_fixture(store: &Store) -> Vec<f32> {
+    for index in 0..6 {
+        let id = format!("semantic-fts-{index:02}");
+        let session = make_session(&id, "test", &format!("raw-fts-{index:02}"), "Semantic FTS");
+        store.insert_session(&session).unwrap();
+        store.insert_messages(&[make_message(&id, Role::User, "semanticstable", 0)]).unwrap();
+    }
+
+    let mut filler =
+        make_session("semantic-vec-fill", "test", "raw-vec-fill", "Semantic vector filler");
+    filler.message_count = 30;
+    store.insert_session(&filler).unwrap();
+    let filler_messages = (0..30)
+        .map(|seq| make_message("semantic-vec-fill", Role::User, "semantic filler", seq))
+        .collect::<Vec<_>>();
+    store.insert_messages(&filler_messages).unwrap();
+
+    let mut stmt = store
+        .conn
+        .prepare("SELECT id FROM messages WHERE session_id = 'semantic-vec-fill' ORDER BY seq")
+        .unwrap();
+    let mut message_ids = stmt
+        .query_map([], |row| row.get::<_, i64>(0))
+        .unwrap()
+        .collect::<rusqlite::Result<Vec<_>>>()
+        .unwrap();
+    message_ids.push(
+        store
+            .conn
+            .query_row("SELECT id FROM messages WHERE session_id = 'semantic-fts-04'", [], |row| {
+                row.get(0)
+            })
+            .unwrap(),
+    );
+
+    let vectors = (1..=message_ids.len())
+        .map(|rank| {
+            let mut vector = vec![0.0f32; 384];
+            vector[0] = rank as f32 / 1_000.0;
+            vector
+        })
+        .collect::<Vec<_>>();
+    let embeddings = message_ids
+        .iter()
+        .zip(&vectors)
+        .map(|(&message_id, vector)| (message_id, vector.as_slice()))
+        .collect::<Vec<_>>();
+    store.upsert_embeddings(&embeddings).unwrap();
+
+    vec![0.0f32; 384]
+}
+
+#[test]
+fn semantic_query_crosses_sqlite_vec_boundary() {
+    let store = setup();
+    seed_semantic_boundary_sessions(&store, 300);
+    let embedding = add_semantic_boundary_embedding(&store);
+    let engine = SearchEngine::new(&store.conn);
+
+    let results =
+        engine.hybrid_search("semanticboundary", Some(&embedding), &no_filters(), 274, 3).unwrap();
+    assert_eq!(results.len(), 274);
+
+    let page = engine
+        .hybrid_search_page("semanticboundary", Some(&embedding), &no_filters(), Some(50), 224)
+        .unwrap();
+    assert_eq!(page.len(), 50);
+}
+
+#[test]
+fn semantic_adjacent_pages_follow_one_global_order() {
+    let store = setup();
+    let embedding = seed_semantic_page_fixture(&store);
+    let engine = SearchEngine::new(&store.conn);
+
+    let first_page = engine
+        .hybrid_search_page("semanticstable", Some(&embedding), &no_filters(), Some(2), 0)
+        .unwrap();
+    let second_page = engine
+        .hybrid_search_page("semanticstable", Some(&embedding), &no_filters(), Some(2), 2)
+        .unwrap();
+    let global = engine
+        .hybrid_search_page("semanticstable", Some(&embedding), &no_filters(), None, 0)
+        .unwrap();
+
+    let paged_ids = first_page
+        .iter()
+        .chain(&second_page)
+        .map(|result| result.session.id.as_str())
+        .collect::<Vec<_>>();
+    let global_prefix =
+        global.iter().take(4).map(|result| result.session.id.as_str()).collect::<Vec<_>>();
+
+    assert_eq!(global_prefix[0], "semantic-fts-04");
+    assert_eq!(paged_ids, global_prefix);
+}
+
+#[test]
+fn semantic_query_all_returns_complete_fts_set() {
+    let store = setup();
+    seed_semantic_boundary_sessions(&store, 10_001);
+    let engine = SearchEngine::new(&store.conn);
+
+    let text_results =
+        engine.hybrid_search_page("semanticboundary", None, &no_filters(), None, 0).unwrap();
+    assert_eq!(text_results.len(), 10_001);
+    assert!(text_results.iter().all(|result| {
+        result.snippet.as_deref().and_then(|snippet| snippet.strip_prefix("semanticboundary "))
+            == Some(result.session.id.as_str())
+    }));
+
+    let embedding = add_semantic_boundary_embedding(&store);
+    let semantic_results = engine
+        .hybrid_search_page("semanticboundary", Some(&embedding), &no_filters(), None, 0)
+        .unwrap();
+    assert_eq!(semantic_results.len(), 10_001);
+}
+
+#[test]
+fn semantic_search_arithmetic_is_saturating() {
+    let store = setup();
+    let engine = SearchEngine::new(&store.conn);
+    let embedding = vec![0.1f32; 384];
+
+    let direct = engine
+        .hybrid_search("semanticboundary", Some(&embedding), &no_filters(), usize::MAX, usize::MAX)
+        .unwrap();
+    assert!(direct.is_empty());
+
+    let page = engine
+        .hybrid_search_page(
+            "semanticboundary",
+            Some(&embedding),
+            &no_filters(),
+            Some(usize::MAX),
+            usize::MAX,
+        )
+        .unwrap();
+    assert!(page.is_empty());
+}
+
 #[test]
 fn search_with_source_filter() {
     let store = setup();

src/session.rs

@@ -356,13 +356,15 @@ pub(crate) fn run_session_list(
             scope: scope.clone(),
             thread_role,
         };
-        let search_limit = effective_limit.unwrap_or(10_000).saturating_add(offset).max(1);
-        let results =
-            engine.hybrid_search(query, embedding.as_deref(), &filters, search_limit, 3)?;
+        let results = engine.hybrid_search_page(
+            query,
+            embedding.as_deref(),
+            &filters,
+            effective_limit,
+            offset,
+        )?;
         results
             .into_iter()
-            .skip(offset)
-            .take(effective_limit.unwrap_or(usize::MAX))
             .map(|result| SessionListRow {
                 session: result.session,
                 match_source: Some(result.match_source),
@@ -849,7 +851,7 @@ fn print_session_list_json(
             "next_offset": if all || rows.len() < limit {
                 serde_json::Value::Null
             } else {
-                serde_json::json!(offset + rows.len())
+                serde_json::json!(offset.saturating_add(rows.len()))
             }
         }))?
     );