1use flate2::Compression;
44use flate2::write::GzEncoder;
45use std::collections::HashMap;
46use std::fmt;
47use std::sync::Arc;
48
49#[cfg(feature = "messagepack")]
50use bytes::BufMut;
51use bytes::{Bytes, BytesMut};
52use serde::{Deserialize, Serialize};
53use uuid::Uuid;
54
55#[cfg(feature = "messagepack")]
56use msgpacker::Packable;
57
58use crate::types::{ContentType, ProtocolVersion, Timestamp};
59use crate::{McpError as Error, Result};
60
61#[cfg(feature = "messagepack")]
63#[derive(Debug, Clone)]
64pub enum JsonValue {
65 Null,
67 Bool(bool),
69 Number(f64),
71 String(String),
73 Array(Vec<JsonValue>),
75 Object(std::collections::HashMap<String, JsonValue>),
77}
78
79#[cfg(feature = "messagepack")]
80impl JsonValue {
81 pub fn from_serde_json(value: &serde_json::Value) -> Self {
83 match value {
84 serde_json::Value::Null => JsonValue::Null,
85 serde_json::Value::Bool(b) => JsonValue::Bool(*b),
86 serde_json::Value::Number(n) => {
87 if let Some(i) = n.as_i64() {
88 JsonValue::Number(i as f64)
89 } else if let Some(u) = n.as_u64() {
90 JsonValue::Number(u as f64)
91 } else if let Some(f) = n.as_f64() {
92 JsonValue::Number(f)
93 } else {
94 JsonValue::Null
95 }
96 }
97 serde_json::Value::String(s) => JsonValue::String(s.clone()),
98 serde_json::Value::Array(arr) => {
99 JsonValue::Array(arr.iter().map(Self::from_serde_json).collect())
100 }
101 serde_json::Value::Object(obj) => {
102 let mut map = std::collections::HashMap::new();
103 for (k, v) in obj {
104 map.insert(k.clone(), Self::from_serde_json(v));
105 }
106 JsonValue::Object(map)
107 }
108 }
109 }
110}
111
112#[cfg(feature = "messagepack")]
113impl msgpacker::Packable for JsonValue {
114 fn pack<T>(&self, buf: &mut T) -> usize
115 where
116 T: BufMut,
117 {
118 match self {
119 JsonValue::Null => {
120 buf.put_u8(0xc0);
122 1
123 }
124 JsonValue::Bool(b) => b.pack(buf),
125 JsonValue::Number(n) => n.pack(buf),
126 JsonValue::String(s) => s.pack(buf),
127 JsonValue::Array(arr) => {
128 let len = arr.len();
130 let mut bytes_written = 0;
131
132 if len <= 15 {
134 buf.put_u8(0x90 + len as u8);
135 bytes_written += 1;
136 } else if len <= u16::MAX as usize {
137 buf.put_u8(0xdc);
138 buf.put_u16(len as u16);
139 bytes_written += 3;
140 } else {
141 buf.put_u8(0xdd);
142 buf.put_u32(len as u32);
143 bytes_written += 5;
144 }
145
146 for item in arr {
148 bytes_written += item.pack(buf);
149 }
150
151 bytes_written
152 }
153 JsonValue::Object(obj) => {
154 let len = obj.len();
156 let mut bytes_written = 0;
157
158 if len <= 15 {
160 buf.put_u8(0x80 + len as u8);
161 bytes_written += 1;
162 } else if len <= u16::MAX as usize {
163 buf.put_u8(0xde);
164 buf.put_u16(len as u16);
165 bytes_written += 3;
166 } else {
167 buf.put_u8(0xdf);
168 buf.put_u32(len as u32);
169 bytes_written += 5;
170 }
171
172 for (k, v) in obj {
174 bytes_written += k.pack(buf);
175 bytes_written += v.pack(buf);
176 }
177
178 bytes_written
179 }
180 }
181 }
182}
183
184#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
186#[serde(untagged)]
187pub enum MessageId {
188 String(String),
190 Number(i64),
192 Uuid(Uuid),
194}
195
196#[derive(Debug, Clone, Serialize, Deserialize)]
198pub struct MessageMetadata {
199 pub created_at: Timestamp,
201
202 pub protocol_version: ProtocolVersion,
204
205 pub encoding: Option<String>,
207
208 pub content_type: ContentType,
210
211 pub size: usize,
213
214 pub correlation_id: Option<String>,
216
217 pub headers: HashMap<String, String>,
219}
220
221#[derive(Debug, Clone)]
223pub struct Message {
224 pub id: MessageId,
226
227 pub metadata: MessageMetadata,
229
230 pub payload: MessagePayload,
232}
233
234#[derive(Debug, Clone)]
236pub enum MessagePayload {
237 Json(JsonPayload),
239
240 Binary(BinaryPayload),
242
243 Text(String),
245
246 Empty,
248}
249
250#[derive(Debug, Clone)]
252pub struct JsonPayload {
253 pub raw: Bytes,
255
256 pub parsed: Option<Arc<serde_json::Value>>,
258
259 pub is_valid: bool,
261}
262
263#[derive(Debug, Clone)]
265pub struct BinaryPayload {
266 pub data: Bytes,
268
269 pub format: BinaryFormat,
271}
272
273#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
275#[serde(rename_all = "lowercase")]
276pub enum BinaryFormat {
277 MessagePack,
279
280 ProtoBuf,
282
283 Cbor,
285
286 Custom,
288}
289
290#[derive(Debug)]
292pub struct MessageSerializer {
293 default_format: SerializationFormat,
295
296 enable_compression: bool,
298
299 compression_threshold: usize,
301}
302
303#[derive(Debug, Clone, Copy, PartialEq, Eq)]
305pub enum SerializationFormat {
306 Json,
308
309 #[cfg(feature = "simd")]
311 SimdJson,
312
313 MessagePack,
315
316 Cbor,
318}
319
320impl Message {
321 pub fn json(id: MessageId, value: impl Serialize) -> Result<Self> {
327 let json_bytes = Self::serialize_json(&value)?;
328 let payload = MessagePayload::Json(JsonPayload {
329 raw: json_bytes.freeze(),
330 parsed: Some(Arc::new(serde_json::to_value(value)?)),
331 is_valid: true,
332 });
333
334 Ok(Self {
335 id,
336 metadata: MessageMetadata::new(ContentType::Json, payload.size()),
337 payload,
338 })
339 }
340
341 pub fn binary(id: MessageId, data: Bytes, format: BinaryFormat) -> Self {
343 let size = data.len();
344 let payload = MessagePayload::Binary(BinaryPayload { data, format });
345
346 Self {
347 id,
348 metadata: MessageMetadata::new(ContentType::Binary, size),
349 payload,
350 }
351 }
352
353 #[must_use]
355 pub fn text(id: MessageId, text: String) -> Self {
356 let size = text.len();
357 let payload = MessagePayload::Text(text);
358
359 Self {
360 id,
361 metadata: MessageMetadata::new(ContentType::Text, size),
362 payload,
363 }
364 }
365
366 #[must_use]
368 pub fn empty(id: MessageId) -> Self {
369 Self {
370 id,
371 metadata: MessageMetadata::new(ContentType::Json, 0),
372 payload: MessagePayload::Empty,
373 }
374 }
375
376 pub const fn size(&self) -> usize {
378 self.metadata.size
379 }
380
381 pub const fn is_empty(&self) -> bool {
383 matches!(self.payload, MessagePayload::Empty)
384 }
385
386 pub fn serialize(&self, format: SerializationFormat) -> Result<Bytes> {
392 match format {
393 SerializationFormat::Json => self.serialize_json_format(),
394 #[cfg(feature = "simd")]
395 SerializationFormat::SimdJson => self.serialize_simd_json(),
396 SerializationFormat::MessagePack => self.serialize_messagepack(),
397 SerializationFormat::Cbor => self.serialize_cbor(),
398 }
399 }
400
401 pub fn deserialize(bytes: Bytes) -> Result<Self> {
407 let format = Self::detect_format(&bytes);
409 Self::deserialize_with_format(bytes, format)
410 }
411
412 pub fn deserialize_with_format(bytes: Bytes, format: SerializationFormat) -> Result<Self> {
414 match format {
415 SerializationFormat::Json => Ok(Self::deserialize_json(bytes)),
416 #[cfg(feature = "simd")]
417 SerializationFormat::SimdJson => Ok(Self::deserialize_simd_json(bytes)),
418 SerializationFormat::MessagePack => Ok(Self::deserialize_messagepack(bytes)),
419 SerializationFormat::Cbor => Self::deserialize_cbor(bytes),
420 }
421 }
422
423 pub fn parse_json<T>(&self) -> Result<T>
425 where
426 T: for<'de> Deserialize<'de>,
427 {
428 match &self.payload {
429 MessagePayload::Json(json_payload) => json_payload.parsed.as_ref().map_or_else(
430 || {
431 #[cfg(feature = "simd")]
432 {
433 let mut json_bytes = json_payload.raw.to_vec();
434 simd_json::from_slice(&mut json_bytes).map_err(|e| {
435 Error::serialization(format!("SIMD JSON parsing failed: {e}"))
436 })
437 }
438 #[cfg(not(feature = "simd"))]
439 {
440 serde_json::from_slice(&json_payload.raw).map_err(|e| {
441 Error::serialization(format!("JSON parsing failed: {}", e))
442 })
443 }
444 },
445 |parsed| {
446 serde_json::from_value((**parsed).clone())
447 .map_err(|e| Error::serialization(format!("JSON parsing failed: {e}")))
448 },
449 ),
450 _ => Err(Error::invalid_params("Message payload is not JSON")),
451 }
452 }
453
454 fn serialize_json(value: &impl Serialize) -> Result<BytesMut> {
457 #[cfg(feature = "simd")]
458 {
459 sonic_rs::to_vec(value)
460 .map(|v| BytesMut::from(v.as_slice()))
461 .map_err(|e| Error::serialization(format!("SIMD JSON serialization failed: {e}")))
462 }
463 #[cfg(not(feature = "simd"))]
464 {
465 serde_json::to_vec(value)
466 .map(|v| BytesMut::from(v.as_slice()))
467 .map_err(|e| Error::serialization(format!("JSON serialization failed: {}", e)))
468 }
469 }
470
471 fn serialize_json_format(&self) -> Result<Bytes> {
472 match &self.payload {
473 MessagePayload::Json(json_payload) => Ok(json_payload.raw.clone()),
474 MessagePayload::Text(text) => Ok(Bytes::from(text.clone())),
475 MessagePayload::Empty => Ok(Bytes::from_static(b"{}")),
476 MessagePayload::Binary(_) => Err(Error::invalid_params(
477 "Cannot serialize non-JSON payload as JSON",
478 )),
479 }
480 }
481
482 #[cfg(feature = "simd")]
483 fn serialize_simd_json(&self) -> Result<Bytes> {
484 match &self.payload {
485 MessagePayload::Json(json_payload) => {
486 if json_payload.is_valid {
487 Ok(json_payload.raw.clone())
488 } else {
489 Err(Error::serialization("Invalid JSON payload"))
490 }
491 }
492 _ => Err(Error::invalid_params(
493 "Cannot serialize non-JSON payload with SIMD JSON",
494 )),
495 }
496 }
497
498 fn serialize_messagepack(&self) -> Result<Bytes> {
499 #[cfg(feature = "messagepack")]
500 {
501 match &self.payload {
502 MessagePayload::Binary(binary) if binary.format == BinaryFormat::MessagePack => {
503 Ok(binary.data.clone())
504 }
505 MessagePayload::Json(json_payload) => json_payload.parsed.as_ref().map_or_else(
506 || {
507 Err(Error::serialization(
508 "Cannot serialize unparsed JSON to MessagePack",
509 ))
510 },
511 |parsed| {
512 let packable_value = JsonValue::from_serde_json(parsed.as_ref());
514 let mut buffer = Vec::new();
515 packable_value.pack(&mut buffer);
516 Ok(Bytes::from(buffer))
517 },
518 ),
519 _ => Err(Error::invalid_params(
520 "Cannot serialize payload as MessagePack",
521 )),
522 }
523 }
524 #[cfg(not(feature = "messagepack"))]
525 {
526 let _ = self; Err(Error::invalid_params(
528 "MessagePack serialization not available",
529 ))
530 }
531 }
532
533 fn serialize_cbor(&self) -> Result<Bytes> {
534 match &self.payload {
535 MessagePayload::Binary(binary) if binary.format == BinaryFormat::Cbor => {
536 Ok(binary.data.clone())
537 }
538 MessagePayload::Json(json_payload) => {
539 if let Some(parsed) = &json_payload.parsed {
540 {
541 let mut buffer = Vec::new();
542 ciborium::into_writer(parsed.as_ref(), &mut buffer)
543 .map(|_| Bytes::from(buffer))
544 .map_err(|e| {
545 Error::serialization(format!("CBOR serialization failed: {e}"))
546 })
547 }
548 } else {
549 #[cfg(feature = "simd")]
551 {
552 let mut json_bytes = json_payload.raw.to_vec();
553 let value: serde_json::Value = simd_json::from_slice(&mut json_bytes)
554 .map_err(|e| {
555 Error::serialization(format!(
556 "SIMD JSON parsing failed before CBOR: {e}"
557 ))
558 })?;
559 {
560 let mut buffer = Vec::new();
561 ciborium::into_writer(&value, &mut buffer)
562 .map(|_| Bytes::from(buffer))
563 .map_err(|e| {
564 Error::serialization(format!("CBOR serialization failed: {e}"))
565 })
566 }
567 }
568 #[cfg(not(feature = "simd"))]
569 {
570 let value: serde_json::Value = serde_json::from_slice(&json_payload.raw)
571 .map_err(|e| {
572 Error::serialization(format!(
573 "JSON parsing failed before CBOR: {}",
574 e
575 ))
576 })?;
577 let mut buf = Vec::new();
578 ciborium::ser::into_writer(&value, &mut buf).map_err(|e| {
579 Error::serialization(format!("CBOR serialization failed: {}", e))
580 })?;
581 Ok(Bytes::from(buf))
582 }
583 }
584 }
585 _ => Err(Error::invalid_params("Cannot serialize payload as CBOR")),
586 }
587 }
588
589 fn deserialize_json(bytes: Bytes) -> Self {
590 let is_valid = serde_json::from_slice::<serde_json::Value>(&bytes).is_ok();
592
593 let payload = MessagePayload::Json(JsonPayload {
594 raw: bytes,
595 parsed: None, is_valid,
597 });
598
599 Self {
600 id: MessageId::Uuid(Uuid::new_v4()),
601 metadata: MessageMetadata::new(ContentType::Json, payload.size()),
602 payload,
603 }
604 }
605
606 #[cfg(feature = "simd")]
607 fn deserialize_simd_json(bytes: Bytes) -> Self {
608 let mut json_bytes = bytes.to_vec();
609 let is_valid = simd_json::from_slice::<serde_json::Value>(&mut json_bytes).is_ok();
610
611 let payload = MessagePayload::Json(JsonPayload {
612 raw: bytes,
613 parsed: None,
614 is_valid,
615 });
616
617 Self {
618 id: MessageId::Uuid(Uuid::new_v4()),
619 metadata: MessageMetadata::new(ContentType::Json, payload.size()),
620 payload,
621 }
622 }
623
624 fn deserialize_messagepack(bytes: Bytes) -> Self {
625 let payload = MessagePayload::Binary(BinaryPayload {
626 data: bytes,
627 format: BinaryFormat::MessagePack,
628 });
629
630 Self {
631 id: MessageId::Uuid(Uuid::new_v4()),
632 metadata: MessageMetadata::new(ContentType::Binary, payload.size()),
633 payload,
634 }
635 }
636
637 fn deserialize_cbor(bytes: Bytes) -> Result<Self> {
638 if let Ok(value) = ciborium::from_reader::<serde_json::Value, _>(&bytes[..]) {
640 let raw = serde_json::to_vec(&value)
641 .map(Bytes::from)
642 .map_err(|e| Error::serialization(format!("JSON re-encode failed: {e}")))?;
643 let payload = MessagePayload::Json(JsonPayload {
644 raw,
645 parsed: Some(Arc::new(value)),
646 is_valid: true,
647 });
648 return Ok(Self {
649 id: MessageId::Uuid(Uuid::new_v4()),
650 metadata: MessageMetadata::new(ContentType::Json, payload.size()),
651 payload,
652 });
653 }
654
655 let payload = MessagePayload::Binary(BinaryPayload {
657 data: bytes,
658 format: BinaryFormat::Cbor,
659 });
660 Ok(Self {
661 id: MessageId::Uuid(Uuid::new_v4()),
662 metadata: MessageMetadata::new(ContentType::Binary, payload.size()),
663 payload,
664 })
665 }
666
667 fn detect_format(bytes: &[u8]) -> SerializationFormat {
668 if bytes.is_empty() {
669 return SerializationFormat::Json;
670 }
671
672 if matches!(bytes[0], b'{' | b'[') {
674 #[cfg(feature = "simd")]
675 {
676 return SerializationFormat::SimdJson;
677 }
678 #[cfg(not(feature = "simd"))]
679 {
680 return SerializationFormat::Json;
681 }
682 }
683
684 if bytes.len() >= 2 && (bytes[0] == 0x82 || bytes[0] == 0x83) {
686 return SerializationFormat::MessagePack;
687 }
688
689 #[cfg(feature = "simd")]
691 {
692 SerializationFormat::SimdJson
693 }
694 #[cfg(not(feature = "simd"))]
695 {
696 SerializationFormat::Json
697 }
698 }
699}
700
701impl MessagePayload {
702 pub const fn size(&self) -> usize {
704 match self {
705 Self::Json(json) => json.raw.len(),
706 Self::Binary(binary) => binary.data.len(),
707 Self::Text(text) => text.len(),
708 Self::Empty => 0,
709 }
710 }
711}
712
713impl MessageMetadata {
714 #[must_use]
716 pub fn new(content_type: ContentType, size: usize) -> Self {
717 Self {
718 created_at: Timestamp::now(),
719 protocol_version: ProtocolVersion::LATEST.clone(),
720 encoding: None,
721 content_type,
722 size,
723 correlation_id: None,
724 headers: HashMap::new(),
725 }
726 }
727
728 #[must_use]
730 pub fn with_header(mut self, key: String, value: String) -> Self {
731 self.headers.insert(key, value);
732 self
733 }
734
735 #[must_use]
737 pub fn with_correlation_id(mut self, correlation_id: String) -> Self {
738 self.correlation_id = Some(correlation_id);
739 self
740 }
741
742 #[must_use]
744 pub fn with_encoding(mut self, encoding: String) -> Self {
745 self.encoding = Some(encoding);
746 self
747 }
748}
749
750impl MessageSerializer {
751 #[must_use]
753 pub const fn new() -> Self {
754 Self {
755 default_format: SerializationFormat::Json,
756 enable_compression: false,
757 compression_threshold: 1024, }
759 }
760
761 #[must_use]
763 pub const fn with_format(mut self, format: SerializationFormat) -> Self {
764 self.default_format = format;
765 self
766 }
767
768 #[must_use]
770 pub const fn with_compression(mut self, enable: bool, threshold: usize) -> Self {
771 self.enable_compression = enable;
772 self.compression_threshold = threshold;
773 self
774 }
775
776 pub fn serialize(&self, message: &mut Message) -> Result<Bytes> {
778 let serialized = message.serialize(self.default_format)?;
779
780 if self.enable_compression && serialized.len() > self.compression_threshold {
782 message.metadata.encoding = Some("gzip".to_string()); Ok(self.compress(serialized))
784 } else {
785 Ok(serialized)
786 }
787 }
788
789 fn compress(&self, data: Bytes) -> Bytes {
792 let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
793 if let Err(e) = std::io::Write::write_all(&mut encoder, &data) {
794 tracing::warn!(error = %e, "Failed to compress message; falling back to original payload");
795 return data; }
797 match encoder.finish() {
798 Ok(compressed_data) => Bytes::from(compressed_data),
799 Err(e) => {
800 tracing::warn!(error = %e, "Failed to finish compression; falling back to original payload");
801 data }
803 }
804 }
805}
806
807impl Default for MessageSerializer {
808 fn default() -> Self {
809 Self::new()
810 }
811}
812
813impl fmt::Display for MessageId {
814 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
815 match self {
816 Self::String(s) => write!(f, "{s}"),
817 Self::Number(n) => write!(f, "{n}"),
818 Self::Uuid(u) => write!(f, "{u}"),
819 }
820 }
821}
822
823impl From<String> for MessageId {
824 fn from(s: String) -> Self {
825 Self::String(s)
826 }
827}
828
829impl From<&str> for MessageId {
830 fn from(s: &str) -> Self {
831 Self::String(s.to_string())
832 }
833}
834
835impl From<i64> for MessageId {
836 fn from(n: i64) -> Self {
837 Self::Number(n)
838 }
839}
840
841impl From<Uuid> for MessageId {
842 fn from(u: Uuid) -> Self {
843 Self::Uuid(u)
844 }
845}
846
847#[cfg(test)]
848mod tests {
849 use super::*;
850 use serde_json::json;
851
852 #[test]
853 fn test_message_creation() {
854 let message = Message::json(MessageId::from("test"), json!({"key": "value"})).unwrap();
855 assert_eq!(message.id.to_string(), "test");
856 assert!(!message.is_empty());
857 }
858
859 #[test]
860 fn test_message_serialization() {
861 let message = Message::json(MessageId::from(1), json!({"test": true})).unwrap();
862 let serialized = message.serialize(SerializationFormat::Json).unwrap();
863 assert!(!serialized.is_empty());
864 }
865
866 #[derive(Deserialize, PartialEq, Debug)]
867 struct TestData {
868 number: i32,
869 }
870
871 #[test]
872 fn test_message_parsing() {
873 let message = Message::json(MessageId::from("test"), json!({"number": 42})).unwrap();
874
875 let parsed: TestData = message.parse_json().unwrap();
876 assert_eq!(parsed.number, 42);
877 }
878
879 #[test]
880 fn test_format_detection() {
881 let json_bytes = Bytes::from(r#"{"test": true}"#);
882 let format = Message::detect_format(&json_bytes);
883
884 #[cfg(feature = "simd")]
885 assert_eq!(format, SerializationFormat::SimdJson);
886 #[cfg(not(feature = "simd"))]
887 assert_eq!(format, SerializationFormat::Json);
888 }
889
890 #[test]
891 fn test_message_metadata() {
892 let metadata = MessageMetadata::new(ContentType::Json, 100)
893 .with_header("custom".to_string(), "value".to_string())
894 .with_correlation_id("corr-123".to_string());
895
896 assert_eq!(metadata.size, 100);
897 assert_eq!(metadata.headers.get("custom"), Some(&"value".to_string()));
898 assert_eq!(metadata.correlation_id, Some("corr-123".to_string()));
899 }
900
901 #[test]
902 fn test_message_serializer_compression() {
903 use flate2::read::GzDecoder;
904 use std::io::Read;
905
906 let serializer = MessageSerializer::new().with_compression(true, 10); let large_json = json!({
909 "data": "a".repeat(100), });
911 let mut message =
912 Message::json(MessageId::from("compressed_test"), large_json.clone()).unwrap();
913
914 let original_size = message.size();
915 assert!(
916 original_size > 10,
917 "Original message size should be greater than compression threshold"
918 );
919
920 let compressed_bytes = serializer.serialize(&mut message).unwrap();
921
922 assert_eq!(message.metadata.encoding, Some("gzip".to_string()));
924
925 assert!(
927 compressed_bytes.len() < original_size,
928 "Compressed size should be smaller than original"
929 );
930
931 let mut decoder = GzDecoder::new(&compressed_bytes[..]);
933 let mut decompressed_data = Vec::new();
934 decoder.read_to_end(&mut decompressed_data).unwrap();
935
936 let decompressed_message = Message::deserialize(Bytes::from(decompressed_data)).unwrap();
937 let parsed_json: serde_json::Value = decompressed_message.parse_json().unwrap();
938
939 assert_eq!(parsed_json, large_json);
940 }
941}