RustAst/src/ast/mcp_server.rs

use rmcp::{
    ServerHandler,
    handler::server::{router::tool::ToolRouter, wrapper::Parameters},
    model::{ServerCapabilities, ServerInfo},
    schemars, tool, tool_handler, tool_router,
    transport::stdio,
};

use crate::ast::{environment::Environment, rtl};

const LANGUAGE_SPEC: &str = include_str!("../../docs/BNF.md");

/// Input struct for tools that take a Myc source code string.
#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct CodeInput {
    /// The Myc source code to process.
    pub code: String,
}

/// Input struct for tools that look up a single RTL symbol by name.
#[derive(Debug, serde::Deserialize, schemars::JsonSchema)]
pub struct SymbolInput {
    /// The RTL symbol name to look up (e.g. "+", "push", "series").
    pub name: String,
}

/// MCP server exposing the Myc compiler as a set of tools for LLMs.
#[derive(Clone)]
pub struct MycMcpServer {
    // Used by the generated `ServerHandler` impl from `#[tool_router]`.
    #[allow(dead_code)]
    tool_router: ToolRouter<Self>,
}

impl MycMcpServer {
    pub fn new() -> Self {
        Self {
            tool_router: Self::tool_router(),
        }
    }

    /// Creates a fresh, fully initialized Myc environment.
    /// A new environment is created per tool call because `Environment` uses
    /// `Rc<RefCell<...>>` internally and is intentionally single-threaded.
    fn make_env() -> Environment {
        let mut env = Environment::new();
        env.optimization = true;
        env
    }
}

impl Default for MycMcpServer {
    fn default() -> Self {
        Self::new()
    }
}

#[tool_router]
impl MycMcpServer {
    /// Evaluates a Myc expression and returns the result as a string.
    #[tool(description = "Evaluate a Myc expression or script and return the result. \
        The Myc language is a Lisp-like DSL for financial analysis with first-class ASTs, \
        closures, and series (time-series queues). \
        Use get_language_spec to learn the syntax first.")]
    fn eval_myc(&self, Parameters(CodeInput { code }): Parameters<CodeInput>) -> String {
        let env = Self::make_env();
        match env.run_script(&code) {
            Ok(value) => format!("{}", value),
            Err(e) => format!("Error: {}", e),
        }
    }

    /// Dumps the compiled and optimized AST for a Myc expression (verbose debug format).
    #[tool(description = "Compile a Myc expression and return a human-readable dump of \
        the compiled AST. Useful for understanding how the compiler transforms code \
        through its optimization and lowering passes. \
        Returns the full debug format including binding addresses and stack offsets. \
        Use dump_ast_compact for a cleaner view focused on types and purity.")]
    fn dump_ast(&self, Parameters(CodeInput { code }): Parameters<CodeInput>) -> String {
        let env = Self::make_env();
        match env.dump_ast(&code) {
            Ok(dump) => dump,
            Err(e) => format!("Error: {}", e),
        }
    }

    /// Dumps the compiled AST in compact, human-readable format.
    #[tool(description = "Compile a Myc expression and return a compact AST dump showing \
        inferred types (e.g. 'fn([int]) -> bool') and purity markers ('!' for impure, \
        '~' for side-effect-free). '[tail]' marks tail-call positions. \
        Omits internal details like stack offsets and binding addresses. \
        Prefer this over dump_ast when you want to understand the type structure of an expression.")]
    fn dump_ast_compact(&self, Parameters(CodeInput { code }): Parameters<CodeInput>) -> String {
        let env = Self::make_env();
        match env.dump_ast_compact(&code) {
            Ok(dump) => dump,
            Err(e) => format!("Error: {}", e),
        }
    }

    /// Checks the syntax and types of a Myc expression without executing it.
    #[tool(description = "Check the syntax and types of a Myc expression without running it. \
        Returns 'OK' if the code is valid, or a list of compiler diagnostics on error. \
        Note: accepts a single expression only. Wrap multiple expressions in (do ...): \
        e.g. (do (def x 42) (+ x 1)).")]
    fn check_syntax(&self, Parameters(CodeInput { code }): Parameters<CodeInput>) -> String {
        let env = Self::make_env();
        let result = env.compile(&code);
        if result.diagnostics.has_errors() {
            let errors: Vec<String> = result
                .diagnostics
                .items
                .into_iter()
                .map(|d| format!("{:?}: {}", d.level, d.message))
                .collect();
            errors.join("\n")
        } else {
            "OK".to_string()
        }
    }

    /// Lists all built-in functions and constants available in the Myc RTL.
    #[tool(description = "List all built-in functions and constants available in the Myc \
        runtime library (RTL). Returns a sorted list of all registered identifiers \
        such as +, -, push, series, print, etc.")]
    fn list_builtins(&self) -> String {
        let env = Self::make_env();
        // Only RTL bindings are in the fixed scope (scope 0), which is what
        // list_bindings() returns. This excludes user-defined symbols.
        let names = env.list_bindings();
        names.join("\n")
    }

    /// Returns the full Myc language specification (BNF grammar and semantics).
    #[tool(description = "Return the complete Myc language specification, including BNF grammar, \
        core semantics, data types, special forms, macro system, and standard library overview. \
        Read this first to understand how to write valid Myc code.")]
    fn get_language_spec(&self) -> String {
        LANGUAGE_SPEC.to_string()
    }

    /// Returns documentation for all documented RTL symbols.
    #[tool(description = "Return documentation for ALL built-in RTL functions and constants at once. \
        Prefer get_symbol_doc for individual lookups to save tokens. \
        Use this only when you need a full overview of the entire standard library.")]
    fn get_rtl_docs(&self) -> String {
        let env = Self::make_env();
        let docs = env.list_rtl_docs();
        if docs.is_empty() {
            "No documented RTL symbols found.".to_string()
        } else {
            docs.join("\n\n")
        }
    }

    /// Returns documentation for a single RTL symbol by name.
    #[tool(description = "Return the type signature, description, and examples for a single \
        built-in RTL symbol. Use list_builtins to see all available names, then call this \
        for each symbol you need details on. More token-efficient than get_rtl_docs.")]
    fn get_symbol_doc(&self, Parameters(SymbolInput { name }): Parameters<SymbolInput>) -> String {
        let env = Self::make_env();
        match env.get_rtl_doc(&name) {
            Some(doc) => doc,
            None => format!("No documentation found for symbol '{}'.", name),
        }
    }
}

#[tool_handler]
impl ServerHandler for MycMcpServer {
    fn get_info(&self) -> ServerInfo {
        ServerInfo::new(
            ServerCapabilities::builder()
                .enable_tools()
                .build(),
        )
        .with_instructions(
            "This server exposes the Myc compiler — a Lisp-like DSL for financial analysis. \
            Use `get_language_spec` to learn the language syntax and semantics, \
            `list_builtins` to see all available built-in functions, \
            `check_syntax` to validate code, \
            `eval_myc` to run expressions, and \
            `dump_ast` to inspect the compiled AST.",
        )
    }
}

/// Starts the MCP server on stdio (JSON-RPC over stdin/stdout).
///
/// Uses a single-threaded Tokio runtime because `Environment` relies on
/// `Rc<RefCell<...>>` and is not `Send`. The server handles one request
/// at a time, which matches the single-threaded design of the Myc runtime.
pub fn run() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
    // Register RTL once just to warm the prelude path — the actual environment
    // used per tool call is created fresh in `make_env()`.
    // This also validates that the embedded system library compiles correctly.
    let _ = rtl::register as fn(&Environment);

    tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?
        .block_on(async {
            let server = MycMcpServer::new();
            let service = rmcp::serve_server(server, stdio()).await?;
            service.waiting().await?;
            Ok(())
        })
}