fixup! 💩(mcp) add a local MCP server configuration

Author: qbey

src/mcp_server/docs_mcp_server/mcp_server.py

@@ -1,21 +1,32 @@
 """The core of the MCP server for the Docs API."""
 
 import asyncio
+from contextlib import asynccontextmanager
 
 import httpx
 from fastmcp import FastMCP
+from starlette.requests import Request
+from starlette.responses import PlainTextResponse
 
 from . import settings, utils
 from .auth.forwarder import HeaderForwarderAuthentication
 from .auth.token import UserTokenAuthentication
 
-# Create a server instance from the OpenAPI spec
-mcp = FastMCP(name="Docs MCP Server")
+# Global variable for the API client
+api_client = None
 
 
-def setup_mcp_server():
-    """Configure the tools and resources for the MCP server."""
-    print("#" * 70)
+@asynccontextmanager
+async def lifespan(app: FastMCP):
+    """
+    Manage the lifecycle of resources for the MCP server.
+
+    This context manager initializes resources on startup and
+    ensures proper cleanup on shutdown.
+    """
+    global api_client
+
+    # Initialize API client on startup
     if settings.DOCS_API_TOKEN:
         print("Setting up Docs MCP Server with user token authentication")
         api_client = httpx.AsyncClient(
@@ -24,8 +35,23 @@ def setup_mcp_server():
         )
     else:
         print("Setting up Docs MCP Server with header forwarder authentication")
-        api_client = httpx.AsyncClient(base_url=settings.DOCS_API_URL, auth=HeaderForwarderAuthentication())
-    print("#" * 70)
+        api_client = httpx.AsyncClient(
+            base_url=settings.DOCS_API_URL, auth=HeaderForwarderAuthentication()
+        )
+
+    print("API client initialized")
+    yield
+    # Cleanup resources on shutdown
+    print("Closing API client")
+    await api_client.aclose()
+
+
+# Create a server instance from the OpenAPI spec with lifespan
+mcp = FastMCP(name="Docs MCP Server", lifespan=lifespan)
+
+
+def setup_mcp_server():
+    """Configure the tools and resources for the MCP server."""
 
     @mcp.tool()
     async def create_document_tool(document_title: str, document_content: str) -> None:
@@ -37,6 +63,9 @@ async def create_document_tool(document_title: str, document_content: str) -> No
             document_content: The content of the document (required)
 
         """
+        if api_client is None:
+            raise RuntimeError("API client is not initialized")
+
         # Get current user information
         user_response = await api_client.get("/api/v1.0/users/me/")
         user_response.raise_for_status()
@@ -57,6 +86,11 @@ async def create_document_tool(document_title: str, document_content: str) -> No
         )
         create_response.raise_for_status()
 
+    # Add health check endpoint
+    @mcp.custom_route("/health", methods=["GET"])
+    async def health_check(request: Request) -> PlainTextResponse:
+        return PlainTextResponse("OK")
+
 
 if __name__ == "__main__":
     setup_mcp_server()