mcp-memory/server/src/handlers/notes.rs

use crate::models::*;
use crate::router::McpTool;
use crate::state::MemoryState;
use crate::tools::*;
use async_trait::async_trait;
use serde_json::Value;
use std::collections::HashSet;
use std::sync::Arc;

pub struct HandoffMemosHandler;

#[async_trait]
impl McpTool for HandoffMemosHandler {
    fn name(&self) -> &'static str {
        "handoff_memos"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<HandoffMemosTool>(
            "handoff_memos",
            "Handoff memos across sessions (leave, read, clear)",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: HandoffMemosTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        match req.action {
            HandoffMemoAction::Leave => {
                let content = req.content.ok_or_else(|| {
                    crate::error::AppError::Internal("Missing required parameter 'content' for action 'leave'. Next step: Provide non-empty 'content' string in request and retry.".to_string())
                })?;
                let ns = req
                    .namespace
                    .unwrap_or_else(crate::models::default_namespace);
                state.telemetry.handoff_memos.modify(|memos| {
                    memos.push(crate::models::HandoffMemo {
                        id: uuid::Uuid::new_v4().to_string(),
                        author: "agy".to_string(),
                        content,
                        namespace: ns,
                        timestamp: crate::handlers::utils::now_secs(),
                        expires_at: None,
                        ..Default::default()
                    });
                    if memos.len() > 200 {
                        memos.remove(0);
                    }
                });
                Ok("Handoff memo left.".to_string())
            }
            HandoffMemoAction::Read => {
                let data = state.telemetry.handoff_memos.read_with(|items| {
                    let filtered: Vec<_> = items
                        .iter()
                        .filter(|i| {
                            if let Some(ns) = &req.namespace {
                                &i.namespace == ns
                            } else {
                                true
                            }
                        })
                        .collect();
                    Ok::<String, crate::error::AppError>(serde_yaml::to_string(&filtered)?)
                })?;
                Ok(data)
            }
            HandoffMemoAction::Clear => {
                let ids: HashSet<_> = req.ids.unwrap_or_default().into_iter().collect();
                state.telemetry.handoff_memos.modify(|memos| {
                    if ids.is_empty() {
                        memos.clear();
                    } else {
                        memos.retain(|m| !ids.contains(&m.id));
                    }
                });
                Ok("Handoff memos cleared.".to_string())
            }
        }
    }
}

pub struct AddSessionSummaryHandler;

#[async_trait]
impl McpTool for AddSessionSummaryHandler {
    fn name(&self) -> &'static str {
        "add_session_summary"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<AddSessionSummaryTool>(
            "add_session_summary",
            "Execute add_session_summary",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: AddSessionSummaryTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.telemetry.session_summaries.modify(|summaries| {
            summaries.push(crate::models::SessionSummary {
                summary: req.summary,
                namespace: req.namespace,
                timestamp: crate::handlers::utils::now_secs(),
                expires_at: None,
                session_id: req.session_id,
                repo_name: req.repo_name,
                git_branch: req.git_branch,
                git_commit: req.git_commit,
                tasks_completed: req.tasks_completed.unwrap_or_default(),
                next_steps: req.next_steps.unwrap_or_default(),
            });
            if summaries.len() > 200 {
                let now = crate::handlers::utils::now_secs();
                if let Some((idx_to_remove, _)) =
                    summaries.iter().enumerate().min_by_key(|(_, s)| {
                        let is_expired = s.expires_at.map(|exp| exp < now).unwrap_or(false);
                        let status_score = if is_expired { 0 } else { 1 };
                        (status_score, s.timestamp)
                    })
                {
                    summaries.remove(idx_to_remove);
                }
            }
        });
        Ok("Session summary added".to_string())
    }
}

pub struct GenerateStandupReportHandler;

#[async_trait]
impl McpTool for GenerateStandupReportHandler {
    fn name(&self) -> &'static str {
        "generate_standup_report"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<GenerateStandupReportTool>(
            "generate_standup_report",
            "Execute generate_standup_report",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: GenerateStandupReportTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        let cutoff = crate::handlers::utils::now_secs().saturating_sub(req.hours_lookback * 3600);

        let filtered_tasks = state.project.tasks.read_with(|items| {
            items
                .iter()
                .filter(|t| t.updated_at >= cutoff)
                .cloned()
                .collect::<Vec<_>>()
        });
        let filtered_changes = state.code.ledger.read_with(|changes| {
            changes
                .iter()
                .filter(|c| c.timestamp >= cutoff)
                .cloned()
                .collect::<Vec<_>>()
        });
        let filtered_summaries = state.telemetry.session_summaries.read_with(|summaries| {
            summaries
                .iter()
                .filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff)
                .cloned()
                .collect::<Vec<_>>()
        });

        let report = serde_json::json!({
            "tasks_updated": filtered_tasks,
            "code_changes": filtered_changes,
            "session_summaries": filtered_summaries
        });
        let report_str = report.to_string();
        Ok(report_str)
    }
}

pub struct PromoteToEntityHandler;

#[async_trait]
impl McpTool for PromoteToEntityHandler {
    fn name(&self) -> &'static str {
        "promote_to_entity"
    }

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<PromoteToEntityTool>(
            "promote_to_entity",
            "Promote a task observation or finding into a permanent Knowledge Graph entity.",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: PromoteToEntityTool = serde_json::from_value(args).map_err(|e| e.to_string())?;

        let entity_name = req.entity_name.clone();
        let obs = req.content.clone();
        let entity_type = req.entity_type.clone();
        let namespace = req.namespace.clone();

        state.modify_graph(|g| {
            let entity = g
                .entities
                .entry(entity_name.clone())
                .or_insert_with(|| Entity {
                    name: entity_name.clone(),
                    entity_type,
                    observations: Vec::new(),
                    namespace,
                    git_branch: None,
                    ..Default::default()
                });
            if !entity.observations.contains(&obs) {
                entity.observations.push(obs.clone());
            }
        });

        Ok(format!(
            "Promoted transient note into permanent entity '{}'.",
            req.entity_name
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;
    use tempfile::tempdir;

    #[tokio::test]
    async fn test_handoff_and_summaries() {
        let dir = tempdir().unwrap();
        let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));

        let handoff_handler = HandoffMemosHandler;
        let args = json!({
            "action": "leave",
            "content": "Finished implementing graph tests",
            "namespace": "global"
        });

        let res = handoff_handler
            .execute(args, state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert_eq!(res, "Handoff memo left.");

        let res2 = handoff_handler
            .execute(
                json!({"action": "read", "namespace": "global"}),
                state.clone(),
            )
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert!(res2.contains("Finished implementing graph tests"));

        let summary_handler = AddSessionSummaryHandler;
        let args_sum = json!({
            "summary": "Completed a bunch of tests",
            "namespace": "global"
        });
        let res3 = summary_handler
            .execute(args_sum, state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert_eq!(res3, "Session summary added");

        let standup_handler = GenerateStandupReportHandler;
        let res4 = standup_handler
            .execute(
                json!({"namespace": "global", "hours_lookback": 24}),
                state.clone(),
            )
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert!(!res4.is_empty());
    }
}