mod cache;
mod config;
mod renderer;
pub mod server;
pub mod types;
use crate::agent::extension::{Extension, ToolDefinition};
use cache::{has_valid_cache, load_cache, update_cache_entry};
use renderer::{McpProxyToolRenderer, McpToolRenderer};
use server::ServerManager;
use std::borrow::Cow;
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use tokio::sync::Mutex;
use types::format_tool_name;
use yoagent::mcp::types::{McpContent, McpToolInfo};
use yoagent::types::{AgentTool, Content, ToolContext, ToolError, ToolResult};
pub use cache::{load_cache as load_metadata_cache, save_cache as save_metadata_cache};
pub use config::load_mcp_config;
pub use types::{CachedTool, McpConfig, McpSettings, MetadataCache, ServerCacheEntry, ServerEntry};
const MAX_SEARCH_RESULTS: usize = 30;
pub struct McpExtension {
config: McpConfig,
manager: Arc<Mutex<ServerManager>>,
tool_cache: Arc<Mutex<HashMap<String, Vec<McpToolInfo>>>>,
}
impl McpExtension {
pub fn new(config: McpConfig) -> Self {
let idle_timeout = config
.settings
.as_ref()
.map(|s| s.idle_timeout)
.unwrap_or(10);
let mut manager = ServerManager::new(idle_timeout);
for (name, entry) in &config.mcp_servers {
let config_hash = config::compute_server_config_hash(entry);
manager.register(name, entry.clone(), config_hash);
}
Self {
config,
manager: Arc::new(Mutex::new(manager)),
tool_cache: Arc::new(Mutex::new(HashMap::new())),
}
}
pub fn from_cwd(cwd: &Path) -> Self {
let config = load_mcp_config(cwd);
Self::new(config)
}
pub async fn restore_cache(&self) {
let cache = load_cache();
let mut tool_cache = self.tool_cache.lock().await;
for (server_name, entry) in &cache.servers {
let def = self.config.mcp_servers.get(server_name);
let ch = def.map(config::compute_server_config_hash).unwrap_or(0);
if entry.config_hash != ch {
continue;
}
let tools: Vec<McpToolInfo> = entry
.tools
.iter()
.map(|t| McpToolInfo {
name: t.name.clone(),
description: t.description.clone(),
input_schema: if t.input_schema.is_null() {
serde_json::json!({"type": "object", "properties": {}})
} else {
t.input_schema.clone()
},
})
.collect();
if !tools.is_empty() {
tool_cache.insert(server_name.clone(), tools);
}
}
}
pub async fn bootstrap_direct_tools(&self) {
let global_direct_tools = self
.config
.settings
.as_ref()
.is_some_and(|s| s.direct_tools);
let missing_cache: Vec<String> = self
.config
.mcp_servers
.iter()
.filter(|(server_name, entry)| {
let has_direct = match entry.direct_tools.as_ref() {
Some(v) if v.is_boolean() => v.as_bool().unwrap_or(false),
Some(v) if v.is_array() => true,
None => global_direct_tools,
Some(_) => false,
};
if !has_direct {
return false;
}
let config_hash = config::compute_server_config_hash(entry);
!has_valid_cache(server_name, config_hash)
})
.map(|(name, _)| name.clone())
.collect();
if !missing_cache.is_empty() {
eprintln!(
"MCP: direct tools configured for {} but no cached metadata yet. \
Connect once via the mcp proxy tool, then restart.",
missing_cache.join(", ")
);
}
}
}
impl Extension for McpExtension {
fn name(&self) -> Cow<'static, str> {
"mcp".into()
}
fn tools(&self) -> Vec<ToolDefinition> {
let mut tools: Vec<ToolDefinition> = Vec::new();
tools.push(ToolDefinition {
tool: Box::new(McpProxyTool {
config: self.config.clone(),
manager: self.manager.clone(),
tool_cache: self.tool_cache.clone(),
}),
snippet: "MCP gateway - connect to MCP servers and call their tools. Non-MCP Pi tools should be called directly, not through mcp.",
guidelines: &[
"Use mcp to connect to external MCP tool servers",
"Direct tools for configured servers can be called directly without mcp",
"The proxy tool handles connect, list, search, describe, and call operations",
],
prepare_arguments: None,
before_tool_call: None,
after_tool_call: None,
renderer: Some(std::sync::Arc::new(McpProxyToolRenderer)),
});
let global_direct_tools = self
.config
.settings
.as_ref()
.is_some_and(|s| s.direct_tools);
let cache = load_cache();
let prefix_mode = self
.config
.settings
.as_ref()
.map(|s| s.tool_prefix.as_str())
.unwrap_or("server");
for (server_name, entry) in &self.config.mcp_servers {
let direct = entry.direct_tools.as_ref();
let has_direct = match direct {
Some(v) if v.is_boolean() => v.as_bool().unwrap_or(false),
Some(v) if v.is_array() => true,
None => global_direct_tools,
Some(_) => false,
};
if !has_direct {
continue;
}
let tool_names: Vec<String> = match direct {
Some(v) if v.is_array() => v
.as_array()
.unwrap()
.iter()
.filter_map(|s| s.as_str().map(String::from))
.collect(),
_ => {
let config_hash = config::compute_server_config_hash(entry);
if !has_valid_cache(server_name, config_hash) {
continue;
}
cache
.servers
.get(server_name)
.map(|s| &s.tools)
.into_iter()
.flatten()
.map(|ct| ct.name.clone())
.collect()
}
};
if tool_names.is_empty() {
continue;
}
let cached_tools: Vec<&CachedTool> = cache
.servers
.get(server_name)
.map(|s| s.tools.iter().collect())
.unwrap_or_default();
for tool_name in &tool_names {
let prefixed = format_tool_name(tool_name, server_name, prefix_mode);
let (description, input_schema) = cached_tools
.iter()
.find(|ct| ct.name == *tool_name)
.map(|ct| {
let desc = ct
.description
.clone()
.unwrap_or_else(|| "MCP tool".to_string());
let schema = if ct.input_schema.is_null() {
serde_json::json!({"type": "object", "properties": {}})
} else {
ct.input_schema.clone()
};
(desc, schema)
})
.unwrap_or_else(|| {
(
format!("MCP tool: {} on {}", tool_name, server_name),
serde_json::json!({"type": "object", "properties": {}}),
)
});
tools.push(ToolDefinition {
tool: Box::new(McpDirectTool {
server_name: server_name.clone(),
original_name: tool_name.clone(),
display_name: prefixed.clone(),
description,
input_schema,
manager: self.manager.clone(),
}),
snippet: "MCP direct tool",
guidelines: &[],
prepare_arguments: None,
before_tool_call: None,
after_tool_call: None,
renderer: Some(std::sync::Arc::new(McpToolRenderer::new(&prefixed))),
});
}
}
tools
}
}
struct McpProxyTool {
config: McpConfig,
manager: Arc<Mutex<ServerManager>>,
tool_cache: Arc<Mutex<HashMap<String, Vec<McpToolInfo>>>>,
}
impl McpProxyTool {
async fn ensure_connected(&self, name: &str) -> bool {
let mut manager = self.manager.lock().await;
manager.ensure_connected(name).await
}
async fn cache_tools(&self, server_name: &str) {
let manager = self.manager.lock().await;
let client = manager.get_client(server_name);
drop(manager);
if let Some(client) = client {
let client = client.lock().await;
if let Ok(tools) = client.list_tools().await {
let config_hash = self
.config
.mcp_servers
.get(server_name)
.map(config::compute_server_config_hash)
.unwrap_or(0);
let mut tool_cache = self.tool_cache.lock().await;
tool_cache.insert(server_name.to_string(), tools.clone());
drop(tool_cache);
update_cache_entry(server_name, config_hash, &tools);
}
}
}
async fn call_tool(
&self,
server_name: &str,
tool_name: &str,
args: serde_json::Value,
) -> Result<(Vec<Content>, bool), String> {
let manager = self.manager.lock().await;
let client = manager.get_client(server_name);
drop(manager);
let client = match client {
Some(c) => c,
None => return Err(format!("Server '{}' is not connected", server_name)),
};
let client = client.lock().await;
let result = client
.call_tool(tool_name, args)
.await
.map_err(|e| format!("MCP call failed: {}", e))?;
let is_error = result.is_error;
let content: Vec<Content> = result
.content
.into_iter()
.map(|c| match c {
McpContent::Text { text } => Content::Text { text },
McpContent::Image { data, mime_type } => Content::Image { data, mime_type },
})
.collect();
Ok((content, is_error))
}
fn format_search_results(query: &str, matches: &[(String, McpToolInfo)]) -> String {
let mut text = format!(
"Found {} tool{} matching \"{}\":\n\n",
matches.len(),
if matches.len() == 1 { "" } else { "s" },
query
);
for (server_name, tool) in matches {
text.push_str(&format!(
"{} @ {}\n {}\n",
tool.name,
server_name,
tool.description.as_deref().unwrap_or("(no description)")
));
let schema = &tool.input_schema;
if !schema.is_null()
&& schema.is_object()
&& let Some(props) = schema.get("properties").and_then(|p| p.as_object())
&& !props.is_empty()
{
let required: std::collections::HashSet<&str> = schema
.get("required")
.and_then(|r| r.as_array())
.map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
.unwrap_or_default();
text.push_str(" Parameters:\n");
for (prop_name, prop_schema) in props {
let is_req = required.contains(prop_name.as_str());
let type_str = prop_schema
.get("type")
.and_then(|t| t.as_str())
.unwrap_or("any");
let desc = prop_schema
.get("description")
.and_then(|d| d.as_str())
.unwrap_or("");
text.push_str(&format!(
" - {} ({}){} {}\n",
prop_name,
type_str,
if is_req { " *required*" } else { "" },
if desc.is_empty() {
String::new()
} else {
format!("- {}", desc)
}
));
}
}
text.push('\n');
}
text.trim().to_string()
}
async fn execute_status(&self) -> ToolResult {
let manager = self.manager.lock().await;
let tool_cache = self.tool_cache.lock().await;
let mut lines = vec![format!(
"MCP: {} servers configured",
self.config.mcp_servers.len()
)];
lines.push(String::new());
for name in self.config.mcp_servers.keys() {
let status = manager.status(name);
let tool_count = tool_cache.get(name).map(|v| v.len()).unwrap_or(0);
let status_str = match status {
Some(server::ConnectionStatus::Connected) => "✓ connected",
Some(server::ConnectionStatus::Idle) => "○ idle",
Some(server::ConnectionStatus::Failed) => "✗ failed",
None => "○ not connected",
};
lines.push(format!("{} {} ({} tools)", status_str, name, tool_count));
}
if !self.config.mcp_servers.is_empty() {
lines.push(String::new());
lines.push(
"mcp({ server: \"name\" }) to list tools, mcp({ search: \"...\" }) to search"
.to_string(),
);
}
ToolResult {
content: vec![Content::Text {
text: lines.join("\n"),
}],
details: serde_json::json!({"mode": "status"}),
}
}
async fn execute_list(&self, server_name: &str) -> ToolResult {
let connected = {
let mut manager = self.manager.lock().await;
manager.ensure_connected(server_name).await
};
if connected {
let manager = self.manager.lock().await;
let client = manager.get_client(server_name);
drop(manager);
if let Some(client) = client {
let client = client.lock().await;
if let Ok(tools) = client.list_tools().await {
let config_hash = self
.config
.mcp_servers
.get(server_name)
.map(config::compute_server_config_hash)
.unwrap_or(0);
let mut tool_cache = self.tool_cache.lock().await;
tool_cache.insert(server_name.to_string(), tools.clone());
drop(tool_cache);
update_cache_entry(server_name, config_hash, &tools);
}
}
}
let tool_cache = self.tool_cache.lock().await;
let tools = tool_cache.get(server_name);
match tools {
Some(tool_list) if !tool_list.is_empty() => {
let mut lines = vec![format!("{} ({} tools):", server_name, tool_list.len())];
lines.push(String::new());
for tool in tool_list {
let desc = tool.description.as_deref().unwrap_or("");
let truncated = if desc.len() > 80 {
format!("{}...", &desc[..77])
} else {
desc.to_string()
};
lines.push(format!("- {}", tool.name));
if !truncated.is_empty() {
lines.push(format!(" {}", truncated));
}
}
ToolResult {
content: vec![Content::Text {
text: lines.join("\n"),
}],
details: serde_json::json!({"mode": "list", "server": server_name, "tools": tool_list.len()}),
}
}
_ => {
if self.config.mcp_servers.contains_key(server_name) {
ToolResult {
content: vec![Content::Text {
text: format!(
"Server \"{}\" has no tools (or hasn't been connected yet). Use mcp({{ connect: \"{}\" }}) to connect.",
server_name, server_name
),
}],
details: serde_json::json!({"mode": "list", "error": "no_tools", "server": server_name}),
}
} else {
ToolResult {
content: vec![Content::Text {
text: format!(
"Server \"{}\" not found. Use mcp({{}}) to see available servers.",
server_name
),
}],
details: serde_json::json!({"mode": "list", "error": "not_found"}),
}
}
}
}
}
async fn execute_search(
&self,
query: &str,
regex: bool,
filter_server: Option<&str>,
) -> ToolResult {
let tool_cache = self.tool_cache.lock().await;
let query_lower = query.to_lowercase();
let matches: Vec<(String, McpToolInfo)> = tool_cache
.iter()
.filter(|(server_name, _)| filter_server.is_none_or(|s| server_name.as_str() == s))
.flat_map(|(server_name, tools)| {
let ql = query_lower.clone();
tools.iter().filter_map(move |tool| {
let name_match = if regex {
regex::Regex::new(query)
.ok()
.is_some_and(|re| re.is_match(&tool.name))
} else {
tool.name.to_lowercase().contains(&ql)
};
let desc_match = tool.description.as_ref().is_some_and(|desc| {
if regex {
regex::Regex::new(query)
.ok()
.is_some_and(|re| re.is_match(desc))
} else {
desc.to_lowercase().contains(&ql)
}
});
if name_match || desc_match {
Some((server_name.clone(), tool.clone()))
} else {
None
}
})
})
.take(MAX_SEARCH_RESULTS)
.collect();
drop(tool_cache);
if matches.is_empty() {
return ToolResult {
content: vec![Content::Text {
text: format!("No tools matching \"{}\"", query),
}],
details: serde_json::json!({"mode": "search", "matches": [], "query": query}),
};
}
let text = McpProxyTool::format_search_results(query, &matches);
let count = matches.len();
ToolResult {
content: vec![Content::Text { text }],
details: serde_json::json!({
"mode": "search",
"matches": matches.iter().map(|(s, t)| serde_json::json!({"server": s, "tool": t.name})).collect::<Vec<_>>(),
"count": count,
"query": query,
}),
}
}
async fn execute_describe(&self, tool_name: &str) -> ToolResult {
let tool_cache = self.tool_cache.lock().await;
for (server_name, tools) in tool_cache.iter() {
for tool in tools {
if tool.name == tool_name {
let prefix = self
.config
.settings
.as_ref()
.map(|s| s.tool_prefix.as_str())
.unwrap_or("server");
let full_name = format_tool_name(&tool.name, server_name, prefix);
let mut lines = vec![
full_name,
format!("Server: {}", server_name),
String::new(),
tool.description
.clone()
.unwrap_or_else(|| "(no description)".to_string()),
String::new(),
];
let schema = &tool.input_schema;
if !schema.is_null() && schema.is_object() {
if let Some(props) = schema.get("properties").and_then(|p| p.as_object()) {
if props.is_empty() {
lines.push("Parameters: (none)".to_string());
} else {
lines.push("Parameters:".to_string());
let required: std::collections::HashSet<&str> = schema
.get("required")
.and_then(|r| r.as_array())
.map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
.unwrap_or_default();
for (prop_name, prop_schema) in props {
let type_str = prop_schema
.get("type")
.and_then(|t| t.as_str())
.unwrap_or("any");
let desc = prop_schema
.get("description")
.and_then(|d| d.as_str())
.unwrap_or("");
let req = if required.contains(prop_name.as_str()) {
" *required*"
} else {
""
};
lines.push(format!(
" - {} ({}){}{}",
prop_name,
type_str,
req,
if desc.is_empty() {
String::new()
} else {
format!(" - {}", desc)
}
));
}
}
} else {
lines.push("Parameters: (empty schema)".to_string());
}
} else {
lines.push("Parameters: (none)".to_string());
}
return ToolResult {
content: vec![Content::Text {
text: lines.join("\n"),
}],
details: serde_json::json!({
"mode": "describe",
"server": server_name,
"tool": tool_name,
}),
};
}
}
}
ToolResult {
content: vec![Content::Text {
text: format!(
"Tool \"{}\" not found. Use mcp({{ search: \"...\" }}) to find tools.",
tool_name
),
}],
details: serde_json::json!({"mode": "describe", "error": "not_found"}),
}
}
async fn execute_connect(&self, server_name: &str) -> ToolResult {
if !self.config.mcp_servers.contains_key(server_name) {
return ToolResult {
content: vec![Content::Text {
text: format!(
"Server \"{}\" not found. Use mcp({{}}) to see available servers.",
server_name
),
}],
details: serde_json::json!({"mode": "connect", "error": "not_found"}),
};
}
let connected = self.ensure_connected(server_name).await;
if connected {
self.cache_tools(server_name).await;
let mut manager = self.manager.lock().await;
manager.touch(server_name);
drop(manager);
self.execute_list(server_name).await
} else {
ToolResult {
content: vec![Content::Text {
text: format!(
"Failed to connect to \"{}\". Check the server config.",
server_name
),
}],
details: serde_json::json!({"mode": "connect", "error": "connect_failed", "server": server_name}),
}
}
}
async fn execute_call(
&self,
tool_name: &str,
args_str: Option<&str>,
server_override: Option<&str>,
) -> ToolResult {
let parsed_args: serde_json::Value = args_str
.and_then(|s| serde_json::from_str(s).ok())
.unwrap_or(serde_json::json!({}));
let prefix_mode = self
.config
.settings
.as_ref()
.map(|s| s.tool_prefix.as_str())
.unwrap_or("server");
let (server_name, original_name) = if let Some(srv) = server_override {
(srv.to_string(), tool_name.to_string())
} else {
let tool_cache = self.tool_cache.lock().await;
let mut found = None;
for (srv, tools) in tool_cache.iter() {
for tool in tools {
let prefixed = format_tool_name(&tool.name, srv, prefix_mode);
if prefixed == tool_name || tool.name == tool_name {
found = Some((srv.clone(), tool.name.clone()));
break;
}
}
if found.is_some() {
break;
}
}
match found {
Some(f) => f,
None => {
return ToolResult {
content: vec![Content::Text {
text: format!(
"Tool \"{}\" not found. Use mcp({{ search: \"...\" }}) to find tools.",
tool_name
),
}],
details: serde_json::json!({"mode": "call", "error": "tool_not_found"}),
};
}
}
};
if !self.ensure_connected(&server_name).await {
return ToolResult {
content: vec![Content::Text {
text: format!(
"Server \"{}\" is not available. Use mcp({{ connect: \"{}\" }}) to connect.",
server_name, server_name
),
}],
details: serde_json::json!({"mode": "call", "error": "server_unavailable"}),
};
}
{
let mut manager = self.manager.lock().await;
manager.touch(&server_name);
}
match self
.call_tool(&server_name, &original_name, parsed_args)
.await
{
Ok((content, is_error)) => {
let text: String = content
.iter()
.filter_map(|c| {
if let Content::Text { text } = c {
Some(text.as_str())
} else {
None
}
})
.collect::<Vec<_>>()
.join("\n");
if is_error {
ToolResult {
content: vec![Content::Text {
text: format!("Error: {}", text),
}],
details: serde_json::json!({"mode": "call", "error": "tool_error", "server": server_name}),
}
} else {
ToolResult {
content: vec![Content::Text { text }],
details: serde_json::json!({"mode": "call", "server": server_name, "tool": original_name}),
}
}
}
Err(e) => ToolResult {
content: vec![Content::Text {
text: format!("Failed to call tool: {}", e),
}],
details: serde_json::json!({"mode": "call", "error": "call_failed", "server": server_name}),
},
}
}
}
#[async_trait::async_trait]
impl AgentTool for McpProxyTool {
fn name(&self) -> &str {
"mcp"
}
fn label(&self) -> &str {
"mcp"
}
fn description(&self) -> &str {
"MCP gateway - connect to MCP servers and call their tools. Non-MCP Pi tools should be called directly, not through mcp.\n\n\
Direct tools available (call as normal tools): varies by configuration\n\n\
Servers: varies by configuration\n\n\
Usage:\n\
mcp({ }) → Show server status\n\
mcp({ server: \"name\" }) → List tools from server\n\
mcp({ search: \"query\" }) → Search MCP tools by name/description\n\
mcp({ describe: \"tool_name\" }) → Show tool details and parameters\n\
mcp({ connect: \"server-name\" }) → Connect to a server and refresh metadata\n\
mcp({ tool: \"name\", args: '{\"key\": \"value\"}' }) → Call a tool (args is JSON string)\n\
mcp({ action: \"ui-messages\" }) → Retrieve accumulated messages from completed UI sessions\n\
mcp({ action: \"auth-start\", server: \"name\" }) → Start manual OAuth and get a browser URL\n\
mcp({ action: \"auth-complete\", server: \"name\", args: '{\"redirectUrl\":\"...\"}' }) → Complete OAuth\n\n\
Mode: action > tool (call) > connect > describe > search > server (list) > nothing (status)"
}
fn parameters_schema(&self) -> serde_json::Value {
serde_json::json!({
"type": "object",
"properties": {
"server": {
"type": "string",
"description": "Server name for listing tools"
},
"tool": {
"type": "string",
"description": "Tool name to call"
},
"args": {
"type": "string",
"description": "JSON string of arguments for the tool call"
},
"connect": {
"type": "string",
"description": "Server name to connect to"
},
"describe": {
"type": "string",
"description": "Tool name to describe"
},
"search": {
"type": "string",
"description": "Search query for finding tools"
},
"regex": {
"type": "boolean",
"description": "Treat search query as regex (default: false)"
},
"includeSchemas": {
"type": "boolean",
"description": "Include parameter schemas in search results (default: true)"
},
"action": {
"type": "string",
"description": "Action: 'ui-messages', 'auth-start', or 'auth-complete'"
}
}
})
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: ToolContext,
) -> Result<ToolResult, ToolError> {
if let Some(action) = params.get("action").and_then(|v| v.as_str()) {
match action {
"ui-messages" => {
return Ok(ToolResult {
content: vec![Content::Text {
text: "No UI session messages available. (UI sessions not yet implemented in this version.)".to_string(),
}],
details: serde_json::json!({"mode": "ui-messages", "sessions": 0}),
});
}
"auth-start" => {
let server_name = params.get("server").and_then(|v| v.as_str()).unwrap_or("");
if server_name.is_empty() {
return Err(ToolError::InvalidArgs(
"Missing 'server' argument for auth-start action".into(),
));
}
return Ok(ToolResult {
content: vec![Content::Text {
text: format!(
"OAuth authentication for \"{}\" is not yet implemented in this version. \
Please start the server manually and configure authentication.",
server_name
),
}],
details: serde_json::json!({"mode": "auth-start", "error": "not_implemented"}),
});
}
"auth-complete" => {
let server_name = params.get("server").and_then(|v| v.as_str()).unwrap_or("");
if server_name.is_empty() {
return Err(ToolError::InvalidArgs(
"Missing 'server' argument for auth-complete action".into(),
));
}
return Ok(ToolResult {
content: vec![Content::Text {
text: format!(
"OAuth completion for \"{}\" is not yet implemented in this version.",
server_name
),
}],
details: serde_json::json!({"mode": "auth-complete", "error": "not_implemented"}),
});
}
_ => {
return Err(ToolError::InvalidArgs(format!(
"Unknown action '{}'. Supported: ui-messages, auth-start, auth-complete",
action
)));
}
}
}
if let Some(tool_name) = params.get("tool").and_then(|v| v.as_str()) {
let args_str = params.get("args").and_then(|v| v.as_str());
let server_override = params.get("server").and_then(|v| v.as_str());
return Ok(self
.execute_call(tool_name, args_str, server_override)
.await);
}
if let Some(server_name) = params.get("connect").and_then(|v| v.as_str()) {
return Ok(self.execute_connect(server_name).await);
}
if let Some(tool_name) = params.get("describe").and_then(|v| v.as_str()) {
return Ok(self.execute_describe(tool_name).await);
}
if let Some(query) = params.get("search").and_then(|v| v.as_str()) {
let regex = params
.get("regex")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let server_filter = params.get("server").and_then(|v| v.as_str());
return Ok(self.execute_search(query, regex, server_filter).await);
}
if let Some(server_name) = params.get("server").and_then(|v| v.as_str()) {
return Ok(self.execute_list(server_name).await);
}
Ok(self.execute_status().await)
}
}
struct McpDirectTool {
server_name: String,
original_name: String,
display_name: String,
description: String,
input_schema: serde_json::Value,
manager: Arc<Mutex<ServerManager>>,
}
#[async_trait::async_trait]
impl AgentTool for McpDirectTool {
fn name(&self) -> &str {
&self.display_name
}
fn label(&self) -> &str {
&self.display_name
}
fn description(&self) -> &str {
&self.description
}
fn parameters_schema(&self) -> serde_json::Value {
self.input_schema.clone()
}
async fn execute(
&self,
params: serde_json::Value,
_ctx: ToolContext,
) -> Result<ToolResult, ToolError> {
let mut manager = self.manager.lock().await;
if !manager.ensure_connected(&self.server_name).await {
return Err(ToolError::Failed(format!(
"Server '{}' is not available",
self.server_name
)));
}
manager.touch(&self.server_name);
let client = manager.get_client(&self.server_name).ok_or_else(|| {
ToolError::Failed(format!("Server '{}' has no client", self.server_name))
})?;
drop(manager);
let client = client.lock().await;
let result = client
.call_tool(&self.original_name, params)
.await
.map_err(|e| ToolError::Failed(format!("MCP call failed: {}", e)))?;
if result.is_error {
let error_text = result
.content
.iter()
.filter_map(|c| match c {
McpContent::Text { text } => Some(text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("\n");
return Err(ToolError::Failed(error_text));
}
let content: Vec<Content> = result
.content
.into_iter()
.map(|c| match c {
McpContent::Text { text } => Content::Text { text },
McpContent::Image { data, mime_type } => Content::Image { data, mime_type },
})
.collect();
Ok(ToolResult {
content,
details: serde_json::Value::Null,
})
}
}