1use std::fs;
2use std::path::{Path, PathBuf};
3use std::sync::Mutex;
4
5use chrono::{DateTime, Utc};
6use codei_config::expand_tilde;
7use rusqlite::{params, Connection};
8
9use crate::error::SessionError;
10use crate::model::{MessageContent, Role, Session, StoredMessage};
11
12pub struct SessionStore {
13 conn: Mutex<Connection>,
14}
15
16impl SessionStore {
17 fn conn(&self) -> Result<std::sync::MutexGuard<'_, Connection>, SessionError> {
18 self.conn.lock().map_err(|_| SessionError::LockPoisoned)
19 }
20 pub fn open_default() -> Result<Self, SessionError> {
21 let path = default_db_path();
22 if let Some(parent) = path.parent() {
23 fs::create_dir_all(parent)?;
24 }
25 let conn = Connection::open(path)?;
26 let store = Self {
27 conn: Mutex::new(conn),
28 };
29 store.init_schema()?;
30 Ok(store)
31 }
32
33 pub fn open(path: &Path) -> Result<Self, SessionError> {
34 if let Some(parent) = path.parent() {
35 fs::create_dir_all(parent)?;
36 }
37 let conn = Connection::open(path)?;
38 let store = Self {
39 conn: Mutex::new(conn),
40 };
41 store.init_schema()?;
42 Ok(store)
43 }
44
45 fn init_schema(&self) -> Result<(), SessionError> {
46 let conn = self.conn()?;
47 conn.execute_batch(
48 r#"
49 CREATE TABLE IF NOT EXISTS sessions (
50 id TEXT PRIMARY KEY,
51 title TEXT,
52 cwd TEXT NOT NULL,
53 created_at TEXT NOT NULL,
54 updated_at TEXT NOT NULL
55 );
56 CREATE TABLE IF NOT EXISTS messages (
57 id TEXT PRIMARY KEY,
58 session_id TEXT NOT NULL,
59 role TEXT NOT NULL,
60 content TEXT NOT NULL,
61 tool_calls TEXT,
62 tool_call_id TEXT,
63 created_at TEXT NOT NULL,
64 FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE
65 );
66 CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id);
67 "#,
68 )?;
69 Ok(())
70 }
71
72 pub fn save(&self, session: &Session) -> Result<(), SessionError> {
73 let conn = self.conn()?;
74 let tx = conn.unchecked_transaction()?;
75 tx.execute(
76 "INSERT INTO sessions (id, title, cwd, created_at, updated_at)
77 VALUES (?1, ?2, ?3, ?4, ?5)
78 ON CONFLICT(id) DO UPDATE SET
79 title = excluded.title,
80 cwd = excluded.cwd,
81 updated_at = excluded.updated_at",
82 params![
83 session.id,
84 session.title,
85 session.cwd.to_string_lossy().to_string(),
86 session.created_at.to_rfc3339(),
87 session.updated_at.to_rfc3339(),
88 ],
89 )?;
90 tx.execute(
91 "DELETE FROM messages WHERE session_id = ?1",
92 params![session.id],
93 )?;
94 for msg in &session.messages {
95 let tool_calls = msg
96 .tool_calls
97 .as_ref()
98 .map(serde_json::to_string)
99 .transpose()?;
100 let content = match &msg.content {
101 MessageContent::Text(s) => s.clone(),
102 };
103 tx.execute(
104 "INSERT INTO messages (id, session_id, role, content, tool_calls, tool_call_id, created_at)
105 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
106 params![
107 msg.id,
108 session.id,
109 role_to_str(msg.role),
110 content,
111 tool_calls,
112 msg.tool_call_id,
113 msg.created_at.to_rfc3339(),
114 ],
115 )?;
116 }
117 tx.commit()?;
118 Ok(())
119 }
120
121 pub fn load(&self, id: &str) -> Result<Session, SessionError> {
122 let conn = self.conn()?;
123 Self::load_with_conn(&conn, id)
124 }
125
126 fn load_with_conn(conn: &Connection, id: &str) -> Result<Session, SessionError> {
127 let mut stmt = conn
128 .prepare("SELECT id, title, cwd, created_at, updated_at FROM sessions WHERE id = ?1")?;
129 let session_row = stmt
130 .query_row(params![id], |row| {
131 Ok((
132 row.get::<_, String>(0)?,
133 row.get::<_, Option<String>>(1)?,
134 row.get::<_, String>(2)?,
135 row.get::<_, String>(3)?,
136 row.get::<_, String>(4)?,
137 ))
138 })
139 .map_err(|_| SessionError::NotFound(id.to_string()))?;
140
141 let messages = Self::load_messages(conn, id)?;
142
143 Ok(Session {
144 id: session_row.0,
145 title: session_row.1,
146 cwd: PathBuf::from(session_row.2),
147 created_at: parse_ts(&session_row.3)?,
148 updated_at: parse_ts(&session_row.4)?,
149 messages,
150 })
151 }
152
153 pub fn list(&self, limit: usize) -> Result<Vec<Session>, SessionError> {
154 let conn = self.conn()?;
155 let mut stmt = conn.prepare("SELECT id FROM sessions ORDER BY updated_at DESC LIMIT ?1")?;
156 let ids: Vec<String> = stmt
157 .query_map(params![limit as i64], |row| row.get(0))?
158 .collect::<Result<Vec<_>, _>>()?;
159 ids.iter()
160 .map(|id| Self::load_with_conn(&conn, id))
161 .collect()
162 }
163
164 pub fn latest(&self) -> Result<Option<Session>, SessionError> {
165 let conn = self.conn()?;
166 let mut stmt = conn.prepare("SELECT id FROM sessions ORDER BY updated_at DESC LIMIT 1")?;
167 let mut rows = stmt.query([])?;
168 if let Some(row) = rows.next()? {
169 let id: String = row.get(0)?;
170 Ok(Some(Self::load_with_conn(&conn, &id)?))
171 } else {
172 Ok(None)
173 }
174 }
175
176 pub fn delete(&self, id: &str) -> Result<(), SessionError> {
177 let conn = self.conn()?;
178 let changed = conn.execute("DELETE FROM sessions WHERE id = ?1", params![id])?;
179 if changed == 0 {
180 return Err(SessionError::NotFound(id.to_string()));
181 }
182 Ok(())
183 }
184
185 pub fn export_jsonl(&self, id: &str) -> Result<String, SessionError> {
186 let session = self.load(id)?;
187 let mut lines = Vec::new();
188 for msg in &session.messages {
189 let line = serde_json::json!({
190 "id": msg.id,
191 "role": msg.role,
192 "content": msg.content,
193 "tool_calls": msg.tool_calls,
194 "tool_call_id": msg.tool_call_id,
195 "created_at": msg.created_at,
196 });
197 lines.push(serde_json::to_string(&line)?);
198 }
199 Ok(lines.join("\n"))
200 }
201
202 fn load_messages(
203 conn: &Connection,
204 session_id: &str,
205 ) -> Result<Vec<StoredMessage>, SessionError> {
206 let mut stmt = conn.prepare(
207 "SELECT id, role, content, tool_calls, tool_call_id, created_at
208 FROM messages WHERE session_id = ?1 ORDER BY created_at ASC",
209 )?;
210 let rows = stmt.query_map(params![session_id], |row| {
211 let role: String = row.get(1)?;
212 let tool_calls: Option<String> = row.get(3)?;
213 let parsed_tool_calls = tool_calls
214 .as_ref()
215 .map(|s| serde_json::from_str(s))
216 .transpose()
217 .map_err(|e| rusqlite::Error::InvalidParameterName(e.to_string()))?;
218 Ok(StoredMessage {
219 id: row.get(0)?,
220 role: str_to_role(&role),
221 content: MessageContent::Text(row.get(2)?),
222 tool_calls: parsed_tool_calls,
223 tool_call_id: row.get(4)?,
224 created_at: parse_ts(&row.get::<_, String>(5)?).unwrap_or_else(|_| Utc::now()),
225 })
226 })?;
227 rows.collect::<Result<Vec<_>, _>>()
228 .map_err(SessionError::from)
229 }
230}
231
232fn default_db_path() -> PathBuf {
233 expand_tilde("~/.local/share/codei/sessions.db")
234}
235
236fn role_to_str(role: Role) -> &'static str {
237 match role {
238 Role::System => "system",
239 Role::User => "user",
240 Role::Assistant => "assistant",
241 Role::Tool => "tool",
242 }
243}
244
245fn str_to_role(role: &str) -> Role {
246 match role {
247 "system" => Role::System,
248 "assistant" => Role::Assistant,
249 "tool" => Role::Tool,
250 _ => Role::User,
251 }
252}
253
254fn parse_ts(value: &str) -> Result<DateTime<Utc>, SessionError> {
255 DateTime::parse_from_rfc3339(value)
256 .map(|dt| dt.with_timezone(&Utc))
257 .map_err(|e| SessionError::Database(rusqlite::Error::InvalidParameterName(e.to_string())))
258}
259
260#[cfg(test)]
261mod tests {
262 use super::*;
263 use crate::model::Session;
264
265 #[test]
266 fn roundtrip_session() {
267 let dir = tempfile::tempdir().unwrap();
268 let db = dir.path().join("test.db");
269 let store = SessionStore::open(&db).unwrap();
270
271 let mut session = Session::new(PathBuf::from("/tmp/project"));
272 session.push_user("hello");
273 session.push_assistant("hi".into(), None);
274
275 store.save(&session).unwrap();
276 let loaded = store.load(&session.id).unwrap();
277 assert_eq!(loaded.messages.len(), 2);
278 assert_eq!(loaded.messages[0].text(), Some("hello"));
279 }
280
281 #[test]
282 fn list_does_not_deadlock() {
283 let dir = tempfile::tempdir().unwrap();
284 let store = SessionStore::open(&dir.path().join("test.db")).unwrap();
285
286 let mut session = Session::new(PathBuf::from("/tmp/project"));
287 session.push_user("hello");
288 store.save(&session).unwrap();
289
290 let sessions = store.list(10).unwrap();
291 assert_eq!(sessions.len(), 1);
292 assert_eq!(sessions[0].id, session.id);
293 }
294}