zesdex/crates/zesdex-backend/src/app/mcp/manager.rs

//! MCP server connection management: spawning/talking to stdio child
//! processes and HTTP endpoints, and adapting their advertised tools to
//! the crate's `Tool` trait.
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use std::sync::{Arc, Mutex};

use super::transport::{call_via_http, call_via_stdio, mcp_static_str, spawn_stdio_child};

// ---------------------------------------------------------------------------
// MCP server descriptor
// ---------------------------------------------------------------------------

/// A connected MCP server: its transport, advertised tools, and (for stdio)
/// a live handle to the child process.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct McpServer {
    pub name: String,
    pub transport: McpTransport,
    pub tools: Vec<McpToolInfo>,
    /// Held child-process handle so subsequent tool calls reuse the same
    /// connection instead of spawning a new child each time.  Not serialized
    /// because the child only lives in this process.
    #[serde(skip)]
    pub child_handle: Option<Arc<Mutex<StdioChild>>>,
}

// ---------------------------------------------------------------------------
// Tool adapter
// ---------------------------------------------------------------------------

/// Adapts a single MCP-advertised tool to the crate's `Tool` trait so it can
/// be dispatched through the same execution path as built-in tools.
pub struct McpToolAdapter {
    pub tool_name: String,
    pub server_name: String,
    pub transport: McpTransport,
    pub description: String,
    pub parameters: Value,
    /// Shared handle to a persistent child process (stdio transport only).
    pub child_handle: Option<Arc<Mutex<StdioChild>>>,
}

impl crate::tool::Tool for McpToolAdapter {
    fn name(&self) -> &'static str {
        mcp_static_str(&format!("mcp__{}__{}", self.server_name, self.tool_name))
    }

    fn description(&self) -> &'static str {
        mcp_static_str(&self.description)
    }

    fn parameters(&self) -> Value {
        self.parameters.clone()
    }

    fn run(&self, _ctx: &crate::tool::ToolCtx, args: &Value) -> anyhow::Result<String> {
        match &self.transport {
            McpTransport::Stdio {
                command,
                args: extra_args,
            } => call_via_stdio(
                self.child_handle.as_ref().map(std::convert::AsRef::as_ref),
                command,
                extra_args,
                &self.tool_name,
                args,
            ),
            McpTransport::StreamableHttp { url } => call_via_http(url, &self.tool_name, args),
        }
    }
}

// ---------------------------------------------------------------------------
// Manager
// ---------------------------------------------------------------------------

/// Registry of connected MCP servers and their tools for the current session.
#[derive(Debug, Clone)]
pub struct McpManager {
    pub servers: Vec<McpServer>,
}

impl McpManager {
    /// Create an empty manager with no connected servers.
    pub fn new() -> Self {
        McpManager {
            servers: Vec::new(),
        }
    }

    /// Flatten all connected servers' tools into a single list of `Tool` trait objects.
    ///
    /// Flow: for each server, clone its child handle → wrap each of its
    /// `McpToolInfo` entries in an `McpToolAdapter` sharing that handle.
    ///
    /// Why: the handle is cloned (Arc) per tool so every adapter for a given
    /// stdio server reuses the same persistent child process/connection.
    ///
    /// Return: boxed `Tool` trait objects ready to merge into the harness's tool list.
    pub fn as_tools(&self) -> Vec<Box<dyn crate::tool::Tool>> {
        self.servers
            .iter()
            .flat_map(|server| {
                let handle = server.child_handle.clone();
                server.tools.iter().map(move |info| {
                    let adapter: Box<dyn crate::tool::Tool> = Box::new(McpToolAdapter {
                        tool_name: info.name.clone(),
                        server_name: server.name.clone(),
                        transport: server.transport.clone(),
                        description: info.description.clone(),
                        parameters: info.input_schema.clone(),
                        child_handle: handle.clone(),
                    });
                    adapter
                })
            })
            .collect()
    }

    /// Connects to an MCP server via stdio by spawning the child process, running
    /// the `initialize` handshake, calling `tools/list`, and registering the server
    /// with its advertised tools in `self.servers`. The child process stays alive
    /// for subsequent `tools/call` invocations via the stored `McpServer.tools`.
    pub fn connect_stdio(
        &mut self,
        name: &str,
        command: &str,
        extra_args: &[String],
    ) -> anyhow::Result<()> {
        let transport = McpTransport::Stdio {
            command: command.to_string(),
            args: extra_args.to_vec(),
        };

        let mut child = spawn_stdio_child(command, extra_args)?;
        let result = child.call("tools/list", &json!({}))?;

        let tools = if let Some(tool_list) = result.get("tools").and_then(|v| v.as_array()) {
            tool_list
                .iter()
                .filter_map(|t| {
                    Some(McpToolInfo {
                        name: t.get("name")?.as_str()?.to_string(),
                        description: t
                            .get("description")
                            .and_then(|v| v.as_str())
                            .unwrap_or_else(|| {
                                tracing::warn!(
                                    "[mcp] tool {} missing description",
                                    t.get("name").and_then(|n| n.as_str()).unwrap_or("?")
                                );
                                ""
                            })
                            .to_string(),
                        input_schema: t.get("inputSchema").cloned().unwrap_or_else(|| {
                            tracing::warn!(
                                "[mcp] tool {} missing inputSchema",
                                t.get("name").and_then(|n| n.as_str()).unwrap_or("?")
                            );
                            serde_json::Value::Null
                        }),
                    })
                })
                .collect()
        } else {
            Vec::new()
        };

        let handle = Arc::new(Mutex::new(child));

        self.servers.push(McpServer {
            name: name.to_string(),
            transport,
            tools,
            child_handle: Some(handle),
        });

        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Re-exports
// ---------------------------------------------------------------------------

pub use super::transport::{McpTransport, McpToolInfo, StdioChild};