Skip to main content

memory_core/store/
session.rs

1use rusqlite::params;
2
3use crate::error::{Error, Result};
4use crate::store::Store;
5use crate::types::Session;
6
7impl Store {
8    pub fn session_start(&self, project: &str, directory: Option<&str>) -> Result<Session> {
9        // Reuse existing active session for this project (handles --resume, --continue, Ctrl+C)
10        let existing: Option<Session> = self
11            .conn()
12            .query_row(
13                "SELECT id, project, directory, started_at, ended_at, summary, status
14                 FROM sessions WHERE project = ?1 AND directory IS ?2 AND status = 'active'
15                 ORDER BY started_at DESC LIMIT 1",
16                params![project, directory],
17                |row| {
18                    Ok(Session {
19                        id: row.get(0)?,
20                        project: row.get(1)?,
21                        directory: row.get(2)?,
22                        started_at: row.get(3)?,
23                        ended_at: row.get(4)?,
24                        summary: row.get(5)?,
25                        status: row.get(6)?,
26                    })
27                },
28            )
29            .ok();
30
31        if let Some(session) = existing {
32            return Ok(session);
33        }
34
35        // Reactivate a completed session for the same project+directory
36        let completed: Option<String> = self
37            .conn()
38            .query_row(
39                "SELECT id FROM sessions
40                 WHERE project = ?1 AND directory IS ?2 AND status = 'completed'
41                 ORDER BY started_at DESC LIMIT 1",
42                params![project, directory],
43                |row| row.get(0),
44            )
45            .ok();
46
47        if let Some(existing_id) = completed {
48            self.conn().execute(
49                "UPDATE sessions SET status = 'active', ended_at = NULL, summary = NULL
50                 WHERE id = ?1",
51                params![existing_id],
52            )?;
53            return self
54                .conn()
55                .query_row(
56                    "SELECT id, project, directory, started_at, ended_at, summary, status
57                     FROM sessions WHERE id = ?1",
58                    params![existing_id],
59                    |row| {
60                        Ok(Session {
61                            id: row.get(0)?,
62                            project: row.get(1)?,
63                            directory: row.get(2)?,
64                            started_at: row.get(3)?,
65                            ended_at: row.get(4)?,
66                            summary: row.get(5)?,
67                            status: row.get(6)?,
68                        })
69                    },
70                )
71                .map_err(Error::Database);
72        }
73
74        let id = generate_session_id();
75        self.conn().execute(
76            "INSERT INTO sessions (id, project, directory) VALUES (?1, ?2, ?3)",
77            params![id, project, directory],
78        )?;
79
80        self.conn()
81            .query_row(
82                "SELECT id, project, directory, started_at, ended_at, summary, status
83                 FROM sessions WHERE id = ?1",
84                params![id],
85                |row| {
86                    Ok(Session {
87                        id: row.get(0)?,
88                        project: row.get(1)?,
89                        directory: row.get(2)?,
90                        started_at: row.get(3)?,
91                        ended_at: row.get(4)?,
92                        summary: row.get(5)?,
93                        status: row.get(6)?,
94                    })
95                },
96            )
97            .map_err(Error::Database)
98    }
99
100    pub fn session_end(
101        &self,
102        session_id: &str,
103        summary: Option<&str>,
104        tokens_input: Option<i32>,
105        tokens_output: Option<i32>,
106    ) -> Result<Session> {
107        // Check existence first (to distinguish not-found from already-ended).
108        let exists: bool = self
109            .conn()
110            .query_row(
111                "SELECT COUNT(*) FROM sessions WHERE id = ?1",
112                params![session_id],
113                |row| row.get::<_, i64>(0),
114            )
115            .map(|n| n > 0)
116            .unwrap_or(false);
117        if !exists {
118            return Err(Error::SessionNotFound(session_id.to_string()));
119        }
120
121        // Atomic check-and-update: only update if still active.
122        let rows_changed = self.conn().execute(
123            "UPDATE sessions SET ended_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
124             summary = ?1, status = 'completed',
125             tokens_used_input = ?3, tokens_used_output = ?4
126             WHERE id = ?2 AND status = 'active'",
127            params![summary, session_id, tokens_input, tokens_output],
128        )?;
129
130        if rows_changed == 0 {
131            return Err(Error::SessionAlreadyEnded(session_id.to_string()));
132        }
133
134        self.conn()
135            .query_row(
136                "SELECT id, project, directory, started_at, ended_at, summary, status
137                 FROM sessions WHERE id = ?1",
138                params![session_id],
139                |row| {
140                    Ok(Session {
141                        id: row.get(0)?,
142                        project: row.get(1)?,
143                        directory: row.get(2)?,
144                        started_at: row.get(3)?,
145                        ended_at: row.get(4)?,
146                        summary: row.get(5)?,
147                        status: row.get(6)?,
148                    })
149                },
150            )
151            .map_err(Error::Database)
152    }
153
154    pub fn recent_sessions(&self, limit: i32) -> Result<Vec<Session>> {
155        let mut stmt = self.conn().prepare(
156            "SELECT id, project, directory, started_at, ended_at, summary, status
157             FROM sessions ORDER BY started_at DESC LIMIT ?1",
158        )?;
159
160        let results = stmt
161            .query_map(params![limit], |row| {
162                Ok(Session {
163                    id: row.get(0)?,
164                    project: row.get(1)?,
165                    directory: row.get(2)?,
166                    started_at: row.get(3)?,
167                    ended_at: row.get(4)?,
168                    summary: row.get(5)?,
169                    status: row.get(6)?,
170                })
171            })?
172            .collect::<std::result::Result<Vec<_>, _>>()?;
173
174        Ok(results)
175    }
176
177    /// End all active sessions, optionally filtered by project name.
178    /// Returns the number of sessions ended.
179    pub fn end_active_sessions(
180        &self,
181        project: Option<&str>,
182        summary: Option<&str>,
183    ) -> Result<usize> {
184        let active = self.recent_sessions(10000)?;
185        let mut ended = 0;
186        for session in &active {
187            if session.status != "active" {
188                continue;
189            }
190            if let Some(proj) = project {
191                if session.project != proj {
192                    continue;
193                }
194            }
195            self.session_end(&session.id, summary, None, None)?;
196            ended += 1;
197        }
198        Ok(ended)
199    }
200}
201
202fn generate_session_id() -> String {
203    use std::time::{SystemTime, UNIX_EPOCH};
204    let ts = SystemTime::now()
205        .duration_since(UNIX_EPOCH)
206        .unwrap()
207        .as_nanos();
208    format!("ses_{ts:x}")
209}