use std::collections::{BTreeMap, BTreeSet};
use mcp_conformance_core::trace::{Direction, TraceEvent};
use serde_json::{Map, Value};
use super::super::FindingSink;
use crate::context::TraceContext;
#[cfg(test)]
mod tests;
const LISTEN: &str = "subscriptions/listen";
const SUBSCRIPTION_ID: &str = "io.modelcontextprotocol/subscriptionId";
const ACKNOWLEDGED: &str = "notifications/subscriptions/acknowledged";
const FILTERED: &[(&str, &str)] = &[
("notifications/tools/list_changed", "toolsListChanged"),
("notifications/prompts/list_changed", "promptsListChanged"),
(
"notifications/resources/list_changed",
"resourcesListChanged",
),
];
const RESOURCE_SUBSCRIPTIONS: &str = "resourceSubscriptions";
const RESOURCE_UPDATED: &str = "notifications/resources/updated";
#[derive(Debug, Clone, Copy)]
struct Subscription<'a> {
seq: u64,
filter: Option<&'a Map<String, Value>>,
}
fn subscriptions<'a>(context: &'a TraceContext<'_>) -> BTreeMap<String, Subscription<'a>> {
context
.messages()
.filter_map(|(event, _, _)| {
if event.direction != Direction::ClientToServer {
return None;
}
let payload = event.message_payload()?;
if payload.get("method")?.as_str()? != LISTEN {
return None;
}
let id = payload.get("id").filter(|id| !id.is_null())?;
Some((
id.to_string(),
Subscription {
seq: event.seq,
filter: payload
.get("params")
.and_then(|params| params.get("notifications"))
.and_then(Value::as_object),
},
))
})
.collect()
}
fn tagged<'a>(context: &'a TraceContext<'_>) -> Vec<(u64, String, Option<&'a str>, &'a Value)> {
context
.messages()
.filter_map(|(event, _, _)| tagged_message(event))
.collect()
}
fn tagged_message(event: &TraceEvent) -> Option<(u64, String, Option<&str>, &Value)> {
if event.direction != Direction::ServerToClient {
return None;
}
let payload = event.message_payload()?;
let params = payload.get("params").or_else(|| payload.get("result"))?;
let id = params.get("_meta")?.get(SUBSCRIPTION_ID)?;
Some((
event.seq,
id.to_string(),
payload.get("method").and_then(Value::as_str),
params,
))
}
pub(in crate::checks) fn only_requested_notifications(
context: &TraceContext<'_>,
sink: &mut FindingSink,
) {
let subscriptions = subscriptions(context);
for (seq, id, method, params) in tagged(context) {
let (Some(method), Some(subscription)) = (method, subscriptions.get(&id)) else {
continue;
};
if method == ACKNOWLEDGED {
continue;
}
sink.examined();
if let Some(reason) = unrequested(subscription.filter, method, params) {
sink.push(
Some(seq),
format!("subscription {id} was sent `{method}`, which {reason}"),
);
}
}
}
fn unrequested(
filter: Option<&Map<String, Value>>,
method: &str,
params: &Value,
) -> Option<String> {
if let Some((_, field)) = FILTERED.iter().find(|(name, _)| *name == method) {
let asked = filter
.and_then(|filter| filter.get(*field))
.is_some_and(|value| value.as_bool() == Some(true));
return (!asked).then(|| format!("its filter did not set `{field}`"));
}
if method == RESOURCE_UPDATED {
let uri = params
.get("uri")
.and_then(Value::as_str)
.unwrap_or_default();
let listed = filter
.and_then(|filter| filter.get(RESOURCE_SUBSCRIPTIONS))
.and_then(Value::as_array)
.is_some_and(|uris| uris.iter().any(|value| value.as_str() == Some(uri)));
return (!listed)
.then(|| format!("its `{RESOURCE_SUBSCRIPTIONS}` does not list the URI {uri:?}"));
}
Some("is not one of the notification types the filter can request".to_owned())
}
pub(in crate::checks) fn acknowledgment_first(context: &TraceContext<'_>, sink: &mut FindingSink) {
let subscriptions = subscriptions(context);
let mut open: BTreeMap<String, u64> = BTreeMap::new();
let mut decided: BTreeSet<String> = BTreeSet::new();
for (event, _, _) in context.messages() {
if let Some((id, subscription)) = subscriptions
.iter()
.find(|(_, subscription)| subscription.seq == event.seq)
{
open.insert(id.clone(), subscription.seq);
continue;
}
let Some((seq, id, method, _)) = tagged_message(event) else {
continue;
};
if !open.contains_key(&id) || !decided.insert(id.clone()) {
continue;
}
sink.examined();
match method {
Some(ACKNOWLEDGED) => {}
Some(other) => sink.push(
Some(seq),
format!(
"subscription {id} opened with `{other}`; `{ACKNOWLEDGED}` must be its \
first message"
),
),
None => sink.push(
Some(seq),
format!(
"subscription {id} was closed by its `{LISTEN}` response before any \
`{ACKNOWLEDGED}` was sent"
),
),
}
}
}
pub(in crate::checks) fn graceful_close_result_shape(
context: &TraceContext<'_>,
sink: &mut FindingSink,
) {
for exchange in context.exchanges_for(LISTEN) {
let Some(result) = exchange.result.and_then(Value::as_object) else {
continue;
};
sink.examined();
let extra: Vec<&String> = result
.keys()
.filter(|key| *key != "resultType" && *key != "_meta")
.collect();
if !extra.is_empty() {
sink.push(
Some(exchange.response.seq),
format!(
"the `{LISTEN}` response carries {}; a graceful closure's result \
carries no method-specific data beyond `resultType` and `_meta`",
extra
.iter()
.map(|key| format!("`{key}`"))
.collect::<Vec<_>>()
.join(", ")
),
);
}
}
}