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)>), PromptNames { server: Arc, }, ArgumentKeys { server: Arc, prompt: String, typed_keys: Vec, }, } #[derive(Default, Clone)] pub struct McpRuntime { pub servers: HashMap>, /// Per-server effective tool allowlists; a server absent here is unfiltered. pub tool_filters: HashMap, } 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) { self.servers.insert(name, handle); } pub fn get(&self, name: &str) -> Option<&Arc> { self.servers.get(name) } pub fn server_features(&self) -> Vec { let mut features: Vec = 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> { 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 = 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> { 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::::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> { 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 { 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 = 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 { 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 { 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> { 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, ) -> Result { 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 { 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::>() .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::>() .join(", ") } pub fn resolve_prompt_args( declared: &[PromptArgument], provided: HashMap, ) -> (HashMap, Vec) { 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) -> Option> { 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, new_items: impl IntoIterator, ) { 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 { 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, pub(crate) tool_result: Option, pub(crate) list_resources_calls: Arc, pub(crate) list_prompts_calls: Arc, pub(crate) get_prompt_calls: Arc, pub(crate) call_tool_calls: Arc, } 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, _context: RequestContext, ) -> Result { 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, ) -> Result { 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, _context: RequestContext, ) -> Result { 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, _context: RequestContext, ) -> Result { 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, ) -> Result { 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, _context: RequestContext, ) -> Result { 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, ) -> Result { 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 { 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) { 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>> = OnceLock::new(); fn warn_messages() -> &'static Mutex> { 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!["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 = 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 = 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::>(), ); 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:?}" ); } }