1848 lines
63 KiB
Rust
1848 lines
63 KiB
Rust
use super::mcp_tool_policy::ToolFilter;
|
|
use crate::function::{Functions, ToolCallTracker};
|
|
use crate::mcp::{CatalogItem, CatalogItemKind, ConnectedServer, McpRegistry, McpServerFeatures};
|
|
|
|
use anyhow::{Context, Result, anyhow};
|
|
use bm25::{Document, Language, SearchEngineBuilder};
|
|
use rmcp::model::{
|
|
Annotations, CallToolRequestParams, CallToolResult, ContentBlock, GetPromptRequestParams,
|
|
GetPromptResult, Prompt, PromptArgument, PromptMessage, ReadResourceRequestParams,
|
|
ReadResourceResult, Resource, ResourceTemplate, Role, Tool,
|
|
};
|
|
use serde_json::{Value, json};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
|
|
#[derive(Clone)]
|
|
pub struct ToolScope {
|
|
pub functions: Functions,
|
|
pub mcp_runtime: McpRuntime,
|
|
pub tool_tracker: ToolCallTracker,
|
|
}
|
|
|
|
impl Default for ToolScope {
|
|
fn default() -> Self {
|
|
Self {
|
|
functions: Functions::default(),
|
|
mcp_runtime: McpRuntime::default(),
|
|
tool_tracker: ToolCallTracker::default(),
|
|
}
|
|
}
|
|
}
|
|
|
|
pub enum McpPromptCompletion {
|
|
Ready(Vec<(String, Option<String>)>),
|
|
PromptNames {
|
|
server: Arc<ConnectedServer>,
|
|
},
|
|
ArgumentKeys {
|
|
server: Arc<ConnectedServer>,
|
|
prompt: String,
|
|
typed_keys: Vec<String>,
|
|
},
|
|
}
|
|
|
|
#[derive(Default, Clone)]
|
|
pub struct McpRuntime {
|
|
pub servers: HashMap<String, Arc<ConnectedServer>>,
|
|
/// Per-server effective tool allowlists; a server absent here is unfiltered.
|
|
pub tool_filters: HashMap<String, ToolFilter>,
|
|
}
|
|
|
|
impl McpRuntime {
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
pub fn is_empty(&self) -> bool {
|
|
self.servers.is_empty()
|
|
}
|
|
|
|
pub fn insert(&mut self, name: String, handle: Arc<ConnectedServer>) {
|
|
self.servers.insert(name, handle);
|
|
}
|
|
|
|
pub fn get(&self, name: &str) -> Option<&Arc<ConnectedServer>> {
|
|
self.servers.get(name)
|
|
}
|
|
|
|
pub fn server_features(&self) -> Vec<McpServerFeatures> {
|
|
let mut features: Vec<McpServerFeatures> = self
|
|
.servers
|
|
.iter()
|
|
.map(|(name, handle)| {
|
|
let info = handle.peer_info();
|
|
McpServerFeatures::from_capabilities(
|
|
name.as_str(),
|
|
info.as_ref().map(|info| &info.capabilities),
|
|
)
|
|
})
|
|
.collect();
|
|
features.sort_by(|a, b| a.name.cmp(&b.name));
|
|
features
|
|
}
|
|
|
|
pub fn sync_from_registry(&mut self, registry: &McpRegistry) {
|
|
self.servers.clear();
|
|
for (name, handle) in registry.running_servers() {
|
|
self.servers.insert(name.clone(), Arc::clone(handle));
|
|
}
|
|
}
|
|
|
|
async fn catalog_items(&self, server: &str) -> Result<HashMap<String, CatalogItem>> {
|
|
let server_handle = self
|
|
.get(server)
|
|
.cloned()
|
|
.with_context(|| format!("{server} MCP server not found in runtime"))?;
|
|
let info = server_handle.peer_info();
|
|
let features = McpServerFeatures::from_capabilities(
|
|
server,
|
|
info.as_ref().map(|info| &info.capabilities),
|
|
);
|
|
let mut items = HashMap::new();
|
|
|
|
if features.tools {
|
|
match server_handle.list_all_tools().await {
|
|
Ok(mut tools) => {
|
|
if let Some(filter) = self.tool_filters.get(server) {
|
|
let advertised: Vec<String> =
|
|
tools.iter().map(|tool| tool.name.to_string()).collect();
|
|
|
|
for (source, pattern) in filter.dead_context_patterns(&advertised) {
|
|
warn!(
|
|
"MCP tool pattern '{pattern}' from {source} matches no allowed tools on server '{server}'"
|
|
);
|
|
}
|
|
|
|
tools.retain(|tool| filter.allows(&tool.name));
|
|
}
|
|
|
|
merge_catalog_items(
|
|
&mut items,
|
|
tools
|
|
.into_iter()
|
|
.map(|tool| tool_catalog_item(server, tool)),
|
|
)
|
|
}
|
|
Err(e) => warn!("Failed to list tools on MCP server {server}: {e}"),
|
|
}
|
|
}
|
|
|
|
if features.resources {
|
|
match server_handle.list_all_resources().await {
|
|
Ok(resources) => merge_catalog_items(
|
|
&mut items,
|
|
resources
|
|
.into_iter()
|
|
.map(|resource| resource_catalog_item(server, resource)),
|
|
),
|
|
Err(e) => warn!("Failed to list resources on MCP server {server}: {e}"),
|
|
}
|
|
match server_handle.list_all_resource_templates().await {
|
|
Ok(templates) => merge_catalog_items(
|
|
&mut items,
|
|
templates
|
|
.into_iter()
|
|
.map(|template| resource_template_catalog_item(server, template)),
|
|
),
|
|
Err(e) => {
|
|
warn!("Failed to list resource templates on MCP server {server}: {e}")
|
|
}
|
|
}
|
|
}
|
|
|
|
if features.prompts {
|
|
match server_handle.list_all_prompts().await {
|
|
Ok(prompts) => merge_catalog_items(
|
|
&mut items,
|
|
prompts
|
|
.into_iter()
|
|
.map(|prompt| prompt_catalog_item(server, prompt)),
|
|
),
|
|
Err(e) => warn!("Failed to list prompts on MCP server {server}: {e}"),
|
|
}
|
|
}
|
|
|
|
Ok(items)
|
|
}
|
|
|
|
pub async fn search(
|
|
&self,
|
|
server: &str,
|
|
query: &str,
|
|
top_k: usize,
|
|
) -> Result<Vec<CatalogItem>> {
|
|
let items = self.catalog_items(server).await?;
|
|
let docs = items.iter().map(|(key, item)| Document {
|
|
id: key.clone(),
|
|
contents: format!(
|
|
"{}\n{}\n{}\n{}\nserver:{}",
|
|
item.name,
|
|
item.description,
|
|
item.kind,
|
|
item.uri.as_deref().unwrap_or_default(),
|
|
item.server
|
|
),
|
|
});
|
|
let engine = SearchEngineBuilder::<String>::with_documents(Language::English, docs).build();
|
|
|
|
Ok(engine
|
|
.search(query, top_k.min(20))
|
|
.into_iter()
|
|
.filter_map(|result| items.get(&result.document.id))
|
|
.take(top_k)
|
|
.cloned()
|
|
.collect())
|
|
}
|
|
|
|
/// Best-effort audience lookup for a catalog resource: a listing failure
|
|
/// or an unknown uri (e.g. template-expanded) yields `None`.
|
|
pub async fn resource_audience(&self, server: &str, uri: &str) -> Option<Vec<String>> {
|
|
let items = self.catalog_items(server).await.ok()?;
|
|
items
|
|
.into_values()
|
|
.find(|item| item.uri.as_deref() == Some(uri))
|
|
.and_then(|item| item.audience)
|
|
}
|
|
|
|
pub async fn describe(&self, server: &str, kind: &str, tool: &str) -> Result<Value> {
|
|
let server_handle = self
|
|
.get(server)
|
|
.cloned()
|
|
.with_context(|| format!("{server} MCP server not found in runtime"))?;
|
|
|
|
match kind {
|
|
"tool" => {
|
|
if let Some(filter) = self.tool_filters.get(server)
|
|
&& !filter.allows(tool)
|
|
{
|
|
return Err(anyhow!("{tool} not found in {server} MCP server catalog"));
|
|
}
|
|
let tool_schema = server_handle
|
|
.list_all_tools()
|
|
.await?
|
|
.into_iter()
|
|
.find(|item| item.name == tool)
|
|
.ok_or_else(|| anyhow!("{tool} not found in {server} MCP server catalog"))?
|
|
.input_schema;
|
|
|
|
Ok(json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"tool": {
|
|
"type": "string",
|
|
},
|
|
"arguments": tool_schema
|
|
}
|
|
}))
|
|
}
|
|
"resource" => {
|
|
let resource = server_handle
|
|
.list_all_resources()
|
|
.await?
|
|
.into_iter()
|
|
.find(|item| item.uri == tool)
|
|
.ok_or_else(|| {
|
|
anyhow!("{tool} not found in {server} MCP server resource catalog")
|
|
})?;
|
|
|
|
Ok(json!({
|
|
"uri": resource.uri,
|
|
"name": resource.name,
|
|
"title": resource.title,
|
|
"description": resource.description,
|
|
"mime_type": resource.mime_type,
|
|
"size": resource.size,
|
|
}))
|
|
}
|
|
"resource_template" => {
|
|
let template = server_handle
|
|
.list_all_resource_templates()
|
|
.await?
|
|
.into_iter()
|
|
.find(|item| item.uri_template == tool)
|
|
.ok_or_else(|| {
|
|
anyhow!("{tool} not found in {server} MCP server resource template catalog")
|
|
})?;
|
|
|
|
Ok(json!({
|
|
"uri_template": template.uri_template,
|
|
"name": template.name,
|
|
"title": template.title,
|
|
"description": template.description,
|
|
"mime_type": template.mime_type,
|
|
"variables": uri_template_variables(&template.uri_template),
|
|
}))
|
|
}
|
|
"prompt" => {
|
|
let prompt = server_handle
|
|
.list_all_prompts()
|
|
.await?
|
|
.into_iter()
|
|
.find(|item| item.name == tool)
|
|
.ok_or_else(|| {
|
|
anyhow!("{tool} not found in {server} MCP server prompt catalog")
|
|
})?;
|
|
|
|
let arguments: Vec<Value> = prompt
|
|
.arguments
|
|
.unwrap_or_default()
|
|
.into_iter()
|
|
.map(|arg| {
|
|
json!({
|
|
"name": arg.name,
|
|
"description": arg.description,
|
|
"required": arg.required,
|
|
})
|
|
})
|
|
.collect();
|
|
|
|
Ok(json!({
|
|
"name": prompt.name,
|
|
"description": prompt.description,
|
|
"arguments": arguments,
|
|
}))
|
|
}
|
|
other => Err(anyhow!(
|
|
"Unknown kind '{other}'. Valid kinds: tool, resource, resource_template, prompt"
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub async fn invoke(
|
|
&self,
|
|
server: &str,
|
|
tool: &str,
|
|
arguments: Value,
|
|
) -> Result<CallToolResult> {
|
|
let server_handle = self
|
|
.get(server)
|
|
.cloned()
|
|
.with_context(|| format!("Invoked MCP server does not exist: {server}"))?;
|
|
|
|
if let Some(filter) = self.tool_filters.get(server)
|
|
&& !filter.allows(tool)
|
|
{
|
|
return Err(anyhow!("{tool} not found in {server} MCP server catalog"));
|
|
}
|
|
|
|
let mut request = CallToolRequestParams::new(tool.to_owned());
|
|
request.arguments = arguments.as_object().cloned();
|
|
|
|
server_handle.call_tool(request).await.map_err(Into::into)
|
|
}
|
|
|
|
pub async fn read(&self, server: &str, uri: &str) -> Result<ReadResourceResult> {
|
|
let server_handle = self
|
|
.get(server)
|
|
.cloned()
|
|
.with_context(|| format!("Read MCP server does not exist: {server}"))?;
|
|
|
|
server_handle
|
|
.read_resource(ReadResourceRequestParams::new(uri))
|
|
.await
|
|
.map_err(Into::into)
|
|
}
|
|
|
|
pub async fn list_prompts(&self, server: &str) -> Result<Vec<Prompt>> {
|
|
let server_handle = self
|
|
.get(server)
|
|
.cloned()
|
|
.with_context(|| format!("Prompt MCP server does not exist: {server}"))?;
|
|
|
|
server_handle.list_all_prompts().await.map_err(Into::into)
|
|
}
|
|
|
|
pub async fn prompt(
|
|
&self,
|
|
server: &str,
|
|
name: &str,
|
|
arguments: HashMap<String, String>,
|
|
) -> Result<GetPromptResult> {
|
|
let server_handle = self
|
|
.get(server)
|
|
.cloned()
|
|
.with_context(|| format!("Prompt MCP server does not exist: {server}"))?;
|
|
|
|
let mut request = GetPromptRequestParams::new(name.to_owned());
|
|
if !arguments.is_empty() {
|
|
request.arguments = Some(
|
|
arguments
|
|
.into_iter()
|
|
.map(|(key, value)| (key, Value::String(value)))
|
|
.collect(),
|
|
);
|
|
}
|
|
|
|
server_handle.get_prompt(request).await.map_err(Into::into)
|
|
}
|
|
|
|
pub async fn prompt_catalog(&self) -> Vec<CatalogItem> {
|
|
let mut items = Vec::new();
|
|
for features in self.server_features() {
|
|
if !features.prompts {
|
|
continue;
|
|
}
|
|
let Some(server_handle) = self.get(&features.name) else {
|
|
continue;
|
|
};
|
|
match server_handle.list_all_prompts().await {
|
|
Ok(prompts) => items.extend(
|
|
prompts
|
|
.into_iter()
|
|
.map(|prompt| prompt_catalog_item(&features.name, prompt)),
|
|
),
|
|
Err(e) => warn!(
|
|
"Failed to list prompts on MCP server {}: {e}",
|
|
features.name
|
|
),
|
|
}
|
|
}
|
|
items
|
|
}
|
|
|
|
pub fn prompt_completion(
|
|
&self,
|
|
enabled_servers: &[String],
|
|
args: &[&str],
|
|
) -> McpPromptCompletion {
|
|
match args {
|
|
[] | [_] => McpPromptCompletion::Ready(
|
|
self.server_features()
|
|
.into_iter()
|
|
.filter(|features| features.prompts && enabled_servers.contains(&features.name))
|
|
.map(|features| (features.name, None))
|
|
.collect(),
|
|
),
|
|
[server, _] => match self.get(server) {
|
|
Some(handle) => McpPromptCompletion::PromptNames {
|
|
server: Arc::clone(handle),
|
|
},
|
|
None => McpPromptCompletion::Ready(vec![]),
|
|
},
|
|
[server, prompt, rest @ ..] => match self.get(server) {
|
|
Some(handle) => McpPromptCompletion::ArgumentKeys {
|
|
server: Arc::clone(handle),
|
|
prompt: (*prompt).to_string(),
|
|
typed_keys: rest
|
|
.iter()
|
|
.filter_map(|arg| arg.split_once('=').map(|(key, _)| key.to_string()))
|
|
.collect(),
|
|
},
|
|
None => McpPromptCompletion::Ready(vec![]),
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn flatten_prompt_messages(messages: &[PromptMessage]) -> String {
|
|
messages
|
|
.iter()
|
|
.map(|message| {
|
|
let label = match message.role {
|
|
Role::User => "[user]",
|
|
Role::Assistant => "[assistant]",
|
|
};
|
|
let content = match &message.content {
|
|
ContentBlock::Text(text) => text.text.clone(),
|
|
ContentBlock::Image(_) => "[image content omitted]".to_string(),
|
|
ContentBlock::Audio(_) => "[audio content omitted]".to_string(),
|
|
_ => "[resource content omitted]".to_string(),
|
|
};
|
|
format!("{label}\n{content}")
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("\n\n")
|
|
}
|
|
|
|
pub fn format_prompt_arguments(arguments: &[PromptArgument]) -> String {
|
|
arguments
|
|
.iter()
|
|
.map(|arg| {
|
|
let name = sanitize_display_text(&arg.name);
|
|
if arg.required == Some(true) {
|
|
format!("{name} (required)")
|
|
} else {
|
|
name
|
|
}
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join(", ")
|
|
}
|
|
|
|
pub fn resolve_prompt_args(
|
|
declared: &[PromptArgument],
|
|
provided: HashMap<String, String>,
|
|
) -> (HashMap<String, String>, Vec<String>) {
|
|
let missing = declared
|
|
.iter()
|
|
.filter(|arg| arg.required == Some(true) && !provided.contains_key(&arg.name))
|
|
.map(|arg| arg.name.clone())
|
|
.collect();
|
|
(provided, missing)
|
|
}
|
|
|
|
pub fn sanitize_display_text(text: &str) -> String {
|
|
let mut sanitized = String::with_capacity(text.len());
|
|
let mut chars = text.chars().peekable();
|
|
while let Some(ch) = chars.next() {
|
|
if ch == '\u{1b}' {
|
|
match chars.next() {
|
|
// CSI: skip everything up to and including the final byte.
|
|
Some('[') => {
|
|
for next in chars.by_ref() {
|
|
if matches!(next, '\u{40}'..='\u{7e}') {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
// OSC: skip until BEL or the ESC \ string terminator.
|
|
Some(']') => {
|
|
while let Some(next) = chars.next() {
|
|
if next == '\u{07}' {
|
|
break;
|
|
}
|
|
if next == '\u{1b}' {
|
|
if chars.peek() == Some(&'\\') {
|
|
chars.next();
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
} else if ch.is_control() {
|
|
sanitized.push(' ');
|
|
} else {
|
|
sanitized.push(ch);
|
|
}
|
|
}
|
|
sanitized
|
|
}
|
|
|
|
fn audience_strings(annotations: Option<Annotations>) -> Option<Vec<String>> {
|
|
let audience = annotations?.audience?;
|
|
Some(
|
|
audience
|
|
.into_iter()
|
|
.map(|role| match role {
|
|
Role::User => "user".to_string(),
|
|
Role::Assistant => "assistant".to_string(),
|
|
})
|
|
.collect(),
|
|
)
|
|
}
|
|
|
|
fn catalog_key(item: &CatalogItem) -> String {
|
|
let id = item.uri.as_deref().unwrap_or(&item.name);
|
|
format!("{}:{id}", item.kind)
|
|
}
|
|
|
|
fn merge_catalog_items(
|
|
items: &mut HashMap<String, CatalogItem>,
|
|
new_items: impl IntoIterator<Item = CatalogItem>,
|
|
) {
|
|
for item in new_items {
|
|
items.insert(catalog_key(&item), item);
|
|
}
|
|
}
|
|
|
|
fn tool_catalog_item(server: &str, tool: Tool) -> CatalogItem {
|
|
CatalogItem {
|
|
kind: CatalogItemKind::Tool,
|
|
name: tool.name.to_string(),
|
|
server: server.to_string(),
|
|
description: tool.description.unwrap_or_default().to_string(),
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
fn resource_catalog_item(server: &str, resource: Resource) -> CatalogItem {
|
|
CatalogItem {
|
|
kind: CatalogItemKind::Resource,
|
|
name: resource.name,
|
|
server: server.to_string(),
|
|
description: resource.description.unwrap_or_default(),
|
|
uri: Some(resource.uri),
|
|
mime_type: resource.mime_type,
|
|
size: resource.size,
|
|
arguments: None,
|
|
audience: audience_strings(resource.annotations),
|
|
}
|
|
}
|
|
|
|
fn resource_template_catalog_item(server: &str, template: ResourceTemplate) -> CatalogItem {
|
|
CatalogItem {
|
|
kind: CatalogItemKind::ResourceTemplate,
|
|
name: template.name,
|
|
server: server.to_string(),
|
|
description: template.description.unwrap_or_default(),
|
|
uri: Some(template.uri_template),
|
|
mime_type: template.mime_type,
|
|
size: None,
|
|
arguments: None,
|
|
audience: audience_strings(template.annotations),
|
|
}
|
|
}
|
|
|
|
fn prompt_catalog_item(server: &str, prompt: Prompt) -> CatalogItem {
|
|
CatalogItem {
|
|
kind: CatalogItemKind::Prompt,
|
|
name: sanitize_display_text(&prompt.name),
|
|
server: server.to_string(),
|
|
description: sanitize_display_text(&prompt.description.unwrap_or_default()),
|
|
arguments: prompt.arguments,
|
|
..Default::default()
|
|
}
|
|
}
|
|
|
|
fn uri_template_variables(template: &str) -> Vec<String> {
|
|
let mut variables = Vec::new();
|
|
let mut rest = template;
|
|
while let Some(start) = rest.find('{') {
|
|
let Some(len) = rest[start + 1..].find('}') else {
|
|
break;
|
|
};
|
|
variables.push(rest[start + 1..start + 1 + len].to_string());
|
|
rest = &rest[start + len + 2..];
|
|
}
|
|
variables
|
|
}
|
|
|
|
#[cfg(test)]
|
|
pub(crate) mod test_fixtures {
|
|
use super::*;
|
|
use base64::Engine;
|
|
use base64::engine::general_purpose::STANDARD;
|
|
use rmcp::model::{
|
|
CallToolResponse, ErrorData, GetPromptResponse, ListPromptsResult,
|
|
ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, PaginatedRequestParams,
|
|
PromptsCapability, ReadResourceResponse, ResourceContents, ResourcesCapability,
|
|
ServerCapabilities, ServerInfo,
|
|
};
|
|
use rmcp::service::{RequestContext, RunningService};
|
|
use rmcp::{RoleServer, ServerHandler, ServiceExt};
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::time::Duration;
|
|
|
|
pub(crate) const FIXTURE_LOG_URI: &str = "file:///app.log";
|
|
pub(crate) const FIXTURE_LOG_TEXT: &str = "début of the log\n\
|
|
second line\n\
|
|
ERROR: disk full\n\
|
|
fourth line\n\
|
|
fifth line\n\
|
|
ERROR: café overheated\n\
|
|
seventh line\n\
|
|
eighth line";
|
|
pub(crate) const FIXTURE_BLOB_URI: &str = "file:///report.pdf";
|
|
pub(crate) const FIXTURE_BLOB_BYTES: &[u8] = &[0xff, 0xfe, 0x00, 0x88, 0x01];
|
|
pub(crate) const FIXTURE_ANNOTATED_URI: &str = "file:///annotated";
|
|
pub(crate) const FIXTURE_ANNOTATED_TEXT: &str = "annotated body";
|
|
|
|
#[derive(Clone)]
|
|
pub(crate) struct FixtureServer {
|
|
pub(crate) tools_capability: bool,
|
|
pub(crate) tool_names: Vec<&'static str>,
|
|
pub(crate) resources_capability: bool,
|
|
pub(crate) prompts_capability: bool,
|
|
pub(crate) hostile_prompt: bool,
|
|
pub(crate) fail_resource_listings: bool,
|
|
pub(crate) fail_template_listings: bool,
|
|
pub(crate) fail_prompt_listings: bool,
|
|
pub(crate) empty_resource_listings: bool,
|
|
pub(crate) empty_prompt_listings: bool,
|
|
pub(crate) fail_get_prompt: bool,
|
|
pub(crate) prompt_delay: Option<Duration>,
|
|
pub(crate) tool_result: Option<CallToolResult>,
|
|
pub(crate) list_resources_calls: Arc<AtomicUsize>,
|
|
pub(crate) list_prompts_calls: Arc<AtomicUsize>,
|
|
pub(crate) get_prompt_calls: Arc<AtomicUsize>,
|
|
pub(crate) call_tool_calls: Arc<AtomicUsize>,
|
|
}
|
|
|
|
impl Default for FixtureServer {
|
|
fn default() -> Self {
|
|
Self {
|
|
tools_capability: true,
|
|
tool_names: vec!["dup"],
|
|
resources_capability: false,
|
|
prompts_capability: false,
|
|
hostile_prompt: false,
|
|
fail_resource_listings: false,
|
|
fail_template_listings: false,
|
|
fail_prompt_listings: false,
|
|
empty_resource_listings: false,
|
|
empty_prompt_listings: false,
|
|
fail_get_prompt: false,
|
|
prompt_delay: None,
|
|
tool_result: None,
|
|
list_resources_calls: Arc::default(),
|
|
list_prompts_calls: Arc::default(),
|
|
get_prompt_calls: Arc::default(),
|
|
call_tool_calls: Arc::default(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl ServerHandler for FixtureServer {
|
|
fn get_info(&self) -> ServerInfo {
|
|
let mut capabilities = if self.tools_capability {
|
|
ServerCapabilities::builder().enable_tools().build()
|
|
} else {
|
|
ServerCapabilities::builder().build()
|
|
};
|
|
capabilities.resources = self.resources_capability.then(ResourcesCapability::default);
|
|
capabilities.prompts = self.prompts_capability.then(PromptsCapability::default);
|
|
ServerInfo::new(capabilities)
|
|
}
|
|
|
|
async fn list_tools(
|
|
&self,
|
|
_request: Option<PaginatedRequestParams>,
|
|
_context: RequestContext<RoleServer>,
|
|
) -> Result<ListToolsResult, ErrorData> {
|
|
let schema = json!({
|
|
"type": "object",
|
|
"properties": { "q": { "type": "string" } }
|
|
})
|
|
.as_object()
|
|
.cloned()
|
|
.unwrap();
|
|
Ok(ListToolsResult::with_all_items(
|
|
self.tool_names
|
|
.iter()
|
|
.map(|name| Tool::new(*name, "Fixture tool", schema.clone()))
|
|
.collect(),
|
|
))
|
|
}
|
|
|
|
async fn call_tool(
|
|
&self,
|
|
_request: CallToolRequestParams,
|
|
_context: RequestContext<RoleServer>,
|
|
) -> Result<CallToolResponse, ErrorData> {
|
|
self.call_tool_calls.fetch_add(1, Ordering::SeqCst);
|
|
match &self.tool_result {
|
|
Some(result) => Ok(CallToolResponse::Complete(result.clone())),
|
|
None => Err(ErrorData::internal_error(
|
|
"call_tool should not be reached",
|
|
None,
|
|
)),
|
|
}
|
|
}
|
|
|
|
async fn list_resources(
|
|
&self,
|
|
_request: Option<PaginatedRequestParams>,
|
|
_context: RequestContext<RoleServer>,
|
|
) -> Result<ListResourcesResult, ErrorData> {
|
|
self.list_resources_calls.fetch_add(1, Ordering::SeqCst);
|
|
if self.fail_resource_listings {
|
|
return Err(ErrorData::internal_error("resource listing exploded", None));
|
|
}
|
|
if self.empty_resource_listings {
|
|
return Ok(ListResourcesResult::with_all_items(vec![]));
|
|
}
|
|
Ok(ListResourcesResult::with_all_items(vec![
|
|
Resource::new("dup", "dup-resource")
|
|
.with_description("Duplicate-named resource")
|
|
.with_mime_type("text/plain")
|
|
.with_size(42),
|
|
Resource::new(FIXTURE_ANNOTATED_URI, "annotated-notes")
|
|
.with_description("Notes with an audience annotation")
|
|
.with_mime_type("text/plain")
|
|
.with_annotations(Annotations::default().with_audience(vec![Role::User])),
|
|
]))
|
|
}
|
|
|
|
async fn list_resource_templates(
|
|
&self,
|
|
_request: Option<PaginatedRequestParams>,
|
|
_context: RequestContext<RoleServer>,
|
|
) -> Result<ListResourceTemplatesResult, ErrorData> {
|
|
if self.fail_resource_listings || self.fail_template_listings {
|
|
return Err(ErrorData::internal_error("template listing exploded", None));
|
|
}
|
|
if self.empty_resource_listings {
|
|
return Ok(ListResourceTemplatesResult::with_all_items(vec![]));
|
|
}
|
|
Ok(ListResourceTemplatesResult::with_all_items(vec![
|
|
ResourceTemplate::new("file:///{path}/{name}", "file-template")
|
|
.with_description("Read a file")
|
|
.with_mime_type("text/plain"),
|
|
]))
|
|
}
|
|
|
|
async fn read_resource(
|
|
&self,
|
|
request: ReadResourceRequestParams,
|
|
_context: RequestContext<RoleServer>,
|
|
) -> Result<ReadResourceResponse, ErrorData> {
|
|
let uri = request.uri.as_str();
|
|
let contents = match uri {
|
|
FIXTURE_LOG_URI => vec![ResourceContents::text(FIXTURE_LOG_TEXT, uri)],
|
|
FIXTURE_BLOB_URI => vec![
|
|
ResourceContents::blob(STANDARD.encode(FIXTURE_BLOB_BYTES), uri)
|
|
.with_mime_type("application/pdf"),
|
|
],
|
|
FIXTURE_ANNOTATED_URI => vec![ResourceContents::text(FIXTURE_ANNOTATED_TEXT, uri)],
|
|
"file:///multi" => vec![
|
|
ResourceContents::text("first", "file:///multi/0"),
|
|
ResourceContents::text("second", "file:///multi/1"),
|
|
ResourceContents::text("third", "file:///multi/2"),
|
|
],
|
|
"file:///huge" => (0..3)
|
|
.map(|i| {
|
|
ResourceContents::text("x".repeat(150 * 1024), format!("file:///huge/{i}"))
|
|
})
|
|
.collect(),
|
|
"file:///docs/readme" => vec![ResourceContents::text("readme body", uri)],
|
|
_ => {
|
|
return Err(ErrorData::resource_not_found(
|
|
format!("Unknown resource: {uri}"),
|
|
None,
|
|
));
|
|
}
|
|
};
|
|
Ok(ReadResourceResponse::Complete(ReadResourceResult::new(
|
|
contents,
|
|
)))
|
|
}
|
|
|
|
async fn list_prompts(
|
|
&self,
|
|
_request: Option<PaginatedRequestParams>,
|
|
_context: RequestContext<RoleServer>,
|
|
) -> Result<ListPromptsResult, ErrorData> {
|
|
self.list_prompts_calls.fetch_add(1, Ordering::SeqCst);
|
|
if let Some(delay) = self.prompt_delay {
|
|
tokio::time::sleep(delay).await;
|
|
}
|
|
if self.fail_prompt_listings {
|
|
return Err(ErrorData::internal_error("prompt listing exploded", None));
|
|
}
|
|
if self.empty_prompt_listings {
|
|
return Ok(ListPromptsResult::with_all_items(vec![]));
|
|
}
|
|
let mut prompts = vec![Prompt::new(
|
|
"summarize",
|
|
Some("Summarize a document"),
|
|
Some(vec![
|
|
PromptArgument::new("path")
|
|
.with_description("Document path")
|
|
.with_required(true),
|
|
PromptArgument::new("style"),
|
|
]),
|
|
)];
|
|
if self.hostile_prompt {
|
|
prompts.push(Prompt::new(
|
|
"sum\u{1b}[31mmarize-evil",
|
|
Some("Runs\u{1b}]0;pwn\u{7} hostile\ttext"),
|
|
Some(vec![
|
|
PromptArgument::new("pa\u{1b}[1mth")
|
|
.with_description("Doc\u{1b}[4m path")
|
|
.with_required(true),
|
|
]),
|
|
));
|
|
}
|
|
Ok(ListPromptsResult::with_all_items(prompts))
|
|
}
|
|
|
|
async fn get_prompt(
|
|
&self,
|
|
request: GetPromptRequestParams,
|
|
_context: RequestContext<RoleServer>,
|
|
) -> Result<GetPromptResponse, ErrorData> {
|
|
self.get_prompt_calls.fetch_add(1, Ordering::SeqCst);
|
|
if let Some(delay) = self.prompt_delay {
|
|
tokio::time::sleep(delay).await;
|
|
}
|
|
if self.fail_get_prompt {
|
|
return Err(ErrorData::internal_error("get_prompt exploded", None));
|
|
}
|
|
let messages = match request.name.as_str() {
|
|
"summarize" => {
|
|
let path = request
|
|
.arguments
|
|
.as_ref()
|
|
.and_then(|args| args.get("path"))
|
|
.and_then(|value| value.as_str())
|
|
.unwrap_or_default();
|
|
vec![
|
|
PromptMessage::new_text(Role::User, format!("Summarize {path}")),
|
|
PromptMessage::new_text(Role::Assistant, "In which style?"),
|
|
PromptMessage::new_text(Role::User, "Concise."),
|
|
]
|
|
}
|
|
"hostile" => vec![
|
|
PromptMessage::new_text(Role::User, "!rm -rf /"),
|
|
PromptMessage::new_text(Role::User, ".session hijack"),
|
|
],
|
|
other => {
|
|
return Err(ErrorData::invalid_params(
|
|
format!("Unknown prompt: {other}"),
|
|
None,
|
|
));
|
|
}
|
|
};
|
|
Ok(GetPromptResponse::Complete(GetPromptResult::new(messages)))
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn add_fixture_server(
|
|
runtime: &mut McpRuntime,
|
|
name: &str,
|
|
fixture: FixtureServer,
|
|
) -> RunningService<RoleServer, FixtureServer> {
|
|
let (client_io, server_io) = tokio::io::duplex(4096);
|
|
let (server, client) = tokio::join!(fixture.serve(server_io), ().serve(client_io));
|
|
runtime.insert(name.to_string(), Arc::new(client.unwrap()));
|
|
server.unwrap()
|
|
}
|
|
|
|
pub(crate) async fn fixture_runtime(
|
|
fixture: FixtureServer,
|
|
) -> (McpRuntime, RunningService<RoleServer, FixtureServer>) {
|
|
let mut runtime = McpRuntime::new();
|
|
let server = add_fixture_server(&mut runtime, "fixture", fixture).await;
|
|
(runtime, server)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::test_fixtures::{
|
|
FIXTURE_ANNOTATED_URI, FixtureServer, add_fixture_server, fixture_runtime,
|
|
};
|
|
use super::*;
|
|
use crate::config::mcp_tool_policy::LayerSource;
|
|
use crate::function::ToolCall;
|
|
use log::{Level, LevelFilter, Log, Metadata, Record};
|
|
use std::sync::atomic::Ordering;
|
|
use std::sync::{Mutex, Once, OnceLock};
|
|
|
|
struct WarnCollector;
|
|
|
|
static WARN_MESSAGES: OnceLock<Mutex<Vec<String>>> = OnceLock::new();
|
|
|
|
fn warn_messages() -> &'static Mutex<Vec<String>> {
|
|
WARN_MESSAGES.get_or_init(Mutex::default)
|
|
}
|
|
|
|
impl Log for WarnCollector {
|
|
fn enabled(&self, metadata: &Metadata) -> bool {
|
|
metadata.level() <= Level::Warn
|
|
}
|
|
|
|
fn log(&self, record: &Record) {
|
|
if self.enabled(record.metadata()) {
|
|
warn_messages()
|
|
.lock()
|
|
.unwrap()
|
|
.push(record.args().to_string());
|
|
}
|
|
}
|
|
|
|
fn flush(&self) {}
|
|
}
|
|
|
|
fn install_warn_collector() {
|
|
static INSTALL: Once = Once::new();
|
|
INSTALL.call_once(|| {
|
|
log::set_logger(&WarnCollector).expect("no other logger should be installed");
|
|
log::set_max_level(LevelFilter::Warn);
|
|
});
|
|
}
|
|
|
|
#[test]
|
|
fn mcp_runtime_new_is_empty() {
|
|
let runtime = McpRuntime::new();
|
|
assert!(runtime.is_empty());
|
|
assert!(runtime.server_features().is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn mcp_runtime_default_is_empty() {
|
|
let runtime = McpRuntime::default();
|
|
assert!(runtime.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn mcp_runtime_get_returns_none_for_missing_server() {
|
|
let runtime = McpRuntime::new();
|
|
assert!(runtime.get("nonexistent").is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn server_features_reports_fixture_capabilities() {
|
|
let (runtime, _server) = fixture_runtime(FixtureServer {
|
|
resources_capability: true,
|
|
..Default::default()
|
|
})
|
|
.await;
|
|
|
|
let features = runtime.server_features();
|
|
assert_eq!(
|
|
features,
|
|
vec![McpServerFeatures {
|
|
name: "fixture".to_string(),
|
|
tools: true,
|
|
resources: true,
|
|
prompts: false,
|
|
}]
|
|
);
|
|
|
|
let mut functions = Functions::default();
|
|
functions.append_mcp_meta_functions(features);
|
|
assert_eq!(functions.declarations().len(), 4);
|
|
assert!(functions.contains("mcp_invoke_fixture"));
|
|
assert!(functions.contains("mcp_search_fixture"));
|
|
assert!(functions.contains("mcp_describe_fixture"));
|
|
assert!(functions.contains("mcp_read_fixture"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn server_features_without_tools_capability_gates_invoke() {
|
|
let (runtime, _server) = fixture_runtime(FixtureServer {
|
|
tools_capability: false,
|
|
resources_capability: true,
|
|
..Default::default()
|
|
})
|
|
.await;
|
|
|
|
let features = runtime.server_features();
|
|
assert_eq!(
|
|
features,
|
|
vec![McpServerFeatures {
|
|
name: "fixture".to_string(),
|
|
tools: false,
|
|
resources: true,
|
|
prompts: false,
|
|
}]
|
|
);
|
|
|
|
let mut functions = Functions::default();
|
|
functions.append_mcp_meta_functions(features);
|
|
assert_eq!(functions.declarations().len(), 3);
|
|
assert!(!functions.contains("mcp_invoke_fixture"));
|
|
assert!(functions.contains("mcp_search_fixture"));
|
|
assert!(functions.contains("mcp_describe_fixture"));
|
|
assert!(functions.contains("mcp_read_fixture"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn server_features_with_prompts_capability_emits_prompt_meta_function() {
|
|
let (runtime, _server) = fixture_runtime(FixtureServer {
|
|
resources_capability: true,
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
})
|
|
.await;
|
|
|
|
let features = runtime.server_features();
|
|
assert_eq!(
|
|
features,
|
|
vec![McpServerFeatures {
|
|
name: "fixture".to_string(),
|
|
tools: true,
|
|
resources: true,
|
|
prompts: true,
|
|
}]
|
|
);
|
|
|
|
let mut functions = Functions::default();
|
|
functions.append_mcp_meta_functions(features);
|
|
assert_eq!(functions.declarations().len(), 5);
|
|
assert!(functions.contains("mcp_prompt_fixture"));
|
|
}
|
|
|
|
#[test]
|
|
fn tool_scope_default_has_empty_mcp_runtime() {
|
|
let scope = ToolScope::default();
|
|
assert!(scope.mcp_runtime.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn tool_scope_default_has_empty_functions() {
|
|
let scope = ToolScope::default();
|
|
assert!(scope.functions.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn tool_scope_default_tracker_has_no_loops() {
|
|
let scope = ToolScope::default();
|
|
let dummy_call = ToolCall::default();
|
|
assert!(scope.tool_tracker.check_loop(&dummy_call).is_none());
|
|
}
|
|
|
|
#[test]
|
|
fn uri_template_variables_extracts_placeholders() {
|
|
assert_eq!(
|
|
uri_template_variables("file:///{path}/{name}"),
|
|
vec!["path", "name"]
|
|
);
|
|
assert!(uri_template_variables("file:///static").is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn catalog_items_keeps_tool_and_resource_with_same_id() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let items = runtime.catalog_items("fixture").await.unwrap();
|
|
|
|
assert!(items.contains_key("tool:dup"));
|
|
assert!(items.contains_key("resource:dup"));
|
|
assert!(items.contains_key("resource_template:file:///{path}/{name}"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn catalog_items_degrades_when_resource_listing_fails() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
fail_resource_listings: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let items = runtime.catalog_items("fixture").await.unwrap();
|
|
|
|
assert!(items.contains_key("tool:dup"));
|
|
assert!(!items.keys().any(|key| key.starts_with("resource")));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn catalog_items_warns_when_resource_listing_fails() {
|
|
install_warn_collector();
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
fail_resource_listings: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
runtime.catalog_items("fixture").await.unwrap();
|
|
|
|
let messages = warn_messages().lock().unwrap();
|
|
assert!(
|
|
messages
|
|
.iter()
|
|
.any(|msg| msg.contains("Failed to list resources on MCP server fixture")),
|
|
"missing resource-listing warning in: {messages:?}"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn catalog_items_skips_unadvertised_capabilities() {
|
|
let fixture = FixtureServer::default();
|
|
let resources_calls = Arc::clone(&fixture.list_resources_calls);
|
|
let prompts_calls = Arc::clone(&fixture.list_prompts_calls);
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let items = runtime.catalog_items("fixture").await.unwrap();
|
|
|
|
assert_eq!(items.keys().collect::<Vec<_>>(), vec!["tool:dup"]);
|
|
assert_eq!(resources_calls.load(Ordering::SeqCst), 0);
|
|
assert_eq!(prompts_calls.load(Ordering::SeqCst), 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn catalog_items_includes_prompts_when_advertised() {
|
|
let fixture = FixtureServer {
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let items = runtime.catalog_items("fixture").await.unwrap();
|
|
|
|
let prompt = items.get("prompt:summarize").unwrap();
|
|
assert_eq!(prompt.kind, CatalogItemKind::Prompt);
|
|
assert_eq!(prompt.description, "Summarize a document");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn search_results_carry_kind_and_uri() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let results = runtime.search("fixture", "dup", 10).await.unwrap();
|
|
|
|
let values: Vec<Value> = results
|
|
.iter()
|
|
.map(|item| serde_json::to_value(item).unwrap())
|
|
.collect();
|
|
let tool = values.iter().find(|v| v["kind"] == "tool").unwrap();
|
|
assert!(tool.get("uri").is_none());
|
|
let resource = values.iter().find(|v| v["kind"] == "resource").unwrap();
|
|
assert_eq!(resource["uri"], "dup");
|
|
assert_eq!(resource["mime_type"], "text/plain");
|
|
assert_eq!(resource["size"], 42);
|
|
assert!(resource.get("audience").is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn search_results_carry_resource_audience() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let results = runtime
|
|
.search("fixture", "annotated notes", 10)
|
|
.await
|
|
.unwrap();
|
|
|
|
let values: Vec<Value> = results
|
|
.iter()
|
|
.map(|item| serde_json::to_value(item).unwrap())
|
|
.collect();
|
|
let resource = values
|
|
.iter()
|
|
.find(|v| v["uri"] == FIXTURE_ANNOTATED_URI)
|
|
.unwrap();
|
|
assert_eq!(resource["audience"], json!(["user"]));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn resource_audience_returns_annotated_roles() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let audience = runtime
|
|
.resource_audience("fixture", FIXTURE_ANNOTATED_URI)
|
|
.await;
|
|
|
|
assert_eq!(audience, Some(vec!["user".to_string()]));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn resource_audience_is_none_for_unknown_uri_or_server() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
assert!(runtime.resource_audience("fixture", "dup").await.is_none());
|
|
assert!(
|
|
runtime
|
|
.resource_audience("fixture", "file:///unknown")
|
|
.await
|
|
.is_none()
|
|
);
|
|
assert!(runtime.resource_audience("ghost", "dup").await.is_none());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn describe_tool_keeps_existing_schema_shape() {
|
|
let fixture = FixtureServer::default();
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let result = runtime.describe("fixture", "tool", "dup").await.unwrap();
|
|
|
|
assert_eq!(
|
|
result,
|
|
json!({
|
|
"type": "object",
|
|
"properties": {
|
|
"tool": { "type": "string" },
|
|
"arguments": {
|
|
"type": "object",
|
|
"properties": { "q": { "type": "string" } }
|
|
}
|
|
}
|
|
})
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn describe_resource_returns_metadata() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let result = runtime
|
|
.describe("fixture", "resource", "dup")
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(result["uri"], "dup");
|
|
assert_eq!(result["name"], "dup-resource");
|
|
assert_eq!(result["description"], "Duplicate-named resource");
|
|
assert_eq!(result["mime_type"], "text/plain");
|
|
assert_eq!(result["size"], 42);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn describe_resource_template_returns_variables() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let result = runtime
|
|
.describe("fixture", "resource_template", "file:///{path}/{name}")
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(result["uri_template"], "file:///{path}/{name}");
|
|
assert_eq!(result["name"], "file-template");
|
|
assert_eq!(result["variables"], json!(["path", "name"]));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn describe_prompt_returns_arguments() {
|
|
let fixture = FixtureServer {
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let result = runtime
|
|
.describe("fixture", "prompt", "summarize")
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(result["name"], "summarize");
|
|
assert_eq!(result["description"], "Summarize a document");
|
|
assert_eq!(result["arguments"][0]["name"], "path");
|
|
assert_eq!(result["arguments"][0]["description"], "Document path");
|
|
assert_eq!(result["arguments"][0]["required"], true);
|
|
assert_eq!(result["arguments"][1]["name"], "style");
|
|
assert_eq!(result["arguments"][1]["required"], Value::Null);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn describe_unknown_kind_lists_valid_kinds() {
|
|
let fixture = FixtureServer::default();
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let err = runtime
|
|
.describe("fixture", "widget", "dup")
|
|
.await
|
|
.unwrap_err()
|
|
.to_string();
|
|
|
|
assert!(err.contains("widget"));
|
|
for kind in ["tool", "resource", "resource_template", "prompt"] {
|
|
assert!(err.contains(kind), "missing {kind} in: {err}");
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn describe_missing_resource_names_kind_in_error() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let err = runtime
|
|
.describe("fixture", "resource", "file:///missing")
|
|
.await
|
|
.unwrap_err()
|
|
.to_string();
|
|
|
|
assert_eq!(
|
|
err,
|
|
"file:///missing not found in fixture MCP server resource catalog"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_returns_result_and_counts_calls() {
|
|
let fixture = FixtureServer {
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
};
|
|
let get_prompt_calls = Arc::clone(&fixture.get_prompt_calls);
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let result = runtime
|
|
.prompt(
|
|
"fixture",
|
|
"summarize",
|
|
HashMap::from([("path".to_string(), "notes.txt".to_string())]),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(result.messages.len(), 3);
|
|
assert_eq!(
|
|
result.messages[0]
|
|
.content
|
|
.as_text()
|
|
.map(|t| t.text.as_str()),
|
|
Some("Summarize notes.txt")
|
|
);
|
|
assert_eq!(get_prompt_calls.load(Ordering::SeqCst), 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_missing_server_errors() {
|
|
let runtime = McpRuntime::new();
|
|
|
|
let err = runtime
|
|
.prompt("ghost", "summarize", HashMap::new())
|
|
.await
|
|
.unwrap_err()
|
|
.to_string();
|
|
|
|
assert_eq!(err, "Prompt MCP server does not exist: ghost");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_surfaces_server_failure() {
|
|
let fixture = FixtureServer {
|
|
prompts_capability: true,
|
|
fail_get_prompt: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let err = runtime
|
|
.prompt("fixture", "summarize", HashMap::new())
|
|
.await
|
|
.unwrap_err()
|
|
.to_string();
|
|
|
|
assert!(
|
|
err.contains("get_prompt exploded"),
|
|
"unexpected error: {err}"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn flatten_prompt_messages_labels_every_message() {
|
|
let messages = vec![
|
|
PromptMessage::new_text(Role::User, "Summarize notes.txt"),
|
|
PromptMessage::new_text(Role::Assistant, "In which style?"),
|
|
PromptMessage::new_text(Role::User, "Concise."),
|
|
];
|
|
|
|
assert_eq!(
|
|
flatten_prompt_messages(&messages),
|
|
"[user]\nSummarize notes.txt\n\n[assistant]\nIn which style?\n\n[user]\nConcise."
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn flatten_prompt_messages_labels_single_message() {
|
|
let messages = vec![PromptMessage::new_text(Role::User, "hello")];
|
|
|
|
assert_eq!(flatten_prompt_messages(&messages), "[user]\nhello");
|
|
}
|
|
|
|
#[test]
|
|
fn flatten_prompt_messages_replaces_non_text_content() {
|
|
let messages = vec![PromptMessage::new(
|
|
Role::User,
|
|
ContentBlock::image("aGk=", "image/png"),
|
|
)];
|
|
|
|
assert_eq!(
|
|
flatten_prompt_messages(&messages),
|
|
"[user]\n[image content omitted]"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn flattened_hostile_prompt_cannot_start_a_command_line() {
|
|
let fixture = FixtureServer {
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let result = runtime
|
|
.prompt("fixture", "hostile", HashMap::new())
|
|
.await
|
|
.unwrap();
|
|
let flattened = flatten_prompt_messages(&result.messages);
|
|
|
|
assert!(flattened.starts_with("[user]"));
|
|
assert!(!flattened.starts_with('!'));
|
|
assert!(!flattened.starts_with('.'));
|
|
assert!(flattened.contains("!rm -rf /"));
|
|
assert!(flattened.contains(".session hijack"));
|
|
}
|
|
|
|
#[test]
|
|
fn format_prompt_arguments_marks_required() {
|
|
let arguments = vec![
|
|
PromptArgument::new("path").with_required(true),
|
|
PromptArgument::new("style"),
|
|
];
|
|
|
|
assert_eq!(
|
|
format_prompt_arguments(&arguments),
|
|
"path (required), style"
|
|
);
|
|
assert_eq!(format_prompt_arguments(&[]), "");
|
|
}
|
|
|
|
#[test]
|
|
fn resolve_prompt_args_reports_missing_required_only() {
|
|
let declared = vec![
|
|
PromptArgument::new("path").with_required(true),
|
|
PromptArgument::new("style"),
|
|
];
|
|
|
|
let (resolved, missing) = resolve_prompt_args(&declared, HashMap::new());
|
|
assert!(resolved.is_empty());
|
|
assert_eq!(missing, vec!["path".to_string()]);
|
|
|
|
let provided = HashMap::from([("path".to_string(), "notes.txt".to_string())]);
|
|
let (resolved, missing) = resolve_prompt_args(&declared, provided);
|
|
assert_eq!(resolved["path"], "notes.txt");
|
|
assert!(missing.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn sanitize_display_text_strips_csi_sequences_entirely() {
|
|
assert_eq!(sanitize_display_text("a\u{1b}[31mred\u{1b}[0mb"), "aredb");
|
|
}
|
|
|
|
#[test]
|
|
fn sanitize_display_text_strips_osc_sequences_entirely() {
|
|
assert_eq!(sanitize_display_text("a\u{1b}]0;title\u{7}b"), "ab");
|
|
assert_eq!(sanitize_display_text("a\u{1b}]0;title\u{1b}\\b"), "ab");
|
|
}
|
|
|
|
#[test]
|
|
fn sanitize_display_text_keeps_plain_text() {
|
|
assert_eq!(
|
|
sanitize_display_text("path (required), café"),
|
|
"path (required), café"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn sanitize_display_text_maps_control_chars_to_spaces() {
|
|
assert_eq!(sanitize_display_text("a\nb\tc\rd\u{7}e"), "a b c d e");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_catalog_carries_arguments() {
|
|
let fixture = FixtureServer {
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let items = runtime.prompt_catalog().await;
|
|
|
|
assert_eq!(items.len(), 1);
|
|
let item = &items[0];
|
|
assert_eq!(item.kind, CatalogItemKind::Prompt);
|
|
assert_eq!(item.server, "fixture");
|
|
assert_eq!(item.name, "summarize");
|
|
assert_eq!(item.description, "Summarize a document");
|
|
let arguments = item.arguments.as_deref().unwrap();
|
|
assert_eq!(format_prompt_arguments(arguments), "path (required), style");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_catalog_sanitizes_hostile_display_strings() {
|
|
let fixture = FixtureServer {
|
|
prompts_capability: true,
|
|
hostile_prompt: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let items = runtime.prompt_catalog().await;
|
|
|
|
let item = items
|
|
.iter()
|
|
.find(|item| item.name == "summarize-evil")
|
|
.unwrap();
|
|
assert_eq!(item.description, "Runs hostile text");
|
|
assert_eq!(
|
|
format_prompt_arguments(item.arguments.as_deref().unwrap()),
|
|
"path (required)"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_catalog_degrades_when_one_server_fails() {
|
|
install_warn_collector();
|
|
let healthy = FixtureServer {
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
};
|
|
let failing = FixtureServer {
|
|
prompts_capability: true,
|
|
fail_prompt_listings: true,
|
|
..Default::default()
|
|
};
|
|
let (mut runtime, _server) = fixture_runtime(healthy).await;
|
|
let _failing_server = add_fixture_server(&mut runtime, "broken", failing).await;
|
|
|
|
let items = runtime.prompt_catalog().await;
|
|
|
|
assert_eq!(items.len(), 1);
|
|
assert_eq!(items[0].server, "fixture");
|
|
let messages = warn_messages().lock().unwrap();
|
|
assert!(
|
|
messages
|
|
.iter()
|
|
.any(|msg| msg.contains("Failed to list prompts on MCP server broken")),
|
|
"missing prompt-listing warning in: {messages:?}"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_catalog_skips_servers_without_prompts_capability() {
|
|
let fixture = FixtureServer::default();
|
|
let prompts_calls = Arc::clone(&fixture.list_prompts_calls);
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let items = runtime.prompt_catalog().await;
|
|
|
|
assert!(items.is_empty());
|
|
assert_eq!(prompts_calls.load(Ordering::SeqCst), 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_completion_stage_one_uses_local_state_only() {
|
|
let fixture = FixtureServer {
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
};
|
|
let list_prompts_calls = Arc::clone(&fixture.list_prompts_calls);
|
|
let get_prompt_calls = Arc::clone(&fixture.get_prompt_calls);
|
|
let (mut runtime, _server) = fixture_runtime(fixture).await;
|
|
let _tools_only =
|
|
add_fixture_server(&mut runtime, "tools-only", FixtureServer::default()).await;
|
|
|
|
let stage =
|
|
runtime.prompt_completion(&["fixture".to_string(), "tools-only".to_string()], &[""]);
|
|
|
|
let McpPromptCompletion::Ready(values) = stage else {
|
|
panic!("stage one must not require an RPC");
|
|
};
|
|
assert_eq!(values, vec![("fixture".to_string(), None)]);
|
|
assert_eq!(list_prompts_calls.load(Ordering::SeqCst), 0);
|
|
assert_eq!(get_prompt_calls.load(Ordering::SeqCst), 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_completion_stage_one_excludes_disabled_servers() {
|
|
let fixture = FixtureServer {
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let McpPromptCompletion::Ready(values) = runtime.prompt_completion(&[], &[""]) else {
|
|
panic!("stage one must not require an RPC");
|
|
};
|
|
|
|
assert!(values.is_empty());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_completion_unknown_server_is_empty() {
|
|
let (runtime, _server) = fixture_runtime(FixtureServer::default()).await;
|
|
|
|
for args in [
|
|
["ghost", ""].as_slice(),
|
|
["ghost", "summarize", ""].as_slice(),
|
|
] {
|
|
let McpPromptCompletion::Ready(values) = runtime.prompt_completion(&[], args) else {
|
|
panic!("unknown server must degrade to empty suggestions");
|
|
};
|
|
assert!(values.is_empty());
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn prompt_completion_later_stages_carry_typed_keys() {
|
|
let fixture = FixtureServer {
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (runtime, _server) = fixture_runtime(fixture).await;
|
|
|
|
let stage = runtime.prompt_completion(&[], &["fixture", "summarize", "path=x", "sty"]);
|
|
|
|
let McpPromptCompletion::ArgumentKeys {
|
|
prompt, typed_keys, ..
|
|
} = stage
|
|
else {
|
|
panic!("expected the argument-key stage");
|
|
};
|
|
assert_eq!(prompt, "summarize");
|
|
assert_eq!(typed_keys, vec!["path".to_string()]);
|
|
}
|
|
|
|
fn single_layer_filter(source: LayerSource, patterns: &[&str]) -> ToolFilter {
|
|
let mut filter = ToolFilter::default();
|
|
filter.push_layer(
|
|
source,
|
|
&patterns.iter().map(|s| s.to_string()).collect::<Vec<_>>(),
|
|
);
|
|
filter
|
|
}
|
|
|
|
fn deny_all_filter() -> ToolFilter {
|
|
single_layer_filter(LayerSource::Global, &[])
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn catalog_items_drops_filtered_tools_but_keeps_other_kinds() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
prompts_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (mut runtime, _server) = fixture_runtime(fixture).await;
|
|
runtime
|
|
.tool_filters
|
|
.insert("fixture".to_string(), deny_all_filter());
|
|
|
|
let items = runtime.catalog_items("fixture").await.unwrap();
|
|
|
|
assert!(!items.contains_key("tool:dup"));
|
|
assert!(items.contains_key("resource:dup"));
|
|
assert!(items.contains_key("resource_template:file:///{path}/{name}"));
|
|
assert!(items.contains_key("prompt:summarize"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn catalog_items_keeps_tools_the_filter_allows() {
|
|
let (mut runtime, _server) = fixture_runtime(FixtureServer::default()).await;
|
|
runtime.tool_filters.insert(
|
|
"fixture".to_string(),
|
|
single_layer_filter(LayerSource::Global, &["d*"]),
|
|
);
|
|
|
|
let items = runtime.catalog_items("fixture").await.unwrap();
|
|
|
|
assert!(items.contains_key("tool:dup"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn search_never_surfaces_filtered_tools() {
|
|
let fixture = FixtureServer {
|
|
resources_capability: true,
|
|
..Default::default()
|
|
};
|
|
let (mut runtime, _server) = fixture_runtime(fixture).await;
|
|
runtime
|
|
.tool_filters
|
|
.insert("fixture".to_string(), deny_all_filter());
|
|
|
|
let results = runtime.search("fixture", "dup", 10).await.unwrap();
|
|
|
|
assert!(
|
|
results
|
|
.iter()
|
|
.all(|item| item.kind != CatalogItemKind::Tool)
|
|
);
|
|
assert!(
|
|
results
|
|
.iter()
|
|
.any(|item| item.kind == CatalogItemKind::Resource)
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn describe_blocked_tool_is_indistinguishable_from_missing() {
|
|
let (unfiltered, _server) = fixture_runtime(FixtureServer::default()).await;
|
|
let missing = unfiltered
|
|
.describe("fixture", "tool", "ghost")
|
|
.await
|
|
.unwrap_err()
|
|
.to_string();
|
|
|
|
let (mut filtered, _other_server) = fixture_runtime(FixtureServer::default()).await;
|
|
filtered
|
|
.tool_filters
|
|
.insert("fixture".to_string(), deny_all_filter());
|
|
let blocked = filtered
|
|
.describe("fixture", "tool", "dup")
|
|
.await
|
|
.unwrap_err()
|
|
.to_string();
|
|
|
|
assert_eq!(blocked, "dup not found in fixture MCP server catalog");
|
|
assert_eq!(blocked, missing.replace("ghost", "dup"));
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn invoke_blocked_tool_errors_like_describe_and_never_reaches_server() {
|
|
let fixture = FixtureServer::default();
|
|
let call_tool_calls = Arc::clone(&fixture.call_tool_calls);
|
|
let (mut runtime, _server) = fixture_runtime(fixture).await;
|
|
runtime
|
|
.tool_filters
|
|
.insert("fixture".to_string(), deny_all_filter());
|
|
|
|
let invoke_err = runtime
|
|
.invoke("fixture", "dup", json!({}))
|
|
.await
|
|
.unwrap_err()
|
|
.to_string();
|
|
let describe_err = runtime
|
|
.describe("fixture", "tool", "dup")
|
|
.await
|
|
.unwrap_err()
|
|
.to_string();
|
|
|
|
assert_eq!(invoke_err, "dup not found in fixture MCP server catalog");
|
|
assert_eq!(invoke_err, describe_err);
|
|
assert_eq!(call_tool_calls.load(Ordering::SeqCst), 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn invoke_allowed_tool_still_reaches_the_server() {
|
|
let fixture = FixtureServer::default();
|
|
let call_tool_calls = Arc::clone(&fixture.call_tool_calls);
|
|
let (mut runtime, _server) = fixture_runtime(fixture).await;
|
|
runtime.tool_filters.insert(
|
|
"fixture".to_string(),
|
|
single_layer_filter(LayerSource::Global, &["d*"]),
|
|
);
|
|
|
|
let _ = runtime.invoke("fixture", "dup", json!({})).await;
|
|
|
|
assert_eq!(call_tool_calls.load(Ordering::SeqCst), 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn catalog_items_warns_on_dead_context_patterns() {
|
|
install_warn_collector();
|
|
let (mut runtime, _server) = fixture_runtime(FixtureServer::default()).await;
|
|
let mut filter = single_layer_filter(LayerSource::Global, &["*"]);
|
|
filter.push_layer(LayerSource::Session, &["zzz_*".to_string()]);
|
|
runtime.tool_filters.insert("fixture".to_string(), filter);
|
|
|
|
runtime.catalog_items("fixture").await.unwrap();
|
|
|
|
let messages = warn_messages().lock().unwrap();
|
|
assert!(
|
|
messages.iter().any(|msg| msg.contains("'zzz_*'")
|
|
&& msg.contains("session (.set)")
|
|
&& msg.contains("'fixture'")),
|
|
"missing dead-pattern warning in: {messages:?}"
|
|
);
|
|
}
|
|
}
|