Skip to main content

rmcp/model/
annotated.rs

1use std::ops::{Deref, DerefMut};
2
3use chrono::{DateTime, Utc};
4use serde::{Deserialize, Serialize};
5
6use super::{
7    RawAudioContent, RawContent, RawEmbeddedResource, RawImageContent, RawResource,
8    RawResourceTemplate, RawTextContent, Role,
9};
10
11#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
12#[serde(rename_all = "camelCase")]
13#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
14#[non_exhaustive]
15pub struct Annotations {
16    #[serde(skip_serializing_if = "Option::is_none")]
17    pub audience: Option<Vec<Role>>,
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub priority: Option<f32>,
20    #[serde(skip_serializing_if = "Option::is_none", rename = "lastModified")]
21    pub last_modified: Option<DateTime<Utc>>,
22}
23
24impl Annotations {
25    /// Creates a new Annotations instance specifically for resources
26    /// optional priority, and a timestamp (defaults to now if None)
27    pub fn for_resource(priority: f32, timestamp: DateTime<Utc>) -> Self {
28        assert!(
29            (0.0..=1.0).contains(&priority),
30            "Priority {priority} must be between 0.0 and 1.0"
31        );
32        Annotations {
33            priority: Some(priority),
34            last_modified: Some(timestamp),
35            audience: None,
36        }
37    }
38}
39
40#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
41#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
42#[expect(clippy::exhaustive_structs, reason = "intentionally exhaustive")]
43pub struct Annotated<T: AnnotateAble> {
44    #[serde(flatten)]
45    pub raw: T,
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub annotations: Option<Annotations>,
48}
49
50impl<T: AnnotateAble> Deref for Annotated<T> {
51    type Target = T;
52    fn deref(&self) -> &Self::Target {
53        &self.raw
54    }
55}
56
57impl<T: AnnotateAble> DerefMut for Annotated<T> {
58    fn deref_mut(&mut self) -> &mut Self::Target {
59        &mut self.raw
60    }
61}
62
63impl<T: AnnotateAble> Annotated<T> {
64    pub fn new(raw: T, annotations: Option<Annotations>) -> Self {
65        Self { raw, annotations }
66    }
67    pub fn remove_annotation(&mut self) -> Option<Annotations> {
68        self.annotations.take()
69    }
70    pub fn audience(&self) -> Option<&Vec<Role>> {
71        self.annotations.as_ref().and_then(|a| a.audience.as_ref())
72    }
73    pub fn priority(&self) -> Option<f32> {
74        self.annotations.as_ref().and_then(|a| a.priority)
75    }
76    pub fn timestamp(&self) -> Option<DateTime<Utc>> {
77        self.annotations.as_ref().and_then(|a| a.last_modified)
78    }
79    pub fn with_audience(self, audience: Vec<Role>) -> Annotated<T>
80    where
81        Self: Sized,
82    {
83        if let Some(annotations) = self.annotations {
84            Annotated {
85                raw: self.raw,
86                annotations: Some(Annotations {
87                    audience: Some(audience),
88                    ..annotations
89                }),
90            }
91        } else {
92            Annotated {
93                raw: self.raw,
94                annotations: Some(Annotations {
95                    audience: Some(audience),
96                    priority: None,
97                    last_modified: None,
98                }),
99            }
100        }
101    }
102    pub fn with_priority(self, priority: f32) -> Annotated<T>
103    where
104        Self: Sized,
105    {
106        if let Some(annotations) = self.annotations {
107            Annotated {
108                raw: self.raw,
109                annotations: Some(Annotations {
110                    priority: Some(priority),
111                    ..annotations
112                }),
113            }
114        } else {
115            Annotated {
116                raw: self.raw,
117                annotations: Some(Annotations {
118                    priority: Some(priority),
119                    last_modified: None,
120                    audience: None,
121                }),
122            }
123        }
124    }
125    pub fn with_timestamp(self, timestamp: DateTime<Utc>) -> Annotated<T>
126    where
127        Self: Sized,
128    {
129        if let Some(annotations) = self.annotations {
130            Annotated {
131                raw: self.raw,
132                annotations: Some(Annotations {
133                    last_modified: Some(timestamp),
134                    ..annotations
135                }),
136            }
137        } else {
138            Annotated {
139                raw: self.raw,
140                annotations: Some(Annotations {
141                    last_modified: Some(timestamp),
142                    priority: None,
143                    audience: None,
144                }),
145            }
146        }
147    }
148    pub fn with_timestamp_now(self) -> Annotated<T>
149    where
150        Self: Sized,
151    {
152        self.with_timestamp(Utc::now())
153    }
154}
155
156mod sealed {
157    pub trait Sealed {}
158}
159macro_rules! annotate {
160    ($T: ident) => {
161        impl sealed::Sealed for $T {}
162        impl AnnotateAble for $T {}
163    };
164}
165
166annotate!(RawContent);
167annotate!(RawTextContent);
168annotate!(RawImageContent);
169annotate!(RawAudioContent);
170annotate!(RawEmbeddedResource);
171annotate!(RawResource);
172annotate!(RawResourceTemplate);
173
174pub trait AnnotateAble: sealed::Sealed {
175    fn optional_annotate(self, annotations: Option<Annotations>) -> Annotated<Self>
176    where
177        Self: Sized,
178    {
179        Annotated::new(self, annotations)
180    }
181    fn annotate(self, annotations: Annotations) -> Annotated<Self>
182    where
183        Self: Sized,
184    {
185        Annotated::new(self, Some(annotations))
186    }
187    fn no_annotation(self) -> Annotated<Self>
188    where
189        Self: Sized,
190    {
191        Annotated::new(self, None)
192    }
193    fn with_audience(self, audience: Vec<Role>) -> Annotated<Self>
194    where
195        Self: Sized,
196    {
197        self.annotate(Annotations {
198            audience: Some(audience),
199            ..Default::default()
200        })
201    }
202    fn with_priority(self, priority: f32) -> Annotated<Self>
203    where
204        Self: Sized,
205    {
206        self.annotate(Annotations {
207            priority: Some(priority),
208            ..Default::default()
209        })
210    }
211    fn with_timestamp(self, timestamp: DateTime<Utc>) -> Annotated<Self>
212    where
213        Self: Sized,
214    {
215        self.annotate(Annotations {
216            last_modified: Some(timestamp),
217            ..Default::default()
218        })
219    }
220    fn with_timestamp_now(self) -> Annotated<Self>
221    where
222        Self: Sized,
223    {
224        self.with_timestamp(Utc::now())
225    }
226}