Infinity-API/src/api.py at 65004757e2ad9078541c6ef64a2f5b3e1dd2e62a · RocketPy-Team/Infinity-API · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
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 JSONResponse, RedirectResponse

from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.requests import RequestsInstrumentor

from src import logger, parse_error
from src.mcp.server import build_mcp
from src.routes import environment, flight, motor, rocket
from src.utils import RocketPyGZipMiddleware

log = logging.getLogger(__name__)


# --- REST application -------------------------------------------------------

rest_app = FastAPI(
    title="Infinity API",
    swagger_ui_parameters={
        "defaultModelsExpandDepth": 0,
        "syntaxHighlight.theme": "obsidian",
    },
)

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(rest_app)
RequestsInstrumentor().instrument()

# Compress responses above 1KB
rest_app.add_middleware(RocketPyGZipMiddleware, minimum_size=1000)


def custom_openapi():
    if rest_app.openapi_schema:
        return rest_app.openapi_schema
    openapi_schema = get_openapi(
        title="RocketPy Infinity-API",
        version="3.0.0",
        description=(
            "<p style='font-size: 18px;'>RocketPy Infinity-API is a RESTful Open API for RocketPy, a rocket flight simulator.</p>"
            "<br/>"
            "<button style='background-color: #4CAF50; color: white; border: none; padding: 10px 20px; text-align: center; text-decoration: none; display: inline-block; font-size: 16px; margin: 4px 2px; cursor: pointer;'>"
            "<a href='https://api.rocketpy.org/docs' style='color: white; text-decoration: none;'>Swagger UI</a>"
            "</button>"
            "<br/>"
            "<button style='background-color: #008CBA; color: white; border: none; padding: 10px 20px; text-align: center; text-decoration: none; display: inline-block; font-size: 16px; margin: 4px 2px; cursor: pointer;'>"
            "<a href='https://api.rocketpy.org/redoc' style='color: white; text-decoration: none;'>ReDoc</a>"
            "</button>"
            "<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=rest_app.routes,
    )
    openapi_schema["info"]["x-logo"] = {
        "url": "https://raw.githubusercontent.com/RocketPy-Team/RocketPy/master/docs/static/RocketPy_Logo_black.png"
    }
    rest_app.openapi_schema = openapi_schema
    return rest_app.openapi_schema


rest_app.openapi = custom_openapi


# Main
@rest_app.get("/", include_in_schema=False)
async def main_page():
    """Redirect to API docs."""
    return RedirectResponse(url="/redoc")


# Additional routes
@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
@rest_app.exception_handler(RequestValidationError)
async def validation_exception_handler(
    request: Request, exc: RequestValidationError
):
    exc_str = parse_error(exc)
    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"]