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
222pub 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 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 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}