futu_mcp/
transport.rs

1//! Resilient stdio transport for MCP server (v1.4.90 P0-A).
2//!
3//! ## Why this exists
4//!
5//! `rmcp::transport::stdio()` (i.e. the default `(Stdin, Stdout)` adapter via
6//! `AsyncRwTransport` + `JsonRpcMessageCodec`) treats *any* JSON parse error
7//! as a fatal stream error. Concretely:
8//!
9//! 1. `JsonRpcMessageCodec::decode()` returns `Err(JsonRpcMessageCodecError::Serde(_))`
10//!    when a line is malformed (e.g. `{"price": Infinity}` — JSON spec forbids
11//!    `Infinity` / `NaN` literals, but LLM clients emit them occasionally).
12//! 2. `FramedRead` yields `Some(Err(_))`.
13//! 3. `AsyncRwTransport::receive()` does `next.await.and_then(|e| e.ok())` —
14//!    converting `Err` to `None`.
15//! 4. The rmcp service loop interprets `None` as "input stream closed" and
16//!    breaks with `QuitReason::Closed`, terminating the entire MCP server.
17//!
18//! Result: a *single* malformed JSON line silently kills the whole server,
19//! disconnecting every client (multi-version sweep proven across v1.4.47 →
20//! v1.4.86 — 11 versions all vulnerable).
21//!
22//! Per JSON-RPC 2.0 §5.1, the correct behavior is to return a `-32700 Parse
23//! error` response and keep the connection alive. This module implements that
24//! behavior as a drop-in replacement for `rmcp::transport::stdio()`.
25//!
26//! ## Design
27//!
28//! - `ResilientStdioTransport` implements `rmcp::transport::Transport<RoleServer>`.
29//! - A background **reader task** owns stdin, reads newline-delimited frames,
30//!   and parses each into `RxJsonRpcMessage<RoleServer>`. Successful parses go
31//!   into an inbound mpsc channel for `receive()`. Parse failures cause a
32//!   synthetic `JsonRpcError(-32700)` to be enqueued onto the **outbound**
33//!   channel directly (bypassing `receive()` so the service never sees an
34//!   error event), and the loop continues.
35//! - A background **writer task** owns stdout and drains the bounded outbound
36//!   channel, serialising messages as one-line JSON each.
37//! - `send()` enqueues onto the outbound channel; `receive()` polls the
38//!   inbound channel; `close()` drops the senders so both tasks exit cleanly.
39//!
40//! ## What this is NOT
41//!
42//! This is a stdio-only fix. The HTTP transport (`StreamableHttpService`)
43//! has its own per-request HTTP body parsing — a malformed request there
44//! returns 4xx without killing the server, so it's not affected by this bug.
45//! Future work: upstream PR to rmcp so all transports share resilient parsing.
46
47use std::{collections::HashSet, io, sync::Arc, time::Duration};
48
49use rmcp::RoleServer;
50use rmcp::model::{ErrorCode, ErrorData, JsonRpcError, NumberOrString};
51use rmcp::service::{RxJsonRpcMessage, TxJsonRpcMessage};
52use rmcp::transport::Transport;
53use serde_json::Value;
54use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
55use tokio::sync::{Mutex, Notify, mpsc};
56
57/// Bound for the inbound channel. 64 frames buffered is plenty — the rmcp
58/// service loop drains promptly. If the channel ever fills up, `receive()` /
59/// the reader task will simply backpressure on stdin, which is fine.
60const INBOUND_BUFFER: usize = 64;
61
62/// Max bytes accepted for one inbound stdio JSONL frame.
63///
64/// This caps only client -> server request lines. Server -> client responses
65/// are still written normally through stdout. 10 MiB matches the REST strict
66/// JSON body order of magnitude and leaves plenty of room for large tool args
67/// while preventing an unbounded `Vec` growth on malformed/no-newline input.
68const MAX_STDIO_JSONL_LINE_BYTES: usize = 10 * 1024 * 1024;
69
70/// One-shot stdio clients commonly write initialize + tool call JSONL and then
71/// close stdin. rmcp treats stdin EOF as service shutdown and starts draining
72/// in-flight responses with its own short timeout; slow backend tools can lose
73/// their result before the response is written. Keep EOF hidden from rmcp until
74/// requests already accepted by this transport either respond or this guard
75/// expires. Normal persistent MCP clients never hit this path.
76const EOF_PENDING_DRAIN_GRACE: Duration = Duration::from_secs(15);
77
78/// Bound for server -> client JSON-RPC messages.
79///
80/// Tool results can be large. Keeping stdout writes behind a bounded channel
81/// makes slow readers apply backpressure instead of accumulating unbounded
82/// response bodies in memory.
83const OUTBOUND_BUFFER: usize = 64;
84
85type OutboundTx = mpsc::Sender<TxJsonRpcMessage<RoleServer>>;
86type OutboundRx = mpsc::Receiver<TxJsonRpcMessage<RoleServer>>;
87type PendingRequestIds = Arc<PendingRequests>;
88
89#[derive(Clone)]
90pub(crate) struct PendingResponseDrain {
91    pending: PendingRequestIds,
92}
93
94impl PendingResponseDrain {
95    pub(crate) async fn wait_pending_responses(&self) -> bool {
96        self.pending
97            .wait_empty_or_timeout(EOF_PENDING_DRAIN_GRACE)
98            .await
99    }
100}
101
102#[derive(Default)]
103struct PendingRequests {
104    ids: Mutex<HashSet<NumberOrString>>,
105    notify: Notify,
106}
107
108impl PendingRequests {
109    async fn insert(&self, id: NumberOrString) {
110        self.ids.lock().await.insert(id);
111    }
112
113    async fn remove(&self, id: &NumberOrString) {
114        let mut ids = self.ids.lock().await;
115        if ids.remove(id) {
116            self.notify.notify_waiters();
117        }
118    }
119
120    async fn wait_empty_or_timeout(&self, timeout: Duration) -> bool {
121        let deadline = tokio::time::Instant::now() + timeout;
122        loop {
123            let notified = self.notify.notified();
124            tokio::pin!(notified);
125            notified.as_mut().enable();
126            if self.ids.lock().await.is_empty() {
127                return true;
128            }
129            let now = tokio::time::Instant::now();
130            if now >= deadline {
131                return false;
132            }
133            tokio::select! {
134                _ = &mut notified => {}
135                _ = tokio::time::sleep_until(deadline) => return false,
136            }
137        }
138    }
139}
140
141/// Resilient stdio transport — see module docs.
142pub struct ResilientStdioTransport {
143    inbound_rx: mpsc::Receiver<RxJsonRpcMessage<RoleServer>>,
144    /// Wrapped in `Arc<Mutex<>>` so `send()` can return a `'static` future
145    /// per the `Transport` trait contract.
146    outbound_tx: Arc<Mutex<Option<OutboundTx>>>,
147    pending: PendingRequestIds,
148}
149
150impl ResilientStdioTransport {
151    /// Spawn reader + writer tasks bound to the supplied I/O handles.
152    ///
153    /// Generic over `R` / `W` so tests can inject in-memory pipes; production
154    /// callers use [`resilient_stdio()`].
155    pub fn new<R, W>(read: R, write: W) -> Self
156    where
157        R: AsyncRead + Send + Unpin + 'static,
158        W: AsyncWrite + Send + Unpin + 'static,
159    {
160        let (inbound_tx, inbound_rx) =
161            mpsc::channel::<RxJsonRpcMessage<RoleServer>>(INBOUND_BUFFER);
162        let (outbound_tx, outbound_rx) =
163            mpsc::channel::<TxJsonRpcMessage<RoleServer>>(OUTBOUND_BUFFER);
164        let pending = Arc::new(PendingRequests::default());
165
166        // Reader task — owns stdin, parses lines, recovers from parse errors.
167        let outbound_tx_for_reader = outbound_tx.clone();
168        tokio::spawn(reader_task(
169            read,
170            inbound_tx,
171            outbound_tx_for_reader,
172            Arc::clone(&pending),
173        ));
174
175        // Writer task — owns stdout, drains outbound queue.
176        tokio::spawn(writer_task(write, outbound_rx, Arc::clone(&pending)));
177
178        Self {
179            inbound_rx,
180            outbound_tx: Arc::new(Mutex::new(Some(outbound_tx))),
181            pending,
182        }
183    }
184
185    pub(crate) fn pending_response_drain(&self) -> PendingResponseDrain {
186        PendingResponseDrain {
187            pending: Arc::clone(&self.pending),
188        }
189    }
190}
191
192/// Drop-in replacement for `rmcp::transport::stdio()`. Returns a transport
193/// that survives malformed JSON instead of `exit(0)`-ing.
194pub fn resilient_stdio() -> ResilientStdioTransport {
195    ResilientStdioTransport::new(tokio::io::stdin(), tokio::io::stdout())
196}
197
198async fn enqueue_parse_error_response(
199    outbound_tx: &OutboundTx,
200    err_msg: JsonRpcError,
201    context: &'static str,
202) {
203    if outbound_tx
204        .send(rmcp::model::JsonRpcMessage::Error(err_msg))
205        .await
206        .is_err()
207    {
208        tracing::debug!(
209            context,
210            "parse error response dropped because writer is gone"
211        );
212    }
213}
214
215async fn enqueue_message_too_large_response(outbound_tx: &OutboundTx) {
216    let err_msg = JsonRpcError::new(
217        None,
218        ErrorData::new(
219            ErrorCode::INVALID_REQUEST,
220            format!("stdio JSON-RPC message too large: exceeds {MAX_STDIO_JSONL_LINE_BYTES} bytes"),
221            None,
222        ),
223    );
224    enqueue_parse_error_response(outbound_tx, err_msg, "message_too_large").await;
225}
226
227#[derive(Debug)]
228pub enum TransportError {
229    Closed,
230}
231
232impl std::fmt::Display for TransportError {
233    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
234        match self {
235            Self::Closed => f.write_str("transport closed"),
236        }
237    }
238}
239
240impl std::error::Error for TransportError {}
241
242impl Transport<RoleServer> for ResilientStdioTransport {
243    type Error = TransportError;
244
245    fn send(
246        &mut self,
247        item: TxJsonRpcMessage<RoleServer>,
248    ) -> impl Future<Output = Result<(), Self::Error>> + Send + 'static {
249        let lock = self.outbound_tx.clone();
250        async move {
251            let tx = { lock.lock().await.as_ref().cloned() };
252            match tx {
253                Some(tx) => tx.send(item).await.map_err(|_| TransportError::Closed),
254                None => Err(TransportError::Closed),
255            }
256        }
257    }
258
259    async fn receive(&mut self) -> Option<RxJsonRpcMessage<RoleServer>> {
260        self.inbound_rx.recv().await
261    }
262
263    async fn close(&mut self) -> Result<(), Self::Error> {
264        let mut guard = self.outbound_tx.lock().await;
265        // Dropping the sender signals the writer task to exit. The reader
266        // task exits naturally on stdin EOF (or when its inbound_tx half is
267        // dropped, which happens when this struct drops).
268        guard.take();
269        self.inbound_rx.close();
270        Ok(())
271    }
272}
273
274// `ResilientStdioTransport: Transport<RoleServer>` automatically gives us
275// `IntoTransport<RoleServer, TransportError, TransportAdapterIdentity>` via
276// rmcp's blanket impl, so no explicit `IntoTransport` impl is needed here.
277
278// ---------------------------------------------------------------------------
279// Reader / writer tasks
280// ---------------------------------------------------------------------------
281
282#[derive(Debug, Clone, Copy, PartialEq, Eq)]
283enum LimitedLineRead {
284    Eof,
285    Line,
286    TooLarge,
287}
288
289async fn read_until_newline_limited<R>(
290    reader: &mut R,
291    buf: &mut Vec<u8>,
292    max_bytes: usize,
293) -> io::Result<LimitedLineRead>
294where
295    R: AsyncBufRead + Unpin,
296{
297    let mut saw_any = false;
298    loop {
299        let available = reader.fill_buf().await?;
300        if available.is_empty() {
301            return if saw_any {
302                Ok(LimitedLineRead::Line)
303            } else {
304                Ok(LimitedLineRead::Eof)
305            };
306        }
307
308        let newline_pos = available.iter().position(|b| *b == b'\n');
309        let found_newline = newline_pos.is_some();
310        let take = newline_pos.map_or(available.len(), |idx| idx + 1);
311
312        if buf.len().saturating_add(take) > max_bytes {
313            reader.consume(take);
314            if !found_newline {
315                discard_until_newline(reader).await?;
316            }
317            return Ok(LimitedLineRead::TooLarge);
318        }
319
320        buf.extend_from_slice(&available[..take]);
321        reader.consume(take);
322        saw_any = true;
323
324        if found_newline {
325            return Ok(LimitedLineRead::Line);
326        }
327    }
328}
329
330async fn discard_until_newline<R>(reader: &mut R) -> io::Result<()>
331where
332    R: AsyncBufRead + Unpin,
333{
334    loop {
335        let available = reader.fill_buf().await?;
336        if available.is_empty() {
337            return Ok(());
338        }
339
340        let newline_pos = available.iter().position(|b| *b == b'\n');
341        let found_newline = newline_pos.is_some();
342        let take = newline_pos.map_or(available.len(), |idx| idx + 1);
343        reader.consume(take);
344
345        if found_newline {
346            return Ok(());
347        }
348    }
349}
350
351async fn reader_task<R>(
352    read: R,
353    inbound_tx: mpsc::Sender<RxJsonRpcMessage<RoleServer>>,
354    outbound_tx: OutboundTx,
355    pending: PendingRequestIds,
356) where
357    R: AsyncRead + Send + Unpin + 'static,
358{
359    let mut reader = BufReader::new(read);
360    // v1.4.93 P0-3 (BUG-003): read raw bytes instead of String to gracefully
361    // handle UTF-8 invalid sequences (e.g. `\xfe\xfe`, UTF-16 BOM, mixed
362    // binary). `read_line` into String returns Err(InvalidData) on bad UTF-8
363    // and the previous code matched `Err => break;` -> server terminated
364    // mid-session. Per JSON-RPC 2.0 §5.1 we should return -32700 Parse error
365    // and keep the connection alive (same as v1.4.90 P0-A but for the
366    // pre-string-conversion stage).
367    let mut line_bytes = Vec::<u8>::new();
368
369    loop {
370        line_bytes.clear();
371        match read_until_newline_limited(&mut reader, &mut line_bytes, MAX_STDIO_JSONL_LINE_BYTES)
372            .await
373        {
374            Ok(LimitedLineRead::Eof) => {
375                // True EOF — stdin closed by client. Give already accepted
376                // requests a chance to emit their JSON-RPC response before
377                // rmcp observes EOF and starts shutdown/drain.
378                if !pending.wait_empty_or_timeout(EOF_PENDING_DRAIN_GRACE).await {
379                    tracing::warn!(
380                        "resilient stdio: pending responses did not drain before EOF grace expired"
381                    );
382                }
383                // Let the inbound channel close so the service loop sees
384                // `receive() -> None` and shuts down cleanly.
385                tracing::debug!("resilient stdio: stdin EOF, closing");
386                break;
387            }
388            Ok(LimitedLineRead::TooLarge) => {
389                tracing::warn!(
390                    max_bytes = MAX_STDIO_JSONL_LINE_BYTES,
391                    "resilient stdio: inbound message too large, returning -32600 (server stays alive)"
392                );
393                enqueue_message_too_large_response(&outbound_tx).await;
394                continue;
395            }
396            Ok(LimitedLineRead::Line) => {
397                // v1.4.93 P0-3: try UTF-8 conversion; on failure emit -32700
398                // and continue reading instead of terminating the reader task.
399                let line_str = match std::str::from_utf8(&line_bytes) {
400                    Ok(s) => s,
401                    Err(utf8_err) => {
402                        // Build a small ASCII preview of the offending bytes
403                        // for the error message (escape non-ASCII as `\xNN`).
404                        let preview_bytes: String = line_bytes
405                            .iter()
406                            .take(64)
407                            .map(|b| {
408                                if (0x20..=0x7e).contains(b) {
409                                    (*b as char).to_string()
410                                } else {
411                                    format!("\\x{b:02x}")
412                                }
413                            })
414                            .collect();
415                        let err_msg = JsonRpcError::new(
416                            None, // no id recoverable from non-UTF8
417                            ErrorData::new(
418                                ErrorCode::PARSE_ERROR,
419                                format!(
420                                    "Parse error: invalid UTF-8 at byte {}: {}",
421                                    utf8_err.valid_up_to(),
422                                    utf8_err
423                                ),
424                                None,
425                            ),
426                        );
427                        tracing::warn!(
428                            error = %utf8_err,
429                            line_preview = %preview_bytes,
430                            "resilient stdio: invalid UTF-8 input, returning -32700 (server stays alive)"
431                        );
432                        enqueue_parse_error_response(&outbound_tx, err_msg, "invalid_utf8").await;
433                        continue;
434                    }
435                };
436                let trimmed =
437                    line_str.trim_matches(|c| c == '\n' || c == '\r' || c == ' ' || c == '\t');
438                if trimmed.is_empty() {
439                    continue;
440                }
441                match serde_json::from_str::<RxJsonRpcMessage<RoleServer>>(trimmed) {
442                    Ok(msg) => {
443                        let pending_id = request_id(&msg).cloned();
444                        if let Some(id) = pending_id.clone() {
445                            pending.insert(id).await;
446                        }
447                        if inbound_tx.send(msg).await.is_err() {
448                            if let Some(id) = pending_id {
449                                pending.remove(&id).await;
450                            }
451                            // Receiver dropped — transport is closing.
452                            break;
453                        }
454                    }
455                    Err(parse_err) => {
456                        // Per JSON-RPC 2.0 §5.1, return -32700 Parse error
457                        // and keep the connection alive. Try to extract the
458                        // request id from the malformed payload (best-effort
459                        // — the spec says id should be `null` if it can't
460                        // be determined, but rmcp's `JsonRpcError` requires
461                        // a `RequestId` so we synthesise one as 0 / "" when
462                        // missing).
463                        let id = recover_request_id(trimmed);
464                        let err_msg = JsonRpcError::new(
465                            id,
466                            ErrorData::new(
467                                ErrorCode::PARSE_ERROR,
468                                format!("Parse error: {parse_err}"),
469                                None,
470                            ),
471                        );
472                        tracing::warn!(
473                            error = %parse_err,
474                            line_preview = %preview(trimmed),
475                            "resilient stdio: parse error, returning -32700 (server stays alive)"
476                        );
477                        enqueue_parse_error_response(&outbound_tx, err_msg, "json_parse_error")
478                            .await;
479                    }
480                }
481            }
482            Err(io_err) => {
483                tracing::warn!(error = %io_err, "resilient stdio: read error, terminating");
484                break;
485            }
486        }
487    }
488
489    // Drop the inbound sender so receive() returns None and the service
490    // loop exits cleanly.
491    drop(inbound_tx);
492}
493
494async fn writer_task<W>(write: W, mut outbound_rx: OutboundRx, pending: PendingRequestIds)
495where
496    W: AsyncWrite + Send + Unpin + 'static,
497{
498    let mut write = write;
499    while let Some(msg) = outbound_rx.recv().await {
500        let response_id = response_id(&msg).cloned();
501        match serde_json::to_vec(&msg) {
502            Ok(mut bytes) => {
503                bytes.push(b'\n');
504                if let Err(io_err) = write.write_all(&bytes).await {
505                    if let Some(id) = response_id.as_ref() {
506                        pending.remove(id).await;
507                    }
508                    tracing::warn!(error = %io_err, "resilient stdio: write error, terminating");
509                    break;
510                }
511                if let Err(io_err) = write.flush().await {
512                    if let Some(id) = response_id.as_ref() {
513                        pending.remove(id).await;
514                    }
515                    tracing::warn!(error = %io_err, "resilient stdio: flush error, terminating");
516                    break;
517                }
518                if let Some(id) = response_id.as_ref() {
519                    pending.remove(id).await;
520                }
521            }
522            Err(serde_err) => {
523                // This should never happen — TxJsonRpcMessage<RoleServer>
524                // is always serialisable. If it does, log and skip.
525                tracing::error!(
526                    error = %serde_err,
527                    "resilient stdio: failed to serialise outbound message (BUG)"
528                );
529            }
530        }
531    }
532}
533
534fn request_id(msg: &RxJsonRpcMessage<RoleServer>) -> Option<&NumberOrString> {
535    match msg {
536        rmcp::model::JsonRpcMessage::Request(req) => Some(&req.id),
537        _ => None,
538    }
539}
540
541fn response_id(msg: &TxJsonRpcMessage<RoleServer>) -> Option<&NumberOrString> {
542    match msg {
543        rmcp::model::JsonRpcMessage::Response(resp) => Some(&resp.id),
544        rmcp::model::JsonRpcMessage::Error(err) => err.id.as_ref(),
545        _ => None,
546    }
547}
548
549/// Best-effort recovery of the `id` field from a malformed JSON-RPC payload.
550/// Returns `None` when extraction fails so rmcp 3.x serializes the JSON-RPC
551/// specification's canonical `null` error id.
552fn recover_request_id(line: &str) -> Option<NumberOrString> {
553    if let Ok(value) = serde_json::from_str::<Value>(line)
554        && let Some(id) = value.get("id")
555    {
556        if let Some(n) = id.as_i64() {
557            return Some(NumberOrString::Number(n));
558        }
559        if let Some(s) = id.as_str() {
560            return Some(NumberOrString::String(s.into()));
561        }
562    }
563    None
564}
565
566/// Truncate a line for log output (avoid dumping arbitrary client input
567/// into the audit log unbounded).
568fn preview(s: &str) -> String {
569    const MAX: usize = 200;
570    if s.len() <= MAX {
571        s.to_string()
572    } else {
573        format!("{}…(+{} bytes)", &s[..MAX], s.len() - MAX)
574    }
575}
576
577// ---------------------------------------------------------------------------
578// Tests
579// ---------------------------------------------------------------------------
580
581#[cfg(test)]
582mod tests;
futu_mcp/transport.rs

futu_mcp/
transport.rs
