mcp-memory/server/src/handlers/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::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).map_err(|e| e.to_string())?;
            Ok(data)
        } 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::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).map_err(|e| e.to_string())
        })?;
        Ok(data)
    }
}

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

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

        let update_handler = UpdateEnvFingerprintHandler;
        let args = json!({
            "namespace": "global",
            "tool_versions": {
                "rustc": "1.70.0"
            }
        });

        let res = update_handler.execute(args, state.clone()).await.unwrap();
        assert_eq!(res, "Env fingerprint updated");

        let read_handler = ReadEnvFingerprintHandler;
        let res2 = read_handler
            .execute(json!({"namespace": "global"}), state.clone())
            .await
            .unwrap();
        assert!(res2.contains("rustc"));
        assert!(res2.contains("1.70.0"));
    }

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

        // Ensure namespace is present in test setup
        state.environments.modify(|e| {
            e.push(crate::models::EnvironmentDetail {
                namespace: "global".to_string(),
                name: "test".to_string(),
                url: "http://localhost".to_string(),
                description: "Test env".to_string(),
                requires_vpn: false,
                updated_at: crate::handlers::utils::now_secs(),
            });
        });

        let handler = GetEnvironmentDetailsHandler;
        let res = handler
            .execute(json!({"namespace": "global"}), state.clone())
            .await
            .unwrap();

        assert!(res.contains("global"));
    }

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

        let req_handler = LogEnvRequirementHandler;
        let args_req = json!({
            "key": "vpn",
            "description": "Requires VPN",
            "is_secret": false,
            "requirement": "Needs specific VPN",
            "context": "For database access",
            "namespace": "global"
        });
        let res1 = req_handler.execute(args_req, state.clone()).await.unwrap();
        assert_eq!(res1, "Env requirement logged");

        let reg_handler = RegisterEnvironmentHandler;
        let args_reg = json!({
            "name": "prod",
            "url": "https://prod.local",
            "description": "Production",
            "requires_vpn": true,
            "namespace": "global"
        });
        let res2 = reg_handler.execute(args_reg, state.clone()).await.unwrap();
        assert_eq!(res2, "Environment registered");

        let get_handler = GetEnvironmentDetailsHandler;
        let res3 = get_handler
            .execute(json!({"namespace": "global"}), state.clone())
            .await
            .unwrap();
        assert!(res3.contains("prod.local"));
        assert!(!res3.is_empty());
    }
}