memvid_cli/commands/

session.rs

//! CLI commands for time-travel replay session management.
//!
//! These commands enable recording and managing agent sessions for
//! debugging and testing purposes.

use anyhow::Result;
use clap::{Args, Subcommand};
use colored::Colorize;
use memvid_core::Memvid;

/// Session management commands for time-travel replay
#[derive(Args, Debug)]
pub struct SessionArgs {
    #[command(subcommand)]
    pub command: SessionCommand,
}

#[derive(Subcommand, Debug)]
pub enum SessionCommand {
    /// Start recording a new session
    Start(SessionStartArgs),
    /// End the current recording session
    End(SessionEndArgs),
    /// List all recorded sessions
    List(SessionListArgs),
    /// View details of a specific session
    View(SessionViewArgs),
    /// Create a checkpoint in the current session
    Checkpoint(SessionCheckpointArgs),
    /// Delete a recorded session
    Delete(SessionDeleteArgs),
    /// Save sessions to the memory file
    Save(SessionSaveArgs),
    /// Load sessions from the memory file
    Load(SessionLoadArgs),
    /// Replay a recorded session and verify consistency
    Replay(SessionReplayArgs),
    /// Compare two recorded sessions
    Compare(SessionCompareArgs),
}

#[derive(Args, Debug)]
pub struct SessionStartArgs {
    /// Path to the memory file
    pub file: String,

    /// Optional name for the session
    #[arg(short, long)]
    pub name: Option<String>,

    /// Auto-checkpoint interval (number of actions between checkpoints, 0 = disabled)
    #[arg(long, default_value = "0")]
    pub auto_checkpoint: u64,
}

#[derive(Args, Debug)]
pub struct SessionEndArgs {
    /// Path to the memory file
    pub file: String,
}

#[derive(Args, Debug)]
pub struct SessionListArgs {
    /// Path to the memory file
    pub file: String,

    /// Output format (text, json)
    #[arg(short, long, default_value = "text")]
    pub format: String,
}

#[derive(Args, Debug)]
pub struct SessionViewArgs {
    /// Path to the memory file
    pub file: String,

    /// Session ID to view
    #[arg(short, long)]
    pub session: String,

    /// Output format (text, json)
    #[arg(short, long, default_value = "text")]
    pub format: String,
}

#[derive(Args, Debug)]
pub struct SessionCheckpointArgs {
    /// Path to the memory file
    pub file: String,
}

#[derive(Args, Debug)]
pub struct SessionDeleteArgs {
    /// Path to the memory file
    pub file: String,

    /// Session ID to delete
    #[arg(short, long)]
    pub session: String,
}

#[derive(Args, Debug)]
pub struct SessionSaveArgs {
    /// Path to the memory file
    pub file: String,
}

#[derive(Args, Debug)]
pub struct SessionLoadArgs {
    /// Path to the memory file
    pub file: String,
}

#[derive(Args, Debug)]
pub struct SessionReplayArgs {
    /// Path to the memory file
    pub file: String,

    /// Session ID to replay
    #[arg(short, long)]
    pub session: String,

    /// Start replay from a specific checkpoint
    #[arg(long)]
    pub from_checkpoint: Option<u64>,

    /// Number of results to retrieve during replay (default: same as original)
    /// Use higher values to discover documents that were missed due to low top-k
    #[arg(short = 'k', long)]
    pub top_k: Option<usize>,

    /// Use adaptive retrieval that adjusts top-k based on score distribution
    #[arg(long)]
    pub adaptive: bool,

    /// Minimum relevancy score for adaptive retrieval (0.0-1.0)
    #[arg(long, default_value = "0.5")]
    pub min_relevancy: f32,

    /// Skip put/update/delete actions
    #[arg(long)]
    pub skip_mutations: bool,

    /// Skip find/search actions
    #[arg(long)]
    pub skip_finds: bool,

    /// Skip ask actions (LLM calls)
    #[arg(long)]
    pub skip_asks: bool,

    /// Stop on first mismatch
    #[arg(long)]
    pub stop_on_mismatch: bool,

    /// Output format (text, json)
    #[arg(short, long, default_value = "text")]
    pub format: String,

    /// Launch the Time Machine web UI for visual replay
    #[cfg(feature = "web")]
    #[arg(long)]
    pub web: bool,

    /// Port for the Time Machine web server (default: 4242)
    #[cfg(feature = "web")]
    #[arg(long, default_value = "4242")]
    pub port: u16,

    /// Don't auto-open browser when starting web UI
    #[cfg(feature = "web")]
    #[arg(long)]
    pub no_open: bool,
}

#[derive(Args, Debug)]
pub struct SessionCompareArgs {
    /// Path to the memory file
    pub file: String,

    /// First session ID
    #[arg(short = 'a', long)]
    pub session_a: String,

    /// Second session ID
    #[arg(short = 'b', long)]
    pub session_b: String,

    /// Output format (text, json)
    #[arg(short, long, default_value = "text")]
    pub format: String,
}

/// Handle session subcommands
#[cfg(feature = "replay")]
pub fn handle_session(args: SessionArgs) -> Result<()> {
    match args.command {
        SessionCommand::Start(args) => handle_session_start(args),
        SessionCommand::End(args) => handle_session_end(args),
        SessionCommand::List(args) => handle_session_list(args),
        SessionCommand::View(args) => handle_session_view(args),
        SessionCommand::Checkpoint(args) => handle_session_checkpoint(args),
        SessionCommand::Delete(args) => handle_session_delete(args),
        SessionCommand::Save(args) => handle_session_save(args),
        SessionCommand::Load(args) => handle_session_load(args),
        SessionCommand::Replay(args) => handle_session_replay(args),
        SessionCommand::Compare(args) => handle_session_compare(args),
    }
}

#[cfg(not(feature = "replay"))]
pub fn handle_session(_args: SessionArgs) -> Result<()> {
    bail!("Session recording requires the 'replay' feature. Rebuild with --features replay")
}

#[cfg(feature = "replay")]
fn handle_session_start(args: SessionStartArgs) -> Result<()> {
    use memvid_core::replay::ReplayConfig;

    let mut mem = Memvid::open(&args.file)?;

    // Check if there's already an active session
    if mem.load_active_session()? {
        let session_id = mem.active_session_id().unwrap();
        anyhow::bail!(
            "Session {} is already active. End it first with `session end`.",
            session_id
        );
    }

    let config = ReplayConfig {
        auto_checkpoint_interval: args.auto_checkpoint,
        ..Default::default()
    };

    let session_id = mem.start_session(args.name, Some(config))?;

    // Save the active session to a sidecar file so it persists across commands
    mem.save_active_session()?;

    println!("Started session: {}", session_id);
    println!("Recording is now active. Run operations like put, find, ask.");
    println!("End the session with: memvid session end {}", args.file);

    Ok(())
}

#[cfg(feature = "replay")]
fn handle_session_end(args: SessionEndArgs) -> Result<()> {
    let mut mem = Memvid::open(&args.file)?;

    // Load the active session from the sidecar file
    if !mem.load_active_session()? {
        anyhow::bail!("No active session found. Start one with `session start`.");
    }

    let session = mem.end_session()?;
    println!("Ended session: {}", session.session_id);
    println!("  Actions recorded: {}", session.actions.len());
    println!("  Checkpoints: {}", session.checkpoints.len());
    println!("  Duration: {}s", session.duration_secs());

    // First commit to finalize any pending Tantivy/index data
    // This ensures data_end is accurate before saving replay segment
    mem.commit()?;

    // Now save the replay sessions after all other data is finalized
    mem.save_replay_sessions()?;

    // Commit again to persist the replay manifest in TOC
    mem.commit()?;

    // Remove the sidecar file
    mem.clear_active_session_file()?;

    println!("Session saved to file.");
    Ok(())
}

#[cfg(feature = "replay")]
fn handle_session_list(args: SessionListArgs) -> Result<()> {
    let mut mem = Memvid::open(&args.file)?;

    // Check for active session
    let has_active = mem.load_active_session()?;

    // Load completed sessions from file
    mem.load_replay_sessions()?;

    let sessions = mem.list_sessions();

    if !has_active && sessions.is_empty() {
        println!("No recorded sessions.");
        return Ok(());
    }

    if args.format == "json" {
        let json = serde_json::to_string_pretty(&sessions)?;
        println!("{}", json);
    } else {
        // Show active session first if present
        if has_active {
            if let Some(session_id) = mem.active_session_id() {
                println!("Active session (recording):");
                println!("  {} - currently recording", session_id);
                println!();
            }
        }

        if !sessions.is_empty() {
            println!("Completed sessions ({}):", sessions.len());
            for summary in sessions {
                let name = summary.name.as_deref().unwrap_or("(unnamed)");
                let status = if summary.ended_secs.is_some() {
                    "completed"
                } else {
                    "recording"
                };
                println!(
                    "  {} - {} [{}, {} actions, {} checkpoints]",
                    summary.session_id, name, status, summary.action_count, summary.checkpoint_count
                );
            }
        }
    }

    Ok(())
}

#[cfg(feature = "replay")]
fn handle_session_view(args: SessionViewArgs) -> Result<()> {
    let mut mem = Memvid::open(&args.file)?;

    // Load sessions from file
    mem.load_replay_sessions()?;

    let session_id: uuid::Uuid = args.session.parse()?;
    let session = mem
        .get_session(session_id)
        .ok_or_else(|| anyhow::anyhow!("Session {} not found", session_id))?;

    if args.format == "json" {
        let json = serde_json::to_string_pretty(&session)?;
        println!("{}", json);
    } else {
        let name = session.name.as_deref().unwrap_or("(unnamed)");
        println!("Session: {} - {}", session.session_id, name);
        println!("  Created: {}", session.created_secs);
        if let Some(ended) = session.ended_secs {
            println!("  Ended: {}", ended);
            println!("  Duration: {}s", session.duration_secs());
        } else {
            println!("  Status: recording");
        }
        println!("  Actions: {}", session.actions.len());
        println!("  Checkpoints: {}", session.checkpoints.len());

        if !session.actions.is_empty() {
            println!("\nActions:");
            for (i, action) in session.actions.iter().enumerate().take(20) {
                println!(
                    "  [{}] {} - {:?}",
                    i, action.action_type.name(), action.action_type
                );
            }
            if session.actions.len() > 20 {
                println!("  ... and {} more", session.actions.len() - 20);
            }
        }
    }

    Ok(())
}

#[cfg(feature = "replay")]
fn handle_session_checkpoint(args: SessionCheckpointArgs) -> Result<()> {
    let mut mem = Memvid::open(&args.file)?;

    let checkpoint_id = mem.create_checkpoint()?;
    println!("Created checkpoint: {}", checkpoint_id);

    Ok(())
}

#[cfg(feature = "replay")]
fn handle_session_delete(args: SessionDeleteArgs) -> Result<()> {
    let mut mem = Memvid::open(&args.file)?;

    // Load sessions from file
    mem.load_replay_sessions()?;

    let session_id: uuid::Uuid = args.session.parse()?;
    mem.delete_session(session_id)?;

    // Save updated sessions
    mem.save_replay_sessions()?;
    mem.commit()?;

    println!("Deleted session: {}", session_id);
    Ok(())
}

#[cfg(feature = "replay")]
fn handle_session_save(args: SessionSaveArgs) -> Result<()> {
    let mut mem = Memvid::open(&args.file)?;

    // Load any existing sessions from file first
    mem.load_replay_sessions()?;

    // Check if there are any sessions to save
    let session_count = mem.list_sessions().len();
    if session_count == 0 {
        println!("No sessions to save. Sessions are automatically saved when you run `session end`.");
        return Ok(());
    }

    // Re-save all sessions (useful if you added sessions via API or need to re-persist)
    mem.save_replay_sessions()?;
    mem.commit()?;

    println!("Saved {} session(s) to file.", session_count);
    Ok(())
}

#[cfg(feature = "replay")]
fn handle_session_load(args: SessionLoadArgs) -> Result<()> {
    let mut mem = Memvid::open(&args.file)?;

    mem.load_replay_sessions()?;

    let sessions = mem.list_sessions();
    println!("Loaded {} sessions from file.", sessions.len());

    Ok(())
}

#[cfg(feature = "replay")]
fn handle_session_replay(args: SessionReplayArgs) -> Result<()> {
    // Check if web UI is requested
    #[cfg(feature = "web")]
    if args.web {
        use crate::commands::session_web::start_web_server;
        use std::path::PathBuf;

        let rt = tokio::runtime::Runtime::new()?;
        return rt.block_on(start_web_server(
            args.session.clone(),
            PathBuf::from(&args.file),
            args.port,
            !args.no_open,
        ));
    }

    use memvid_core::replay::{ReplayEngine, ReplayExecutionConfig};

    // Use open_read_only to match how find command opens files
    // This ensures lex_enabled is properly set from has_lex_index()
    let mut mem = Memvid::open_read_only(&args.file)?;

    // Load sessions from file
    mem.load_replay_sessions()?;

    let session_id: uuid::Uuid = args.session.parse()?;
    let session = mem
        .get_session(session_id)
        .ok_or_else(|| anyhow::anyhow!("Session {} not found", session_id))?
        .clone();

    let config = ReplayExecutionConfig {
        skip_puts: args.skip_mutations,
        skip_finds: args.skip_finds,
        skip_asks: args.skip_asks,
        stop_on_mismatch: args.stop_on_mismatch,
        verbose: true,
        top_k: args.top_k,
        adaptive: args.adaptive,
        min_relevancy: args.min_relevancy,
    };

    let mut engine = ReplayEngine::new(&mut mem, config);
    let result = engine.replay_session_from(&session, args.from_checkpoint)?;

    if args.format == "json" {
        let json = serde_json::to_string_pretty(&result)?;
        println!("{}", json);
    } else {
        println!("{} {}", "Replay Results for Session:".bold(), session_id.to_string().cyan());
        println!("  Total actions: {}", result.total_actions.to_string().white());
        println!("  Matched: {}", result.matched_actions.to_string().green());
        println!("  Mismatched: {}", result.mismatched_actions.to_string().red());
        println!("  Skipped: {}", result.skipped_actions.to_string().yellow());
        println!("  Match rate: {}", format!("{:.1}%", result.match_rate()).cyan());
        println!("  Duration: {}ms", result.total_duration_ms);

        if let Some(cp) = result.from_checkpoint {
            println!("  Started from checkpoint: {}", cp);
        }

        if result.is_success() {
            println!("\n{}", "✓ Replay successful - all actions matched".green().bold());
        } else {
            println!("\n{}", "✗ Replay found mismatches:".red().bold());
            for action_result in result
                .action_results
                .iter()
                .filter(|r| !r.matched && r.diff.as_ref().map_or(true, |d| d != "skipped"))
            {
                let diff = action_result.diff.as_deref().unwrap_or("unknown");
                // Color the diff based on type
                let colored_diff = if diff.starts_with("FILTER:") {
                    diff.replace("FILTER:", &"FILTER:".red().bold().to_string())
                        .replace("would be MISSED", &"would be MISSED".red().to_string())
                } else if diff.starts_with("DISCOVERY:") {
                    diff.replace("DISCOVERY:", &"DISCOVERY:".green().bold().to_string())
                        .replace("[NEW]", &"[NEW]".green().to_string())
                } else {
                    diff.to_string()
                };
                println!(
                    "  [{}] {} - {}",
                    action_result.sequence.to_string().yellow(),
                    action_result.action_type.cyan(),
                    colored_diff
                );
            }
        }
    }

    Ok(())
}

#[cfg(feature = "replay")]
fn handle_session_compare(args: SessionCompareArgs) -> Result<()> {
    use memvid_core::replay::ReplayEngine;

    let mut mem = Memvid::open(&args.file)?;

    // Load sessions from file
    mem.load_replay_sessions()?;

    let session_a_id: uuid::Uuid = args.session_a.parse()?;
    let session_b_id: uuid::Uuid = args.session_b.parse()?;

    let session_a = mem
        .get_session(session_a_id)
        .ok_or_else(|| anyhow::anyhow!("Session {} not found", session_a_id))?
        .clone();

    let session_b = mem
        .get_session(session_b_id)
        .ok_or_else(|| anyhow::anyhow!("Session {} not found", session_b_id))?
        .clone();

    let comparison = ReplayEngine::compare_sessions(&session_a, &session_b);

    if args.format == "json" {
        let json = serde_json::to_string_pretty(&comparison)?;
        println!("{}", json);
    } else {
        println!("Session Comparison:");
        println!("  Session A: {}", session_a_id);
        println!("  Session B: {}", session_b_id);
        println!();

        if comparison.is_identical() {
            println!("✓ Sessions are identical");
            println!("  Matching actions: {}", comparison.matching_actions);
        } else {
            println!("✗ Sessions differ:");
            println!("  Matching actions: {}", comparison.matching_actions);

            if !comparison.actions_only_in_a.is_empty() {
                println!(
                    "  Actions only in A: {:?}",
                    comparison.actions_only_in_a
                );
            }

            if !comparison.actions_only_in_b.is_empty() {
                println!(
                    "  Actions only in B: {:?}",
                    comparison.actions_only_in_b
                );
            }

            if !comparison.differing_actions.is_empty() {
                println!("  Differing actions:");
                for diff in &comparison.differing_actions {
                    println!(
                        "    [{}] {} vs {} - {}",
                        diff.sequence, diff.action_type_a, diff.action_type_b, diff.description
                    );
                }
            }
        }
    }

    Ok(())
}