enki-runtime 0.1.4

//! MCP tool configuration from TOML files.

use serde::{Deserialize, Serialize};

/// Transport type for MCP connections.
///
/// Specifies how to connect to an MCP server.
#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum McpTransportType {
    /// STDIO transport - spawns a child process and communicates via stdin/stdout
    #[default]
    Stdio,
    /// HTTP/SSE transport - connects to a remote MCP server
    Http,
}

/// Configuration for a single MCP tool source.
///
/// Defines how to connect to an MCP server and use its tools.
///
/// # Examples
///
/// ## STDIO Transport (npx package)
/// ```toml
/// [[agents.mcp_tools]]
/// type = "stdio"
/// command = "npx"
/// args = ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
/// ```
///
/// ## STDIO Transport (Python MCP server)
/// ```toml
/// [[agents.mcp_tools]]
/// type = "stdio"
/// command = "python"
/// args = ["-m", "mcp_server_time"]
/// ```
///
/// ## HTTP Transport
/// ```toml
/// [[agents.mcp_tools]]
/// type = "http"
/// url = "http://localhost:3000/mcp"
/// auth_token = "${MCP_AUTH_TOKEN}"
/// ```
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct McpToolConfig {
    /// Transport type: "stdio" or "http"
    #[serde(rename = "type", default)]
    pub transport_type: McpTransportType,

    /// Command to execute (for stdio transport).
    /// Examples: "npx", "python", "node", or path to an executable.
    #[serde(default)]
    pub command: Option<String>,

    /// Arguments for the command (for stdio transport).
    #[serde(default)]
    pub args: Vec<String>,

    /// Environment variables for the child process (for stdio transport).
    /// Format: [["KEY", "VALUE"], ["KEY2", "VALUE2"]]
    #[serde(default)]
    pub env: Option<Vec<(String, String)>>,

    /// URL for HTTP transport.
    /// Example: "http://localhost:3000/mcp"
    #[serde(default)]
    pub url: Option<String>,

    /// Authentication token for HTTP transport (optional).
    /// Can use environment variable syntax: "${MCP_AUTH_TOKEN}"
    #[serde(default)]
    pub auth_token: Option<String>,

    /// Optional name for this tool source (for logging/debugging).
    #[serde(default)]
    pub name: Option<String>,
}

impl McpToolConfig {
    /// Create a new STDIO transport configuration.
    ///
    /// # Arguments
    /// * `command` - Command to execute (e.g., "npx", "python")
    /// * `args` - Arguments to pass to the command
    pub fn stdio(command: impl Into<String>, args: Vec<String>) -> Self {
        Self {
            transport_type: McpTransportType::Stdio,
            command: Some(command.into()),
            args,
            ..Default::default()
        }
    }

    /// Create an npx-based MCP tool configuration.
    ///
    /// # Arguments
    /// * `package` - NPM package name (e.g., "@modelcontextprotocol/server-filesystem")
    /// * `extra_args` - Additional arguments after the package name
    pub fn npx(package: impl Into<String>, extra_args: Vec<String>) -> Self {
        let mut args = vec!["-y".to_string(), package.into()];
        args.extend(extra_args);
        Self::stdio("npx", args)
    }

    /// Create an HTTP transport configuration.
    ///
    /// # Arguments
    /// * `url` - URL of the MCP server
    pub fn http(url: impl Into<String>) -> Self {
        Self {
            transport_type: McpTransportType::Http,
            url: Some(url.into()),
            ..Default::default()
        }
    }

    /// Set the name for this tool source.
    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = Some(name.into());
        self
    }

    /// Set environment variables for the process.
    pub fn with_env(mut self, env: Vec<(String, String)>) -> Self {
        self.env = Some(env);
        self
    }

    /// Set authentication token for HTTP transport.
    pub fn with_auth_token(mut self, token: impl Into<String>) -> Self {
        self.auth_token = Some(token.into());
        self
    }

    /// Validate the configuration based on transport type.
    pub fn validate(&self) -> Result<(), String> {
        match self.transport_type {
            McpTransportType::Stdio => {
                if self.command.is_none() {
                    return Err("STDIO transport requires 'command' field".to_string());
                }
            }
            McpTransportType::Http => {
                if self.url.is_none() {
                    return Err("HTTP transport requires 'url' field".to_string());
                }
            }
        }
        Ok(())
    }
}

// Conditional implementation when mcp feature is enabled
#[cfg(feature = "mcp")]
impl McpToolConfig {
    /// Convert this configuration to an MCP transport for runtime use.
    ///
    /// This method is only available when the `mcp` feature is enabled.
    pub fn to_transport(&self) -> Result<enki_mcp::McpTransport, String> {
        self.validate()?;

        match self.transport_type {
            McpTransportType::Stdio => Ok(enki_mcp::McpTransport::Stdio {
                command: self.command.clone().unwrap(),
                args: self.args.clone(),
                env: self.env.clone(),
            }),
            McpTransportType::Http => Ok(enki_mcp::McpTransport::Http {
                url: self.url.clone().unwrap(),
                auth_token: self.auth_token.clone(),
            }),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_stdio_config() {
        let toml = r#"
            type = "stdio"
            command = "npx"
            args = ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
        "#;

        let config: McpToolConfig = toml::from_str(toml).unwrap();
        assert_eq!(config.transport_type, McpTransportType::Stdio);
        assert_eq!(config.command, Some("npx".to_string()));
        assert_eq!(config.args.len(), 3);
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_parse_http_config() {
        let toml = r#"
            type = "http"
            url = "http://localhost:3000/mcp"
            auth_token = "secret-token"
        "#;

        let config: McpToolConfig = toml::from_str(toml).unwrap();
        assert_eq!(config.transport_type, McpTransportType::Http);
        assert_eq!(config.url, Some("http://localhost:3000/mcp".to_string()));
        assert_eq!(config.auth_token, Some("secret-token".to_string()));
        assert!(config.validate().is_ok());
    }

    #[test]
    fn test_builder_methods() {
        let config = McpToolConfig::npx(
            "@modelcontextprotocol/server-filesystem",
            vec!["/home".into()],
        )
        .with_name("filesystem-tools");

        assert_eq!(config.transport_type, McpTransportType::Stdio);
        assert_eq!(config.command, Some("npx".to_string()));
        assert_eq!(
            config.args,
            vec!["-y", "@modelcontextprotocol/server-filesystem", "/home"]
        );
        assert_eq!(config.name, Some("filesystem-tools".to_string()));
    }

    #[test]
    fn test_validation_fails_without_command() {
        let config = McpToolConfig {
            transport_type: McpTransportType::Stdio,
            command: None,
            ..Default::default()
        };
        assert!(config.validate().is_err());
    }

    #[test]
    fn test_validation_fails_without_url() {
        let config = McpToolConfig {
            transport_type: McpTransportType::Http,
            url: None,
            ..Default::default()
        };
        assert!(config.validate().is_err());
    }
}