use async_trait::async_trait;
use std::sync::Arc;
use super::client::McpClient;
use crate::tools::{Tool, ToolCategory, ToolContext, ToolOutput};
pub struct McpToolWrapper {
tool_name: String,
description: String,
input_schema: serde_json::Value,
remote_name: String,
client: Arc<McpClient>,
}
impl McpToolWrapper {
pub fn new(
server_name: &str,
remote_name: &str,
description: &str,
input_schema: serde_json::Value,
client: Arc<McpClient>,
) -> Self {
Self {
tool_name: format!("{}_{}", server_name, remote_name),
description: description.to_string(),
input_schema,
remote_name: remote_name.to_string(),
client,
}
}
pub fn tool_name(&self) -> &str {
&self.tool_name
}
pub fn remote_name(&self) -> &str {
&self.remote_name
}
}
#[async_trait]
impl Tool for McpToolWrapper {
fn name(&self) -> &str {
&self.tool_name
}
fn description(&self) -> &str {
&self.description
}
fn compact_description(&self) -> &str {
self.description()
}
fn category(&self) -> ToolCategory {
ToolCategory::NetworkWrite
}
fn parameters(&self) -> serde_json::Value {
self.input_schema.clone()
}
async fn execute(
&self,
args: serde_json::Value,
_context: &ToolContext,
) -> Result<ToolOutput, crate::error::ZeptoError> {
let result = self
.client
.call_tool(&self.remote_name, args)
.await
.map_err(crate::error::ZeptoError::Mcp)?;
let text: String = result
.content
.iter()
.filter_map(|block| block.as_text())
.collect::<Vec<_>>()
.join("\n");
if result.is_error {
Ok(ToolOutput::error(if text.is_empty() {
"MCP tool returned error".to_string()
} else {
text
}))
} else {
Ok(ToolOutput::llm_only(if text.is_empty() {
"(no output)".to_string()
} else {
text
}))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn make_client() -> Arc<McpClient> {
Arc::new(McpClient::new("testserver", "http://127.0.0.1:1", 5))
}
#[test]
fn test_wrapper_tool_name_prefixed() {
let client = make_client();
let wrapper = McpToolWrapper::new(
"myserver",
"read_file",
"Read a file",
json!({"type": "object"}),
client,
);
assert_eq!(wrapper.tool_name(), "myserver_read_file");
}
#[test]
fn test_wrapper_description() {
let client = make_client();
let wrapper = McpToolWrapper::new("srv", "tool1", "A useful tool", json!({}), client);
assert_eq!(wrapper.description(), "A useful tool");
}
#[test]
fn test_wrapper_parameters() {
let client = make_client();
let schema = json!({
"type": "object",
"properties": {
"path": {"type": "string"}
},
"required": ["path"]
});
let wrapper = McpToolWrapper::new("srv", "tool1", "desc", schema.clone(), client);
assert_eq!(wrapper.parameters(), schema);
}
#[test]
fn test_wrapper_remote_name() {
let client = make_client();
let wrapper = McpToolWrapper::new("prefix", "original_name", "desc", json!({}), client);
assert_eq!(wrapper.remote_name(), "original_name");
}
#[test]
fn test_wrapper_name_trait_method() {
let client = make_client();
let wrapper = McpToolWrapper::new("server", "tool", "desc", json!({}), client);
assert_eq!(Tool::name(&wrapper), wrapper.tool_name());
}
#[test]
fn test_wrapper_tool_name_special_chars() {
let client = make_client();
let wrapper = McpToolWrapper::new("my-server.v2", "read-file", "desc", json!({}), client);
assert_eq!(wrapper.tool_name(), "my-server.v2_read-file");
}
#[test]
fn test_wrapper_empty_description() {
let client = make_client();
let wrapper = McpToolWrapper::new("srv", "tool", "", json!({}), client);
assert_eq!(wrapper.description(), "");
}
#[test]
fn test_wrapper_complex_schema() {
let client = make_client();
let schema = json!({
"type": "object",
"properties": {
"command": {"type": "string", "description": "Shell command"},
"timeout": {"type": "integer", "minimum": 1, "maximum": 600},
"env": {
"type": "object",
"additionalProperties": {"type": "string"}
}
},
"required": ["command"]
});
let wrapper = McpToolWrapper::new(
"srv",
"shell",
"Run a shell command",
schema.clone(),
client,
);
assert_eq!(wrapper.parameters(), schema);
assert_eq!(wrapper.parameters()["properties"]["timeout"]["minimum"], 1);
}
#[tokio::test]
async fn test_wrapper_execute_no_server() {
let client = make_client();
let wrapper = McpToolWrapper::new(
"srv",
"some_tool",
"desc",
json!({"type": "object"}),
client,
);
let ctx = ToolContext::new();
let result = wrapper.execute(json!({"key": "value"}), &ctx).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(
matches!(err, crate::error::ZeptoError::Mcp(_)),
"Expected Mcp error variant, got: {:?}",
err
);
}
#[test]
fn test_wrapper_tool_definition_matches() {
let client = make_client();
let schema = json!({
"type": "object",
"properties": {"path": {"type": "string"}},
"required": ["path"]
});
let wrapper = McpToolWrapper::new(
"files",
"read",
"Read a file from disk",
schema.clone(),
client,
);
let def = crate::providers::ToolDefinition {
name: wrapper.name().to_string(),
description: wrapper.description().to_string(),
parameters: wrapper.parameters(),
};
assert_eq!(def.name, "files_read");
assert_eq!(def.description, "Read a file from disk");
assert_eq!(def.parameters, schema);
}
}