diff --git a/crates/mcp-server/src/lib.rs b/crates/mcp-server/src/lib.rs index de63666789..975945239b 100644 --- a/crates/mcp-server/src/lib.rs +++ b/crates/mcp-server/src/lib.rs @@ -18,8 +18,11 @@ pub use router::Router; /// A transport layer that handles JSON-RPC messages over byte #[pin_project] pub struct ByteTransport { + // Reader is a BufReader on the underlying stream (stdin or similar) buffering + // the underlying data across poll calls, we clear one line (\n) during each + // iteration of poll_next from this buffer #[pin] - reader: R, + reader: BufReader, #[pin] writer: W, } @@ -30,7 +33,12 @@ where W: AsyncWrite, { pub fn new(reader: R, writer: W) -> Self { - Self { reader, writer } + Self { + // Default BufReader capacity is 8 * 1024, increase this to 2MB to the file size limit + // allows the buffer to have the capacity to read very large calls + reader: BufReader::with_capacity(2 * 1024 * 1024, reader), + writer, + } } } @@ -44,10 +52,8 @@ where fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut this = self.project(); let mut buf = Vec::new(); - // Default BufReader capacity is 8 * 1024, increase this to 2MB to the file size limit - // allows the buffer to have the capacity to read very large calls - let mut reader = BufReader::with_capacity(2 * 1024 * 1024, &mut this.reader); + let mut reader = this.reader.as_mut(); let mut read_future = Box::pin(reader.read_until(b'\n', &mut buf)); match read_future.as_mut().poll(cx) { Poll::Ready(Ok(0)) => Poll::Ready(None), // EOF @@ -70,7 +76,6 @@ where "Message must be a JSON object".into(), )))); } - let obj = value.as_object().unwrap(); // Safe due to check above // Check jsonrpc version field