Address request body limit review feedback

Author: Kludex

src/mcp/server/streamable_http_manager.py

@@ -100,7 +100,7 @@ def __init__(
         self.retry_interval = retry_interval
         self.session_idle_timeout = session_idle_timeout
         self.max_request_body_size = max_request_body_size
-        self.asgi_app = RequestBodyLimitMiddleware(self.handle_request, max_request_body_size)
+        self.asgi_app = RequestBodyLimitMiddleware(self._handle_request, max_request_body_size)
 
         # Session tracking (only used if not stateless)
         self._session_creation_lock = anyio.Lock()
@@ -167,6 +167,9 @@ async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> No
 
         Dispatches to the appropriate handler based on stateless mode.
         """
+        await self.asgi_app(scope, receive, send)
+
+    async def _handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
         if self._task_group is None:
             raise RuntimeError("Task group is not initialized. Make sure to use run().")
 
@@ -392,22 +395,34 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
                     response = Response("Request body too large", status_code=413)
                     return await response(scope, receive, send)
 
-        cached_messages: deque[Message] = deque()
-        received_size = 0
+        received_body = bytearray()
+        received_request = False
+        body_complete = False
+        trailing_message: Message | None = None
         while True:
             message = await receive()
             if message["type"] != "http.request":
-                cached_messages.append(message)
+                trailing_message = message
                 break
 
-            received_size += len(message.get("body", b""))
-            if received_size > self.max_body_size:
+            received_request = True
+            body = message.get("body", b"")
+            if len(received_body) + len(body) > self.max_body_size:
                 response = Response("Request body too large", status_code=413)
                 return await response(scope, receive, send)
-            cached_messages.append(message)
+            received_body.extend(body)
             if not message.get("more_body", False):
+                body_complete = True
                 break
 
+        cached_messages: deque[Message] = deque()
+        if received_request:
+            cached_messages.append(
+                {"type": "http.request", "body": bytes(received_body), "more_body": not body_complete}
+            )
+        if trailing_message is not None:
+            cached_messages.append(trailing_message)
+
         async def replay() -> Message:
             if cached_messages:
                 return cached_messages.popleft()

tests/server/test_streamable_http_manager.py

@@ -76,9 +76,9 @@ async def test_handle_request_without_run_raises_error():
     manager = StreamableHTTPSessionManager(app=app)
 
     # Mock ASGI parameters
-    scope = {"type": "http", "method": "POST", "path": "/test"}
+    scope: Scope = {"type": "http", "method": "POST", "path": "/test", "headers": []}
 
-    async def receive():  # pragma: no cover
+    async def receive() -> Message:
         return {"type": "http.request", "body": b""}
 
     async def send(message: Message):  # pragma: no cover
@@ -108,7 +108,7 @@ async def send(message: Message) -> None:
         "headers": [(b"content-length", b"9")],
     }
     async with manager.run():
-        await manager.asgi_app(scope, receive, send)
+        await manager.handle_request(scope, receive, send)
         assert manager._server_instances == {}
 
     response_start = next(message for message in sent_messages if message["type"] == "http.response.start")
@@ -150,6 +150,33 @@ async def send(message: Message) -> None:
 async def test_client_disconnect_while_streaming_request_body_is_replayed() -> None:
     """SDK-defined: raw ASGI is required to prove a disconnect before body completion reaches the transport."""
     disconnect: Message = {"type": "http.disconnect"}
+    request_messages: Iterator[Message] = iter(
+        [{"type": "http.request", "body": b"1234", "more_body": True}, disconnect]
+    )
+    received_messages: list[Message] = []
+
+    async def receive() -> Message:
+        return next(request_messages)
+
+    async def app(scope: Scope, receive: Receive, send: Send) -> None:
+        received_messages.append(await receive())
+        received_messages.append(await receive())
+
+    scope: Scope = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
+    middleware = RequestBodyLimitMiddleware(app, max_body_size=8)
+
+    await middleware(scope, receive, AsyncMock())
+
+    assert received_messages == [
+        {"type": "http.request", "body": b"1234", "more_body": True},
+        disconnect,
+    ]
+
+
+@pytest.mark.anyio
+async def test_client_disconnect_before_request_body_is_replayed() -> None:
+    """SDK-defined: raw ASGI proves a disconnect before the first body message reaches the transport."""
+    disconnect: Message = {"type": "http.disconnect"}
     received_messages: list[Message] = []
 
     async def receive() -> Message:
@@ -166,6 +193,32 @@ async def app(scope: Scope, receive: Receive, send: Send) -> None:
     assert received_messages == [disconnect]
 
 
+@pytest.mark.anyio
+async def test_request_body_chunks_are_replayed_as_one_message() -> None:
+    """SDK-defined: raw ASGI proves chunk overhead is discarded before the body reaches the transport."""
+    request_messages: Iterator[Message] = iter(
+        [
+            {"type": "http.request", "body": b"12", "more_body": True},
+            {"type": "http.request", "body": b"34", "more_body": True},
+            {"type": "http.request", "body": b"56", "more_body": False},
+        ]
+    )
+    received_messages: list[Message] = []
+
+    async def receive() -> Message:
+        return next(request_messages)
+
+    async def app(scope: Scope, receive: Receive, send: Send) -> None:
+        received_messages.append(await receive())
+
+    scope: Scope = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
+    middleware = RequestBodyLimitMiddleware(app, max_body_size=8)
+
+    await middleware(scope, receive, AsyncMock())
+
+    assert received_messages == [{"type": "http.request", "body": b"123456", "more_body": False}]
+
+
 def test_request_body_limit_defaults_to_four_mib() -> None:
     """SDK-defined: Streamable HTTP request bodies are limited to 4 MiB by default."""
     manager = StreamableHTTPSessionManager(app=Server("test-default-size-limit"))
@@ -216,7 +269,7 @@ async def mock_send(message: Message):
         "headers": [(b"content-type", b"application/json")],
     }
 
-    async def mock_receive():  # pragma: no cover
+    async def mock_receive():
         return {"type": "http.request", "body": b"", "more_body": False}
 
     # Trigger session creation
@@ -274,7 +327,7 @@ async def mock_send(message: Message):
         "headers": [(b"content-type", b"application/json")],
     }
 
-    async def mock_receive():  # pragma: no cover
+    async def mock_receive():
         return {"type": "http.request", "body": b"", "more_body": False}
 
     # Trigger session creation
@@ -392,7 +445,7 @@ async def mock_send(message: Message):
         }
 
         async def mock_receive():
-            return {"type": "http.request", "body": b"{}", "more_body": False}  # pragma: no cover
+            return {"type": "http.request", "body": b"{}", "more_body": False}
 
         with caplog.at_level(logging.INFO):
             await manager.handle_request(scope, mock_receive, mock_send)
@@ -472,7 +525,7 @@ async def mock_send(message: Message):
             "headers": [(b"content-type", b"application/json")],
         }
 
-        async def mock_receive():  # pragma: no cover
+        async def mock_receive():
             return {"type": "http.request", "body": b"", "more_body": False}
 
         await manager.handle_request(scope, mock_receive, mock_send)