fix mcp<>rest bridge

Author: GabrielBarberini

README.md

@@ -31,7 +31,8 @@ $ touch .env && echo MONGODB_CONNECTION_STRING="$ConnectionString" > .env
 - Prod: `gunicorn -k uvicorn.workers.UvicornWorker src:app -b 0.0.0.0:3000`
 
 ## MCP Server
-- Infinity API automatically serves an MCP bridge at `/mcp` alongside the REST endpoints.
+- The MCP bridge is mounted directly on the FastAPI app and is available at `/mcp` alongside the REST API.
+- No extra process is required: `uvicorn src:app` serves both the REST routes and the MCP transport.
 
 ## Project structure
 ```

src/api.py

@@ -1,41 +1,49 @@
+from __future__ import annotations
+
+import logging
+
+from contextlib import asynccontextmanager
 from fastapi import FastAPI, Request, status
 from fastapi.exceptions import RequestValidationError
 from fastapi.openapi.utils import get_openapi
-from fastapi.responses import RedirectResponse, JSONResponse
+from fastapi.responses import JSONResponse, RedirectResponse
 
 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 from opentelemetry.instrumentation.requests import RequestsInstrumentor
 
 from src import logger, parse_error
-from src.routes import flight, environment, motor, rocket
-from src.utils import RocketPyGZipMiddleware
 from src.mcp.server import build_mcp
+from src.routes import environment, flight, motor, rocket
+from src.utils import RocketPyGZipMiddleware
 
-app = FastAPI(
+log = logging.getLogger(__name__)
+
+
+# --- REST application -------------------------------------------------------
+
+rest_app = FastAPI(
     title="Infinity API",
     swagger_ui_parameters={
         "defaultModelsExpandDepth": 0,
         "syntaxHighlight.theme": "obsidian",
     },
 )
-app.include_router(flight.router)
-app.include_router(environment.router)
-app.include_router(motor.router)
-app.include_router(rocket.router)
 
-_mcp_server = build_mcp(app)
-app.mount('/mcp', _mcp_server.http_app())
+rest_app.include_router(flight.router)
+rest_app.include_router(environment.router)
+rest_app.include_router(motor.router)
+rest_app.include_router(rocket.router)
 
-FastAPIInstrumentor.instrument_app(app)
+FastAPIInstrumentor.instrument_app(rest_app)
 RequestsInstrumentor().instrument()
 
 # Compress responses above 1KB
-app.add_middleware(RocketPyGZipMiddleware, minimum_size=1000)
+rest_app.add_middleware(RocketPyGZipMiddleware, minimum_size=1000)
 
 
 def custom_openapi():
-    if app.openapi_schema:
-        return app.openapi_schema
+    if rest_app.openapi_schema:
+        return rest_app.openapi_schema
     openapi_schema = get_openapi(
         title="RocketPy Infinity-API",
         version="3.0.0",
@@ -52,40 +60,70 @@ def custom_openapi():
             "<p>Create, manage, and simulate rocket flights, environments, rockets, and motors.</p>"
             "<p>Please report any bugs at <a href='https://github.com/RocketPy-Team/infinity-api/issues/new/choose' style='text-decoration: none; color: #008CBA;'>GitHub Issues</a></p>"
         ),
-        routes=app.routes,
+        routes=rest_app.routes,
     )
     openapi_schema["info"]["x-logo"] = {
         "url": "https://raw.githubusercontent.com/RocketPy-Team/RocketPy/master/docs/static/RocketPy_Logo_black.png"
     }
-    app.openapi_schema = openapi_schema
-    return app.openapi_schema
+    rest_app.openapi_schema = openapi_schema
+    return rest_app.openapi_schema
 
 
-app.openapi = custom_openapi
+rest_app.openapi = custom_openapi
 
 
 # Main
-@app.get("/", include_in_schema=False)
+@rest_app.get("/", include_in_schema=False)
 async def main_page():
-    """
-    Redirects to API docs.
-    """
+    """Redirect to API docs."""
     return RedirectResponse(url="/redoc")
 
 
 # Additional routes
-@app.get("/health", status_code=status.HTTP_200_OK, include_in_schema=False)
+@rest_app.get(
+    "/health", status_code=status.HTTP_200_OK, include_in_schema=False
+)
 async def __perform_healthcheck():
     return {"health": "Everything OK!"}
 
 
 # Global exception handler
-@app.exception_handler(RequestValidationError)
+@rest_app.exception_handler(RequestValidationError)
 async def validation_exception_handler(
     request: Request, exc: RequestValidationError
 ):
     exc_str = parse_error(exc)
-    logger.error(f"{request}: {exc_str}")
+    logger.error("%s: %s", request, exc_str)
     return JSONResponse(
         content=exc_str, status_code=status.HTTP_422_UNPROCESSABLE_ENTITY
     )
+
+
+# --- MCP server mounted under /mcp ------------------------------------------
+mcp_app = build_mcp(rest_app).http_app(path="/")
+
+
+def _combine_lifespans(rest_lifespan, mcp_lifespan):
+    """Combine FastAPI and MCP lifespans."""
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI):
+        async with rest_lifespan(app):
+            async with mcp_lifespan(app):
+                yield
+
+    return lifespan
+
+
+app = FastAPI(
+    docs_url=None,
+    redoc_url=None,
+    openapi_url=None,
+    lifespan=_combine_lifespans(
+        rest_app.router.lifespan_context, mcp_app.lifespan
+    ),
+)
+app.mount("/mcp", mcp_app)
+app.mount("/", rest_app)
+
+__all__ = ["app", "rest_app"]

tests/unit/test_mcp/test_mcp_server.py

@@ -1,51 +1,48 @@
 from __future__ import annotations
 
-
 from unittest.mock import MagicMock, patch
 
 import pytest
-
 from fastmcp.client import Client
 from fastapi.routing import APIRoute
 
-from src.api import app
+from src.api import app, rest_app
 from src.mcp.server import build_mcp
 
 
 @pytest.fixture(autouse=True)
 def reset_mcp_state():
-    if hasattr(app.state, 'mcp'):
-        delattr(app.state, 'mcp')
+    if hasattr(rest_app.state, 'mcp'):
+        delattr(rest_app.state, 'mcp')
     yield
-    if hasattr(app.state, 'mcp'):
-        delattr(app.state, 'mcp')
+    if hasattr(rest_app.state, 'mcp'):
+        delattr(rest_app.state, 'mcp')
 
 
 def test_build_mcp_uses_fastapi_adapter():
     mock_mcp = MagicMock()
     with patch(
         'src.mcp.server.FastMCP.from_fastapi', return_value=mock_mcp
     ) as mock_factory:
-        result = build_mcp(app)
+        result = build_mcp(rest_app)
         assert result is mock_mcp
-        mock_factory.assert_called_once_with(app, name=app.title)
-        # Subsequent calls reuse cached server
-        again = build_mcp(app)
+        mock_factory.assert_called_once_with(rest_app, name=rest_app.title)
+        again = build_mcp(rest_app)
         assert again is mock_mcp
         mock_factory.assert_called_once()
 
 
 @pytest.mark.asyncio
 async def test_mcp_tools_cover_registered_routes():
-    mcp_server = build_mcp(app)
+    mcp_server = build_mcp(rest_app)
 
     async with Client(mcp_server) as client:
         tools = await client.list_tools()
 
     tool_by_name = {tool.name: tool for tool in tools}
 
     expected = {}
-    for route in app.routes:
+    for route in rest_app.routes:
         if not isinstance(route, APIRoute) or not route.include_in_schema:
             continue
         tag = route.tags[0].lower()
@@ -60,7 +57,23 @@ async def test_mcp_tools_cover_registered_routes():
         schema = tool_by_name[tool_name].inputSchema or {}
         required = set(schema.get('required', []))
         path_params = {param.name for param in route.dependant.path_params}
-        # Path parameters must be represented as required MCP tool arguments
         assert path_params.issubset(
             required
         ), f"{tool_name} missing path params {path_params - required}"
+
+
+@pytest.mark.asyncio
+async def test_combined_app_serves_rest_and_mcp(monkeypatch):
+    monkeypatch.setattr('src.mcp.server.FastMCP.from_fastapi', MagicMock())
+    build_mcp(rest_app)
+
+    from httpx import ASGITransport, AsyncClient
+
+    transport = ASGITransport(app=app)
+    async with AsyncClient(
+        transport=transport, base_url='http://test'
+    ) as client:
+        resp_rest = await client.get('/health')
+        assert resp_rest.status_code == 200
+        resp_docs = await client.get('/docs')
+        assert resp_docs.status_code == 200