@@ -83,19 +83,24 @@ fn request_version_headers(
8383
8484fn cache_tools_from_response (
8585 cache : & mut HashMap < String , Arc < JsonObject > > ,
86- message : & ServerJsonRpcMessage ,
86+ message : & mut ServerJsonRpcMessage ,
87+ protocol_version : & ProtocolVersion ,
8788) {
89+ if protocol_version < & ProtocolVersion :: STANDARD_HEADERS {
90+ return ;
91+ }
8892 if let ServerJsonRpcMessage :: Response ( response) = message {
89- if let ServerResult :: ListToolsResult ( list) = & response. result {
90- for tool in & list. tools {
91- if let Err ( reason) =
93+ if let ServerResult :: ListToolsResult ( list) = & mut response. result {
94+ list. tools . retain ( |tool| {
95+ let Err ( reason) =
9296 mcp_headers:: validate_param_header_annotations ( & tool. input_schema )
93- {
94- tracing:: warn!( tool = %tool. name, "ignoring x-mcp-header annotations: {reason}" ) ;
95- continue ;
96- }
97- cache. insert ( tool. name . to_string ( ) , tool. input_schema . clone ( ) ) ;
98- }
97+ else {
98+ cache. insert ( tool. name . to_string ( ) , tool. input_schema . clone ( ) ) ;
99+ return true ;
100+ } ;
101+ tracing:: warn!( tool = %tool. name, "rejecting invalid x-mcp-header annotations: {reason}" ) ;
102+ false
103+ } ) ;
99104 }
100105 }
101106}
@@ -1213,10 +1218,11 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
12131218 ) ;
12141219 Ok ( ( ) )
12151220 }
1216- Ok ( StreamableHttpPostResponse :: Json ( msg, ..) ) => {
1221+ Ok ( StreamableHttpPostResponse :: Json ( mut msg, ..) ) => {
12171222 cache_tools_from_response (
12181223 & mut tool_header_cache,
1219- & msg,
1224+ & mut msg,
1225+ & negotiated_version,
12201226 ) ;
12211227 context. send_to_handler ( msg) . await ?;
12221228 Ok ( ( ) )
@@ -1262,8 +1268,12 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
12621268 tracing:: trace!( "client message accepted" ) ;
12631269 Ok ( ( ) )
12641270 }
1265- Ok ( StreamableHttpPostResponse :: Json ( message, ..) ) => {
1266- cache_tools_from_response ( & mut tool_header_cache, & message) ;
1271+ Ok ( StreamableHttpPostResponse :: Json ( mut message, ..) ) => {
1272+ cache_tools_from_response (
1273+ & mut tool_header_cache,
1274+ & mut message,
1275+ & negotiated_version,
1276+ ) ;
12671277 context. send_to_handler ( message) . await ?;
12681278 Ok ( ( ) )
12691279 }
@@ -1311,12 +1321,16 @@ impl<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
13111321 }
13121322 let _ = responder. send ( send_result) ;
13131323 }
1314- Event :: ServerMessage ( json_rpc_message) => {
1324+ Event :: ServerMessage ( mut json_rpc_message) => {
13151325 Self :: clear_stream_response_pending (
13161326 & mut pending_stream_response_ids,
13171327 & json_rpc_message,
13181328 ) ;
1319- cache_tools_from_response ( & mut tool_header_cache, & json_rpc_message) ;
1329+ cache_tools_from_response (
1330+ & mut tool_header_cache,
1331+ & mut json_rpc_message,
1332+ & negotiated_version,
1333+ ) ;
13201334 // send the message to the handler
13211335 if let Err ( e) = context. send_to_handler ( json_rpc_message) . await {
13221336 break ' main_loop Err ( e) ;
@@ -1669,3 +1683,86 @@ impl Default for StreamableHttpClientTransportConfig {
16691683 }
16701684 }
16711685}
1686+
1687+ #[ cfg( test) ]
1688+ mod tests {
1689+ use serde_json:: json;
1690+
1691+ use super :: * ;
1692+ use crate :: model:: { ListToolsResult , NumberOrString , ServerResult , Tool } ;
1693+
1694+ fn tool ( name : & ' static str , annotation : serde_json:: Value ) -> Tool {
1695+ let schema = json ! ( {
1696+ "type" : "object" ,
1697+ "properties" : {
1698+ "value" : annotation,
1699+ } ,
1700+ } ) ;
1701+ Tool :: new (
1702+ name,
1703+ name,
1704+ Arc :: new ( schema. as_object ( ) . expect ( "object schema" ) . clone ( ) ) ,
1705+ )
1706+ }
1707+
1708+ #[ test]
1709+ fn cache_tools_removes_invalid_header_annotations ( ) {
1710+ let valid = tool (
1711+ "valid" ,
1712+ json ! ( { "type" : "string" , "x-mcp-header" : "Value" } ) ,
1713+ ) ;
1714+ let invalid = tool ( "invalid" , json ! ( { "type" : "string" , "x-mcp-header" : "" } ) ) ;
1715+ let mut message = ServerJsonRpcMessage :: response (
1716+ ServerResult :: ListToolsResult ( ListToolsResult :: with_all_items ( vec ! [ valid, invalid] ) ) ,
1717+ NumberOrString :: Number ( 1 ) ,
1718+ ) ;
1719+ let mut cache = HashMap :: new ( ) ;
1720+
1721+ cache_tools_from_response ( & mut cache, & mut message, & ProtocolVersion :: V_2026_07_28 ) ;
1722+
1723+ let ServerJsonRpcMessage :: Response ( response) = & mut message else {
1724+ panic ! ( "expected tools/list response" ) ;
1725+ } ;
1726+ let ServerResult :: ListToolsResult ( result) = & mut response. result else {
1727+ panic ! ( "expected tools/list result" ) ;
1728+ } ;
1729+ assert_eq ! (
1730+ (
1731+ result
1732+ . tools
1733+ . iter( )
1734+ . map( |tool| tool. name. as_ref( ) )
1735+ . collect:: <Vec <_>>( ) ,
1736+ cache. keys( ) . map( String :: as_str) . collect:: <Vec <_>>( ) ,
1737+ ) ,
1738+ ( vec![ "valid" ] , vec![ "valid" ] )
1739+ ) ;
1740+ }
1741+
1742+ #[ test]
1743+ fn cache_tools_preserves_pre_standard_header_results ( ) {
1744+ let invalid = tool ( "legacy" , json ! ( { "type" : "string" , "x-mcp-header" : "" } ) ) ;
1745+ let mut message = ServerJsonRpcMessage :: response (
1746+ ServerResult :: ListToolsResult ( ListToolsResult :: with_all_items ( vec ! [ invalid] ) ) ,
1747+ NumberOrString :: Number ( 1 ) ,
1748+ ) ;
1749+ let mut cache = HashMap :: new ( ) ;
1750+
1751+ cache_tools_from_response ( & mut cache, & mut message, & ProtocolVersion :: V_2025_11_25 ) ;
1752+
1753+ let ServerJsonRpcMessage :: Response ( response) = message else {
1754+ panic ! ( "expected tools/list response" ) ;
1755+ } ;
1756+ let ServerResult :: ListToolsResult ( result) = response. result else {
1757+ panic ! ( "expected tools/list result" ) ;
1758+ } ;
1759+ assert_eq ! (
1760+ result
1761+ . tools
1762+ . iter( )
1763+ . map( |tool| tool. name. as_ref( ) )
1764+ . collect:: <Vec <_>>( ) ,
1765+ vec![ "legacy" ]
1766+ ) ;
1767+ }
1768+ }
0 commit comments