Skip to content

Commit 297deb9

Browse files
author
nightcityblade
committed
fix(model): preserve elicitation property order metadata
1 parent b5cf34e commit 297deb9

2 files changed

Lines changed: 82 additions & 3 deletions

File tree

crates/rmcp/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ serde_json = "1.0"
5151
thiserror = "2"
5252
tokio = { version = "1", features = ["sync", "macros", "rt", "time"] }
5353
futures = "0.3"
54+
indexmap = { version = "2", features = ["serde"] }
5455
tracing = { version = "0.1" }
5556
tokio-util = { version = "0.7" }
5657
pin-project-lite = "0.2"

crates/rmcp/src/model/elicitation_schema.rs

Lines changed: 81 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@
1818
1919
use std::{borrow::Cow, collections::BTreeMap, marker::PhantomData};
2020

21-
use serde::{Deserialize, Serialize};
21+
use indexmap::IndexMap;
22+
use serde::{Deserialize, Deserializer, Serialize};
2223

2324
use crate::{const_string, model::ConstString};
2425

@@ -1109,9 +1110,10 @@ impl EnumSchema {
11091110
/// .optional_bool("newsletter", false)
11101111
/// .build();
11111112
/// ```
1112-
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
1113+
#[derive(Debug, Clone, PartialEq, Serialize)]
11131114
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
1114-
#[serde(rename_all = "camelCase")]
1115+
#[cfg_attr(feature = "schemars", schemars(!into))]
1116+
#[serde(rename_all = "camelCase", into = "ElicitationSchemaWire")]
11151117
#[non_exhaustive]
11161118
pub struct ElicitationSchema {
11171119
/// Always "object" for elicitation schemas
@@ -1125,6 +1127,11 @@ pub struct ElicitationSchema {
11251127
/// Property definitions (must be primitive types)
11261128
pub properties: BTreeMap<String, PrimitiveSchemaDefinition>,
11271129

1130+
/// Property names in wire order. Schemas constructed from a `BTreeMap`
1131+
/// use the map's sorted key order.
1132+
#[serde(skip)]
1133+
pub property_order: Option<Vec<String>>,
1134+
11281135
/// List of required property names
11291136
#[serde(skip_serializing_if = "Option::is_none")]
11301137
pub required: Option<Vec<String>>,
@@ -1134,13 +1141,75 @@ pub struct ElicitationSchema {
11341141
pub description: Option<Cow<'static, str>>,
11351142
}
11361143

1144+
#[derive(Deserialize, Serialize)]
1145+
#[serde(rename_all = "camelCase")]
1146+
struct ElicitationSchemaWire {
1147+
#[serde(rename = "type")]
1148+
type_: ObjectTypeConst,
1149+
#[serde(skip_serializing_if = "Option::is_none")]
1150+
title: Option<Cow<'static, str>>,
1151+
properties: IndexMap<String, PrimitiveSchemaDefinition>,
1152+
#[serde(skip_serializing_if = "Option::is_none")]
1153+
required: Option<Vec<String>>,
1154+
#[serde(skip_serializing_if = "Option::is_none")]
1155+
description: Option<Cow<'static, str>>,
1156+
}
1157+
1158+
impl From<ElicitationSchemaWire> for ElicitationSchema {
1159+
fn from(schema: ElicitationSchemaWire) -> Self {
1160+
Self {
1161+
type_: schema.type_,
1162+
title: schema.title,
1163+
property_order: Some(schema.properties.keys().cloned().collect()),
1164+
properties: schema.properties.into_iter().collect(),
1165+
required: schema.required,
1166+
description: schema.description,
1167+
}
1168+
}
1169+
}
1170+
1171+
impl From<ElicitationSchema> for ElicitationSchemaWire {
1172+
fn from(schema: ElicitationSchema) -> Self {
1173+
let mut remaining = schema.properties;
1174+
let mut properties = IndexMap::with_capacity(remaining.len());
1175+
1176+
if let Some(property_order) = schema.property_order {
1177+
for name in property_order {
1178+
if let Some(definition) = remaining.remove(&name) {
1179+
properties.insert(name, definition);
1180+
}
1181+
}
1182+
}
1183+
properties.extend(remaining);
1184+
1185+
Self {
1186+
type_: schema.type_,
1187+
title: schema.title,
1188+
properties,
1189+
required: schema.required,
1190+
description: schema.description,
1191+
}
1192+
}
1193+
}
1194+
1195+
impl<'de> Deserialize<'de> for ElicitationSchema {
1196+
fn deserialize<__D>(__deserializer: __D) -> Result<Self, __D::Error>
1197+
where
1198+
__D: Deserializer<'de>,
1199+
{
1200+
ElicitationSchemaWire::deserialize(__deserializer).map(Into::into)
1201+
}
1202+
}
1203+
11371204
impl ElicitationSchema {
11381205
/// Create a new elicitation schema with the given properties
11391206
pub fn new(properties: BTreeMap<String, PrimitiveSchemaDefinition>) -> Self {
1207+
let property_order = Some(properties.keys().cloned().collect());
11401208
Self {
11411209
type_: ObjectTypeConst,
11421210
title: None,
11431211
properties,
1212+
property_order,
11441213
required: None,
11451214
description: None,
11461215
}
@@ -1632,10 +1701,12 @@ impl ElicitationSchemaBuilder {
16321701
}
16331702
}
16341703

1704+
let property_order = Some(self.properties.keys().cloned().collect());
16351705
Ok(ElicitationSchema {
16361706
type_: ObjectTypeConst,
16371707
title: self.title,
16381708
properties: self.properties,
1709+
property_order,
16391710
required: if self.required.is_empty() {
16401711
None
16411712
} else {
@@ -1821,6 +1892,13 @@ mod tests {
18211892
output["properties"]["choice"]["enumNames"],
18221893
serde_json::json!(["Option One", "Option Two", "Option Three"]),
18231894
);
1895+
let input = r#"{"type":"object","properties":{"firstName":{"type":"string"},"lastName":{"type":"string"},"email":{"type":"string"}}}"#;
1896+
let ordered: ElicitationSchema = serde_json::from_str(input)?;
1897+
assert_eq!(
1898+
ordered.property_order.as_ref().unwrap().join(","),
1899+
"firstName,lastName,email",
1900+
);
1901+
assert_eq!(serde_json::to_string(&ordered)?, input);
18241902
Ok(())
18251903
}
18261904

0 commit comments

Comments
 (0)