lean_ctx/
mcp_stdio.rs

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)]
30struct 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}
62
63impl<Role: ServiceRole, R, W> HybridStdioTransport<Role, R, W>
64where
65    R: Send + AsyncRead + Unpin,
66    W: Send + AsyncWrite + Unpin + 'static,
67{
68    pub fn new(read: R, write: W) -> Self {
69        let protocol = SharedProtocol::new();
70        let read = FramedRead::new(
71            read,
72            HybridJsonRpcMessageCodec::<RxJsonRpcMessage<Role>>::new(protocol.clone()),
73        );
74        let write = Arc::new(AsyncMutex::new(Some(FramedWrite::new(
75            write,
76            HybridJsonRpcMessageCodec::<TxJsonRpcMessage<Role>>::new(protocol),
77        ))));
78        Self { read, write }
79    }
80}
81
82impl<R, W> HybridStdioTransport<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 HybridStdioTransport<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>>> + Send {
118        let next = self.read.next();
119        async {
120            next.await.and_then(|result| {
121                result
122                    .inspect_err(|error| {
123                        tracing::error!("Error reading from stream: {}", error);
124                    })
125                    .ok()
126            })
127        }
128    }
129
130    async fn close(&mut self) -> Result<(), Self::Error> {
131        let mut write = self.write.lock().await;
132        drop(write.take());
133        Ok(())
134    }
135}
136
137#[derive(Debug, Clone)]
138pub struct HybridJsonRpcMessageCodec<T> {
139    _marker: PhantomData<fn() -> T>,
140    next_index: usize,
141    max_length: usize,
142    is_discarding: bool,
143    protocol: SharedProtocol,
144}
145
146impl<T> HybridJsonRpcMessageCodec<T> {
147    fn new(protocol: SharedProtocol) -> Self {
148        Self {
149            _marker: PhantomData,
150            next_index: 0,
151            max_length: 32 * 1024 * 1024, // 32 MiB — prevents OOM from oversized messages
152            is_discarding: false,
153            protocol,
154        }
155    }
156}
157
158fn without_carriage_return(s: &[u8]) -> &[u8] {
159    if let Some(&b'\r') = s.last() {
160        &s[..s.len() - 1]
161    } else {
162        s
163    }
164}
165
166fn is_standard_method(method: &str) -> bool {
167    matches!(
168        method,
169        "initialize"
170            | "ping"
171            | "prompts/get"
172            | "prompts/list"
173            | "resources/list"
174            | "resources/read"
175            | "resources/subscribe"
176            | "resources/unsubscribe"
177            | "resources/templates/list"
178            | "tools/call"
179            | "tools/list"
180            | "completion/complete"
181            | "logging/setLevel"
182            | "roots/list"
183            | "sampling/createMessage"
184    ) || is_standard_notification(method)
185}
186
187fn is_standard_notification(method: &str) -> bool {
188    matches!(
189        method,
190        "notifications/cancelled"
191            | "notifications/initialized"
192            | "notifications/message"
193            | "notifications/progress"
194            | "notifications/prompts/list_changed"
195            | "notifications/resources/list_changed"
196            | "notifications/resources/updated"
197            | "notifications/roots/list_changed"
198            | "notifications/tools/list_changed"
199    )
200}
201
202fn should_ignore_notification(json_value: &serde_json::Value, method: &str) -> bool {
203    let is_notification = json_value.get("id").is_none();
204    if is_notification && !is_standard_method(method) {
205        tracing::trace!(
206            "Ignoring non-MCP notification '{}' for compatibility",
207            method
208        );
209        return true;
210    }
211
212    matches!(
213        (
214            method.starts_with("notifications/"),
215            is_standard_notification(method)
216        ),
217        (true, false)
218    )
219}
220
221/// GH #1434: Reply -32601 for unknown JSON-RPC requests so compliant
222/// clients (e.g. Go SDK >= 1.7 sending `server/discover`) fall back
223/// gracefully instead of seeing EOF.
224fn write_method_not_found(id: &serde_json::Value, method: &str, protocol: Option<WireProtocol>) {
225    let response = serde_json::json!({
226        "jsonrpc": "2.0",
227        "id": id,
228        "error": {
229            "code": -32601,
230            "message": format!("Method not found: {method}")
231        }
232    });
233    let body = serde_json::to_string(&response).unwrap_or_default();
234    let mut out = std::io::stdout().lock();
235    use std::io::Write;
236    let _ = if let Some(WireProtocol::ContentLength) = protocol {
237        write!(out, "Content-Length: {}\r\n\r\n{}", body.len(), body)
238    } else {
239        writeln!(out, "{body}")
240    };
241    let _ = std::io::Write::flush(&mut out);
242    tracing::debug!("Replied -32601 MethodNotFound for '{method}'");
243}
244
245fn try_parse_with_compatibility<T: DeserializeOwned>(
246    payload: &[u8],
247    context: &str,
248    protocol: Option<WireProtocol>,
249) -> Result<Option<T>, HybridCodecError> {
250    if let Ok(line_str) = std::str::from_utf8(payload) {
251        match serde_json::from_slice(payload) {
252            Ok(item) => Ok(Some(item)),
253            Err(error) => {
254                if let Ok(json_value) = serde_json::from_str::<serde_json::Value>(line_str)
255                    && let Some(method) =
256                        json_value.get("method").and_then(serde_json::Value::as_str)
257                {
258                    if should_ignore_notification(&json_value, method) {
259                        return Ok(None);
260                    }
261                    // GH #1434: unknown request (has id) → reply MethodNotFound
262                    if let Some(id) = json_value.get("id") {
263                        if !is_standard_method(method) {
264                            write_method_not_found(id, method, protocol);
265                            return Ok(None);
266                        }
267                    }
268                }
269
270                tracing::debug!(
271                    "Failed to parse message {}: {} | Error: {}",
272                    context,
273                    line_str,
274                    error
275                );
276                Err(HybridCodecError::Serde(error))
277            }
278        }
279    } else {
280        serde_json::from_slice(payload)
281            .map(Some)
282            .map_err(HybridCodecError::Serde)
283    }
284}
285
286#[derive(Debug, Error)]
287pub enum HybridCodecError {
288    #[error("max line length exceeded")]
289    MaxLineLengthExceeded,
290    #[error("missing Content-Length header")]
291    MissingContentLength,
292    #[error("invalid Content-Length value: {0}")]
293    InvalidContentLength(String),
294    #[error("invalid header frame: {0}")]
295    InvalidHeaderFrame(String),
296    #[error("serde error {0}")]
297    Serde(#[from] serde_json::Error),
298    #[error("io error {0}")]
299    Io(#[from] std::io::Error),
300}
301
302impl From<HybridCodecError> for std::io::Error {
303    fn from(value: HybridCodecError) -> Self {
304        match value {
305            HybridCodecError::MaxLineLengthExceeded
306            | HybridCodecError::MissingContentLength
307            | HybridCodecError::InvalidContentLength(_)
308            | HybridCodecError::InvalidHeaderFrame(_) => {
309                std::io::Error::new(std::io::ErrorKind::InvalidData, value)
310            }
311            HybridCodecError::Serde(error) => error.into(),
312            HybridCodecError::Io(error) => error,
313        }
314    }
315}
316
317fn looks_like_content_length_frame(buf: &BytesMut) -> bool {
318    let prefix = &buf[..buf.len().min(32)];
319    prefix
320        .windows(b"content-length".len())
321        .next()
322        .is_some_and(|candidate| candidate.eq_ignore_ascii_case(b"content-length"))
323}
324
325fn find_header_terminator(buf: &BytesMut) -> Option<(usize, usize)> {
326    if let Some(index) = buf.windows(4).position(|window| window == b"\r\n\r\n") {
327        return Some((index, 4));
328    }
329    buf.windows(2)
330        .position(|window| window == b"\n\n")
331        .map(|index| (index, 2))
332}
333
334fn parse_content_length(header: &str) -> Result<usize, HybridCodecError> {
335    for raw_line in header.lines() {
336        let line = raw_line.trim_end_matches('\r');
337        let Some((name, value)) = line.split_once(':') else {
338            continue;
339        };
340        if name.trim().eq_ignore_ascii_case("content-length") {
341            return value
342                .trim()
343                .parse::<usize>()
344                .map_err(|_| HybridCodecError::InvalidContentLength(value.trim().to_string()));
345        }
346    }
347
348    Err(HybridCodecError::MissingContentLength)
349}
350
351impl<T: DeserializeOwned> HybridJsonRpcMessageCodec<T> {
352    fn decode_content_length(&mut self, buf: &mut BytesMut) -> Result<Option<T>, HybridCodecError> {
353        // Loop so a skipped (malformed) frame body resyncs onto the next framed
354        // message already buffered, instead of stalling until more bytes arrive.
355        loop {
356            let Some((header_end, delimiter_len)) = find_header_terminator(buf) else {
357                return Ok(None);
358            };
359
360            let header = std::str::from_utf8(&buf[..header_end])
361                .map_err(|error| HybridCodecError::InvalidHeaderFrame(error.to_string()))?;
362            let content_length = parse_content_length(header)?;
363            if content_length > self.max_length {
364                return Err(HybridCodecError::MaxLineLengthExceeded);
365            }
366            let body_start = header_end + delimiter_len;
367            let frame_len = body_start
368                .checked_add(content_length)
369                .ok_or(HybridCodecError::MaxLineLengthExceeded)?;
370            if buf.len() < frame_len {
371                return Ok(None);
372            }
373
374            let frame = buf.split_to(frame_len);
375            let payload = &frame[body_start..];
376            self.protocol.set_if_unset(WireProtocol::ContentLength);
377
378            match try_parse_with_compatibility(
379                payload,
380                "decode_content_length",
381                self.protocol.get(),
382            ) {
383                Ok(Some(item)) => return Ok(Some(item)),
384                // An ignored notification — fall through and scan the next frame.
385                Ok(None) => {}
386                // Skip a malformed body instead of fusing the transport (see
387                // `decode_json_line` / #453). The Content-Length framing stayed
388                // intact, so the next frame in the buffer is still recoverable.
389                Err(err) => {
390                    tracing::warn!("skipping malformed Content-Length frame: {err}");
391                }
392            }
393        }
394    }
395
396    fn decode_json_line(&mut self, buf: &mut BytesMut) -> Result<Option<T>, HybridCodecError> {
397        loop {
398            let read_to = std::cmp::min(self.max_length.saturating_add(1), buf.len());
399            let newline_offset = buf[self.next_index..read_to]
400                .iter()
401                .position(|byte| *byte == b'\n');
402
403            match (self.is_discarding, newline_offset) {
404                (true, Some(offset)) => {
405                    buf.advance(offset + self.next_index + 1);
406                    self.is_discarding = false;
407                    self.next_index = 0;
408                }
409                (true, None) => {
410                    buf.advance(read_to);
411                    self.next_index = 0;
412                    if buf.is_empty() {
413                        return Ok(None);
414                    }
415                }
416                (false, Some(offset)) => {
417                    let newline_index = offset + self.next_index;
418                    self.next_index = 0;
419                    let line = buf.split_to(newline_index + 1);
420                    let line = &line[..line.len() - 1];
421                    let payload = without_carriage_return(line);
422                    self.protocol.set_if_unset(WireProtocol::JsonLine);
423
424                    match try_parse_with_compatibility(
425                        payload,
426                        "decode_json_line",
427                        self.protocol.get(),
428                    ) {
429                        Ok(Some(item)) => return Ok(Some(item)),
430                        // An ignored notification — keep scanning for a real frame.
431                        Ok(None) => {}
432                        // A single malformed line must NOT tear down the transport.
433                        // Returning Err here makes `FramedRead` fuse the stream
434                        // (`has_errored`), so the next poll yields `None`; rmcp then
435                        // exits with `QuitReason::Closed`, the agent respawns us, and
436                        // the fresh process pays another index build — the respawn /
437                        // idle-CPU churn behind #453. The bad line is already consumed
438                        // (`split_to` above), so we just skip it and resync on the
439                        // next newline-delimited frame.
440                        Err(err) => {
441                            tracing::warn!("skipping malformed JSON-RPC line: {err}");
442                        }
443                    }
444                }
445                (false, None) if buf.len() > self.max_length => {
446                    self.is_discarding = true;
447                    return Err(HybridCodecError::MaxLineLengthExceeded);
448                }
449                (false, None) => {
450                    self.next_index = read_to;
451                    return Ok(None);
452                }
453            }
454        }
455    }
456}
457
458impl<T: DeserializeOwned> Decoder for HybridJsonRpcMessageCodec<T> {
459    type Item = T;
460    type Error = HybridCodecError;
461
462    fn decode(&mut self, buf: &mut BytesMut) -> Result<Option<T>, HybridCodecError> {
463        match self.protocol.get() {
464            Some(WireProtocol::ContentLength) => self.decode_content_length(buf),
465            Some(WireProtocol::JsonLine) => self.decode_json_line(buf),
466            None => {
467                if looks_like_content_length_frame(buf) {
468                    self.decode_content_length(buf)
469                } else {
470                    self.decode_json_line(buf)
471                }
472            }
473        }
474    }
475
476    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<T>, HybridCodecError> {
477        match self.protocol.get() {
478            Some(WireProtocol::ContentLength) if !buf.is_empty() => self.decode_content_length(buf),
479            _ => Ok(if let Some(frame) = self.decode(buf)? {
480                Some(frame)
481            } else {
482                self.next_index = 0;
483                if buf.is_empty() || buf == &b"\r"[..] {
484                    None
485                } else {
486                    let line = buf.split_to(buf.len());
487                    let payload = without_carriage_return(&line);
488                    // At true stream end a malformed trailing frame is discarded
489                    // (clean close) rather than surfaced as a transport error.
490                    match try_parse_with_compatibility(payload, "decode_eof", self.protocol.get()) {
491                        Ok(item) => item,
492                        Err(err) => {
493                            tracing::warn!("discarding malformed trailing frame at EOF: {err}");
494                            None
495                        }
496                    }
497                }
498            }),
499        }
500    }
501}
502
503impl<T: Serialize> Encoder<T> for HybridJsonRpcMessageCodec<T> {
504    type Error = HybridCodecError;
505
506    fn encode(&mut self, item: T, buf: &mut BytesMut) -> Result<(), HybridCodecError> {
507        let payload = serde_json::to_vec(&item)?;
508
509        match self.protocol.get().unwrap_or(WireProtocol::ContentLength) {
510            WireProtocol::ContentLength => {
511                buf.extend_from_slice(
512                    format!("Content-Length: {}\r\n\r\n", payload.len()).as_bytes(),
513                );
514                buf.extend_from_slice(&payload);
515            }
516            WireProtocol::JsonLine => {
517                buf.extend_from_slice(&payload);
518                buf.put_u8(b'\n');
519            }
520        }
521
522        Ok(())
523    }
524}
525
526#[cfg(test)]
527mod tests {
528    use super::*;
529
530    use tokio_util::bytes::BytesMut;
531
532    fn sample_message() -> serde_json::Value {
533        serde_json::json!({
534            "jsonrpc": "2.0",
535            "id": 1,
536            "method": "initialize",
537            "params": {
538                "protocolVersion": "2024-11-05",
539                "capabilities": {},
540                "clientInfo": {
541                    "name": "probe",
542                    "version": "0.0.0"
543                }
544            }
545        })
546    }
547
548    #[test]
549    fn decodes_json_line_and_marks_protocol() {
550        let protocol = SharedProtocol::new();
551        let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol.clone());
552        let payload = serde_json::to_vec(&sample_message()).unwrap();
553        let mut buf = BytesMut::from(&payload[..]);
554        buf.put_u8(b'\n');
555
556        let item = codec.decode(&mut buf).unwrap();
557        assert!(item.is_some());
558        assert_eq!(protocol.get(), Some(WireProtocol::JsonLine));
559    }
560
561    #[test]
562    fn decodes_content_length_and_marks_protocol() {
563        let protocol = SharedProtocol::new();
564        let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol.clone());
565        let payload = serde_json::to_vec(&sample_message()).unwrap();
566        let mut frame = BytesMut::new();
567        frame.extend_from_slice(format!("Content-Length: {}\r\n\r\n", payload.len()).as_bytes());
568        frame.extend_from_slice(&payload);
569
570        let item = codec.decode(&mut frame).unwrap();
571        assert!(item.is_some());
572        assert_eq!(protocol.get(), Some(WireProtocol::ContentLength));
573    }
574
575    #[test]
576    fn encodes_using_content_length_when_protocol_is_detected() {
577        let protocol = SharedProtocol::new();
578        protocol.set_if_unset(WireProtocol::ContentLength);
579        let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol);
580        let mut buf = BytesMut::new();
581        codec
582            .encode(
583                serde_json::json!({"jsonrpc":"2.0","id":1,"result":{"ok":true}}),
584                &mut buf,
585            )
586            .unwrap();
587
588        assert!(
589            std::str::from_utf8(&buf)
590                .unwrap()
591                .starts_with("Content-Length: ")
592        );
593    }
594
595    #[test]
596    fn decode_skips_malformed_json_line_and_recovers() {
597        // A bad line followed by a valid frame: the codec must skip the bad one
598        // and return the valid frame in the same decode pass — never an Err that
599        // would fuse the transport and trigger an MCP-server respawn (#453).
600        let protocol = SharedProtocol::new();
601        let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol);
602        let good = serde_json::to_vec(&sample_message()).unwrap();
603        let mut buf = BytesMut::new();
604        buf.extend_from_slice(b"{ this is : not valid json }\n");
605        buf.extend_from_slice(&good);
606        buf.put_u8(b'\n');
607
608        let item = codec
609            .decode(&mut buf)
610            .expect("a malformed line must not be a hard transport error");
611        assert!(item.is_some(), "valid frame after a bad line must decode");
612    }
613
614    #[test]
615    fn decode_malformed_json_line_alone_is_not_a_transport_error() {
616        // A lone malformed line yields Ok(None) (stream stays alive), not Err.
617        let protocol = SharedProtocol::new();
618        let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol);
619        let mut buf = BytesMut::from(&b"{ not valid json }\n"[..]);
620
621        let item = codec
622            .decode(&mut buf)
623            .expect("a malformed line must not be a hard transport error");
624        assert!(item.is_none());
625    }
626
627    #[test]
628    fn decode_skips_malformed_content_length_frame_and_recovers() {
629        // Same guarantee for the Content-Length wire protocol: a bad body is
630        // skipped and the next well-framed message is still delivered.
631        let protocol = SharedProtocol::new();
632        protocol.set_if_unset(WireProtocol::ContentLength);
633        let mut codec = HybridJsonRpcMessageCodec::<serde_json::Value>::new(protocol);
634
635        let bad_body = b"{ not json }";
636        let good_body = serde_json::to_vec(&sample_message()).unwrap();
637        let mut buf = BytesMut::new();
638        buf.extend_from_slice(format!("Content-Length: {}\r\n\r\n", bad_body.len()).as_bytes());
639        buf.extend_from_slice(bad_body);
640        buf.extend_from_slice(format!("Content-Length: {}\r\n\r\n", good_body.len()).as_bytes());
641        buf.extend_from_slice(&good_body);
642
643        let item = codec
644            .decode(&mut buf)
645            .expect("a malformed CL frame must not be a hard transport error");
646        assert!(item.is_some(), "valid CL frame after a bad one must decode");
647    }
648}
lean_ctx/mcp_stdio.rs

lean_ctx/
mcp_stdio.rs
