diff --git a/crates/mcp/src/core/mod.rs b/crates/mcp/src/core/mod.rs index d323e4b90..99f8ef834 100644 --- a/crates/mcp/src/core/mod.rs +++ b/crates/mcp/src/core/mod.rs @@ -27,4 +27,4 @@ pub use orchestrator::{ }; pub use pool::{McpConnectionPool, PoolKey}; pub use reconnect::ReconnectionManager; -pub use session::{McpServerBinding, McpToolSession, DEFAULT_SERVER_LABEL}; +pub use session::{McpServerBinding, McpToolExposureFilter, McpToolSession, DEFAULT_SERVER_LABEL}; diff --git a/crates/mcp/src/core/session.rs b/crates/mcp/src/core/session.rs index 8e72544e5..c73940e6c 100644 --- a/crates/mcp/src/core/session.rs +++ b/crates/mcp/src/core/session.rs @@ -45,11 +45,62 @@ pub struct McpServerBinding { pub label: String, /// Internal key used to look up the server in the orchestrator. pub server_key: String, - /// Optional per-server tool allowlist. + /// Optional per-server tool exposure filter. /// - /// When `Some`, only the listed tool names are exposed for this server. + /// When `Some`, only tools matching the filter are exposed for this server. /// When `None`, all tools from the server are exposed. - pub allowed_tools: Option>, + pub allowed_tools: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct McpToolExposureFilter { + pub tool_names: Option>, + pub read_only: Option, +} + +impl McpToolExposureFilter { + pub fn tool_names(tool_names: Vec) -> Self { + Self { + tool_names: Some(tool_names), + read_only: None, + } + } + + pub fn read_only(read_only: bool) -> Self { + Self { + tool_names: None, + read_only: Some(read_only), + } + } + + pub fn new(tool_names: Option>, read_only: Option) -> Self { + Self { + tool_names, + read_only, + } + } +} + +struct PreparedMcpToolExposureFilter<'a> { + tool_names: Option>, + read_only: Option, +} + +impl<'a> PreparedMcpToolExposureFilter<'a> { + fn from_filter(filter: &'a McpToolExposureFilter) -> Self { + let tool_names = filter.tool_names.as_ref().map(|names| { + names + .iter() + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .collect() + }); + + Self { + tool_names, + read_only: filter.read_only, + } + } } #[derive(Debug, Clone)] @@ -120,26 +171,25 @@ impl<'a> McpToolSession<'a> { let server_keys: Vec = mcp_servers.iter().map(|b| b.server_key.clone()).collect(); let mut mcp_tools = Self::collect_visible_mcp_tools(orchestrator, &server_keys); - // Build per-server allowlists from bindings that specify allowed_tools. - let allowed_tools_by_server_key: HashMap<&str, HashSet<&str>> = mcp_servers - .iter() - .filter_map(|b| { - b.allowed_tools.as_ref().map(|tools| { - let set: HashSet<&str> = tools - .iter() - .map(|s| s.trim()) - .filter(|s| !s.is_empty()) - .collect(); - (b.server_key.as_str(), set) + // Build per-server exposure filters from bindings that specify allowed_tools. + let allowed_tools_by_server_key: HashMap<&str, PreparedMcpToolExposureFilter<'_>> = + mcp_servers + .iter() + .filter_map(|b| { + b.allowed_tools.as_ref().map(|filter| { + ( + b.server_key.as_str(), + PreparedMcpToolExposureFilter::from_filter(filter), + ) + }) }) - }) - .collect(); + .collect(); if !allowed_tools_by_server_key.is_empty() { mcp_tools.retain(|entry| { match allowed_tools_by_server_key.get(Self::associated_server_key(entry)) { None => true, - Some(allowed) => Self::matches_allowed_tool_name(entry, allowed), + Some(filter) => Self::matches_allowed_tool_filter(entry, filter), } }); } @@ -830,6 +880,9 @@ impl<'a> McpToolSession<'a> { for direct_entry in direct_tools { if let Some(mut alias_entries) = aliases_by_target.remove(&direct_entry.qualified_name) { + for alias_entry in &mut alias_entries { + alias_entry.annotations = direct_entry.annotations.clone(); + } visible_tools.append(&mut alias_entries); } else { visible_tools.push(direct_entry); @@ -851,7 +904,20 @@ impl<'a> McpToolSession<'a> { .unwrap_or_else(|| entry.server_key()) } - fn matches_allowed_tool_name(entry: &ToolEntry, allowed: &HashSet<&str>) -> bool { + fn matches_allowed_tool_filter( + entry: &ToolEntry, + filter: &PreparedMcpToolExposureFilter<'_>, + ) -> bool { + if let Some(read_only) = filter.read_only { + if entry.annotations.read_only != read_only { + return false; + } + } + + let Some(allowed) = filter.tool_names.as_ref() else { + return true; + }; + allowed.contains(entry.tool_name()) || entry .alias_target @@ -917,7 +983,7 @@ mod tests { use serde_json::json; use super::*; - use crate::core::config::Tool as McpTool; + use crate::{annotations::ToolAnnotations, core::config::Tool as McpTool}; #[test] fn test_session_creation_keeps_servers() { @@ -1338,7 +1404,9 @@ mod tests { vec![McpServerBinding { label: "mock".to_string(), server_key: "server1".to_string(), - allowed_tools: Some(vec!["brave_web_search".to_string()]), + allowed_tools: Some(McpToolExposureFilter::tool_names(vec![ + "brave_web_search".to_string() + ])), }], "test-request", ); @@ -1352,6 +1420,87 @@ mod tests { assert_eq!(listed[0].tool_name(), "brave_web_search"); } + #[test] + fn test_allowed_tools_read_only_filter_uses_tool_annotations() { + let orchestrator = McpOrchestrator::new_test(); + + orchestrator.tool_inventory().insert_entry( + ToolEntry::from_server_tool("server1", create_test_tool("read_tool")) + .with_annotations(ToolAnnotations::new().with_read_only(true)), + ); + orchestrator.tool_inventory().insert_entry( + ToolEntry::from_server_tool("server1", create_test_tool("write_tool")) + .with_annotations(ToolAnnotations::new().with_read_only(false)), + ); + orchestrator + .tool_inventory() + .insert_entry(ToolEntry::from_server_tool( + "server1", + create_test_tool("missing_hint_tool"), + )); + + let session = McpToolSession::new( + &orchestrator, + vec![McpServerBinding { + label: "mock".to_string(), + server_key: "server1".to_string(), + allowed_tools: Some(McpToolExposureFilter::read_only(true)), + }], + "test-request", + ); + + let names: HashSet<&str> = session + .mcp_tools() + .iter() + .map(|entry| entry.tool_name()) + .collect(); + assert_eq!(names, HashSet::from(["read_tool"])); + assert!(session.has_exposed_tool("read_tool")); + assert!(!session.has_exposed_tool("write_tool")); + assert!(!session.has_exposed_tool("missing_hint_tool")); + } + + #[test] + fn test_allowed_tools_read_only_filter_combines_with_tool_names() { + let orchestrator = McpOrchestrator::new_test(); + + orchestrator.tool_inventory().insert_entry( + ToolEntry::from_server_tool("server1", create_test_tool("allowed_read_tool")) + .with_annotations(ToolAnnotations::new().with_read_only(true)), + ); + orchestrator.tool_inventory().insert_entry( + ToolEntry::from_server_tool("server1", create_test_tool("other_read_tool")) + .with_annotations(ToolAnnotations::new().with_read_only(true)), + ); + orchestrator.tool_inventory().insert_entry( + ToolEntry::from_server_tool("server1", create_test_tool("allowed_write_tool")) + .with_annotations(ToolAnnotations::new().with_read_only(false)), + ); + + let session = McpToolSession::new( + &orchestrator, + vec![McpServerBinding { + label: "mock".to_string(), + server_key: "server1".to_string(), + allowed_tools: Some(McpToolExposureFilter::new( + Some(vec![ + "allowed_read_tool".to_string(), + "allowed_write_tool".to_string(), + ]), + Some(true), + )), + }], + "test-request", + ); + + let names: HashSet<&str> = session + .mcp_tools() + .iter() + .map(|entry| entry.tool_name()) + .collect(); + assert_eq!(names, HashSet::from(["allowed_read_tool"])); + } + #[test] fn test_alias_tools_replace_target_tool_in_session_inventory() { let orchestrator = McpOrchestrator::new_test(); @@ -1427,7 +1576,9 @@ mod tests { vec![McpServerBinding { label: "brave".to_string(), server_key: "server1".to_string(), - allowed_tools: Some(vec!["web_search".to_string()]), + allowed_tools: Some(McpToolExposureFilter::tool_names(vec![ + "web_search".to_string() + ])), }], "test-request", ); @@ -1437,6 +1588,53 @@ mod tests { assert_eq!(session.mcp_tools()[0].tool_name(), "web_search"); } + #[test] + fn test_allowed_tools_read_only_filter_uses_alias_target_annotations() { + let orchestrator = McpOrchestrator::new_test(); + + orchestrator.tool_inventory().insert_entry( + ToolEntry::from_server_tool("server1", create_test_tool("brave_web_search")) + .with_annotations(ToolAnnotations::new().with_read_only(true)), + ); + + orchestrator + .register_alias( + "web_search", + "server1", + "brave_web_search", + None, + ResponseFormat::WebSearchCall, + ) + .expect("alias registration should succeed"); + + let read_only_session = McpToolSession::new( + &orchestrator, + vec![McpServerBinding { + label: "brave".to_string(), + server_key: "server1".to_string(), + allowed_tools: Some(McpToolExposureFilter::read_only(true)), + }], + "test-request", + ); + + assert!(read_only_session.has_exposed_tool("web_search")); + assert_eq!(read_only_session.mcp_tools().len(), 1); + assert_eq!(read_only_session.mcp_tools()[0].tool_name(), "web_search"); + + let write_session = McpToolSession::new( + &orchestrator, + vec![McpServerBinding { + label: "brave".to_string(), + server_key: "server1".to_string(), + allowed_tools: Some(McpToolExposureFilter::read_only(false)), + }], + "test-request", + ); + + assert!(!write_session.has_exposed_tool("web_search")); + assert!(write_session.mcp_tools().is_empty()); + } + #[test] fn test_is_internal_tool_for_internal_server() { use crate::core::config::{McpConfig, McpServerConfig, McpTransport}; @@ -1752,7 +1950,9 @@ mod tests { McpServerBinding { label: "brave".to_string(), server_key: "server1".to_string(), - allowed_tools: Some(vec!["brave_web_search".to_string()]), + allowed_tools: Some(McpToolExposureFilter::tool_names(vec![ + "brave_web_search".to_string(), + ])), }, McpServerBinding { label: "deepwiki".to_string(), diff --git a/crates/mcp/src/lib.rs b/crates/mcp/src/lib.rs index 352cd95ec..f5fb4bd1b 100644 --- a/crates/mcp/src/lib.rs +++ b/crates/mcp/src/lib.rs @@ -30,10 +30,11 @@ pub use core::{config, pool as connection_pool}; pub use core::{ ArgMappingConfig, BuiltinToolType, ConfigValidationError, HandlerRequestContext, LatencySnapshot, McpConfig, McpMetrics, McpOrchestrator, McpRequestContext, McpServerBinding, - McpServerConfig, McpToolSession, McpTransport, MetricsSnapshot, PendingToolExecution, - PolicyConfig, PolicyDecisionConfig, PoolKey, RefreshRequest, ResponseFormatConfig, - ServerPolicyConfig, SmgClientHandler, Tool, ToolCallResult, ToolConfig, ToolExecutionInput, - ToolExecutionOutput, ToolExecutionResult, TrustLevelConfig, DEFAULT_SERVER_LABEL, + McpServerConfig, McpToolExposureFilter, McpToolSession, McpTransport, MetricsSnapshot, + PendingToolExecution, PolicyConfig, PolicyDecisionConfig, PoolKey, RefreshRequest, + ResponseFormatConfig, ServerPolicyConfig, SmgClientHandler, Tool, ToolCallResult, ToolConfig, + ToolExecutionInput, ToolExecutionOutput, ToolExecutionResult, TrustLevelConfig, + DEFAULT_SERVER_LABEL, }; // Re-export shared types diff --git a/model_gateway/src/routers/anthropic/router.rs b/model_gateway/src/routers/anthropic/router.rs index 2b1cae62d..95dfe1f38 100644 --- a/model_gateway/src/routers/anthropic/router.rs +++ b/model_gateway/src/routers/anthropic/router.rs @@ -112,7 +112,9 @@ impl RouterTrait for AnthropicRouter { url: Some(server.url.clone()), authorization: server.authorization_token.clone(), headers: HashMap::new(), - allowed_tools: toolset_allowed.get(&server.name).and_then(|v| v.clone()), + allowed_tools: toolset_allowed + .get(&server.name) + .and_then(|v| v.clone().map(smg_mcp::McpToolExposureFilter::tool_names)), }) .collect(); diff --git a/model_gateway/src/routers/common/mcp_utils.rs b/model_gateway/src/routers/common/mcp_utils.rs index dd872f51b..ba0a21910 100644 --- a/model_gateway/src/routers/common/mcp_utils.rs +++ b/model_gateway/src/routers/common/mcp_utils.rs @@ -9,37 +9,33 @@ use openai_protocol::responses::{McpAllowedTools, ResponseTool, ResponsesRequest use serde_json::{json, Value}; use smg_mcp::{ apply_hosted_tool_overrides, extract_hosted_tool_overrides, BuiltinToolType, McpOrchestrator, - McpServerBinding, McpServerConfig, McpTransport, ResponseFormat, + McpServerBinding, McpServerConfig, McpToolExposureFilter, McpTransport, ResponseFormat, }; use tracing::{debug, warn}; /// Default maximum tool loop iterations (safety limit). pub const DEFAULT_MAX_ITERATIONS: usize = 10; -/// Project the T11 `McpAllowedTools` union into the flat name list consumed by +/// Project the T11 `McpAllowedTools` union into the exposure filter consumed by /// the router-side `McpServerInput` and `McpServerBinding` allowlist paths. /// /// Mapping (documented on the calling sites): /// * `None` → `None` (no constraint; all tools exposed) -/// * `Some(List(names))` → `Some(names)` -/// * `Some(Filter { tool_names: Some(v), read_only: None })` → `Some(v)` -/// * `Some(Filter { read_only: Some(_), .. })` → `Some(vec![])` — -/// fail-closed whenever the caller specified a `read_only` restriction, -/// regardless of `tool_names`. smg has no `readOnlyHint`-based filter -/// implementation yet, so honoring only the `tool_names` half would -/// silently broaden exposure past caller intent (e.g. `{read_only: true, -/// tool_names: ["mutating_tool"]}` must NOT expose `mutating_tool`). -/// Narrow to nothing so the downstream retain path exposes none. +/// * `Some(List(names))` → `Some({ tool_names: names })` +/// * `Some(Filter { tool_names, read_only })` → preserve both constraints. +/// `read_only` is evaluated against MCP `readOnlyHint` annotations after +/// listTools. /// * `Some(Filter { tool_names: None, read_only: None })` → `None` -pub(crate) fn project_allowed_tools(value: Option<&McpAllowedTools>) -> Option> { +pub(crate) fn project_allowed_tools( + value: Option<&McpAllowedTools>, +) -> Option { value.and_then(|at| match at { - McpAllowedTools::List(names) => Some(names.clone()), + McpAllowedTools::List(names) => Some(McpToolExposureFilter::tool_names(names.clone())), McpAllowedTools::Filter(filter) => match (&filter.tool_names, &filter.read_only) { - // Any `read_only` restriction fails closed: readOnlyHint-based - // filtering is unimplemented, so we cannot safely project a subset. - (_, Some(_)) => Some(Vec::new()), - (Some(names), None) => Some(names.clone()), (None, None) => None, + (tool_names, read_only) => { + Some(McpToolExposureFilter::new(tool_names.clone(), *read_only)) + } }, }) } @@ -53,8 +49,8 @@ pub struct McpServerInput { pub url: Option, pub authorization: Option, pub headers: HashMap, - /// Optional per-server tool allowlist. - pub allowed_tools: Option>, + /// Optional per-server tool exposure filter. + pub allowed_tools: Option, } /// Connect to MCP servers described by protocol-agnostic inputs. @@ -309,9 +305,7 @@ pub async fn ensure_request_mcp_client( authorization: mcp.authorization.clone(), headers: mcp.headers.clone().unwrap_or_default(), // T11: project the `allowed_tools` union (List | Filter) into - // the flat name list `McpServerInput` still expects. See - // [`project_allowed_tools`] for the mapping table and the - // rationale for the `read_only`-only fail-closed case. + // the router exposure filter used by MCP session construction. allowed_tools: project_allowed_tools(mcp.allowed_tools.as_ref()), }), _ => None, @@ -802,10 +796,12 @@ mod tests { #[test] fn test_project_allowed_tools_list_variant() { let value = McpAllowedTools::List(vec!["a".to_string(), "b".to_string()]); + let filter = project_allowed_tools(Some(&value)).expect("filter"); assert_eq!( - project_allowed_tools(Some(&value)), + filter.tool_names, Some(vec!["a".to_string(), "b".to_string()]) ); + assert_eq!(filter.read_only, None); } #[test] @@ -814,23 +810,20 @@ mod tests { read_only: None, tool_names: Some(vec!["x".to_string()]), }); - assert_eq!( - project_allowed_tools(Some(&value)), - Some(vec!["x".to_string()]) - ); + let filter = project_allowed_tools(Some(&value)).expect("filter"); + assert_eq!(filter.tool_names, Some(vec!["x".to_string()])); + assert_eq!(filter.read_only, None); } - /// `Filter { read_only: Some(true) }` with no `tool_names` must project - /// to `Some(vec![])` (fail-closed) so downstream does not expose the full - /// tool surface when the caller explicitly asked to restrict. See - /// [`project_allowed_tools`] rationale. #[test] - fn test_project_allowed_tools_filter_read_only_only_is_fail_closed() { + fn test_project_allowed_tools_filter_read_only_only_is_preserved() { let value = McpAllowedTools::Filter(McpToolFilter { read_only: Some(true), tool_names: None, }); - assert_eq!(project_allowed_tools(Some(&value)), Some(Vec::new())); + let filter = project_allowed_tools(Some(&value)).expect("filter"); + assert_eq!(filter.tool_names, None); + assert_eq!(filter.read_only, Some(true)); } /// `Filter { read_only: None, tool_names: None }` — both sub-fields @@ -842,19 +835,15 @@ mod tests { assert_eq!(project_allowed_tools(Some(&value)), None); } - /// `Filter { read_only: Some(_), tool_names: Some(_) }` — both set — must - /// still fail-closed. Any `read_only` restriction disables the normal - /// name-list projection because `readOnlyHint`-based filtering is - /// unimplemented; honoring only the `tool_names` half would broaden - /// exposure past caller intent (e.g. `tool_names: ["mutating_tool"]` with - /// `read_only: true` must not expose `mutating_tool`). #[test] - fn test_project_allowed_tools_filter_read_only_plus_names_is_fail_closed() { + fn test_project_allowed_tools_filter_read_only_plus_names_is_preserved() { let value = McpAllowedTools::Filter(McpToolFilter { read_only: Some(true), tool_names: Some(vec!["mutating_tool".to_string()]), }); - assert_eq!(project_allowed_tools(Some(&value)), Some(Vec::new())); + let filter = project_allowed_tools(Some(&value)).expect("filter"); + assert_eq!(filter.tool_names, Some(vec!["mutating_tool".to_string()])); + assert_eq!(filter.read_only, Some(true)); } /// When the synthesized function tool exposes `user` as a parameter, diff --git a/model_gateway/src/routers/openai/mcp/tool_loop.rs b/model_gateway/src/routers/openai/mcp/tool_loop.rs index 730cdd430..e8bbc02f7 100644 --- a/model_gateway/src/routers/openai/mcp/tool_loop.rs +++ b/model_gateway/src/routers/openai/mcp/tool_loop.rs @@ -1447,12 +1447,16 @@ mod tests { McpServerBinding { label: "deepwiki_ask".to_string(), server_key: "server-ask".to_string(), - allowed_tools: Some(vec!["ask_question".to_string()]), + allowed_tools: Some(smg_mcp::McpToolExposureFilter::tool_names(vec![ + "ask_question".to_string(), + ])), }, McpServerBinding { label: "deepwiki_read".to_string(), server_key: "server-read".to_string(), - allowed_tools: Some(vec!["read_wiki_structure".to_string()]), + allowed_tools: Some(smg_mcp::McpToolExposureFilter::tool_names(vec![ + "read_wiki_structure".to_string(), + ])), }, ]; @@ -1470,7 +1474,9 @@ mod tests { let bindings = vec![McpServerBinding { label: "deepwiki_ask".to_string(), server_key: "server-ask".to_string(), - allowed_tools: Some(vec!["ask_question".to_string()]), + allowed_tools: Some(smg_mcp::McpToolExposureFilter::tool_names(vec![ + "ask_question".to_string(), + ])), }]; let bindings_to_emit = mcp_list_tools_bindings_to_emit(&existing_labels, &bindings);