Skip to content

Commit 619a260

Browse files
DaleSeoDaleSeo
committed
fix: pass client header conformance
1 parent 61b40af commit 619a260

2 files changed

Lines changed: 167 additions & 19 deletions

File tree

‎conformance/src/bin/client.rs‎

Lines changed: 54 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ struct ConformanceToolCall {
2222

2323
#[derive(Debug, Default, serde::Deserialize)]
2424
struct ConformanceContext {
25-
#[serde(default)]
25+
#[serde(default, alias = "toolCalls")]
2626
tool_calls: Vec<ConformanceToolCall>,
2727
#[serde(default)]
2828
client_id: Option<String>,
@@ -859,8 +859,32 @@ async fn run_basic_client(server_url: &str) -> anyhow::Result<()> {
859859
}
860860

861861
async fn run_tools_call_client(server_url: &str, ctx: &ConformanceContext) -> anyhow::Result<()> {
862+
run_tools_call_client_with_lifecycle(server_url, ctx, ClientLifecycleMode::Initialize).await
863+
}
864+
865+
async fn run_discover_tools_call_client(
866+
server_url: &str,
867+
ctx: &ConformanceContext,
868+
) -> anyhow::Result<()> {
869+
run_tools_call_client_with_lifecycle(
870+
server_url,
871+
ctx,
872+
ClientLifecycleMode::Discover {
873+
preferred_versions: preferred_protocol_versions(),
874+
},
875+
)
876+
.await
877+
}
878+
879+
async fn run_tools_call_client_with_lifecycle(
880+
server_url: &str,
881+
ctx: &ConformanceContext,
882+
lifecycle: ClientLifecycleMode,
883+
) -> anyhow::Result<()> {
862884
let transport = StreamableHttpClientTransport::from_uri(server_url);
863-
let client = FullClientHandler.serve(transport).await?;
885+
let client = FullClientHandler
886+
.serve_with_lifecycle(transport, lifecycle)
887+
.await?;
864888
let tools = client.list_tools(Default::default()).await?;
865889

866890
if ctx.tool_calls.is_empty() {
@@ -1000,7 +1024,7 @@ async fn main() -> anyhow::Result<()> {
10001024
run_discover_client(&server_url).await?
10011025
}
10021026
"http-standard-headers" | "http-custom-headers" | "http-invalid-tool-headers" => {
1003-
run_tools_call_client(&server_url, &ctx).await?
1027+
run_discover_tools_call_client(&server_url, &ctx).await?
10041028
}
10051029

10061030
// Auth scenarios - standard OAuth flow
@@ -1083,3 +1107,30 @@ async fn main() -> anyhow::Result<()> {
10831107

10841108
Ok(())
10851109
}
1110+
1111+
#[cfg(test)]
1112+
mod tests {
1113+
use super::*;
1114+
1115+
#[test]
1116+
fn conformance_context_accepts_camel_case_tool_calls() {
1117+
let context: ConformanceContext = serde_json::from_value(json!({
1118+
"toolCalls": [{
1119+
"name": "test_custom_headers",
1120+
"arguments": { "region": "us-west1" }
1121+
}]
1122+
}))
1123+
.expect("valid conformance context");
1124+
1125+
assert_eq!(
1126+
context.tool_calls.first().map(|tool_call| (
1127+
tool_call.name.as_str(),
1128+
tool_call
1129+
.arguments
1130+
.as_ref()
1131+
.and_then(|arguments| arguments.get("region")),
1132+
)),
1133+
Some(("test_custom_headers", Some(&json!("us-west1"))))
1134+
);
1135+
}
1136+
}

‎crates/rmcp/src/transport/streamable_http_client.rs‎

Lines changed: 113 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -83,19 +83,24 @@ fn request_version_headers(
8383

8484
fn 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

Comments
 (0)