mcp-compressor-core 0.21.3

use crate::ffi::FfiClientArtifactKind;
use crate::llm_assist::{ModelRef, DEFAULT_MODEL_REF};
use std::path::PathBuf;
use std::str::FromStr;

use clap::{ArgAction, Parser, Subcommand, ValueEnum};

use crate::compression::CompressionLevel;
use crate::server::{BackendServerConfig, ProxyTransformMode};

#[derive(Debug, Parser)]
#[command(
    name = "mcp-compressor",
    about = "Standalone Rust MCP compressor core binary",
    disable_help_subcommand = true,
    version,
    override_usage = "mcp-compressor [OPTIONS] [COMMAND] -- <URL_OR_COMMAND> [BACKEND_ARGS]..."
)]
pub struct CliOptions {
    #[command(subcommand)]
    pub command_kind: Option<CliCommand>,

    /// Compression level: low, medium, high, or max.
    #[arg(short = 'c', long, value_enum, default_value = "medium")]
    compression: CompressionLevelArg,

    /// MCP config JSON file.
    #[arg(long = "config")]
    pub config_path: Option<PathBuf>,

    /// Frontend server name/prefix.
    #[arg(short = 'n', long)]
    pub server_name: Option<String>,

    /// Frontend transform mode.
    #[arg(long, value_enum, default_value = "compressed-tools")]
    transform_mode: TransformModeArg,

    /// Alias for --transform-mode cli.
    #[arg(long, action = ArgAction::SetTrue)]
    cli_mode: bool,

    /// Alias for --transform-mode just-bash.
    #[arg(long, action = ArgAction::SetTrue)]
    just_bash: bool,

    /// Alias for --transform-mode just-bash.
    #[arg(long = "just-bash-mode", action = ArgAction::SetTrue)]
    just_bash_mode: bool,

    /// Generate a Python or TypeScript code client that talks to the local proxy.
    #[arg(long = "code-mode", value_enum, value_name = "LANGUAGE")]
    code_mode: Option<CodeModeArg>,

    /// Deprecated alias for --code-mode python.
    #[arg(long = "python-mode", action = ArgAction::SetTrue, hide = true)]
    python_mode: bool,

    /// Deprecated alias for --code-mode typescript.
    #[arg(long = "typescript-mode", action = ArgAction::SetTrue, hide = true)]
    typescript_mode: bool,

    /// Comma-separated backend tool names to include.
    #[arg(long, value_delimiter = ',')]
    pub include_tools: Vec<String>,

    /// Comma-separated backend tool names to exclude.
    #[arg(long, value_delimiter = ',')]
    pub exclude_tools: Vec<String>,

    /// Convert JSON text outputs to TOON where possible.
    #[arg(long, action = ArgAction::SetTrue)]
    pub toonify: bool,

    /// Output directory for generated Python/TypeScript code clients.
    #[arg(long = "output-dir")]
    pub output_dir: Option<PathBuf>,

    /// Multi-server backend spec: name=command [args...]. Repeat for each backend.
    #[arg(long = "multi-server", value_name = "NAME=COMMAND [ARGS...]", action = ArgAction::Append)]
    pub multi_server: Vec<MultiServerArg>,

    /// Frontend transport.
    #[arg(long, value_enum, default_value = "stdio")]
    pub transport: FrontendTransport,

    /// Port for streamable-http frontend; 0 chooses an available port.
    #[arg(long, default_value_t = 8000)]
    pub port: u16,

    /// Backend URL or command plus backend arguments. All backend server arguments belong after `--`.
    #[arg(value_name = "URL_OR_COMMAND", allow_hyphen_values = true, last = true)]
    pub command: Vec<String>,
}

impl CliOptions {
    pub fn validate(&self) -> Result<(), String> {
        if self.config_path.is_some() && self.server_name.is_some() {
            return Err("--server-name cannot be used with --config; MCP config server names come from mcpServers keys".to_string());
        }

        let mode_aliases = [
            self.cli_mode,
            self.just_bash || self.just_bash_mode,
            self.code_mode.is_some() || self.python_mode || self.typescript_mode,
        ]
        .into_iter()
        .filter(|enabled| *enabled)
        .count();
        if mode_aliases > 1 {
            return Err(
                "choose only one of --cli-mode, --just-bash-mode, or --code-mode".to_string(),
            );
        }

        let code_mode_aliases = [
            self.code_mode.is_some(),
            self.python_mode,
            self.typescript_mode,
        ]
        .into_iter()
        .filter(|enabled| *enabled)
        .count();
        if code_mode_aliases > 1 {
            return Err(
                "choose only one code mode: --code-mode python or --code-mode typescript"
                    .to_string(),
            );
        }
        Ok(())
    }

    pub fn compression(&self) -> CompressionLevel {
        self.compression.into()
    }

    pub fn transform_mode(&self) -> ProxyTransformMode {
        if self.just_bash || self.just_bash_mode {
            ProxyTransformMode::JustBash
        } else if self.cli_mode
            || self.python_mode
            || self.typescript_mode
            || self.code_mode.is_some()
        {
            ProxyTransformMode::Cli
        } else {
            self.transform_mode.into()
        }
    }

    pub fn client_artifact_kind(&self) -> Option<FfiClientArtifactKind> {
        if let Some(code_mode) = self.code_mode {
            return Some(code_mode.into());
        }
        if self.python_mode {
            Some(FfiClientArtifactKind::Python)
        } else if self.typescript_mode {
            Some(FfiClientArtifactKind::TypeScript)
        } else {
            None
        }
    }
}

#[derive(Debug, Clone, Subcommand)]
pub enum CliCommand {
    /// Clear stored OAuth credentials.
    ClearOauth {
        /// Backend server name or URL to clear. If omitted, all Rust OAuth state is removed.
        target: Option<String>,
    },

    /// Manage optional local LLM assistance assets.
    Llm {
        #[command(subcommand)]
        command: LlmCommand,
    },
}

impl CliCommand {
    pub fn clear_oauth_target(&self) -> Option<&str> {
        match self {
            Self::ClearOauth { target } => target.as_deref(),
            Self::Llm { .. } => None,
        }
    }

    pub fn llm_command(&self) -> Option<&LlmCommand> {
        match self {
            Self::Llm { command } => Some(command),
            Self::ClearOauth { .. } => None,
        }
    }
}

#[derive(Debug, Clone, Subcommand)]
pub enum LlmCommand {
    /// Show local llama-server/model installation status.
    Status(LlmCommandOptions),
    /// Download/install llama-server and the configured model.
    Pull(LlmCommandOptions),
    /// Remove managed LLM runtime and model assets from the cache.
    Remove(LlmCommandOptions),
    /// Download/install assets and run a test prompt.
    Test(LlmTestOptions),
}

#[derive(Debug, Clone, Parser)]
pub struct LlmCommandOptions {
    /// Model reference, for example LiquidAI/LFM2.5-350M-GGUF:Q4_K_M.
    #[arg(long = "model", default_value = DEFAULT_MODEL_REF)]
    model: String,

    /// Override mcp-compressor LLM cache directory.
    #[arg(long = "cache-dir")]
    cache_dir: Option<PathBuf>,

    /// Explicit llama-server binary path.
    #[arg(long = "llama-server")]
    llama_server_path: Option<PathBuf>,
}

impl LlmCommandOptions {
    pub fn config(
        &self,
        allow_download: bool,
    ) -> Result<crate::llm_assist::LlmRuntimeConfig, String> {
        let mut config = crate::llm_assist::LlmRuntimeConfig::local_default(allow_download)
            .with_model(ModelRef::parse(&self.model).map_err(|error| error.to_string())?);
        if let Some(cache_dir) = &self.cache_dir {
            config = config.with_cache_dir(cache_dir.clone());
        }
        if let Some(path) = &self.llama_server_path {
            config = config.with_llama_server_path(path.clone());
        }
        Ok(config)
    }

    pub fn cache_dir(&self) -> Option<PathBuf> {
        self.cache_dir.clone()
    }
}

#[derive(Debug, Clone, Parser)]
pub struct LlmTestOptions {
    #[command(flatten)]
    options: LlmCommandOptions,

    /// Prompt to send to the local model.
    #[arg(long = "prompt", default_value = "Say hello in one short sentence.")]
    prompt: String,
}

impl LlmTestOptions {
    pub fn config(
        &self,
        allow_download: bool,
    ) -> Result<crate::llm_assist::LlmRuntimeConfig, String> {
        self.options.config(allow_download)
    }

    pub fn prompt(&self) -> &str {
        &self.prompt
    }
}

#[derive(Debug, Clone)]
pub struct MultiServerArg {
    name: String,
    command: String,
    args: Vec<String>,
}

impl MultiServerArg {
    pub fn into_backend(self) -> BackendServerConfig {
        BackendServerConfig::new(self.name, self.command, self.args)
    }
}

impl FromStr for MultiServerArg {
    type Err = String;

    fn from_str(value: &str) -> Result<Self, Self::Err> {
        let mut parts = value.split_whitespace();
        let spec = parts
            .next()
            .ok_or_else(|| "expected name=command".to_string())?;
        let (name, command) = spec
            .split_once('=')
            .filter(|(name, command)| !name.is_empty() && !command.is_empty())
            .ok_or_else(|| "expected name=command".to_string())?;
        Ok(Self {
            name: name.to_string(),
            command: command.to_string(),
            args: parts.map(ToString::to_string).collect(),
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
enum CompressionLevelArg {
    Low,
    Medium,
    High,
    Max,
}

impl std::fmt::Display for CompressionLevelArg {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(match self {
            Self::Low => "low",
            Self::Medium => "medium",
            Self::High => "high",
            Self::Max => "max",
        })
    }
}

impl From<CompressionLevelArg> for CompressionLevel {
    fn from(value: CompressionLevelArg) -> Self {
        match value {
            CompressionLevelArg::Low => CompressionLevel::Low,
            CompressionLevelArg::Medium => CompressionLevel::Medium,
            CompressionLevelArg::High => CompressionLevel::High,
            CompressionLevelArg::Max => CompressionLevel::Max,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
enum CodeModeArg {
    Python,
    #[value(name = "typescript")]
    TypeScript,
}

impl From<CodeModeArg> for FfiClientArtifactKind {
    fn from(value: CodeModeArg) -> Self {
        match value {
            CodeModeArg::Python => FfiClientArtifactKind::Python,
            CodeModeArg::TypeScript => FfiClientArtifactKind::TypeScript,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
enum TransformModeArg {
    CompressedTools,
    Cli,
    JustBash,
}

impl std::fmt::Display for TransformModeArg {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(match self {
            Self::CompressedTools => "compressed-tools",
            Self::Cli => "cli",
            Self::JustBash => "just-bash",
        })
    }
}

impl From<TransformModeArg> for ProxyTransformMode {
    fn from(value: TransformModeArg) -> Self {
        match value {
            TransformModeArg::CompressedTools => ProxyTransformMode::CompressedTools,
            TransformModeArg::Cli => ProxyTransformMode::Cli,
            TransformModeArg::JustBash => ProxyTransformMode::JustBash,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
pub enum FrontendTransport {
    Stdio,
    StreamableHttp,
}

impl std::fmt::Display for FrontendTransport {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(match self {
            Self::Stdio => "stdio",
            Self::StreamableHttp => "streamable-http",
        })
    }
}