fix: pass client header conformance (#1012)

This commit is contained in:
Dale Seo 2026-07-20 14:18:53 -04:00 • committed by GitHub
commit 7abbf29e93
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 167 additions and 19 deletions

View file

@ -22,7 +22,7 @@ struct ConformanceToolCall {
#[derive(Debug, Default, serde::Deserialize)]
struct ConformanceContext {
#[serde(default)]
#[serde(default, alias = "toolCalls")]
tool_calls: Vec<ConformanceToolCall>,
#[serde(default)]
client_id: Option<String>,
@ -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"))))
);
}
}

View file

@ -83,19 +83,24 @@ fn request_version_headers(
fn cache_tools_from_response(
cache: &mut HashMap<String, Arc<JsonObject>>,
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<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
);
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<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
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<C: StreamableHttpClient> Worker for StreamableHttpClientWorker<C> {
}
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::<Vec<_>>(),
cache.keys().map(String::as_str).collect::<Vec<_>>(),
),
(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<_>>(),
vec!["legacy"]
);
}
}