recursive/

mcp_server.rs

//! MCP server lifecycle management.
//!
//! This module provides [`McpServerManager`], which handles spawning MCP
//! servers (stdio or HTTP+SSE), discovering their tools, and registering
//! them into the agent's [`ToolRegistry`]. Each MCP tool is wrapped in a
//! thin adapter that delegates execution to the corresponding server.
//!
//! # Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────┐
//! │              McpServerManager               │
//! │  ┌──────────┐  ┌──────────┐  ┌──────────┐  │
//! │  │ Server A │  │ Server B │  │ Server C │  │
//! │  │ (stdio)  │  │ (SSE)    │  │ (stdio)  │  │
//! │  └────┬─────┘  └────┬─────┘  └────┬─────┘  │
//! │       │              │              │        │
//! │  ┌────▼──────────────▼──────────────▼────┐  │
//! │  │          ToolRegistry                 │  │
//! │  │  mcp__A__tool1  mcp__B__tool2  ...    │  │
//! │  └───────────────────────────────────────┘  │
//! └─────────────────────────────────────────────┘
//! ```

use std::collections::HashMap;
use std::sync::Arc;

use tokio::sync::Mutex;
use tracing::{info, instrument};

use crate::error::{Error, Result};
use crate::llm::ToolSpec;
use crate::mcp::{McpClient, McpServer, McpTool};
use crate::tools::ToolRegistry;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};

// ---------------------------------------------------------------------------
// JSON-RPC 2.0 protocol types (MCP server side)
// ---------------------------------------------------------------------------

/// A JSON-RPC 2.0 request.
#[derive(Debug, Deserialize)]
pub struct JsonRpcRequest {
    pub jsonrpc: String,
    #[serde(default)]
    pub id: Option<Value>,
    pub method: String,
    #[serde(default)]
    pub params: Value,
}

/// A JSON-RPC 2.0 response.
#[derive(Debug, Serialize)]
pub struct JsonRpcResponse {
    pub jsonrpc: String,
    pub id: Value,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<JsonRpcError>,
}

/// A JSON-RPC 2.0 error object.
#[derive(Debug, Serialize)]
pub struct JsonRpcError {
    pub code: i32,
    pub message: String,
}

impl JsonRpcResponse {
    /// Create a success response.
    pub fn success(id: Value, result: Value) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id,
            result: Some(result),
            error: None,
        }
    }

    /// Create an error response.
    pub fn error(id: Value, code: i32, message: String) -> Self {
        Self {
            jsonrpc: "2.0".to_string(),
            id,
            result: None,
            error: Some(JsonRpcError { code, message }),
        }
    }

    /// Create a parse error response (-32700).
    pub fn parse_error() -> Self {
        Self::error(Value::Null, -32700, "Parse error".to_string())
    }

    /// Create a method not found error response (-32601).
    pub fn method_not_found(id: Value) -> Self {
        Self::error(id, -32601, "Method not found".to_string())
    }
}

// ---------------------------------------------------------------------------
// Dispatch
// ---------------------------------------------------------------------------

/// Dispatch a parsed JSON-RPC request and return a response.
/// Returns `None` for notifications (no response needed).
pub async fn dispatch_request(
    request: &JsonRpcRequest,
    tool_specs: &[ToolSpec],
    tools: &ToolRegistry,
) -> Option<JsonRpcResponse> {
    match request.method.as_str() {
        "initialize" => Some(handle_initialize(request)),
        "notifications/initialized" => None,
        "tools/list" => Some(handle_tools_list(request, tool_specs)),
        "tools/call" => Some(handle_tools_call(request, tools).await),
        _ => {
            let id = request.id.clone().unwrap_or(Value::Null);
            Some(JsonRpcResponse::method_not_found(id))
        }
    }
}

/// Handle `initialize` — return server capabilities and info.
fn handle_initialize(request: &JsonRpcRequest) -> JsonRpcResponse {
    let id = request.id.clone().unwrap_or(Value::Null);
    JsonRpcResponse::success(
        id,
        serde_json::json!({
            "protocolVersion": "2024-11-05",
            "capabilities": {
                "tools": {}
            },
            "serverInfo": {
                "name": "recursive-mcp-server",
                "version": "0.1.0"
            }
        }),
    )
}

/// Handle `tools/list` — return the list of tool specs in MCP format.
fn handle_tools_list(request: &JsonRpcRequest, tool_specs: &[ToolSpec]) -> JsonRpcResponse {
    let id = request.id.clone().unwrap_or(Value::Null);
    let tools: Vec<Value> = tool_specs
        .iter()
        .map(|spec| {
            serde_json::json!({
                "name": spec.name,
                "description": spec.description,
                "inputSchema": spec.parameters,
            })
        })
        .collect();
    JsonRpcResponse::success(id, serde_json::json!({ "tools": tools }))
}

/// Handle `tools/call` — execute a tool and return the result.
async fn handle_tools_call(request: &JsonRpcRequest, tools: &ToolRegistry) -> JsonRpcResponse {
    let id = request.id.clone().unwrap_or(Value::Null);
    let tool_name = request
        .params
        .get("name")
        .and_then(|v| v.as_str())
        .unwrap_or("");
    let arguments = request
        .params
        .get("arguments")
        .cloned()
        .unwrap_or(serde_json::json!({}));

    if tool_name.is_empty() {
        return JsonRpcResponse::error(id, -32602, "Missing tool name".to_string());
    }

    match tools.get(tool_name) {
        Some(tool) => match tool.execute(arguments).await {
            Ok(text) => JsonRpcResponse::success(
                id,
                serde_json::json!({
                    "content": [{"type": "text", "text": text}]
                }),
            ),
            Err(e) => JsonRpcResponse::success(
                id,
                serde_json::json!({
                    "isError": true,
                    "content": [{"type": "text", "text": e.to_string()}]
                }),
            ),
        },
        None => JsonRpcResponse::error(id, -32602, format!("Tool not found: {tool_name}")),
    }
}

// ---------------------------------------------------------------------------
// McpServerRunner — stdio transport loop
// ---------------------------------------------------------------------------

/// Runs the MCP server stdio loop: reads JSON-RPC from stdin, dispatches,
/// and writes responses to stdout.
pub struct McpServerRunner {
    tool_specs: Vec<ToolSpec>,
    tools: ToolRegistry,
}

impl McpServerRunner {
    /// Create a new runner from a [`ToolRegistry`].
    ///
    /// The tool specs are extracted immediately so they can be served
    /// without holding a borrow on the registry.
    pub fn new(tools: ToolRegistry) -> Self {
        let tool_specs = tools.specs();
        Self { tool_specs, tools }
    }

    /// Run the stdio server loop until EOF on stdin.
    pub async fn run(&self) -> Result<()> {
        let stdin = tokio::io::stdin();
        let stdout = tokio::io::stdout();
        self.run_on(BufReader::new(stdin), stdout).await
    }

    /// Run the server loop on generic reader/writer (testable).
    pub async fn run_on<R, W>(&self, reader: R, writer: W) -> Result<()>
    where
        R: tokio::io::AsyncBufRead + Unpin,
        W: tokio::io::AsyncWrite + Unpin,
    {
        let mut lines = BufReader::new(reader).lines();
        let mut out = writer;

        while let Some(line) = lines.next_line().await? {
            let line = line.trim().to_string();
            if line.is_empty() {
                continue;
            }

            // Parse the request
            let request: JsonRpcRequest = match serde_json::from_str(&line) {
                Ok(req) => req,
                Err(_) => {
                    let resp = JsonRpcResponse::parse_error();
                    let json = serde_json::to_string(&resp)?;
                    out.write_all(json.as_bytes()).await?;
                    out.write_all(b"\n").await?;
                    out.flush().await?;
                    continue;
                }
            };

            // Dispatch
            if let Some(response) = dispatch_request(&request, &self.tool_specs, &self.tools).await
            {
                let json = serde_json::to_string(&response)?;
                out.write_all(json.as_bytes()).await?;
                out.write_all(b"\n").await?;
                out.flush().await?;
            }
        }

        Ok(())
    }
}

/// Manages the lifecycle of one or more MCP servers and their tools.
///
/// Call [`McpServerManager::register_all`] to spawn every configured server,
/// discover its tools, and register them into a [`ToolRegistry`]. The manager
/// keeps the underlying [`McpClient`]s alive so they can handle tool calls.
pub struct McpServerManager {
    /// Configured servers.
    servers: Vec<McpServer>,
    /// Running clients, keyed by server name.
    clients: Mutex<HashMap<String, Arc<Mutex<McpClient>>>>,
}

impl McpServerManager {
    /// Create a new manager from a list of server configurations.
    ///
    /// The servers are not spawned until [`register_all`](Self::register_all) is called.
    pub fn new(servers: Vec<McpServer>) -> Self {
        Self {
            servers,
            clients: Mutex::new(HashMap::new()),
        }
    }

    /// Spawn all configured servers, discover their tools, and register them
    /// into the given [`ToolRegistry`].
    ///
    /// Returns a list of `(server_name, tool_count)` pairs for logging.
    ///
    /// # Errors
    ///
    /// Returns an error if any server fails to start or if tool discovery
    /// fails. Servers are started sequentially; a failure stops the process.
    #[instrument(skip_all, name = "mcp.register_all")]
    pub async fn register_all(&self, registry: &mut ToolRegistry) -> Result<Vec<(String, usize)>> {
        let mut results = Vec::new();

        for server in &self.servers {
            let name = server.name.clone();
            info!(server = %name, "Starting MCP server");

            let client = McpClient::spawn(server).await.map_err(|e| Error::Tool {
                name: format!("mcp_server:{name}"),
                call_id: None,
                message: format!("Failed to start MCP server: {e}"),
            })?;

            let client = Arc::new(Mutex::new(client));

            // Discover tools from this server.
            let tool_specs = client
                .lock()
                .await
                .list_tools()
                .await
                .map_err(|e| Error::Tool {
                    name: format!("mcp_server:{name}"),
                    call_id: None,
                    message: format!("Failed to discover tools from MCP server: {e}"),
                })?;

            let tool_count = tool_specs.len();
            info!(
                server = %name,
                count = tool_count,
                "Discovered MCP tools"
            );

            // Wrap each tool spec in an McpTool and register it.
            for spec in &tool_specs {
                let tool = McpTool::new(client.clone(), spec.clone(), &name);
                registry.register_mut(Arc::new(tool));
                info!(
                    server = %name,
                    tool = %spec.name,
                    "Registered MCP tool"
                );
            }

            // Store the client so it stays alive.
            self.clients.lock().await.insert(name.clone(), client);
            results.push((name, tool_count));
        }

        Ok(results)
    }

    /// Shut down all running MCP clients.
    ///
    /// This drops the clients, which causes their background tasks (stdio
    /// reader/writer, SSE listener) to be cancelled.
    pub async fn shutdown(&self) {
        let mut clients = self.clients.lock().await;
        let names: Vec<String> = clients.keys().cloned().collect();
        for name in &names {
            info!(server = %name, "Shutting down MCP server");
            clients.remove(name);
        }
    }

    /// Return the number of currently running clients.
    pub async fn running_count(&self) -> usize {
        self.clients.lock().await.len()
    }
}

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

    #[test]
    fn success_response_has_result_and_no_error() {
        // kills function-level replacement of success() or field swaps
        use serde_json::{json, Value};
        let r = JsonRpcResponse::success(Value::from(1i64), json!({"ok": true}));
        assert!(r.result.is_some(), "success must have a result");
        assert!(r.error.is_none(), "success must not have an error");
        assert_eq!(r.jsonrpc, "2.0");
    }

    #[test]
    fn parse_error_uses_negative_code() {
        // kills `delete -` mutation in -32700
        let r = JsonRpcResponse::parse_error();
        let err = r.error.expect("parse_error must have an error");
        assert_eq!(err.code, -32700, "parse error code must be -32700");
    }

    #[test]
    fn method_not_found_uses_correct_negative_code() {
        // kills `delete -` mutation in -32601
        use serde_json::Value;
        let r = JsonRpcResponse::method_not_found(Value::from(42i64));
        let err = r.error.expect("method_not_found must have an error");
        assert_eq!(err.code, -32601, "method not found code must be -32601");
    }

    /// Test that an empty server list produces no registrations.
    #[tokio::test]
    async fn empty_servers_registers_nothing() {
        let manager = McpServerManager::new(vec![]);
        let mut registry = ToolRegistry::local();
        let results = manager.register_all(&mut registry).await.unwrap();
        assert!(results.is_empty());
        assert!(registry.names().is_empty());
    }

    /// Test that tool names are correctly namespaced.
    #[test]
    fn tool_name_format() {
        let server = "my-server";
        let tool = "Read";
        let namespaced = format!("mcp__{}__{}", server, tool);
        assert_eq!(namespaced, "mcp__my-server__Read");
    }

    /// Test that a server config with a URL creates an SSE-based client
    /// (as opposed to stdio). This is a config-level test only.
    #[test]
    fn sse_server_config_detection() {
        let server = McpServer {
            name: "test-sse".into(),
            command: String::new(),
            args: vec![],
            url: Some("http://localhost:8080/sse".into()),
            env: None,
        };
        // The McpClient::spawn method checks server.url.is_some()
        // to decide transport. We verify the config is wired correctly.
        assert!(server.url.is_some());
        assert!(server.command.is_empty());
    }

    /// Test that a server config with a command creates a stdio-based client.
    #[test]
    fn stdio_server_config_detection() {
        let server = McpServer {
            name: "test-stdio".into(),
            command: "echo".into(),
            args: vec!["hello".into()],
            url: None,
            env: None,
        };
        assert!(!server.command.is_empty());
        assert!(server.url.is_none());
    }

    #[test]
    fn error_response_has_error_and_no_result() {
        // kills field swap mutations in JsonRpcResponse::error()
        let r = JsonRpcResponse::error(Value::from(1i64), -32600, "bad request".to_string());
        assert!(r.error.is_some(), "error response must have error field");
        assert!(
            r.result.is_none(),
            "error response must not have result field"
        );
        let err = r.error.unwrap();
        assert_eq!(err.code, -32600);
        assert_eq!(err.message, "bad request");
    }

    #[tokio::test]
    async fn dispatch_notifications_returns_none() {
        // kills `=> None` → `=> Some(...)` mutation for notifications/initialized
        let reg = ToolRegistry::local();
        let specs = reg.specs();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: None,
            method: "notifications/initialized".to_string(),
            params: Value::Null,
        };
        let result = dispatch_request(&request, &specs, &reg).await;
        assert!(
            result.is_none(),
            "notifications/initialized must return None (no response)"
        );
    }

    #[tokio::test]
    async fn dispatch_unknown_method_returns_method_not_found() {
        // kills `_ => Some(JsonRpcResponse::method_not_found(...))` mutation
        let reg = ToolRegistry::local();
        let specs = reg.specs();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::from(99i64)),
            method: "totally/unknown".to_string(),
            params: Value::Null,
        };
        let result = dispatch_request(&request, &specs, &reg).await;
        assert!(result.is_some(), "unknown method must return Some");
        let r = result.unwrap();
        let err = r.error.expect("unknown method must return error");
        assert_eq!(err.code, -32601, "must use -32601 for method not found");
    }

    #[tokio::test]
    async fn dispatch_initialize_returns_protocol_version() {
        // kills `"initialize" => Some(handle_initialize(...))` removal mutation
        let reg = ToolRegistry::local();
        let specs = reg.specs();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::from(1i64)),
            method: "initialize".to_string(),
            params: Value::Null,
        };
        let result = dispatch_request(&request, &specs, &reg).await;
        assert!(result.is_some(), "initialize must return Some response");
        let r = result.unwrap();
        assert!(r.error.is_none(), "initialize must not return error");
        let res = r.result.expect("initialize must have result");
        assert_eq!(
            res["protocolVersion"], "2024-11-05",
            "initialize result must contain protocolVersion"
        );
    }

    #[tokio::test]
    async fn dispatch_tools_list_returns_tool_names() {
        // kills `"tools/list" => Some(handle_tools_list(...))` removal mutation
        // ToolRegistry::local() is empty; build a real registry via the standard builder.
        let tmp = tempfile::TempDir::new().unwrap();
        let reg = crate::tools::build_standard_tools(tmp.path(), &[], 30);
        let specs = reg.specs();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::from(2i64)),
            method: "tools/list".to_string(),
            params: Value::Null,
        };
        let result = dispatch_request(&request, &specs, &reg).await;
        assert!(result.is_some(), "tools/list must return Some response");
        let r = result.unwrap();
        let res = r.result.expect("tools/list must have result");
        let tools = res["tools"]
            .as_array()
            .expect("result must have tools array");
        assert!(
            !tools.is_empty(),
            "tools/list must return at least one tool"
        );
        let has_read = tools.iter().any(|t| t["name"].as_str() == Some("Read"));
        assert!(has_read, "tools/list must include 'Read' tool");
    }

    #[tokio::test]
    async fn handle_tools_call_missing_name_returns_error() {
        // kills `if tool_name.is_empty()` guard removal mutation
        let reg = ToolRegistry::local();
        let specs = reg.specs();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::from(3i64)),
            method: "tools/call".to_string(),
            params: serde_json::json!({ "arguments": {} }), // no "name" field
        };
        let result = dispatch_request(&request, &specs, &reg).await;
        let r = result.expect("tools/call must return Some");
        let err = r.error.expect("missing tool name must return error");
        assert_eq!(err.code, -32602, "missing name must use -32602");
    }

    #[tokio::test]
    async fn handle_tools_call_unknown_tool_uses_negative_error_code() {
        // kills `delete - in handle_tools_call` mutation (changing -32602 → 32602)
        let tmp = tempfile::TempDir::new().unwrap();
        let reg = crate::tools::build_standard_tools(tmp.path(), &[], 30);
        let specs = reg.specs();
        let request = JsonRpcRequest {
            jsonrpc: "2.0".to_string(),
            id: Some(Value::from(4i64)),
            method: "tools/call".to_string(),
            params: serde_json::json!({ "name": "no_such_tool_xyz", "arguments": {} }),
        };
        let result = dispatch_request(&request, &specs, &reg).await;
        let r = result.expect("tools/call must return Some for unknown tool");
        let err = r.error.expect("unknown tool must return an error response");
        assert_eq!(
            err.code, -32602,
            "unknown tool error code must be -32602 (negative); got: {}",
            err.code
        );
        assert!(
            err.message.contains("no_such_tool_xyz"),
            "error message must mention the tool name; got: {}",
            err.message
        );
    }
}