streamable_http_manager

class StreamableHTTPSessionManager:
    """Manages StreamableHTTP sessions with optional resumability via event store.

    This class abstracts away the complexity of session management, event storage,
    and request handling for StreamableHTTP transports. It handles:

    1. Session tracking for clients
    2. Resumability via an optional event store
    3. Connection management and lifecycle
    4. Request handling and transport setup
    5. Idle session cleanup via optional timeout

    Important: Only one StreamableHTTPSessionManager instance should be created
    per application. The instance cannot be reused after its run() context has
    completed. If you need to restart the manager, create a new instance.

    Args:
        app: The MCP server instance
        event_store: Optional event store for resumability support. If provided, enables resumable connections
            where clients can reconnect and receive missed events. If None, sessions are still tracked but not
            resumable.
        json_response: Whether to use JSON responses instead of SSE streams
        stateless: If True, creates a completely fresh transport for each request with no session tracking or
            state persistence between requests.
        security_settings: Optional transport security settings.
        retry_interval: Retry interval in milliseconds to suggest to clients in SSE retry field. Used for SSE
            polling behavior.
        session_idle_timeout: Optional idle timeout in seconds for stateful sessions. If set, sessions that
            receive no HTTP requests for this duration will be automatically terminated and removed. When
            retry_interval is also configured, ensure the idle timeout comfortably exceeds the retry interval to
            avoid reaping sessions during normal SSE polling gaps. Default is None (no timeout). A value of 1800
            (30 minutes) is recommended for most deployments.
        max_request_body_size: Maximum size in bytes for Streamable HTTP POST request bodies. Requests that
            exceed this limit receive a 413 response before parsing or session creation. Defaults to 4 MiB.
    """

    def __init__(
        self,
        app: Server[Any],
        event_store: EventStore | None = None,
        json_response: bool = False,
        stateless: bool = False,
        security_settings: TransportSecuritySettings | None = None,
        retry_interval: int | None = None,
        session_idle_timeout: float | None = None,
        max_request_body_size: int = DEFAULT_MAX_REQUEST_BODY_SIZE,
    ):
        if session_idle_timeout is not None and session_idle_timeout <= 0:
            raise ValueError("session_idle_timeout must be a positive number of seconds")
        if stateless and session_idle_timeout is not None:
            raise RuntimeError("session_idle_timeout is not supported in stateless mode")
        if max_request_body_size <= 0:
            raise ValueError("max_request_body_size must be a positive number of bytes")

        self.app = app
        self.event_store = event_store
        self.json_response = json_response
        self.stateless = stateless
        self.security_settings = security_settings
        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)

        # Session tracking (only used if not stateless)
        self._session_creation_lock = anyio.Lock()
        self._server_instances: dict[str, StreamableHTTPServerTransport] = {}
        # Identity of the credential that created each session; requests for a
        # session must present the same credential.
        self._session_owners: dict[str, AuthorizationContext] = {}

        # The task group and lifespan state are set during run()
        self._task_group = None
        self._lifespan_state: Any = None
        # Thread-safe tracking of run() calls
        self._run_lock = anyio.Lock()
        self._has_started = False

    @contextlib.asynccontextmanager
    async def run(self) -> AsyncIterator[None]:
        """Run the session manager with proper lifecycle management.

        This creates and manages the task group for all session operations.

        Important: This method can only be called once per instance. The same
        StreamableHTTPSessionManager instance cannot be reused after this
        context manager exits. Create a new instance if you need to restart.

        Use this in the lifespan context manager of your Starlette app:

        @contextlib.asynccontextmanager
        async def lifespan(app: Starlette) -> AsyncIterator[None]:
            async with session_manager.run():
                yield
        """
        # Thread-safe check to ensure run() is only called once
        async with self._run_lock:
            if self._has_started:
                raise RuntimeError(
                    "StreamableHTTPSessionManager .run() can only be called "
                    "once per instance. Create a new instance if you need to run again."
                )
            self._has_started = True

        async with self.app.lifespan(self.app) as lifespan_state, anyio.create_task_group() as tg:
            # Store for handle_request: lifespan is entered once for the
            # manager's lifetime, not per request (per-connection cleanup
            # belongs on `connection.exit_stack`).
            self._lifespan_state = lifespan_state
            self._task_group = tg
            logger.info("StreamableHTTP session manager started")
            try:
                yield  # Let the application run
            finally:
                logger.info("StreamableHTTP session manager shutting down")
                # Cancel task group to stop all spawned tasks
                tg.cancel_scope.cancel()
                self._task_group = None
                self._lifespan_state = None
                # Clear any remaining server instances
                self._server_instances.clear()
                self._session_owners.clear()
        await resync_tracer()

    async def handle_request(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Process ASGI request with proper session handling and transport setup.

        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().")

        # TODO(L49): header-only era-routing for now; body-primary classification
        # is a follow-up. The legacy paths below own only the known
        # initialize-handshake versions; anything else (including unknown
        # values) goes to the modern entry so the classifier can validate it
        # and return a structured rejection. 2025 paths below remain unchanged.
        header = MCP_PROTOCOL_VERSION_HEADER.encode("ascii")
        pv = next((v.decode("latin-1") for k, v in scope["headers"] if k == header), None)
        if pv is not None and pv not in HANDSHAKE_PROTOCOL_VERSIONS:
            await handle_modern_request(
                self.app, self.security_settings, self.json_response, self._lifespan_state, scope, receive, send
            )
            return

        # Dispatch to the appropriate handler
        if self.stateless:
            await self._handle_stateless_request(scope, receive, send)
        else:
            await self._handle_stateful_request(scope, receive, send)

    async def _handle_stateless_request(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Process request in stateless mode - creating a new transport for each request."""
        logger.debug("Stateless mode: Creating new transport for this request")
        # No session ID needed in stateless mode: the transport serves this one
        # request with a born-ready connection (no `initialize`, no standalone
        # GET stream) and no back-channel for server-to-client requests.
        http_transport = StreamableHTTPServerTransport(
            mcp_session_id=None,  # No session tracking in stateless mode
            is_json_response_enabled=self.json_response,
            event_store=None,  # No event store in stateless mode
            security_settings=self.security_settings,
            app=self.app,
            lifespan_state=self._lifespan_state,
        )

        # Handle the HTTP request and return the response
        await http_transport.handle_request(scope, receive, send)

        # Terminate the transport after the request is handled
        await http_transport.terminate()

    async def _handle_stateful_request(self, scope: Scope, receive: Receive, send: Send) -> None:
        """Process request in stateful mode - maintaining session state between requests."""
        request = Request(scope, receive)
        request_mcp_session_id = request.headers.get(MCP_SESSION_ID_HEADER)

        user = scope.get("user")
        requestor = authorization_context(user) if isinstance(user, AuthenticatedUser) else None

        # Existing session case
        if request_mcp_session_id is not None and request_mcp_session_id in self._server_instances:
            transport = self._server_instances[request_mcp_session_id]
            if requestor != self._session_owners.get(request_mcp_session_id):
                # A session can only be used with the credential that created
                # it. Respond exactly as if the session did not exist.
                logger.warning(
                    "Rejecting request for session %s: credential does not match the one that created the session",
                    request_mcp_session_id[:64],
                )
                body = JSONRPCError(
                    jsonrpc="2.0", id=None, error=ErrorData(code=INVALID_REQUEST, message="Session not found")
                )
                response = Response(
                    body.model_dump_json(by_alias=True, exclude_unset=True),
                    status_code=404,
                    media_type="application/json",
                )
                await response(scope, receive, send)
                return
            logger.debug("Session already exists, handling request directly")
            # Push back idle deadline on activity
            if transport.idle_scope is not None and self.session_idle_timeout is not None:
                transport.idle_scope.deadline = anyio.current_time() + self.session_idle_timeout  # pragma: no cover
            await transport.handle_request(scope, receive, send)
            return

        if request_mcp_session_id is None:
            # New session case
            logger.debug("Creating new transport")
            async with self._session_creation_lock:
                new_session_id = uuid4().hex
                http_transport = StreamableHTTPServerTransport(
                    mcp_session_id=new_session_id,
                    is_json_response_enabled=self.json_response,
                    event_store=self.event_store,  # May be None (no resumability)
                    security_settings=self.security_settings,
                    retry_interval=self.retry_interval,
                    # The manager owns the lifespan (entered once in `run()`),
                    # so the transport serves every request off that state.
                    app=self.app,
                    lifespan_state=self._lifespan_state,
                )

                assert http_transport.mcp_session_id is not None
                if requestor is not None:
                    self._session_owners[http_transport.mcp_session_id] = requestor
                self._server_instances[http_transport.mcp_session_id] = http_transport
                logger.info(f"Created new transport with session ID: {new_session_id}")

                # Define the session task: hosts the session's request handlers,
                # which outlive the HTTP requests that started them.
                async def run_server(*, task_status: TaskStatus[None] = anyio.TASK_STATUS_IGNORED) -> None:
                    try:
                        # Use a cancel scope for idle timeout — when the
                        # deadline passes the scope cancels the session and
                        # execution continues after the ``with`` block.
                        # Incoming requests push the deadline forward.
                        idle_scope = anyio.CancelScope()
                        if self.session_idle_timeout is not None:
                            idle_scope.deadline = anyio.current_time() + self.session_idle_timeout
                            http_transport.idle_scope = idle_scope

                        with idle_scope:
                            await http_transport.run(task_status=task_status)

                        if idle_scope.cancelled_caught:
                            assert http_transport.mcp_session_id is not None
                            logger.info(f"Session {http_transport.mcp_session_id} idle timeout")
                            self._server_instances.pop(http_transport.mcp_session_id, None)
                            self._session_owners.pop(http_transport.mcp_session_id, None)
                            await http_transport.terminate()
                    except Exception:
                        logger.exception(f"Session {http_transport.mcp_session_id} crashed")
                    finally:
                        if (  # pragma: no branch
                            http_transport.mcp_session_id
                            and http_transport.mcp_session_id in self._server_instances
                            and not http_transport.is_terminated
                        ):
                            logger.info(
                                f"Cleaning up crashed session {http_transport.mcp_session_id} from active instances."
                            )
                            del self._server_instances[http_transport.mcp_session_id]
                            self._session_owners.pop(http_transport.mcp_session_id, None)

                # Assert task group is not None for type checking
                assert self._task_group is not None
                # Start the server task
                await self._task_group.start(run_server)

                # Handle the HTTP request and return the response
                await http_transport.handle_request(scope, receive, send)
        else:
            # Unknown or expired session ID - return 404 per MCP spec
            # TODO(L62): Align error code once spec clarifies
            # See: https://github.com/modelcontextprotocol/python-sdk/issues/1821
            logger.info(f"Rejected request with unknown or expired session ID: {request_mcp_session_id[:64]}")
            body = JSONRPCError(
                jsonrpc="2.0", id=None, error=ErrorData(code=INVALID_REQUEST, message="Session not found")
            )
            response = Response(
                body.model_dump_json(by_alias=True, exclude_unset=True), status_code=404, media_type="application/json"
            )
            await response(scope, receive, send)

class RequestBodyLimitMiddleware:
    """Reject oversized HTTP request bodies before invoking an ASGI application."""

    def __init__(self, app: ASGIApp, max_body_size: int) -> None:
        self.app = app
        self.max_body_size = max_body_size

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http" or scope["method"] != "POST":
            await self.app(scope, receive, send)
            return

        headers = Headers(scope=scope)
        content_length = headers.get("content-length")
        if content_length is not None:
            try:
                declared_size = int(content_length)
            except ValueError:
                pass
            else:
                if declared_size > self.max_body_size:
                    response = Response("Request body too large", status_code=413)
                    return await response(scope, receive, send)

        received_body = bytearray()
        received_request = False
        body_complete = False
        trailing_message: Message | None = None
        while True:
            message = await receive()
            if message["type"] != "http.request":
                trailing_message = message
                break

            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)
            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()
            return await receive()

        await self.app(scope, replay, send)

class StreamableHTTPASGIApp:
    """ASGI application for Streamable HTTP server transport."""

    def __init__(self, session_manager: StreamableHTTPSessionManager):
        self.session_manager = session_manager

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        await self.session_manager.asgi_app(scope, receive, send)