use std::path::PathBuf;
use std::sync::Arc;
use aws_lc_rs::digest::{SHA256, digest};
use base64::Engine as _;
use futures_util::future::BoxFuture;
use serde_json::{Map, Value, json};
use tokio_util::sync::CancellationToken;
use yapi_agent::{Tool, UpdateSink};
use yapi_types::event::ToolResult;
use yapi_types::message::{ContentBlock, ToolDeclaration, blocks_text};
use super::client::{RequestOptions, ResourceKind};
use super::connection::Connection;
use super::content::{block_to_llm, to_llm_content};
use crate::tools::truncate::{format_size, truncate_middle};
const MAX_TOOL_NAME_LENGTH: usize = 64;
pub const OUTPUT_MAX_BYTES: usize = 20 * 1024;
pub const READ_MCP_RESOURCE_TOOL: &str = "read_mcp_resource";
pub const LIST_MCP_RESOURCES_TOOL: &str = "list_mcp_resources";
pub const LIST_MCP_RESOURCE_TEMPLATES_TOOL: &str = "list_mcp_resource_templates";
pub fn tool_name(server: &str, tool: &str, is_taken: impl Fn(&str) -> bool) -> String {
let name: String = format!("mcp__{server}__{tool}")
.chars()
.map(|c| {
if c.is_ascii_alphanumeric() || c == '_' {
c
} else {
'_'
}
})
.collect();
if name.len() <= MAX_TOOL_NAME_LENGTH && !is_taken(&name) {
return name;
}
let digest = digest(&SHA256, format!("{server}\0{tool}").as_bytes());
let hash = crate::time::hex(&digest.as_ref()[..4]);
format!(
"{}_{hash}",
&name[..name.len().min(MAX_TOOL_NAME_LENGTH - hash.len() - 1)]
)
}
pub fn server_tool_label(tools: &crate::extensions::Tools, name: &str) -> Option<String> {
tools.with(|registry| {
let tool = registry.get(name)?;
let namespace = tool.namespace.as_ref()?;
namespace
.name
.starts_with(&super::config::namespace(""))
.then(|| tool.tool.label().to_owned())
})
}
async fn save_temp(data: &[u8], extension: &str) -> Result<PathBuf, String> {
let path = std::env::temp_dir().join(format!(
"yapi-mcp-{}{extension}",
crate::time::random_hex(8)
));
let mut options = tokio::fs::OpenOptions::new();
options.write(true).create_new(true);
#[cfg(unix)]
options.mode(0o600);
let mut file = options
.open(&path)
.await
.map_err(|error| error.to_string())?;
tokio::io::AsyncWriteExt::write_all(&mut file, data)
.await
.map_err(|error| error.to_string())?;
Ok(path)
}
pub async fn limit_content(content: Vec<ContentBlock>) -> (Vec<ContentBlock>, Option<PathBuf>) {
let combined = blocks_text(&content, "\n");
let truncation = truncate_middle(&combined, OUTPUT_MAX_BYTES);
if !truncation.truncated {
return (content, None);
}
let (where_, path) = match save_temp(combined.as_bytes(), ".txt").await {
Ok(path) => (
format!(
"[Full output: {} (read it with offset/limit)]",
path.display()
),
Some(path),
),
Err(error) => (format!("[Could not save the full output: {error}]"), None),
};
let tokens = truncation.total_bytes.div_ceil(4);
let mut limited = vec![ContentBlock::text(format!(
"Warning: truncated output (original token count: {tokens})\nTotal output lines: {}\n\n{}\n\n{where_}",
truncation.total_lines, truncation.content
))];
limited.extend(
content
.into_iter()
.filter(|block| matches!(block, ContentBlock::Image(_))),
);
(limited, path)
}
fn extension_of(uri: &str) -> String {
let path = reqwest::Url::parse(uri)
.map(|url| url.path().to_owned())
.unwrap_or_else(|_| uri.to_owned());
let Some(dot) = path.rfind('.') else {
return ".bin".into();
};
let suffix = &path[dot + 1..];
if (1..=8).contains(&suffix.len()) && suffix.chars().all(|c| c.is_ascii_alphanumeric()) {
format!(".{suffix}")
} else {
".bin".into()
}
}
fn is_text_mime(mime: Option<&str>) -> bool {
let Some(mime) = mime else {
return false;
};
let kind = mime
.split(';')
.next()
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
kind.starts_with("text/")
|| kind == "application/json"
|| kind.ends_with("+json")
|| kind.ends_with("+xml")
}
fn field<'a>(value: &'a Value, key: &str) -> Option<&'a str> {
value.get(key).and_then(Value::as_str)
}
async fn block_content(server: &str, block: &Value, readable_resources: bool) -> ContentBlock {
match field(block, "type") {
Some("resource_link") => {
let mut details = Vec::new();
if let Some(mime) = field(block, "mimeType").filter(|mime| !mime.is_empty()) {
details.push(mime.to_owned());
}
if let Some(size) = block.get("size").and_then(Value::as_f64) {
details.push(format_size(size as usize));
}
let details = if details.is_empty() {
String::new()
} else {
format!(" ({})", details.join(", "))
};
let description = field(block, "description")
.filter(|text| !text.is_empty())
.map(|text| format!(": {text}"))
.unwrap_or_default();
let read = if readable_resources {
format!(". Read it with {READ_MCP_RESOURCE_TOOL} (server \"{server}\")")
} else {
String::new()
};
let title = field(block, "title")
.or(field(block, "name"))
.unwrap_or_default();
ContentBlock::text(format!(
"[Resource {} \"{title}\"{details}{description}{read}]",
field(block, "uri").unwrap_or_default()
))
}
Some("resource")
if block["resource"].get("blob").is_some()
&& !field(&block["resource"], "mimeType")
.is_some_and(|mime| mime.starts_with("image/")) =>
{
let resource = &block["resource"];
let uri = field(resource, "uri").unwrap_or_default();
let mime = field(resource, "mimeType");
let data = node_base64(field(resource, "blob").unwrap_or_default());
if is_text_mime(mime) {
return ContentBlock::text(String::from_utf8_lossy(&data));
}
let kind = format!(
"{}, {}",
mime.unwrap_or("unknown type"),
format_size(data.len())
);
match save_temp(&data, &extension_of(uri)).await {
Ok(path) => ContentBlock::text(format!(
"[Binary resource {uri} ({kind}) saved to {}]",
path.display()
)),
Err(error) => ContentBlock::text(format!(
"[Binary resource {uri} ({kind}) could not be saved: {error}]"
)),
}
}
_ => block_to_llm(block),
}
}
fn node_base64(text: &str) -> Vec<u8> {
const LENIENT: base64::engine::GeneralPurpose = base64::engine::GeneralPurpose::new(
&base64::alphabet::STANDARD,
base64::engine::GeneralPurposeConfig::new()
.with_decode_padding_mode(base64::engine::DecodePaddingMode::RequireNone)
.with_decode_allow_trailing_bits(true),
);
let mut digits: Vec<u8> = text
.bytes()
.take_while(|&byte| byte != b'=')
.filter_map(|byte| match byte {
b'-' => Some(b'+'),
b'_' => Some(b'/'),
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'+' | b'/' => Some(byte),
_ => None,
})
.collect();
if digits.len() % 4 == 1 {
digits.pop();
}
LENIENT.decode(digits).unwrap_or_default()
}
pub async fn model_content(
server: &str,
blocks: &[Value],
readable_resources: bool,
) -> Vec<ContentBlock> {
let mut content = Vec::new();
for block in blocks {
content.push(block_content(server, block, readable_resources).await);
}
content
}
async fn limit_with_details(
content: Vec<ContentBlock>,
server: &str,
tool: &str,
) -> (Vec<ContentBlock>, Value) {
let (content, path) = limit_content(content).await;
let mut details = json!({"server": server, "tool": tool});
if let Some(path) = path {
details["fullOutputPath"] = json!(path.display().to_string());
}
(content, details)
}
pub async fn convert_result(
server: &str,
tool: &str,
result: Value,
readable_resources: bool,
) -> ToolResult {
let blocks = result
.get("content")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
let mut converted = if blocks.is_empty() {
to_llm_content(&result)
} else {
model_content(server, &blocks, readable_resources).await
};
let is_error = result.get("isError").and_then(Value::as_bool) == Some(true);
if is_error && blocks_text(&converted, "\n").is_empty() {
converted.push(ContentBlock::text(format!(
"MCP tool {server}/{tool} returned an error"
)));
}
let (content, details) = limit_with_details(converted, server, tool).await;
let mut structured = result;
if let Some(object) = structured.as_object_mut() {
object.shift_remove("_meta");
}
ToolResult {
content,
details: Some(details),
structured_content: Some(structured),
is_error: is_error.then_some(true),
..ToolResult::default()
}
}
fn parameters(schema: &Map<String, Value>) -> Value {
let mut parameters = schema.clone();
if !parameters.contains_key("type") {
parameters.insert("type".into(), json!("object"));
}
if !parameters.contains_key("properties") {
parameters.insert("properties".into(), json!({}));
}
Value::Object(parameters)
}
pub fn annotations(tool: &super::Tool) -> Option<Value> {
const HINTS: [&str; 4] = [
"readOnlyHint",
"destructiveHint",
"idempotentHint",
"openWorldHint",
];
let given = tool.annotations.as_ref()?;
let hints: serde_json::Map<String, Value> = HINTS
.iter()
.filter_map(|hint| {
Some((
(*hint).to_owned(),
given.get(*hint).filter(|value| value.is_boolean())?.clone(),
))
})
.collect();
(!hints.is_empty()).then_some(Value::Object(hints))
}
pub struct McpTool {
declaration: ToolDeclaration,
output_schema: Value,
label: String,
server: String,
tool: String,
connection: Arc<Connection>,
}
impl McpTool {
pub fn new(connection: Arc<Connection>, tool: &super::Tool, name: String) -> McpTool {
let server = connection.name().to_owned();
let title = tool.title.clone().or_else(|| {
tool.annotations
.as_ref()
.and_then(|annotations| annotations.get("title"))
.and_then(Value::as_str)
.map(str::to_owned)
});
let description = tool
.description
.as_deref()
.map(str::trim)
.filter(|text| !text.is_empty())
.map(str::to_owned)
.or(title)
.unwrap_or_else(|| format!("MCP tool {} from server {server}", tool.name));
McpTool {
declaration: ToolDeclaration {
name,
description,
parameters: parameters(&tool.input_schema),
constrained_sampling: None,
},
output_schema: result_schema(tool.output_schema.as_ref()),
label: format!("{server}/{}", tool.name),
server,
tool: tool.name.clone(),
connection,
}
}
}
fn result_schema(structured: Option<&Map<String, Value>>) -> Value {
let mut properties = Map::new();
properties.insert(
"content".into(),
json!({"type": "array", "items": {"type": "object"}}),
);
if let Some(structured) = structured {
properties.insert(
"structuredContent".into(),
Value::Object(structured.clone()),
);
}
properties.insert("isError".into(), json!({"type": "boolean"}));
properties.insert("_meta".into(), json!({"type": "object"}));
json!({"type": "object", "properties": properties, "required": ["content"]})
}
impl Tool for McpTool {
fn declaration(&self) -> &ToolDeclaration {
&self.declaration
}
fn output_schema(&self) -> Option<&Value> {
Some(&self.output_schema)
}
fn label(&self) -> &str {
&self.label
}
fn execute(
&self,
_call_id: String,
args: Value,
cancel: CancellationToken,
updates: UpdateSink,
) -> BoxFuture<'_, Result<ToolResult, String>> {
Box::pin(async move {
let (server, tool) = (self.server.clone(), self.tool.clone());
let options = RequestOptions {
cancel: Some(cancel),
timeout: Some(self.connection.timeout()),
on_progress: Some(Arc::new(move |progress: &Value| {
let total = progress
.get("total")
.filter(|total| !total.is_null())
.map(|total| format!("/{total}"))
.unwrap_or_default();
let message = progress
.get("message")
.and_then(Value::as_str)
.map(str::to_owned)
.unwrap_or_else(|| format!("Progress {}{total}", progress["progress"]));
updates(ToolResult {
content: vec![ContentBlock::text(message)],
details: Some(json!({"server": server, "tool": tool})),
..ToolResult::default()
});
})),
};
let args = if args.is_null() { json!({}) } else { args };
let result = self
.connection
.call_tool(&self.tool, args, options)
.await
.map_err(|error| error.to_string())?;
Ok(convert_result(
&self.server,
&self.tool,
result,
self.connection.readable_resources(),
)
.await)
})
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ResourceTool {
List(ResourceKind),
Read,
}
pub type ResourceServers = Arc<dyn Fn() -> Vec<Arc<Connection>> + Send + Sync>;
pub struct McpResourceTool {
declaration: ToolDeclaration,
output_schema: Value,
kind: ResourceTool,
servers: ResourceServers,
}
pub fn is_app_resource(item: &Value) -> bool {
let uri = field(item, "uri")
.or_else(|| field(item, "uriTemplate"))
.unwrap_or_default();
let mime = field(item, "mimeType")
.unwrap_or_default()
.to_ascii_lowercase();
let profile = mime.split(';').skip(1).any(|parameter| {
parameter.split_once('=').is_some_and(|(key, value)| {
key.trim() == "profile" && value.trim().trim_matches('"') == "mcp-app"
})
});
uri.starts_with("ui://") || profile
}
fn listed(server: &str, item: &Value) -> Value {
let mut out = Map::new();
out.insert("server".into(), json!(server));
if let Some(object) = item.as_object() {
for (key, value) in object {
if key != "_meta" && key != "icons" {
out.insert(key.clone(), value.clone());
}
}
}
Value::Object(out)
}
fn string_argument(args: &Value, key: &str) -> Result<Option<String>, String> {
match args.get(key) {
None | Some(Value::Null) => Ok(None),
Some(Value::String(value)) => {
Ok(Some(value.trim().to_owned()).filter(|value| !value.is_empty()))
}
Some(_) => Err(format!("{key} must be a string")),
}
}
pub fn resource_tools(servers: ResourceServers) -> Vec<McpResourceTool> {
let string = |description: &str| json!({"type": "string", "description": description});
let list_parameters = json!({
"type": "object",
"properties": {
"server": string("MCP server name. Omit to list every server with resources."),
"cursor": string("Opaque cursor from a previous call with the same server; omit for the first page."),
},
"additionalProperties": false,
});
let read_parameters = json!({
"type": "object",
"properties": {
"server": string("MCP server name exactly as configured. Must match the 'server' field returned by list_mcp_resources."),
"uri": string("Resource URI to read. Must be one of the URIs returned by list_mcp_resources."),
},
"required": ["server", "uri"],
"additionalProperties": false,
});
let declare = |name: &str, description: &str, parameters: Value| ToolDeclaration {
name: name.into(),
description: description.into(),
parameters,
constrained_sampling: None,
};
let optional_string = json!({"type": "string"});
let listing_errors = json!({
"type": "array",
"description": "Servers that could not be listed",
"items": {
"type": "object",
"properties": {"server": {"type": "string"}, "error": {"type": "string"}},
"required": ["server", "error"],
},
});
let list_output = json!({
"type": "object",
"properties": {
"server": optional_string,
"resources": {
"type": "array",
"items": {
"type": "object",
"properties": {
"server": {"type": "string"},
"uri": {"type": "string"},
"name": {"type": "string"},
"title": optional_string,
"description": optional_string,
"mimeType": optional_string,
"size": {"type": "number"},
},
"required": ["server", "uri", "name"],
},
},
"nextCursor": optional_string,
"errors": listing_errors,
},
"required": ["resources"],
});
let list_templates_output = json!({
"type": "object",
"properties": {
"server": optional_string,
"resourceTemplates": {
"type": "array",
"items": {
"type": "object",
"properties": {
"server": {"type": "string"},
"uriTemplate": {"type": "string", "description": "RFC 6570 URI template"},
"name": {"type": "string"},
"title": optional_string,
"description": optional_string,
"mimeType": optional_string,
},
"required": ["server", "uriTemplate", "name"],
},
},
"nextCursor": optional_string,
"errors": listing_errors,
},
"required": ["resourceTemplates"],
});
let read_output = json!({
"type": "object",
"properties": {
"server": {"type": "string"},
"uri": {"type": "string"},
"contents": {
"type": "array",
"items": {
"anyOf": [
{
"type": "object",
"properties": {"uri": {"type": "string"}, "mimeType": optional_string, "text": {"type": "string"}},
"required": ["uri", "text"],
},
{
"type": "object",
"properties": {
"uri": {"type": "string"},
"mimeType": optional_string,
"blob": {"type": "string", "description": "base64"},
},
"required": ["uri", "blob"],
},
],
},
},
},
"required": ["server", "uri", "contents"],
});
vec![
McpResourceTool {
declaration: declare(
LIST_MCP_RESOURCES_TOOL,
"Lists resources provided by MCP servers. Resources allow servers to share data that provides context to language models, such as files, database schemas, or application-specific information. Prefer resources over web search when possible.",
list_parameters.clone(),
),
output_schema: list_output,
kind: ResourceTool::List(ResourceKind::Resources),
servers: Arc::clone(&servers),
},
McpResourceTool {
declaration: declare(
LIST_MCP_RESOURCE_TEMPLATES_TOOL,
"Lists resource templates provided by MCP servers. Parameterized resource templates allow servers to share data that takes parameters and provides context to language models, such as files, database schemas, or application-specific information. Prefer resource templates over web search when possible.",
list_parameters,
),
output_schema: list_templates_output,
kind: ResourceTool::List(ResourceKind::Templates),
servers: Arc::clone(&servers),
},
McpResourceTool {
declaration: declare(
READ_MCP_RESOURCE_TOOL,
"Read a specific resource from an MCP server given the server name and resource URI.",
read_parameters,
),
output_schema: read_output,
kind: ResourceTool::Read,
servers,
},
]
}
impl McpResourceTool {
fn find(&self, name: &str) -> Result<Arc<Connection>, String> {
let servers = (self.servers)();
if let Some(server) = servers.iter().find(|server| server.name() == name) {
return Ok(Arc::clone(server));
}
let available: Vec<&str> = servers.iter().map(|server| server.name()).collect();
Err(if available.is_empty() {
format!("MCP server \"{name}\" has no resources")
} else {
format!(
"MCP server \"{name}\" has no resources. Servers with resources: {}",
available.join(", ")
)
})
}
async fn list(
&self,
kind: ResourceKind,
args: &Value,
cancel: &CancellationToken,
) -> Result<Value, String> {
let key = kind.key();
let server_name = string_argument(args, "server")?;
let cursor = string_argument(args, "cursor")?;
let options = |server: &Connection| RequestOptions {
cancel: Some(cancel.clone()),
timeout: Some(server.timeout()),
on_progress: None,
};
if let Some(name) = server_name {
let server = self.find(&name)?;
let (items, next) = server
.resources_page(kind, cursor, options(&server))
.await
.map_err(|error| error.to_string())?;
let items: Vec<Value> = items
.iter()
.filter(|item| !is_app_resource(item))
.map(|item| listed(server.name(), item))
.collect();
let mut payload = json!({ "server": server.name(), key: items });
if let Some(next) = next {
payload["nextCursor"] = json!(next);
}
return Ok(payload);
}
if cursor.is_some() {
return Err("cursor can only be used when a server is specified".into());
}
let mut servers = (self.servers)();
servers.sort_by(|a, b| yapi_types::collate::locale_compare(a.name(), b.name()));
let mut items = Vec::new();
let mut errors = Vec::new();
for server in &servers {
match server.all_resources(kind, options(server)).await {
Ok(found) => items.extend(
found
.iter()
.filter(|item| !is_app_resource(item))
.map(|item| listed(server.name(), item)),
),
Err(error) => {
errors.push(json!({"server": server.name(), "error": error.to_string()}))
}
}
}
let mut payload = json!({ key: items });
if !errors.is_empty() {
payload["errors"] = json!(errors);
}
Ok(payload)
}
async fn read(&self, args: &Value, cancel: &CancellationToken) -> Result<ToolResult, String> {
let name = string_argument(args, "server")?.ok_or("server must be provided")?;
let uri = string_argument(args, "uri")?.ok_or("uri must be provided")?;
let server = self.find(&name)?;
let result = server
.read_resource(
&uri,
RequestOptions {
cancel: Some(cancel.clone()),
timeout: Some(server.timeout()),
on_progress: None,
},
)
.await
.map_err(|error| error.to_string())?;
let contents = result["contents"].as_array().cloned().unwrap_or_default();
let mut blocks = Vec::new();
for item in &contents {
if contents.len() > 1 {
blocks.push(json!({"type": "text", "text": format!("{}:", field(item, "uri").unwrap_or_default())}));
}
blocks.push(json!({"type": "resource", "resource": item}));
}
let converted = model_content(server.name(), &blocks, false).await;
let converted = if converted.is_empty() {
vec![ContentBlock::text(format!("Resource {uri} is empty."))]
} else {
converted
};
let (content, details) =
limit_with_details(converted, server.name(), READ_MCP_RESOURCE_TOOL).await;
let contents: Vec<Value> = contents
.into_iter()
.map(|mut item| {
if let Some(object) = item.as_object_mut() {
object.shift_remove("_meta");
}
item
})
.collect();
Ok(ToolResult {
content,
details: Some(details),
structured_content: Some(
json!({"server": server.name(), "uri": uri, "contents": contents}),
),
..ToolResult::default()
})
}
}
impl Tool for McpResourceTool {
fn declaration(&self) -> &ToolDeclaration {
&self.declaration
}
fn output_schema(&self) -> Option<&Value> {
Some(&self.output_schema)
}
fn execute(
&self,
_call_id: String,
args: Value,
cancel: CancellationToken,
_updates: UpdateSink,
) -> BoxFuture<'_, Result<ToolResult, String>> {
Box::pin(async move {
let kind = match self.kind {
ResourceTool::Read => return self.read(&args, &cancel).await,
ResourceTool::List(kind) => kind,
};
let payload = self.list(kind, &args, &cancel).await?;
let json = yapi_types::json::stringify(&payload);
let server = string_argument(&args, "server")?.unwrap_or_default();
let (content, details) = limit_with_details(
vec![ContentBlock::text(json)],
&server,
&self.declaration.name,
)
.await;
Ok(ToolResult {
content,
details: Some(details),
structured_content: Some(payload),
..ToolResult::default()
})
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn names_tools_like_pi() {
assert_eq!(
tool_name("docs", "search-pages", |_| false),
"mcp__docs__search_pages"
);
let long = tool_name("server", &"x".repeat(80), |_| false);
assert_eq!(long.len(), 64);
assert!(long.starts_with("mcp__server__xxx"));
let taken = tool_name("a", "b", |name| name == "mcp__a__b");
assert_eq!(&taken[..10], "mcp__a__b_");
assert_eq!(taken.len(), 18);
assert_eq!(extension_of("file:///tmp/archive.tar.gz"), ".gz");
assert_eq!(extension_of("blob:thing"), ".bin");
assert!(is_app_resource(&json!({"uri": "ui://widget"})));
assert!(is_app_resource(
&json!({"uri": "a", "mimeType": "text/html;profile=mcp-app"})
));
}
#[tokio::test]
async fn lenient_resource_blob() {
let text = |blob: &str| {
let block = json!({
"type": "resource",
"resource": {"uri": "file:///a.txt", "mimeType": "text/plain", "blob": blob},
});
async move { block_content("docs", &block, false).await }
};
let wrapped = "aGVsbG8g".repeat(10) + "\n" + &"aGVsbG8g".repeat(2) + "\n";
for (blob, expected) in [
("aGVsbG8", "hello".to_owned()),
(&wrapped, "hello ".repeat(12)),
("_-8=QQ==", "\u{fffd}\u{fffd}".to_owned()),
] {
assert_eq!(text(blob).await, ContentBlock::text(expected), "{blob:?}");
}
}
}