talos_session/
sqlite.rs

1//! SQLite-based full-text search index for session messages.
2//!
3//! This module provides a [`SessionIndex`] that wraps a SQLite database with FTS5
4//! support, enabling efficient full-text search across session messages. The index
5//! is supplementary to the primary JSONL storage — it is created lazily on first use
6//! and updated whenever sessions are saved.
7//!
8//! # Schema
9//!
10//! - `sessions` table: metadata about indexed sessions
11//! - `messages_fts` FTS5 virtual table: full-text searchable message content
12
13use chrono::{DateTime, Utc};
14use rusqlite::{Connection, OptionalExtension, Result as RusqliteResult, params};
15use std::path::{Path, PathBuf};
16use thiserror::Error;
17use uuid::Uuid;
18
19use crate::{Session, SessionInfo};
20
21/// Errors that can occur during SQLite index operations.
22#[derive(Debug, Error)]
23pub enum IndexError {
24    /// An index store operation failed.
25    #[error("index store error: {0}")]
26    Store(#[from] IndexStoreError),
27
28    /// An I/O error occurred.
29    #[error("I/O error: {0}")]
30    IoError(#[from] std::io::Error),
31
32    /// Failed to parse a UUID.
33    #[error("invalid UUID: {0}")]
34    InvalidUuid(String),
35}
36
37/// Errors originating from the SQLite index store.
38#[derive(Debug, Error)]
39pub enum IndexStoreError {
40    /// A database operation failed.
41    #[error("database operation failed: {0}")]
42    Database(String),
43}
44
45impl From<rusqlite::Error> for IndexStoreError {
46    fn from(err: rusqlite::Error) -> Self {
47        IndexStoreError::Database(err.to_string())
48    }
49}
50
51impl From<rusqlite::Error> for IndexError {
52    fn from(err: rusqlite::Error) -> Self {
53        IndexError::Store(IndexStoreError::from(err))
54    }
55}
56
57/// A search result from the FTS5 index.
58#[derive(Debug, Clone)]
59pub struct SearchResult {
60    /// The session ID that contains the matching message.
61    pub session_id: String,
62
63    /// The project name associated with the session.
64    pub project: String,
65
66    /// The stable workspace identity (canonical absolute path).
67    pub workspace_root: String,
68
69    /// A snippet of the matching content with highlighted matches.
70    pub snippet: String,
71
72    /// The timestamp of the matching message.
73    pub timestamp: DateTime<Utc>,
74
75    /// The BM25 rank of the match (lower is more relevant).
76    pub rank: f64,
77}
78
79/// Metadata about a session fork relationship.
80///
81/// Captures the source session, the forked session, and the entry point
82/// from which the fork was created.
83#[derive(Debug, Clone)]
84pub struct ForkInfo {
85    /// The session ID of the forked (child) session.
86    pub forked_session_id: String,
87
88    /// The entry ID in the source session from which the fork was created.
89    pub fork_entry_id: String,
90
91    /// When the fork was created.
92    pub forked_at: DateTime<Utc>,
93}
94
95/// SQLite-based full-text search index for session messages.
96///
97/// Wraps a [`Connection`] to a SQLite database with FTS5 support. The index is
98/// created lazily on first use and must be explicitly updated when sessions change.
99#[derive(Debug)]
100pub struct SessionIndex {
101    conn: Connection,
102    db_path: PathBuf,
103}
104
105impl SessionIndex {
106    /// Open or create a SQLite database at the given path.
107    ///
108    /// The database file is created if it does not exist. The schema is NOT
109    /// automatically initialized — call [`SessionIndex::init_schema`] after creation.
110    ///
111    /// # Arguments
112    ///
113    /// * `path` — Path to the SQLite database file.
114    ///
115    /// # Errors
116    ///
117    /// Returns an error if the database cannot be opened.
118    pub fn new(path: &Path) -> Result<Self, IndexError> {
119        if let Some(parent) = path.parent() {
120            std::fs::create_dir_all(parent)?;
121        }
122
123        let conn = Connection::open(path)?;
124
125        conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA synchronous=NORMAL;")?;
126
127        Ok(Self {
128            conn,
129            db_path: path.to_path_buf(),
130        })
131    }
132
133    /// Initialize the database schema.
134    ///
135    /// Creates the `sessions` table and the `messages_fts` FTS5 virtual table
136    /// if they do not already exist. Safe to call multiple times.
137    ///
138    /// # Errors
139    ///
140    /// Returns an error if the schema cannot be created.
141    pub fn init_schema(&self) -> Result<(), IndexError> {
142        self.conn.execute_batch(
143            r#"
144            CREATE TABLE IF NOT EXISTS sessions (
145                id TEXT PRIMARY KEY,
146                project TEXT NOT NULL,
147                workspace_root TEXT NOT NULL DEFAULT '',
148                created_at TEXT NOT NULL,
149                updated_at TEXT NOT NULL,
150                message_count INTEGER NOT NULL DEFAULT 0
151            );
152
153            CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
154                session_id,
155                role,
156                content,
157                timestamp
158            );
159
160            CREATE TABLE IF NOT EXISTS forks (
161                source_session_id TEXT NOT NULL,
162                forked_session_id TEXT PRIMARY KEY,
163                fork_entry_id TEXT NOT NULL,
164                forked_at TEXT NOT NULL
165            );
166            "#,
167        )?;
168
169        // Migration: add workspace_root column to existing databases.
170        self.conn
171            .execute_batch(
172                "ALTER TABLE sessions ADD COLUMN workspace_root TEXT NOT NULL DEFAULT '';",
173            )
174            .ok();
175
176        Ok(())
177    }
178
179    /// Index a session and all its messages.
180    ///
181    /// Upserts the session metadata into the `sessions` table and inserts all
182    /// entries into the `messages_fts` table. Existing entries for the same session
183    /// are deleted before re-indexing to avoid duplicates.
184    ///
185    /// # Arguments
186    ///
187    /// * `session` — The session to index.
188    ///
189    /// # Errors
190    ///
191    /// Returns an error if the index cannot be updated.
192    pub fn index_session(&mut self, session: &Session) -> Result<(), IndexError> {
193        let entries = session
194            .read_entries()
195            .map_err(|e| IndexError::IoError(std::io::Error::other(e.to_string())))?;
196
197        let tx = self.conn.transaction()?;
198
199        tx.execute(
200            "DELETE FROM messages_fts WHERE session_id = ?1",
201            params![session.id.to_string()],
202        )?;
203
204        let created_at = session.created_at.to_rfc3339();
205        let updated_at = entries
206            .last()
207            .map(|e| e.timestamp.to_rfc3339())
208            .unwrap_or_else(|| Utc::now().to_rfc3339());
209        let message_count = entries.len() as i64;
210
211        tx.execute(
212            r#"
213            INSERT INTO sessions (id, project, workspace_root, created_at, updated_at, message_count)
214            VALUES (?1, ?2, ?3, ?4, ?5, ?6)
215            ON CONFLICT(id) DO UPDATE SET
216                project = excluded.project,
217                workspace_root = excluded.workspace_root,
218                updated_at = excluded.updated_at,
219                message_count = excluded.message_count
220            "#,
221            params![
222                session.id.to_string(),
223                session.project,
224                session.workspace_root,
225                created_at,
226                updated_at,
227                message_count,
228            ],
229        )?;
230
231        {
232            let mut stmt = tx.prepare(
233                "INSERT INTO messages_fts (session_id, role, content, timestamp) VALUES (?1, ?2, ?3, ?4)",
234            )?;
235
236            for entry in &entries {
237                stmt.execute(params![
238                    session.id.to_string(),
239                    entry.role,
240                    entry.content,
241                    entry.timestamp.to_rfc3339(),
242                ])?;
243            }
244        }
245
246        tx.commit()?;
247
248        Ok(())
249    }
250
251    /// Perform a full-text search across all indexed messages.
252    ///
253    /// Returns results ranked by relevance (BM25), with the most relevant first.
254    /// Each result includes a snippet with matching terms highlighted.
255    ///
256    /// # Arguments
257    ///
258    /// * `query` — The FTS5 search query. Supports FTS5 syntax (e.g., `"exact phrase"`, `term1 OR term2`).
259    /// * `limit` — Maximum number of results to return.
260    ///
261    /// # Errors
262    ///
263    /// Returns an error if the search cannot be executed.
264    pub fn search(&self, query: &str, limit: usize) -> Result<Vec<SearchResult>, IndexError> {
265        let mut stmt = self.conn.prepare(
266            r#"
267            SELECT
268                f.session_id,
269                s.project,
270                s.workspace_root,
271                snippet(messages_fts, 2, '<b>', '</b>', '...', 32) AS snippet,
272                f.timestamp,
273                bm25(messages_fts) AS rank
274            FROM messages_fts f
275            LEFT JOIN sessions s ON f.session_id = s.id
276            WHERE messages_fts MATCH ?1
277            ORDER BY rank ASC
278            LIMIT ?2
279            "#,
280        )?;
281
282        let results = stmt
283            .query_map(params![query, limit as i64], |row| {
284                let session_id: String = row.get(0)?;
285                let project: String = row.get(1).unwrap_or_else(|_| "unknown".to_string());
286                let workspace_root: String = row.get(2).unwrap_or_else(|_| String::new());
287                let snippet: String = row.get(3)?;
288                let timestamp_str: String = row.get(4)?;
289                let rank: f64 = row.get(5)?;
290
291                let timestamp = DateTime::parse_from_rfc3339(&timestamp_str)
292                    .map(|dt| dt.with_timezone(&Utc))
293                    .unwrap_or_else(|_| Utc::now());
294
295                Ok(SearchResult {
296                    session_id,
297                    project,
298                    workspace_root,
299                    snippet,
300                    timestamp,
301                    rank,
302                })
303            })?
304            .collect::<RusqliteResult<Vec<_>>>()?;
305
306        Ok(results)
307    }
308
309    /// List the most recently updated sessions.
310    ///
311    /// Returns sessions ordered by `updated_at` descending.
312    ///
313    /// # Arguments
314    ///
315    /// * `limit` — Maximum number of sessions to return.
316    ///
317    /// # Errors
318    ///
319    /// Returns an error if the query fails.
320    pub fn list_recent(&self, limit: usize) -> Result<Vec<SessionInfo>, IndexError> {
321        let mut stmt = self.conn.prepare(
322            r#"
323            SELECT id, project, workspace_root, message_count, updated_at
324            FROM sessions
325            ORDER BY updated_at DESC
326            LIMIT ?1
327            "#,
328        )?;
329
330        let results = stmt
331            .query_map(params![limit as i64], |row| {
332                let id_str: String = row.get(0)?;
333                let project: String = row.get(1)?;
334                let workspace_root: String = row.get(2).unwrap_or_default();
335                let message_count: i64 = row.get(3)?;
336                let updated_at_str: String = row.get(4)?;
337
338                let id = Uuid::parse_str(&id_str).map_err(|_| {
339                    rusqlite::Error::InvalidColumnType(
340                        0,
341                        id_str.clone(),
342                        rusqlite::types::Type::Text,
343                    )
344                })?;
345
346                let timestamp = DateTime::parse_from_rfc3339(&updated_at_str)
347                    .map(|dt| dt.with_timezone(&Utc))
348                    .unwrap_or_else(|_| Utc::now());
349
350                Ok(SessionInfo {
351                    id,
352                    project,
353                    workspace_root,
354                    last_message_preview: String::new(),
355                    timestamp,
356                    message_count: message_count as usize,
357                })
358            })?
359            .collect::<RusqliteResult<Vec<_>>>()?;
360
361        Ok(results)
362    }
363
364    /// Get metadata for a specific session.
365    ///
366    /// Returns `None` if the session is not in the index.
367    ///
368    /// # Arguments
369    ///
370    /// * `session_id` — The session ID to look up.
371    ///
372    /// # Errors
373    ///
374    /// Returns an error if the query fails.
375    pub fn get_session_info(&self, session_id: &str) -> Result<Option<SessionInfo>, IndexError> {
376        let mut stmt = self.conn.prepare(
377            r#"
378            SELECT id, project, workspace_root, message_count, updated_at
379            FROM sessions
380            WHERE id = ?1
381            "#,
382        )?;
383
384        let result = stmt
385            .query_row(params![session_id], |row| {
386                let id_str: String = row.get(0)?;
387                let project: String = row.get(1)?;
388                let workspace_root: String = row.get(2).unwrap_or_default();
389                let message_count: i64 = row.get(3)?;
390                let updated_at_str: String = row.get(4)?;
391
392                let id = Uuid::parse_str(&id_str).map_err(|_| {
393                    rusqlite::Error::InvalidColumnType(
394                        0,
395                        id_str.clone(),
396                        rusqlite::types::Type::Text,
397                    )
398                })?;
399
400                let timestamp = DateTime::parse_from_rfc3339(&updated_at_str)
401                    .map(|dt| dt.with_timezone(&Utc))
402                    .unwrap_or_else(|_| Utc::now());
403
404                Ok(SessionInfo {
405                    id,
406                    project,
407                    workspace_root,
408                    last_message_preview: String::new(),
409                    timestamp,
410                    message_count: message_count as usize,
411                })
412            })
413            .optional()?;
414
415        Ok(result)
416    }
417
418    /// Return the path to the SQLite database file.
419    pub fn db_path(&self) -> &Path {
420        &self.db_path
421    }
422
423    pub fn list_all_session_ids(&self) -> Result<Vec<String>, IndexError> {
424        let mut stmt = self.conn.prepare("SELECT id FROM sessions")?;
425        let rows = stmt.query_map([], |row| row.get::<_, String>(0))?;
426        let mut ids = Vec::new();
427        for row in rows {
428            ids.push(row?);
429        }
430        Ok(ids)
431    }
432
433    pub fn delete_session(&mut self, session_id: &str) -> Result<(), IndexError> {
434        let tx = self.conn.transaction()?;
435        tx.execute(
436            "DELETE FROM messages_fts WHERE session_id = ?1",
437            params![session_id],
438        )?;
439        tx.execute("DELETE FROM sessions WHERE id = ?1", params![session_id])?;
440        tx.execute(
441            "DELETE FROM forks WHERE source_session_id = ?1 OR forked_session_id = ?1",
442            params![session_id],
443        )?;
444        tx.commit()?;
445        Ok(())
446    }
447
448    /// Checkpoint the write-ahead log and truncate it where SQLite can do so safely.
449    pub fn checkpoint_truncate(&self) -> Result<(), IndexError> {
450        self.conn
451            .execute_batch("PRAGMA wal_checkpoint(TRUNCATE);")?;
452        Ok(())
453    }
454
455    /// Rebuild the database file to reclaim free pages after explicit cleanup.
456    pub fn vacuum(&self) -> Result<(), IndexError> {
457        self.conn.execute_batch("VACUUM;")?;
458        Ok(())
459    }
460
461    /// Record a fork relationship in the index.
462    ///
463    /// Inserts a row into the `forks` table linking the source session to the
464    /// newly forked session at the specified entry point.
465    ///
466    /// # Arguments
467    ///
468    /// * `source_session_id` — The session ID being forked from.
469    /// * `forked_session_id` — The new session ID created by the fork.
470    /// * `fork_entry_id` — The entry ID in the source session where the fork occurred.
471    ///
472    /// # Errors
473    ///
474    /// Returns an error if the fork relationship cannot be recorded.
475    pub fn record_fork(
476        &mut self,
477        source_session_id: &str,
478        forked_session_id: &str,
479        fork_entry_id: &str,
480    ) -> Result<(), IndexError> {
481        let forked_at = Utc::now().to_rfc3339();
482
483        self.conn.execute(
484            "INSERT INTO forks (source_session_id, forked_session_id, fork_entry_id, forked_at) VALUES (?1, ?2, ?3, ?4)",
485            params![source_session_id, forked_session_id, fork_entry_id, forked_at],
486        )?;
487
488        Ok(())
489    }
490
491    /// Get all forks originating from a specific session.
492    ///
493    /// Returns fork metadata for all sessions that were forked from the given
494    /// source session, ordered by fork time descending.
495    ///
496    /// # Arguments
497    ///
498    /// * `session_id` — The source session ID to query forks for.
499    ///
500    /// # Errors
501    ///
502    /// Returns an error if the query fails.
503    pub fn get_forks(&self, session_id: &str) -> Result<Vec<ForkInfo>, IndexError> {
504        let mut stmt = self.conn.prepare(
505            r#"
506            SELECT forked_session_id, fork_entry_id, forked_at
507            FROM forks
508            WHERE source_session_id = ?1
509            ORDER BY forked_at DESC
510            "#,
511        )?;
512
513        let results = stmt
514            .query_map(params![session_id], |row| {
515                let forked_session_id: String = row.get(0)?;
516                let fork_entry_id: String = row.get(1)?;
517                let forked_at_str: String = row.get(2)?;
518
519                let forked_at = DateTime::parse_from_rfc3339(&forked_at_str)
520                    .map(|dt| dt.with_timezone(&Utc))
521                    .unwrap_or_else(|_| Utc::now());
522
523                Ok(ForkInfo {
524                    forked_session_id,
525                    fork_entry_id,
526                    forked_at,
527                })
528            })?
529            .collect::<RusqliteResult<Vec<_>>>()?;
530
531        Ok(results)
532    }
533}
534
535#[cfg(test)]
536#[allow(warnings)]
537mod tests {
538    use super::*;
539    use crate::{Session, SessionManager};
540    use chrono::Datelike;
541    use talos_core::message::Message;
542
543    fn temp_index() -> (SessionIndex, tempfile::TempDir) {
544        let dir = tempfile::tempdir().unwrap();
545        let db_path = dir.path().join("test_index.db");
546        let index = SessionIndex::new(&db_path).unwrap();
547        index.init_schema().unwrap();
548        (index, dir)
549    }
550
551    fn test_session(manager: &SessionManager) -> Session {
552        let session = manager.create_session("test-project", "").unwrap();
553        session
554            .append(&Message::User {
555                content: "Hello, how do I implement full-text search in Rust?".into(),
556            })
557            .unwrap();
558        session
559            .append(&Message::Assistant {
560                content: "You can use SQLite with FTS5 extension. It provides efficient full-text indexing.".into(),
561                tool_calls: vec![],
562            })
563            .unwrap();
564        session
565            .append(&Message::User {
566                content: "What about ranking and relevance?".into(),
567            })
568            .unwrap();
569        session
570            .append(&Message::Assistant {
571                content: "FTS5 uses BM25 ranking by default. Lower scores indicate more relevant matches.".into(),
572                tool_calls: vec![],
573            })
574            .unwrap();
575        session
576    }
577
578    #[test]
579    fn test_schema_creation() {
580        let (index, _dir) = temp_index();
581
582        let table_count: i64 = index
583            .conn
584            .query_row(
585                "SELECT COUNT(*) FROM sqlite_master WHERE type IN ('table', 'view') AND name IN ('sessions', 'messages_fts')",
586                [],
587                |row| row.get(0),
588            )
589            .unwrap();
590
591        assert_eq!(
592            table_count, 2,
593            "Both sessions and messages_fts tables should exist"
594        );
595    }
596
597    #[test]
598    fn test_session_indexing() {
599        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
600        let session = test_session(&manager);
601
602        let (mut index, _dir) = temp_index();
603        index.index_session(&session).unwrap();
604
605        let info = index.get_session_info(&session.id.to_string()).unwrap();
606        assert!(info.is_some());
607        let info = info.unwrap();
608        assert_eq!(info.project, "test-project");
609        assert_eq!(info.message_count, 4);
610    }
611
612    #[test]
613    fn test_full_text_search_basic() {
614        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
615        let session = test_session(&manager);
616
617        let (mut index, _dir) = temp_index();
618        index.index_session(&session).unwrap();
619
620        let results = index.search("SQLite", 10).unwrap();
621        assert!(!results.is_empty(), "Should find results for 'SQLite'");
622
623        for result in &results {
624            assert_eq!(result.session_id, session.id.to_string());
625        }
626    }
627
628    #[test]
629    fn test_full_text_search_bm25() {
630        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
631        let session = test_session(&manager);
632
633        let (mut index, _dir) = temp_index();
634        index.index_session(&session).unwrap();
635
636        let results = index.search("ranking BM25", 10).unwrap();
637        assert!(
638            !results.is_empty(),
639            "Should find results for 'ranking BM25'"
640        );
641
642        for i in 1..results.len() {
643            assert!(
644                results[i].rank >= results[i - 1].rank,
645                "Results should be ordered by rank ascending"
646            );
647        }
648    }
649
650    #[test]
651    fn test_full_text_search_no_results() {
652        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
653        let session = test_session(&manager);
654
655        let (mut index, _dir) = temp_index();
656        index.index_session(&session).unwrap();
657
658        let results = index.search("nonexistent_term_xyz_12345", 10).unwrap();
659        assert!(
660            results.is_empty(),
661            "Should return empty results for non-matching query"
662        );
663    }
664
665    #[test]
666    fn test_full_text_search_with_limit() {
667        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
668        let session = test_session(&manager);
669
670        let (mut index, _dir) = temp_index();
671        index.index_session(&session).unwrap();
672
673        let results_all = index.search("session", 100).unwrap();
674        let results_limited = index.search("session", 2).unwrap();
675
676        assert!(results_limited.len() <= 2, "Should respect limit");
677        assert!(
678            results_limited.len() <= results_all.len(),
679            "Limited results should not exceed all results"
680        );
681    }
682
683    #[test]
684    fn test_search_result_snippet() {
685        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
686        let session = test_session(&manager);
687
688        let (mut index, _dir) = temp_index();
689        index.index_session(&session).unwrap();
690
691        let results = index.search("Rust", 10).unwrap();
692        assert!(!results.is_empty(), "Should find results for 'Rust'");
693
694        let found_highlight = results.iter().any(|r| r.snippet.contains("<b>"));
695        assert!(
696            found_highlight,
697            "Snippet should contain highlighted matches"
698        );
699    }
700
701    #[test]
702    fn test_list_recent_ordering() {
703        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
704
705        let session1 = manager.create_session("project-alpha", "").unwrap();
706        session1
707            .append(&Message::User {
708                content: "First session message".into(),
709            })
710            .unwrap();
711
712        std::thread::sleep(std::time::Duration::from_millis(10));
713
714        let session2 = manager.create_session("project-beta", "").unwrap();
715        session2
716            .append(&Message::User {
717                content: "Second session message".into(),
718            })
719            .unwrap();
720
721        let (mut index, _dir) = temp_index();
722        index.index_session(&session1).unwrap();
723        index.index_session(&session2).unwrap();
724
725        let recent = index.list_recent(10).unwrap();
726        assert_eq!(recent.len(), 2, "Should return both sessions");
727
728        assert_eq!(
729            recent[0].id, session2.id,
730            "Most recent session should be first"
731        );
732        assert_eq!(recent[1].id, session1.id, "Older session should be second");
733    }
734
735    #[test]
736    fn test_list_recent_with_limit() {
737        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
738        let db_path = tempfile::tempdir().unwrap().path().join("test_limit.db");
739        let mut index = SessionIndex::new(&db_path).unwrap();
740        index.init_schema().unwrap();
741
742        for i in 0..5 {
743            let session = manager
744                .create_session(&format!("project-limit-{i}"), "")
745                .unwrap();
746            session
747                .append(&Message::User {
748                    content: format!("Message in project {i}"),
749                })
750                .unwrap();
751            index.index_session(&session).unwrap();
752            std::thread::sleep(std::time::Duration::from_millis(5));
753        }
754
755        let limited = index.list_recent(3).unwrap();
756        assert_eq!(limited.len(), 3, "Should respect limit");
757    }
758
759    #[test]
760    fn test_get_session_info_existing() {
761        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
762        let session = test_session(&manager);
763
764        let (mut index, _dir) = temp_index();
765        index.index_session(&session).unwrap();
766
767        let info = index.get_session_info(&session.id.to_string()).unwrap();
768        assert!(info.is_some());
769        let info = info.unwrap();
770        assert_eq!(info.id, session.id);
771        assert_eq!(info.project, "test-project");
772        assert_eq!(info.message_count, 4);
773    }
774
775    #[test]
776    fn test_get_session_info_nonexistent() {
777        let (index, _dir) = temp_index();
778
779        let info = index.get_session_info("nonexistent-id").unwrap();
780        assert!(
781            info.is_none(),
782            "Should return None for non-existent session"
783        );
784    }
785
786    #[test]
787    fn test_list_all_session_ids() {
788        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
789        let s1 = test_session(&manager);
790        let s2 = manager.create_session("test-project", "").unwrap();
791        let (mut index, _dir) = temp_index();
792        index.index_session(&s1).unwrap();
793        index.index_session(&s2).unwrap();
794
795        let mut ids = index.list_all_session_ids().unwrap();
796        ids.sort();
797        let mut expected = vec![s1.id.to_string(), s2.id.to_string()];
798        expected.sort();
799        assert_eq!(ids, expected);
800    }
801
802    #[test]
803    fn test_delete_session_removes_index_entries() {
804        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
805        let session = test_session(&manager);
806        let (mut index, _dir) = temp_index();
807        index.index_session(&session).unwrap();
808
809        index
810            .delete_session(&session.id.to_string())
811            .expect("delete should succeed");
812        let info = index.get_session_info(&session.id.to_string()).unwrap();
813        assert!(info.is_none(), "index entry should be removed");
814        let ids = index.list_all_session_ids().unwrap();
815        assert!(!ids.contains(&session.id.to_string()));
816    }
817
818    #[test]
819    fn test_index_session_upsert() {
820        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
821        let session = manager.create_session("test-project", "").unwrap();
822
823        session
824            .append(&Message::User {
825                content: "Initial message".into(),
826            })
827            .unwrap();
828
829        let (mut index, _dir) = temp_index();
830        index.index_session(&session).unwrap();
831
832        session
833            .append(&Message::Assistant {
834                content: "Reply to initial message".into(),
835                tool_calls: vec![],
836            })
837            .unwrap();
838        index.index_session(&session).unwrap();
839
840        let info = index.get_session_info(&session.id.to_string()).unwrap();
841        assert!(info.is_some());
842        assert_eq!(
843            info.unwrap().message_count,
844            2,
845            "Message count should be updated"
846        );
847
848        let results = index.search("message", 100).unwrap();
849        let unique_count = results
850            .iter()
851            .filter(|r| r.session_id == session.id.to_string())
852            .count();
853        assert_eq!(unique_count, 2, "Should not have duplicate entries");
854    }
855
856    #[test]
857    fn test_search_with_phrase_query() {
858        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
859        let session = test_session(&manager);
860
861        let (mut index, _dir) = temp_index();
862        index.index_session(&session).unwrap();
863
864        let results = index.search("\"BM25 ranking\"", 10).unwrap();
865        assert!(!results.is_empty(), "Should find exact phrase match");
866    }
867
868    #[test]
869    fn test_search_result_timestamp() {
870        let manager = SessionManager::with_dir(tempfile::tempdir().unwrap().path().to_path_buf());
871        let session = test_session(&manager);
872
873        let (mut index, _dir) = temp_index();
874        index.index_session(&session).unwrap();
875
876        let results = index.search("Rust", 10).unwrap();
877        assert!(!results.is_empty());
878
879        for result in &results {
880            assert!(
881                result.timestamp.year() > 2020,
882                "Timestamp should be reasonable"
883            );
884        }
885    }
886
887    #[test]
888    fn test_db_path() {
889        let dir = tempfile::tempdir().unwrap();
890        let db_path = dir.path().join("test_path.db");
891        let index = SessionIndex::new(&db_path).unwrap();
892
893        assert_eq!(index.db_path(), db_path);
894    }
895
896    // === Fork Metadata Tests ===
897
898    #[test]
899    fn test_fork_schema_creation() {
900        let (index, _dir) = temp_index();
901
902        let table_count: i64 = index
903            .conn
904            .query_row(
905                "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='forks'",
906                [],
907                |row| row.get(0),
908            )
909            .unwrap();
910
911        assert_eq!(table_count, 1, "forks table should exist");
912    }
913
914    #[test]
915    fn test_record_fork() {
916        let (mut index, _dir) = temp_index();
917
918        index
919            .record_fork("source-session-1", "forked-session-1", "entry-abc")
920            .unwrap();
921
922        let count: i64 = index
923            .conn
924            .query_row("SELECT COUNT(*) FROM forks", [], |row| row.get(0))
925            .unwrap();
926
927        assert_eq!(count, 1, "One fork relationship should be recorded");
928    }
929
930    #[test]
931    fn test_get_forks_returns_correct_relationships() {
932        let (mut index, _dir) = temp_index();
933
934        index
935            .record_fork("source-1", "forked-1", "entry-a")
936            .unwrap();
937        index
938            .record_fork("source-1", "forked-2", "entry-b")
939            .unwrap();
940        index
941            .record_fork("source-2", "forked-3", "entry-c")
942            .unwrap();
943
944        let forks = index.get_forks("source-1").unwrap();
945        assert_eq!(forks.len(), 2, "Should return 2 forks for source-1");
946
947        let fork_ids: Vec<&str> = forks.iter().map(|f| f.forked_session_id.as_str()).collect();
948        assert!(fork_ids.contains(&"forked-1"));
949        assert!(fork_ids.contains(&"forked-2"));
950
951        let forks_source2 = index.get_forks("source-2").unwrap();
952        assert_eq!(forks_source2.len(), 1);
953        assert_eq!(forks_source2[0].forked_session_id, "forked-3");
954    }
955
956    #[test]
957    fn test_get_forks_empty_for_unknown_session() {
958        let (index, _dir) = temp_index();
959
960        let forks = index.get_forks("unknown-session").unwrap();
961        assert!(forks.is_empty(), "Should return empty for unknown session");
962    }
963
964    #[test]
965    fn test_fork_info_contains_entry_id_and_timestamp() {
966        let (mut index, _dir) = temp_index();
967
968        index
969            .record_fork("source-1", "forked-1", "entry-123")
970            .unwrap();
971
972        let forks = index.get_forks("source-1").unwrap();
973        assert_eq!(forks.len(), 1);
974
975        let fork = &forks[0];
976        assert_eq!(fork.fork_entry_id, "entry-123");
977        assert_eq!(fork.forked_session_id, "forked-1");
978        assert!(
979            fork.forked_at.year() > 2020,
980            "Timestamp should be reasonable"
981        );
982    }
983}
talos_session/sqlite.rs

talos_session/
sqlite.rs
