reagent-rs 0.2.9

use std::sync::Arc;

use serde_json::Value;
use tokio::{
    process::Command,
    sync::{mpsc::Sender, Mutex},
};
use tracing::{info, trace};

use crate::{
    notifications::{Notification, NotificationContent},
    Tool, ToolBuilder, ToolExecutionError,
};

use super::error::McpIntegrationError;
use crate::AsyncToolFn;
use rmcp::{
    model::{CallToolRequestParam, CallToolResult, JsonObject},
    service::RunningService,
    transport::{
        ConfigureCommandExt, SseClientTransport, StreamableHttpClientTransport, TokioChildProcess,
    },
    ClientHandler, ServiceExt,
};

/// Defines the type of MCP server the agent can connect to.
///
/// This allows the agent to communicate with external tools
/// over different transports:
/// - `Sse` for Server-Sent Events
/// - `Stdio` for a child process over standard I/O
/// - `StreamableHttp` for HTTP transport with streaming
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum McpServerType {
    /// Connect via Server-Sent Events at the provided URL.
    Sse(String),
    /// Spawn and connect to a process over stdin/stdout pipes.
    Stdio(String),
    /// Connect via a streaming HTTP endpoint.
    StreamableHttp(String),
}

impl McpServerType {
    /// Creates an SSE-based MCP server type with the given URL.
    pub fn sse<S: Into<String>>(url: S) -> Self {
        McpServerType::Sse(url.into())
    }

    /// Creates an stdio-based MCP server type with the given command.
    pub fn stdio<S: Into<String>>(cmd: S) -> Self {
        McpServerType::Stdio(cmd.into())
    }

    /// Creates a streamable HTTP-based MCP server type with the given URL.
    pub fn streamable_http<S: Into<String>>(url: S) -> Self {
        McpServerType::StreamableHttp(url.into())
    }
}

/// A handle to a running MCP client instance, wrapped in an async lock.
///
/// This client manages communication with a remote MCP server.
pub type McpClient = Arc<Mutex<RunningService<rmcp::RoleClient, AgentMcpHandler>>>;

#[derive(Clone)]
pub struct AgentMcpHandler {
    /// The channel to send notifications back to the main agent.
    agent_notification_tx: Option<Sender<Notification>>,
}

impl ClientHandler for AgentMcpHandler {
    async fn on_progress(
        &self,
        params: rmcp::model::ProgressNotificationParam,
        _context: rmcp::service::NotificationContext<rmcp::RoleClient>,
    ) {
        trace!("Received progress notification: {:?}", params);
        if self.agent_notification_tx.is_none() {
            return;
        }
        let tx = self.agent_notification_tx.clone().unwrap();
        let notification_string = serde_json::to_string(&params)
            .unwrap_or_else(|e| format!("Failed to serialize MCP notification: {e}"));

        let agent_notification = Notification::new(
            "MCP".to_string(),
            NotificationContent::McpToolNotification(notification_string),
        )
        .unwrap();

        if tx.send(agent_notification).await.is_err() {
            tracing::warn!("Agent notification channel closed. Cannot forward MCP notification.");
        }
    }
}

/// Discover and register MCP tools from a given server type.
///
/// This function connects to the specified MCP server, fetches
/// the available tool definitions, and converts them into
/// agent-compatible [`Tool`] instances.
///
/// # Arguments
/// - `mcp_server_type` - The transport type and connection info for the MCP server.
/// - `notification_channel` - Optional channel to forward MCP notifications back to the agent.
///
/// # Errors
/// Returns [`McpIntegrationError`] if the connection, discovery, or tool conversion fails.
pub async fn get_mcp_tools(
    mcp_server_type: McpServerType,
    notification_channel: Option<Sender<Notification>>,
) -> Result<Vec<Tool>, McpIntegrationError> {
    let (mcp_client, mcp_raw_tools) = match mcp_server_type {
        McpServerType::Sse(url) => get_mcp_sse_tools(url, notification_channel).await?,
        McpServerType::StreamableHttp(url) => {
            get_mcp_streamable_http_tools(url, notification_channel).await?
        }
        McpServerType::Stdio(command) => get_mcp_stdio_tools(command, notification_channel).await?,
    };

    trace!(
        "Discovered {} raw tools from MCP server. Converting...",
        mcp_raw_tools.len()
    );
    let mut agent_tools = Vec::new();

    for mcp_tool_def in mcp_raw_tools {
        let arc_mcp_client = Arc::clone(&mcp_client);
        let tool_namme_cow = mcp_tool_def.name.clone();

        // MCP tool executor closure
        let executor: AsyncToolFn = Arc::new(move |args: Value| {
            let mcp_client_ref = Arc::clone(&arc_mcp_client);
            let tool_name = tool_namme_cow.clone().into_owned();

            Box::pin(async move {
                let inner_mcp_client = mcp_client_ref.lock().await;

                // call remote tool
                let result = match inner_mcp_client
                    .call_tool(CallToolRequestParam {
                        name: tool_name.clone().into(),
                        arguments: serde_json::json!(args).as_object().cloned(),
                    })
                    .await
                {
                    Ok(result) => result,
                    Err(e) => {
                        return Err(ToolExecutionError::ExecutionFailed(format!(
                            "MCP tool '{tool_name}' execution failed: {e}"
                        )))
                    }
                };

                let CallToolResult { content, is_error } = result;

                if let Some(true) = is_error {
                    return Err(ToolExecutionError::ExecutionFailed(format!(
                        "tool call failed, mcp call error: {content:#?}"
                    )));
                }

                let mut out_result = "".to_string();
                for content in content.iter() {
                    if let Some(content_text) = content.as_text() {
                        out_result = format!("{}\n{}", out_result, content_text.text);
                    }
                }
                Ok(out_result.to_string())
            })
        });

        let tool_name = mcp_tool_def.name.clone().into_owned();
        let tool_desciption = match mcp_tool_def.description {
            Some(d) => d.into_owned(),
            None => {
                return Err(McpIntegrationError::ToolConversion(
                    "No description".to_string(),
                ))
            }
        };

        let mut tool_builer = ToolBuilder::new()
            .function_name(tool_name)
            .function_description(tool_desciption)
            .executor(executor);

        // create tool description

        let input_schema_json_obj: &JsonObject = &mcp_tool_def.input_schema;

        if let Some(Value::Object(properties_map)) = input_schema_json_obj.get("properties") {
            for (prop_name, prop_schema_value) in properties_map {
                if let Value::Object(prop_details) = prop_schema_value {
                    let property_type = prop_details
                        .get("type")
                        .and_then(Value::as_str)
                        .unwrap_or("string")
                        .to_string();

                    let description = prop_details
                        .get("description")
                        .and_then(Value::as_str)
                        .unwrap_or("")
                        .to_string();

                    tool_builer = tool_builer.add_required_property(
                        prop_name.clone(),
                        property_type,
                        description,
                    );
                }
            }
        }

        let created_tool = match tool_builer.build() {
            Ok(tool) => tool,
            Err(e) => return Err(McpIntegrationError::ToolConversion(e.to_string())),
        };

        agent_tools.push(created_tool);
        info!(
            "[MCP] Registered agent tool for MCP action: {}",
            mcp_tool_def.name
        );
    }

    Ok(agent_tools)
}

/// Connect to an MCP server over SSE and fetch its tools.
///
/// Returns both a [`McpClient`] and the raw tool definitions discovered.
///
/// # Errors
/// Returns [`McpIntegrationError`] if the SSE connection or tool discovery fails.
pub async fn get_mcp_sse_tools<T>(
    url: T,
    notification_channel: Option<Sender<Notification>>,
) -> Result<(McpClient, Vec<rmcp::model::Tool>), McpIntegrationError>
where
    T: Into<String>,
{
    let transport = match SseClientTransport::start(url.into()).await {
        Ok(t) => t,
        Err(e) => return Err(McpIntegrationError::Connection(e.to_string())),
    };

    let handler = AgentMcpHandler {
        agent_notification_tx: notification_channel,
    };

    let client = match handler.serve(transport).await {
        Ok(c) => c,
        Err(e) => return Err(McpIntegrationError::Connection(e.to_string())),
    };

    let tool_list = match client.list_tools(Default::default()).await {
        Ok(l) => l,
        Err(e) => return Err(McpIntegrationError::Discovery(e.to_string())),
    };
    Ok((Arc::new(Mutex::new(client)), tool_list.tools))
}

/// Connect to an MCP server over Streamable HTTP and fetch its tools.
///
/// Returns both a [`McpClient`] and the raw tool definitions discovered.
///
/// # Errors
/// Returns [`McpIntegrationError`] if the connection or tool discovery fails.
pub async fn get_mcp_streamable_http_tools<T>(
    url: T,
    notification_channel: Option<Sender<Notification>>,
) -> Result<(McpClient, Vec<rmcp::model::Tool>), McpIntegrationError>
where
    T: Into<String>,
{
    let transport = StreamableHttpClientTransport::from_uri(url.into());

    let handler = AgentMcpHandler {
        agent_notification_tx: notification_channel,
    };

    let client = handler
        .serve(transport)
        .await
        .map_err(|e| McpIntegrationError::Connection(e.to_string()))?;
    let tool_list = client
        .list_tools(Default::default())
        .await
        .map_err(|e| McpIntegrationError::Discovery(e.to_string()))?;

    Ok((Arc::new(Mutex::new(client)), tool_list.tools))
}

/// Connect to an MCP server spawned as a child process and fetch its tools.
///
/// The child process is started with the given command string, and the
/// connection is established over stdin/stdout pipes.
///
/// Returns both a [`McpClient`] and the raw tool definitions discovered.
///
/// # Errors
/// Returns [`McpIntegrationError`] if the process fails to start or tool discovery fails.
pub async fn get_mcp_stdio_tools<T>(
    full_command: T,
    notification_channel: Option<Sender<Notification>>,
) -> Result<(McpClient, Vec<rmcp::model::Tool>), McpIntegrationError>
where
    T: Into<String>,
{
    let full_command_string = full_command.into();
    let mut command_args = full_command_string.split(" ");
    let first = command_args.next();
    if first.is_none() {
        return Err(McpIntegrationError::Connection("Invalid command.".into()));
    }
    let transport = match TokioChildProcess::new(Command::new(first.unwrap()).configure(|cmd| {
        for arg in command_args {
            cmd.arg(arg);
        }
    })) {
        Ok(t) => t,
        Err(e) => return Err(McpIntegrationError::Connection(e.to_string())),
    };

    let handler = AgentMcpHandler {
        agent_notification_tx: notification_channel,
    };

    let client = match handler.serve(transport).await {
        Ok(c) => c,
        Err(e) => return Err(McpIntegrationError::Connection(e.to_string())),
    };

    let tool_list = match client.list_tools(Default::default()).await {
        Ok(l) => l,
        Err(e) => return Err(McpIntegrationError::Discovery(e.to_string())),
    };
    Ok((Arc::new(Mutex::new(client)), tool_list.tools))
}