1use serde::{
4 Deserialize, Deserializer, Serialize,
5 de::{Error as DeError, MapAccess, Visitor},
6 ser::{SerializeMap, Serializer},
7};
8use serde_json::Value;
9use std::{collections::BTreeMap, error::Error, fmt, ops::Index};
10
11pub const V1_SCHEMA_VERSION: u32 = 1;
13
14pub const MAX_IDENTIFIER_LENGTH: usize = 256;
16
17pub const MAX_PROTOCOL_ITEMS: usize = 256;
19
20pub const MAX_EXTENSION_FIELDS: usize = MAX_PROTOCOL_ITEMS;
22pub const MAX_EXTENSION_VALUE_BYTES: usize = 64 * 1024;
24pub const MAX_EXTENSION_DEPTH: usize = 8;
26
27#[derive(Debug, Clone, Default, PartialEq, Eq)]
33pub struct ExtensionsV1(BTreeMap<String, Value>);
34
35impl ExtensionsV1 {
36 pub fn insert(
38 &mut self,
39 key: impl Into<String>,
40 value: Value,
41 ) -> Result<Option<Value>, ValidationError> {
42 let key = key.into();
43 validate_extension_key(&key)?;
44 validate_extension_value(&value, 0)?;
45 if !self.0.contains_key(&key) && self.0.len() >= MAX_EXTENSION_FIELDS {
46 return Err(ValidationError::new(format!(
47 "extensions exceeds the {MAX_EXTENSION_FIELDS} field limit"
48 )));
49 }
50 Ok(self.0.insert(key, value))
51 }
52
53 pub fn get(&self, key: &str) -> Option<&Value> {
55 self.0.get(key)
56 }
57
58 pub fn contains_key(&self, key: &str) -> bool {
60 self.0.contains_key(key)
61 }
62
63 pub fn iter(&self) -> impl Iterator<Item = (&String, &Value)> {
65 self.0.iter()
66 }
67
68 pub fn as_map(&self) -> &BTreeMap<String, Value> {
70 &self.0
71 }
72
73 pub fn into_inner(self) -> BTreeMap<String, Value> {
75 self.0
76 }
77}
78
79impl Index<&str> for ExtensionsV1 {
80 type Output = Value;
81
82 fn index(&self, key: &str) -> &Self::Output {
83 &self.0[key]
84 }
85}
86
87impl Serialize for ExtensionsV1 {
88 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
89 where
90 S: Serializer,
91 {
92 let mut map = serializer.serialize_map(Some(self.0.len()))?;
93 for (key, value) in &self.0 {
94 map.serialize_entry(key, value)?;
95 }
96 map.end()
97 }
98}
99
100struct ExtensionsVisitor;
101
102impl<'de> Visitor<'de> for ExtensionsVisitor {
103 type Value = ExtensionsV1;
104
105 fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
106 formatter.write_str("a bounded JSON object of additive fields")
107 }
108
109 fn visit_map<A>(self, mut map: A) -> Result<Self::Value, A::Error>
110 where
111 A: MapAccess<'de>,
112 {
113 let mut extensions = ExtensionsV1::default();
114 while let Some((key, value)) = map.next_entry::<String, Value>()? {
115 if extensions.contains_key(&key) {
116 return Err(A::Error::custom(format!("duplicate extension key {key:?}")));
117 }
118 extensions.insert(key, value).map_err(A::Error::custom)?;
119 }
120 Ok(extensions)
121 }
122}
123
124impl<'de> Deserialize<'de> for ExtensionsV1 {
125 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
126 where
127 D: Deserializer<'de>,
128 {
129 deserializer.deserialize_map(ExtensionsVisitor)
130 }
131}
132
133fn validate_extension_key(key: &str) -> Result<(), ValidationError> {
134 validate_bounded_opaque_identifier(key, "extension key")
135}
136
137impl ExtensionsV1 {
138 pub(crate) fn validate_reserved(&self, reserved: &[&str]) -> Result<(), ValidationError> {
140 if let Some(key) = self.0.keys().find(|key| reserved.contains(&key.as_str())) {
141 return Err(ValidationError::new(format!(
142 "extension key {key:?} collides with a reserved field"
143 )));
144 }
145 Ok(())
146 }
147}
148
149fn validate_extension_value(value: &Value, depth: usize) -> Result<(), ValidationError> {
150 if depth > MAX_EXTENSION_DEPTH {
151 return Err(ValidationError::new(format!(
152 "extension value exceeds nesting depth {MAX_EXTENSION_DEPTH}"
153 )));
154 }
155 match value {
156 Value::Array(values) => {
157 if values.len() > MAX_PROTOCOL_ITEMS {
158 return Err(ValidationError::new("extension array exceeds item limit"));
159 }
160 for value in values {
161 validate_extension_value(value, depth + 1)?;
162 }
163 }
164 Value::Object(values) => {
165 if values.len() > MAX_PROTOCOL_ITEMS {
166 return Err(ValidationError::new("extension object exceeds field limit"));
167 }
168 for (key, value) in values {
169 validate_bounded_opaque_identifier(key, "extension object key")?;
170 validate_extension_value(value, depth + 1)?;
171 }
172 }
173 Value::String(value) if value.len() > MAX_EXTENSION_VALUE_BYTES => {
174 return Err(ValidationError::new("extension string exceeds byte limit"));
175 }
176 Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {}
177 }
178 let encoded = serde_json::to_vec(value)
179 .map_err(|error| ValidationError::new(format!("invalid extension value: {error}")))?;
180 if encoded.len() > MAX_EXTENSION_VALUE_BYTES {
181 return Err(ValidationError::new(
182 "extension value exceeds serialized byte limit",
183 ));
184 }
185 Ok(())
186}
187
188pub(crate) fn validate_bounded_opaque_identifier(
190 value: &str,
191 type_name: &str,
192) -> Result<(), ValidationError> {
193 if value.trim().is_empty() {
194 return Err(ValidationError::new(format!(
195 "{type_name} must not be empty"
196 )));
197 }
198 if value.len() > MAX_IDENTIFIER_LENGTH {
199 return Err(ValidationError::new(format!(
200 "{type_name} exceeds the {MAX_IDENTIFIER_LENGTH} byte limit"
201 )));
202 }
203 if value.chars().any(char::is_control) {
204 return Err(ValidationError::new(format!(
205 "{type_name} must not contain control characters"
206 )));
207 }
208 Ok(())
209}
210
211#[derive(Debug, Clone, PartialEq, Eq)]
213pub struct ValidationError(pub(crate) String);
214
215impl ValidationError {
216 pub(crate) fn new(message: impl Into<String>) -> Self {
217 Self(message.into())
218 }
219}
220
221impl fmt::Display for ValidationError {
222 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
223 formatter.write_str(&self.0)
224 }
225}
226
227impl Error for ValidationError {}
228
229macro_rules! bounded_identifier {
230 ($name:ident) => {
231 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize)]
233 #[serde(transparent)]
234 pub struct $name(String);
235
236 impl $name {
237 pub fn new(value: impl Into<String>) -> Result<Self, ValidationError> {
239 let value = value.into();
240 $crate::common::validate_bounded_opaque_identifier(&value, stringify!($name))?;
241 Ok(Self(value))
242 }
243
244 pub fn as_str(&self) -> &str {
246 &self.0
247 }
248
249 pub fn into_inner(self) -> String {
251 self.0
252 }
253 }
254
255 impl AsRef<str> for $name {
256 fn as_ref(&self) -> &str {
257 self.as_str()
258 }
259 }
260
261 impl std::str::FromStr for $name {
262 type Err = ValidationError;
263
264 fn from_str(value: &str) -> Result<Self, Self::Err> {
265 Self::new(value)
266 }
267 }
268
269 impl TryFrom<String> for $name {
270 type Error = ValidationError;
271
272 fn try_from(value: String) -> Result<Self, Self::Error> {
273 Self::new(value)
274 }
275 }
276
277 impl TryFrom<&str> for $name {
278 type Error = ValidationError;
279
280 fn try_from(value: &str) -> Result<Self, Self::Error> {
281 Self::new(value)
282 }
283 }
284
285 impl From<$name> for String {
286 fn from(value: $name) -> Self {
287 value.0
288 }
289 }
290
291 impl<'de> Deserialize<'de> for $name {
292 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
293 where
294 D: Deserializer<'de>,
295 {
296 let value = String::deserialize(deserializer)?;
297 Self::new(value).map_err(DeError::custom)
298 }
299 }
300 };
301}
302
303bounded_identifier!(TaskId);
304bounded_identifier!(TraceId);
305bounded_identifier!(ProjectId);
306bounded_identifier!(SessionId);
307bounded_identifier!(AgentId);
308bounded_identifier!(TenantId);
309bounded_identifier!(PlanId);
310bounded_identifier!(ReceiptId);
311bounded_identifier!(OutcomeId);
312bounded_identifier!(CapabilityId);
313bounded_identifier!(DecisionId);
314bounded_identifier!(ContextPlanId);
315bounded_identifier!(AttemptId);
316
317pub fn validate_bounded_string(value: &str, field: &str) -> Result<(), ValidationError> {
319 validate_bounded_opaque_identifier(value, field)
320}
321
322pub fn validate_unique_strings<T>(values: &[T], field: &str) -> Result<(), ValidationError>
324where
325 T: AsRef<str>,
326{
327 if values.len() > MAX_PROTOCOL_ITEMS {
328 return Err(ValidationError::new(format!(
329 "{field} exceeds the {MAX_PROTOCOL_ITEMS} item limit"
330 )));
331 }
332 let mut seen = std::collections::BTreeSet::new();
333 for value in values {
334 validate_bounded_string(value.as_ref(), field)?;
335 if !seen.insert(value.as_ref()) {
336 return Err(ValidationError::new(format!(
337 "{field} contains duplicate values"
338 )));
339 }
340 }
341 Ok(())
342}
343
344pub fn deserialize_schema_version<'de, D>(deserializer: D) -> Result<u32, D::Error>
346where
347 D: Deserializer<'de>,
348{
349 let version = u32::deserialize(deserializer)?;
350 if version == V1_SCHEMA_VERSION {
351 Ok(version)
352 } else {
353 Err(DeError::custom(format!(
354 "unsupported schema_version {version}; expected {V1_SCHEMA_VERSION}"
355 )))
356 }
357}
358
359pub fn deserialize_milliunit<'de, D>(deserializer: D) -> Result<u16, D::Error>
361where
362 D: Deserializer<'de>,
363{
364 let value = u16::deserialize(deserializer)?;
365 if value <= 1000 {
366 Ok(value)
367 } else {
368 Err(DeError::custom("milliunit must be between 0 and 1000"))
369 }
370}
371
372pub fn deserialize_optional_milliunit<'de, D>(deserializer: D) -> Result<Option<u16>, D::Error>
374where
375 D: Deserializer<'de>,
376{
377 let value = Option::<u16>::deserialize(deserializer)?;
378 value.map_or(Ok(None), |value| {
379 if value <= 1000 {
380 Ok(Some(value))
381 } else {
382 Err(DeError::custom("milliunit must be between 0 and 1000"))
383 }
384 })
385}
386
387pub fn validate_schema_version(version: u32) -> Result<(), ValidationError> {
389 if version == V1_SCHEMA_VERSION {
390 Ok(())
391 } else {
392 Err(ValidationError::new(format!(
393 "unsupported schema_version {version}; expected {V1_SCHEMA_VERSION}"
394 )))
395 }
396}
397
398pub fn validate_milliunit(value: u16, field: &str) -> Result<(), ValidationError> {
400 if value <= 1000 {
401 Ok(())
402 } else {
403 Err(ValidationError::new(format!(
404 "{field} must be between 0 and 1000"
405 )))
406 }
407}
408
409#[cfg(test)]
410mod tests {
411 use super::*;
412 use serde_json::json;
413
414 #[test]
415 fn extensions_retain_unknown_fields_and_reject_reserved_names() {
416 let value = json!({
417 "future_flag": true,
418 "future_nested": {"stable": ["value", 7]}
419 });
420 let extensions: ExtensionsV1 =
421 serde_json::from_value(value.clone()).expect("extensions should deserialize");
422 assert_eq!(
423 serde_json::to_value(&extensions).expect("extensions should serialize"),
424 value
425 );
426 let mut direct = ExtensionsV1::default();
427 assert!(direct.insert("extensions", Value::from(true)).is_ok());
428 assert!(direct.insert("schema_version", Value::from(1)).is_ok());
429 assert!(direct.insert("future_flag", Value::from(true)).is_ok());
430 assert!(direct.validate_reserved(&["schema_version"]).is_err());
431 assert!(direct.validate_reserved(&["provider"]).is_ok());
432 }
433
434 #[test]
435 fn extensions_are_bounded_by_count_and_value_shape() {
436 let mut extensions = ExtensionsV1::default();
437 for index in 0..MAX_EXTENSION_FIELDS {
438 extensions
439 .insert(format!("future_{index}"), Value::from(index))
440 .expect("entry should fit");
441 }
442 assert!(
443 extensions
444 .insert("future_overflow", Value::from(true))
445 .is_err()
446 );
447 assert!(
448 extensions
449 .insert(
450 "future_large",
451 Value::String("x".repeat(MAX_EXTENSION_VALUE_BYTES + 1))
452 )
453 .is_err()
454 );
455 }
456
457 #[test]
458 fn extensions_reject_duplicate_wire_keys() {
459 let json = r#"{"future_key":1,"future_key":2}"#;
460 assert!(serde_json::from_str::<ExtensionsV1>(json).is_err());
461 }
462}