rmcp/transport/
async_rw.rs

1use std::{marker::PhantomData, sync::Arc};
2
3// use crate::schema::*;
4use futures::{SinkExt, StreamExt};
5use serde::{Serialize, de::DeserializeOwned};
6use thiserror::Error;
7use tokio::{
8    io::{AsyncRead, AsyncWrite},
9    sync::Mutex,
10};
11use tokio_util::{
12    bytes::{Buf, BufMut, BytesMut},
13    codec::{Decoder, Encoder, FramedRead, FramedWrite},
14};
15
16use super::{IntoTransport, Transport};
17use crate::service::{RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage};
18
19pub enum TransportAdapterAsyncRW {}
20
21impl<Role, R, W> IntoTransport<Role, std::io::Error, TransportAdapterAsyncRW> for (R, W)
22where
23    Role: ServiceRole,
24    R: AsyncRead + Send + 'static + Unpin,
25    W: AsyncWrite + Send + 'static + Unpin,
26{
27    fn into_transport(self) -> impl Transport<Role, Error = std::io::Error> + 'static {
28        AsyncRwTransport::new(self.0, self.1)
29    }
30}
31
32pub enum TransportAdapterAsyncCombinedRW {}
33impl<Role, S> IntoTransport<Role, std::io::Error, TransportAdapterAsyncCombinedRW> for S
34where
35    Role: ServiceRole,
36    S: AsyncRead + AsyncWrite + Send + 'static,
37{
38    fn into_transport(self) -> impl Transport<Role, Error = std::io::Error> + 'static {
39        IntoTransport::<Role, std::io::Error, TransportAdapterAsyncRW>::into_transport(
40            tokio::io::split(self),
41        )
42    }
43}
44
45pub type TransportWriter<Role, W> = FramedWrite<W, JsonRpcMessageCodec<TxJsonRpcMessage<Role>>>;
46
47pub struct AsyncRwTransport<Role: ServiceRole, R: AsyncRead, W: AsyncWrite> {
48    read: FramedRead<R, JsonRpcMessageCodec<RxJsonRpcMessage<Role>>>,
49    write: Arc<Mutex<Option<TransportWriter<Role, W>>>>,
50}
51
52impl<Role: ServiceRole, R, W> AsyncRwTransport<Role, R, W>
53where
54    R: Send + AsyncRead + Unpin,
55    W: Send + AsyncWrite + Unpin + 'static,
56{
57    pub fn new(read: R, write: W) -> Self {
58        let read = FramedRead::new(
59            read,
60            JsonRpcMessageCodec::<RxJsonRpcMessage<Role>>::default(),
61        );
62        let write = Arc::new(Mutex::new(Some(FramedWrite::new(
63            write,
64            JsonRpcMessageCodec::<TxJsonRpcMessage<Role>>::default(),
65        ))));
66        Self { read, write }
67    }
68}
69
70#[cfg(feature = "client")]
71impl<R, W> AsyncRwTransport<crate::RoleClient, R, W>
72where
73    R: Send + AsyncRead + Unpin,
74    W: Send + AsyncWrite + Unpin + 'static,
75{
76    pub fn new_client(read: R, write: W) -> Self {
77        Self::new(read, write)
78    }
79}
80
81#[cfg(feature = "server")]
82impl<R, W> AsyncRwTransport<crate::RoleServer, R, W>
83where
84    R: Send + AsyncRead + Unpin,
85    W: Send + AsyncWrite + Unpin + 'static,
86{
87    pub fn new_server(read: R, write: W) -> Self {
88        Self::new(read, write)
89    }
90}
91
92impl<Role: ServiceRole, R, W> Transport<Role> for AsyncRwTransport<Role, R, W>
93where
94    R: Send + AsyncRead + Unpin,
95    W: Send + AsyncWrite + Unpin + 'static,
96{
97    type Error = std::io::Error;
98
99    fn send(
100        &mut self,
101        item: TxJsonRpcMessage<Role>,
102    ) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
103        let lock = self.write.clone();
104        async move {
105            let mut write = lock.lock().await;
106            if let Some(ref mut write) = *write {
107                write.send(item).await.map_err(Into::into)
108            } else {
109                Err(std::io::Error::new(
110                    std::io::ErrorKind::NotConnected,
111                    "Transport is closed",
112                ))
113            }
114        }
115    }
116
117    fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<Role>>> {
118        let next = self.read.next();
119        async {
120            next.await.and_then(|e| {
121                e.inspect_err(|e| {
122                    tracing::error!("Error reading from stream: {}", e);
123                })
124                .ok()
125            })
126        }
127    }
128
129    async fn close(&mut self) -> Result<(), Self::Error> {
130        let mut write = self.write.lock().await;
131        drop(write.take());
132        Ok(())
133    }
134}
135
136#[derive(Debug, Clone)]
137pub struct JsonRpcMessageCodec<T> {
138    _marker: PhantomData<fn() -> T>,
139    next_index: usize,
140    max_length: usize,
141    is_discarding: bool,
142}
143
144impl<T> Default for JsonRpcMessageCodec<T> {
145    fn default() -> Self {
146        Self::new()
147    }
148}
149
150impl<T> JsonRpcMessageCodec<T> {
151    pub fn new() -> Self {
152        Self {
153            _marker: PhantomData,
154            next_index: 0,
155            max_length: usize::MAX,
156            is_discarding: false,
157        }
158    }
159
160    pub fn new_with_max_length(max_length: usize) -> Self {
161        Self {
162            max_length,
163            ..Self::new()
164        }
165    }
166
167    pub fn max_length(&self) -> usize {
168        self.max_length
169    }
170}
171
172fn without_carriage_return(s: &[u8]) -> &[u8] {
173    if let Some(&b'\r') = s.last() {
174        &s[..s.len() - 1]
175    } else {
176        s
177    }
178}
179
180/// Check if a method is a standard MCP method (request, response, or notification).
181/// This includes both requests and notifications defined in the MCP specification.
182///
183/// Based on MCP specification 2025-06-18: https://modelcontextprotocol.io/specification/2025-06-18
184fn is_standard_method(method: &str) -> bool {
185    matches!(
186        method,
187        "initialize"
188            | "ping"
189            | "prompts/get"
190            | "prompts/list"
191            | "resources/list"
192            | "resources/read"
193            | "resources/subscribe"
194            | "resources/unsubscribe"
195            | "resources/templates/list"
196            | "tools/call"
197            | "tools/list"
198            | "completion/complete"
199            | "logging/setLevel"
200            | "roots/list"
201            | "sampling/createMessage"
202    ) || is_standard_notification(method)
203}
204
205fn is_standard_notification(method: &str) -> bool {
206    matches!(
207        method,
208        "notifications/cancelled"
209            | "notifications/initialized"
210            | "notifications/message"
211            | "notifications/progress"
212            | "notifications/prompts/list_changed"
213            | "notifications/resources/list_changed"
214            | "notifications/resources/updated"
215            | "notifications/roots/list_changed"
216            | "notifications/tools/list_changed"
217    )
218}
219
220/// Determines if a notification should be ignored for compatibility.
221fn should_ignore_notification(json_value: &serde_json::Value, method: &str) -> bool {
222    let is_notification = json_value.get("id").is_none();
223
224    // Ignore non-MCP notifications (like LSP messages) for compatibility
225    if is_notification && !is_standard_method(method) {
226        tracing::trace!(
227            "Ignoring non-MCP notification '{}' for compatibility",
228            method
229        );
230        return true;
231    }
232
233    // Ignore non-standard MCP notifications
234    matches!(
235        (
236            method.starts_with("notifications/"),
237            is_standard_notification(method)
238        ),
239        (true, false)
240    )
241}
242
243/// Try to parse a message with compatibility handling for non-standard notifications
244fn try_parse_with_compatibility<T: serde::de::DeserializeOwned>(
245    line: &[u8],
246    context: &str,
247) -> Result<Option<T>, JsonRpcMessageCodecError> {
248    if let Ok(line_str) = std::str::from_utf8(line) {
249        match serde_json::from_slice(line) {
250            Ok(item) => Ok(Some(item)),
251            Err(e) => {
252                // Check if this is a notification that should be ignored for compatibility
253                if let Ok(json_value) = serde_json::from_str::<serde_json::Value>(line_str) {
254                    if let Some(method) =
255                        json_value.get("method").and_then(serde_json::Value::as_str)
256                    {
257                        if should_ignore_notification(&json_value, method) {
258                            return Ok(None);
259                        }
260                    }
261                }
262
263                tracing::debug!(
264                    "Failed to parse message {}: {} | Error: {}",
265                    context,
266                    line_str,
267                    e
268                );
269                Err(JsonRpcMessageCodecError::Serde(e))
270            }
271        }
272    } else {
273        serde_json::from_slice(line)
274            .map(Some)
275            .map_err(JsonRpcMessageCodecError::Serde)
276    }
277}
278
279#[derive(Debug, Error)]
280pub enum JsonRpcMessageCodecError {
281    #[error("max line length exceeded")]
282    MaxLineLengthExceeded,
283    #[error("serde error {0}")]
284    Serde(#[from] serde_json::Error),
285    #[error("io error {0}")]
286    Io(#[from] std::io::Error),
287}
288
289impl From<JsonRpcMessageCodecError> for std::io::Error {
290    fn from(value: JsonRpcMessageCodecError) -> Self {
291        match value {
292            JsonRpcMessageCodecError::MaxLineLengthExceeded => {
293                std::io::Error::new(std::io::ErrorKind::InvalidData, value)
294            }
295            JsonRpcMessageCodecError::Serde(e) => e.into(),
296            JsonRpcMessageCodecError::Io(e) => e,
297        }
298    }
299}
300
301impl<T: DeserializeOwned> Decoder for JsonRpcMessageCodec<T> {
302    type Item = T;
303
304    type Error = JsonRpcMessageCodecError;
305
306    fn decode(
307        &mut self,
308        buf: &mut BytesMut,
309    ) -> Result<Option<Self::Item>, JsonRpcMessageCodecError> {
310        loop {
311            // Determine how far into the buffer we'll search for a newline. If
312            // there's no max_length set, we'll read to the end of the buffer.
313            let read_to = std::cmp::min(self.max_length.saturating_add(1), buf.len());
314
315            let newline_offset = buf[self.next_index..read_to]
316                .iter()
317                .position(|b| *b == b'\n');
318
319            match (self.is_discarding, newline_offset) {
320                (true, Some(offset)) => {
321                    // If we found a newline, discard up to that offset and
322                    // then stop discarding. On the next iteration, we'll try
323                    // to read a line normally.
324                    buf.advance(offset + self.next_index + 1);
325                    self.is_discarding = false;
326                    self.next_index = 0;
327                }
328                (true, None) => {
329                    // Otherwise, we didn't find a newline, so we'll discard
330                    // everything we read. On the next iteration, we'll continue
331                    // discarding up to max_len bytes unless we find a newline.
332                    buf.advance(read_to);
333                    self.next_index = 0;
334                    if buf.is_empty() {
335                        return Ok(None);
336                    }
337                }
338                (false, Some(offset)) => {
339                    // Found a line!
340                    let newline_index = offset + self.next_index;
341                    self.next_index = 0;
342                    let line = buf.split_to(newline_index + 1);
343                    let line = &line[..line.len() - 1];
344                    let line = without_carriage_return(line);
345
346                    // Use compatibility handling function
347                    let item = match try_parse_with_compatibility(line, "decode")? {
348                        Some(item) => item,
349                        None => return Ok(None), // Skip non-standard message
350                    };
351                    return Ok(Some(item));
352                }
353                (false, None) if buf.len() > self.max_length => {
354                    // Reached the maximum length without finding a
355                    // newline, return an error and start discarding on the
356                    // next call.
357                    self.is_discarding = true;
358                    return Err(JsonRpcMessageCodecError::MaxLineLengthExceeded);
359                }
360                (false, None) => {
361                    // We didn't find a line or reach the length limit, so the next
362                    // call will resume searching at the current offset.
363                    self.next_index = read_to;
364                    return Ok(None);
365                }
366            }
367        }
368    }
369
370    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<T>, JsonRpcMessageCodecError> {
371        Ok(match self.decode(buf)? {
372            Some(frame) => Some(frame),
373            None => {
374                self.next_index = 0;
375                // No terminating newline - return remaining data, if any
376                if buf.is_empty() || buf == &b"\r"[..] {
377                    None
378                } else {
379                    let line = buf.split_to(buf.len());
380                    let line = without_carriage_return(&line);
381
382                    // Use compatibility handling function
383                    let item = match try_parse_with_compatibility(line, "decode_eof")? {
384                        Some(item) => item,
385                        None => return Ok(None), // Skip non-standard message
386                    };
387                    Some(item)
388                }
389            }
390        })
391    }
392}
393
394impl<T: Serialize> Encoder<T> for JsonRpcMessageCodec<T> {
395    type Error = JsonRpcMessageCodecError;
396
397    fn encode(&mut self, item: T, buf: &mut BytesMut) -> Result<(), JsonRpcMessageCodecError> {
398        serde_json::to_writer(buf.writer(), &item)?;
399        buf.put_u8(b'\n');
400        Ok(())
401    }
402}
403
404#[cfg(test)]
405mod test {
406    use futures::{Sink, Stream};
407
408    use super::*;
409    fn from_async_read<T: DeserializeOwned, R: AsyncRead>(reader: R) -> impl Stream<Item = T> {
410        FramedRead::new(reader, JsonRpcMessageCodec::<T>::default()).filter_map(|result| {
411            if let Err(e) = &result {
412                tracing::error!("Error reading from stream: {}", e);
413            }
414            futures::future::ready(result.ok())
415        })
416    }
417
418    fn from_async_write<T: Serialize, W: AsyncWrite + Send>(
419        writer: W,
420    ) -> impl Sink<T, Error = std::io::Error> {
421        FramedWrite::new(writer, JsonRpcMessageCodec::<T>::default()).sink_map_err(Into::into)
422    }
423    #[tokio::test]
424    async fn test_decode() {
425        use futures::StreamExt;
426        use tokio::io::BufReader;
427
428        let data = r#"{"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":1}
429    {"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":2}
430    {"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":3}
431    {"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":4}
432    {"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":5}
433    {"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":6}
434    {"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":7}
435    {"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":8}
436    {"jsonrpc":"2.0","method":"subtract","params":[42,23],"id":9}
437    {"jsonrpc":"2.0","method":"subtract","params":[23,42],"id":10}
438
439    "#;
440
441        let mut cursor = BufReader::new(data.as_bytes());
442        let mut stream = from_async_read::<serde_json::Value, _>(&mut cursor);
443
444        for i in 1..=10 {
445            let item = stream.next().await.unwrap();
446            assert_eq!(
447                item,
448                serde_json::json!({
449                    "jsonrpc": "2.0",
450                    "method": "subtract",
451                    "params": if i % 2 != 0 { [42, 23] } else { [23, 42] },
452                    "id": i,
453                })
454            );
455        }
456    }
457
458    #[tokio::test]
459    async fn test_encode() {
460        let test_messages = vec![
461            serde_json::json!({
462                "jsonrpc": "2.0",
463                "method": "subtract",
464                "params": [42, 23],
465                "id": 1,
466            }),
467            serde_json::json!({
468                "jsonrpc": "2.0",
469                "method": "subtract",
470                "params": [23, 42],
471                "id": 2,
472            }),
473        ];
474
475        // Create a buffer to write to
476        let mut buffer = Vec::new();
477        let mut writer = from_async_write(&mut buffer);
478
479        // Write the test messages
480        for message in test_messages.iter() {
481            writer.send(message.clone()).await.unwrap();
482        }
483        writer.close().await.unwrap();
484        drop(writer);
485        // Parse the buffer back into lines and check each one
486        let output = String::from_utf8_lossy(&buffer);
487        let mut lines = output.lines();
488
489        for expected_message in test_messages {
490            let line = lines.next().unwrap();
491            let parsed_message: serde_json::Value = serde_json::from_str(line).unwrap();
492            assert_eq!(parsed_message, expected_message);
493        }
494
495        // Make sure there are no extra lines
496        assert!(lines.next().is_none());
497    }
498
499    #[test]
500    fn test_standard_notification_check() {
501        // Test that all standard notifications are recognized
502        assert!(is_standard_notification("notifications/cancelled"));
503        assert!(is_standard_notification("notifications/initialized"));
504        assert!(is_standard_notification("notifications/progress"));
505        assert!(is_standard_notification(
506            "notifications/resources/list_changed"
507        ));
508        assert!(is_standard_notification("notifications/resources/updated"));
509        assert!(is_standard_notification(
510            "notifications/prompts/list_changed"
511        ));
512        assert!(is_standard_notification("notifications/tools/list_changed"));
513        assert!(is_standard_notification("notifications/message"));
514        assert!(is_standard_notification("notifications/roots/list_changed"));
515
516        // Test that non-standard notifications are not recognized
517        assert!(!is_standard_notification("notifications/stderr"));
518        assert!(!is_standard_notification("notifications/custom"));
519        assert!(!is_standard_notification("notifications/debug"));
520        assert!(!is_standard_notification("some/other/method"));
521    }
522
523    #[test]
524    fn test_compatibility_function() {
525        // Test the compatibility function directly
526        let stderr_message =
527            r#"{"method":"notifications/stderr","params":{"content":"stderr message"}}"#;
528        let custom_message = r#"{"method":"notifications/custom","params":{"data":"custom"}}"#;
529        let standard_message =
530            r#"{"method":"notifications/message","params":{"level":"info","data":"standard"}}"#;
531        let progress_message = r#"{"method":"notifications/progress","params":{"progressToken":"token","progress":50}}"#;
532
533        // Test with valid JSON - all should parse successfully
534        let result1 =
535            try_parse_with_compatibility::<serde_json::Value>(stderr_message.as_bytes(), "test");
536        let result2 =
537            try_parse_with_compatibility::<serde_json::Value>(custom_message.as_bytes(), "test");
538        let result3 =
539            try_parse_with_compatibility::<serde_json::Value>(standard_message.as_bytes(), "test");
540        let result4 =
541            try_parse_with_compatibility::<serde_json::Value>(progress_message.as_bytes(), "test");
542
543        // All should parse successfully since they're valid JSON
544        assert!(result1.is_ok());
545        assert!(result2.is_ok());
546        assert!(result3.is_ok());
547        assert!(result4.is_ok());
548
549        // Standard notifications should return Some(value)
550        assert!(result3.unwrap().is_some());
551        assert!(result4.unwrap().is_some());
552
553        println!("Standard notifications are preserved, non-standard are handled gracefully");
554    }
555}
rmcp/transport/async_rw.rs

rmcp/transport/
async_rw.rs
