fix: re-use BufReader across poll_next invocations, instead of creating a new one (#1093)

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<R, W> {
+    // 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<R>,
     #[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<Option<Self::Item>> {
         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