zeph_memory/

snapshot.rs

// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Memory snapshot export and import.
//!
//! A [`MemorySnapshot`] serializes the full set of conversations, messages, and
//! summaries from `SQLite` into a JSON document suitable for backup or migration.
//!
//! Import is idempotent: conversations and messages already present in the target
//! database are skipped and counted in [`ImportStats::skipped`].

use serde::{Deserialize, Serialize};
#[allow(unused_imports)]
use zeph_db::sql;

use zeph_db::ActiveDialect;

use crate::error::MemoryError;
use crate::store::SqliteStore;
use crate::types::ConversationId;

/// Complete serializable snapshot of the `SQLite` memory store.
#[derive(Debug, Serialize, Deserialize)]
pub struct MemorySnapshot {
    /// Schema version for forward-compatibility checks.
    pub version: u32,
    /// RFC 3339 timestamp when the snapshot was created.
    pub exported_at: String,
    /// All conversations with their messages and summaries.
    pub conversations: Vec<ConversationSnapshot>,
}

/// One conversation and all its associated data in a [`MemorySnapshot`].
#[derive(Debug, Serialize, Deserialize)]
pub struct ConversationSnapshot {
    /// `SQLite` row ID of the conversation.
    pub id: i64,
    /// All messages in this conversation.
    pub messages: Vec<MessageSnapshot>,
    /// Compression summaries for this conversation.
    pub summaries: Vec<SummarySnapshot>,
}

/// A single message as stored in a [`MemorySnapshot`].
#[derive(Debug, Serialize, Deserialize)]
pub struct MessageSnapshot {
    /// `SQLite` row ID.
    pub id: i64,
    /// Parent conversation ID.
    pub conversation_id: i64,
    /// Message role (`"user"`, `"assistant"`, `"system"`).
    pub role: String,
    /// Flattened text content.
    pub content: String,
    /// JSON-encoded `Vec<MessagePart>` payload.
    pub parts_json: String,
    /// Unix timestamp (seconds since epoch).
    pub created_at: i64,
}

/// A compression summary record in a [`MemorySnapshot`].
#[derive(Debug, Serialize, Deserialize)]
pub struct SummarySnapshot {
    /// `SQLite` row ID.
    pub id: i64,
    /// Parent conversation ID.
    pub conversation_id: i64,
    /// Summary text.
    pub content: String,
    /// Inclusive lower bound of the summarised message range.
    pub first_message_id: Option<i64>,
    /// Inclusive upper bound of the summarised message range.
    pub last_message_id: Option<i64>,
    /// Estimated token count of the summary content.
    pub token_estimate: i64,
}

/// Counters returned by [`import_snapshot`].
#[derive(Debug, Default)]
pub struct ImportStats {
    /// Number of conversations inserted.
    pub conversations_imported: usize,
    /// Number of messages inserted.
    pub messages_imported: usize,
    /// Number of summaries inserted.
    pub summaries_imported: usize,
    /// Number of rows skipped because they already existed.
    pub skipped: usize,
}

/// Export all conversations, messages and summaries from `SQLite` into a snapshot.
///
/// # Errors
///
/// Returns an error if any database query fails.
pub async fn export_snapshot(sqlite: &SqliteStore) -> Result<MemorySnapshot, MemoryError> {
    let conv_ids: Vec<(i64,)> =
        zeph_db::query_as(sql!("SELECT id FROM conversations ORDER BY id ASC"))
            .fetch_all(sqlite.pool())
            .await?;

    let exported_at = chrono_now();
    let mut conversations = Vec::with_capacity(conv_ids.len());

    for (cid_raw,) in conv_ids {
        let cid = ConversationId(cid_raw);

        let epoch_expr = <ActiveDialect as zeph_db::dialect::Dialect>::epoch_from_col("created_at");
        let parts_select = <ActiveDialect as zeph_db::dialect::Dialect>::select_as_text("parts");
        let msg_sql = zeph_db::rewrite_placeholders(&format!(
            "SELECT id, role, content, {parts_select} AS parts, {epoch_expr} \
            FROM messages WHERE conversation_id = ? ORDER BY id ASC"
        ));
        let msg_rows: Vec<(i64, String, String, String, i64)> =
            zeph_db::query_as(sqlx::AssertSqlSafe(msg_sql))
                .bind(cid)
                .fetch_all(sqlite.pool())
                .await?;

        let messages = msg_rows
            .into_iter()
            .map(
                |(id, role, content, parts_json, created_at)| MessageSnapshot {
                    id,
                    conversation_id: cid_raw,
                    role,
                    content,
                    parts_json,
                    created_at,
                },
            )
            .collect();

        let sum_rows = sqlite.load_summaries(cid).await?;
        let summaries = sum_rows
            .into_iter()
            .map(
                |(
                    id,
                    conversation_id,
                    content,
                    first_message_id,
                    last_message_id,
                    token_estimate,
                )| {
                    SummarySnapshot {
                        id,
                        conversation_id: conversation_id.0,
                        content,
                        first_message_id: first_message_id.map(|m| m.0),
                        last_message_id: last_message_id.map(|m| m.0),
                        token_estimate,
                    }
                },
            )
            .collect();

        conversations.push(ConversationSnapshot {
            id: cid_raw,
            messages,
            summaries,
        });
    }

    Ok(MemorySnapshot {
        version: 1,
        exported_at,
        conversations,
    })
}

/// Import a snapshot into `SQLite`, skipping duplicate entries.
///
/// Returns stats about what was imported.
///
/// # Errors
///
/// Returns an error if any database operation fails.
pub async fn import_snapshot(
    sqlite: &SqliteStore,
    snapshot: MemorySnapshot,
) -> Result<ImportStats, MemoryError> {
    if snapshot.version != 1 {
        return Err(MemoryError::Snapshot(format!(
            "unsupported snapshot version {}: only version 1 is supported",
            snapshot.version
        )));
    }
    let mut stats = ImportStats::default();

    for conv in snapshot.conversations {
        let exists: Option<(i64,)> =
            zeph_db::query_as(sql!("SELECT id FROM conversations WHERE id = ?"))
                .bind(conv.id)
                .fetch_optional(sqlite.pool())
                .await?;

        if exists.is_none() {
            zeph_db::query(sql!("INSERT INTO conversations (id) VALUES (?)"))
                .bind(conv.id)
                .execute(sqlite.pool())
                .await?;
            stats.conversations_imported += 1;
        } else {
            stats.skipped += 1;
        }

        for msg in conv.messages {
            let json_cast = <ActiveDialect as zeph_db::dialect::Dialect>::JSON_CAST;
            let msg_sql = zeph_db::rewrite_placeholders(&format!(
                "{} INTO messages (id, conversation_id, role, content, parts) VALUES (?, ?, ?, ?, ?{json_cast}){}",
                <ActiveDialect as zeph_db::dialect::Dialect>::INSERT_IGNORE,
                <ActiveDialect as zeph_db::dialect::Dialect>::CONFLICT_NOTHING,
            ));
            let result = zeph_db::query(sqlx::AssertSqlSafe(msg_sql))
                .bind(msg.id)
                .bind(msg.conversation_id)
                .bind(&msg.role)
                .bind(&msg.content)
                .bind(&msg.parts_json)
                .execute(sqlite.pool())
                .await?;

            if result.rows_affected() > 0 {
                stats.messages_imported += 1;
            } else {
                stats.skipped += 1;
            }
        }

        for sum in conv.summaries {
            let sum_sql = zeph_db::rewrite_placeholders(&format!(
                "{} INTO summaries (id, conversation_id, content, first_message_id, last_message_id, token_estimate) VALUES (?, ?, ?, ?, ?, ?){}",
                <ActiveDialect as zeph_db::dialect::Dialect>::INSERT_IGNORE,
                <ActiveDialect as zeph_db::dialect::Dialect>::CONFLICT_NOTHING,
            ));
            let result = zeph_db::query(sqlx::AssertSqlSafe(sum_sql))
                .bind(sum.id)
                .bind(sum.conversation_id)
                .bind(&sum.content)
                .bind(sum.first_message_id)
                .bind(sum.last_message_id)
                .bind(sum.token_estimate)
                .execute(sqlite.pool())
                .await?;

            if result.rows_affected() > 0 {
                stats.summaries_imported += 1;
            } else {
                stats.skipped += 1;
            }
        }
    }

    Ok(stats)
}

fn chrono_now() -> String {
    zeph_common::timestamp::utc_now_rfc3339()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn export_empty_database() {
        let store = SqliteStore::new(":memory:").await.unwrap();
        let snapshot = export_snapshot(&store).await.unwrap();
        assert_eq!(snapshot.version, 1);
        assert!(snapshot.conversations.is_empty());
        assert!(!snapshot.exported_at.is_empty());
    }

    #[tokio::test]
    async fn export_import_roundtrip() {
        let src = SqliteStore::new(":memory:").await.unwrap();
        let cid = src.create_conversation().await.unwrap();
        src.save_message(cid, "user", "hello export").await.unwrap();
        src.save_message(cid, "assistant", "hi import")
            .await
            .unwrap();

        let snapshot = export_snapshot(&src).await.unwrap();
        assert_eq!(snapshot.conversations.len(), 1);
        assert_eq!(snapshot.conversations[0].messages.len(), 2);

        let dst = SqliteStore::new(":memory:").await.unwrap();
        let stats = import_snapshot(&dst, snapshot).await.unwrap();
        assert_eq!(stats.conversations_imported, 1);
        assert_eq!(stats.messages_imported, 2);
        assert_eq!(stats.skipped, 0);

        let history = dst.load_history(cid, 50).await.unwrap();
        assert_eq!(history.len(), 2);
        assert_eq!(history[0].content, "hello export");
        assert_eq!(history[1].content, "hi import");
    }

    #[tokio::test]
    async fn import_duplicate_skips() {
        let src = SqliteStore::new(":memory:").await.unwrap();
        let cid = src.create_conversation().await.unwrap();
        src.save_message(cid, "user", "msg").await.unwrap();

        let snapshot = export_snapshot(&src).await.unwrap();

        let dst = SqliteStore::new(":memory:").await.unwrap();
        let stats1 = import_snapshot(&dst, snapshot).await.unwrap();
        assert_eq!(stats1.messages_imported, 1);

        let snapshot2 = export_snapshot(&src).await.unwrap();
        let stats2 = import_snapshot(&dst, snapshot2).await.unwrap();
        assert_eq!(stats2.messages_imported, 0);
        assert!(stats2.skipped > 0);
    }

    #[tokio::test]
    async fn import_existing_conversation_increments_skipped_not_imported() {
        let src = SqliteStore::new(":memory:").await.unwrap();
        let cid = src.create_conversation().await.unwrap();
        src.save_message(cid, "user", "only message").await.unwrap();

        let snapshot = export_snapshot(&src).await.unwrap();

        // Import once — conversation is new.
        let dst = SqliteStore::new(":memory:").await.unwrap();
        let stats1 = import_snapshot(&dst, snapshot).await.unwrap();
        assert_eq!(stats1.conversations_imported, 1);
        assert_eq!(stats1.messages_imported, 1);

        // Import again with no new messages — conversation already exists, must be counted as skipped.
        let snapshot2 = export_snapshot(&src).await.unwrap();
        let stats2 = import_snapshot(&dst, snapshot2).await.unwrap();
        assert_eq!(
            stats2.conversations_imported, 0,
            "existing conversation must not be counted as imported"
        );
        // The conversation itself contributes one skipped, plus the duplicate message.
        assert!(
            stats2.skipped >= 1,
            "re-importing an existing conversation must increment skipped"
        );
    }

    #[tokio::test]
    async fn export_includes_summaries() {
        let store = SqliteStore::new(":memory:").await.unwrap();
        let cid = store.create_conversation().await.unwrap();
        let m1 = store.save_message(cid, "user", "a").await.unwrap();
        let m2 = store.save_message(cid, "assistant", "b").await.unwrap();
        store
            .save_summary(cid, "summary", Some(m1), Some(m2), 5)
            .await
            .unwrap();

        let snapshot = export_snapshot(&store).await.unwrap();
        assert_eq!(snapshot.conversations[0].summaries.len(), 1);
        assert_eq!(snapshot.conversations[0].summaries[0].content, "summary");
    }

    #[test]
    fn chrono_now_not_empty() {
        let ts = chrono_now();
        assert!(ts.contains('T'));
        assert!(ts.ends_with('Z'));
    }

    #[test]
    fn import_corrupt_json_returns_error() {
        let result = serde_json::from_str::<MemorySnapshot>("not valid json at all {{{");
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn import_unsupported_version_returns_error() {
        let store = SqliteStore::new(":memory:").await.unwrap();
        let snapshot = MemorySnapshot {
            version: 99,
            exported_at: "2026-01-01T00:00:00Z".into(),
            conversations: vec![],
        };
        let err = import_snapshot(&store, snapshot).await.unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("unsupported snapshot version 99"));
    }

    #[tokio::test]
    async fn import_partial_overlap_adds_new_messages() {
        let src = SqliteStore::new(":memory:").await.unwrap();
        let cid = src.create_conversation().await.unwrap();
        src.save_message(cid, "user", "existing message")
            .await
            .unwrap();

        let snapshot1 = export_snapshot(&src).await.unwrap();

        let dst = SqliteStore::new(":memory:").await.unwrap();
        let stats1 = import_snapshot(&dst, snapshot1).await.unwrap();
        assert_eq!(stats1.messages_imported, 1);

        src.save_message(cid, "assistant", "new reply")
            .await
            .unwrap();
        let snapshot2 = export_snapshot(&src).await.unwrap();
        let stats2 = import_snapshot(&dst, snapshot2).await.unwrap();

        assert_eq!(
            stats2.messages_imported, 1,
            "only the new message should be imported"
        );
        // skipped includes the existing conversation (1) plus the duplicate message (1).
        assert_eq!(
            stats2.skipped, 2,
            "existing conversation and duplicate message should be skipped"
        );

        let history = dst.load_history(cid, 50).await.unwrap();
        assert_eq!(history.len(), 2);
        assert_eq!(history[1].content, "new reply");
    }
}