invariant-gateway/gateway/mcp/streamable.py

"""Gateway service to forward requests to the MCP Streamable HTTP servers"""

import json
from typing import Any

import httpx
from httpx_sse import aconnect_sse
from fastapi import APIRouter, HTTPException, Request, Response
from fastapi.responses import StreamingResponse

from gateway.common.constants import (
    CLIENT_TIMEOUT,
    CONTENT_TYPE_HEADER,
    CONTENT_TYPE_JSON,
    CONTENT_TYPE_EVENT_STREAM,
)
from gateway.mcp.constants import (
    INVARIANT_SESSION_ID_PREFIX,
    MCP_CUSTOM_HEADER_PREFIX,
    UTF_8,
)
from gateway.mcp.mcp_sessions_manager import (
    McpSessionsManager,
    McpAttributes,
)
from gateway.mcp.mcp_transport_base import McpTransportBase

gateway = APIRouter()
mcp_sessions_manager = McpSessionsManager()

MCP_SESSION_ID_HEADER = "mcp-session-id"
MCP_SERVER_POST_AND_DELETE_HEADERS = {
    "connection",
    "accept",
    CONTENT_TYPE_HEADER,
    MCP_SESSION_ID_HEADER,
}
MCP_SERVER_GET_HEADERS = {
    "connection",
    "accept",
    "cache-control",
    MCP_SESSION_ID_HEADER,
}


@gateway.post("/mcp/streamable")
async def mcp_post_streamable_gateway(
    request: Request,
):
    """Forward a POST request to the MCP Streamable server using transport strategy."""
    return await create_streamable_transport_and_handle_request(
        request, "POST", mcp_sessions_manager
    )


@gateway.get("/mcp/streamable")
async def mcp_get_streamable_gateway(request: Request) -> StreamingResponse:
    """Forward a GET request to the MCP Streamable server using transport strategy."""
    return await create_streamable_transport_and_handle_request(
        request, "GET", mcp_sessions_manager
    )


@gateway.delete("/mcp/streamable")
async def mcp_delete_streamable_gateway(request: Request) -> Response:
    """Forward a DELETE request to the MCP Streamable server using transport strategy."""
    return await create_streamable_transport_and_handle_request(
        request, "DELETE", mcp_sessions_manager
    )


async def create_streamable_transport_and_handle_request(
    request: Request, method: str, session_store: McpSessionsManager
) -> Response | StreamingResponse:
    """Integration function for streamable routes."""
    streamable_transport = StreamableTransport(session_store)
    return await streamable_transport.handle_communication(
        request=request, method=method
    )


class StreamableTransport(McpTransportBase):
    """
    Streamable HTTP transport implementation for MCP communication.
    Handles HTTP POST/GET/DELETE requests with JSON and streaming responses.
    """

    async def initialize_session(
        self,
        **kwargs,
    ) -> str:
        """Initialize streamable HTTP session."""
        session_id: str | None = kwargs.get("session_id", None)
        session_attributes: McpAttributes | None = kwargs.get(
            "session_attributes", None
        )
        is_initialization_request: bool = kwargs.get("is_initialization_request", False)
        if session_id and self.session_store.session_exists(session_id):
            return session_id

        if is_initialization_request and not session_id:
            session_id = self.generate_session_id()

        if (
            session_id
            and not self.session_store.session_exists(session_id)
            and session_attributes
        ):
            await self.session_store.initialize_session(session_id, session_attributes)

        return session_id

    async def handle_post_request(
        self, request: Request, request_body: dict[str, Any]
    ) -> Response | StreamingResponse:
        """Handle POST request to streamable endpoint."""
        session_attributes = McpAttributes.from_request_headers(request.headers)
        session_id = request.headers.get(MCP_SESSION_ID_HEADER)
        is_initialization_request = self._is_initialization_request(request_body)

        # Handle session initialization
        if session_id:
            self.update_tool_call_id_in_session(
                self.session_store.get_session(session_id), request_body
            )
        elif is_initialization_request:
            session_id = await self.initialize_session(
                session_attributes=session_attributes, is_initialization_request=True
            )

        # Process request if not initialization
        if not is_initialization_request:
            request_interception_result = await self._process_non_init_request(
                session_id, request_body
            )
            if request_interception_result:
                return request_interception_result

        # Forward to MCP server
        return await self._forward_to_mcp_server(
            request,
            request_body,
            session_id,
            session_attributes,
            is_initialization_request,
        )

    async def handle_get_request(self, request: Request) -> StreamingResponse:
        """Handle GET request for server-initiated communication."""
        mcp_server_endpoint = self._get_mcp_server_endpoint(request)
        response_headers = {}

        filtered_headers = {}
        for k, v in request.headers.items():
            if k.startswith(MCP_CUSTOM_HEADER_PREFIX):
                filtered_headers[k.removeprefix(MCP_CUSTOM_HEADER_PREFIX)] = v
            if k.lower() in MCP_SERVER_GET_HEADERS:
                filtered_headers[k] = v

        async def event_generator():
            async with httpx.AsyncClient(
                timeout=httpx.Timeout(CLIENT_TIMEOUT)
            ) as client:
                try:
                    async with aconnect_sse(
                        client,
                        "GET",
                        mcp_server_endpoint,
                        headers=filtered_headers,
                    ) as event_source:
                        if event_source.response.status_code != 200:
                            error_content = await event_source.response.aread()
                            raise HTTPException(
                                status_code=event_source.response.status_code,
                                detail=error_content,
                            )

                        response_headers.update(
                            dict(event_source.response.headers.items())
                        )

                        async for sse in event_source.aiter_sse():
                            yield sse

                except httpx.StreamClosed as e:
                    print(f"Server stream closed: {e}")
                except Exception as e:  # pylint: disable=broad-except
                    print(f"Error processing server events: {e}")

        return StreamingResponse(
            event_generator(),
            media_type=CONTENT_TYPE_EVENT_STREAM,
            headers={"X-Proxied-By": "mcp-gateway", **response_headers},
        )

    async def handle_delete_request(self, request: Request) -> Response:
        """Handle DELETE request for session termination."""
        session_id = self._get_session_id(request)

        if not self.session_store.session_exists(session_id):
            raise HTTPException(status_code=400, detail="Session does not exist")

        if session_id.startswith(INVARIANT_SESSION_ID_PREFIX):
            return Response(
                content="", status_code=200, headers={"X-Proxied-By": "mcp-gateway"}
            )

        mcp_server_endpoint = self._get_mcp_server_endpoint(request)

        async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client:
            try:
                response = await client.delete(
                    url=mcp_server_endpoint,
                    headers=self._get_headers_for_mcp_post_and_delete(request),
                )
                await self.session_store.cleanup_session_lock(session_id)
                return Response(
                    content=response.content,
                    status_code=response.status_code,
                    headers={"X-Proxied-By": "mcp-gateway", **response.headers},
                )
            except httpx.RequestError as e:
                print(f"[MCP DELETE] Request error: {str(e)}")
                raise HTTPException(status_code=500, detail="Request error") from e

    async def handle_communication(self, **kwargs) -> Response | StreamingResponse:
        """Main communication handler for streamable transport."""
        request = kwargs.get("request")
        method = kwargs.get("method", "POST")

        if method == "POST":
            request_body = json.loads(await request.body())
            return await self.handle_post_request(request, request_body)
        elif method == "GET":
            return await self.handle_get_request(request)
        elif method == "DELETE":
            return await self.handle_delete_request(request)
        else:
            raise HTTPException(status_code=405, detail="Method not allowed")

    async def _process_non_init_request(
        self, session_id: str, request_body: dict[str, Any]
    ) -> Response | None:
        """Process non-initialization requests for guardrails."""
        processed_request, is_blocked = await self.process_outgoing_request(
            session_id, request_body
        )

        if is_blocked:
            return Response(
                content=json.dumps(processed_request),
                status_code=400,
                media_type=CONTENT_TYPE_JSON,
            )
        return None

    async def _forward_to_mcp_server(
        self,
        request: Request,
        request_body: dict[str, Any],
        session_id: str,
        session_attributes: McpAttributes,
        is_initialization_request: bool,
    ) -> Response | StreamingResponse:
        """Forward request to MCP server and handle response."""
        async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client:
            try:
                response = await client.post(
                    url=self._get_mcp_server_endpoint(request),
                    headers=self._get_headers_for_mcp_post_and_delete(request),
                    content=json.dumps(request_body).encode(),
                    follow_redirects=True,
                )

                # Handle session ID from MCP server response
                resp_session_id = response.headers.get(MCP_SESSION_ID_HEADER)

                if resp_session_id:
                    if not self.session_store.session_exists(resp_session_id):
                        await self.session_store.initialize_session(
                            resp_session_id, session_attributes
                        )
                    session_id = resp_session_id
                elif (
                    is_initialization_request
                    and not self.session_store.session_exists(session_id)
                ):
                    await self.session_store.initialize_session(
                        session_id, session_attributes
                    )

                # Update client info for initialization requests
                if is_initialization_request:
                    self.update_mcp_client_info_in_session(
                        self.session_store.get_session(session_id), request_body
                    )

                # Handle response based on content type
                if response.headers.get(CONTENT_TYPE_HEADER) == CONTENT_TYPE_JSON:
                    return await self._handle_json_response(
                        session_id, is_initialization_request, response
                    )
                else:
                    return await self._handle_streaming_response(
                        session_id, is_initialization_request, response
                    )

            except httpx.RequestError as e:
                print(f"[MCP POST] Request error: {str(e)}")
                raise HTTPException(status_code=500, detail="Request error") from e

    async def _handle_json_response(
        self, session_id: str, is_initialization_request: bool, response: httpx.Response
    ) -> Response:
        """Handle JSON response from MCP server."""
        response_content = response.content
        response_body = (
            json.loads(response_content.decode(UTF_8)) if response_content else {}
        )

        if response_body:
            self._update_mcp_response_info_in_session(session_id, response_body, True)

        response_code = response.status_code

        if not is_initialization_request and response_body:
            processed_response, blocked = await self.process_incoming_response(
                session_id, response_body
            )
            if blocked:
                response_content = json.dumps(processed_response).encode(UTF_8)
                response_code = 400

        # Build response headers
        response_headers = {"X-Proxied-By": "mcp-gateway", **response.headers}
        if MCP_SESSION_ID_HEADER not in response.headers:
            response_headers[MCP_SESSION_ID_HEADER] = session_id

        return Response(
            content=response_content,
            status_code=response_code,
            headers=response_headers,
        )

    async def _handle_streaming_response(
        self, session_id: str, is_initialization_request: bool, response: httpx.Response
    ) -> StreamingResponse:
        """Handle streaming response from MCP server."""

        async def event_generator():
            buffer = ""
            async for line in response.aiter_lines():
                stripped_line = line.strip()
                if not stripped_line:
                    break

                if buffer:
                    response_body = json.loads(stripped_line.split("data: ")[1].strip())

                    if not is_initialization_request:
                        (
                            processed_response,
                            blocked,
                        ) = await self.process_incoming_response(
                            session_id, response_body
                        )
                        if blocked:
                            yield f"{buffer}\ndata: {json.dumps(processed_response)}\n\n"
                            break
                    else:
                        self._update_mcp_response_info_in_session(
                            session_id, response_body, False
                        )

                    yield f"{buffer}\n{stripped_line}\n\n"
                    buffer = ""
                else:
                    buffer = stripped_line

        # Build response headers
        response_headers = {"X-Proxied-By": "mcp-gateway", **response.headers}
        if MCP_SESSION_ID_HEADER not in response.headers:
            response_headers[MCP_SESSION_ID_HEADER] = session_id

        return StreamingResponse(
            event_generator(),
            media_type=CONTENT_TYPE_EVENT_STREAM,
            headers=response_headers,
        )

    def _update_mcp_response_info_in_session(
        self, session_id: str, response_body: dict, is_json_response: bool
    ) -> None:
        """Update MCP response info in session metadata."""
        session = self.session_store.get_session(session_id)
        self.update_mcp_server_in_session_metadata(session, response_body)
        session.attributes.metadata["is_stateless_http_server"] = session_id.startswith(
            INVARIANT_SESSION_ID_PREFIX
        )
        session.attributes.metadata["server_response_type"] = (
            "json" if is_json_response else "sse"
        )

    def _get_headers_for_mcp_post_and_delete(self, request: Request) -> dict:
        """Get filtered headers for MCP server requests."""
        filtered_headers = {}
        for k, v in request.headers.items():
            if k.startswith(MCP_CUSTOM_HEADER_PREFIX):
                filtered_headers[k.removeprefix(MCP_CUSTOM_HEADER_PREFIX)] = v
            if k.lower() in MCP_SERVER_POST_AND_DELETE_HEADERS and not (
                k.lower() == MCP_SESSION_ID_HEADER
                and v.startswith(INVARIANT_SESSION_ID_PREFIX)
            ):
                filtered_headers[k] = v
        return filtered_headers

    def _get_session_id(self, request: Request) -> str:
        """Extract session ID from request headers."""
        session_id = request.headers.get(MCP_SESSION_ID_HEADER)
        if not session_id:
            raise HTTPException(status_code=400, detail="Missing mcp-session-id header")
        return session_id

    def _get_mcp_server_endpoint(self, request: Request) -> str:
        """Get MCP server endpoint URL."""
        return self.get_mcp_server_base_url(request) + "/mcp/"

    def _is_initialization_request(self, request_data: dict[str, Any]) -> bool:
        """Check if request is an initialization request."""
        return (
            request_data.get("method") in ["initialize", "notifications/initialized"]
            and "jsonrpc" in request_data
        )