robit_agent/
storage.rs

1//! Session and message storage helpers.
2
3use std::path::{Path, PathBuf};
4
5use rusqlite::{params, Connection, Result as SqliteResult};
6use serde::{Deserialize, Serialize};
7
8use crate::datetime::current_timestamp;
9
10const ROBIT_DIR: &str = ".robit";
11const MEMORY_DIR: &str = "memory";
12const DB_FILE: &str = "robit.db";
13
14/// Resolve the session database path for a working directory and storage scope.
15pub fn resolve_db_path(working_dir: &Path, global_storage: bool) -> Result<PathBuf, String> {
16    if global_storage {
17        let home = dirs::home_dir().ok_or_else(|| "Cannot determine home directory".to_string())?;
18        Ok(home.join(ROBIT_DIR).join(MEMORY_DIR).join(DB_FILE))
19    } else {
20        Ok(working_dir.join(ROBIT_DIR).join(MEMORY_DIR).join(DB_FILE))
21    }
22}
23
24/// Session metadata returned to frontends.
25#[derive(Debug, Clone, Serialize)]
26pub struct SessionInfo {
27    pub id: String,
28    pub title: String,
29    pub model: String,
30    pub status: String, // "idle" | "ready" | "running"
31    pub created_at: String,
32    pub updated_at: String,
33}
34
35/// Message data returned to frontends.
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct MessageData {
38    pub id: i64,
39    pub role: String,
40    pub content: String,
41    pub tool_name: Option<String>,
42    pub tool_call_id: Option<String>,
43    pub tool_info: Option<serde_json::Value>,
44    pub created_at: String,
45}
46
47/// Tool call info for storage in message.
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct ToolCallInfoData {
50    pub tool_call_id: String,
51    pub name: String,
52    pub arguments: String,
53    pub status: String,
54    pub output: Option<String>,
55    pub requires_confirm: bool,
56}
57
58/// Initialize the database: create tables and indexes if they don't exist.
59pub fn init_db(conn: &Connection) -> SqliteResult<()> {
60    conn.execute_batch(
61        "
62        CREATE TABLE IF NOT EXISTS sessions (
63            id          TEXT PRIMARY KEY,
64            title       TEXT NOT NULL,
65            model       TEXT NOT NULL,
66            created_at  TEXT NOT NULL,
67            updated_at  TEXT NOT NULL,
68            is_active   INTEGER DEFAULT 1
69        );
70
71        CREATE TABLE IF NOT EXISTS messages (
72            id           INTEGER PRIMARY KEY AUTOINCREMENT,
73            session_id   TEXT NOT NULL REFERENCES sessions(id),
74            role         TEXT NOT NULL,
75            content      TEXT NOT NULL,
76            tool_name    TEXT,
77            tool_call_id TEXT,
78            tokens       INTEGER,
79            created_at   TEXT NOT NULL
80        );
81
82        CREATE INDEX IF NOT EXISTS idx_messages_session
83            ON messages(session_id);
84        CREATE INDEX IF NOT EXISTS idx_messages_created
85            ON messages(session_id, created_at);
86        ",
87    )?;
88
89    // Try to add tool_info column, ignore error if already exists
90    let _ = conn.execute("ALTER TABLE messages ADD COLUMN tool_info TEXT", ());
91
92    Ok(())
93}
94
95/// Insert a new session.
96pub fn insert_session(conn: &Connection, id: &str, title: &str, model: &str) -> SqliteResult<()> {
97    let now = current_timestamp();
98    conn.execute(
99        "INSERT INTO sessions (id, title, model, created_at, updated_at) VALUES (?1, ?2, ?3, ?4, ?5)",
100        params![id, title, model, now, now],
101    )?;
102    Ok(())
103}
104
105/// List all active sessions, ordered by most recently updated.
106pub fn list_sessions(conn: &Connection) -> SqliteResult<Vec<SessionInfo>> {
107    let mut stmt = conn.prepare(
108        "SELECT id, title, model, created_at, updated_at FROM sessions WHERE is_active = 1 ORDER BY updated_at DESC"
109    )?;
110    let rows = stmt.query_map([], |row| {
111        Ok(SessionInfo {
112            id: row.get(0)?,
113            title: row.get(1)?,
114            model: row.get(2)?,
115            status: "idle".to_string(),
116            created_at: row.get(3)?,
117            updated_at: row.get(4)?,
118        })
119    })?;
120    rows.collect()
121}
122
123/// Get a single session by ID.
124pub fn get_session(conn: &Connection, id: &str) -> SqliteResult<Option<SessionInfo>> {
125    let mut stmt = conn.prepare(
126        "SELECT id, title, model, created_at, updated_at FROM sessions WHERE id = ?1 AND is_active = 1"
127    )?;
128    let mut rows = stmt.query_map(params![id], |row| {
129        Ok(SessionInfo {
130            id: row.get(0)?,
131            title: row.get(1)?,
132            model: row.get(2)?,
133            status: "idle".to_string(),
134            created_at: row.get(3)?,
135            updated_at: row.get(4)?,
136        })
137    })?;
138    match rows.next() {
139        Some(Ok(session)) => Ok(Some(session)),
140        _ => Ok(None),
141    }
142}
143
144/// Update a session's title.
145pub fn update_session_title(conn: &Connection, id: &str, title: &str) -> SqliteResult<()> {
146    let now = current_timestamp();
147    conn.execute(
148        "UPDATE sessions SET title = ?1, updated_at = ?2 WHERE id = ?3",
149        params![title, now, id],
150    )?;
151    Ok(())
152}
153
154/// Update a session's updated_at timestamp.
155pub fn touch_session(conn: &Connection, id: &str) -> SqliteResult<()> {
156    let now = current_timestamp();
157    conn.execute(
158        "UPDATE sessions SET updated_at = ?1 WHERE id = ?2",
159        params![now, id],
160    )?;
161    Ok(())
162}
163
164/// Soft-delete a session.
165pub fn delete_session(conn: &Connection, id: &str) -> SqliteResult<()> {
166    conn.execute(
167        "UPDATE sessions SET is_active = 0 WHERE id = ?1",
168        params![id],
169    )?;
170    Ok(())
171}
172
173/// Insert a message into a session.
174pub fn insert_message(
175    conn: &Connection,
176    session_id: &str,
177    role: &str,
178    content: &str,
179    tool_name: Option<&str>,
180    tool_call_id: Option<&str>,
181    tool_info: Option<&str>,
182) -> SqliteResult<i64> {
183    let now = current_timestamp();
184    conn.execute(
185        "INSERT INTO messages (session_id, role, content, tool_name, tool_call_id, tool_info, created_at) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
186        params![session_id, role, content, tool_name, tool_call_id, tool_info, now],
187    )?;
188    Ok(conn.last_insert_rowid())
189}
190
191/// Get all messages for a session, ordered by creation time.
192pub fn get_messages(conn: &Connection, session_id: &str) -> SqliteResult<Vec<MessageData>> {
193    let mut stmt = conn.prepare(
194        "SELECT id, role, content, tool_name, tool_call_id, tool_info, created_at FROM messages WHERE session_id = ?1 ORDER BY id ASC"
195    )?;
196    let rows = stmt.query_map(params![session_id], |row| {
197        let tool_info_str: Option<String> = row.get(5)?;
198        let tool_info = tool_info_str.and_then(|s| serde_json::from_str(&s).ok());
199        Ok(MessageData {
200            id: row.get(0)?,
201            role: row.get(1)?,
202            content: row.get(2)?,
203            tool_name: row.get(3)?,
204            tool_call_id: row.get(4)?,
205            tool_info,
206            created_at: row.get(6)?,
207        })
208    })?;
209    rows.collect()
210}
211
212/// Update a tool message with output and status.
213pub fn update_tool_message(
214    conn: &Connection,
215    session_id: &str,
216    tool_call_id: &str,
217    tool_info: &str,
218) -> SqliteResult<()> {
219    conn.execute(
220        "UPDATE messages SET tool_info = ?1 WHERE session_id = ?2 AND tool_call_id = ?3",
221        params![tool_info, session_id, tool_call_id],
222    )?;
223    Ok(())
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    #[test]
231    fn resolves_local_db_path() {
232        let working_dir = PathBuf::from("project");
233        let path = resolve_db_path(&working_dir, false).unwrap();
234        assert_eq!(
235            path,
236            working_dir.join(ROBIT_DIR).join(MEMORY_DIR).join(DB_FILE)
237        );
238    }
239
240    #[test]
241    fn resolves_global_db_path() {
242        let path = resolve_db_path(&PathBuf::from("project"), true).unwrap();
243        assert_eq!(
244            path.file_name().and_then(|name| name.to_str()),
245            Some(DB_FILE)
246        );
247        assert!(path.ends_with(PathBuf::from(ROBIT_DIR).join(MEMORY_DIR).join(DB_FILE)));
248    }
249
250    #[test]
251    fn session_crud() {
252        let conn = Connection::open_in_memory().unwrap();
253        init_db(&conn).unwrap();
254
255        insert_session(&conn, "test-123", "Test Session", "deepseek/deepseek-chat").unwrap();
256
257        let sessions = list_sessions(&conn).unwrap();
258        assert_eq!(sessions.len(), 1);
259        assert_eq!(sessions[0].id, "test-123");
260        assert_eq!(sessions[0].title, "Test Session");
261        assert_eq!(sessions[0].status, "idle");
262
263        let session = get_session(&conn, "test-123").unwrap().unwrap();
264        assert_eq!(session.title, "Test Session");
265
266        update_session_title(&conn, "test-123", "Updated Title").unwrap();
267        let updated = get_session(&conn, "test-123").unwrap().unwrap();
268        assert_eq!(updated.title, "Updated Title");
269
270        delete_session(&conn, "test-123").unwrap();
271        assert!(get_session(&conn, "test-123").unwrap().is_none());
272        assert!(list_sessions(&conn).unwrap().is_empty());
273    }
274
275    #[test]
276    fn message_operations() {
277        let conn = Connection::open_in_memory().unwrap();
278        init_db(&conn).unwrap();
279
280        insert_session(&conn, "session-msg", "Chat Session", "model").unwrap();
281        let user_id = insert_message(
282            &conn,
283            "session-msg",
284            "user",
285            "Hello Robit",
286            None,
287            None,
288            None,
289        )
290        .unwrap();
291        let assistant_id = insert_message(
292            &conn,
293            "session-msg",
294            "assistant",
295            "Hello! How can I help?",
296            None,
297            None,
298            None,
299        )
300        .unwrap();
301
302        let messages = get_messages(&conn, "session-msg").unwrap();
303        assert_eq!(messages.len(), 2);
304        assert_eq!(messages[0].id, user_id);
305        assert_eq!(messages[0].role, "user");
306        assert_eq!(messages[0].content, "Hello Robit");
307        assert_eq!(messages[1].id, assistant_id);
308        assert_eq!(messages[1].role, "assistant");
309        assert_eq!(messages[1].content, "Hello! How can I help?");
310    }
311
312    #[test]
313    fn empty_sessions() {
314        let conn = Connection::open_in_memory().unwrap();
315        init_db(&conn).unwrap();
316
317        let sessions = list_sessions(&conn).unwrap();
318        assert_eq!(sessions.len(), 0);
319    }
320
321    #[test]
322    fn get_nonexistent_session() {
323        let conn = Connection::open_in_memory().unwrap();
324        init_db(&conn).unwrap();
325
326        let session = get_session(&conn, "nonexistent").unwrap();
327        assert!(session.is_none());
328    }
329
330    #[test]
331    fn tool_message_update() {
332        let conn = Connection::open_in_memory().unwrap();
333        init_db(&conn).unwrap();
334
335        insert_session(&conn, "session-tool", "Tool Session", "model").unwrap();
336        let initial = serde_json::json!({
337            "tool_call_id": "tool-1",
338            "name": "bash",
339            "arguments": "{}",
340            "status": "pending",
341            "requires_confirm": true
342        })
343        .to_string();
344        insert_message(
345            &conn,
346            "session-tool",
347            "tool",
348            "{}",
349            Some("bash"),
350            Some("tool-1"),
351            Some(&initial),
352        )
353        .unwrap();
354
355        let updated = serde_json::json!({
356            "tool_call_id": "tool-1",
357            "status": "success",
358            "output": "done"
359        })
360        .to_string();
361        update_tool_message(&conn, "session-tool", "tool-1", &updated).unwrap();
362
363        let messages = get_messages(&conn, "session-tool").unwrap();
364        assert_eq!(messages.len(), 1);
365        assert_eq!(messages[0].tool_name.as_deref(), Some("bash"));
366        assert_eq!(messages[0].tool_call_id.as_deref(), Some("tool-1"));
367        assert_eq!(messages[0].tool_info.as_ref().unwrap()["status"], "success");
368        assert_eq!(messages[0].tool_info.as_ref().unwrap()["output"], "done");
369    }
370}
robit_agent/storage.rs

robit_agent/
storage.rs
