DS4Server/src/database.rs

use diesel::prelude::*;
use diesel_migrations::{EmbeddedMigrations, MigrationHarness, embed_migrations};
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use time::OffsetDateTime;

use crate::config::PermissionMode;
use crate::schema::{a2ui_messages, messages, projects, sessions};

pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!("migrations");

#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
#[diesel(table_name = projects)]
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
pub struct Project {
    pub id: i32,
    pub name: String,
    pub path: String,
    /// Sidebar state: hides this project's session rows.
    pub collapsed: bool,
}

#[derive(Insertable)]
#[diesel(table_name = projects)]
struct NewProject<'a> {
    name: &'a str,
    path: &'a str,
}

#[derive(Associations, Clone, Debug, Identifiable, Queryable, Selectable)]
#[diesel(belongs_to(Project))]
#[diesel(table_name = sessions)]
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
pub struct Session {
    pub id: i32,
    pub project_id: i32,
    pub title: String,
    pub context_used: i32,
    pub context_limit: i32,
    pub last_tokens_per_second: Option<f32>,
    /// Raw column value; read it through [`Session::state`].
    state: String,
    pub compacted_summary: Option<String>,
    last_used: i64,
    /// Raw column value; read it through [`Session::permission_mode`].
    permission_mode: String,
}

impl Session {
    pub fn state(&self) -> SessionState {
        SessionState::from_id(&self.state).unwrap_or_default()
    }

    pub fn permission_mode(&self) -> PermissionMode {
        PermissionMode::from_id(&self.permission_mode).unwrap_or_default()
    }

    #[cfg(test)]
    pub fn fixture(id: i32, project_id: i32, title: &str, state: SessionState) -> Self {
        Self {
            id,
            project_id,
            title: title.to_owned(),
            context_used: 0,
            context_limit: 0,
            last_tokens_per_second: None,
            state: state.as_id().to_owned(),
            compacted_summary: None,
            last_used: 0,
            permission_mode: PermissionMode::default().as_id().to_owned(),
        }
    }
}

/// Where a session sits in the sidebar. These are lifecycle states, not flags:
/// a session is exactly one of them.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum SessionState {
    #[default]
    Normal,
    Pinned,
    Archived,
}

impl SessionState {
    pub fn as_id(self) -> &'static str {
        match self {
            Self::Normal => "normal",
            Self::Pinned => "pinned",
            Self::Archived => "archived",
        }
    }

    pub fn from_id(id: &str) -> Option<Self> {
        match id {
            "normal" => Some(Self::Normal),
            "pinned" => Some(Self::Pinned),
            "archived" => Some(Self::Archived),
            _ => None,
        }
    }

    /// Sidebar grouping order: pinned, then ordinary, then archived.
    pub fn rank(self) -> u8 {
        match self {
            Self::Pinned => 0,
            Self::Normal => 1,
            Self::Archived => 2,
        }
    }
}

#[derive(Insertable)]
#[diesel(table_name = sessions)]
struct NewSession<'a> {
    project_id: i32,
    title: &'a str,
    last_used: i64,
    permission_mode: &'a str,
}

#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
#[diesel(table_name = messages)]
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
pub struct StoredMessage {
    pub id: i32,
    pub session_id: i32,
    pub user: bool,
    pub tool: bool,
    pub reasoning: Option<String>,
    pub reasoning_complete: bool,
    pub content: String,
    pub model_content: Option<String>,
    pub system: bool,
    pub compaction: bool,
    pub compaction_tail_start: Option<i32>,
    pub generation_duration_ms: Option<i32>,
    pub input_tokens: Option<i32>,
    pub cached_tokens: Option<i32>,
    pub output_tokens: Option<i32>,
}

#[derive(Insertable)]
#[diesel(table_name = messages)]
struct NewMessage<'a> {
    session_id: i32,
    user: bool,
    tool: bool,
    reasoning: Option<&'a str>,
    reasoning_complete: bool,
    content: &'a str,
    model_content: Option<&'a str>,
    system: bool,
    compaction: bool,
    compaction_tail_start: Option<i32>,
}

impl<'a> NewMessage<'a> {
    fn system(session_id: i32, content: &'a str) -> Self {
        Self::new(session_id, content, false, false, true)
    }

    fn user(session_id: i32, content: &'a str, model_content: Option<&'a str>) -> Self {
        Self {
            model_content,
            ..Self::new(session_id, content, true, false, false)
        }
    }

    fn tool_result(session_id: i32, content: &'a str) -> Self {
        Self::new(session_id, content, false, true, false)
    }

    fn assistant(session_id: i32, reasoning: bool) -> Self {
        Self {
            reasoning: reasoning.then_some(""),
            reasoning_complete: !reasoning,
            ..Self::new(session_id, "", false, false, false)
        }
    }

    fn compaction(session_id: i32, content: &'a str, tail_start: Option<i32>) -> Self {
        Self {
            compaction: true,
            compaction_tail_start: tail_start,
            ..Self::system(session_id, content)
        }
    }

    fn new(session_id: i32, content: &'a str, user: bool, tool: bool, system: bool) -> Self {
        Self {
            session_id,
            user,
            tool,
            reasoning: None,
            reasoning_complete: true,
            content,
            model_content: None,
            system,
            compaction: false,
            compaction_tail_start: None,
        }
    }
}

#[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
#[diesel(table_name = a2ui_messages)]
#[diesel(check_for_backend(diesel::sqlite::Sqlite))]
pub struct StoredA2uiMessage {
    pub id: i32,
    pub session_id: i32,
    pub message_id: i32,
    pub json: String,
    pub dismissed: bool,
}

#[derive(Insertable)]
#[diesel(table_name = a2ui_messages)]
struct NewA2uiMessage<'a> {
    session_id: i32,
    message_id: i32,
    json: &'a str,
    dismissed: bool,
}

#[derive(Debug)]
pub struct ProjectWithSessions {
    pub project: Project,
    pub sessions: Vec<Session>,
}

pub struct Database {
    connection: SqliteConnection,
}

impl Database {
    pub fn open(path: &Path) -> Result<Self, String> {
        let parent = path
            .parent()
            .ok_or_else(|| "database path has no parent directory".to_owned())?;
        fs::create_dir_all(parent).map_err(|error| error.to_string())?;
        let url = path
            .to_str()
            .ok_or_else(|| "database path is not valid UTF-8".to_owned())?;
        let mut connection = SqliteConnection::establish(url).map_err(|error| error.to_string())?;
        connection
            .run_pending_migrations(MIGRATIONS)
            .map_err(|error| error.to_string())?;
        Ok(Self { connection })
    }

    pub fn load_projects(&mut self) -> Result<Vec<ProjectWithSessions>, String> {
        let project_rows = projects::table
            .order(projects::id.asc())
            .select(Project::as_select())
            .load(&mut self.connection)
            .map_err(|error| error.to_string())?;
        let session_rows = sessions::table
            .order(sessions::id.asc())
            .select(Session::as_select())
            .load::<Session>(&mut self.connection)
            .map_err(|error| error.to_string())?;
        let mut sessions_by_project: HashMap<i32, Vec<Session>> = HashMap::new();
        for session in session_rows {
            sessions_by_project
                .entry(session.project_id)
                .or_default()
                .push(session);
        }
        // Preserve sidebar groups while matching DS4's newest-used-first order.
        for sessions in sessions_by_project.values_mut() {
            sessions.sort_by(|left, right| {
                left.state()
                    .rank()
                    .cmp(&right.state().rank())
                    .then_with(|| right.last_used.cmp(&left.last_used))
                    .then_with(|| right.id.cmp(&left.id))
            });
        }

        Ok(project_rows
            .into_iter()
            .map(|project| ProjectWithSessions {
                sessions: sessions_by_project.remove(&project.id).unwrap_or_default(),
                project,
            })
            .collect())
    }

    pub fn create_project(&mut self, name: &str, path: &str) -> Result<Project, String> {
        diesel::insert_into(projects::table)
            .values(NewProject { name, path })
            .returning(Project::as_returning())
            .get_result(&mut self.connection)
            .map_err(|error| error.to_string())
    }

    pub fn set_project_collapsed(
        &mut self,
        project_id: i32,
        collapsed: bool,
    ) -> Result<(), String> {
        diesel::update(projects::table.find(project_id))
            .set(projects::collapsed.eq(collapsed))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn delete_project(&mut self, project_id: i32) -> Result<(), String> {
        self.connection
            .transaction(|connection| {
                diesel::delete(
                    a2ui_messages::table.filter(
                        a2ui_messages::session_id.eq_any(
                            sessions::table
                                .filter(sessions::project_id.eq(project_id))
                                .select(sessions::id),
                        ),
                    ),
                )
                .execute(connection)?;
                diesel::delete(
                    messages::table.filter(
                        messages::session_id.eq_any(
                            sessions::table
                                .filter(sessions::project_id.eq(project_id))
                                .select(sessions::id),
                        ),
                    ),
                )
                .execute(connection)?;
                diesel::delete(sessions::table.filter(sessions::project_id.eq(project_id)))
                    .execute(connection)?;
                diesel::delete(projects::table.find(project_id)).execute(connection)?;
                Ok(())
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }

    pub fn create_session(
        &mut self,
        project_id: i32,
        title: &str,
        permission_mode: PermissionMode,
    ) -> Result<Session, String> {
        diesel::insert_into(sessions::table)
            .values(NewSession {
                project_id,
                title,
                last_used: OffsetDateTime::now_utc().unix_timestamp(),
                permission_mode: permission_mode.as_id(),
            })
            .returning(Session::as_returning())
            .get_result(&mut self.connection)
            .map_err(|error| error.to_string())
    }

    pub fn set_session_permission_mode(
        &mut self,
        session_id: i32,
        permission_mode: PermissionMode,
    ) -> Result<(), String> {
        diesel::update(sessions::table.find(session_id))
            .set(sessions::permission_mode.eq(permission_mode.as_id()))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn touch_session(&mut self, session_id: i32) -> Result<(), String> {
        touch_session(&mut self.connection, session_id).map_err(|error| error.to_string())
    }

    pub fn rename_session(&mut self, session_id: i32, title: &str) -> Result<(), String> {
        let title = title.trim();
        if title.is_empty() {
            return Err("Session title cannot be empty.".into());
        }
        diesel::update(sessions::table.find(session_id))
            .set(sessions::title.eq(title))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn set_session_state(
        &mut self,
        session_id: i32,
        state: SessionState,
    ) -> Result<(), String> {
        diesel::update(sessions::table.find(session_id))
            .set(sessions::state.eq(state.as_id()))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn delete_session(&mut self, session_id: i32) -> Result<(), String> {
        self.connection
            .transaction(|connection| {
                diesel::delete(
                    a2ui_messages::table.filter(a2ui_messages::session_id.eq(session_id)),
                )
                .execute(connection)?;
                diesel::delete(messages::table.filter(messages::session_id.eq(session_id)))
                    .execute(connection)?;
                diesel::delete(sessions::table.find(session_id)).execute(connection)?;
                Ok(())
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }

    pub fn load_messages(&mut self, session_id: i32) -> Result<Vec<StoredMessage>, String> {
        messages::table
            .filter(messages::session_id.eq(session_id))
            .order(messages::id.asc())
            .select(StoredMessage::as_select())
            .load(&mut self.connection)
            .map_err(|error| error.to_string())
    }

    pub fn load_a2ui_messages(
        &mut self,
        session_id: i32,
    ) -> Result<Vec<StoredA2uiMessage>, String> {
        a2ui_messages::table
            .filter(a2ui_messages::session_id.eq(session_id))
            .order(a2ui_messages::id.asc())
            .select(StoredA2uiMessage::as_select())
            .load(&mut self.connection)
            .map_err(|error| error.to_string())
    }

    pub fn insert_a2ui_message(
        &mut self,
        session_id: i32,
        message_id: i32,
        json: &str,
    ) -> Result<StoredA2uiMessage, String> {
        diesel::insert_into(a2ui_messages::table)
            .values(NewA2uiMessage {
                session_id,
                message_id,
                json,
                dismissed: false,
            })
            .returning(StoredA2uiMessage::as_returning())
            .get_result(&mut self.connection)
            .map_err(|error| error.to_string())
    }

    pub fn dismiss_a2ui_surface(
        &mut self,
        session_id: i32,
        surface_id: &str,
    ) -> Result<StoredMessage, String> {
        let content = format!(
            "A2UI surface `{surface_id}` was dismissed by the user. There is no active A2UI surface. If the user asks for UI again, create a complete new surface with createSurface and a new surfaceId; do not update any earlier surface."
        );
        let json = serde_json::json!({"surfaceId": surface_id}).to_string();
        self.connection
            .transaction(|connection| {
                let message = diesel::insert_into(messages::table)
                    .values(NewMessage::system(session_id, &content))
                    .returning(StoredMessage::as_returning())
                    .get_result(connection)?;
                diesel::insert_into(a2ui_messages::table)
                    .values(NewA2uiMessage {
                        session_id,
                        message_id: message.id,
                        json: &json,
                        dismissed: true,
                    })
                    .execute(connection)?;
                Ok(message)
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }

    pub fn update_session_context(
        &mut self,
        session_id: i32,
        used: u32,
        limit: u32,
        tokens_per_second: Option<f32>,
    ) -> Result<(), String> {
        let used = i32::try_from(used).map_err(|_| "Used context is too large to save")?;
        let limit = i32::try_from(limit).map_err(|_| "Context limit is too large to save")?;
        if tokens_per_second.is_some_and(|speed| !speed.is_finite() || speed < 0.0) {
            return Err("Generation speed is invalid".into());
        }
        diesel::update(sessions::table.find(session_id))
            .set((
                sessions::context_used.eq(used),
                sessions::context_limit.eq(limit),
                sessions::last_tokens_per_second.eq(tokens_per_second),
            ))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn start_chat_turn(
        &mut self,
        session_id: i32,
        prompt: &str,
        model_prompt: Option<&str>,
        system_messages: &[String],
        reasoning: bool,
    ) -> Result<Vec<StoredMessage>, String> {
        self.connection
            .transaction(|connection| {
                reactivate_archived_session(connection, session_id)?;
                touch_session(connection, session_id)?;
                let mut stored = Vec::with_capacity(system_messages.len() + 2);
                for content in system_messages {
                    stored.push(
                        diesel::insert_into(messages::table)
                            .values(NewMessage::system(session_id, content))
                            .returning(StoredMessage::as_returning())
                            .get_result(connection)?,
                    );
                }
                let user = diesel::insert_into(messages::table)
                    .values(NewMessage::user(session_id, prompt, model_prompt))
                    .returning(StoredMessage::as_returning())
                    .get_result(connection)?;
                let assistant = diesel::insert_into(messages::table)
                    .values(NewMessage::assistant(session_id, reasoning))
                    .returning(StoredMessage::as_returning())
                    .get_result(connection)?;
                stored.extend([user, assistant]);
                Ok(stored)
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }

    pub fn continue_tool_turn(
        &mut self,
        session_id: i32,
        result: &str,
        queued_user: Option<&str>,
        system_messages: &[String],
        reasoning: bool,
    ) -> Result<Vec<StoredMessage>, String> {
        self.connection
            .transaction(|connection| {
                touch_session(connection, session_id)?;
                let mut stored = Vec::with_capacity(system_messages.len() + 3);
                stored.push(
                    diesel::insert_into(messages::table)
                        .values(NewMessage::tool_result(session_id, result))
                        .returning(StoredMessage::as_returning())
                        .get_result(connection)?,
                );
                if let Some(content) = queued_user {
                    stored.push(
                        diesel::insert_into(messages::table)
                            .values(NewMessage::user(session_id, content, None))
                            .returning(StoredMessage::as_returning())
                            .get_result(connection)?,
                    );
                }
                for content in system_messages {
                    stored.push(
                        diesel::insert_into(messages::table)
                            .values(NewMessage::system(session_id, content))
                            .returning(StoredMessage::as_returning())
                            .get_result(connection)?,
                    );
                }
                let assistant = diesel::insert_into(messages::table)
                    .values(NewMessage::assistant(session_id, reasoning))
                    .returning(StoredMessage::as_returning())
                    .get_result(connection)?;
                stored.push(assistant);
                Ok(stored)
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }

    pub fn update_message(
        &mut self,
        id: i32,
        reasoning: Option<&str>,
        reasoning_complete: bool,
        content: &str,
    ) -> Result<(), String> {
        diesel::update(messages::table.find(id))
            .set((
                messages::reasoning.eq(reasoning),
                messages::reasoning_complete.eq(reasoning_complete),
                messages::content.eq(content),
            ))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn update_message_generation_stats(
        &mut self,
        id: i32,
        duration_ms: u64,
        input_tokens: u32,
        cached_tokens: u32,
        output_tokens: u32,
    ) -> Result<(), String> {
        let duration_ms =
            i32::try_from(duration_ms).map_err(|_| "Generation duration is too large to save")?;
        let input_tokens =
            i32::try_from(input_tokens).map_err(|_| "Input token count is too large to save")?;
        let cached_tokens =
            i32::try_from(cached_tokens).map_err(|_| "Cached token count is too large to save")?;
        let output_tokens =
            i32::try_from(output_tokens).map_err(|_| "Output token count is too large to save")?;
        diesel::update(messages::table.find(id))
            .set((
                messages::generation_duration_ms.eq(duration_ms),
                messages::input_tokens.eq(input_tokens),
                messages::cached_tokens.eq(cached_tokens),
                messages::output_tokens.eq(output_tokens),
            ))
            .execute(&mut self.connection)
            .map(|_| ())
            .map_err(|error| error.to_string())
    }

    pub fn record_compaction(
        &mut self,
        session_id: i32,
        summary: &str,
        tail_start: Option<i32>,
        running_jobs: Option<&str>,
        context_used: u32,
        context_limit: u32,
    ) -> Result<Vec<StoredMessage>, String> {
        let context_used = i32::try_from(context_used)
            .map_err(|_| "Compacted context is too large to save".to_owned())?;
        let context_limit = i32::try_from(context_limit)
            .map_err(|_| "Context limit is too large to save".to_owned())?;
        self.connection
            .transaction(|connection| {
                reactivate_archived_session(connection, session_id)?;
                diesel::update(sessions::table.find(session_id))
                    .set((
                        sessions::compacted_summary.eq(Some(summary)),
                        sessions::context_used.eq(context_used),
                        sessions::context_limit.eq(context_limit),
                        sessions::last_tokens_per_second.eq(None::<f32>),
                    ))
                    .execute(connection)?;
                let mut stored = Vec::with_capacity(2);
                stored.push(
                    diesel::insert_into(messages::table)
                        .values(NewMessage::compaction(session_id, summary, tail_start))
                        .returning(StoredMessage::as_returning())
                        .get_result(connection)?,
                );
                if let Some(content) = running_jobs {
                    stored.push(
                        diesel::insert_into(messages::table)
                            .values(NewMessage::tool_result(session_id, content))
                            .returning(StoredMessage::as_returning())
                            .get_result(connection)?,
                    );
                }
                Ok(stored)
            })
            .map_err(|error: diesel::result::Error| error.to_string())
    }
}

fn touch_session(
    connection: &mut SqliteConnection,
    session_id: i32,
) -> Result<(), diesel::result::Error> {
    diesel::update(sessions::table.find(session_id))
        .set(sessions::last_used.eq(OffsetDateTime::now_utc().unix_timestamp()))
        .execute(connection)
        .map(|_| ())
}

fn reactivate_archived_session(
    connection: &mut SqliteConnection,
    session_id: i32,
) -> Result<(), diesel::result::Error> {
    diesel::update(
        sessions::table
            .find(session_id)
            .filter(sessions::state.eq(SessionState::Archived.as_id())),
    )
    .set(sessions::state.eq(SessionState::Normal.as_id()))
    .execute(connection)
    .map(|_| ())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::time::{SystemTime, UNIX_EPOCH};

    #[test]
    fn migrations_and_crud_work() {
        let id = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("ds4-server-{id}.sqlite3"));
        let mut database = Database::open(&path).unwrap();

        let project = database.create_project("DS4", "/tmp/ds4").unwrap();
        let first = database
            .create_session(project.id, "First session", PermissionMode::Heuristic)
            .unwrap();
        let second = database
            .create_session(project.id, "Second session", PermissionMode::Ai)
            .unwrap();
        database.delete_session(first.id).unwrap();
        let loaded = database.load_projects().unwrap();
        assert_eq!(loaded[0].sessions[0].title, "Second session");
        assert_eq!(loaded[0].sessions[0].state(), SessionState::Normal);
        assert_eq!(loaded[0].sessions[0].permission_mode(), PermissionMode::Ai);
        database
            .set_session_permission_mode(second.id, PermissionMode::Heuristic)
            .unwrap();
        assert_eq!(
            database.load_projects().unwrap()[0].sessions[0].permission_mode(),
            PermissionMode::Heuristic
        );

        let ordinary = loaded[0].sessions[0].id;
        let pinned = database
            .create_session(project.id, "Pinned", PermissionMode::Heuristic)
            .unwrap();
        let archived = database
            .create_session(project.id, "Archived", PermissionMode::Heuristic)
            .unwrap();
        database
            .set_session_state(pinned.id, SessionState::Pinned)
            .unwrap();
        database
            .set_session_state(archived.id, SessionState::Archived)
            .unwrap();
        database.rename_session(pinned.id, "  Renamed  ").unwrap();
        assert!(database.rename_session(pinned.id, "   ").is_err());
        let loaded = database.load_projects().unwrap();
        assert_eq!(
            loaded[0]
                .sessions
                .iter()
                .map(|session| (session.id, session.state()))
                .collect::<Vec<_>>(),
            [
                (pinned.id, SessionState::Pinned),
                (ordinary, SessionState::Normal),
                (archived.id, SessionState::Archived),
            ]
        );
        assert_eq!(loaded[0].sessions[0].title, "Renamed");
        // Archiving a pinned session moves it out of the pinned group entirely.
        database
            .set_session_state(pinned.id, SessionState::Archived)
            .unwrap();
        let loaded = database.load_projects().unwrap();
        assert_eq!(loaded[0].sessions[0].id, ordinary);

        assert!(!loaded[0].project.collapsed);
        database.set_project_collapsed(project.id, true).unwrap();
        assert!(database.load_projects().unwrap()[0].project.collapsed);

        database.delete_project(project.id).unwrap();
        assert!(database.load_projects().unwrap().is_empty());
        drop(database);

        let mut reopened = Database::open(&path).unwrap();
        assert!(reopened.load_projects().unwrap().is_empty());
        drop(reopened);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn sessions_are_ordered_by_recent_use_within_sidebar_groups() {
        let id = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("ds4-session-order-{id}.sqlite3"));
        let mut database = Database::open(&path).unwrap();
        let project = database
            .create_project("DS4", "/tmp/ds4-session-order")
            .unwrap();
        let older = database
            .create_session(project.id, "Older", PermissionMode::Heuristic)
            .unwrap();
        let newer = database
            .create_session(project.id, "Newer", PermissionMode::Heuristic)
            .unwrap();
        let pinned = database
            .create_session(project.id, "Pinned", PermissionMode::Heuristic)
            .unwrap();
        database
            .set_session_state(pinned.id, SessionState::Pinned)
            .unwrap();
        diesel::update(sessions::table.find(older.id))
            .set(sessions::last_used.eq(100_i64))
            .execute(&mut database.connection)
            .unwrap();
        diesel::update(sessions::table.find(newer.id))
            .set(sessions::last_used.eq(200_i64))
            .execute(&mut database.connection)
            .unwrap();

        let loaded = database.load_projects().unwrap();
        assert_eq!(
            loaded[0]
                .sessions
                .iter()
                .map(|session| session.id)
                .collect::<Vec<_>>(),
            [pinned.id, newer.id, older.id]
        );

        drop(database);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn context_work_reactivates_an_archived_session() {
        let id = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("ds4-reactivate-{id}.sqlite3"));
        let mut database = Database::open(&path).unwrap();
        let project = database
            .create_project("DS4", "/tmp/ds4-reactivate")
            .unwrap();
        let session = database
            .create_session(project.id, "Archived", PermissionMode::Heuristic)
            .unwrap();
        database
            .set_session_state(session.id, SessionState::Archived)
            .unwrap();

        database
            .start_chat_turn(session.id, "Resume", None, &[], false)
            .unwrap();

        assert_eq!(
            database.load_projects().unwrap()[0].sessions[0].state(),
            SessionState::Normal
        );
        database
            .set_session_state(session.id, SessionState::Archived)
            .unwrap();
        database
            .record_compaction(session.id, "Summary", None, None, 100, 1_000)
            .unwrap();
        assert_eq!(
            database.load_projects().unwrap()[0].sessions[0].state(),
            SessionState::Normal
        );
        drop(database);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn a2ui_dismissal_persists_a_fresh_surface_boundary() {
        let id = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("ds4-a2ui-dismiss-{id}.sqlite3"));
        let mut database = Database::open(&path).unwrap();
        let project = database
            .create_project("DS4", "/tmp/ds4-a2ui-dismiss")
            .unwrap();
        let session = database
            .create_session(project.id, "A2UI", PermissionMode::Heuristic)
            .unwrap();
        let first = database
            .start_chat_turn(session.id, "First", None, &[], false)
            .unwrap()
            .pop()
            .unwrap();
        database
            .insert_a2ui_message(
                session.id,
                first.id,
                r#"{"version":"v1.0","createSurface":{"surfaceId":"first","catalogId":"https://ds4server.local/a2ui/v1_0/catalog.json"}}"#,
            )
            .unwrap();
        let dismissal = database.dismiss_a2ui_surface(session.id, "first").unwrap();
        assert!(dismissal.system);
        assert!(dismissal.content.contains("createSurface"));
        let protocol = database.load_a2ui_messages(session.id).unwrap();
        let (_, active, errors) = crate::a2ui::replay_epochs(protocol.iter().map(|message| {
            (
                message.id,
                message.message_id,
                message.dismissed,
                message.json.as_str(),
            )
        }));
        assert!(errors.is_empty());
        assert!(active.active_surface().is_none());

        let second = database
            .start_chat_turn(session.id, "Second", None, &[], false)
            .unwrap()
            .pop()
            .unwrap();
        database
            .insert_a2ui_message(
                session.id,
                second.id,
                r#"{"version":"v1.0","createSurface":{"surfaceId":"second","catalogId":"https://ds4server.local/a2ui/v1_0/catalog.json"}}"#,
            )
            .unwrap();
        drop(database);

        let mut reopened = Database::open(&path).unwrap();
        let protocol = reopened.load_a2ui_messages(session.id).unwrap();
        assert_eq!(protocol.len(), 3);
        assert!(protocol[1].dismissed);
        assert_eq!(protocol[1].message_id, dismissal.id);
        let (_, active, errors) = crate::a2ui::replay_epochs(protocol.iter().map(|message| {
            (
                message.id,
                message.message_id,
                message.dismissed,
                message.json.as_str(),
            )
        }));
        assert!(errors.is_empty());
        assert_eq!(active.active_surface().unwrap().id, "second");
        let messages = reopened.load_messages(session.id).unwrap();
        assert!(
            messages
                .iter()
                .position(|message| message.id == dismissal.id)
                .is_some_and(|index| messages[index + 1].content == "Second")
        );
        reopened.delete_project(project.id).unwrap();
        drop(reopened);
        fs::remove_file(path).unwrap();
    }

    #[test]
    fn chat_and_compaction_history_survive_reopen_and_session_deletion() {
        let id = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos();
        let path = std::env::temp_dir().join(format!("ds4-chat-{id}.sqlite3"));
        let mut database = Database::open(&path).unwrap();
        let project = database.create_project("DS4", "/tmp/ds4-chat").unwrap();
        let session = database
            .create_session(project.id, "Chat", PermissionMode::Heuristic)
            .unwrap();
        let mut opening = database
            .start_chat_turn(
                session.id,
                "Question",
                Some("Question\n\nhidden metadata"),
                &["Date context".into()],
                true,
            )
            .unwrap();
        let assistant = opening.pop().unwrap();
        database
            .update_message(assistant.id, Some("Reasoning"), true, "Answer")
            .unwrap();
        database
            .update_message_generation_stats(assistant.id, 12_345, 1_024, 768, 256)
            .unwrap();
        database
            .insert_a2ui_message(
                session.id,
                assistant.id,
                r#"{"version":"v1.0","createSurface":{"surfaceId":"saved","catalogId":"https://ds4server.local/a2ui/v1_0/catalog.json"}}"#,
            )
            .unwrap();
        database
            .continue_tool_turn(
                session.id,
                "Tool result",
                Some("Queued correction"),
                &["Tool reminder".into()],
                false,
            )
            .unwrap();
        database
            .update_session_context(session.id, 1_234, 65_536, Some(12.5))
            .unwrap();
        drop(database);

        let mut reopened = Database::open(&path).unwrap();
        let projects = reopened.load_projects().unwrap();
        assert_eq!(projects[0].sessions[0].context_used, 1_234);
        assert_eq!(projects[0].sessions[0].context_limit, 65_536);
        assert_eq!(projects[0].sessions[0].last_tokens_per_second, Some(12.5));
        let messages = reopened.load_messages(session.id).unwrap();
        assert_eq!(messages.len(), 7);
        assert!(messages[0].system);
        assert_eq!(messages[1].content, "Question");
        assert_eq!(
            messages[1].model_content.as_deref(),
            Some("Question\n\nhidden metadata")
        );
        assert_eq!(messages[2].reasoning.as_deref(), Some("Reasoning"));
        assert!(messages[2].reasoning_complete);
        assert_eq!(messages[2].content, "Answer");
        assert_eq!(messages[2].generation_duration_ms, Some(12_345));
        assert_eq!(messages[2].input_tokens, Some(1_024));
        assert_eq!(messages[2].cached_tokens, Some(768));
        assert_eq!(messages[2].output_tokens, Some(256));
        assert!(messages[3].tool);
        assert_eq!(messages[3].content, "Tool result");
        assert!(messages[4].user);
        assert_eq!(messages[4].content, "Queued correction");
        assert!(messages[5].system);
        assert_eq!(messages[5].content, "Tool reminder");
        assert!(!messages[6].user);
        assert!(!messages[6].tool);
        let a2ui = reopened.load_a2ui_messages(session.id).unwrap();
        assert_eq!(a2ui.len(), 1);
        assert_eq!(a2ui[0].message_id, messages[2].id);
        assert!(a2ui[0].json.contains("createSurface"));
        let first = reopened
            .record_compaction(
                session.id,
                "First durable state.",
                Some(messages[4].id),
                Some("bash job=1 status=running"),
                321,
                32_768,
            )
            .unwrap();
        assert_eq!(first.len(), 2);
        assert!(first[0].compaction);
        assert_eq!(first[0].compaction_tail_start, Some(messages[4].id));
        drop(reopened);
        let mut reopened = Database::open(&path).unwrap();
        assert_eq!(
            reopened.load_projects().unwrap()[0].sessions[0]
                .compacted_summary
                .as_deref(),
            Some("First durable state.")
        );
        assert_eq!(
            reopened.load_projects().unwrap()[0].sessions[0].context_used,
            321
        );
        let messages = reopened.load_messages(session.id).unwrap();
        assert_eq!(messages.len(), 9);
        assert_eq!(messages[1].content, "Question");
        assert!(messages[7].compaction);
        assert_eq!(messages[7].content, "First durable state.");
        assert_eq!(messages[8].content, "bash job=1 status=running");
        reopened
            .continue_tool_turn(session.id, "Reloaded tool result", None, &[], false)
            .unwrap();
        let continued = reopened.load_messages(session.id).unwrap();
        let second_tail = continued[9].id;
        reopened
            .record_compaction(
                session.id,
                "Second durable state.",
                Some(second_tail),
                None,
                222,
                32_768,
            )
            .unwrap();
        reopened
            .continue_tool_turn(session.id, "Final tool result", None, &[], false)
            .unwrap();
        let continued = reopened.load_messages(session.id).unwrap();
        let third_tail = continued[12].id;
        reopened
            .record_compaction(
                session.id,
                "Third durable state.",
                Some(third_tail),
                None,
                111,
                32_768,
            )
            .unwrap();
        reopened
            .start_chat_turn(session.id, "After third compaction", None, &[], false)
            .unwrap();
        drop(reopened);

        let mut reopened = Database::open(&path).unwrap();
        let history = reopened.load_messages(session.id).unwrap();
        assert_eq!(history.len(), 17);
        assert_eq!(history[1].content, "Question");
        let markers = history
            .iter()
            .filter(|message| message.compaction)
            .collect::<Vec<_>>();
        assert_eq!(markers.len(), 3);
        assert_eq!(markers[0].content, "First durable state.");
        assert_eq!(markers[1].content, "Second durable state.");
        assert_eq!(markers[2].content, "Third durable state.");
        assert_eq!(markers[2].compaction_tail_start, Some(third_tail));
        assert_eq!(history[15].content, "After third compaction");
        reopened.delete_session(session.id).unwrap();
        assert!(reopened.load_messages(session.id).unwrap().is_empty());
        assert!(reopened.load_a2ui_messages(session.id).unwrap().is_empty());
        drop(reopened);
        fs::remove_file(path).unwrap();
    }
}