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 EnvironmentHandler;

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

    fn schema(&self) -> Value {
        crate::mcp::tool_def::<EnvironmentTool>(
            "environment",
            "Consolidated environment & infrastructure management",
        )
    }

    async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
        let req: EnvironmentTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
        let ns = req
            .namespace
            .unwrap_or_else(crate::models::default_namespace);
        match req.action {
            EnvAction::UpdateFingerprint => {
                let tool_versions = req.tool_versions.unwrap_or_default();
                let detected_os = req.os.unwrap_or_else(|| std::env::consts::OS.to_string());
                let detected_shell = req
                    .shell
                    .or_else(|| std::env::var("SHELL").ok())
                    .or_else(|| std::env::var("STARSHIP_SHELL").ok())
                    .or_else(|| {
                        if std::env::var("PSModulePath").is_ok() {
                            Some("pwsh".to_string())
                        } else {
                            None
                        }
                    })
                    .or_else(|| {
                        if std::env::var("COMSPEC").is_ok() {
                            Some("cmd".to_string())
                        } else {
                            None
                        }
                    })
                    .unwrap_or_else(|| "unknown".to_string());

                state.env.env_fingerprints.modify(|fps| {
                    fps.insert(
                        ns.clone(),
                        crate::models::EnvFingerprint {
                            namespace: ns.clone(),
                            os: detected_os,
                            shell: detected_shell,
                            tool_versions,
                            updated_at: crate::handlers::utils::now_secs(),
                            ..Default::default()
                        },
                    );
                });
                Ok("Env fingerprint updated".to_string())
            }
            EnvAction::ReadFingerprint => {
                let data = state
                    .env
                    .env_fingerprints
                    .read_with(|fps| fps.get(&ns).cloned());
                if let Some(fp) = data {
                    let data = Ok::<String, crate::error::AppError>(serde_yaml::to_string(&fp)?)?;
                    Ok(data)
                } else {
                    Ok("{}".to_string())
                }
            }
            EnvAction::LogRequirement => {
                let key = req.key.ok_or_else(|| {
                    crate::error::AppError::Internal("Missing required parameter 'key' for action 'log_requirement'. Next step: Provide environment variable 'key' string in request and retry.".to_string())
                })?;
                let desc = req.description.unwrap_or_default();
                let secret = req.is_secret.unwrap_or(false);
                state.env.env_requirements.modify(|reqs| {
                    reqs.retain(|r| !(r.namespace == ns && r.key == key));
                    reqs.push(crate::models::EnvRequirement {
                        namespace: ns,
                        key,
                        description: desc,
                        is_secret: secret,
                        ..Default::default()
                    });
                });
                Ok("Env requirement logged".to_string())
            }
            EnvAction::Register => {
                let name = req.name.ok_or_else(|| {
                    crate::error::AppError::Internal("Missing required parameter 'name' for action 'register'. Next step: Provide environment 'name' string in request and retry.".to_string())
                })?;
                let url = req.url.unwrap_or_default();
                let desc = req.description.unwrap_or_default();
                let vpn = req.requires_vpn.unwrap_or(false);
                state.env.environments.modify(|envs| {
                    envs.retain(|e| !(e.namespace == ns && e.name == name));
                    envs.push(crate::models::EnvironmentDetail {
                        namespace: ns,
                        name,
                        url,
                        description: desc,
                        requires_vpn: vpn,
                        updated_at: crate::handlers::utils::now_secs(),
                        ..Default::default()
                    });
                });
                Ok("Environment registered".to_string())
            }
            EnvAction::GetDetails => {
                let data = state.env.environments.read_with(|envs| {
                    let filtered: Vec<_> = envs.iter().filter(|e| e.namespace == ns).collect();
                    Ok::<String, crate::error::AppError>(serde_yaml::to_string(&filtered)?)
                })?;
                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 handler = EnvironmentHandler;

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

        let res = handler
            .execute(args, state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert_eq!(res, "Env fingerprint updated");

        let res2 = handler
            .execute(
                json!({"action": "read_fingerprint", "namespace": "global"}),
                state.clone(),
            )
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .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.env.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(),
                ..Default::default()
            });
        });

        let handler = EnvironmentHandler;
        let res = handler
            .execute(
                json!({"action": "get_details", "namespace": "global"}),
                state.clone(),
            )
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .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 handler = EnvironmentHandler;

        let args_req = json!({
            "action": "log_requirement",
            "key": "vpn",
            "description": "Requires VPN",
            "is_secret": false,
            "namespace": "global"
        });
        let res1 = handler
            .execute(args_req, state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert_eq!(res1, "Env requirement logged");

        let args_reg = json!({
            "action": "register",
            "name": "prod",
            "url": "https://prod.local",
            "description": "Production",
            "requires_vpn": true,
            "namespace": "global"
        });
        let res2 = handler
            .execute(args_reg, state.clone())
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert_eq!(res2, "Environment registered");

        let res3 = handler
            .execute(
                json!({"action": "get_details", "namespace": "global"}),
                state.clone(),
            )
            .await
            .map_err(|e| crate::error::AppError::Internal(e.to_string()))
            .unwrap();
        assert!(res3.contains("prod.local"));
        assert!(!res3.is_empty());
    }

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

        let args = serde_json::json!({
            "action": "update_fingerprint",
            "namespace": "",
            "tool_versions": {
                "rustc": "1.70.0"
            }
        });

        let res = handler.execute(args, state.clone()).await;

        assert!(res.is_ok());
    }
}