trait WithMeta

Author: =

crates/rmcp/src/handler/server/tool.rs

@@ -65,7 +65,11 @@ pub struct ToolCallContext<'service, S> {
 impl<'service, S> ToolCallContext<'service, S> {
     pub fn new(
         service: &'service S,
-        CallToolRequestParam { name, arguments }: CallToolRequestParam,
+        CallToolRequestParam {
+            name,
+            arguments,
+            _meta,
+        }: CallToolRequestParam,
         request_context: RequestContext<RoleServer>,
     ) -> Self {
         Self {

crates/rmcp/src/model.rs

@@ -2,13 +2,15 @@ use std::{borrow::Cow, sync::Arc};
 mod annotated;
 mod capabilities;
 mod content;
+mod meta;
 mod prompt;
 mod resource;
 mod tool;
 
 pub use annotated::*;
 pub use capabilities::*;
 pub use content::*;
+pub use meta::*;
 pub use prompt::*;
 pub use resource::*;
 use serde::{Deserialize, Serialize};
@@ -190,22 +192,15 @@ impl<'de> Deserialize<'de> for NumberOrString {
 
 pub type RequestId = NumberOrString;
 pub type ProgressToken = NumberOrString;
-#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
-pub struct WithMeta<P = JsonObject, M = ()> {
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub _meta: Option<M>,
-    #[serde(flatten)]
-    pub inner: P,
-}
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 #[serde(rename_all = "camelCase")]
 pub struct RequestMeta {
-    progress_token: ProgressToken,
+    pub progress_token: ProgressToken,
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
-pub struct Request<M = String, P = Option<WithMeta<JsonObject, RequestMeta>>> {
+pub struct Request<M = String, P = JsonObject> {
     pub method: M,
     // #[serde(skip_serializing_if = "Option::is_none")]
     pub params: P,
@@ -216,7 +211,7 @@ pub struct RequestNoParam<M = String> {
 }
 
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
-pub struct Notification<M = String, P = Option<WithMeta<JsonObject, JsonObject>>> {
+pub struct Notification<M = String, P = JsonObject> {
     pub method: M,
     pub params: P,
 }
@@ -233,9 +228,9 @@ pub struct JsonRpcRequest<R = Request> {
     #[serde(flatten)]
     pub request: R,
 }
-type DefaultResponse = WithMeta<JsonObject, JsonObject>;
+type DefaultResponse = JsonObject;
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
-pub struct JsonRpcResponse<R = DefaultResponse> {
+pub struct JsonRpcResponse<R = JsonObject> {
     pub jsonrpc: JsonRpcVersion2_0,
     pub id: RequestId,
     pub result: R,
@@ -846,6 +841,8 @@ const_string!(CallToolRequestMethod = "tools/call");
 #[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
 #[serde(rename_all = "camelCase")]
 pub struct CallToolRequestParam {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub _meta: Option<RequestMeta>,
     pub name: Cow<'static, str>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub arguments: Option<JsonObject>,
@@ -1041,7 +1038,7 @@ mod tests {
                 assert_eq!(r.id, RequestId::Number(1));
                 assert_eq!(r.request.method, "request");
                 assert_eq!(
-                    &r.request.params.as_ref().unwrap().inner,
+                    &r.request.params,
                     json!({"key": "value"})
                         .as_object()
                         .expect("should be an object")

crates/rmcp/src/model/meta.rs

@@ -0,0 +1,34 @@
+use super::{ClientRequest, RequestMeta, ServerRequest};
+
+pub trait WithMeta<M> {
+    fn set_meta(&mut self, meta: Option<M>);
+    fn get_meta(&self) -> Option<&M>;
+}
+
+impl WithMeta<RequestMeta> for ClientRequest {
+    fn set_meta(&mut self, meta: Option<RequestMeta>) {
+        #[allow(clippy::single_match)]
+        match self {
+            ClientRequest::CallToolRequest(req) => {
+                req.params._meta = meta;
+            }
+            _ => {}
+        }
+    }
+
+    fn get_meta(&self) -> Option<&RequestMeta> {
+        #[allow(clippy::single_match)]
+        match self {
+            ClientRequest::CallToolRequest(req) => req.params._meta.as_ref(),
+            _ => None,
+        }
+    }
+}
+
+impl WithMeta<RequestMeta> for ServerRequest {
+    fn set_meta(&mut self, _meta: Option<RequestMeta>) {}
+
+    fn get_meta(&self) -> Option<&RequestMeta> {
+        None
+    }
+}

crates/rmcp/src/service.rs

@@ -6,7 +6,7 @@ use crate::{
     model::{
         CancelledNotification, CancelledNotificationParam, JsonRpcBatchRequestItem,
         JsonRpcBatchResponseItem, JsonRpcError, JsonRpcMessage, JsonRpcNotification,
-        JsonRpcRequest, JsonRpcResponse, RequestId,
+        JsonRpcRequest, JsonRpcResponse, ProgressToken, RequestId, RequestMeta, WithMeta,
     },
     transport::IntoTransport,
 };
@@ -58,12 +58,12 @@ impl<T> TransferObject for T where
 
 #[allow(private_bounds, reason = "there's no the third implementation")]
 pub trait ServiceRole: std::fmt::Debug + Send + Sync + 'static + Copy + Clone {
-    type Req: TransferObject;
+    type Req: TransferObject + WithMeta<RequestMeta>;
     type Resp: TransferObject;
     type Not: TryInto<CancelledNotification, Error = Self::Not>
         + From<CancelledNotification>
         + TransferObject;
-    type PeerReq: TransferObject;
+    type PeerReq: TransferObject + WithMeta<RequestMeta>;
     type PeerResp: TransferObject;
     type PeerNot: TryInto<CancelledNotification, Error = Self::PeerNot>
         + From<CancelledNotification>
@@ -201,17 +201,30 @@ pub trait RequestIdProvider: Send + Sync + 'static {
     fn next_request_id(&self) -> RequestId;
 }
 
+pub trait ProgressTokenProvider: Send + Sync + 'static {
+    fn next_progress_token(&self) -> ProgressToken;
+}
+
+pub type AtomicU32RequestIdProvider = AtomicU32Provider;
+pub type AtomicU32ProgressTokenProvider = AtomicU32Provider;
+
 #[derive(Debug, Default)]
-pub struct AtomicU32RequestIdProvider {
+pub struct AtomicU32Provider {
     id: AtomicU32,
 }
 
-impl RequestIdProvider for AtomicU32RequestIdProvider {
+impl RequestIdProvider for AtomicU32Provider {
     fn next_request_id(&self) -> RequestId {
         RequestId::Number(self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst))
     }
 }
 
+impl ProgressTokenProvider for AtomicU32Provider {
+    fn next_progress_token(&self) -> RequestId {
+        RequestId::Number(self.id.fetch_add(1, std::sync::atomic::Ordering::SeqCst))
+    }
+}
+
 type Responder<T> = tokio::sync::oneshot::Sender<T>;
 
 /// A handle to a remote request
@@ -225,6 +238,7 @@ pub struct RequestHandle<R: ServiceRole> {
     pub options: PeerRequestOptions,
     pub peer: Peer<R>,
     pub id: RequestId,
+    pub progress_token: ProgressToken,
 }
 
 impl<R: ServiceRole> RequestHandle<R> {
@@ -275,13 +289,16 @@ impl<R: ServiceRole> RequestHandle<R> {
 }
 
 #[derive(Debug)]
-pub enum PeerSinkMessage<R: ServiceRole> {
-    Request(
-        R::Req,
-        RequestId,
-        Responder<Result<R::PeerResp, ServiceError>>,
-    ),
-    Notification(R::Not, Responder<Result<(), ServiceError>>),
+pub(crate) enum PeerSinkMessage<R: ServiceRole> {
+    Request {
+        request: R::Req,
+        id: RequestId,
+        responder: Responder<Result<R::PeerResp, ServiceError>>,
+    },
+    Notification {
+        notification: R::Not,
+        responder: Responder<Result<(), ServiceError>>,
+    },
 }
 
 /// An interface to fetch the remote client or server
@@ -293,6 +310,7 @@ pub enum PeerSinkMessage<R: ServiceRole> {
 pub struct Peer<R: ServiceRole> {
     tx: mpsc::Sender<PeerSinkMessage<R>>,
     request_id_provider: Arc<dyn RequestIdProvider>,
+    progress_token_provider: Arc<dyn ProgressTokenProvider>,
     info: Arc<R::PeerInfo>,
 }
 
@@ -320,7 +338,7 @@ impl PeerRequestOptions {
 
 impl<R: ServiceRole> Peer<R> {
     const CLIENT_CHANNEL_BUFFER_SIZE: usize = 1024;
-    pub fn new(
+    pub(crate) fn new(
         request_id_provider: Arc<dyn RequestIdProvider>,
         peer_info: R::PeerInfo,
     ) -> (Peer<R>, ProxyOutbound<R>) {
@@ -329,6 +347,7 @@ impl<R: ServiceRole> Peer<R> {
             Self {
                 tx,
                 request_id_provider,
+                progress_token_provider: Arc::new(AtomicU32ProgressTokenProvider::default()),
                 info: peer_info.into(),
             },
             rx,
@@ -337,7 +356,10 @@ impl<R: ServiceRole> Peer<R> {
     pub async fn send_notification(&self, notification: R::Not) -> Result<(), ServiceError> {
         let (responder, receiver) = tokio::sync::oneshot::channel();
         self.tx
-            .send(PeerSinkMessage::Notification(notification, responder))
+            .send(PeerSinkMessage::Notification {
+                notification,
+                responder,
+            })
             .await
             .map_err(|_m| ServiceError::Transport(std::io::Error::other("disconnected")))?;
         receiver
@@ -352,18 +374,27 @@ impl<R: ServiceRole> Peer<R> {
     }
     pub async fn send_cancellable_request(
         &self,
-        request: R::Req,
+        mut request: R::Req,
         options: PeerRequestOptions,
     ) -> Result<RequestHandle<R>, ServiceError> {
         let id = self.request_id_provider.next_request_id();
+        let progress_token = self.progress_token_provider.next_progress_token();
+        request.set_meta(Some(RequestMeta {
+            progress_token: progress_token.clone(),
+        }));
         let (responder, receiver) = tokio::sync::oneshot::channel();
         self.tx
-            .send(PeerSinkMessage::Request(request, id.clone(), responder))
+            .send(PeerSinkMessage::Request {
+                request,
+                id: id.clone(),
+                responder,
+            })
             .await
             .map_err(|_m| ServiceError::Transport(std::io::Error::other("disconnected")))?;
         Ok(RequestHandle {
             id,
             rx: receiver,
+            progress_token,
             options,
             peer: self.clone(),
         })
@@ -419,6 +450,7 @@ pub struct RequestContext<R: ServiceRole> {
     /// this token will be cancelled when the [`CancelledNotification`] is received.
     pub ct: CancellationToken,
     pub id: RequestId,
+    pub meta: Option<RequestMeta>,
     /// An interface to fetch the remote client or server
     pub peer: Peer<R>,
 }
@@ -459,7 +491,7 @@ async fn serve_inner<R, S, T, E, A>(
     mut service: S,
     transport: T,
     peer_info: R::PeerInfo,
-    id_provider: Arc<AtomicU32RequestIdProvider>,
+    id_provider: Arc<AtomicU32Provider>,
     ct: CancellationToken,
 ) -> Result<RunningService<R, S>, E>
 where
@@ -555,7 +587,11 @@ where
                         }
                     }
                 }
-                Event::ProxyMessage(PeerSinkMessage::Request(request, id, responder)) => {
+                Event::ProxyMessage(PeerSinkMessage::Request {
+                    request,
+                    id,
+                    responder,
+                }) => {
                     local_responder_pool.insert(id.clone(), responder);
                     let send_result = sink
                         .send(JsonRpcMessage::request(request, id.clone()))
@@ -567,7 +603,10 @@ where
                         }
                     }
                 }
-                Event::ProxyMessage(PeerSinkMessage::Notification(notification, responder)) => {
+                Event::ProxyMessage(PeerSinkMessage::Notification {
+                    notification,
+                    responder,
+                }) => {
                     // catch cancellation notification
                     let mut cancellation_param = None;
                     let notification = match notification.try_into() {
@@ -605,6 +644,7 @@ where
                             ct: context_ct,
                             id: id.clone(),
                             peer: peer.clone(),
+                            meta: request.get_meta().cloned(),
                         };
                         tokio::spawn(async move {
                             let result = service.handle_request(request, context).await;

crates/rmcp/src/service/client.rs

@@ -73,7 +73,7 @@ where
     let (sink, stream) = transport.into_transport();
     let mut sink = Box::pin(sink);
     let mut stream = Box::pin(stream);
-    let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
+    let id_provider = <Arc<AtomicU32Provider>>::default();
     // service
     let id = id_provider.next_request_id();
     let init_request = InitializeRequest {

crates/rmcp/src/service/server.rs

@@ -6,10 +6,9 @@ use crate::model::{
     ClientRequest, ClientResult, CreateMessageRequest, CreateMessageRequestParam,
     CreateMessageResult, ListRootsRequest, ListRootsResult, LoggingMessageNotification,
     LoggingMessageNotificationParam, ProgressNotification, ProgressNotificationParam,
-    PromptListChangedNotification, ProtocolVersion, ResourceListChangedNotification,
-    ResourceUpdatedNotification, ResourceUpdatedNotificationParam, ServerInfo,
-    ServerJsonRpcMessage, ServerNotification, ServerRequest, ServerResult,
-    ToolListChangedNotification,
+    PromptListChangedNotification, ResourceListChangedNotification, ResourceUpdatedNotification,
+    ResourceUpdatedNotificationParam, ServerInfo, ServerJsonRpcMessage, ServerNotification,
+    ServerRequest, ServerResult, ToolListChangedNotification,
 };
 
 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
@@ -69,7 +68,7 @@ where
     let (sink, stream) = transport.into_transport();
     let mut sink = Box::pin(sink);
     let mut stream = Box::pin(stream);
-    let id_provider = <Arc<AtomicU32RequestIdProvider>>::default();
+    let id_provider = <Arc<AtomicU32Provider>>::default();
     // service
     let (request, id) = stream
         .next()

crates/rmcp/src/transport.rs

@@ -67,6 +67,7 @@ pub use sse_server::SseServer;
 // #[cfg(feature = "transport-ws")]
 // pub mod ws;
 
+pub mod streamable_http;
 pub trait IntoTransport<R, E, A>: Send + 'static
 where
     R: ServiceRole,

crates/rmcp/src/transport/streamable_http.rs

@@ -0,0 +1 @@
+pub mod session;