imp_core/
session_index.rs

1use std::path::Path;
2
3use imp_llm::truncate_chars_with_suffix;
4use rusqlite::{params, Connection, OptionalExtension};
5
6use crate::error::Result;
7use crate::session::{SessionEntry, SessionManager};
8
9/// SQLite FTS5 index over past sessions for cross-session search.
10pub struct SessionIndex {
11    db: Connection,
12}
13
14/// A search hit from the session index.
15#[derive(Debug, Clone)]
16pub struct SessionSearchHit {
17    pub session_id: String,
18    pub cwd: String,
19    pub created_at: u64,
20    pub snippet: String,
21    pub message_count: usize,
22    pub first_message: Option<String>,
23}
24
25impl SessionIndex {
26    /// Open or create the session index database.
27    pub fn open(path: &Path) -> Result<Self> {
28        if let Some(parent) = path.parent() {
29            std::fs::create_dir_all(parent)?;
30        }
31
32        let db = Connection::open(path)?;
33        db.execute_batch(
34            "CREATE TABLE IF NOT EXISTS sessions (
35                id TEXT PRIMARY KEY,
36                cwd TEXT NOT NULL,
37                created_at INTEGER NOT NULL,
38                message_count INTEGER NOT NULL,
39                first_message TEXT
40            );
41
42            CREATE VIRTUAL TABLE IF NOT EXISTS session_content USING fts5(
43                session_id,
44                content,
45                tokenize='porter unicode61'
46            );",
47        )?;
48        Ok(Self { db })
49    }
50
51    /// Index a session's content. Extracts user messages and assistant text
52    /// (not tool results — too noisy), plus compaction summaries.
53    ///
54    /// Idempotent: re-indexing the same session updates the existing entry.
55    pub fn index_session(&self, session: &SessionManager) -> Result<()> {
56        let session_id = session
57            .path()
58            .and_then(|p| p.file_stem())
59            .map(|s| s.to_string_lossy().to_string())
60            .unwrap_or_else(|| "unknown".to_string());
61
62        let mut cwd = String::new();
63        let mut created_at: u64 = 0;
64        let mut message_count: usize = 0;
65        let mut first_message: Option<String> = None;
66        let mut content_parts: Vec<String> = Vec::new();
67
68        for entry in session.entries() {
69            match entry {
70                SessionEntry::Header {
71                    cwd: c,
72                    created_at: t,
73                    ..
74                } => {
75                    cwd = c.clone();
76                    created_at = *t;
77                }
78                SessionEntry::Message { message, .. } => {
79                    message_count += 1;
80                    let text = extract_message_text(message);
81                    if !text.is_empty() {
82                        if first_message.is_none() {
83                            first_message = Some(truncate(&text, 200));
84                        }
85                        content_parts.push(text);
86                    }
87                }
88                SessionEntry::Compaction { summary, .. } => {
89                    content_parts.push(summary.clone());
90                }
91                _ => {}
92            }
93        }
94
95        if content_parts.is_empty() {
96            return Ok(());
97        }
98
99        let content = content_parts.join("\n");
100
101        // Upsert session metadata
102        self.db.execute(
103            "INSERT INTO sessions (id, cwd, created_at, message_count, first_message)
104             VALUES (?1, ?2, ?3, ?4, ?5)
105             ON CONFLICT(id) DO UPDATE SET
106                message_count = excluded.message_count,
107                first_message = excluded.first_message",
108            params![
109                session_id,
110                cwd,
111                created_at as i64,
112                message_count as i64,
113                first_message
114            ],
115        )?;
116
117        // Delete old FTS content and re-insert
118        self.db.execute(
119            "DELETE FROM session_content WHERE session_id = ?1",
120            params![session_id],
121        )?;
122        self.db.execute(
123            "INSERT INTO session_content (session_id, content) VALUES (?1, ?2)",
124            params![session_id, content],
125        )?;
126
127        Ok(())
128    }
129
130    /// Full-text search across indexed sessions.
131    pub fn search(&self, query: &str, limit: usize) -> Result<Vec<SessionSearchHit>> {
132        let mut stmt = self.db.prepare(
133            "SELECT
134                sc.session_id,
135                s.cwd,
136                s.created_at,
137                snippet(session_content, 1, '>>>', '<<<', '...', 40) as snippet,
138                s.message_count,
139                s.first_message
140             FROM session_content sc
141             JOIN sessions s ON s.id = sc.session_id
142             WHERE session_content MATCH ?1
143             ORDER BY rank, s.created_at DESC
144             LIMIT ?2",
145        )?;
146
147        let rows = stmt.query_map(params![query, limit as i64], |row| {
148            Ok(SessionSearchHit {
149                session_id: row.get(0)?,
150                cwd: row.get(1)?,
151                created_at: row.get::<_, i64>(2)? as u64,
152                snippet: row.get(3)?,
153                message_count: row.get::<_, i64>(4)? as usize,
154                first_message: row.get(5)?,
155            })
156        })?;
157
158        let mut results = Vec::new();
159        for row in rows {
160            results.push(row?);
161        }
162        Ok(results)
163    }
164
165    /// Check if a session is already indexed.
166    pub fn is_indexed(&self, session_id: &str) -> bool {
167        self.db
168            .query_row(
169                "SELECT 1 FROM sessions WHERE id = ?1",
170                params![session_id],
171                |_| Ok(()),
172            )
173            .optional()
174            .ok()
175            .flatten()
176            .is_some()
177    }
178}
179
180/// Extract searchable text from a message. Skips tool results (too noisy).
181fn extract_message_text(message: &imp_llm::Message) -> String {
182    let blocks = match message {
183        imp_llm::Message::User(u) => &u.content,
184        imp_llm::Message::Assistant(a) => &a.content,
185        imp_llm::Message::ToolResult(_) => return String::new(),
186    };
187
188    blocks
189        .iter()
190        .filter_map(|b| match b {
191            imp_llm::ContentBlock::Text { text } => Some(text.as_str()),
192            _ => None,
193        })
194        .collect::<Vec<_>>()
195        .join(" ")
196}
197
198fn truncate(s: &str, max: usize) -> String {
199    truncate_chars_with_suffix(s, max, "...")
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205    use crate::session::SessionManager;
206    use tempfile::TempDir;
207
208    fn make_session_with_messages(dir: &std::path::Path, texts: &[&str]) -> SessionManager {
209        let session_dir = dir.join("sessions");
210        let cwd = dir.join("project");
211        let mut mgr = SessionManager::new(&cwd, &session_dir).unwrap();
212
213        for (i, text) in texts.iter().enumerate() {
214            let entry = SessionEntry::Message {
215                id: format!("m{i}"),
216                parent_id: None,
217                message: imp_llm::Message::user(*text),
218            };
219            mgr.append(entry).unwrap();
220
221            // Add an assistant response
222            let reply = SessionEntry::Message {
223                id: format!("a{i}"),
224                parent_id: None,
225                message: imp_llm::Message::Assistant(imp_llm::AssistantMessage {
226                    content: vec![imp_llm::ContentBlock::Text {
227                        text: format!("Response to: {text}"),
228                    }],
229                    usage: None,
230                    stop_reason: imp_llm::StopReason::EndTurn,
231                    timestamp: 0,
232                }),
233            };
234            mgr.append(reply).unwrap();
235        }
236
237        mgr
238    }
239
240    #[test]
241    fn session_index_create_and_search() {
242        let dir = TempDir::new().unwrap();
243        let db_path = dir.path().join("index.db");
244        let index = SessionIndex::open(&db_path).unwrap();
245
246        let session = make_session_with_messages(
247            dir.path(),
248            &["Help me deploy to kubernetes", "Show me the docker config"],
249        );
250        index.index_session(&session).unwrap();
251
252        let results = index.search("kubernetes", 10).unwrap();
253        assert_eq!(results.len(), 1);
254        assert!(results[0].snippet.contains("kubernetes"));
255    }
256
257    #[test]
258    fn session_index_no_results() {
259        let dir = TempDir::new().unwrap();
260        let db_path = dir.path().join("index.db");
261        let index = SessionIndex::open(&db_path).unwrap();
262
263        let session = make_session_with_messages(dir.path(), &["Hello world"]);
264        index.index_session(&session).unwrap();
265
266        let results = index.search("kubernetes", 10).unwrap();
267        assert!(results.is_empty());
268    }
269
270    #[test]
271    fn session_index_multiple_sessions() {
272        let dir = TempDir::new().unwrap();
273        let db_path = dir.path().join("index.db");
274        let index = SessionIndex::open(&db_path).unwrap();
275
276        let s1 = make_session_with_messages(dir.path(), &["Deploy to kubernetes cluster"]);
277        index.index_session(&s1).unwrap();
278
279        // Create second session in a different subdir to get a different session file
280        let dir2 = dir.path().join("other");
281        std::fs::create_dir_all(&dir2).unwrap();
282        let s2 = make_session_with_messages(&dir2, &["Fix the kubernetes ingress"]);
283        index.index_session(&s2).unwrap();
284
285        let results = index.search("kubernetes", 10).unwrap();
286        assert_eq!(results.len(), 2);
287    }
288
289    #[test]
290    fn session_index_idempotent() {
291        let dir = TempDir::new().unwrap();
292        let db_path = dir.path().join("index.db");
293        let index = SessionIndex::open(&db_path).unwrap();
294
295        let session = make_session_with_messages(dir.path(), &["test content"]);
296        index.index_session(&session).unwrap();
297        index.index_session(&session).unwrap(); // Re-index
298
299        let results = index.search("test", 10).unwrap();
300        assert_eq!(results.len(), 1, "should not duplicate on re-index");
301    }
302
303    #[test]
304    fn session_index_is_indexed() {
305        let dir = TempDir::new().unwrap();
306        let db_path = dir.path().join("index.db");
307        let index = SessionIndex::open(&db_path).unwrap();
308
309        assert!(!index.is_indexed("nonexistent"));
310
311        let session = make_session_with_messages(dir.path(), &["hello"]);
312        index.index_session(&session).unwrap();
313
314        let session_id = session
315            .path()
316            .unwrap()
317            .file_stem()
318            .unwrap()
319            .to_string_lossy()
320            .to_string();
321        assert!(index.is_indexed(&session_id));
322    }
323
324    #[test]
325    fn session_index_fts5_and_or_not() {
326        let dir = TempDir::new().unwrap();
327        let db_path = dir.path().join("index.db");
328        let index = SessionIndex::open(&db_path).unwrap();
329
330        let session = make_session_with_messages(
331            dir.path(),
332            &["Deploy kubernetes cluster", "Configure docker networking"],
333        );
334        index.index_session(&session).unwrap();
335
336        // AND query
337        let results = index.search("kubernetes AND cluster", 10).unwrap();
338        assert_eq!(results.len(), 1);
339
340        // OR query
341        let results = index.search("kubernetes OR docker", 10).unwrap();
342        assert_eq!(results.len(), 1); // same session has both
343
344        // NOT query
345        let results = index.search("kubernetes NOT docker", 10).unwrap();
346        // FTS5 NOT: matches docs containing kubernetes but not docker
347        // Since both are in the same session content, this should return 0
348        assert_eq!(results.len(), 0);
349    }
350
351    #[test]
352    fn session_index_snippet_highlights() {
353        let dir = TempDir::new().unwrap();
354        let db_path = dir.path().join("index.db");
355        let index = SessionIndex::open(&db_path).unwrap();
356
357        let session =
358            make_session_with_messages(dir.path(), &["The kubernetes deployment is broken"]);
359        index.index_session(&session).unwrap();
360
361        let results = index.search("kubernetes", 10).unwrap();
362        assert_eq!(results.len(), 1);
363        // Snippet should contain >>> and <<< markers
364        assert!(
365            results[0].snippet.contains(">>>") && results[0].snippet.contains("<<<"),
366            "snippet should have highlight markers: {}",
367            results[0].snippet
368        );
369    }
370}
imp_core/session_index.rs

imp_core/
session_index.rs
