klieo-tools-mcp 3.4.0

//! `McpClient` — `list_tools` + `call_tool` over an `rmcp` peer.
//!
//! Wraps a connected `rmcp` [`RunningService`] (which derefs to a
//! [`Peer<RoleClient>`]). The MCP `initialize` handshake completes
//! during `serve`, so this layer holds no handshake state of its own.

use klieo_core::error::ToolError;
use rmcp::model::{CallToolRequestParams, CallToolResult, RawContent, Tool};
use rmcp::service::{RoleClient, RunningService, ServiceError};
use std::sync::Arc;

const TEXT_BLOCK_SEPARATOR: &str = "\n";
const ERROR_BLOCK_SEPARATOR: &str = " / ";

/// MCP client over a connected `rmcp` peer.
pub(crate) struct McpClient {
    inner: Arc<RunningService<RoleClient, ()>>,
}

impl McpClient {
    pub(crate) fn new(inner: Arc<RunningService<RoleClient, ()>>) -> Self {
        Self { inner }
    }

    pub(crate) async fn list_tools(&self) -> Result<Vec<Tool>, ToolError> {
        self.inner.list_all_tools().await.map_err(map_service_error)
    }

    pub(crate) async fn call_tool(
        &self,
        name: &str,
        args: serde_json::Value,
    ) -> Result<serde_json::Value, ToolError> {
        let mut params = CallToolRequestParams::new(name.to_string());
        params.arguments = wrap_arguments(args);
        let result = self
            .inner
            .call_tool(params)
            .await
            .map_err(map_service_error)?;
        interpret_tool_result(result)
    }
}

/// The MCP `arguments` field is defined only over JSON objects; any
/// non-object shape (including the "no arguments" null) maps to `None`.
fn wrap_arguments(args: serde_json::Value) -> Option<serde_json::Map<String, serde_json::Value>> {
    match args {
        serde_json::Value::Object(map) => Some(map),
        _ => None,
    }
}

fn interpret_tool_result(result: CallToolResult) -> Result<serde_json::Value, ToolError> {
    if result.is_error.unwrap_or(false) {
        let message = join_text_blocks(&result, ERROR_BLOCK_SEPARATOR);
        return Err(ToolError::Permanent(format!(
            "MCP tool returned isError: {message}"
        )));
    }
    let text = join_text_blocks(&result, TEXT_BLOCK_SEPARATOR);
    Ok(serde_json::Value::String(text))
}

fn join_text_blocks(result: &CallToolResult, separator: &str) -> String {
    result
        .content
        .iter()
        .filter_map(|block| match &block.raw {
            RawContent::Text(text) => Some(text.text.as_str()),
            _ => None,
        })
        .collect::<Vec<_>>()
        .join(separator)
}

/// Infrastructure-failure class of an `rmcp` [`ServiceError`]. None of
/// rmcp's failures are client-side argument validation, so the class
/// space deliberately excludes any `InvalidArgs`-equivalent.
enum ServiceFailure {
    Cancelled,
    Timeout,
    Transport,
}

fn map_service_error(error: ServiceError) -> ToolError {
    let class = match &error {
        ServiceError::Cancelled { .. } => ServiceFailure::Cancelled,
        ServiceError::Timeout { .. } => ServiceFailure::Timeout,
        _ => ServiceFailure::Transport,
    };
    failure_to_tool_error(class, &error.to_string())
}

/// No rmcp failure is client-side argument validation, so this never
/// yields [`ToolError::InvalidArgs`] (reserved for that).
fn failure_to_tool_error(class: ServiceFailure, display: &str) -> ToolError {
    match class {
        ServiceFailure::Cancelled => ToolError::Cancelled,
        ServiceFailure::Timeout => ToolError::Timeout,
        ServiceFailure::Transport => ToolError::Permanent(format!("MCP service error: {display}")),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rmcp::model::{Annotated, Content, RawTextContent};

    fn text_block(text: &str) -> Content {
        Annotated::new(
            RawContent::Text(RawTextContent {
                text: text.to_string(),
                meta: None,
            }),
            None,
        )
    }

    fn image_block() -> Content {
        Annotated::new(
            RawContent::Image(rmcp::model::RawImageContent {
                data: "base64".to_string(),
                mime_type: "image/png".to_string(),
                meta: None,
            }),
            None,
        )
    }

    fn success_result(content: Vec<Content>) -> CallToolResult {
        CallToolResult::success(content)
    }

    fn error_result(content: Vec<Content>) -> CallToolResult {
        CallToolResult::error(content)
    }

    #[test]
    fn object_arguments_wrap_to_some_map() {
        let wrapped = wrap_arguments(serde_json::json!({ "key": "value" }));
        let map = wrapped.expect("an object must wrap to Some");
        assert_eq!(map.get("key").and_then(|v| v.as_str()), Some("value"));
    }

    #[test]
    fn null_arguments_wrap_to_none() {
        assert!(wrap_arguments(serde_json::Value::Null).is_none());
    }

    #[test]
    fn non_object_arguments_wrap_to_none() {
        assert!(wrap_arguments(serde_json::json!("a string")).is_none());
        assert!(wrap_arguments(serde_json::json!(42)).is_none());
    }

    #[test]
    fn success_result_joins_text_blocks_with_newline() {
        let r = success_result(vec![text_block("first"), text_block("second")]);
        let value = interpret_tool_result(r).expect("success must yield a string");
        assert_eq!(value.as_str(), Some("first\nsecond"));
    }

    #[test]
    fn success_result_drops_non_text_blocks() {
        let r = success_result(vec![text_block("kept"), image_block()]);
        let value = interpret_tool_result(r).expect("success must yield a string");
        assert_eq!(value.as_str(), Some("kept"));
    }

    #[test]
    fn error_result_joins_text_blocks_with_slash_and_is_permanent() {
        let r = error_result(vec![text_block("bad input"), text_block("retry later")]);
        match interpret_tool_result(r) {
            Err(ToolError::Permanent(message)) => {
                assert!(message.contains("isError"));
                assert!(message.contains("bad input / retry later"));
            }
            other => panic!("expected Permanent, got {other:?}"),
        }
    }

    #[test]
    fn cancelled_failure_maps_to_cancelled() {
        let mapped = failure_to_tool_error(ServiceFailure::Cancelled, "task cancelled");
        assert!(matches!(mapped, ToolError::Cancelled));
    }

    #[test]
    fn timeout_failure_maps_to_timeout() {
        let mapped = failure_to_tool_error(ServiceFailure::Timeout, "request timeout");
        assert!(matches!(mapped, ToolError::Timeout));
    }

    #[test]
    fn transport_failure_maps_to_permanent_not_invalid_args() {
        let mapped = failure_to_tool_error(ServiceFailure::Transport, "Transport closed");
        match mapped {
            ToolError::Permanent(message) => {
                assert!(message.contains("MCP service error"));
                assert!(message.contains("Transport closed"));
            }
            other => panic!("transport failure must be Permanent, got {other:?}"),
        }
    }
}