diff --git a/conformance/src/bin/client.rs b/conformance/src/bin/client.rs index 7f481f6bc..e82cc65af 100644 --- a/conformance/src/bin/client.rs +++ b/conformance/src/bin/client.rs @@ -22,7 +22,7 @@ struct ConformanceToolCall { #[derive(Debug, Default, serde::Deserialize)] struct ConformanceContext { - #[serde(default)] + #[serde(default, alias = "toolCalls")] tool_calls: Vec, #[serde(default)] client_id: Option, @@ -859,8 +859,32 @@ async fn run_basic_client(server_url: &str) -> anyhow::Result<()> { } async fn run_tools_call_client(server_url: &str, ctx: &ConformanceContext) -> anyhow::Result<()> { + run_tools_call_client_with_lifecycle(server_url, ctx, ClientLifecycleMode::Initialize).await +} + +async fn run_discover_tools_call_client( + server_url: &str, + ctx: &ConformanceContext, +) -> anyhow::Result<()> { + run_tools_call_client_with_lifecycle( + server_url, + ctx, + ClientLifecycleMode::Discover { + preferred_versions: preferred_protocol_versions(), + }, + ) + .await +} + +async fn run_tools_call_client_with_lifecycle( + server_url: &str, + ctx: &ConformanceContext, + lifecycle: ClientLifecycleMode, +) -> anyhow::Result<()> { let transport = StreamableHttpClientTransport::from_uri(server_url); - let client = FullClientHandler.serve(transport).await?; + let client = FullClientHandler + .serve_with_lifecycle(transport, lifecycle) + .await?; let tools = client.list_tools(Default::default()).await?; if ctx.tool_calls.is_empty() { @@ -1000,7 +1024,7 @@ async fn main() -> anyhow::Result<()> { run_discover_client(&server_url).await? } "http-standard-headers" | "http-custom-headers" | "http-invalid-tool-headers" => { - run_tools_call_client(&server_url, &ctx).await? + run_discover_tools_call_client(&server_url, &ctx).await? } // Auth scenarios - standard OAuth flow @@ -1083,3 +1107,30 @@ async fn main() -> anyhow::Result<()> { Ok(()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn conformance_context_accepts_camel_case_tool_calls() { + let context: ConformanceContext = serde_json::from_value(json!({ + "toolCalls": [{ + "name": "test_custom_headers", + "arguments": { "region": "us-west1" } + }] + })) + .expect("valid conformance context"); + + assert_eq!( + context.tool_calls.first().map(|tool_call| ( + tool_call.name.as_str(), + tool_call + .arguments + .as_ref() + .and_then(|arguments| arguments.get("region")), + )), + Some(("test_custom_headers", Some(&json!("us-west1")))) + ); + } +} diff --git a/crates/rmcp/src/transport/streamable_http_client.rs b/crates/rmcp/src/transport/streamable_http_client.rs index 8ba3e423c..e54a17373 100644 --- a/crates/rmcp/src/transport/streamable_http_client.rs +++ b/crates/rmcp/src/transport/streamable_http_client.rs @@ -83,19 +83,24 @@ fn request_version_headers( fn cache_tools_from_response( cache: &mut HashMap>, - message: &ServerJsonRpcMessage, + message: &mut ServerJsonRpcMessage, + protocol_version: &ProtocolVersion, ) { + if protocol_version < &ProtocolVersion::STANDARD_HEADERS { + return; + } if let ServerJsonRpcMessage::Response(response) = message { - if let ServerResult::ListToolsResult(list) = &response.result { - for tool in &list.tools { - if let Err(reason) = + if let ServerResult::ListToolsResult(list) = &mut response.result { + list.tools.retain(|tool| { + let Err(reason) = mcp_headers::validate_param_header_annotations(&tool.input_schema) - { - tracing::warn!(tool = %tool.name, "ignoring x-mcp-header annotations: {reason}"); - continue; - } - cache.insert(tool.name.to_string(), tool.input_schema.clone()); - } + else { + cache.insert(tool.name.to_string(), tool.input_schema.clone()); + return true; + }; + tracing::warn!(tool = %tool.name, "rejecting invalid x-mcp-header annotations: {reason}"); + false + }); } } } @@ -1213,10 +1218,11 @@ impl Worker for StreamableHttpClientWorker { ); Ok(()) } - Ok(StreamableHttpPostResponse::Json(msg, ..)) => { + Ok(StreamableHttpPostResponse::Json(mut msg, ..)) => { cache_tools_from_response( &mut tool_header_cache, - &msg, + &mut msg, + &negotiated_version, ); context.send_to_handler(msg).await?; Ok(()) @@ -1262,8 +1268,12 @@ impl Worker for StreamableHttpClientWorker { tracing::trace!("client message accepted"); Ok(()) } - Ok(StreamableHttpPostResponse::Json(message, ..)) => { - cache_tools_from_response(&mut tool_header_cache, &message); + Ok(StreamableHttpPostResponse::Json(mut message, ..)) => { + cache_tools_from_response( + &mut tool_header_cache, + &mut message, + &negotiated_version, + ); context.send_to_handler(message).await?; Ok(()) } @@ -1311,12 +1321,16 @@ impl Worker for StreamableHttpClientWorker { } let _ = responder.send(send_result); } - Event::ServerMessage(json_rpc_message) => { + Event::ServerMessage(mut json_rpc_message) => { Self::clear_stream_response_pending( &mut pending_stream_response_ids, &json_rpc_message, ); - cache_tools_from_response(&mut tool_header_cache, &json_rpc_message); + cache_tools_from_response( + &mut tool_header_cache, + &mut json_rpc_message, + &negotiated_version, + ); // send the message to the handler if let Err(e) = context.send_to_handler(json_rpc_message).await { break 'main_loop Err(e); @@ -1669,3 +1683,86 @@ impl Default for StreamableHttpClientTransportConfig { } } } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::model::{ListToolsResult, NumberOrString, ServerResult, Tool}; + + fn tool(name: &'static str, annotation: serde_json::Value) -> Tool { + let schema = json!({ + "type": "object", + "properties": { + "value": annotation, + }, + }); + Tool::new( + name, + name, + Arc::new(schema.as_object().expect("object schema").clone()), + ) + } + + #[test] + fn cache_tools_removes_invalid_header_annotations() { + let valid = tool( + "valid", + json!({ "type": "string", "x-mcp-header": "Value" }), + ); + let invalid = tool("invalid", json!({ "type": "string", "x-mcp-header": "" })); + let mut message = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::with_all_items(vec![valid, invalid])), + NumberOrString::Number(1), + ); + let mut cache = HashMap::new(); + + cache_tools_from_response(&mut cache, &mut message, &ProtocolVersion::V_2026_07_28); + + let ServerJsonRpcMessage::Response(response) = &mut message else { + panic!("expected tools/list response"); + }; + let ServerResult::ListToolsResult(result) = &mut response.result else { + panic!("expected tools/list result"); + }; + assert_eq!( + ( + result + .tools + .iter() + .map(|tool| tool.name.as_ref()) + .collect::>(), + cache.keys().map(String::as_str).collect::>(), + ), + (vec!["valid"], vec!["valid"]) + ); + } + + #[test] + fn cache_tools_preserves_pre_standard_header_results() { + let invalid = tool("legacy", json!({ "type": "string", "x-mcp-header": "" })); + let mut message = ServerJsonRpcMessage::response( + ServerResult::ListToolsResult(ListToolsResult::with_all_items(vec![invalid])), + NumberOrString::Number(1), + ); + let mut cache = HashMap::new(); + + cache_tools_from_response(&mut cache, &mut message, &ProtocolVersion::V_2025_11_25); + + let ServerJsonRpcMessage::Response(response) = message else { + panic!("expected tools/list response"); + }; + let ServerResult::ListToolsResult(result) = response.result else { + panic!("expected tools/list result"); + }; + assert_eq!( + result + .tools + .iter() + .map(|tool| tool.name.as_ref()) + .collect::>(), + vec!["legacy"] + ); + } +}