use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::sync::Arc;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use turbomcp_core::context::RequestContext;
use turbomcp_core::error::{McpError, McpResult};
use turbomcp_core::handler::McpHandler;
use turbomcp_types::{
ComponentFilter, ComponentMeta, Prompt, PromptResult, Resource, ResourceResult,
ResourceTemplate, Tool, ToolResult,
};
type SessionVisibilityMap = Arc<dashmap::DashMap<String, HashSet<String>>>;
type SharedComponentRegistry = Arc<RwLock<ComponentRegistryCache>>;
#[derive(Debug, Clone, Default)]
struct ComponentRegistryCache {
tools: Option<BTreeMap<String, Tool>>,
resources_by_uri: Option<BTreeMap<String, Resource>>,
resource_templates_by_uri_template: Option<BTreeMap<String, ResourceTemplate>>,
prompts: Option<BTreeMap<String, Prompt>>,
}
enum RegistryLookup<T> {
Uninitialized,
Found(T),
Missing,
}
impl ComponentRegistryCache {
fn replace_tools(&mut self, tools: Vec<Tool>) {
let mut registry = BTreeMap::new();
for tool in tools {
registry.entry(tool.name.clone()).or_insert(tool);
}
self.tools = Some(registry);
}
fn replace_resources(&mut self, resources: Vec<Resource>) {
let mut registry = BTreeMap::new();
for resource in resources {
registry.entry(resource.uri.clone()).or_insert(resource);
}
self.resources_by_uri = Some(registry);
}
fn replace_resource_templates(&mut self, templates: Vec<ResourceTemplate>) {
let mut registry = BTreeMap::new();
for template in templates {
registry
.entry(template.uri_template.clone())
.or_insert(template);
}
self.resource_templates_by_uri_template = Some(registry);
}
fn replace_prompts(&mut self, prompts: Vec<Prompt>) {
let mut registry = BTreeMap::new();
for prompt in prompts {
registry.entry(prompt.name.clone()).or_insert(prompt);
}
self.prompts = Some(registry);
}
fn tool(&self, name: &str) -> RegistryLookup<Tool> {
match &self.tools {
Some(tools) => tools
.get(name)
.cloned()
.map_or(RegistryLookup::Missing, RegistryLookup::Found),
None => RegistryLookup::Uninitialized,
}
}
fn resource_by_uri(&self, uri: &str) -> RegistryLookup<Resource> {
match &self.resources_by_uri {
Some(resources) => resources
.get(uri)
.cloned()
.map_or(RegistryLookup::Missing, RegistryLookup::Found),
None => RegistryLookup::Uninitialized,
}
}
fn prompt(&self, name: &str) -> RegistryLookup<Prompt> {
match &self.prompts {
Some(prompts) => prompts
.get(name)
.cloned()
.map_or(RegistryLookup::Missing, RegistryLookup::Found),
None => RegistryLookup::Uninitialized,
}
}
fn clear(&mut self) {
*self = Self::default();
}
fn is_initialized(&self) -> bool {
self.tools.is_some()
|| self.resources_by_uri.is_some()
|| self.resource_templates_by_uri_template.is_some()
|| self.prompts.is_some()
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default)]
pub struct ComponentVisibilityRules {
#[serde(skip_serializing_if = "Option::is_none")]
pub allow: Option<BTreeSet<String>>,
#[serde(skip_serializing_if = "BTreeSet::is_empty")]
pub deny: BTreeSet<String>,
#[serde(skip_serializing_if = "BTreeSet::is_empty")]
pub hide: BTreeSet<String>,
}
impl ComponentVisibilityRules {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn allow<I, S>(names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
Self {
allow: Some(collect_names(names)),
deny: BTreeSet::new(),
hide: BTreeSet::new(),
}
}
#[must_use]
pub fn deny<I, S>(names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
Self {
allow: None,
deny: collect_names(names),
hide: BTreeSet::new(),
}
}
#[must_use]
pub fn hide<I, S>(names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
Self {
allow: None,
deny: BTreeSet::new(),
hide: collect_names(names),
}
}
#[must_use]
pub fn with_allowed<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.allow = Some(collect_names(names));
self
}
#[must_use]
pub fn with_disabled<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.deny = collect_names(names);
self
}
#[must_use]
pub fn with_hidden<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.hide = collect_names(names);
self
}
#[must_use]
pub fn is_enabled(&self, identifier: &str) -> bool {
self.is_enabled_any([identifier])
}
#[must_use]
pub fn is_enabled_any<'a, I>(&self, identifiers: I) -> bool
where
I: IntoIterator<Item = &'a str>,
{
let identifiers = identifiers.into_iter().collect::<Vec<_>>();
if identifiers
.iter()
.any(|identifier| self.deny.contains(*identifier))
{
return false;
}
self.allow.as_ref().is_none_or(|allow| {
identifiers
.iter()
.any(|identifier| allow.contains(*identifier))
})
}
#[must_use]
pub fn is_listed(&self, identifier: &str) -> bool {
self.is_listed_any([identifier])
}
#[must_use]
pub fn is_listed_any<'a, I>(&self, identifiers: I) -> bool
where
I: IntoIterator<Item = &'a str>,
{
let identifiers = identifiers.into_iter().collect::<Vec<_>>();
self.is_enabled_any(identifiers.iter().copied())
&& !identifiers
.iter()
.any(|identifier| self.hide.contains(*identifier))
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default)]
pub struct VisibilityConfig {
pub tools: ComponentVisibilityRules,
pub resources: ComponentVisibilityRules,
pub resource_templates: ComponentVisibilityRules,
pub prompts: ComponentVisibilityRules,
#[serde(skip_serializing_if = "is_false")]
pub require_read_only_tools: bool,
}
impl VisibilityConfig {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_allowed_tools<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.tools = self.tools.with_allowed(names);
self
}
#[must_use]
pub fn with_disabled_tools<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.tools = self.tools.with_disabled(names);
self
}
#[must_use]
pub fn with_hidden_tools<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.tools = self.tools.with_hidden(names);
self
}
#[must_use]
pub fn with_allowed_resources<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resources = self.resources.with_allowed(identifiers);
self
}
#[must_use]
pub fn with_disabled_resources<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resources = self.resources.with_disabled(identifiers);
self
}
#[must_use]
pub fn with_hidden_resources<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resources = self.resources.with_hidden(identifiers);
self
}
#[must_use]
pub fn with_allowed_resource_templates<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resource_templates = self.resource_templates.with_allowed(identifiers);
self
}
#[must_use]
pub fn with_disabled_resource_templates<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resource_templates = self.resource_templates.with_disabled(identifiers);
self
}
#[must_use]
pub fn with_hidden_resource_templates<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resource_templates = self.resource_templates.with_hidden(identifiers);
self
}
#[must_use]
pub fn with_allowed_prompts<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.prompts = self.prompts.with_allowed(names);
self
}
#[must_use]
pub fn with_disabled_prompts<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.prompts = self.prompts.with_disabled(names);
self
}
#[must_use]
pub fn with_hidden_prompts<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.prompts = self.prompts.with_hidden(names);
self
}
#[must_use]
pub fn require_read_only_tools(mut self) -> Self {
self.require_read_only_tools = true;
self
}
}
fn collect_names<I, S>(names: I) -> BTreeSet<String>
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
names.into_iter().map(Into::into).collect()
}
fn is_false(value: &bool) -> bool {
!*value
}
fn is_explicit_read_only_tool(tool: &Tool) -> bool {
tool.annotations.as_ref().is_some_and(|annotations| {
annotations.read_only_hint == Some(true) && annotations.destructive_hint != Some(true)
})
}
#[derive(Debug)]
pub struct VisibilitySessionGuard {
session_id: String,
session_enabled: SessionVisibilityMap,
session_disabled: SessionVisibilityMap,
}
impl VisibilitySessionGuard {
pub fn session_id(&self) -> &str {
&self.session_id
}
}
impl Drop for VisibilitySessionGuard {
fn drop(&mut self) {
self.session_enabled.remove(&self.session_id);
self.session_disabled.remove(&self.session_id);
}
}
pub struct VisibilityLayer<H> {
inner: H,
global_disabled: Arc<RwLock<Vec<ComponentFilter>>>,
tool_rules: ComponentVisibilityRules,
resource_rules: ComponentVisibilityRules,
resource_template_rules: ComponentVisibilityRules,
prompt_rules: ComponentVisibilityRules,
read_only_tools_required: bool,
component_registry: SharedComponentRegistry,
session_enabled: SessionVisibilityMap,
session_disabled: SessionVisibilityMap,
}
impl<H: Clone> Clone for VisibilityLayer<H> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
global_disabled: Arc::new(RwLock::new(self.global_disabled.read().clone())),
tool_rules: self.tool_rules.clone(),
resource_rules: self.resource_rules.clone(),
resource_template_rules: self.resource_template_rules.clone(),
prompt_rules: self.prompt_rules.clone(),
read_only_tools_required: self.read_only_tools_required,
component_registry: Arc::clone(&self.component_registry),
session_enabled: Arc::clone(&self.session_enabled),
session_disabled: Arc::clone(&self.session_disabled),
}
}
}
impl<H: std::fmt::Debug> std::fmt::Debug for VisibilityLayer<H> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("VisibilityLayer")
.field("inner", &self.inner)
.field("global_disabled_count", &self.global_disabled.read().len())
.field(
"tool_allow_count",
&self.tool_rules.allow.as_ref().map(BTreeSet::len),
)
.field("tool_deny_count", &self.tool_rules.deny.len())
.field("tool_hide_count", &self.tool_rules.hide.len())
.field("read_only_tools_required", &self.read_only_tools_required)
.field(
"component_registry_initialized",
&self.component_registry.read().is_initialized(),
)
.field("session_enabled_count", &self.session_enabled.len())
.field("session_disabled_count", &self.session_disabled.len())
.finish()
}
}
impl<H: McpHandler> VisibilityLayer<H> {
pub fn new(inner: H) -> Self {
Self {
inner,
global_disabled: Arc::new(RwLock::new(Vec::new())),
tool_rules: ComponentVisibilityRules::new(),
resource_rules: ComponentVisibilityRules::new(),
resource_template_rules: ComponentVisibilityRules::new(),
prompt_rules: ComponentVisibilityRules::new(),
read_only_tools_required: false,
component_registry: Arc::new(RwLock::new(ComponentRegistryCache::default())),
session_enabled: Arc::new(dashmap::DashMap::new()),
session_disabled: Arc::new(dashmap::DashMap::new()),
}
}
#[must_use]
pub fn with_disabled(self, filter: ComponentFilter) -> Self {
self.global_disabled.write().push(filter);
self
}
#[must_use]
pub fn disable_tags<I, S>(self, tags: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.with_disabled(ComponentFilter::with_tags(tags))
}
#[must_use]
pub fn with_visibility_config(mut self, config: VisibilityConfig) -> Self {
self.tool_rules = config.tools;
self.resource_rules = config.resources;
self.resource_template_rules = config.resource_templates;
self.prompt_rules = config.prompts;
self.read_only_tools_required = config.require_read_only_tools;
self
}
#[must_use]
pub fn visibility_config(&self) -> VisibilityConfig {
VisibilityConfig {
tools: self.tool_rules.clone(),
resources: self.resource_rules.clone(),
resource_templates: self.resource_template_rules.clone(),
prompts: self.prompt_rules.clone(),
require_read_only_tools: self.read_only_tools_required,
}
}
pub fn refresh_component_registry(&self) {
let tools = self.inner.list_tools();
let resources = self.inner.list_resources();
let resource_templates = self.inner.list_resource_templates();
let prompts = self.inner.list_prompts();
let mut registry = self.component_registry.write();
registry.replace_tools(tools);
registry.replace_resources(resources);
registry.replace_resource_templates(resource_templates);
registry.replace_prompts(prompts);
}
pub fn clear_component_registry(&self) {
self.component_registry.write().clear();
}
#[must_use]
pub fn with_allowed_tools<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.tool_rules = self.tool_rules.with_allowed(names);
self
}
#[must_use]
pub fn with_disabled_tools<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.tool_rules = self.tool_rules.with_disabled(names);
self
}
#[must_use]
pub fn with_hidden_tools<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.tool_rules = self.tool_rules.with_hidden(names);
self
}
#[must_use]
pub fn with_allowed_resources<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resource_rules = self.resource_rules.with_allowed(identifiers);
self
}
#[must_use]
pub fn with_disabled_resources<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resource_rules = self.resource_rules.with_disabled(identifiers);
self
}
#[must_use]
pub fn with_hidden_resources<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resource_rules = self.resource_rules.with_hidden(identifiers);
self
}
#[must_use]
pub fn with_allowed_resource_templates<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resource_template_rules = self.resource_template_rules.with_allowed(identifiers);
self
}
#[must_use]
pub fn with_disabled_resource_templates<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resource_template_rules = self.resource_template_rules.with_disabled(identifiers);
self
}
#[must_use]
pub fn with_hidden_resource_templates<I, S>(mut self, identifiers: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.resource_template_rules = self.resource_template_rules.with_hidden(identifiers);
self
}
#[must_use]
pub fn with_allowed_prompts<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.prompt_rules = self.prompt_rules.with_allowed(names);
self
}
#[must_use]
pub fn with_disabled_prompts<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.prompt_rules = self.prompt_rules.with_disabled(names);
self
}
#[must_use]
pub fn with_hidden_prompts<I, S>(mut self, names: I) -> Self
where
I: IntoIterator<Item = S>,
S: Into<String>,
{
self.prompt_rules = self.prompt_rules.with_hidden(names);
self
}
#[must_use]
pub fn require_read_only_tools(mut self) -> Self {
self.read_only_tools_required = true;
self
}
fn register_tools(&self, tools: Vec<Tool>) {
self.component_registry.write().replace_tools(tools);
}
fn register_resources(&self, resources: Vec<Resource>) {
self.component_registry.write().replace_resources(resources);
}
fn register_resource_templates(&self, templates: Vec<ResourceTemplate>) {
self.component_registry
.write()
.replace_resource_templates(templates);
}
fn register_prompts(&self, prompts: Vec<Prompt>) {
self.component_registry.write().replace_prompts(prompts);
}
fn registered_tool(&self, name: &str) -> Option<Tool> {
let lookup = { self.component_registry.read().tool(name) };
match lookup {
RegistryLookup::Found(tool) => return Some(tool),
RegistryLookup::Missing => return None,
RegistryLookup::Uninitialized => {}
}
let tools = self.inner.list_tools();
let tool = tools.iter().find(|tool| tool.name == name).cloned();
self.register_tools(tools);
tool
}
fn registered_resource(&self, uri: &str) -> Option<Resource> {
let lookup = { self.component_registry.read().resource_by_uri(uri) };
match lookup {
RegistryLookup::Found(resource) => return Some(resource),
RegistryLookup::Missing => return None,
RegistryLookup::Uninitialized => {}
}
let resources = self.inner.list_resources();
let resource = resources
.iter()
.find(|resource| resource.uri == uri)
.cloned();
self.register_resources(resources);
resource
}
fn registered_prompt(&self, name: &str) -> Option<Prompt> {
let lookup = { self.component_registry.read().prompt(name) };
match lookup {
RegistryLookup::Found(prompt) => return Some(prompt),
RegistryLookup::Missing => return None,
RegistryLookup::Uninitialized => {}
}
let prompts = self.inner.list_prompts();
let prompt = prompts.iter().find(|prompt| prompt.name == name).cloned();
self.register_prompts(prompts);
prompt
}
fn is_visible(&self, meta: &ComponentMeta, session_id: Option<&str>) -> bool {
let global_disabled = self.global_disabled.read();
let globally_hidden = global_disabled.iter().any(|filter| filter.matches(meta));
if !globally_hidden {
if let Some(sid) = session_id
&& let Some(disabled) = self.session_disabled.get(sid)
&& meta.tags.iter().any(|t| disabled.contains(t))
{
return false;
}
return true;
}
if let Some(sid) = session_id
&& let Some(enabled) = self.session_enabled.get(sid)
&& meta.tags.iter().any(|t| enabled.contains(t))
{
return true;
}
false
}
fn is_tool_enabled(&self, tool: &Tool, session_id: Option<&str>) -> bool {
if !self.tool_rules.is_enabled(&tool.name) {
return false;
}
if self.read_only_tools_required && !is_explicit_read_only_tool(tool) {
return false;
}
let meta = ComponentMeta::from_meta_value(tool.meta.as_ref());
self.is_visible(&meta, session_id)
}
fn is_tool_listed(&self, tool: &Tool, session_id: Option<&str>) -> bool {
self.is_tool_enabled(tool, session_id) && self.tool_rules.is_listed(&tool.name)
}
fn is_unregistered_tool_callable(&self, name: &str) -> bool {
self.tool_rules.is_enabled(name) && !self.read_only_tools_required
}
fn is_resource_enabled(&self, resource: &Resource, session_id: Option<&str>) -> bool {
if !self
.resource_rules
.is_enabled_any([resource.name.as_str(), resource.uri.as_str()])
{
return false;
}
let meta = ComponentMeta::from_meta_value(resource.meta.as_ref());
self.is_visible(&meta, session_id)
}
fn is_resource_listed(&self, resource: &Resource, session_id: Option<&str>) -> bool {
self.is_resource_enabled(resource, session_id)
&& self
.resource_rules
.is_listed_any([resource.name.as_str(), resource.uri.as_str()])
}
fn is_unregistered_resource_readable(&self, uri: &str) -> bool {
self.resource_rules.is_enabled(uri)
}
fn is_resource_template_listed(
&self,
template: &ResourceTemplate,
session_id: Option<&str>,
) -> bool {
if !self
.resource_template_rules
.is_listed_any([template.name.as_str(), template.uri_template.as_str()])
{
return false;
}
let meta = ComponentMeta::from_meta_value(template.meta.as_ref());
self.is_visible(&meta, session_id)
}
fn is_prompt_enabled(&self, prompt: &Prompt, session_id: Option<&str>) -> bool {
if !self.prompt_rules.is_enabled(&prompt.name) {
return false;
}
let meta = ComponentMeta::from_meta_value(prompt.meta.as_ref());
self.is_visible(&meta, session_id)
}
fn is_prompt_listed(&self, prompt: &Prompt, session_id: Option<&str>) -> bool {
self.is_prompt_enabled(prompt, session_id) && self.prompt_rules.is_listed(&prompt.name)
}
fn is_unregistered_prompt_gettable(&self, name: &str) -> bool {
self.prompt_rules.is_enabled(name)
}
pub fn enable_for_session(&self, session_id: &str, tags: &[String]) {
let mut entry = self
.session_enabled
.entry(session_id.to_string())
.or_default();
entry.extend(tags.iter().cloned());
if let Some(mut disabled) = self.session_disabled.get_mut(session_id) {
for tag in tags {
disabled.remove(tag);
}
}
}
pub fn disable_for_session(&self, session_id: &str, tags: &[String]) {
let mut entry = self
.session_disabled
.entry(session_id.to_string())
.or_default();
entry.extend(tags.iter().cloned());
if let Some(mut enabled) = self.session_enabled.get_mut(session_id) {
for tag in tags {
enabled.remove(tag);
}
}
}
pub fn clear_session(&self, session_id: &str) {
self.session_enabled.remove(session_id);
self.session_disabled.remove(session_id);
}
pub fn session_guard(&self, session_id: impl Into<String>) -> VisibilitySessionGuard {
VisibilitySessionGuard {
session_id: session_id.into(),
session_enabled: Arc::clone(&self.session_enabled),
session_disabled: Arc::clone(&self.session_disabled),
}
}
pub fn active_sessions_count(&self) -> usize {
let mut sessions = HashSet::new();
for entry in self.session_enabled.iter() {
sessions.insert(entry.key().clone());
}
for entry in self.session_disabled.iter() {
sessions.insert(entry.key().clone());
}
sessions.len()
}
pub fn inner(&self) -> &H {
&self.inner
}
pub fn inner_mut(&mut self) -> &mut H {
self.clear_component_registry();
&mut self.inner
}
pub fn into_inner(self) -> H {
self.inner
}
}
#[allow(clippy::manual_async_fn)]
impl<H: McpHandler> McpHandler for VisibilityLayer<H> {
fn server_info(&self) -> turbomcp_types::ServerInfo {
self.inner.server_info()
}
fn server_capabilities(&self) -> turbomcp_types::ServerCapabilities {
self.inner.server_capabilities()
}
fn list_tools(&self) -> Vec<Tool> {
let tools = self.inner.list_tools();
self.register_tools(tools.clone());
tools
.into_iter()
.filter(|tool| self.is_tool_listed(tool, None))
.collect()
}
fn list_resources(&self) -> Vec<Resource> {
let resources = self.inner.list_resources();
self.register_resources(resources.clone());
resources
.into_iter()
.filter(|resource| self.is_resource_listed(resource, None))
.collect()
}
fn list_resource_templates(&self) -> Vec<ResourceTemplate> {
let templates = self.inner.list_resource_templates();
self.register_resource_templates(templates.clone());
templates
.into_iter()
.filter(|template| self.is_resource_template_listed(template, None))
.collect()
}
fn list_prompts(&self) -> Vec<Prompt> {
let prompts = self.inner.list_prompts();
self.register_prompts(prompts.clone());
prompts
.into_iter()
.filter(|prompt| self.is_prompt_listed(prompt, None))
.collect()
}
fn call_tool<'a>(
&'a self,
name: &'a str,
args: serde_json::Value,
ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ToolResult>> + turbomcp_core::marker::MaybeSend + 'a
{
async move {
if let Some(tool) = self.registered_tool(name) {
if !self.is_tool_enabled(&tool, ctx.session_id()) {
return Err(McpError::tool_not_found(name));
}
} else if !self.is_unregistered_tool_callable(name) {
return Err(McpError::tool_not_found(name));
}
self.inner.call_tool(name, args, ctx).await
}
}
fn read_resource<'a>(
&'a self,
uri: &'a str,
ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ResourceResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move {
if let Some(resource) = self.registered_resource(uri) {
if !self.is_resource_enabled(&resource, ctx.session_id()) {
return Err(McpError::resource_not_found(uri));
}
} else if !self.is_unregistered_resource_readable(uri) {
return Err(McpError::resource_not_found(uri));
}
self.inner.read_resource(uri, ctx).await
}
}
fn get_prompt<'a>(
&'a self,
name: &'a str,
args: Option<serde_json::Value>,
ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<PromptResult>> + turbomcp_core::marker::MaybeSend + 'a
{
async move {
if let Some(prompt) = self.registered_prompt(name) {
if !self.is_prompt_enabled(&prompt, ctx.session_id()) {
return Err(McpError::prompt_not_found(name));
}
} else if !self.is_unregistered_prompt_gettable(name) {
return Err(McpError::prompt_not_found(name));
}
self.inner.get_prompt(name, args, ctx).await
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use turbomcp_types::ToolAnnotations;
#[derive(Clone, Debug)]
struct MockHandler;
#[allow(clippy::manual_async_fn)]
impl McpHandler for MockHandler {
fn server_info(&self) -> turbomcp_types::ServerInfo {
turbomcp_types::ServerInfo::new("test", "1.0.0")
}
fn list_tools(&self) -> Vec<Tool> {
vec![
Tool {
name: "public_tool".to_string(),
description: Some("Public tool".to_string()),
annotations: Some(ToolAnnotations::default().with_read_only(true)),
meta: Some({
let mut m = std::collections::HashMap::new();
m.insert("tags".to_string(), serde_json::json!(["public"]));
m
}),
..Default::default()
},
Tool {
name: "admin_tool".to_string(),
description: Some("Admin tool".to_string()),
annotations: Some(ToolAnnotations::default().with_destructive(true)),
meta: Some({
let mut m = std::collections::HashMap::new();
m.insert("tags".to_string(), serde_json::json!(["admin"]));
m
}),
..Default::default()
},
]
}
fn list_resources(&self) -> Vec<Resource> {
vec![
Resource {
uri: "vault://public".to_string(),
name: "public_resource".to_string(),
meta: Some({
let mut m = std::collections::HashMap::new();
m.insert("tags".to_string(), serde_json::json!(["public"]));
m
}),
..Default::default()
},
Resource {
uri: "vault://admin".to_string(),
name: "admin_resource".to_string(),
meta: Some({
let mut m = std::collections::HashMap::new();
m.insert("tags".to_string(), serde_json::json!(["admin"]));
m
}),
..Default::default()
},
]
}
fn list_resource_templates(&self) -> Vec<ResourceTemplate> {
vec![ResourceTemplate {
uri_template: "vault://notes/{id}".to_string(),
name: "note_template".to_string(),
meta: Some({
let mut m = std::collections::HashMap::new();
m.insert("tags".to_string(), serde_json::json!(["public"]));
m
}),
..Default::default()
}]
}
fn list_prompts(&self) -> Vec<Prompt> {
vec![
Prompt {
name: "public_prompt".to_string(),
meta: Some({
let mut m = std::collections::HashMap::new();
m.insert("tags".to_string(), serde_json::json!(["public"]));
m
}),
..Default::default()
},
Prompt {
name: "admin_prompt".to_string(),
meta: Some({
let mut m = std::collections::HashMap::new();
m.insert("tags".to_string(), serde_json::json!(["admin"]));
m
}),
..Default::default()
},
]
}
fn call_tool<'a>(
&'a self,
name: &'a str,
_args: serde_json::Value,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ToolResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(ToolResult::text(format!("Called {}", name))) }
}
fn read_resource<'a>(
&'a self,
uri: &'a str,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ResourceResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(ResourceResult::text(uri, format!("Read {}", uri))) }
}
fn get_prompt<'a>(
&'a self,
name: &'a str,
_args: Option<serde_json::Value>,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<PromptResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(PromptResult::user(format!("Prompt {}", name))) }
}
}
#[derive(Clone, Debug)]
struct DynamicHandler;
#[allow(clippy::manual_async_fn)]
impl McpHandler for DynamicHandler {
fn server_info(&self) -> turbomcp_types::ServerInfo {
turbomcp_types::ServerInfo::new("dynamic", "1.0.0")
}
fn list_tools(&self) -> Vec<Tool> {
vec![]
}
fn list_resources(&self) -> Vec<Resource> {
vec![]
}
fn list_prompts(&self) -> Vec<Prompt> {
vec![]
}
fn call_tool<'a>(
&'a self,
name: &'a str,
_args: serde_json::Value,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ToolResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(ToolResult::text(format!("Dynamic {}", name))) }
}
fn read_resource<'a>(
&'a self,
uri: &'a str,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ResourceResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(ResourceResult::text(uri, format!("Dynamic {}", uri))) }
}
fn get_prompt<'a>(
&'a self,
name: &'a str,
_args: Option<serde_json::Value>,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<PromptResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(PromptResult::user(format!("Dynamic {}", name))) }
}
}
#[derive(Debug, Default)]
struct CountingState {
tool_lists: AtomicUsize,
resource_lists: AtomicUsize,
prompt_lists: AtomicUsize,
}
#[derive(Clone, Debug, Default)]
struct CountingHandler {
state: Arc<CountingState>,
}
#[allow(clippy::manual_async_fn)]
impl McpHandler for CountingHandler {
fn server_info(&self) -> turbomcp_types::ServerInfo {
turbomcp_types::ServerInfo::new("counting", "1.0.0")
}
fn list_tools(&self) -> Vec<Tool> {
self.state.tool_lists.fetch_add(1, Ordering::SeqCst);
vec![Tool {
name: "counted_tool".to_string(),
annotations: Some(ToolAnnotations::default().with_read_only(true)),
..Default::default()
}]
}
fn list_resources(&self) -> Vec<Resource> {
self.state.resource_lists.fetch_add(1, Ordering::SeqCst);
vec![Resource::new("counted://resource", "counted_resource")]
}
fn list_prompts(&self) -> Vec<Prompt> {
self.state.prompt_lists.fetch_add(1, Ordering::SeqCst);
vec![Prompt::new("counted_prompt", "Counted prompt")]
}
fn call_tool<'a>(
&'a self,
name: &'a str,
_args: serde_json::Value,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ToolResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(ToolResult::text(format!("Called {}", name))) }
}
fn read_resource<'a>(
&'a self,
uri: &'a str,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ResourceResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(ResourceResult::text(uri, format!("Read {}", uri))) }
}
fn get_prompt<'a>(
&'a self,
name: &'a str,
_args: Option<serde_json::Value>,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<PromptResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(PromptResult::user(format!("Prompt {}", name))) }
}
}
#[derive(Clone, Debug)]
struct MutableToolHandler {
read_only: Arc<AtomicBool>,
}
#[allow(clippy::manual_async_fn)]
impl McpHandler for MutableToolHandler {
fn server_info(&self) -> turbomcp_types::ServerInfo {
turbomcp_types::ServerInfo::new("mutable", "1.0.0")
}
fn list_tools(&self) -> Vec<Tool> {
let annotation = if self.read_only.load(Ordering::SeqCst) {
ToolAnnotations::default().with_read_only(true)
} else {
ToolAnnotations::default().with_destructive(true)
};
vec![Tool {
name: "mutable_tool".to_string(),
annotations: Some(annotation),
..Default::default()
}]
}
fn list_resources(&self) -> Vec<Resource> {
Vec::new()
}
fn list_prompts(&self) -> Vec<Prompt> {
Vec::new()
}
fn call_tool<'a>(
&'a self,
name: &'a str,
_args: serde_json::Value,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ToolResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(ToolResult::text(format!("Called {}", name))) }
}
fn read_resource<'a>(
&'a self,
uri: &'a str,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ResourceResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Err(McpError::resource_not_found(uri)) }
}
fn get_prompt<'a>(
&'a self,
name: &'a str,
_args: Option<serde_json::Value>,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<PromptResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Err(McpError::prompt_not_found(name)) }
}
}
#[derive(Clone, Debug)]
struct DuplicateToolMetadataHandler;
#[allow(clippy::manual_async_fn)]
impl McpHandler for DuplicateToolMetadataHandler {
fn server_info(&self) -> turbomcp_types::ServerInfo {
turbomcp_types::ServerInfo::new("duplicates", "1.0.0")
}
fn list_tools(&self) -> Vec<Tool> {
vec![
Tool {
name: "duplicate_tool".to_string(),
annotations: Some(ToolAnnotations::default().with_read_only(true)),
..Default::default()
},
Tool {
name: "duplicate_tool".to_string(),
annotations: Some(ToolAnnotations::default().with_destructive(true)),
..Default::default()
},
]
}
fn list_resources(&self) -> Vec<Resource> {
Vec::new()
}
fn list_prompts(&self) -> Vec<Prompt> {
Vec::new()
}
fn call_tool<'a>(
&'a self,
name: &'a str,
_args: serde_json::Value,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ToolResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Ok(ToolResult::text(format!("Called {}", name))) }
}
fn read_resource<'a>(
&'a self,
uri: &'a str,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<ResourceResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Err(McpError::resource_not_found(uri)) }
}
fn get_prompt<'a>(
&'a self,
name: &'a str,
_args: Option<serde_json::Value>,
_ctx: &'a RequestContext,
) -> impl std::future::Future<Output = McpResult<PromptResult>>
+ turbomcp_core::marker::MaybeSend
+ 'a {
async move { Err(McpError::prompt_not_found(name)) }
}
}
fn tool_names<H: McpHandler>(layer: &VisibilityLayer<H>) -> Vec<String> {
layer
.list_tools()
.into_iter()
.map(|tool| tool.name)
.collect()
}
#[test]
fn test_component_visibility_rules_deny_wins() {
let rules = ComponentVisibilityRules::allow(["search", "delete"]).with_disabled(["delete"]);
assert!(rules.is_enabled("search"));
assert!(rules.is_listed("search"));
assert!(!rules.is_enabled("delete"));
assert!(!rules.is_listed("delete"));
assert!(!rules.is_enabled("unknown"));
}
#[test]
fn test_component_visibility_rules_match_aliases() {
let rules = ComponentVisibilityRules::allow(["vault://public"]);
assert!(rules.is_enabled_any(["public_resource", "vault://public"]));
assert!(!rules.is_enabled_any(["public_resource", "vault://private"]));
}
#[test]
fn test_component_visibility_rules_can_hide_without_disabling() {
let rules = ComponentVisibilityRules::hide(["advanced_tool"]);
assert!(rules.is_enabled("advanced_tool"));
assert!(!rules.is_listed("advanced_tool"));
assert!(rules.is_enabled("public_tool"));
assert!(rules.is_listed("public_tool"));
}
#[test]
fn test_visibility_config_round_trips_serialization() {
let config = VisibilityConfig::new()
.with_allowed_tools(["search", "read_note"])
.with_disabled_tools(["delete_note"])
.with_hidden_tools(["advanced_graph"])
.with_allowed_resources(["vault://public"])
.with_allowed_prompts(["summarize"])
.require_read_only_tools();
let json = serde_json::to_string(&config).expect("visibility config serializes");
let decoded: VisibilityConfig =
serde_json::from_str(&json).expect("visibility config deserializes");
assert_eq!(decoded, config);
}
#[test]
fn test_empty_tool_allowlist_hides_all_tools() {
let layer =
VisibilityLayer::new(MockHandler).with_allowed_tools(std::iter::empty::<&str>());
assert!(layer.list_tools().is_empty());
}
#[test]
fn test_conflicting_read_only_and_destructive_hints_fail_closed() {
let tool = Tool {
name: "conflicting_tool".to_string(),
annotations: Some(
ToolAnnotations::default()
.with_read_only(true)
.with_destructive(true),
),
..Default::default()
};
assert!(!is_explicit_read_only_tool(&tool));
}
#[test]
fn test_visibility_layer_hides_admin() {
let layer = VisibilityLayer::new(MockHandler).disable_tags(["admin"]);
let tools = layer.list_tools();
assert_eq!(tools.len(), 1);
assert_eq!(tools[0].name, "public_tool");
}
#[test]
fn test_visibility_layer_shows_all_by_default() {
let layer = VisibilityLayer::new(MockHandler);
let tools = layer.list_tools();
assert_eq!(tools.len(), 2);
}
#[test]
fn test_exact_tool_allowlist_reduces_list_surface() {
let layer = VisibilityLayer::new(MockHandler).with_allowed_tools(["public_tool"]);
assert_eq!(tool_names(&layer), vec!["public_tool"]);
}
#[test]
fn test_exact_tool_denylist_wins_over_allowlist() {
let layer = VisibilityLayer::new(MockHandler)
.with_allowed_tools(["public_tool", "admin_tool"])
.with_disabled_tools(["public_tool"]);
assert_eq!(tool_names(&layer), vec!["admin_tool"]);
}
#[test]
fn test_layer_clone_has_independent_exact_rules() {
let base = VisibilityLayer::new(MockHandler);
let narrowed = base.clone().with_allowed_tools(["public_tool"]);
assert_eq!(base.list_tools().len(), 2);
assert_eq!(tool_names(&narrowed), vec!["public_tool"]);
}
#[test]
fn test_layer_clone_has_independent_tag_filters() {
let base = VisibilityLayer::new(MockHandler);
let narrowed = base.clone().disable_tags(["admin"]);
assert_eq!(base.list_tools().len(), 2);
assert_eq!(tool_names(&narrowed), vec!["public_tool"]);
}
#[test]
fn test_with_disabled_tools_replaces_previous_denylist() {
let layer = VisibilityLayer::new(MockHandler)
.with_disabled_tools(["public_tool"])
.with_disabled_tools(["admin_tool"]);
assert_eq!(tool_names(&layer), vec!["public_tool"]);
}
#[tokio::test]
async fn test_hidden_tool_is_not_listed_but_remains_callable() {
let layer = VisibilityLayer::new(MockHandler).with_hidden_tools(["public_tool"]);
let ctx = RequestContext::default();
assert_eq!(tool_names(&layer), vec!["admin_tool"]);
let result = layer
.call_tool("public_tool", serde_json::json!({}), &ctx)
.await
.expect("hidden but enabled tool should remain callable");
assert_eq!(result.first_text(), Some("Called public_tool"));
}
#[tokio::test]
async fn test_hidden_resource_and_prompt_are_not_listed_but_remain_callable() {
let layer = VisibilityLayer::new(MockHandler)
.with_hidden_resources(["vault://public"])
.with_hidden_prompts(["public_prompt"]);
let ctx = RequestContext::default();
assert_eq!(layer.list_resources().len(), 1);
assert_eq!(layer.list_prompts().len(), 1);
let resource = layer
.read_resource("vault://public", &ctx)
.await
.expect("hidden but enabled resource should remain readable");
assert_eq!(resource.first_text(), Some("Read vault://public"));
let prompt = layer
.get_prompt("public_prompt", None, &ctx)
.await
.expect("hidden but enabled prompt should remain gettable");
assert_eq!(
prompt.messages[0].content.as_text(),
Some("Prompt public_prompt")
);
}
#[test]
fn test_hidden_only_profile_still_advertises_operation_capabilities() {
let layer = VisibilityLayer::new(MockHandler)
.with_hidden_tools(["public_tool", "admin_tool"])
.with_hidden_resources(["vault://public", "vault://admin"])
.with_hidden_resource_templates(["vault://notes/{id}"])
.with_hidden_prompts(["public_prompt", "admin_prompt"]);
assert!(layer.list_tools().is_empty());
assert!(layer.list_resources().is_empty());
assert!(layer.list_resource_templates().is_empty());
assert!(layer.list_prompts().is_empty());
let capabilities = layer.server_capabilities();
assert!(
capabilities.tools.is_some(),
"hidden-but-callable tools still require the tools capability"
);
assert!(
capabilities.resources.is_some(),
"hidden-but-readable resources still require the resources capability"
);
assert!(
capabilities.prompts.is_some(),
"hidden-but-gettable prompts still require the prompts capability"
);
}
#[tokio::test]
async fn test_disabled_tool_wins_over_hidden_tool() {
let layer = VisibilityLayer::new(MockHandler)
.with_hidden_tools(["public_tool"])
.with_disabled_tools(["public_tool"]);
let ctx = RequestContext::default();
assert!(!tool_names(&layer).contains(&"public_tool".to_string()));
let err = layer
.call_tool("public_tool", serde_json::json!({}), &ctx)
.await
.expect_err("disabled tool should not be callable even if hidden");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::ToolNotFound);
}
#[test]
fn test_tag_disabled_tool_stays_hidden_despite_name_allowlist() {
let layer = VisibilityLayer::new(MockHandler)
.with_allowed_tools(["admin_tool"])
.disable_tags(["admin"]);
assert!(layer.list_tools().is_empty());
}
#[test]
fn test_visibility_config_getter_reflects_builder_methods() {
let layer = VisibilityLayer::new(MockHandler)
.with_allowed_tools(["public_tool"])
.with_disabled_resources(["vault://admin"])
.with_hidden_prompts(["admin_prompt"])
.require_read_only_tools();
let config = layer.visibility_config();
assert!(config.tools.allow.unwrap().contains("public_tool"));
assert!(config.resources.deny.contains("vault://admin"));
assert!(config.prompts.hide.contains("admin_prompt"));
assert!(config.require_read_only_tools);
}
#[tokio::test]
async fn test_disabled_tool_call_returns_not_found() {
let layer = VisibilityLayer::new(MockHandler).with_disabled_tools(["public_tool"]);
let ctx = RequestContext::default();
let err = layer
.call_tool("public_tool", serde_json::json!({}), &ctx)
.await
.expect_err("hidden tool calls should be rejected");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::ToolNotFound);
}
#[tokio::test]
async fn test_session_enable_allows_hidden_tagged_tool_call() {
let layer = VisibilityLayer::new(MockHandler).disable_tags(["admin"]);
let ctx = RequestContext::default().with_session_id("session1");
let err = layer
.call_tool("admin_tool", serde_json::json!({}), &ctx)
.await
.expect_err("globally hidden tagged tool should be rejected");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::ToolNotFound);
layer.enable_for_session("session1", &["admin".to_string()]);
let result = layer
.call_tool("admin_tool", serde_json::json!({}), &ctx)
.await
.expect("session-enabled tagged tool should pass through");
assert_eq!(result.first_text(), Some("Called admin_tool"));
}
#[tokio::test]
async fn test_session_disable_blocks_visible_tagged_tool_call() {
let layer = VisibilityLayer::new(MockHandler);
let ctx = RequestContext::default().with_session_id("session1");
let result = layer
.call_tool("public_tool", serde_json::json!({}), &ctx)
.await
.expect("public tool should initially pass through");
assert_eq!(result.first_text(), Some("Called public_tool"));
layer.disable_for_session("session1", &["public".to_string()]);
let err = layer
.call_tool("public_tool", serde_json::json!({}), &ctx)
.await
.expect_err("session-disabled tagged tool should be rejected");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::ToolNotFound);
}
#[tokio::test]
async fn test_dispatch_uses_registered_tool_without_relisting() {
let handler = CountingHandler::default();
let state = Arc::clone(&handler.state);
let layer = VisibilityLayer::new(handler);
let ctx = RequestContext::default();
assert_eq!(tool_names(&layer), vec!["counted_tool"]);
assert_eq!(state.tool_lists.load(Ordering::SeqCst), 1);
layer
.call_tool("counted_tool", serde_json::json!({}), &ctx)
.await
.expect("registered tool should dispatch");
layer
.call_tool("counted_tool", serde_json::json!({}), &ctx)
.await
.expect("registered tool should dispatch again");
assert_eq!(
state.tool_lists.load(Ordering::SeqCst),
1,
"dispatch should use the registry populated by tools/list"
);
}
#[tokio::test]
async fn test_dispatch_lazily_builds_registry_once_per_component_family() {
let handler = CountingHandler::default();
let state = Arc::clone(&handler.state);
let layer = VisibilityLayer::new(handler);
let ctx = RequestContext::default();
layer
.call_tool("counted_tool", serde_json::json!({}), &ctx)
.await
.expect("registered tool should dispatch");
layer
.call_tool("counted_tool", serde_json::json!({}), &ctx)
.await
.expect("registered tool should dispatch again");
assert_eq!(state.tool_lists.load(Ordering::SeqCst), 1);
layer
.read_resource("counted://resource", &ctx)
.await
.expect("registered resource should dispatch");
layer
.read_resource("counted://resource", &ctx)
.await
.expect("registered resource should dispatch again");
assert_eq!(state.resource_lists.load(Ordering::SeqCst), 1);
layer
.get_prompt("counted_prompt", None, &ctx)
.await
.expect("registered prompt should dispatch");
layer
.get_prompt("counted_prompt", None, &ctx)
.await
.expect("registered prompt should dispatch again");
assert_eq!(state.prompt_lists.load(Ordering::SeqCst), 1);
}
#[tokio::test]
async fn test_refresh_component_registry_updates_dispatch_metadata() {
let read_only = Arc::new(AtomicBool::new(false));
let layer = VisibilityLayer::new(MutableToolHandler {
read_only: Arc::clone(&read_only),
})
.require_read_only_tools();
let ctx = RequestContext::default();
let err = layer
.call_tool("mutable_tool", serde_json::json!({}), &ctx)
.await
.expect_err("initial destructive metadata should be rejected");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::ToolNotFound);
read_only.store(true, Ordering::SeqCst);
let err = layer
.call_tool("mutable_tool", serde_json::json!({}), &ctx)
.await
.expect_err("cached destructive metadata should remain fail-closed");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::ToolNotFound);
layer.refresh_component_registry();
let result = layer
.call_tool("mutable_tool", serde_json::json!({}), &ctx)
.await
.expect("refreshed read-only metadata should permit dispatch");
assert_eq!(result.first_text(), Some("Called mutable_tool"));
}
#[tokio::test]
async fn test_registry_preserves_first_duplicate_tool_metadata() {
let layer = VisibilityLayer::new(DuplicateToolMetadataHandler).require_read_only_tools();
let ctx = RequestContext::default();
let result = layer
.call_tool("duplicate_tool", serde_json::json!({}), &ctx)
.await
.expect("first listed read-only metadata should govern dispatch");
assert_eq!(result.first_text(), Some("Called duplicate_tool"));
}
#[tokio::test]
async fn test_exact_tool_policy_blocks_unlisted_dynamic_call() {
let layer = VisibilityLayer::new(DynamicHandler).with_disabled_tools(["dynamic_tool"]);
let ctx = RequestContext::default();
let err = layer
.call_tool("dynamic_tool", serde_json::json!({}), &ctx)
.await
.expect_err("denylisted dynamic tool calls should be rejected");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::ToolNotFound);
}
#[tokio::test]
async fn test_exact_tool_allowlist_can_permit_unlisted_dynamic_call() {
let layer = VisibilityLayer::new(DynamicHandler).with_allowed_tools(["dynamic_tool"]);
let ctx = RequestContext::default();
let result = layer
.call_tool("dynamic_tool", serde_json::json!({}), &ctx)
.await
.expect("allowlisted dynamic tool should pass through");
assert_eq!(result.first_text(), Some("Dynamic dynamic_tool"));
}
#[tokio::test]
async fn test_read_only_policy_blocks_unlisted_dynamic_tool() {
let layer = VisibilityLayer::new(DynamicHandler)
.with_allowed_tools(["dynamic_tool"])
.require_read_only_tools();
let ctx = RequestContext::default();
let err = layer
.call_tool("dynamic_tool", serde_json::json!({}), &ctx)
.await
.expect_err("read-only policy should fail closed without annotations");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::ToolNotFound);
}
#[test]
fn test_require_read_only_tools_hides_mutating_tools() {
let layer = VisibilityLayer::new(MockHandler).require_read_only_tools();
assert_eq!(tool_names(&layer), vec!["public_tool"]);
}
#[tokio::test]
async fn test_disabled_resource_read_returns_not_found() {
let layer = VisibilityLayer::new(MockHandler).with_disabled_resources(["vault://public"]);
let ctx = RequestContext::default();
let err = layer
.read_resource("vault://public", &ctx)
.await
.expect_err("hidden resource reads should be rejected");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::ResourceNotFound);
}
#[tokio::test]
async fn test_resource_allowlist_by_name_allows_uri_read() {
let layer = VisibilityLayer::new(MockHandler).with_allowed_resources(["public_resource"]);
let ctx = RequestContext::default();
let result = layer
.read_resource("vault://public", &ctx)
.await
.expect("allowlisted resource name should permit URI read");
assert_eq!(result.first_text(), Some("Read vault://public"));
}
#[tokio::test]
async fn test_exact_resource_policy_blocks_unlisted_dynamic_read() {
let layer =
VisibilityLayer::new(DynamicHandler).with_disabled_resources(["vault://dynamic"]);
let ctx = RequestContext::default();
let err = layer
.read_resource("vault://dynamic", &ctx)
.await
.expect_err("denylisted dynamic resources should be rejected");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::ResourceNotFound);
}
#[tokio::test]
async fn test_disabled_prompt_get_returns_not_found() {
let layer = VisibilityLayer::new(MockHandler).with_disabled_prompts(["public_prompt"]);
let ctx = RequestContext::default();
let err = layer
.get_prompt("public_prompt", None, &ctx)
.await
.expect_err("hidden prompts should be rejected");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::PromptNotFound);
}
#[tokio::test]
async fn test_exact_prompt_policy_blocks_unlisted_dynamic_get() {
let layer = VisibilityLayer::new(DynamicHandler).with_disabled_prompts(["dynamic_prompt"]);
let ctx = RequestContext::default();
let err = layer
.get_prompt("dynamic_prompt", None, &ctx)
.await
.expect_err("denylisted dynamic prompts should be rejected");
assert_eq!(err.kind, turbomcp_core::error::ErrorKind::PromptNotFound);
}
#[test]
fn test_visibility_config_applies_component_rules() {
let config = VisibilityConfig::new()
.with_allowed_tools(["public_tool"])
.with_disabled_resources(["vault://admin"])
.with_allowed_prompts(["public_prompt"])
.with_allowed_resource_templates(["vault://notes/{id}"]);
let layer = VisibilityLayer::new(MockHandler).with_visibility_config(config);
assert_eq!(tool_names(&layer), vec!["public_tool"]);
let resources = layer.list_resources();
assert_eq!(resources.len(), 1);
assert_eq!(resources[0].name, "public_resource");
let prompts = layer.list_prompts();
assert_eq!(prompts.len(), 1);
assert_eq!(prompts[0].name, "public_prompt");
let templates = layer.list_resource_templates();
assert_eq!(templates.len(), 1);
assert_eq!(templates[0].name, "note_template");
}
#[test]
fn test_session_enable_override() {
let layer = VisibilityLayer::new(MockHandler).disable_tags(["admin"]);
assert_eq!(layer.list_tools().len(), 1);
layer.enable_for_session("session1", &["admin".to_string()]);
assert_eq!(layer.list_tools().len(), 1);
layer.clear_session("session1");
}
#[test]
fn test_session_guard_cleanup() {
let layer = VisibilityLayer::new(MockHandler).disable_tags(["admin"]);
{
let _guard = layer.session_guard("guard-session");
layer.enable_for_session("guard-session", &["admin".to_string()]);
layer.disable_for_session("guard-session", &["public".to_string()]);
assert!(layer.active_sessions_count() > 0);
}
assert_eq!(layer.active_sessions_count(), 0);
}
#[test]
fn test_active_sessions_count() {
let layer = VisibilityLayer::new(MockHandler);
assert_eq!(layer.active_sessions_count(), 0);
layer.enable_for_session("session1", &["tag1".to_string()]);
assert_eq!(layer.active_sessions_count(), 1);
layer.disable_for_session("session2", &["tag2".to_string()]);
assert_eq!(layer.active_sessions_count(), 2);
layer.enable_for_session("session1", &["tag2".to_string()]);
assert_eq!(layer.active_sessions_count(), 2);
layer.clear_session("session1");
assert_eq!(layer.active_sessions_count(), 1);
layer.clear_session("session2");
assert_eq!(layer.active_sessions_count(), 0);
}
}