lash-plugin-mcp 0.1.0-alpha.39

use std::collections::BTreeMap;
use std::path::PathBuf;
use std::time::Duration;

use serde::{Deserialize, Serialize};

use crate::error::McpError;

const DEFAULT_STARTUP_TIMEOUT_MS: u64 = 10_000;
const DEFAULT_CALL_TIMEOUT_MS: u64 = 60_000;

fn default_startup_timeout_ms() -> u64 {
    DEFAULT_STARTUP_TIMEOUT_MS
}

fn default_call_timeout_ms() -> u64 {
    DEFAULT_CALL_TIMEOUT_MS
}

fn is_default_startup_timeout_ms(value: &u64) -> bool {
    *value == DEFAULT_STARTUP_TIMEOUT_MS
}

fn is_default_call_timeout_ms(value: &u64) -> bool {
    *value == DEFAULT_CALL_TIMEOUT_MS
}

/// Connection configuration for one MCP server. Tag (`transport`) selects
/// the wire transport; per-variant fields configure that transport.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "transport", rename_all = "snake_case")]
pub enum McpServerConfig {
    /// Spawn a child process and speak JSON-RPC over stdio.
    Stdio {
        command: String,
        #[serde(default, skip_serializing_if = "Vec::is_empty")]
        args: Vec<String>,
        #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
        env: BTreeMap<String, String>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        cwd: Option<PathBuf>,
        #[serde(
            default = "default_startup_timeout_ms",
            skip_serializing_if = "is_default_startup_timeout_ms"
        )]
        startup_timeout_ms: u64,
        #[serde(
            default = "default_call_timeout_ms",
            skip_serializing_if = "is_default_call_timeout_ms"
        )]
        call_timeout_ms: u64,
    },
    /// Newer MCP spec HTTP/JSON streaming transport.
    StreamableHttp {
        url: String,
        #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
        headers: BTreeMap<String, String>,
        #[serde(
            default = "default_startup_timeout_ms",
            skip_serializing_if = "is_default_startup_timeout_ms"
        )]
        startup_timeout_ms: u64,
        #[serde(
            default = "default_call_timeout_ms",
            skip_serializing_if = "is_default_call_timeout_ms"
        )]
        call_timeout_ms: u64,
    },
    /// Older MCP spec HTTP+SSE transport.
    Sse {
        url: String,
        #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
        headers: BTreeMap<String, String>,
        #[serde(
            default = "default_startup_timeout_ms",
            skip_serializing_if = "is_default_startup_timeout_ms"
        )]
        startup_timeout_ms: u64,
        #[serde(
            default = "default_call_timeout_ms",
            skip_serializing_if = "is_default_call_timeout_ms"
        )]
        call_timeout_ms: u64,
    },
}

impl McpServerConfig {
    /// Convenience constructor for stdio servers.
    pub fn stdio(command: impl Into<String>, args: Vec<String>) -> Self {
        Self::Stdio {
            command: command.into(),
            args,
            env: BTreeMap::new(),
            cwd: None,
            startup_timeout_ms: default_startup_timeout_ms(),
            call_timeout_ms: default_call_timeout_ms(),
        }
    }

    /// Convenience constructor for streamable-HTTP servers.
    pub fn streamable_http(url: impl Into<String>) -> Self {
        Self::StreamableHttp {
            url: url.into(),
            headers: BTreeMap::new(),
            startup_timeout_ms: default_startup_timeout_ms(),
            call_timeout_ms: default_call_timeout_ms(),
        }
    }

    /// Convenience constructor for SSE servers.
    pub fn sse(url: impl Into<String>) -> Self {
        Self::Sse {
            url: url.into(),
            headers: BTreeMap::new(),
            startup_timeout_ms: default_startup_timeout_ms(),
            call_timeout_ms: default_call_timeout_ms(),
        }
    }

    pub fn startup_timeout(&self) -> Duration {
        Duration::from_millis(match self {
            Self::Stdio {
                startup_timeout_ms, ..
            }
            | Self::StreamableHttp {
                startup_timeout_ms, ..
            }
            | Self::Sse {
                startup_timeout_ms, ..
            } => *startup_timeout_ms,
        })
    }

    pub fn call_timeout(&self) -> Duration {
        Duration::from_millis(match self {
            Self::Stdio {
                call_timeout_ms, ..
            }
            | Self::StreamableHttp {
                call_timeout_ms, ..
            }
            | Self::Sse {
                call_timeout_ms, ..
            } => *call_timeout_ms,
        })
    }

    pub(crate) fn validate(&self, server_name: &str) -> Result<(), McpError> {
        if server_name.trim().is_empty() {
            return Err(McpError::Config(
                "MCP server name cannot be empty".to_string(),
            ));
        }
        if server_name.contains("__") {
            return Err(McpError::Config(format!(
                "MCP server `{server_name}` cannot contain `__`"
            )));
        }
        match self {
            Self::Stdio { command, .. } if command.trim().is_empty() => Err(McpError::Config(
                format!("MCP server `{server_name}` command cannot be empty"),
            )),
            Self::StreamableHttp { url, .. } | Self::Sse { url, .. } if url.trim().is_empty() => {
                Err(McpError::Config(format!(
                    "MCP server `{server_name}` URL cannot be empty"
                )))
            }
            _ => Ok(()),
        }
    }
}