1use mcp_conformance_core::message::{MessageKind, classify};
29use mcp_conformance_core::trace::{Direction, TraceEvent};
30use serde_json::Value;
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38#[non_exhaustive]
39pub enum DraftPhase {
40 Active,
43 AwaitingDiscoverResult,
45}
46
47#[derive(Debug, Clone, Copy, Default)]
50#[non_exhaustive]
51pub struct DiscoverExchange<'a> {
52 pub request: Option<(u64, Option<&'a Value>)>,
54 pub result: Option<(u64, &'a Value)>,
56 pub error: Option<u64>,
58}
59
60#[derive(Debug)]
79pub struct DraftLifecycle<'a> {
80 phases: Vec<DraftPhase>,
81 discover: DiscoverExchange<'a>,
82 final_phase: DraftPhase,
83}
84
85impl<'a> DraftLifecycle<'a> {
86 #[must_use]
89 pub fn new(events: &'a [TraceEvent]) -> Self {
90 let mut phases = Vec::with_capacity(events.len());
91 let mut tracker = DraftTracker::start();
92 for event in events {
93 phases.push(tracker.phase);
94 if let Some(kind) = event.message_payload().map(classify) {
95 tracker.step(event, &kind);
96 }
97 }
98 Self {
99 phases,
100 discover: tracker.discover,
101 final_phase: tracker.phase,
102 }
103 }
104
105 #[must_use]
107 pub fn phases(&self) -> &[DraftPhase] {
108 &self.phases
109 }
110
111 #[must_use]
113 pub const fn final_phase(&self) -> DraftPhase {
114 self.final_phase
115 }
116
117 #[must_use]
119 pub const fn discover(&self) -> &DiscoverExchange<'a> {
120 &self.discover
121 }
122
123 #[must_use]
128 pub fn server_capabilities(&self) -> Option<&'a Value> {
129 self.discover
130 .result
131 .and_then(|(_, result)| result.get("capabilities"))
132 }
133}
134
135struct DraftTracker<'a> {
140 phase: DraftPhase,
141 discover: DiscoverExchange<'a>,
142 discover_id: Option<&'a Value>,
143}
144
145impl<'a> DraftTracker<'a> {
146 const fn start() -> Self {
147 Self {
148 phase: DraftPhase::Active,
149 discover: DiscoverExchange {
150 request: None,
151 result: None,
152 error: None,
153 },
154 discover_id: None,
155 }
156 }
157
158 fn step(&mut self, event: &'a TraceEvent, kind: &MessageKind<'a>) {
159 match (self.phase, event.direction, kind) {
160 (
161 DraftPhase::Active,
162 Direction::ClientToServer,
163 MessageKind::Request { method, id },
164 ) if *method == "server/discover" && self.discover.request.is_none() => {
165 self.discover_id = Some(id);
166 self.discover.request = Some((
167 event.seq,
168 event
169 .message_payload()
170 .and_then(|payload| payload.get("params")),
171 ));
172 self.phase = DraftPhase::AwaitingDiscoverResult;
173 }
174 (
175 DraftPhase::AwaitingDiscoverResult,
176 Direction::ServerToClient,
177 MessageKind::Result { id: Some(id) },
178 ) if Some(*id) == self.discover_id => {
179 self.discover.result = event
180 .message_payload()
181 .and_then(|payload| payload.get("result"))
182 .map(|result| (event.seq, result));
183 self.phase = DraftPhase::Active;
184 }
185 (
186 DraftPhase::AwaitingDiscoverResult,
187 Direction::ServerToClient,
188 MessageKind::Error { id: Some(id), .. },
189 ) if Some(*id) == self.discover_id => {
190 self.discover.error = Some(event.seq);
191 self.phase = DraftPhase::Active;
192 }
193 _ => {}
194 }
195 }
196}
197
198#[cfg(test)]
199#[allow(clippy::unwrap_used)]
200mod tests {
201 use super::*;
202 use crate::reader::{Limits, parse_trace};
203
204 fn events(doc: &str) -> Vec<TraceEvent> {
205 parse_trace(doc, &Limits::default()).unwrap()
206 }
207
208 #[test]
209 fn operational_from_the_first_message_without_a_handshake() {
210 let trace = events(
214 r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"tools/list"}}
215{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}}"#,
216 );
217 let lifecycle = DraftLifecycle::new(&trace);
218 assert_eq!(lifecycle.phases(), [DraftPhase::Active, DraftPhase::Active]);
219 assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
220 assert!(lifecycle.discover().request.is_none());
221 assert_eq!(lifecycle.server_capabilities(), None);
222 }
223
224 #[test]
225 fn discover_request_then_result_records_capabilities_and_returns_to_active() {
226 let trace = events(
227 r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"server/discover","params":{"x":1}}}
228{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"result":{"capabilities":{"tools":{}},"serverInfo":{"name":"s","version":"0"}}}}"#,
229 );
230 let lifecycle = DraftLifecycle::new(&trace);
231 assert_eq!(
233 lifecycle.phases(),
234 [DraftPhase::Active, DraftPhase::AwaitingDiscoverResult]
235 );
236 assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
238 assert_eq!(lifecycle.discover().request.unwrap().0, 0);
239 assert!(lifecycle.discover().request.unwrap().1.is_some());
240 assert_eq!(lifecycle.discover().result.unwrap().0, 1);
241 assert!(lifecycle.discover().error.is_none());
242 assert_eq!(
243 lifecycle.server_capabilities(),
244 Some(&serde_json::json!({"tools": {}}))
245 );
246 }
247
248 #[test]
249 fn discover_error_is_an_error_edge_back_to_active() {
250 let trace = events(
251 r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":7,"method":"server/discover"}}
252{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":7,"error":{"code":-32601,"message":"no discover"}}}"#,
253 );
254 let lifecycle = DraftLifecycle::new(&trace);
255 assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
256 assert_eq!(lifecycle.discover().error, Some(1));
257 assert!(lifecycle.discover().result.is_none());
258 assert_eq!(lifecycle.server_capabilities(), None);
259 }
260
261 #[test]
262 fn a_response_with_an_unrelated_id_does_not_complete_discovery() {
263 for body in [
266 r#"{"jsonrpc":"2.0","id":99,"result":{}}"#,
267 r#"{"jsonrpc":"2.0","id":99,"error":{"code":-32600,"message":"x"}}"#,
268 ] {
269 let response = format!(
270 r#"{{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{body}}}"#
271 );
272 let doc = format!(
273 "{}\n{response}",
274 r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"server/discover"}}"#,
275 );
276 let trace = events(&doc);
277 let lifecycle = DraftLifecycle::new(&trace);
278 assert_eq!(
279 lifecycle.final_phase(),
280 DraftPhase::AwaitingDiscoverResult,
281 "{body}"
282 );
283 assert!(lifecycle.discover().result.is_none(), "{body}");
284 assert!(lifecycle.discover().error.is_none(), "{body}");
285 }
286 }
287
288 #[test]
289 fn removed_handshake_methods_are_not_lifecycle_transitions() {
290 let trace = events(
294 r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}}
295{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"result":{}}}
296{"seq":2,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","method":"notifications/initialized"}}"#,
297 );
298 let lifecycle = DraftLifecycle::new(&trace);
299 assert!(lifecycle.phases().iter().all(|p| *p == DraftPhase::Active));
300 assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
301 assert!(lifecycle.discover().request.is_none());
302 }
303
304 #[test]
305 fn empty_trace_is_active_with_no_discovery() {
306 let lifecycle = DraftLifecycle::new(&[]);
307 assert!(lifecycle.phases().is_empty());
308 assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
309 assert!(lifecycle.discover().request.is_none());
310 assert_eq!(lifecycle.server_capabilities(), None);
311 }
312
313 #[test]
314 fn discovery_is_one_shot_a_second_request_while_active_is_ignored() {
315 let trace = events(
318 r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"server/discover"}}
319{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"result":{"capabilities":{}}}}
320{"seq":2,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":2,"method":"server/discover"}}"#,
321 );
322 let lifecycle = DraftLifecycle::new(&trace);
323 assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
325 assert_eq!(lifecycle.discover().request.unwrap().0, 0);
326 assert_eq!(lifecycle.discover().result.unwrap().0, 1);
327 }
328
329 mod properties {
332 use super::*;
333 use proptest::prelude::*;
334 use serde_json::json;
335
336 fn arbitrary_event(seq: u64, choice: u8, direction_bit: bool) -> TraceEvent {
337 let payload = match choice % 6 {
338 0 => json!({"jsonrpc":"2.0","id":1,"method":"server/discover","params":{}}),
339 1 => json!({"jsonrpc":"2.0","id":1,"result":{"capabilities":{}}}),
340 2 => json!({"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"x"}}),
341 3 => json!({"jsonrpc":"2.0","id":99,"result":{}}),
342 4 => json!({"jsonrpc":"2.0","id":2,"method":"tools/list"}),
343 _ => json!({"jsonrpc":"2.0","method":"notifications/cancelled"}),
344 };
345 let direction = if direction_bit {
346 "client-to-server"
347 } else {
348 "server-to-client"
349 };
350 serde_json::from_value(json!({
351 "seq": seq,
352 "direction": direction,
353 "transport": "stdio",
354 "kind": "message",
355 "payload": payload,
356 }))
357 .unwrap()
358 }
359
360 proptest! {
361 #[test]
362 fn invariants_hold_for_arbitrary_sequences(
363 moves in proptest::collection::vec((any::<u8>(), any::<bool>()), 0..32)
364 ) {
365 let events: Vec<TraceEvent> = moves
366 .iter()
367 .enumerate()
368 .map(|(index, (choice, direction))| {
369 arbitrary_event(index as u64, *choice, *direction)
370 })
371 .collect();
372 let lifecycle = DraftLifecycle::new(&events);
373
374 prop_assert_eq!(lifecycle.phases().len(), events.len());
376 if let Some(first) = lifecycle.phases().first() {
377 prop_assert_eq!(*first, DraftPhase::Active);
378 }
379
380 let discover = lifecycle.discover();
381 let outstanding =
383 discover.request.is_some() && discover.result.is_none() && discover.error.is_none();
384 prop_assert_eq!(lifecycle.final_phase() == DraftPhase::AwaitingDiscoverResult, outstanding);
385
386 if discover.result.is_some() || discover.error.is_some() {
388 prop_assert!(discover.request.is_some());
389 }
390 prop_assert!(!(discover.result.is_some() && discover.error.is_some()));
391
392 for (index, pair) in lifecycle.phases().windows(2).enumerate() {
395 if pair[0] == DraftPhase::Active && pair[1] == DraftPhase::AwaitingDiscoverResult {
396 let event = &events[index];
397 prop_assert_eq!(event.direction, Direction::ClientToServer);
398 let kind = event.message_payload().map(classify);
399 let is_discover_request = matches!(
400 kind,
401 Some(MessageKind::Request { method, .. }) if method == "server/discover"
402 );
403 prop_assert!(is_discover_request);
404 }
405 }
406 }
407 }
408 }
409}