memory_core/store/

session.rs

use rusqlite::params;

use crate::error::{Error, Result};
use crate::store::Store;
use crate::types::Session;

impl Store {
    pub fn session_start(&self, project: &str, directory: Option<&str>) -> Result<Session> {
        // Reuse existing active session for this project (handles --resume, --continue, Ctrl+C)
        let existing: Option<Session> = self
            .conn()
            .query_row(
                "SELECT id, project, directory, started_at, ended_at, summary, status
                 FROM sessions WHERE project = ?1 AND directory IS ?2 AND status = 'active'
                 ORDER BY started_at DESC LIMIT 1",
                params![project, directory],
                |row| {
                    Ok(Session {
                        id: row.get(0)?,
                        project: row.get(1)?,
                        directory: row.get(2)?,
                        started_at: row.get(3)?,
                        ended_at: row.get(4)?,
                        summary: row.get(5)?,
                        status: row.get(6)?,
                    })
                },
            )
            .ok();

        if let Some(session) = existing {
            return Ok(session);
        }

        // Reactivate a completed session for the same project+directory
        let completed: Option<String> = self
            .conn()
            .query_row(
                "SELECT id FROM sessions
                 WHERE project = ?1 AND directory IS ?2 AND status = 'completed'
                 ORDER BY started_at DESC LIMIT 1",
                params![project, directory],
                |row| row.get(0),
            )
            .ok();

        if let Some(existing_id) = completed {
            self.conn().execute(
                "UPDATE sessions SET status = 'active', ended_at = NULL, summary = NULL
                 WHERE id = ?1",
                params![existing_id],
            )?;
            return self
                .conn()
                .query_row(
                    "SELECT id, project, directory, started_at, ended_at, summary, status
                     FROM sessions WHERE id = ?1",
                    params![existing_id],
                    |row| {
                        Ok(Session {
                            id: row.get(0)?,
                            project: row.get(1)?,
                            directory: row.get(2)?,
                            started_at: row.get(3)?,
                            ended_at: row.get(4)?,
                            summary: row.get(5)?,
                            status: row.get(6)?,
                        })
                    },
                )
                .map_err(Error::Database);
        }

        let id = generate_session_id();
        self.conn().execute(
            "INSERT INTO sessions (id, project, directory) VALUES (?1, ?2, ?3)",
            params![id, project, directory],
        )?;

        self.conn()
            .query_row(
                "SELECT id, project, directory, started_at, ended_at, summary, status
                 FROM sessions WHERE id = ?1",
                params![id],
                |row| {
                    Ok(Session {
                        id: row.get(0)?,
                        project: row.get(1)?,
                        directory: row.get(2)?,
                        started_at: row.get(3)?,
                        ended_at: row.get(4)?,
                        summary: row.get(5)?,
                        status: row.get(6)?,
                    })
                },
            )
            .map_err(Error::Database)
    }

    pub fn session_end(
        &self,
        session_id: &str,
        summary: Option<&str>,
        tokens_input: Option<i32>,
        tokens_output: Option<i32>,
    ) -> Result<Session> {
        // Check existence first (to distinguish not-found from already-ended).
        let exists: bool = self
            .conn()
            .query_row(
                "SELECT COUNT(*) FROM sessions WHERE id = ?1",
                params![session_id],
                |row| row.get::<_, i64>(0),
            )
            .map(|n| n > 0)
            .unwrap_or(false);
        if !exists {
            return Err(Error::SessionNotFound(session_id.to_string()));
        }

        // Atomic check-and-update: only update if still active.
        let rows_changed = self.conn().execute(
            "UPDATE sessions SET ended_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now'),
             summary = ?1, status = 'completed',
             tokens_used_input = ?3, tokens_used_output = ?4
             WHERE id = ?2 AND status = 'active'",
            params![summary, session_id, tokens_input, tokens_output],
        )?;

        if rows_changed == 0 {
            return Err(Error::SessionAlreadyEnded(session_id.to_string()));
        }

        self.conn()
            .query_row(
                "SELECT id, project, directory, started_at, ended_at, summary, status
                 FROM sessions WHERE id = ?1",
                params![session_id],
                |row| {
                    Ok(Session {
                        id: row.get(0)?,
                        project: row.get(1)?,
                        directory: row.get(2)?,
                        started_at: row.get(3)?,
                        ended_at: row.get(4)?,
                        summary: row.get(5)?,
                        status: row.get(6)?,
                    })
                },
            )
            .map_err(Error::Database)
    }

    /// Returns the ID of the most recently started active session, if any.
    /// Used as a fallback when the MCP server's in-memory session state has not been initialised
    /// (i.e. `memory_session_start` MCP tool was never called, only the CLI hook).
    pub fn most_recent_active_session_id(&self) -> Result<Option<String>> {
        Ok(self
            .conn()
            .query_row(
                "SELECT id FROM sessions WHERE status = 'active' ORDER BY started_at DESC LIMIT 1",
                [],
                |row| row.get(0),
            )
            .ok())
    }

    pub fn recent_sessions(&self, limit: i32) -> Result<Vec<Session>> {
        let mut stmt = self.conn().prepare(
            "SELECT id, project, directory, started_at, ended_at, summary, status
             FROM sessions ORDER BY started_at DESC LIMIT ?1",
        )?;

        let results = stmt
            .query_map(params![limit], |row| {
                Ok(Session {
                    id: row.get(0)?,
                    project: row.get(1)?,
                    directory: row.get(2)?,
                    started_at: row.get(3)?,
                    ended_at: row.get(4)?,
                    summary: row.get(5)?,
                    status: row.get(6)?,
                })
            })?
            .collect::<std::result::Result<Vec<_>, _>>()?;

        Ok(results)
    }

    /// End all active sessions, optionally filtered by project name.
    /// Returns the number of sessions ended.
    pub fn end_active_sessions(
        &self,
        project: Option<&str>,
        summary: Option<&str>,
    ) -> Result<usize> {
        let active = self.recent_sessions(10000)?;
        let mut ended = 0;
        for session in &active {
            if session.status != "active" {
                continue;
            }
            if let Some(proj) = project {
                if session.project != proj {
                    continue;
                }
            }
            self.session_end(&session.id, summary, None, None)?;
            ended += 1;
        }
        Ok(ended)
    }
}

fn generate_session_id() -> String {
    use std::time::{SystemTime, UNIX_EPOCH};
    let ts = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_nanos();
    format!("ses_{ts:x}")
}