rmcp/model/
meta.rs

1use std::ops::{Deref, DerefMut};
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use super::{
7    ClientNotification, ClientRequest, Extensions, JsonObject, JsonRpcMessage, NumberOrString,
8    ProgressToken, ServerNotification, ServerRequest,
9};
10
11pub trait GetMeta {
12    fn get_meta_mut(&mut self) -> &mut Meta;
13    fn get_meta(&self) -> &Meta;
14}
15
16pub trait GetExtensions {
17    fn extensions(&self) -> &Extensions;
18    fn extensions_mut(&mut self) -> &mut Extensions;
19}
20
21macro_rules! variant_extension {
22    (
23        $Enum: ident {
24            $($variant: ident)*
25        }
26    ) => {
27        impl GetExtensions for $Enum {
28            fn extensions(&self) -> &Extensions {
29                match self {
30                    $(
31                        $Enum::$variant(v) => &v.extensions,
32                    )*
33                }
34            }
35            fn extensions_mut(&mut self) -> &mut Extensions {
36                match self {
37                    $(
38                        $Enum::$variant(v) => &mut v.extensions,
39                    )*
40                }
41            }
42        }
43        impl GetMeta for $Enum {
44            fn get_meta_mut(&mut self) -> &mut Meta {
45                self.extensions_mut().get_or_insert_default()
46            }
47            fn get_meta(&self) -> &Meta {
48                self.extensions().get::<Meta>().unwrap_or(Meta::static_empty())
49            }
50        }
51    };
52}
53
54variant_extension! {
55    ClientRequest {
56        PingRequest
57        InitializeRequest
58        CompleteRequest
59        SetLevelRequest
60        GetPromptRequest
61        ListPromptsRequest
62        ListResourcesRequest
63        ListResourceTemplatesRequest
64        ReadResourceRequest
65        SubscribeRequest
66        UnsubscribeRequest
67        CallToolRequest
68        ListToolsRequest
69    }
70}
71
72variant_extension! {
73    ServerRequest {
74        PingRequest
75        CreateMessageRequest
76        ListRootsRequest
77        CreateElicitationRequest
78    }
79}
80
81variant_extension! {
82    ClientNotification {
83        CancelledNotification
84        ProgressNotification
85        InitializedNotification
86        RootsListChangedNotification
87    }
88}
89
90variant_extension! {
91    ServerNotification {
92        CancelledNotification
93        ProgressNotification
94        LoggingMessageNotification
95        ResourceUpdatedNotification
96        ResourceListChangedNotification
97        ToolListChangedNotification
98        PromptListChangedNotification
99    }
100}
101#[derive(Debug, Serialize, Deserialize, Clone, Default)]
102#[serde(transparent)]
103pub struct Meta(pub JsonObject);
104const PROGRESS_TOKEN_FIELD: &str = "progressToken";
105impl Meta {
106    pub fn new() -> Self {
107        Self(JsonObject::new())
108    }
109
110    pub(crate) fn static_empty() -> &'static Self {
111        static EMPTY: std::sync::OnceLock<Meta> = std::sync::OnceLock::new();
112        EMPTY.get_or_init(Default::default)
113    }
114
115    pub fn get_progress_token(&self) -> Option<ProgressToken> {
116        self.0.get(PROGRESS_TOKEN_FIELD).and_then(|v| match v {
117            Value::String(s) => Some(ProgressToken(NumberOrString::String(s.to_string().into()))),
118            Value::Number(n) => n
119                .as_u64()
120                .map(|n| ProgressToken(NumberOrString::Number(n as u32))),
121            _ => None,
122        })
123    }
124
125    pub fn set_progress_token(&mut self, token: ProgressToken) {
126        match token.0 {
127            NumberOrString::String(ref s) => self.0.insert(
128                PROGRESS_TOKEN_FIELD.to_string(),
129                Value::String(s.to_string()),
130            ),
131            NumberOrString::Number(n) => self
132                .0
133                .insert(PROGRESS_TOKEN_FIELD.to_string(), Value::Number(n.into())),
134        };
135    }
136
137    pub fn extend(&mut self, other: Meta) {
138        for (k, v) in other.0.into_iter() {
139            self.0.insert(k, v);
140        }
141    }
142}
143
144impl Deref for Meta {
145    type Target = JsonObject;
146
147    fn deref(&self) -> &Self::Target {
148        &self.0
149    }
150}
151
152impl DerefMut for Meta {
153    fn deref_mut(&mut self) -> &mut Self::Target {
154        &mut self.0
155    }
156}
157
158impl<Req, Resp, Noti> JsonRpcMessage<Req, Resp, Noti>
159where
160    Req: GetExtensions,
161    Noti: GetExtensions,
162{
163    pub fn insert_extension<T: Clone + Send + Sync + 'static>(&mut self, value: T) {
164        match self {
165            JsonRpcMessage::Request(json_rpc_request) => {
166                json_rpc_request.request.extensions_mut().insert(value);
167            }
168            JsonRpcMessage::Notification(json_rpc_notification) => {
169                json_rpc_notification
170                    .notification
171                    .extensions_mut()
172                    .insert(value);
173            }
174            JsonRpcMessage::BatchRequest(json_rpc_batch_request_items) => {
175                for item in json_rpc_batch_request_items {
176                    match item {
177                        super::JsonRpcBatchRequestItem::Request(json_rpc_request) => {
178                            json_rpc_request
179                                .request
180                                .extensions_mut()
181                                .insert(value.clone());
182                        }
183                        super::JsonRpcBatchRequestItem::Notification(json_rpc_notification) => {
184                            json_rpc_notification
185                                .notification
186                                .extensions_mut()
187                                .insert(value.clone());
188                        }
189                    }
190                }
191            }
192            _ => {}
193        }
194    }
195}
rmcp/model/meta.rs

rmcp/model/
meta.rs
