zesdex/src/app/mode/rewind.rs

#![allow(clippy::cast_possible_truncation, clippy::cast_sign_loss, clippy::cast_precision_loss, clippy::cast_possible_wrap)]
//! Rewind mode: restores a file to a pre-edit snapshot stored in the
//! session's `SQLite` blob store.

use crate::app::state::rest::AppStateRest;
use sha2::Digest;

/// Returns the number of stored pre-edit blobs (snapshots) for this session.
pub fn rewind_count(state: &AppStateRest) -> usize {
    let Ok(conn) = open_session_db(&state.session_dir) else { return 0 };
    crate::model::msglog::blobs::list_blob_keys(&conn, &state.session_id)
        .ok()
        .map_or(0, |keys| keys.len())
}

/// Restores a file to its pre-edit state by retrieving the blob stored under index
/// `index` (0 = oldest). Opens a fresh `SQLite` connection so this works outside
/// of a running turn (e.g. from the Rewind overlay).
pub fn rewind_to(state: &mut AppStateRest, index: usize) {
    let conn = match open_session_db(&state.session_dir) {
        Ok(c) => c,
        Err(e) => {
            state.push_toast(crate::app::state::types::Toast::new(
                crate::app::state::types::ToastKind::Error,
                format!("Failed to open session DB: {e}"),
            ));
            state.dirty = true;
            return;
        }
    };

    let keys = match crate::model::msglog::blobs::list_blob_keys(&conn, &state.session_id) {
        Ok(k) => k,
        Err(e) => {
            state.push_toast(crate::app::state::types::Toast::new(
                crate::app::state::types::ToastKind::Error,
                format!("Failed to list snapshots: {e}"),
            ));
            state.dirty = true;
            return;
        }
    };

    if keys.is_empty() || index >= keys.len() {
        state.push_toast(crate::app::state::types::Toast::new(
            crate::app::state::types::ToastKind::Warning,
            "No snapshot available at that index".to_string(),
        ));
        state.dirty = true;
        return;
    }

    let blob_key = &keys[index];
    let bytes = match crate::model::msglog::blobs::retrieve_blob(&conn, &state.session_id, blob_key) {
        Ok(Some(b)) => b,
        Ok(None) => {
            state.push_toast(crate::app::state::types::Toast::new(
                crate::app::state::types::ToastKind::Error,
                "Snapshot data not found".to_string(),
            ));
            state.dirty = true;
            return;
        }
        Err(e) => {
            state.push_toast(crate::app::state::types::Toast::new(
                crate::app::state::types::ToastKind::Error,
                format!("Failed to retrieve snapshot: {e}"),
            ));
            state.dirty = true;
            return;
        }
    };

    // Look up the path from the edit log — the blob key is the tool_call_id.
    // The edit log doesn't store the tool_call_id directly, so fall back to the
    // path from the most recent write/edit entry.
    let restore_path = find_edit_path(state, blob_key)
        .unwrap_or_else(|| state.session_dir.join("snapshot.dat"));

    match std::fs::write(&restore_path, &bytes) {
        Ok(()) => {
            state.push_toast(crate::app::state::types::Toast::new(
                crate::app::state::types::ToastKind::Success,
                format!("Restored {} from snapshot", restore_path.display()),
            ));
        }
        Err(e) => {
            state.push_toast(crate::app::state::types::Toast::new(
                crate::app::state::types::ToastKind::Error,
                format!("Failed to write restored file: {e}"),
            ));
        }
    }

    // Log the rewind itself as an edit entry
    let mut el = crate::model::editlog::EditLog::new(&state.session_dir);
    let entry = crate::model::editlog::EditLogEntry {
        ts: chrono::Utc::now().timestamp_millis(),
        tool: "rewind".to_string(),
        path: restore_path.to_string_lossy().to_string(),
        reason: format!("rewind_to({index})"),
        content_sha256: hex::encode(sha2::Sha256::digest(&bytes)),
        bytes_delta: bytes.len() as i64,
        origin: crate::app::state::types::Origin::Main.tag(),
        session_id: state.session_id.clone(),
    };
    let _ = el.append(entry);

    // Clear the transcript to force a refresh
    state.transcript_cache.dirty = true;
    state.dirty = true;
}

fn open_session_db(session_dir: &std::path::Path) -> anyhow::Result<rusqlite::Connection> {
    let path = session_dir.join("messages.sqlite");
    let conn = rusqlite::Connection::open(&path)?;
    Ok(conn)
}

fn find_edit_path(state: &AppStateRest, _blob_key: &str) -> Option<std::path::PathBuf> {
    let el = crate::model::editlog::EditLog::new(&state.session_dir);
    let entry = el.entries.iter().rev().find(|e| e.tool == "write" || e.tool == "edit")?;
    Some(std::path::PathBuf::from(&entry.path))
}