|
4 | 4 |
|
5 | 5 | import contextlib |
6 | 6 | import logging |
7 | | -from collections import deque |
8 | 7 | from collections.abc import AsyncIterator |
9 | | -from typing import TYPE_CHECKING, Any, Final |
| 8 | +from typing import TYPE_CHECKING, Any |
10 | 9 | from uuid import uuid4 |
11 | 10 |
|
12 | 11 | import anyio |
13 | 12 | from anyio.abc import TaskStatus |
14 | 13 | from mcp_types import DEFAULT_NEGOTIATED_VERSION, INVALID_REQUEST, ErrorData, JSONRPCError |
15 | 14 | from mcp_types.version import HANDSHAKE_PROTOCOL_VERSIONS |
16 | | -from starlette.datastructures import Headers |
17 | 15 | from starlette.requests import Request |
18 | 16 | from starlette.responses import Response |
19 | | -from starlette.types import ASGIApp, Message, Receive, Scope, Send |
| 17 | +from starlette.types import Receive, Scope, Send |
20 | 18 |
|
21 | 19 | from mcp.server._streamable_http_modern import handle_modern_request |
22 | 20 | from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser, AuthorizationContext, authorization_context |
23 | 21 | from mcp.server.connection import Connection |
24 | 22 | from mcp.server.runner import serve_connection, serve_loop |
25 | 23 | from mcp.server.streamable_http import MCP_SESSION_ID_HEADER, EventStore, StreamableHTTPServerTransport |
| 24 | +from mcp.server.transport_security import DEFAULT_MAX_REQUEST_BODY_SIZE as DEFAULT_MAX_REQUEST_BODY_SIZE |
| 25 | +from mcp.server.transport_security import RequestBodyLimitMiddleware as RequestBodyLimitMiddleware |
26 | 26 | from mcp.server.transport_security import TransportSecuritySettings |
27 | 27 | from mcp.shared._compat import resync_tracer |
28 | 28 | from mcp.shared.inbound import MCP_PROTOCOL_VERSION_HEADER |
|
34 | 34 |
|
35 | 35 | logger = logging.getLogger(__name__) |
36 | 36 |
|
37 | | -DEFAULT_MAX_REQUEST_BODY_SIZE: Final = 4 * 1024 * 1024 |
38 | | -"""Default maximum Streamable HTTP request body size in bytes (4 MiB).""" |
39 | | - |
40 | 37 |
|
41 | 38 | class StreamableHTTPSessionManager: |
42 | 39 | """Manages StreamableHTTP sessions with optional resumability via event store. |
@@ -70,7 +67,7 @@ class StreamableHTTPSessionManager: |
70 | 67 | retry_interval is also configured, ensure the idle timeout comfortably exceeds the retry interval to |
71 | 68 | avoid reaping sessions during normal SSE polling gaps. Default is None (no timeout). A value of 1800 |
72 | 69 | (30 minutes) is recommended for most deployments. |
73 | | - max_request_body_size: Maximum size in bytes for Streamable HTTP POST request bodies. Requests that |
| 70 | + max_request_body_size: Maximum size in bytes for Streamable HTTP request bodies. Requests that |
74 | 71 | exceed this limit receive a 413 response before parsing or session creation. Defaults to 4 MiB. |
75 | 72 | """ |
76 | 73 |
|
@@ -371,66 +368,6 @@ async def run_server(*, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORE |
371 | 368 | await response(scope, receive, send) |
372 | 369 |
|
373 | 370 |
|
374 | | -class RequestBodyLimitMiddleware: |
375 | | - """Reject oversized HTTP request bodies before invoking an ASGI application.""" |
376 | | - |
377 | | - def __init__(self, app: ASGIApp, max_body_size: int) -> None: |
378 | | - self.app = app |
379 | | - self.max_body_size = max_body_size |
380 | | - |
381 | | - async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: |
382 | | - if scope["type"] != "http" or scope["method"] != "POST": |
383 | | - await self.app(scope, receive, send) |
384 | | - return |
385 | | - |
386 | | - headers = Headers(scope=scope) |
387 | | - content_length = headers.get("content-length") |
388 | | - if content_length is not None: |
389 | | - try: |
390 | | - declared_size = int(content_length) |
391 | | - except ValueError: |
392 | | - pass |
393 | | - else: |
394 | | - if declared_size > self.max_body_size: |
395 | | - response = Response("Request body too large", status_code=413) |
396 | | - return await response(scope, receive, send) |
397 | | - |
398 | | - received_body = bytearray() |
399 | | - received_request = False |
400 | | - body_complete = False |
401 | | - trailing_message: Message | None = None |
402 | | - while True: |
403 | | - message = await receive() |
404 | | - if message["type"] != "http.request": |
405 | | - trailing_message = message |
406 | | - break |
407 | | - |
408 | | - received_request = True |
409 | | - body = message.get("body", b"") |
410 | | - if len(received_body) + len(body) > self.max_body_size: |
411 | | - response = Response("Request body too large", status_code=413) |
412 | | - return await response(scope, receive, send) |
413 | | - received_body.extend(body) |
414 | | - if not message.get("more_body", False): |
415 | | - body_complete = True |
416 | | - break |
417 | | - |
418 | | - cached_messages: deque[Message] = deque() |
419 | | - if received_request: |
420 | | - cached_messages.append( |
421 | | - {"type": "http.request", "body": bytes(received_body), "more_body": not body_complete} |
422 | | - ) |
423 | | - if trailing_message is not None: |
424 | | - cached_messages.append(trailing_message) |
425 | | - |
426 | | - async def replay() -> Message: |
427 | | - if cached_messages: |
428 | | - return cached_messages.popleft() |
429 | | - return await receive() |
430 | | - |
431 | | - await self.app(scope, replay, send) |
432 | | - |
433 | | - |
434 | 371 | class StreamableHTTPASGIApp: |
435 | 372 | """ASGI application for Streamable HTTP server transport.""" |
436 | 373 |
|
|
0 commit comments