mcp_trace_validator/checks/draft/
mrtr.rs1use std::collections::BTreeMap;
29
30use mcp_conformance_core::trace::Direction;
31use serde_json::{Map, Value};
32
33use super::super::FindingSink;
34use crate::context::TraceContext;
35
36#[cfg(test)]
37mod tests;
38
39const SUPPORTED: &[&str] = &["prompts/get", "resources/read", "tools/call"];
41
42const INPUT_REQUEST_METHODS: &[&str] =
44 &["elicitation/create", "sampling/createMessage", "roots/list"];
45
46const INPUT_REQUIRED: &str = "input_required";
48
49#[derive(Debug, Clone, Copy)]
51struct Round<'a> {
52 seq: u64,
54 origin: (u64, &'a Value, &'a str),
56 requests: Option<&'a Map<String, Value>>,
58 state: Option<&'a Value>,
60}
61
62#[derive(Debug, Clone, Copy)]
64struct Retry<'a> {
65 seq: u64,
66 id: &'a Value,
67 method: &'a str,
68 responses: Option<&'a Map<String, Value>>,
69 state: Option<&'a Value>,
70}
71
72fn rounds<'a>(context: &'a TraceContext<'_>) -> Vec<Round<'a>> {
78 context
79 .exchanges()
80 .filter_map(|exchange| {
81 let result = exchange.result?;
82 if result.get("resultType").and_then(Value::as_str) != Some(INPUT_REQUIRED) {
83 return None;
84 }
85 let id = exchange.request.message_payload()?.get("id")?;
86 Some(Round {
87 seq: exchange.response.seq,
88 origin: (exchange.request.seq, id, exchange.method),
89 requests: result.get("inputRequests").and_then(Value::as_object),
90 state: result.get("requestState"),
91 })
92 })
93 .collect()
94}
95
96fn retries<'a>(context: &'a TraceContext<'_>) -> Vec<Retry<'a>> {
98 context
99 .messages()
100 .filter_map(|(event, _, _)| {
101 if event.direction != Direction::ClientToServer {
102 return None;
103 }
104 let payload = event.message_payload()?;
105 let method = payload.get("method")?.as_str()?;
106 let id = payload.get("id").filter(|id| !id.is_null())?;
107 let params = payload.get("params")?;
108 let responses = params.get("inputResponses").and_then(Value::as_object);
109 let state = params.get("requestState");
110 (responses.is_some() || state.is_some()).then_some(Retry {
111 seq: event.seq,
112 id,
113 method,
114 responses,
115 state,
116 })
117 })
118 .collect()
119}
120
121fn retries_with_rounds<'a>(context: &'a TraceContext<'_>) -> Vec<(Retry<'a>, Option<Round<'a>>)> {
129 let rounds: BTreeMap<u64, Round<'a>> = rounds(context)
130 .into_iter()
131 .map(|round| (round.seq, round))
132 .collect();
133 let retries: BTreeMap<u64, Retry<'a>> = retries(context)
134 .into_iter()
135 .map(|retry| (retry.seq, retry))
136 .collect();
137 let mut latest: Option<Round<'a>> = None;
138 let mut out = Vec::new();
139 for (event, _, _) in context.messages() {
140 if let Some(round) = rounds.get(&event.seq) {
141 latest = Some(*round);
142 } else if let Some(retry) = retries.get(&event.seq) {
143 out.push((*retry, latest));
144 }
145 }
146 out
147}
148
149pub(in crate::checks) fn input_required_supported_methods(
151 context: &TraceContext<'_>,
152 sink: &mut FindingSink,
153) {
154 for round in rounds(context) {
155 sink.examined();
156 let (_, _, method) = round.origin;
157 if !SUPPORTED.contains(&method) {
158 sink.push(
159 Some(round.seq),
160 format!(
161 "`input_required` answers `{method}`; this revision permits it only on \
162 {}",
163 SUPPORTED.join(", ")
164 ),
165 );
166 }
167 }
168}
169
170pub(in crate::checks) fn input_request_methods(context: &TraceContext<'_>, sink: &mut FindingSink) {
172 for round in rounds(context) {
173 let Some(requests) = round.requests else {
174 continue;
175 };
176 for (key, request) in requests {
177 sink.examined();
178 match request.get("method").and_then(Value::as_str) {
179 Some(method) if INPUT_REQUEST_METHODS.contains(&method) => {}
180 Some(method) => sink.push(
181 Some(round.seq),
182 format!(
183 "`inputRequests[{key}]` asks for `{method}`, which is not one of \
184 ElicitRequest, CreateMessageRequest or ListRootsRequest"
185 ),
186 ),
187 None => sink.push(
188 Some(round.seq),
189 format!("`inputRequests[{key}]` is not a request object with a `method`"),
190 ),
191 }
192 }
193 }
194}
195
196pub(in crate::checks) fn input_required_has_content(
202 context: &TraceContext<'_>,
203 sink: &mut FindingSink,
204) {
205 for round in rounds(context) {
206 sink.examined();
207 if round.requests.is_none() && round.state.is_none() {
208 sink.push(
209 Some(round.seq),
210 "`input_required` carries neither `inputRequests` nor `requestState`, so the \
211 round it opens cannot be completed"
212 .to_owned(),
213 );
214 }
215 }
216}
217
218pub(in crate::checks) fn retry_carries_input_responses(
220 context: &TraceContext<'_>,
221 sink: &mut FindingSink,
222) {
223 for (retry, round) in retries_with_rounds(context) {
224 let Some(round) = round else {
225 continue;
226 };
227 if round.requests.is_none_or(Map::is_empty) {
230 continue;
231 }
232 sink.examined();
233 for key in missing_keys(&round, &retry) {
234 sink.push(
235 Some(retry.seq),
236 format!(
237 "the retry carries no `inputResponses[{key}]` for the input the \
238 `input_required` at seq {} asked for",
239 round.seq
240 ),
241 );
242 }
243 }
244}
245
246fn missing_keys(round: &Round<'_>, retry: &Retry<'_>) -> Vec<String> {
248 let Some(requests) = round.requests else {
249 return Vec::new();
250 };
251 requests
252 .keys()
253 .filter(|key| {
254 !retry
255 .responses
256 .is_some_and(|responses| responses.contains_key(*key))
257 })
258 .cloned()
259 .collect()
260}
261
262pub(in crate::checks) fn request_state_echoed(context: &TraceContext<'_>, sink: &mut FindingSink) {
269 for (retry, round) in retries_with_rounds(context) {
270 let Some(round) = round else {
271 continue;
272 };
273 let Some(issued) = round.state else { continue };
274 sink.examined();
275 match retry.state {
276 Some(echoed) if echoed == issued => {}
277 Some(echoed) => sink.push(
278 Some(retry.seq),
279 format!(
280 "the retry echoes `requestState` {echoed} instead of the {issued} the \
281 `input_required` at seq {} issued",
282 round.seq
283 ),
284 ),
285 None => sink.push(
286 Some(retry.seq),
287 format!(
288 "the retry omits the `requestState` the `input_required` at seq {} \
289 issued, which it must echo back exactly",
290 round.seq
291 ),
292 ),
293 }
294 }
295}
296
297pub(in crate::checks) fn no_unsolicited_request_state(
299 context: &TraceContext<'_>,
300 sink: &mut FindingSink,
301) {
302 for (retry, round) in retries_with_rounds(context) {
303 if retry.state.is_none() {
304 continue;
305 }
306 sink.examined();
307 let issued = round.and_then(|round| round.state);
308 if issued.is_none() {
309 sink.push(
310 Some(retry.seq),
311 "the request carries a `requestState` that no `input_required` before it \
312 issued"
313 .to_owned(),
314 );
315 }
316 }
317}
318
319pub(in crate::checks) fn retry_id_differs(context: &TraceContext<'_>, sink: &mut FindingSink) {
321 for (retry, round) in retries_with_rounds(context) {
322 let Some(round) = round else {
323 continue;
324 };
325 sink.examined();
326 let (origin_seq, origin_id, _) = round.origin;
327 if retry.id == origin_id {
328 sink.push(
329 Some(retry.seq),
330 format!(
331 "the retry reuses id {origin_id} from the request at seq {origin_seq}; \
332 the two are independent requests and must not share one"
333 ),
334 );
335 }
336 }
337}
338
339pub(in crate::checks) fn request_state_scoped_to_retry(
347 context: &TraceContext<'_>,
348 sink: &mut FindingSink,
349) {
350 let issued: BTreeMap<String, &str> = rounds(context)
352 .iter()
353 .filter_map(|round| round.state.map(|state| (state.to_string(), round.origin.2)))
354 .collect();
355 for retry in retries(context) {
356 let Some(state) = retry.state else { continue };
357 let Some(&origin_method) = issued.get(&state.to_string()) else {
358 continue;
359 };
360 sink.examined();
361 if retry.method != origin_method {
362 sink.push(
363 Some(retry.seq),
364 format!(
365 "`{}` carries the `requestState` issued for a `{origin_method}` request; \
366 it affects only that request's retry",
367 retry.method
368 ),
369 );
370 }
371 }
372}
373
374pub(in crate::checks) fn missing_input_reasked(context: &TraceContext<'_>, sink: &mut FindingSink) {
384 let paired: BTreeMap<u64, (Retry<'_>, Option<Round<'_>>)> = retries_with_rounds(context)
385 .into_iter()
386 .map(|(retry, round)| (retry.seq, (retry, round)))
387 .collect();
388 for exchange in context.exchanges() {
389 let Some((retry, Some(round))) = paired.get(&exchange.request.seq).copied() else {
390 continue;
391 };
392 let missing = missing_keys(&round, &retry);
393 if missing.is_empty() {
394 continue; }
396 sink.examined();
399 if exchange.result.is_some() {
400 continue;
401 }
402 sink.push(
403 Some(exchange.response.seq),
404 format!(
405 "the retry omitted {} that the `input_required` at seq {} asked for, and the \
406 server answered with an error rather than asking again",
407 missing
408 .iter()
409 .map(|key| format!("`{key}`"))
410 .collect::<Vec<_>>()
411 .join(", "),
412 round.seq
413 ),
414 );
415 }
416}