rustapi_mcp/

server.rs

//! The main `McpServer` type and lifecycle.

use crate::config::{InvocationMode, McpConfig};
use crate::discovery;
use crate::error::{McpError, Result};
use crate::invocation::RequestInvoker;
use crate::types::{McpCapability, McpTool};
use bytes::Bytes;
use http_body_util::{BodyExt, Full};
use hyper::body::Incoming;
use hyper::server::conn::http1;
use hyper::Response;
use hyper_util::rt::TokioIo;
use rustapi_core::{Request as CoreRequest, RustApi};
use rustapi_openapi::OpenApiSpec;
use std::collections::HashMap;
use std::convert::Infallible;
use std::future::Future;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::TcpListener;
use tracing::{error, info};

/// The main handle for the MCP integration.
///
/// `McpServer` can be attached to a `RustApi` instance (via its OpenAPI spec)
/// to automatically discover tools. Tool invocations (in later milestones)
/// will be driven through the normal RustAPI handler pipeline.
#[derive(Debug, Clone)]
pub struct McpServer {
    config: Arc<McpConfig>,
    /// Attached OpenAPI spec used for tool discovery.
    openapi: Option<OpenApiSpec>,
    /// Base URL of the main RustAPI HTTP server (for proxy-style tool invocation over localhost).
    /// Example: "http://127.0.0.1:8080"
    http_base: Option<String>,
    /// Internal mapping from tool name (as advertised to MCP) to the original HTTP route info.
    /// Used to reconstruct the correct request when a tool is called.
    tool_map: HashMap<String, ToolExecutionInfo>,

    /// Optional in-process invoker. Present when created via `from_rustapi`
    /// and the configured `invocation_mode` allows direct dispatch.
    invoker: Option<RequestInvoker>,
}

/// Internal info needed to turn an MCP tool/call into an actual HTTP request
/// against the main API.
#[derive(Debug, Clone)]
struct ToolExecutionInfo {
    /// Original path template, e.g. "/users/{id}"
    path_template: String,
    /// HTTP method as string, e.g. "GET", "POST"
    method: String,
}

impl McpServer {
    /// Create a new MCP server with the given configuration.
    pub fn new(config: McpConfig) -> Self {
        Self {
            config: Arc::new(config),
            openapi: None,
            http_base: None,
            tool_map: HashMap::new(),
            invoker: None,
        }
    }

    /// Create an MCP server pre-attached to a `RustApi` instance.
    ///
    /// This is the most ergonomic way when you already have a built `RustApi`.
    /// If `invocation_mode` is `InProcess` or `Auto`, a direct in-process
    /// dispatcher is captured for zero-overhead tool calls (see `RequestInvoker`).
    pub fn from_rustapi(app: &RustApi, config: McpConfig) -> Self {
        let mut server = Self::new(config);

        let spec = app.openapi_spec().clone();
        server.openapi = Some(spec.clone());
        server.rebuild_tool_map(&spec);

        // Only capture in-process dispatcher when the configured mode can actually use it.
        // This avoids unnecessary work and state cloning when the user explicitly wants Proxy.
        let mode = server.config.invocation_mode;
        if mode == InvocationMode::InProcess || mode == InvocationMode::Auto {
            let dispatcher = app.request_dispatcher();
            server.invoker = Some(RequestInvoker::new(dispatcher, server.config.clone()));
        }

        server
    }

    /// Create from an explicit OpenAPI spec.
    pub fn from_spec(config: McpConfig, spec: &OpenApiSpec) -> Self {
        let mut server = Self::new(config);
        server.openapi = Some(spec.clone());
        server.rebuild_tool_map(spec);
        server
    }

    /// Attach (or replace) the OpenAPI spec used for discovery.
    ///
    /// Call this after `RustApi::auto()` / builder if you built the app first.
    pub fn with_openapi(mut self, spec: OpenApiSpec) -> Self {
        self.rebuild_tool_map(&spec);
        self.openapi = Some(spec);
        self
    }

    /// Configure the base URL of the main RustAPI HTTP server.
    ///
    /// When set, `tools/call` will proxy the call over HTTP to this base
    /// (typically "http://127.0.0.1:8080" when using the concurrent runner).
    /// This guarantees that tool invocations go through the exact same
    /// middleware, auth, validation, and handler code as normal traffic.
    pub fn with_http_base(mut self, base: impl Into<String>) -> Self {
        self.http_base = Some(base.into());
        self
    }

    /// Rebuild the internal tool_name -> execution info map from the OpenAPI spec.
    /// Called automatically when attaching a spec.
    fn rebuild_tool_map(&mut self, spec: &OpenApiSpec) {
        self.tool_map.clear();

        let config = &*self.config;

        let insert_tool = |map: &mut HashMap<String, ToolExecutionInfo>,
                           path: &str,
                           method: &str,
                           op: &rustapi_openapi::Operation| {
            if !config.allowed_tags.is_empty() {
                let has_match = op.tags.iter().any(|t| config.allowed_tags.contains(t));
                if !has_match {
                    return;
                }
            }

            let name = generate_tool_name(method, path, op);
            let info = ToolExecutionInfo {
                path_template: path.to_string(),
                method: method.to_string(),
            };
            map.insert(name, info);
        };

        for (path, path_item) in &spec.paths {
            if !path_matches_prefixes(path, &config.allowed_path_prefixes) {
                continue;
            }

            if let Some(op) = &path_item.get {
                insert_tool(&mut self.tool_map, path, "GET", op);
            }
            if let Some(op) = &path_item.post {
                insert_tool(&mut self.tool_map, path, "POST", op);
            }
            if let Some(op) = &path_item.put {
                insert_tool(&mut self.tool_map, path, "PUT", op);
            }
            if let Some(op) = &path_item.patch {
                insert_tool(&mut self.tool_map, path, "PATCH", op);
            }
            if let Some(op) = &path_item.delete {
                insert_tool(&mut self.tool_map, path, "DELETE", op);
            }
        }
    }

    /// Get a reference to the active configuration.
    pub fn config(&self) -> &McpConfig {
        &self.config
    }

    /// Return the capabilities this server currently advertises.
    pub fn capabilities(&self) -> Vec<McpCapability> {
        let mut caps = vec![];

        if self.config.tools_enabled {
            caps.push(McpCapability::Tools);
        }

        caps
    }

    /// Perform the `initialize` handshake.
    ///
    /// MCP clients call this first. Returns server info + supported capabilities.
    pub fn initialize(&self) -> InitializeResult {
        InitializeResult {
            name: self.config.name.clone(),
            version: self.config.version.clone(),
            description: self.config.description.clone(),
            capabilities: self.capabilities(),
        }
    }

    /// Discover and return the list of tools that should be exposed to MCP clients.
    ///
    /// Tools are derived from the attached OpenAPI spec (from `RustApi`), filtered
    /// according to `McpConfig::allowed_tags` and `allowed_path_prefixes`.
    pub async fn list_tools(&self) -> Result<Vec<McpTool>> {
        if !self.config.tools_enabled {
            return Err(crate::error::McpError::CapabilityNotEnabled(
                "tools".to_string(),
            ));
        }

        if let Some(spec) = &self.openapi {
            let tools = discovery::extract_tools_from_spec(spec, &self.config);
            Ok(tools)
        } else {
            // No spec attached → no tools (safe default)
            Ok(vec![])
        }
    }

    /// Execute a tool call by proxying it as a real HTTP request to the main
    /// RustAPI server (using the configured `http_base`).
    ///
    /// This is the heart of Native MCP: the call goes through your normal
    /// layers, interceptors, extractors, validation, error handling, etc.
    pub async fn call_tool(
        &self,
        req: crate::types::ToolCallRequest,
    ) -> Result<crate::types::ToolCallResponse> {
        let info = self.tool_map.get(&req.name).ok_or_else(|| {
            McpError::ToolNotFound(format!("no tool registered with name '{}'", req.name))
        })?;

        let path = substitute_path_params(&info.path_template, &req.arguments);

        // Decide execution strategy
        let use_inprocess = if let Some(inv) = &self.invoker {
            inv.should_use()
        } else {
            false
        };

        if use_inprocess {
            if let Some(invoker) = &self.invoker {
                let state = invoker.state_ref();
                let core_req = self
                    .build_core_request(info, &req.arguments, &path, state)
                    .await?;
                let core_resp = invoker.invoke(core_req).await;
                return self.core_response_to_tool_response(core_resp).await;
            }
        }

        // Fallback / default: proxy via HTTP (current behavior)
        let base = self.http_base.as_deref().unwrap_or("http://127.0.0.1:8080");

        let url = format!("{}{}", base.trim_end_matches('/'), path);

        let method =
            reqwest::Method::from_bytes(info.method.as_bytes()).unwrap_or(reqwest::Method::GET);

        let client = reqwest::Client::new();

        let mut request_builder = client.request(method.clone(), &url);

        // If this looks like a mutating method, send the arguments as JSON body
        let is_body_method = matches!(info.method.as_str(), "POST" | "PUT" | "PATCH");
        if is_body_method && !req.arguments.is_empty() {
            request_builder = request_builder
                .header("content-type", "application/json")
                .json(&req.arguments);
        } else if !is_body_method && !req.arguments.is_empty() {
            // For GET/DELETE etc, we could turn remaining args into query params.
            // For MVP we rely on path params; extra args are ignored for now.
        }

        let resp = request_builder.send().await.map_err(|e| {
            McpError::ToolExecution(format!("failed to proxy tool call to main API: {}", e))
        })?;

        let status = resp.status();
        let is_error = !status.is_success();

        let content = if let Ok(text) = resp.text().await {
            if text.trim().is_empty() {
                serde_json::Value::Null
            } else {
                serde_json::from_str(&text).unwrap_or(serde_json::Value::String(text))
            }
        } else {
            serde_json::Value::Null
        };

        Ok(crate::types::ToolCallResponse {
            content,
            is_error,
            meta: Some(serde_json::json!({ "proxied_status": status.as_u16() })),
        })
    }

    /// Build a `core::Request` suitable for in-process dispatch.
    async fn build_core_request(
        &self,
        info: &ToolExecutionInfo,
        arguments: &std::collections::HashMap<String, serde_json::Value>,
        substituted_path: &str,
        state: Arc<http::Extensions>,
    ) -> Result<CoreRequest> {
        use bytes::Bytes;
        use http::{Method, Request as HttpRequest, Uri, Version};
        use rustapi_core::{BodyVariant, PathParams};

        let method = Method::from_bytes(info.method.as_bytes()).unwrap_or(Method::GET);
        let uri: Uri = substituted_path
            .parse()
            .unwrap_or_else(|_| "/".parse().expect("valid fallback uri"));

        let mut builder = HttpRequest::builder()
            .method(method.clone())
            .uri(uri)
            .version(Version::HTTP_11);

        let is_body_method = matches!(info.method.as_str(), "POST" | "PUT" | "PATCH");

        let body_bytes = if is_body_method && !arguments.is_empty() {
            builder = builder.header("content-type", "application/json");
            let json = serde_json::to_vec(arguments).unwrap_or_default();
            Bytes::from(json)
        } else {
            Bytes::new()
        };

        let http_req = builder
            .body(())
            .map_err(|e| McpError::ToolExecution(format!("failed to build request: {}", e)))?;

        let (mut parts, _) = http_req.into_parts();

        parts.method = method;
        parts.uri = substituted_path.parse().unwrap_or(parts.uri);

        let path_params = PathParams::new();

        let core_req =
            CoreRequest::new(parts, BodyVariant::Buffered(body_bytes), state, path_params);

        Ok(core_req)
    }

    /// Convert a core Response back into a ToolCallResponse.
    async fn core_response_to_tool_response(
        &self,
        resp: rustapi_core::Response,
    ) -> Result<crate::types::ToolCallResponse> {
        use http_body_util::BodyExt;

        let status = resp.status();
        let is_error = !status.is_success();

        // Collect body
        let body_bytes = resp
            .into_body()
            .collect()
            .await
            .map(|collected| collected.to_bytes())
            .unwrap_or_default();

        let content = if body_bytes.is_empty() {
            serde_json::Value::Null
        } else if let Ok(text) = std::str::from_utf8(&body_bytes) {
            if text.trim().is_empty() {
                serde_json::Value::Null
            } else {
                serde_json::from_str(text).unwrap_or(serde_json::Value::String(text.to_string()))
            }
        } else {
            serde_json::Value::String(String::from_utf8_lossy(&body_bytes).to_string())
        };

        Ok(crate::types::ToolCallResponse {
            content,
            is_error,
            meta: Some(serde_json::json!({ "status": status.as_u16(), "in_process": true })),
        })
    }
}

/// Result of the `initialize` handshake.
#[derive(Debug, Clone)]
pub struct InitializeResult {
    /// Server name (from config).
    pub name: String,
    /// Server version (from config).
    pub version: String,
    /// Optional description.
    pub description: Option<String>,
    /// Advertised capabilities.
    pub capabilities: Vec<McpCapability>,
}

impl McpServer {
    /// Start serving the MCP protocol over HTTP on the given address.
    ///
    /// This starts a sidecar HTTP server (separate from your main RustAPI HTTP server)
    /// that MCP clients (Claude, etc.) can connect to for tool discovery and invocation.
    ///
    /// Supports a minimal JSON-RPC over POST transport for:
    /// - `initialize`
    /// - `tools/list`
    /// - `tools/call` (proxies to the main RustAPI server so full middleware / validation / error handling applies)
    pub async fn serve(self, addr: &str) -> Result<()> {
        self.serve_with_shutdown(addr, std::future::pending()).await
    }

    /// Like `serve`, but with a shutdown signal (e.g. ctrl_c()).
    pub async fn serve_with_shutdown<F>(self, addr: &str, signal: F) -> Result<()>
    where
        F: Future<Output = ()> + Send + 'static,
    {
        let addr: SocketAddr = addr
            .parse()
            .map_err(|e| McpError::Transport(format!("invalid MCP address '{}': {}", addr, e)))?;

        let listener = TcpListener::bind(addr).await.map_err(|e| {
            McpError::Transport(format!("failed to bind MCP listener on {}: {}", addr, e))
        })?;

        info!("🧠 MCP server listening on http://{}", addr);

        let this = Arc::new(self);
        tokio::pin!(signal);

        loop {
            tokio::select! {
                biased;

                accept_result = listener.accept() => {
                    let (stream, remote) = match accept_result {
                        Ok(v) => v,
                        Err(e) => {
                            error!("MCP accept error: {}", e);
                            continue;
                        }
                    };

                    let io = TokioIo::new(stream);
                    let mcp = this.clone();

                    tokio::task::spawn(async move {
                        let service = hyper::service::service_fn(move |req| {
                            let mcp = mcp.clone();
                            async move { handle_mcp_http_request(mcp, req).await }
                        });

                        if let Err(e) = http1::Builder::new()
                            .serve_connection(io, service)
                            .await
                        {
                            error!("MCP connection error from {}: {}", remote, e);
                        }
                    });
                }

                _ = &mut signal => {
                    info!("MCP server received shutdown signal");
                    break;
                }
            }
        }

        Ok(())
    }
}

/// Handle a single HTTP request for the MCP protocol (minimal JSON-RPC over POST).
async fn handle_mcp_http_request(
    mcp: Arc<McpServer>,
    req: hyper::Request<Incoming>,
) -> std::result::Result<Response<Full<Bytes>>, Infallible> {
    if req.method() != hyper::Method::POST {
        let body = Bytes::from_static(b"{\"error\":\"MCP transport expects POST requests\"}");
        return Ok(Response::builder()
            .status(405)
            .header("content-type", "application/json")
            .body(Full::new(body))
            .expect("static response must build"));
    }

    let body_bytes = match req.collect().await {
        Ok(collected) => collected.to_bytes(),
        Err(e) => {
            let err_body = format!("{{\"error\":\"failed to read body: {}\"}}", e);
            return Ok(Response::builder()
                .status(400)
                .header("content-type", "application/json")
                .body(Full::new(Bytes::from(err_body)))
                .expect("error response must build"));
        }
    };

    let json: serde_json::Value = match serde_json::from_slice(&body_bytes) {
        Ok(v) => v,
        Err(_) => {
            return Ok(jsonrpc_error_response(
                serde_json::Value::Null,
                -32700,
                "parse error",
            ));
        }
    };

    let id = json.get("id").cloned().unwrap_or(serde_json::Value::Null);
    let method = json.get("method").and_then(|m| m.as_str()).unwrap_or("");

    match method {
        "initialize" => {
            let init = mcp.initialize();
            let result = serde_json::json!({
                "protocolVersion": "2024-11-05",
                "serverInfo": {
                    "name": init.name,
                    "version": init.version
                },
                "capabilities": {
                    "tools": {}
                }
            });
            Ok(jsonrpc_success_response(id, result))
        }
        "tools/list" => match mcp.list_tools().await {
            Ok(tools) => {
                let tool_defs: Vec<_> = tools
                    .into_iter()
                    .map(|t| {
                        let mut def = serde_json::json!({
                            "name": t.name,
                            "description": t.description,
                            "inputSchema": t.input_schema
                        });

                        // Permission scoping metadata (framework-native)
                        if let Some(perm) = &t.permission {
                            def["permission"] = serde_json::Value::String(perm.clone());
                        }
                        if let Some(confirm) = t.requires_confirmation {
                            if confirm {
                                def["requiresConfirmation"] = serde_json::Value::Bool(true);
                            }
                        }
                        // Future: "confirmationMessage"

                        def
                    })
                    .collect();

                let result = serde_json::json!({ "tools": tool_defs });
                Ok(jsonrpc_success_response(id, result))
            }
            Err(e) => Ok(jsonrpc_error_response(
                id,
                -32603,
                &format!("internal error: {}", e),
            )),
        },
        "tools/call" => {
            let params = json.get("params").cloned().unwrap_or(serde_json::json!({}));
            let name = params
                .get("name")
                .and_then(|v| v.as_str())
                .unwrap_or("")
                .to_string();
            let arguments: std::collections::HashMap<String, serde_json::Value> = params
                .get("arguments")
                .and_then(|v| serde_json::from_value(v.clone()).ok())
                .unwrap_or_default();

            let tool_req = crate::types::ToolCallRequest { name, arguments };

            match mcp.call_tool(tool_req).await {
                Ok(resp) => {
                    // Shape the response per MCP tools/call convention (protocol 2024-11-05):
                    // The result contains a "content" array of content blocks + "isError" flag.
                    // We serialize non-string content as pretty JSON text for broad client compatibility.
                    let text = if resp.content.is_null() {
                        String::new()
                    } else if let Some(s) = resp.content.as_str() {
                        s.to_owned()
                    } else {
                        serde_json::to_string_pretty(&resp.content)
                            .unwrap_or_else(|_| resp.content.to_string())
                    };

                    let result = serde_json::json!({
                        "content": [{
                            "type": "text",
                            "text": text
                        }],
                        "isError": resp.is_error
                    });
                    Ok(jsonrpc_success_response(id, result))
                }
                Err(e) => {
                    // Surface execution errors as a successful JSON-RPC but with isError inside the result
                    // (this is the MCP convention so the client knows it was a tool-level failure).
                    let result = serde_json::json!({
                        "content": [{
                            "type": "text",
                            "text": format!("Tool execution error: {}", e)
                        }],
                        "isError": true
                    });
                    Ok(jsonrpc_success_response(id, result))
                }
            }
        }
        _ => Ok(jsonrpc_error_response(id, -32601, "method not found")),
    }
}

fn jsonrpc_success_response(
    id: serde_json::Value,
    result: serde_json::Value,
) -> Response<Full<Bytes>> {
    let body = serde_json::json!({
        "jsonrpc": "2.0",
        "id": id,
        "result": result
    });
    Response::builder()
        .header("content-type", "application/json")
        .body(Full::new(Bytes::from(
            serde_json::to_vec(&body).expect("json must serialize"),
        )))
        .expect("response must build")
}

fn jsonrpc_error_response(
    id: serde_json::Value,
    code: i32,
    message: &str,
) -> Response<Full<Bytes>> {
    let body = serde_json::json!({
        "jsonrpc": "2.0",
        "id": id,
        "error": {
            "code": code,
            "message": message
        }
    });
    Response::builder()
        .header("content-type", "application/json")
        .body(Full::new(Bytes::from(
            serde_json::to_vec(&body).expect("json must serialize"),
        )))
        .expect("response must build")
}

/// Helpers duplicated from discovery for tool map building (small & self-contained).
fn path_matches_prefixes(path: &str, prefixes: &[String]) -> bool {
    if prefixes.is_empty() {
        return true;
    }
    prefixes.iter().any(|p| path.starts_with(p))
}

fn generate_tool_name(method: &str, path: &str, op: &rustapi_openapi::Operation) -> String {
    if let Some(oid) = &op.operation_id {
        return sanitize_name(oid);
    }
    let mut slug = path
        .trim_start_matches('/')
        .replace(['/', '{', '}', ':'], "_")
        .replace(['-', '.', ' '], "_");
    while slug.contains("__") {
        slug = slug.replace("__", "_");
    }
    let slug = slug.trim_matches('_').to_string();
    let method_lower = method.to_lowercase();
    if slug.is_empty() {
        method_lower
    } else {
        format!("{}_{}", method_lower, slug)
    }
}

fn sanitize_name(s: &str) -> String {
    s.chars()
        .map(|c| {
            if c.is_alphanumeric() || c == '_' {
                c
            } else {
                '_'
            }
        })
        .collect::<String>()
        .trim_matches('_')
        .to_string()
        .to_lowercase()
}

/// Substitute {param} placeholders in the path template using values from the MCP arguments.
fn substitute_path_params(
    template: &str,
    args: &std::collections::HashMap<String, serde_json::Value>,
) -> String {
    let mut result = template.to_string();
    for (key, value) in args {
        let placeholder = format!("{{{}}}", key);
        if result.contains(&placeholder) {
            let val_str = match value {
                serde_json::Value::String(s) => s.clone(),
                other => other.to_string().trim_matches('"').to_string(),
            };
            result = result.replace(&placeholder, &val_str);
        }
    }
    result
}