mcp-memory/server/src/handlers_v2/env.rs

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

pub struct UpdateEnvFingerprintHandler;

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

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

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: UpdateEnvFingerprintTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.env_fingerprints.modify(|fps| {
            fps.insert(
                req.namespace.clone(),
                crate::models::EnvFingerprint {
                    namespace: req.namespace.clone(),
                    os: std::env::consts::OS.to_string(),
                    shell: std::env::var("SHELL").unwrap_or_else(|_| "unknown".to_string()),
                    tool_versions: req.tool_versions,
                    updated_at: crate::handlers_v2::utils::now_secs(),
                },
            );
        });
        Ok("Env fingerprint updated".to_string())
    }
}

pub struct ReadEnvFingerprintHandler;

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

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

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: ReadEnvFingerprintTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state
            .env_fingerprints
            .read_with(|fps| fps.get(&req.namespace).cloned());
        if let Some(fp) = data {
            let data = serde_json::to_string(&fp).unwrap_or_default();
            Ok(data.to_string())
        } else {
            Ok("{}".to_string())
        }
    }
}

pub struct LogEnvRequirementHandler;

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

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

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: LogEnvRequirementTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.env_requirements.modify(|reqs| {
            reqs.retain(|r| !(r.namespace == req.namespace && r.key == req.key));
            reqs.push(crate::models::EnvRequirement {
                namespace: req.namespace,
                key: req.key,
                description: req.description,
                is_secret: req.is_secret,
            });
        });
        Ok("Env requirement logged".to_string())
    }
}

pub struct RegisterEnvironmentHandler;

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

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

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: RegisterEnvironmentTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        state.environments.modify(|envs| {
            envs.retain(|e| !(e.namespace == req.namespace && e.name == req.name));
            envs.push(crate::models::EnvironmentDetail {
                namespace: req.namespace,
                name: req.name,
                url: req.url,
                description: req.description,
                requires_vpn: req.requires_vpn,
                updated_at: crate::handlers_v2::utils::now_secs(),
            });
        });
        Ok("Environment registered".to_string())
    }
}

pub struct GetEnvironmentDetailsHandler;

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

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

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> Result<String, String> {
        let req: GetEnvironmentDetailsTool =
            serde_json::from_value(args).map_err(|e| e.to_string())?;
        let data = state.environments.read_with(|envs| {
            let filtered: Vec<_> = envs
                .iter()
                .filter(|e| e.namespace == req.namespace)
                .collect();
            serde_json::to_string(&filtered).unwrap_or_default()
        });
        Ok(data.to_string())
    }
}