chore: use typed notifications from rmcp (#3653)

crates/mcp-client/src/client.rs

@@ -5,6 +5,7 @@ use mcp_core::protocol::{
 use rmcp::model::{
     GetPromptResult, JsonRpcError, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest,
     JsonRpcResponse, JsonRpcVersion2_0, Notification, NumberOrString, Request, RequestId,
+    ServerNotification,
 };
 use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
@@ -106,7 +107,7 @@ pub trait McpClientTrait: Send + Sync {
 
     async fn get_prompt(&self, name: &str, arguments: Value) -> Result<GetPromptResult, Error>;
 
-    async fn subscribe(&self) -> mpsc::Receiver<JsonRpcMessage>;
+    async fn subscribe(&self) -> mpsc::Receiver<ServerNotification>;
 }
 
 /// The MCP client is the interface for MCP operations.
@@ -118,7 +119,7 @@ where
     next_id_counter: AtomicU64, // Added for atomic ID generation
     server_capabilities: Option<ServerCapabilities>,
     server_info: Option<Implementation>,
-    notification_subscribers: Arc<Mutex<Vec<mpsc::Sender<JsonRpcMessage>>>>,
+    notification_subscribers: Arc<Mutex<Vec<mpsc::Sender<ServerNotification>>>>,
 }
 
 impl<T> McpClient<T>
@@ -129,7 +130,7 @@ where
         let service = McpService::new(transport.clone());
         let service_ptr = service.clone();
         let notification_subscribers =
-            Arc::new(Mutex::new(Vec::<mpsc::Sender<JsonRpcMessage>>::new()));
+            Arc::new(Mutex::new(Vec::<mpsc::Sender<ServerNotification>>::new()));
         let subscribers_ptr = notification_subscribers.clone();
 
         tokio::spawn(async move {
@@ -148,9 +149,22 @@ where
                             }) => {
                                 service_ptr.respond(&id.to_string(), Ok(message)).await;
                             }
-                            _ => {
+                            JsonRpcMessage::Notification(JsonRpcNotification {
+                                notification,
+                                ..
+                            }) => {
                                 let mut subs = subscribers_ptr.lock().await;
-                                subs.retain(|sub| sub.try_send(message.clone()).is_ok());
+                                if let Some(server_notification) = notification.into() {
+                                    subs.retain(|sub| {
+                                        sub.try_send(server_notification.clone()).is_ok()
+                                    });
+                                }
+                            }
+                            _ => {
+                                tracing::warn!(
+                                    "Received unexpected received message type: {:?}",
+                                    message
+                                );
                             }
                         }
                     }
@@ -437,7 +451,7 @@ where
         self.send_request("prompts/get", params).await
     }
 
-    async fn subscribe(&self) -> mpsc::Receiver<JsonRpcMessage> {
+    async fn subscribe(&self) -> mpsc::Receiver<ServerNotification> {
         let (tx, rx) = mpsc::channel(16);
         self.notification_subscribers.lock().await.push(tx);
         rx

crates/mcp-client/src/service.rs

@@ -6,7 +6,7 @@ use std::task::{Context, Poll};
 use tokio::sync::{oneshot, RwLock};
 use tower::{timeout::Timeout, Service, ServiceBuilder};
 
-use crate::transport::{Error, TransportHandle};
+use crate::transport::{Error, TransportHandle, TransportMessageRecv};
 
 /// A wrapper service that implements Tower's Service trait for MCP transport
 #[derive(Clone)]
@@ -23,7 +23,7 @@ impl<T: TransportHandle> McpService<T> {
         }
     }
 
-    pub async fn respond(&self, id: &str, response: Result<JsonRpcMessage, Error>) {
+    pub async fn respond(&self, id: &str, response: Result<TransportMessageRecv, Error>) {
         self.pending_requests.respond(id, response).await
     }
 
@@ -36,7 +36,7 @@ impl<T> Service<JsonRpcMessage> for McpService<T>
 where
     T: TransportHandle + Send + Sync + 'static,
 {
-    type Response = JsonRpcMessage;
+    type Response = TransportMessageRecv;
     type Error = Error;
     type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
 
@@ -63,7 +63,7 @@ where
                     // Handle notifications without waiting for a response
                     transport.send(request).await?;
                     // Return a dummy response for notifications
-                    let dummy_response: JsonRpcMessage =
+                    let dummy_response: Self::Response =
                         JsonRpcMessage::Response(rmcp::model::JsonRpcResponse {
                             jsonrpc: rmcp::model::JsonRpcVersion2_0,
                             id: rmcp::model::RequestId::Number(0),
@@ -91,7 +91,7 @@ where
 
 // A data structure to store pending requests and their response channels
 pub struct PendingRequests {
-    requests: RwLock<HashMap<String, oneshot::Sender<Result<JsonRpcMessage, Error>>>>,
+    requests: RwLock<HashMap<String, oneshot::Sender<Result<TransportMessageRecv, Error>>>>,
 }
 
 impl Default for PendingRequests {
@@ -107,11 +107,15 @@ impl PendingRequests {
         }
     }
 
-    pub async fn insert(&self, id: String, sender: oneshot::Sender<Result<JsonRpcMessage, Error>>) {
+    pub async fn insert(
+        &self,
+        id: String,
+        sender: oneshot::Sender<Result<TransportMessageRecv, Error>>,
+    ) {
         self.requests.write().await.insert(id, sender);
     }
 
-    pub async fn respond(&self, id: &str, response: Result<JsonRpcMessage, Error>) {
+    pub async fn respond(&self, id: &str, response: Result<TransportMessageRecv, Error>) {
         if let Some(tx) = self.requests.write().await.remove(id) {
             let _ = tx.send(response);
         }

crates/mcp-client/src/transport/mod.rs

@@ -1,7 +1,7 @@
 use async_trait::async_trait;
-use rmcp::model::JsonRpcMessage;
+use rmcp::model::{JsonObject, JsonRpcMessage, Request, ServerNotification};
 use thiserror::Error;
-use tokio::sync::{mpsc, oneshot};
+use tokio::sync::mpsc;
 
 pub type BoxError = Box<dyn std::error::Error + Sync + Send>;
 /// A generic error type for transport operations.
@@ -38,15 +38,6 @@ pub enum Error {
     SessionError(String),
 }
 
-/// A message that can be sent through the transport
-#[derive(Debug)]
-pub struct TransportMessage {
-    /// The JSON-RPC message to send
-    pub message: JsonRpcMessage,
-    /// Channel to receive the response on (None for notifications)
-    pub response_tx: Option<oneshot::Sender<Result<JsonRpcMessage, Error>>>,
-}
-
 /// A generic asynchronous transport trait with channel-based communication
 #[async_trait]
 pub trait Transport {
@@ -60,10 +51,12 @@ pub trait Transport {
     async fn close(&self) -> Result<(), Error>;
 }
 
+pub type TransportMessageRecv = JsonRpcMessage<Request, JsonObject, ServerNotification>;
+
 #[async_trait]
 pub trait TransportHandle: Send + Sync + Clone + 'static {
     async fn send(&self, message: JsonRpcMessage) -> Result<(), Error>;
-    async fn receive(&self) -> Result<JsonRpcMessage, Error>;
+    async fn receive(&self) -> Result<TransportMessageRecv, Error>;
 }
 
 pub async fn serialize_and_send(

crates/mcp-client/src/transport/sse.rs

@@ -1,4 +1,4 @@
-use crate::transport::Error;
+use crate::transport::{Error, TransportMessageRecv};
 use async_trait::async_trait;
 use eventsource_client::{Client, SSE};
 use futures::TryStreamExt;
@@ -23,7 +23,7 @@ pub struct SseActor {
     /// Receives messages (requests/notifications) from the handle
     receiver: mpsc::Receiver<String>,
     /// Sends messages (responses) back to the handle
-    sender: mpsc::Sender<JsonRpcMessage>,
+    sender: mpsc::Sender<TransportMessageRecv>,
     /// Base SSE URL
     sse_url: String,
     /// For sending HTTP POST requests
@@ -35,7 +35,7 @@ pub struct SseActor {
 impl SseActor {
     pub fn new(
         receiver: mpsc::Receiver<String>,
-        sender: mpsc::Sender<JsonRpcMessage>,
+        sender: mpsc::Sender<TransportMessageRecv>,
         sse_url: String,
         post_endpoint: Arc<RwLock<Option<String>>>,
     ) -> Self {
@@ -71,7 +71,7 @@ impl SseActor {
     /// - If a `message` event is received, parse it as `JsonRpcMessage`
     ///   and respond to pending requests if it's a `Response`.
     async fn handle_incoming_messages(
-        sender: mpsc::Sender<JsonRpcMessage>,
+        sender: mpsc::Sender<TransportMessageRecv>,
         sse_url: String,
         post_endpoint: Arc<RwLock<Option<String>>>,
     ) {
@@ -109,7 +109,7 @@ impl SseActor {
                     match event {
                         SSE::Event(e) if e.event_type == "message" => {
                             // Attempt to parse the SSE data as a JsonRpcMessage
-                            match serde_json::from_str::<JsonRpcMessage>(&e.data) {
+                            match serde_json::from_str::<TransportMessageRecv>(&e.data) {
                                 Ok(message) => {
                                     let _ = sender.send(message).await;
                                 }
@@ -184,7 +184,7 @@ impl SseActor {
 #[derive(Clone)]
 pub struct SseTransportHandle {
     sender: mpsc::Sender<String>,
-    receiver: Arc<Mutex<mpsc::Receiver<JsonRpcMessage>>>,
+    receiver: Arc<Mutex<mpsc::Receiver<TransportMessageRecv>>>,
 }
 
 #[async_trait::async_trait]
@@ -193,7 +193,7 @@ impl TransportHandle for SseTransportHandle {
         serialize_and_send(&self.sender, message).await
     }
 
-    async fn receive(&self) -> Result<JsonRpcMessage, Error> {
+    async fn receive(&self) -> Result<TransportMessageRecv, Error> {
         let mut receiver = self.receiver.lock().await;
         receiver.recv().await.ok_or(Error::ChannelClosed)
     }

crates/mcp-client/src/transport/stdio.rs

@@ -14,6 +14,8 @@ use nix::sys::signal::{kill, Signal};
 #[cfg(unix)]
 use nix::unistd::{getpgid, Pid};
 
+use crate::transport::TransportMessageRecv;
+
 use super::{serialize_and_send, Error, Transport, TransportHandle};
 
 // Global to track process groups we've created
@@ -24,7 +26,7 @@ static PROCESS_GROUP: AtomicI32 = AtomicI32::new(-1);
 /// It uses channels for message passing and handles responses asynchronously through a background task.
 pub struct StdioActor {
     receiver: Option<mpsc::Receiver<String>>,
-    sender: Option<mpsc::Sender<JsonRpcMessage>>,
+    sender: Option<mpsc::Sender<TransportMessageRecv>>,
     process: Child, // we store the process to keep it alive
     error_sender: mpsc::Sender<Error>,
     stdin: Option<ChildStdin>,
@@ -98,7 +100,7 @@ impl StdioActor {
         }
     }
 
-    async fn handle_proc_output(stdout: ChildStdout, sender: mpsc::Sender<JsonRpcMessage>) {
+    async fn handle_proc_output(stdout: ChildStdout, sender: mpsc::Sender<TransportMessageRecv>) {
         let mut reader = BufReader::new(stdout);
         let mut line = String::new();
         loop {
@@ -108,7 +110,7 @@ impl StdioActor {
                     break;
                 } // EOF
                 Ok(_) => {
-                    if let Ok(message) = serde_json::from_str::<JsonRpcMessage>(&line) {
+                    if let Ok(message) = serde_json::from_str::<TransportMessageRecv>(&line) {
                         tracing::debug!(
                             message = ?message,
                             "Received incoming message"
@@ -149,8 +151,8 @@ impl StdioActor {
 
 #[derive(Clone)]
 pub struct StdioTransportHandle {
-    sender: mpsc::Sender<String>,                         // to process
-    receiver: Arc<Mutex<mpsc::Receiver<JsonRpcMessage>>>, // from process
+    sender: mpsc::Sender<String>,                               // to process
+    receiver: Arc<Mutex<mpsc::Receiver<TransportMessageRecv>>>, // from process
     error_receiver: Arc<Mutex<mpsc::Receiver<Error>>>,
 }
 
@@ -163,7 +165,7 @@ impl TransportHandle for StdioTransportHandle {
         result
     }
 
-    async fn receive(&self) -> Result<JsonRpcMessage, Error> {
+    async fn receive(&self) -> Result<TransportMessageRecv, Error> {
         let mut receiver = self.receiver.lock().await;
         match receiver.recv().await {
             Some(message) => Ok(message),

crates/mcp-client/src/transport/streamable_http.rs

@@ -1,5 +1,5 @@
 use crate::oauth::{authenticate_service, ServiceConfig};
-use crate::transport::Error;
+use crate::transport::{Error, TransportMessageRecv};
 use async_trait::async_trait;
 use eventsource_client::{Client, SSE};
 use futures::TryStreamExt;
@@ -25,7 +25,7 @@ pub struct StreamableHttpActor {
     /// Receives messages (requests/notifications) from the handle
     receiver: mpsc::Receiver<String>,
     /// Sends messages (responses) back to the handle
-    sender: mpsc::Sender<JsonRpcMessage>,
+    sender: mpsc::Sender<TransportMessageRecv>,
     /// MCP endpoint URL
     mcp_endpoint: String,
     /// HTTP client for sending requests
@@ -41,7 +41,7 @@ pub struct StreamableHttpActor {
 impl StreamableHttpActor {
     pub fn new(
         receiver: mpsc::Receiver<String>,
-        sender: mpsc::Sender<JsonRpcMessage>,
+        sender: mpsc::Sender<TransportMessageRecv>,
         mcp_endpoint: String,
         session_id: Arc<RwLock<Option<String>>>,
         env: HashMap<String, String>,
@@ -84,8 +84,8 @@ impl StreamableHttpActor {
         debug!("Sending message to MCP endpoint: {}", message_str);
 
         // Parse the message to determine if it's a request that expects a response
-        let parsed_message: JsonRpcMessage =
-            serde_json::from_str(&message_str).map_err(Error::Serialization)?;
+        let parsed_message = serde_json::from_str::<TransportMessageRecv>(&message_str)
+            .map_err(Error::Serialization)?;
 
         let expects_response = matches!(
             parsed_message,
@@ -196,8 +196,8 @@ impl StreamableHttpActor {
             })?;
 
             if !response_text.is_empty() {
-                let json_message: JsonRpcMessage =
-                    serde_json::from_str(&response_text).map_err(Error::Serialization)?;
+                let json_message = serde_json::from_str::<TransportMessageRecv>(&response_text)
+                    .map_err(Error::Serialization)?;
 
                 let _ = self.sender.send(json_message).await;
             }
@@ -267,7 +267,7 @@ impl StreamableHttpActor {
                 // Empty line indicates end of event
                 if !event_data.is_empty() {
                     // Parse the streamed data as JSON-RPC message
-                    match serde_json::from_str::<JsonRpcMessage>(&event_data) {
+                    match serde_json::from_str::<TransportMessageRecv>(&event_data) {
                         Ok(message) => {
                             debug!("Received streaming HTTP response message: {:?}", message);
                             let _ = self.sender.send(message).await;
@@ -301,7 +301,7 @@ impl StreamableHttpActor {
 #[derive(Clone)]
 pub struct StreamableHttpTransportHandle {
     sender: mpsc::Sender<String>,
-    receiver: Arc<Mutex<mpsc::Receiver<JsonRpcMessage>>>,
+    receiver: Arc<Mutex<mpsc::Receiver<TransportMessageRecv>>>,
     session_id: Arc<RwLock<Option<String>>>,
     mcp_endpoint: String,
     http_client: HttpClient,
@@ -314,7 +314,7 @@ impl TransportHandle for StreamableHttpTransportHandle {
         serialize_and_send(&self.sender, message).await
     }
 
-    async fn receive(&self) -> Result<JsonRpcMessage, Error> {
+    async fn receive(&self) -> Result<TransportMessageRecv, Error> {
         let mut receiver = self.receiver.lock().await;
         receiver.recv().await.ok_or(Error::ChannelClosed)
     }