Skip to main content

rho_core/session/
mod.rs

1use std::path::Path;
2use std::sync::Mutex;
3
4use chrono::Utc;
5use rusqlite::{params, Connection};
6use uuid::Uuid;
7
8use crate::types::Message;
9
10pub struct SessionStore {
11    conn: Mutex<Connection>,
12}
13
14pub struct Session {
15    pub id: String,
16    pub title: String,
17    pub model: String,
18    pub created_at: String,
19    pub cwd: String,
20}
21
22#[derive(Debug, Clone)]
23pub struct SessionSummary {
24    pub id: String,
25    pub title: String,
26    pub model: String,
27    pub updated_at: String,
28    pub message_count: usize,
29}
30
31impl SessionStore {
32    pub fn open(path: &Path) -> Result<Self, rusqlite::Error> {
33        if let Some(parent) = path.parent() {
34            std::fs::create_dir_all(parent).ok();
35        }
36        let conn = Connection::open(path)?;
37        conn.execute_batch(
38            "CREATE TABLE IF NOT EXISTS sessions (
39                id TEXT PRIMARY KEY,
40                title TEXT NOT NULL,
41                model TEXT NOT NULL,
42                created_at TEXT NOT NULL,
43                updated_at TEXT NOT NULL,
44                cwd TEXT NOT NULL
45            );
46
47            CREATE TABLE IF NOT EXISTS messages (
48                id INTEGER PRIMARY KEY AUTOINCREMENT,
49                session_id TEXT NOT NULL REFERENCES sessions(id),
50                role TEXT NOT NULL,
51                content_json TEXT NOT NULL,
52                timestamp INTEGER NOT NULL,
53                seq INTEGER NOT NULL
54            );
55
56            CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id, seq);",
57        )?;
58        Ok(Self {
59            conn: Mutex::new(conn),
60        })
61    }
62
63    pub fn create_session(&self, model: &str, cwd: &Path) -> Result<Session, rusqlite::Error> {
64        let conn = self.conn.lock().unwrap();
65        let id = Uuid::new_v4().to_string();
66        let now = Utc::now().to_rfc3339();
67        let cwd_str = cwd.display().to_string();
68        let title = "New session".to_string();
69
70        conn.execute(
71            "INSERT INTO sessions (id, title, model, created_at, updated_at, cwd)
72             VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
73            params![id, title, model, now, now, cwd_str],
74        )?;
75
76        Ok(Session {
77            id,
78            title,
79            model: model.to_string(),
80            created_at: now.clone(),
81            cwd: cwd_str,
82        })
83    }
84
85    pub fn update_title(&self, id: &str, title: &str) -> Result<(), rusqlite::Error> {
86        let conn = self.conn.lock().unwrap();
87        conn.execute(
88            "UPDATE sessions SET title = ?1, updated_at = ?2 WHERE id = ?3",
89            params![title, Utc::now().to_rfc3339(), id],
90        )?;
91        Ok(())
92    }
93
94    pub fn save_messages(
95        &self,
96        session_id: &str,
97        messages: &[Message],
98    ) -> Result<(), rusqlite::Error> {
99        let conn = self.conn.lock().unwrap();
100        let tx = conn.unchecked_transaction()?;
101
102        tx.execute(
103            "DELETE FROM messages WHERE session_id = ?1",
104            params![session_id],
105        )?;
106
107        let now = Utc::now().to_rfc3339();
108        tx.execute(
109            "UPDATE sessions SET updated_at = ?1 WHERE id = ?2",
110            params![now, session_id],
111        )?;
112
113        for (seq, msg) in messages.iter().enumerate() {
114            let role = match msg {
115                Message::User { .. } => "user",
116                Message::Assistant { .. } => "assistant",
117                Message::ToolResult { .. } => "toolResult",
118            };
119            let content_json = serde_json::to_string(msg).unwrap_or_default();
120            let timestamp = match msg {
121                Message::User { timestamp, .. } => *timestamp,
122                Message::Assistant { timestamp, .. } => *timestamp,
123                Message::ToolResult { timestamp, .. } => *timestamp,
124            };
125
126            tx.execute(
127                "INSERT INTO messages (session_id, role, content_json, timestamp, seq)
128                 VALUES (?1, ?2, ?3, ?4, ?5)",
129                params![session_id, role, content_json, timestamp as i64, seq as i64],
130            )?;
131        }
132
133        tx.commit()?;
134        Ok(())
135    }
136
137    pub fn load_messages(&self, session_id: &str) -> Result<Vec<Message>, rusqlite::Error> {
138        let conn = self.conn.lock().unwrap();
139        let mut stmt =
140            conn.prepare("SELECT content_json FROM messages WHERE session_id = ?1 ORDER BY seq")?;
141        let messages = stmt
142            .query_map(params![session_id], |row| {
143                let json: String = row.get(0)?;
144                Ok(json)
145            })?
146            .filter_map(|r| r.ok())
147            .filter_map(|json| serde_json::from_str::<Message>(&json).ok())
148            .collect();
149        Ok(messages)
150    }
151
152    pub fn session_exists(&self, session_id: &str) -> Result<bool, rusqlite::Error> {
153        let conn = self.conn.lock().unwrap();
154        let mut stmt = conn.prepare("SELECT COUNT(1) FROM sessions WHERE id = ?1")?;
155        let count: i64 = stmt.query_row(params![session_id], |row| row.get(0))?;
156        Ok(count > 0)
157    }
158
159    pub fn list_sessions(&self, limit: usize) -> Result<Vec<SessionSummary>, rusqlite::Error> {
160        let conn = self.conn.lock().unwrap();
161        let mut stmt = conn.prepare(
162            "SELECT s.id, s.title, s.model, s.updated_at,
163                    (SELECT COUNT(*) FROM messages m WHERE m.session_id = s.id) as msg_count
164             FROM sessions s
165             ORDER BY s.updated_at DESC
166             LIMIT ?1",
167        )?;
168        let sessions = stmt
169            .query_map(params![limit as i64], |row| {
170                Ok(SessionSummary {
171                    id: row.get(0)?,
172                    title: row.get(1)?,
173                    model: row.get(2)?,
174                    updated_at: row.get(3)?,
175                    message_count: row.get::<_, i64>(4)? as usize,
176                })
177            })?
178            .filter_map(|r| r.ok())
179            .collect();
180        Ok(sessions)
181    }
182
183    pub fn delete_session(&self, id: &str) -> Result<(), rusqlite::Error> {
184        let conn = self.conn.lock().unwrap();
185        let tx = conn.unchecked_transaction()?;
186        tx.execute("DELETE FROM messages WHERE session_id = ?1", params![id])?;
187        tx.execute("DELETE FROM sessions WHERE id = ?1", params![id])?;
188        tx.commit()?;
189        Ok(())
190    }
191}
192
193impl crate::event_handler::SessionPersistence for SessionStore {
194    fn create_session(
195        &self,
196        model: &str,
197        cwd: &std::path::Path,
198    ) -> Result<String, Box<dyn std::error::Error>> {
199        let session = SessionStore::create_session(self, model, cwd)?;
200        Ok(session.id)
201    }
202
203    fn update_title(
204        &self,
205        session_id: &str,
206        title: &str,
207    ) -> Result<(), Box<dyn std::error::Error>> {
208        SessionStore::update_title(self, session_id, title)?;
209        Ok(())
210    }
211
212    fn save_messages(
213        &self,
214        session_id: &str,
215        messages: &[Message],
216    ) -> Result<(), Box<dyn std::error::Error>> {
217        SessionStore::save_messages(self, session_id, messages)?;
218        Ok(())
219    }
220}
221
222/// Extract a session title from conversation messages.
223/// Uses the first user message text, truncated to 80 chars.
224pub fn extract_title(messages: &[Message]) -> String {
225    for msg in messages {
226        if let Message::User {
227            content: crate::types::UserContent::Text(text),
228            ..
229        } = msg
230        {
231            let trimmed = text.trim();
232            if trimmed.is_empty() {
233                continue;
234            }
235            // Take first line, truncate to 80 chars
236            let first_line = trimmed.lines().next().unwrap_or(trimmed);
237            if first_line.len() > 80 {
238                return format!("{}...", &first_line[..77]);
239            }
240            return first_line.to_string();
241        }
242    }
243    "Untitled session".to_string()
244}
245
246#[cfg(test)]
247mod tests {
248    use super::*;
249    use crate::types::*;
250
251    fn test_store() -> SessionStore {
252        SessionStore::open(Path::new(":memory:")).unwrap()
253    }
254
255    #[test]
256    fn create_and_list_sessions() {
257        let store = test_store();
258        let s1 = store
259            .create_session("claude-sonnet", Path::new("/tmp"))
260            .unwrap();
261        let s2 = store
262            .create_session("claude-opus", Path::new("/tmp"))
263            .unwrap();
264
265        let list = store.list_sessions(50).unwrap();
266        assert_eq!(list.len(), 2);
267        // Most recent first
268        assert_eq!(list[0].id, s2.id);
269        assert_eq!(list[1].id, s1.id);
270    }
271
272    #[test]
273    fn save_and_load_messages() {
274        let store = test_store();
275        let session = store
276            .create_session("claude-sonnet", Path::new("/tmp"))
277            .unwrap();
278
279        let messages = vec![
280            Message::User {
281                content: UserContent::Text("hello".into()),
282                timestamp: 1000,
283            },
284            Message::Assistant {
285                content: vec![Content::Text {
286                    text: "hi there".into(),
287                }],
288                model: "claude-sonnet".into(),
289                usage: Usage::default(),
290                stop_reason: StopReason::Stop,
291                timestamp: 1001,
292            },
293        ];
294
295        store.save_messages(&session.id, &messages).unwrap();
296
297        let loaded = store.load_messages(&session.id).unwrap();
298        assert_eq!(loaded.len(), 2);
299
300        match &loaded[0] {
301            Message::User {
302                content: UserContent::Text(t),
303                ..
304            } => assert_eq!(t, "hello"),
305            _ => panic!("expected User message"),
306        }
307        match &loaded[1] {
308            Message::Assistant { content, .. } => {
309                assert_eq!(content.len(), 1);
310            }
311            _ => panic!("expected Assistant message"),
312        }
313    }
314
315    #[test]
316    fn session_exists_returns_true_for_known_session() {
317        let store = test_store();
318        let session = store
319            .create_session("claude-sonnet", Path::new("/tmp"))
320            .unwrap();
321        assert!(store.session_exists(&session.id).unwrap());
322    }
323
324    #[test]
325    fn session_exists_returns_false_for_unknown_session() {
326        let store = test_store();
327        assert!(!store.session_exists("does-not-exist").unwrap());
328    }
329
330    #[test]
331    fn save_messages_replaces() {
332        let store = test_store();
333        let session = store
334            .create_session("claude-sonnet", Path::new("/tmp"))
335            .unwrap();
336
337        let messages1 = vec![Message::User {
338            content: UserContent::Text("first".into()),
339            timestamp: 1000,
340        }];
341        store.save_messages(&session.id, &messages1).unwrap();
342
343        let messages2 = vec![
344            Message::User {
345                content: UserContent::Text("second".into()),
346                timestamp: 2000,
347            },
348            Message::User {
349                content: UserContent::Text("third".into()),
350                timestamp: 2001,
351            },
352        ];
353        store.save_messages(&session.id, &messages2).unwrap();
354
355        let loaded = store.load_messages(&session.id).unwrap();
356        assert_eq!(loaded.len(), 2);
357    }
358
359    #[test]
360    fn delete_session() {
361        let store = test_store();
362        let session = store
363            .create_session("claude-sonnet", Path::new("/tmp"))
364            .unwrap();
365
366        let messages = vec![Message::User {
367            content: UserContent::Text("hello".into()),
368            timestamp: 1000,
369        }];
370        store.save_messages(&session.id, &messages).unwrap();
371
372        store.delete_session(&session.id).unwrap();
373        let list = store.list_sessions(50).unwrap();
374        assert!(list.is_empty());
375
376        let loaded = store.load_messages(&session.id).unwrap();
377        assert!(loaded.is_empty());
378    }
379
380    #[test]
381    fn update_title() {
382        let store = test_store();
383        let session = store
384            .create_session("claude-sonnet", Path::new("/tmp"))
385            .unwrap();
386
387        store.update_title(&session.id, "My Session").unwrap();
388
389        let list = store.list_sessions(50).unwrap();
390        assert_eq!(list[0].title, "My Session");
391    }
392
393    #[test]
394    fn list_sessions_with_message_count() {
395        let store = test_store();
396        let session = store
397            .create_session("claude-sonnet", Path::new("/tmp"))
398            .unwrap();
399
400        let messages = vec![
401            Message::User {
402                content: UserContent::Text("hello".into()),
403                timestamp: 1000,
404            },
405            Message::Assistant {
406                content: vec![Content::Text { text: "hi".into() }],
407                model: "claude-sonnet".into(),
408                usage: Usage::default(),
409                stop_reason: StopReason::Stop,
410                timestamp: 1001,
411            },
412        ];
413        store.save_messages(&session.id, &messages).unwrap();
414
415        let list = store.list_sessions(50).unwrap();
416        assert_eq!(list[0].message_count, 2);
417    }
418
419    #[test]
420    fn extract_title_from_messages() {
421        let messages = vec![Message::User {
422            content: UserContent::Text("Help me write a function".into()),
423            timestamp: 0,
424        }];
425        assert_eq!(extract_title(&messages), "Help me write a function");
426    }
427
428    #[test]
429    fn extract_title_truncates() {
430        let long = "a".repeat(100);
431        let messages = vec![Message::User {
432            content: UserContent::Text(long),
433            timestamp: 0,
434        }];
435        let title = extract_title(&messages);
436        assert!(title.len() <= 80);
437        assert!(title.ends_with("..."));
438    }
439
440    #[test]
441    fn extract_title_empty() {
442        let messages: Vec<Message> = vec![];
443        assert_eq!(extract_title(&messages), "Untitled session");
444    }
445
446    #[test]
447    fn extract_title_skips_empty_user_messages() {
448        let messages = vec![
449            Message::User {
450                content: UserContent::Text("".into()),
451                timestamp: 0,
452            },
453            Message::User {
454                content: UserContent::Text("Real message".into()),
455                timestamp: 1,
456            },
457        ];
458        assert_eq!(extract_title(&messages), "Real message");
459    }
460}