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().expect("operation should succeed");
545        let db_path = dir.path().join("test_index.db");
546        let index = SessionIndex::new(&db_path).expect("operation should succeed");
547        index.init_schema().expect("operation should succeed");
548        (index, dir)
549    }
550
551    fn test_session(manager: &SessionManager) -> Session {
552        let session = manager
553            .create_session("test-project", "")
554            .expect("operation should succeed");
555        session
556            .append(&Message::User {
557                content: "Hello, how do I implement full-text search in Rust?".into(),
558            })
559            .expect("operation should succeed");
560        session
561            .append(&Message::Assistant {
562                content: "You can use SQLite with FTS5 extension. It provides efficient full-text indexing.".into(),
563                tool_calls: vec![],
564                reasoning: None,
565            })
566            .expect("operation should succeed");
567        session
568            .append(&Message::User {
569                content: "What about ranking and relevance?".into(),
570            })
571            .expect("operation should succeed");
572        session
573            .append(&Message::Assistant {
574                content: "FTS5 uses BM25 ranking by default. Lower scores indicate more relevant matches.".into(),
575                tool_calls: vec![],
576                reasoning: None,
577            })
578            .expect("operation should succeed");
579        session
580    }
581
582    #[test]
583    fn test_schema_creation() {
584        let (index, _dir) = temp_index();
585
586        let table_count: i64 = index
587            .conn
588            .query_row(
589                "SELECT COUNT(*) FROM sqlite_master WHERE type IN ('table', 'view') AND name IN ('sessions', 'messages_fts')",
590                [],
591                |row| row.get(0),
592            )
593            .expect("operation should succeed");
594
595        assert_eq!(
596            table_count, 2,
597            "Both sessions and messages_fts tables should exist"
598        );
599    }
600
601    #[test]
602    fn test_session_indexing() {
603        let manager = SessionManager::with_dir(
604            tempfile::tempdir()
605                .expect("operation should succeed")
606                .path()
607                .to_path_buf(),
608        );
609        let session = test_session(&manager);
610
611        let (mut index, _dir) = temp_index();
612        index
613            .index_session(&session)
614            .expect("operation should succeed");
615
616        let info = index
617            .get_session_info(&session.id.to_string())
618            .expect("operation should succeed");
619        assert!(info.is_some());
620        let info = info.expect("operation should succeed");
621        assert_eq!(info.project, "test-project");
622        assert_eq!(info.message_count, 4);
623    }
624
625    #[test]
626    fn test_full_text_search_basic() {
627        let manager = SessionManager::with_dir(
628            tempfile::tempdir()
629                .expect("operation should succeed")
630                .path()
631                .to_path_buf(),
632        );
633        let session = test_session(&manager);
634
635        let (mut index, _dir) = temp_index();
636        index
637            .index_session(&session)
638            .expect("operation should succeed");
639
640        let results = index
641            .search("SQLite", 10)
642            .expect("operation should succeed");
643        assert!(!results.is_empty(), "Should find results for 'SQLite'");
644
645        for result in &results {
646            assert_eq!(result.session_id, session.id.to_string());
647        }
648    }
649
650    #[test]
651    fn test_full_text_search_bm25() {
652        let manager = SessionManager::with_dir(
653            tempfile::tempdir()
654                .expect("operation should succeed")
655                .path()
656                .to_path_buf(),
657        );
658        let session = test_session(&manager);
659
660        let (mut index, _dir) = temp_index();
661        index
662            .index_session(&session)
663            .expect("operation should succeed");
664
665        let results = index
666            .search("ranking BM25", 10)
667            .expect("operation should succeed");
668        assert!(
669            !results.is_empty(),
670            "Should find results for 'ranking BM25'"
671        );
672
673        for i in 1..results.len() {
674            assert!(
675                results[i].rank >= results[i - 1].rank,
676                "Results should be ordered by rank ascending"
677            );
678        }
679    }
680
681    #[test]
682    fn test_full_text_search_no_results() {
683        let manager = SessionManager::with_dir(
684            tempfile::tempdir()
685                .expect("operation should succeed")
686                .path()
687                .to_path_buf(),
688        );
689        let session = test_session(&manager);
690
691        let (mut index, _dir) = temp_index();
692        index
693            .index_session(&session)
694            .expect("operation should succeed");
695
696        let results = index
697            .search("nonexistent_term_xyz_12345", 10)
698            .expect("operation should succeed");
699        assert!(
700            results.is_empty(),
701            "Should return empty results for non-matching query"
702        );
703    }
704
705    #[test]
706    fn test_full_text_search_with_limit() {
707        let manager = SessionManager::with_dir(
708            tempfile::tempdir()
709                .expect("operation should succeed")
710                .path()
711                .to_path_buf(),
712        );
713        let session = test_session(&manager);
714
715        let (mut index, _dir) = temp_index();
716        index
717            .index_session(&session)
718            .expect("operation should succeed");
719
720        let results_all = index
721            .search("session", 100)
722            .expect("operation should succeed");
723        let results_limited = index
724            .search("session", 2)
725            .expect("operation should succeed");
726
727        assert!(results_limited.len() <= 2, "Should respect limit");
728        assert!(
729            results_limited.len() <= results_all.len(),
730            "Limited results should not exceed all results"
731        );
732    }
733
734    #[test]
735    fn test_search_result_snippet() {
736        let manager = SessionManager::with_dir(
737            tempfile::tempdir()
738                .expect("operation should succeed")
739                .path()
740                .to_path_buf(),
741        );
742        let session = test_session(&manager);
743
744        let (mut index, _dir) = temp_index();
745        index
746            .index_session(&session)
747            .expect("operation should succeed");
748
749        let results = index.search("Rust", 10).expect("operation should succeed");
750        assert!(!results.is_empty(), "Should find results for 'Rust'");
751
752        let found_highlight = results.iter().any(|r| r.snippet.contains("<b>"));
753        assert!(
754            found_highlight,
755            "Snippet should contain highlighted matches"
756        );
757    }
758
759    #[test]
760    fn test_list_recent_ordering() {
761        let manager = SessionManager::with_dir(
762            tempfile::tempdir()
763                .expect("operation should succeed")
764                .path()
765                .to_path_buf(),
766        );
767
768        let session1 = manager
769            .create_session("project-alpha", "")
770            .expect("operation should succeed");
771        session1
772            .append(&Message::User {
773                content: "First session message".into(),
774            })
775            .expect("operation should succeed");
776
777        std::thread::sleep(std::time::Duration::from_millis(10));
778
779        let session2 = manager
780            .create_session("project-beta", "")
781            .expect("operation should succeed");
782        session2
783            .append(&Message::User {
784                content: "Second session message".into(),
785            })
786            .expect("operation should succeed");
787
788        let (mut index, _dir) = temp_index();
789        index
790            .index_session(&session1)
791            .expect("operation should succeed");
792        index
793            .index_session(&session2)
794            .expect("operation should succeed");
795
796        let recent = index.list_recent(10).expect("operation should succeed");
797        assert_eq!(recent.len(), 2, "Should return both sessions");
798
799        assert_eq!(
800            recent[0].id, session2.id,
801            "Most recent session should be first"
802        );
803        assert_eq!(recent[1].id, session1.id, "Older session should be second");
804    }
805
806    #[test]
807    fn test_list_recent_with_limit() {
808        let manager = SessionManager::with_dir(
809            tempfile::tempdir()
810                .expect("operation should succeed")
811                .path()
812                .to_path_buf(),
813        );
814        let db_path = tempfile::tempdir()
815            .expect("operation should succeed")
816            .path()
817            .join("test_limit.db");
818        let mut index = SessionIndex::new(&db_path).expect("operation should succeed");
819        index.init_schema().expect("operation should succeed");
820
821        for i in 0..5 {
822            let session = manager
823                .create_session(&format!("project-limit-{i}"), "")
824                .expect("operation should succeed");
825            session
826                .append(&Message::User {
827                    content: format!("Message in project {i}"),
828                })
829                .expect("operation should succeed");
830            index
831                .index_session(&session)
832                .expect("operation should succeed");
833            std::thread::sleep(std::time::Duration::from_millis(5));
834        }
835
836        let limited = index.list_recent(3).expect("operation should succeed");
837        assert_eq!(limited.len(), 3, "Should respect limit");
838    }
839
840    #[test]
841    fn test_get_session_info_existing() {
842        let manager = SessionManager::with_dir(
843            tempfile::tempdir()
844                .expect("operation should succeed")
845                .path()
846                .to_path_buf(),
847        );
848        let session = test_session(&manager);
849
850        let (mut index, _dir) = temp_index();
851        index
852            .index_session(&session)
853            .expect("operation should succeed");
854
855        let info = index
856            .get_session_info(&session.id.to_string())
857            .expect("operation should succeed");
858        assert!(info.is_some());
859        let info = info.expect("operation should succeed");
860        assert_eq!(info.id, session.id);
861        assert_eq!(info.project, "test-project");
862        assert_eq!(info.message_count, 4);
863    }
864
865    #[test]
866    fn test_get_session_info_nonexistent() {
867        let (index, _dir) = temp_index();
868
869        let info = index
870            .get_session_info("nonexistent-id")
871            .expect("operation should succeed");
872        assert!(
873            info.is_none(),
874            "Should return None for non-existent session"
875        );
876    }
877
878    #[test]
879    fn test_list_all_session_ids() {
880        let manager = SessionManager::with_dir(
881            tempfile::tempdir()
882                .expect("operation should succeed")
883                .path()
884                .to_path_buf(),
885        );
886        let s1 = test_session(&manager);
887        let s2 = manager
888            .create_session("test-project", "")
889            .expect("operation should succeed");
890        let (mut index, _dir) = temp_index();
891        index.index_session(&s1).expect("operation should succeed");
892        index.index_session(&s2).expect("operation should succeed");
893
894        let mut ids = index
895            .list_all_session_ids()
896            .expect("operation should succeed");
897        ids.sort();
898        let mut expected = vec![s1.id.to_string(), s2.id.to_string()];
899        expected.sort();
900        assert_eq!(ids, expected);
901    }
902
903    #[test]
904    fn test_delete_session_removes_index_entries() {
905        let manager = SessionManager::with_dir(
906            tempfile::tempdir()
907                .expect("operation should succeed")
908                .path()
909                .to_path_buf(),
910        );
911        let session = test_session(&manager);
912        let (mut index, _dir) = temp_index();
913        index
914            .index_session(&session)
915            .expect("operation should succeed");
916
917        index
918            .delete_session(&session.id.to_string())
919            .expect("delete should succeed");
920        let info = index
921            .get_session_info(&session.id.to_string())
922            .expect("operation should succeed");
923        assert!(info.is_none(), "index entry should be removed");
924        let ids = index
925            .list_all_session_ids()
926            .expect("operation should succeed");
927        assert!(!ids.contains(&session.id.to_string()));
928    }
929
930    #[test]
931    fn test_index_session_upsert() {
932        let manager = SessionManager::with_dir(
933            tempfile::tempdir()
934                .expect("operation should succeed")
935                .path()
936                .to_path_buf(),
937        );
938        let session = manager
939            .create_session("test-project", "")
940            .expect("operation should succeed");
941
942        session
943            .append(&Message::User {
944                content: "Initial message".into(),
945            })
946            .expect("operation should succeed");
947
948        let (mut index, _dir) = temp_index();
949        index
950            .index_session(&session)
951            .expect("operation should succeed");
952
953        session
954            .append(&Message::Assistant {
955                content: "Reply to initial message".into(),
956                tool_calls: vec![],
957                reasoning: None,
958            })
959            .expect("operation should succeed");
960        index
961            .index_session(&session)
962            .expect("operation should succeed");
963
964        let info = index
965            .get_session_info(&session.id.to_string())
966            .expect("operation should succeed");
967        assert!(info.is_some());
968        assert_eq!(
969            info.expect("operation should succeed").message_count,
970            2,
971            "Message count should be updated"
972        );
973
974        let results = index
975            .search("message", 100)
976            .expect("operation should succeed");
977        let unique_count = results
978            .iter()
979            .filter(|r| r.session_id == session.id.to_string())
980            .count();
981        assert_eq!(unique_count, 2, "Should not have duplicate entries");
982    }
983
984    #[test]
985    fn test_search_with_phrase_query() {
986        let manager = SessionManager::with_dir(
987            tempfile::tempdir()
988                .expect("operation should succeed")
989                .path()
990                .to_path_buf(),
991        );
992        let session = test_session(&manager);
993
994        let (mut index, _dir) = temp_index();
995        index
996            .index_session(&session)
997            .expect("operation should succeed");
998
999        let results = index
1000            .search("\"BM25 ranking\"", 10)
1001            .expect("operation should succeed");
1002        assert!(!results.is_empty(), "Should find exact phrase match");
1003    }
1004
1005    #[test]
1006    fn test_search_result_timestamp() {
1007        let manager = SessionManager::with_dir(
1008            tempfile::tempdir()
1009                .expect("operation should succeed")
1010                .path()
1011                .to_path_buf(),
1012        );
1013        let session = test_session(&manager);
1014
1015        let (mut index, _dir) = temp_index();
1016        index
1017            .index_session(&session)
1018            .expect("operation should succeed");
1019
1020        let results = index.search("Rust", 10).expect("operation should succeed");
1021        assert!(!results.is_empty());
1022
1023        for result in &results {
1024            assert!(
1025                result.timestamp.year() > 2020,
1026                "Timestamp should be reasonable"
1027            );
1028        }
1029    }
1030
1031    #[test]
1032    fn test_db_path() {
1033        let dir = tempfile::tempdir().expect("operation should succeed");
1034        let db_path = dir.path().join("test_path.db");
1035        let index = SessionIndex::new(&db_path).expect("operation should succeed");
1036
1037        assert_eq!(index.db_path(), db_path);
1038    }
1039
1040    // === Fork Metadata Tests ===
1041
1042    #[test]
1043    fn test_fork_schema_creation() {
1044        let (index, _dir) = temp_index();
1045
1046        let table_count: i64 = index
1047            .conn
1048            .query_row(
1049                "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name='forks'",
1050                [],
1051                |row| row.get(0),
1052            )
1053            .expect("operation should succeed");
1054
1055        assert_eq!(table_count, 1, "forks table should exist");
1056    }
1057
1058    #[test]
1059    fn test_record_fork() {
1060        let (mut index, _dir) = temp_index();
1061
1062        index
1063            .record_fork("source-session-1", "forked-session-1", "entry-abc")
1064            .expect("operation should succeed");
1065
1066        let count: i64 = index
1067            .conn
1068            .query_row("SELECT COUNT(*) FROM forks", [], |row| row.get(0))
1069            .expect("operation should succeed");
1070
1071        assert_eq!(count, 1, "One fork relationship should be recorded");
1072    }
1073
1074    #[test]
1075    fn test_get_forks_returns_correct_relationships() {
1076        let (mut index, _dir) = temp_index();
1077
1078        index
1079            .record_fork("source-1", "forked-1", "entry-a")
1080            .expect("operation should succeed");
1081        index
1082            .record_fork("source-1", "forked-2", "entry-b")
1083            .expect("operation should succeed");
1084        index
1085            .record_fork("source-2", "forked-3", "entry-c")
1086            .expect("operation should succeed");
1087
1088        let forks = index
1089            .get_forks("source-1")
1090            .expect("operation should succeed");
1091        assert_eq!(forks.len(), 2, "Should return 2 forks for source-1");
1092
1093        let fork_ids: Vec<&str> = forks.iter().map(|f| f.forked_session_id.as_str()).collect();
1094        assert!(fork_ids.contains(&"forked-1"));
1095        assert!(fork_ids.contains(&"forked-2"));
1096
1097        let forks_source2 = index
1098            .get_forks("source-2")
1099            .expect("operation should succeed");
1100        assert_eq!(forks_source2.len(), 1);
1101        assert_eq!(forks_source2[0].forked_session_id, "forked-3");
1102    }
1103
1104    #[test]
1105    fn test_get_forks_empty_for_unknown_session() {
1106        let (index, _dir) = temp_index();
1107
1108        let forks = index
1109            .get_forks("unknown-session")
1110            .expect("operation should succeed");
1111        assert!(forks.is_empty(), "Should return empty for unknown session");
1112    }
1113
1114    #[test]
1115    fn test_fork_info_contains_entry_id_and_timestamp() {
1116        let (mut index, _dir) = temp_index();
1117
1118        index
1119            .record_fork("source-1", "forked-1", "entry-123")
1120            .expect("operation should succeed");
1121
1122        let forks = index
1123            .get_forks("source-1")
1124            .expect("operation should succeed");
1125        assert_eq!(forks.len(), 1);
1126
1127        let fork = &forks[0];
1128        assert_eq!(fork.fork_entry_id, "entry-123");
1129        assert_eq!(fork.forked_session_id, "forked-1");
1130        assert!(
1131            fork.forked_at.year() > 2020,
1132            "Timestamp should be reasonable"
1133        );
1134    }
1135}
talos_session/sqlite.rs

talos_session/
sqlite.rs
