1use std::{
2 future::Future,
3 marker::PhantomData,
4 sync::{Arc, Mutex},
5};
6
7use futures::{SinkExt, StreamExt};
8use rmcp::{
9 service::{RoleServer, RxJsonRpcMessage, ServiceRole, TxJsonRpcMessage},
10 transport::Transport,
11};
12use serde::{Serialize, de::DeserializeOwned};
13use thiserror::Error;
14use tokio::{
15 io::{AsyncRead, AsyncWrite},
16 sync::Mutex as AsyncMutex,
17};
18use tokio_util::{
19 bytes::{Buf, BufMut, BytesMut},
20 codec::{Decoder, Encoder, FramedRead, FramedWrite},
21};
22
23#[derive(Clone, Copy, Debug, Eq, PartialEq)]
24enum WireProtocol {
25 JsonLine,
26 ContentLength,
27}
28
29#[derive(Debug, Clone)]
30pub struct SharedProtocol(Arc<Mutex<Option<WireProtocol>>>);
31
32impl SharedProtocol {
33 fn new() -> Self {
34 Self(Arc::new(Mutex::new(None)))
35 }
36
37 fn get(&self) -> Option<WireProtocol> {
38 *self
39 .0
40 .lock()
41 .unwrap_or_else(std::sync::PoisonError::into_inner)
42 }
43
44 fn set_if_unset(&self, protocol: WireProtocol) {
45 let mut guard = self
46 .0
47 .lock()
48 .unwrap_or_else(std::sync::PoisonError::into_inner);
49 if guard.is_none() {
50 *guard = Some(protocol);
51 }
52 }
53}
54
55pub type TransportWriter<Role, W> =
56 FramedWrite<W, HybridJsonRpcMessageCodec<TxJsonRpcMessage<Role>>>;
57
58pub struct HybridStdioTransport<Role: ServiceRole, R: AsyncRead, W: AsyncWrite> {
59 read: FramedRead<R, HybridJsonRpcMessageCodec<RxJsonRpcMessage<Role>>>,
60 write: Arc<AsyncMutex<Option<TransportWriter<Role, W>>>>,
61 protocol: SharedProtocol,
62}
63
64impl<Role: ServiceRole, R, W> HybridStdioTransport<Role, R, W>
65where
66 R: Send + AsyncRead + Unpin,
67 W: Send + AsyncWrite + Unpin + 'static,
68{
69 pub fn new(read: R, write: W) -> Self {
70 let protocol = SharedProtocol::new();
71 let read = FramedRead::new(
72 read,
73 HybridJsonRpcMessageCodec::<RxJsonRpcMessage<Role>>::new(protocol.clone()),
74 );
75 let write = Arc::new(AsyncMutex::new(Some(FramedWrite::new(
76 write,
77 HybridJsonRpcMessageCodec::<TxJsonRpcMessage<Role>>::new(protocol.clone()),
78 ))));
79 Self {
80 read,
81 write,
82 protocol,
83 }
84 }
85
86 pub fn protocol(&self) -> SharedProtocol {
89 self.protocol.clone()
90 }
91}
92
93impl<R, W> HybridStdioTransport<RoleServer, R, W>
94where
95 R: Send + AsyncRead + Unpin,
96 W: Send + AsyncWrite + Unpin + 'static,
97{
98 pub fn new_server(read: R, write: W) -> Self {
99 Self::new(read, write)
100 }
101}
102
103impl<Role: ServiceRole, R, W> Transport<Role> for HybridStdioTransport<Role, R, W>
104where
105 R: Send + AsyncRead + Unpin,
106 W: Send + AsyncWrite + Unpin + 'static,
107{
108 type Error = std::io::Error;
109
110 fn send(
111 &mut self,
112 item: TxJsonRpcMessage<Role>,
113 ) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
114 let lock = self.write.clone();
115 async move {
116 let mut write = lock.lock().await;
117 if let Some(ref mut write) = *write {
118 write.send(item).await.map_err(Into::into)
119 } else {
120 Err(std::io::Error::new(
121 std::io::ErrorKind::NotConnected,
122 "Transport is closed",
123 ))
124 }
125 }
126 }
127
128 fn receive(&mut self) -> impl Future<Output = Option<RxJsonRpcMessage<Role>>> + Send {
129 let next = self.read.next();
130 async {
131 next.await.and_then(|result| {
132 result
133 .inspect_err(|error| {
134 tracing::error!("Error reading from stream: {}", error);
135 })
136 .ok()
137 })
138 }
139 }
140
141 async fn close(&mut self) -> Result<(), Self::Error> {
142 let mut write = self.write.lock().await;
143 drop(write.take());
144 Ok(())
145 }
146}
147
148#[derive(Debug, Clone)]
149pub struct HybridJsonRpcMessageCodec<T> {
150 _marker: PhantomData<fn() -> T>,
151 next_index: usize,
152 max_length: usize,
153 is_discarding: bool,
154 protocol: SharedProtocol,
155}
156
157impl<T> HybridJsonRpcMessageCodec<T> {
158 fn new(protocol: SharedProtocol) -> Self {
159 Self {
160 _marker: PhantomData,
161 next_index: 0,
162 max_length: 32 * 1024 * 1024, is_discarding: false,
164 protocol,
165 }
166 }
167}
168
169fn without_carriage_return(s: &[u8]) -> &[u8] {
170 if let Some(&b'\r') = s.last() {
171 &s[..s.len() - 1]
172 } else {
173 s
174 }
175}
176
177fn is_standard_method(method: &str) -> bool {
178 matches!(
179 method,
180 "initialize"
181 | "ping"
182 | "prompts/get"
183 | "prompts/list"
184 | "resources/list"
185 | "resources/read"
186 | "resources/subscribe"
187 | "resources/unsubscribe"
188 | "resources/templates/list"
189 | "tools/call"
190 | "tools/list"
191 | "completion/complete"
192 | "logging/setLevel"
193 | "roots/list"
194 | "sampling/createMessage"
195 ) || is_standard_notification(method)
196}
197
198fn is_standard_notification(method: &str) -> bool {
199 matches!(
200 method,
201 "notifications/cancelled"
202 | "notifications/initialized"
203 | "notifications/message"
204 | "notifications/progress"
205 | "notifications/prompts/list_changed"
206 | "notifications/resources/list_changed"
207 | "notifications/resources/updated"
208 | "notifications/roots/list_changed"
209 | "notifications/tools/list_changed"
210 )
211}
212
213fn should_ignore_notification(json_value: &serde_json::Value, method: &str) -> bool {
214 let is_notification = json_value.get("id").is_none();
215 if is_notification && !is_standard_method(method) {
216 tracing::trace!(
217 "Ignoring non-MCP notification '{}' for compatibility",
218 method
219 );
220 return true;
221 }
222
223 matches!(
224 (
225 method.starts_with("notifications/"),
226 is_standard_notification(method)
227 ),
228 (true, false)
229 )
230}
231
232fn write_method_not_found(id: &serde_json::Value, method: &str, protocol: Option<WireProtocol>) {
236 let response = serde_json::json!({
237 "jsonrpc": "2.0",
238 "id": id,
239 "error": {
240 "code": -32601,
241 "message": format!("Method not found: {method}")
242 }
243 });
244 let body = serde_json::to_string(&response).unwrap_or_default();
245 let mut out = std::io::stdout().lock();
246 use std::io::Write;
247 let _ = if let Some(WireProtocol::ContentLength) = protocol {
248 write!(out, "Content-Length: {}\r\n\r\n{}", body.len(), body)
249 } else {
250 writeln!(out, "{body}")
251 };
252 let _ = std::io::Write::flush(&mut out);
253 tracing::debug!("Replied -32601 MethodNotFound for '{method}'");
254}
255
256pub fn write_method_not_found_pre_init(
270 id: &serde_json::Value,
271 method: &str,
272 protocol: &SharedProtocol,
273) {
274 let response = serde_json::json!({
275 "jsonrpc": "2.0",
276 "id": id,
277 "error": {
278 "code": -32601,
279 "message": concat!(
280 "Method not found. This server supports MCP protocol 2025-11-25 ",
281 "(legacy initialize handshake only)."
282 )
283 }
284 });
285 let body = serde_json::to_string(&response).unwrap_or_default();
286 let mut out = std::io::stdout().lock();
287 use std::io::Write;
288 let _ = if protocol.get() == Some(WireProtocol::ContentLength) {
289 write!(out, "Content-Length: {}\r\n\r\n{}", body.len(), body)
290 } else {
291 writeln!(out, "{body}")
292 };
293 let _ = std::io::Write::flush(&mut out);
294 tracing::debug!("Replied -32601 MethodNotFound for pre-init '{method}'");
295}
296
297fn try_parse_with_compatibility<T: DeserializeOwned>(
298 payload: &[u8],
299 context: &str,
300 protocol: Option<WireProtocol>,
301) -> Result<Option<T>, HybridCodecError> {
302 if let Ok(line_str) = std::str::from_utf8(payload) {
303 match serde_json::from_slice(payload) {
304 Ok(item) => Ok(Some(item)),
305 Err(error) => {
306 if let Ok(json_value) = serde_json::from_str::<serde_json::Value>(line_str)
307 && let Some(method) =
308 json_value.get("method").and_then(serde_json::Value::as_str)
309 {
310 if should_ignore_notification(&json_value, method) {
311 return Ok(None);
312 }
313 if let Some(id) = json_value.get("id") {
315 if !is_standard_method(method) {
316 write_method_not_found(id, method, protocol);
317 return Ok(None);
318 }
319 }
320 }
321
322 tracing::debug!(
323 "Failed to parse message {}: {} | Error: {}",
324 context,
325 line_str,
326 error
327 );
328 Err(HybridCodecError::Serde(error))
329 }
330 }
331 } else {
332 serde_json::from_slice(payload)
333 .map(Some)
334 .map_err(HybridCodecError::Serde)
335 }
336}
337
338#[derive(Debug, Error)]
339pub enum HybridCodecError {
340 #[error("max line length exceeded")]
341 MaxLineLengthExceeded,
342 #[error("missing Content-Length header")]
343 MissingContentLength,
344 #[error("invalid Content-Length value: {0}")]
345 InvalidContentLength(String),
346 #[error("invalid header frame: {0}")]
347 InvalidHeaderFrame(String),
348 #[error("serde error {0}")]
349 Serde(#[from] serde_json::Error),
350 #[error("io error {0}")]
351 Io(#[from] std::io::Error),
352}
353
354impl From<HybridCodecError> for std::io::Error {
355 fn from(value: HybridCodecError) -> Self {
356 match value {
357 HybridCodecError::MaxLineLengthExceeded
358 | HybridCodecError::MissingContentLength
359 | HybridCodecError::InvalidContentLength(_)
360 | HybridCodecError::InvalidHeaderFrame(_) => {
361 std::io::Error::new(std::io::ErrorKind::InvalidData, value)
362 }
363 HybridCodecError::Serde(error) => error.into(),
364 HybridCodecError::Io(error) => error,
365 }
366 }
367}
368
369fn looks_like_content_length_frame(buf: &BytesMut) -> bool {
370 let prefix = &buf[..buf.len().min(32)];
371 prefix
372 .windows(b"content-length".len())
373 .next()
374 .is_some_and(|candidate| candidate.eq_ignore_ascii_case(b"content-length"))
375}
376
377fn find_header_terminator(buf: &BytesMut) -> Option<(usize, usize)> {
378 if let Some(index) = buf.windows(4).position(|window| window == b"\r\n\r\n") {
379 return Some((index, 4));
380 }
381 buf.windows(2)
382 .position(|window| window == b"\n\n")
383 .map(|index| (index, 2))
384}
385
386fn parse_content_length(header: &str) -> Result<usize, HybridCodecError> {
387 for raw_line in header.lines() {
388 let line = raw_line.trim_end_matches('\r');
389 let Some((name, value)) = line.split_once(':') else {
390 continue;
391 };
392 if name.trim().eq_ignore_ascii_case("content-length") {
393 return value
394 .trim()
395 .parse::<usize>()
396 .map_err(|_| HybridCodecError::InvalidContentLength(value.trim().to_string()));
397 }
398 }
399
400 Err(HybridCodecError::MissingContentLength)
401}
402
403impl<T: DeserializeOwned> HybridJsonRpcMessageCodec<T> {
404 fn decode_content_length(&mut self, buf: &mut BytesMut) -> Result<Option<T>, HybridCodecError> {
405 loop {
408 let Some((header_end, delimiter_len)) = find_header_terminator(buf) else {
409 return Ok(None);
410 };
411
412 let header = std::str::from_utf8(&buf[..header_end])
413 .map_err(|error| HybridCodecError::InvalidHeaderFrame(error.to_string()))?;
414 let content_length = parse_content_length(header)?;
415 if content_length > self.max_length {
416 return Err(HybridCodecError::MaxLineLengthExceeded);
417 }
418 let body_start = header_end + delimiter_len;
419 let frame_len = body_start
420 .checked_add(content_length)
421 .ok_or(HybridCodecError::MaxLineLengthExceeded)?;
422 if buf.len() < frame_len {
423 return Ok(None);
424 }
425
426 let frame = buf.split_to(frame_len);
427 let payload = &frame[body_start..];
428 self.protocol.set_if_unset(WireProtocol::ContentLength);
429
430 match try_parse_with_compatibility(
431 payload,
432 "decode_content_length",
433 self.protocol.get(),
434 ) {
435 Ok(Some(item)) => return Ok(Some(item)),
436 Ok(None) => {}
438 Err(err) => {
442 tracing::warn!("skipping malformed Content-Length frame: {err}");
443 }
444 }
445 }
446 }
447
448 fn decode_json_line(&mut self, buf: &mut BytesMut) -> Result<Option<T>, HybridCodecError> {
449 loop {
450 let read_to = std::cmp::min(self.max_length.saturating_add(1), buf.len());
451 let newline_offset = buf[self.next_index..read_to]
452 .iter()
453 .position(|byte| *byte == b'\n');
454
455 match (self.is_discarding, newline_offset) {
456 (true, Some(offset)) => {
457 buf.advance(offset + self.next_index + 1);
458 self.is_discarding = false;
459 self.next_index = 0;
460 }
461 (true, None) => {
462 buf.advance(read_to);
463 self.next_index = 0;
464 if buf.is_empty() {
465 return Ok(None);
466 }
467 }
468 (false, Some(offset)) => {
469 let newline_index = offset + self.next_index;
470 self.next_index = 0;
471 let line = buf.split_to(newline_index + 1);
472 let line = &line[..line.len() - 1];
473 let payload = without_carriage_return(line);
474 self.protocol.set_if_unset(WireProtocol::JsonLine);
475
476 match try_parse_with_compatibility(
477 payload,
478 "decode_json_line",
479 self.protocol.get(),
480 ) {
481 Ok(Some(item)) => return Ok(Some(item)),
482 Ok(None) => {}
484 Err(err) => {
493 tracing::warn!("skipping malformed JSON-RPC line: {err}");
494 }
495 }
496 }
497 (false, None) if buf.len() > self.max_length => {
498 self.is_discarding = true;
499 return Err(HybridCodecError::MaxLineLengthExceeded);
500 }
501 (false, None) => {
502 self.next_index = read_to;
503 return Ok(None);
504 }
505 }
506 }
507 }
508}
509
510impl<T: DeserializeOwned> Decoder for HybridJsonRpcMessageCodec<T> {
511 type Item = T;
512 type Error = HybridCodecError;
513
514 fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<T>, HybridCodecError> {
515 match self.protocol.get() {
516 Some(WireProtocol::ContentLength) => self.decode_content_length(buf),
517 Some(WireProtocol::JsonLine) => self.decode_json_line(buf),
518 None => {
519 if looks_like_content_length_frame(buf) {
520 self.decode_content_length(buf)
521 } else {
522 self.decode_json_line(buf)
523 }
524 }
525 }
526 }
527
528 fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<T>, HybridCodecError> {
529 match self.protocol.get() {
530 Some(WireProtocol::ContentLength) if !buf.is_empty() => self.decode_content_length(buf),
531 _ => Ok(if let Some(frame) = self.decode(buf)? {
532 Some(frame)
533 } else {
534 self.next_index = 0;
535 if buf.is_empty() || buf == &b"\r"[..] {
536 None
537 } else {
538 let line = buf.split_to(buf.len());
539 let payload = without_carriage_return(&line);
540 match try_parse_with_compatibility(payload, "decode_eof", self.protocol.get()) {
543 Ok(item) => item,
544 Err(err) => {
545 tracing::warn!("discarding malformed trailing frame at EOF: {err}");
546 None
547 }
548 }
549 }
550 }),
551 }
552 }
553}
554
555impl<T: Serialize> Encoder<T> for HybridJsonRpcMessageCodec<T> {
556 type Error = HybridCodecError;
557
558 fn encode(&mut self, item: T, buf: &mut BytesMut) -> Result<(), HybridCodecError> {
559 let payload = serde_json::to_vec(&item)?;
560
561 match self.protocol.get().unwrap_or(WireProtocol::ContentLength) {
562 WireProtocol::ContentLength => {
563 buf.extend_from_slice(
564 format!("Content-Length: {}\r\n\r\n", payload.len()).as_bytes(),
565 );
566 buf.extend_from_slice(&payload);
567 }
568 WireProtocol::JsonLine => {
569 buf.extend_from_slice(&payload);
570 buf.put_u8(b'\n');
571 }
572 }
573
574 Ok(())
575 }
576}
577
578#[cfg(test)]
579mod tests {
580 use super::*;
581
582 use tokio_util::bytes::BytesMut;
583
584 fn sample_message() -> serde_json::Value {
585 serde_json::json!({
586 "jsonrpc": "2.0",
587 "id": 1,
588 "method": "initialize",
589 "params": {
590 "protocolVersion": "2024-11-05",
591 "capabilities": {},
592 "clientInfo": {
593 "name": "probe",
594 "version": "0.0.0"
595 }
596 }
597 })
598 }
599
600 #[test]
601 fn decodes_json_line_and_marks_protocol() {
602 let protocol = SharedProtocol::new();
603 let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol.clone());
604 let payload = serde_json::to_vec(&sample_message()).unwrap();
605 let mut buf = BytesMut::from(&payload[..]);
606 buf.put_u8(b'\n');
607
608 let item = codec.decode(&mut buf).unwrap();
609 assert!(item.is_some());
610 assert_eq!(protocol.get(), Some(WireProtocol::JsonLine));
611 }
612
613 #[test]
614 fn decodes_content_length_and_marks_protocol() {
615 let protocol = SharedProtocol::new();
616 let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol.clone());
617 let payload = serde_json::to_vec(&sample_message()).unwrap();
618 let mut frame = BytesMut::new();
619 frame.extend_from_slice(format!("Content-Length: {}\r\n\r\n", payload.len()).as_bytes());
620 frame.extend_from_slice(&payload);
621
622 let item = codec.decode(&mut frame).unwrap();
623 assert!(item.is_some());
624 assert_eq!(protocol.get(), Some(WireProtocol::ContentLength));
625 }
626
627 #[test]
628 fn encodes_using_content_length_when_protocol_is_detected() {
629 let protocol = SharedProtocol::new();
630 protocol.set_if_unset(WireProtocol::ContentLength);
631 let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol);
632 let mut buf = BytesMut::new();
633 codec
634 .encode(
635 serde_json::json!({"jsonrpc":"2.0","id":1,"result":{"ok":true}}),
636 &mut buf,
637 )
638 .unwrap();
639
640 assert!(
641 std::str::from_utf8(&buf)
642 .unwrap()
643 .starts_with("Content-Length: ")
644 );
645 }
646
647 #[test]
648 fn decode_skips_malformed_json_line_and_recovers() {
649 let protocol = SharedProtocol::new();
653 let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol);
654 let good = serde_json::to_vec(&sample_message()).unwrap();
655 let mut buf = BytesMut::new();
656 buf.extend_from_slice(b"{ this is : not valid json }\n");
657 buf.extend_from_slice(&good);
658 buf.put_u8(b'\n');
659
660 let item = codec
661 .decode(&mut buf)
662 .expect("a malformed line must not be a hard transport error");
663 assert!(item.is_some(), "valid frame after a bad line must decode");
664 }
665
666 #[test]
667 fn decode_malformed_json_line_alone_is_not_a_transport_error() {
668 let protocol = SharedProtocol::new();
670 let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol);
671 let mut buf = BytesMut::from(&b"{ not valid json }\n"[..]);
672
673 let item = codec
674 .decode(&mut buf)
675 .expect("a malformed line must not be a hard transport error");
676 assert!(item.is_none());
677 }
678
679 #[test]
680 fn decode_skips_malformed_content_length_frame_and_recovers() {
681 let protocol = SharedProtocol::new();
684 protocol.set_if_unset(WireProtocol::ContentLength);
685 let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol);
686
687 let bad_body = b"{ not json }";
688 let good_body = serde_json::to_vec(&sample_message()).unwrap();
689 let mut buf = BytesMut::new();
690 buf.extend_from_slice(format!("Content-Length: {}\r\n\r\n", bad_body.len()).as_bytes());
691 buf.extend_from_slice(bad_body);
692 buf.extend_from_slice(format!("Content-Length: {}\r\n\r\n", good_body.len()).as_bytes());
693 buf.extend_from_slice(&good_body);
694
695 let item = codec
696 .decode(&mut buf)
697 .expect("a malformed CL frame must not be a hard transport error");
698 assert!(item.is_some(), "valid CL frame after a bad one must decode");
699 }
700}