post_cortex_mcp/

session.rs

//! Session lifecycle: create, load, checkpoint, list, search, and metadata.

use crate::{MCPToolResult, get_memory_system, string_to_anyhow};
use anyhow::Result;
use post_cortex_core::core::timeout_utils::with_storage_timeout;
use post_cortex_core::session::active_session::ActiveSession;
use post_cortex_memory::ConversationMemorySystem;
use post_cortex_storage::rocksdb_storage::SessionCheckpoint;
use std::sync::Arc;
use tracing::{debug, error, instrument};
use uuid::Uuid;

/// Create a session checkpoint using an explicit memory system reference.
pub async fn create_session_checkpoint_with_system(
    session_id: Uuid,
    system: &ConversationMemorySystem,
) -> Result<MCPToolResult> {
    let session_arc = system
        .get_session(session_id)
        .await
        .map_err(|e| anyhow::anyhow!("Failed to load session: {}", e))?;
    let session = session_arc.load();

    let checkpoint = create_comprehensive_checkpoint(&session).await?;

    system
        .storage_actor
        .save_checkpoint(&checkpoint)
        .await
        .map_err(string_to_anyhow)?;

    Ok(MCPToolResult::success(
        "Checkpoint created successfully".to_string(),
        Some(serde_json::json!({ "checkpoint_id": checkpoint.id.to_string() })),
    ))
}

/// Load a session from a previously saved checkpoint.
pub async fn load_session_checkpoint_with_system(
    checkpoint_id: String,
    session_id: Uuid,
    system: &ConversationMemorySystem,
) -> Result<MCPToolResult> {
    eprintln!("Loading checkpoint - step 1: Parsing checkpoint ID");
    let checkpoint_id = Uuid::parse_str(&checkpoint_id)?;

    eprintln!("Loading checkpoint - step 2: Loading checkpoint from storage");
    let checkpoint = system
        .storage_actor
        .load_checkpoint(checkpoint_id)
        .await
        .map_err(string_to_anyhow)?;

    eprintln!("Loading checkpoint - step 3: Checkpoint loaded successfully");

    eprintln!("Loading checkpoint - step 4: Creating session from checkpoint");
    let mut session = ActiveSession::new(session_id, None, None);

    eprintln!("Loading checkpoint - step 5: Restoring current state");
    session.current_state = Arc::new(checkpoint.structured_context);

    eprintln!("Loading checkpoint - step 6: Restoring incremental updates");
    session.set_incremental_updates(checkpoint.recent_updates);

    eprintln!("Loading checkpoint - step 10: Restoring code references");
    session.code_references = Arc::new(checkpoint.code_references);

    eprintln!("Loading checkpoint - step 11: Restoring change history");
    session.set_change_history(checkpoint.change_history);

    eprintln!("Loading checkpoint - step 12: Entity graph restored");

    eprintln!("Loading checkpoint - step 13: Adding session to session manager");
    system.session_manager.install(session_id, session);

    eprintln!("Loading checkpoint - step 14: Updated session manager");

    Ok(MCPToolResult::success(
        "Session loaded from checkpoint successfully".to_string(),
        None,
    ))
}

/// Create a session checkpoint using the global memory system.
pub async fn create_session_checkpoint(session_id: Uuid) -> Result<MCPToolResult> {
    let result = with_storage_timeout(async {
        let system = get_memory_system().await?;
        let session_arc = system
            .get_session(session_id)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to load session: {}", e))?;
        let session = session_arc.load();

        let checkpoint = create_comprehensive_checkpoint(&session).await?;

        system
            .storage_actor
            .save_checkpoint(&checkpoint)
            .await
            .map_err(string_to_anyhow)?;

        Ok(MCPToolResult::success(
            "Checkpoint created successfully".to_string(),
            Some(serde_json::json!({ "checkpoint_id": checkpoint.id.to_string() })),
        ))
    })
    .await;

    match result {
        Ok(success_result) => success_result,
        Err(timeout_error) => {
            error!(
                "TIMEOUT: create_session_checkpoint - session: {}, error: {}",
                session_id, timeout_error
            );
            Ok(MCPToolResult::error(format!(
                "Checkpoint creation timed out: {}",
                timeout_error
            )))
        }
    }
}

/// Load a session checkpoint using the global memory system.
pub async fn load_session_checkpoint(
    checkpoint_id: String,
    session_id: Uuid,
) -> Result<MCPToolResult> {
    let result = with_storage_timeout(async {
        let system = get_memory_system().await?;
        let checkpoint_id = Uuid::parse_str(&checkpoint_id)?;

        let checkpoint = system
            .storage_actor
            .load_checkpoint(checkpoint_id)
            .await
            .map_err(string_to_anyhow)?;

        let mut session = ActiveSession::new(session_id, None, None);
        session.current_state = Arc::new(checkpoint.structured_context);
        session.set_incremental_updates(checkpoint.recent_updates);
        session.code_references = Arc::new(checkpoint.code_references);
        session.set_change_history(checkpoint.change_history);

        system.session_manager.install(session_id, session);

        Ok(MCPToolResult::success(
            "Session loaded from checkpoint successfully".to_string(),
            None,
        ))
    })
    .await;

    match result {
        Ok(success_result) => success_result,
        Err(timeout_error) => {
            error!(
                "TIMEOUT: load_session_checkpoint - session: {}, error: {}",
                session_id, timeout_error
            );
            Ok(MCPToolResult::error(format!(
                "Checkpoint loading timed out: {}",
                timeout_error
            )))
        }
    }
}

/// Mark a specific context update as important within a session.
pub async fn mark_important(session_id: Uuid, update_id: String) -> Result<MCPToolResult> {
    let update_id = Uuid::parse_str(&update_id)?;
    let system = get_memory_system().await?;
    let session_arc = system
        .get_session(session_id)
        .await
        .map_err(|e| anyhow::anyhow!("Failed to load session: {}", e))?;

    let mut found = false;
    session_arc.rcu(|current| {
        let mut updated = (**current).clone();
        let updates = Arc::make_mut(&mut updated.incremental_updates);
        for update in updates.iter_mut() {
            if update.id == update_id {
                Arc::make_mut(update).user_marked_important = true;
                found = true;
                break;
            }
        }
        Arc::new(updated)
    });

    if found {
        Ok(MCPToolResult::success(
            "Update marked as important".to_string(),
            None,
        ))
    } else {
        Ok(MCPToolResult::error("Update not found".to_string()))
    }
}

/// List all persisted sessions using a direct storage reference.
pub async fn list_sessions_with_storage(
    storage: &post_cortex_storage::rocksdb_storage::RealRocksDBStorage,
) -> Result<MCPToolResult> {
    // Metadata records only: listing never loads a session.
    match storage.list_session_summaries().await {
        Ok(summaries) => {
            let sessions_info: Vec<_> = summaries
                .iter()
                .map(|summary| summary_json(summary, None))
                .collect();
            Ok(MCPToolResult::success(
                format!("Found {} sessions", sessions_info.len()),
                Some(serde_json::json!({
                    "sessions": sessions_info
                })),
            ))
        }
        Err(e) => Ok(MCPToolResult::error(format!(
            "Failed to load sessions: {e}"
        ))),
    }
}

/// Listing row for a session; `entity_count` is null when unknown.
fn summary_json(
    summary: &post_cortex_storage::traits::SessionSummary,
    entity_count: Option<usize>,
) -> serde_json::Value {
    serde_json::json!({
        "id": summary.id.to_string(),
        "name": summary.name,
        "description": summary.description,
        "created_at": summary.created_at.to_rfc3339(),
        "last_updated": summary.last_updated.to_rfc3339(),
        "update_count": summary.update_count,
        "entity_count": entity_count
    })
}

/// List all sessions via the global memory system.
pub async fn list_sessions() -> Result<MCPToolResult> {
    debug!("MCP-TOOLS: list_sessions() called");
    let result = with_storage_timeout(async {
        debug!("MCP-TOOLS: Getting memory system for list_sessions");
        let system = get_memory_system().await?;
        debug!("MCP-TOOLS: Got memory system, listing sessions");
        // Metadata records only: listing neither loads sessions nor fills
        // the session cache. Entity counts are only known for sessions
        // that are already in memory.
        let summaries = system
            .list_session_summaries()
            .await
            .map_err(string_to_anyhow)?;
        let sessions_info: Vec<_> = summaries
            .iter()
            .map(|summary| {
                let entity_count = system
                    .session_manager
                    .sessions
                    .get(&summary.id)
                    .map(|s| s.load().entity_graph.entities.len());
                summary_json(summary, entity_count)
            })
            .collect();

        Ok(MCPToolResult::success(
            format!("Found {} sessions", sessions_info.len()),
            Some(serde_json::json!({
                "sessions": sessions_info
            })),
        ))
    })
    .await;

    match result {
        Ok(success_result) => success_result,
        Err(timeout_error) => {
            error!("TIMEOUT: list_sessions - error: {timeout_error}");
            Ok(MCPToolResult::error(format!(
                "Session listing timed out: {timeout_error}"
            )))
        }
    }
}

/// Load a session by ID using an explicit memory system reference.
pub async fn load_session_with_system(
    session_id: Uuid,
    system: &ConversationMemorySystem,
) -> Result<MCPToolResult> {
    match system.get_session(session_id).await {
        Ok(session_arc) => {
            let session = session_arc.load();
            Ok(MCPToolResult::success(
                "Session loaded successfully".to_string(),
                Some(serde_json::json!({
                    "session": {
                        "id": session.id().to_string(),
                        "created_at": session.created_at().to_rfc3339(),
                        "last_updated": session.last_updated.to_rfc3339(),
                        "update_count": session.incremental_updates.len(),
                        "entity_count": session.entity_graph.entities.len(),
                        "hot_context_size": session.hot_context.len(),
                        "warm_context_size": session.warm_context.len(),
                        "cold_context_size": session.cold_context.len(),
                        "code_references": session.code_references.keys().collect::<Vec<_>>(),
                        "change_history_count": session.change_history.len()
                    }
                })),
            ))
        }
        Err(e) => Ok(MCPToolResult::error(
            format!("Failed to load session: {e}",),
        )),
    }
}

/// Load a session by ID via the global memory system.
pub async fn load_session(session_id: Uuid) -> Result<MCPToolResult> {
    let result = with_storage_timeout(async {
        let system = get_memory_system().await?;
        load_session_with_system(session_id, &system).await
    })
    .await;

    match result {
        Ok(success_result) => success_result,
        Err(timeout_error) => {
            error!(
                "TIMEOUT: load_session - session: {}, error: {}",
                session_id, timeout_error
            );
            Ok(MCPToolResult::error(format!(
                "Session loading timed out: {}",
                timeout_error
            )))
        }
    }
}

/// Search sessions by name or description.
pub async fn search_sessions(query: String) -> Result<MCPToolResult> {
    let result = with_storage_timeout(async {
        let system = get_memory_system().await?;
        let sessions: Vec<_> = system
            .find_session_summaries(&query)
            .await
            .map_err(string_to_anyhow)?
            .into_iter()
            .map(|s| {
                serde_json::json!({
                    "id": s.id.to_string(),
                    "name": s.name,
                    "description": s.description
                })
            })
            .collect();

        Ok(MCPToolResult::success(
            format!("Found {} sessions matching '{}'", sessions.len(), query),
            Some(serde_json::json!({
                "sessions": sessions
            })),
        ))
    })
    .await;

    match result {
        Ok(success_result) => success_result,
        Err(timeout_error) => {
            error!(
                "TIMEOUT: search_sessions - query: {}, error: {}",
                query, timeout_error
            );
            Ok(MCPToolResult::error(format!(
                "Session search timed out: {}",
                timeout_error
            )))
        }
    }
}

/// Update the name and/or description of a session.
#[instrument(skip(session_id), fields(session_id = %session_id))]
pub async fn update_session_metadata(
    session_id: Uuid,
    name: Option<String>,
    description: Option<String>,
) -> Result<MCPToolResult> {
    let result = with_storage_timeout(async {
        let system = get_memory_system().await?;
        system
            .update_session_metadata(session_id, name, description)
            .await
            .map_err(string_to_anyhow)?;

        let session_arc = system
            .get_session(session_id)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to load session: {}", e))?;
        let session = session_arc.load();
        let (final_name, final_description) = session.get_metadata();

        Ok(MCPToolResult::success(
            "Session metadata updated successfully".to_string(),
            Some(serde_json::json!({
                "session_id": session_id.to_string(),
                "name": final_name,
                "description": final_description
            })),
        ))
    })
    .await;

    match result {
        Ok(success_result) => success_result,
        Err(timeout_error) => {
            error!(
                "TIMEOUT: update_session_metadata - session_id: {}, error: {}",
                session_id, timeout_error
            );
            Ok(MCPToolResult::error(format!(
                "TIMEOUT: Failed to update session metadata: {}",
                timeout_error
            )))
        }
    }
}

/// Build a comprehensive checkpoint snapshot from an active session.
async fn create_comprehensive_checkpoint(session: &ActiveSession) -> Result<SessionCheckpoint> {
    Ok(SessionCheckpoint {
        id: Uuid::new_v4(),
        session_id: session.id(),
        created_at: chrono::Utc::now(),
        structured_context: (*session.current_state).clone(),
        recent_updates: session.incremental_updates_owned(),
        code_references: (*session.code_references).clone(),
        change_history: session.change_history_owned(),
        total_updates: session.incremental_updates.len(),
        context_quality_score: 1.0,
        compression_ratio: 1.0,
    })
}