use mcp_conformance_core::message::{MessageKind, classify};
use mcp_conformance_core::trace::{Direction, TraceEvent};
use serde_json::Value;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum DraftPhase {
Active,
AwaitingDiscoverResult,
}
#[derive(Debug, Clone, Copy, Default)]
#[non_exhaustive]
pub struct DiscoverExchange<'a> {
pub request: Option<(u64, Option<&'a Value>)>,
pub result: Option<(u64, &'a Value)>,
pub error: Option<u64>,
}
#[derive(Debug)]
pub struct DraftLifecycle<'a> {
phases: Vec<DraftPhase>,
discover: DiscoverExchange<'a>,
final_phase: DraftPhase,
}
impl<'a> DraftLifecycle<'a> {
#[must_use]
pub fn new(events: &'a [TraceEvent]) -> Self {
let mut phases = Vec::with_capacity(events.len());
let mut tracker = DraftTracker::start();
for event in events {
phases.push(tracker.phase);
if let Some(kind) = event.message_payload().map(classify) {
tracker.step(event, &kind);
}
}
Self {
phases,
discover: tracker.discover,
final_phase: tracker.phase,
}
}
#[must_use]
pub fn phases(&self) -> &[DraftPhase] {
&self.phases
}
#[must_use]
pub const fn final_phase(&self) -> DraftPhase {
self.final_phase
}
#[must_use]
pub const fn discover(&self) -> &DiscoverExchange<'a> {
&self.discover
}
#[must_use]
pub fn server_capabilities(&self) -> Option<&'a Value> {
self.discover
.result
.and_then(|(_, result)| result.get("capabilities"))
}
}
struct DraftTracker<'a> {
phase: DraftPhase,
discover: DiscoverExchange<'a>,
discover_id: Option<&'a Value>,
}
impl<'a> DraftTracker<'a> {
const fn start() -> Self {
Self {
phase: DraftPhase::Active,
discover: DiscoverExchange {
request: None,
result: None,
error: None,
},
discover_id: None,
}
}
fn step(&mut self, event: &'a TraceEvent, kind: &MessageKind<'a>) {
match (self.phase, event.direction, kind) {
(
DraftPhase::Active,
Direction::ClientToServer,
MessageKind::Request { method, id },
) if *method == "server/discover" && self.discover.request.is_none() => {
self.discover_id = Some(id);
self.discover.request = Some((
event.seq,
event
.message_payload()
.and_then(|payload| payload.get("params")),
));
self.phase = DraftPhase::AwaitingDiscoverResult;
}
(
DraftPhase::AwaitingDiscoverResult,
Direction::ServerToClient,
MessageKind::Result { id: Some(id) },
) if Some(*id) == self.discover_id => {
self.discover.result = event
.message_payload()
.and_then(|payload| payload.get("result"))
.map(|result| (event.seq, result));
self.phase = DraftPhase::Active;
}
(
DraftPhase::AwaitingDiscoverResult,
Direction::ServerToClient,
MessageKind::Error { id: Some(id), .. },
) if Some(*id) == self.discover_id => {
self.discover.error = Some(event.seq);
self.phase = DraftPhase::Active;
}
_ => {}
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
use super::*;
use crate::reader::{Limits, parse_trace};
fn events(doc: &str) -> Vec<TraceEvent> {
parse_trace(doc, &Limits::default()).unwrap()
}
#[test]
fn operational_from_the_first_message_without_a_handshake() {
let trace = events(
r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"tools/list"}}
{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"result":{"tools":[]}}}"#,
);
let lifecycle = DraftLifecycle::new(&trace);
assert_eq!(lifecycle.phases(), [DraftPhase::Active, DraftPhase::Active]);
assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
assert!(lifecycle.discover().request.is_none());
assert_eq!(lifecycle.server_capabilities(), None);
}
#[test]
fn discover_request_then_result_records_capabilities_and_returns_to_active() {
let trace = events(
r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"server/discover","params":{"x":1}}}
{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"result":{"capabilities":{"tools":{}},"serverInfo":{"name":"s","version":"0"}}}}"#,
);
let lifecycle = DraftLifecycle::new(&trace);
assert_eq!(
lifecycle.phases(),
[DraftPhase::Active, DraftPhase::AwaitingDiscoverResult]
);
assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
assert_eq!(lifecycle.discover().request.unwrap().0, 0);
assert!(lifecycle.discover().request.unwrap().1.is_some());
assert_eq!(lifecycle.discover().result.unwrap().0, 1);
assert!(lifecycle.discover().error.is_none());
assert_eq!(
lifecycle.server_capabilities(),
Some(&serde_json::json!({"tools": {}}))
);
}
#[test]
fn discover_error_is_an_error_edge_back_to_active() {
let trace = events(
r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":7,"method":"server/discover"}}
{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":7,"error":{"code":-32601,"message":"no discover"}}}"#,
);
let lifecycle = DraftLifecycle::new(&trace);
assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
assert_eq!(lifecycle.discover().error, Some(1));
assert!(lifecycle.discover().result.is_none());
assert_eq!(lifecycle.server_capabilities(), None);
}
#[test]
fn a_response_with_an_unrelated_id_does_not_complete_discovery() {
for body in [
r#"{"jsonrpc":"2.0","id":99,"result":{}}"#,
r#"{"jsonrpc":"2.0","id":99,"error":{"code":-32600,"message":"x"}}"#,
] {
let response = format!(
r#"{{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{body}}}"#
);
let doc = format!(
"{}\n{response}",
r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"server/discover"}}"#,
);
let trace = events(&doc);
let lifecycle = DraftLifecycle::new(&trace);
assert_eq!(
lifecycle.final_phase(),
DraftPhase::AwaitingDiscoverResult,
"{body}"
);
assert!(lifecycle.discover().result.is_none(), "{body}");
assert!(lifecycle.discover().error.is_none(), "{body}");
}
}
#[test]
fn removed_handshake_methods_are_not_lifecycle_transitions() {
let trace = events(
r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}}
{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"result":{}}}
{"seq":2,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","method":"notifications/initialized"}}"#,
);
let lifecycle = DraftLifecycle::new(&trace);
assert!(lifecycle.phases().iter().all(|p| *p == DraftPhase::Active));
assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
assert!(lifecycle.discover().request.is_none());
}
#[test]
fn empty_trace_is_active_with_no_discovery() {
let lifecycle = DraftLifecycle::new(&[]);
assert!(lifecycle.phases().is_empty());
assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
assert!(lifecycle.discover().request.is_none());
assert_eq!(lifecycle.server_capabilities(), None);
}
#[test]
fn discovery_is_one_shot_a_second_request_while_active_is_ignored() {
let trace = events(
r#"{"seq":0,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"method":"server/discover"}}
{"seq":1,"direction":"server-to-client","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":1,"result":{"capabilities":{}}}}
{"seq":2,"direction":"client-to-server","transport":"stdio","kind":"message","payload":{"jsonrpc":"2.0","id":2,"method":"server/discover"}}"#,
);
let lifecycle = DraftLifecycle::new(&trace);
assert_eq!(lifecycle.final_phase(), DraftPhase::Active);
assert_eq!(lifecycle.discover().request.unwrap().0, 0);
assert_eq!(lifecycle.discover().result.unwrap().0, 1);
}
mod properties {
use super::*;
use proptest::prelude::*;
use serde_json::json;
fn arbitrary_event(seq: u64, choice: u8, direction_bit: bool) -> TraceEvent {
let payload = match choice % 6 {
0 => json!({"jsonrpc":"2.0","id":1,"method":"server/discover","params":{}}),
1 => json!({"jsonrpc":"2.0","id":1,"result":{"capabilities":{}}}),
2 => json!({"jsonrpc":"2.0","id":1,"error":{"code":-32601,"message":"x"}}),
3 => json!({"jsonrpc":"2.0","id":99,"result":{}}),
4 => json!({"jsonrpc":"2.0","id":2,"method":"tools/list"}),
_ => json!({"jsonrpc":"2.0","method":"notifications/cancelled"}),
};
let direction = if direction_bit {
"client-to-server"
} else {
"server-to-client"
};
serde_json::from_value(json!({
"seq": seq,
"direction": direction,
"transport": "stdio",
"kind": "message",
"payload": payload,
}))
.unwrap()
}
proptest! {
#[test]
fn invariants_hold_for_arbitrary_sequences(
moves in proptest::collection::vec((any::<u8>(), any::<bool>()), 0..32)
) {
let events: Vec<TraceEvent> = moves
.iter()
.enumerate()
.map(|(index, (choice, direction))| {
arbitrary_event(index as u64, *choice, *direction)
})
.collect();
let lifecycle = DraftLifecycle::new(&events);
prop_assert_eq!(lifecycle.phases().len(), events.len());
if let Some(first) = lifecycle.phases().first() {
prop_assert_eq!(*first, DraftPhase::Active);
}
let discover = lifecycle.discover();
let outstanding =
discover.request.is_some() && discover.result.is_none() && discover.error.is_none();
prop_assert_eq!(lifecycle.final_phase() == DraftPhase::AwaitingDiscoverResult, outstanding);
if discover.result.is_some() || discover.error.is_some() {
prop_assert!(discover.request.is_some());
}
prop_assert!(!(discover.result.is_some() && discover.error.is_some()));
for (index, pair) in lifecycle.phases().windows(2).enumerate() {
if pair[0] == DraftPhase::Active && pair[1] == DraftPhase::AwaitingDiscoverResult {
let event = &events[index];
prop_assert_eq!(event.direction, Direction::ClientToServer);
let kind = event.message_payload().map(classify);
let is_discover_request = matches!(
kind,
Some(MessageKind::Request { method, .. }) if method == "server/discover"
);
prop_assert!(is_discover_request);
}
}
}
}
}
}