From 15e9a75bc72b0d1a12f96292a183fca190aafbc2 Mon Sep 17 00:00:00 2001
From: Claude <claude@anthropic.com>
Date: Thu, 25 Jun 2026 12:33:14 +0800
Subject: [PATCH] refactor: make lifespan extensible for additional MCP
 services
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The lifespan was hardcoded to mount just one MCP service
(spark_executor_app at /spark-executor-mcp). To add a second MCP
service, you'd have to edit the lifespan function body — error-prone
and not obvious where to add a new entry.

Refactor:
  - New McpService dataclass (name, app, mount_path)
  - MCP_SERVICES list at module top — adding a new service is one append
  - Lifespan iterates MCP_SERVICES; the per-service mount + tool-list
    logic is now a helper (_log_mcp_tools)
  - Tool-list log now prefixes the service name:
      'MCP tools registered (spark_executor, 14):'
    instead of just '(14):', so multi-service startups are unambiguous

116/116 still pass. Live verified: gunicorn 2-worker startup logs
each service with its tool list and the source app version.

Future services just need a McpService entry — no lifespan edit. The
MCP_SERVICES list has a docstring example showing the one-liner.
---
 main.py | 71 ++++++++++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 55 insertions(+), 16 deletions(-)

diff --git a/main.py b/main.py
index 323f4e6..b88fc3f 100644
--- a/main.py
+++ b/main.py
@@ -4,6 +4,7 @@
 @Author :tao.chen
 """
 from contextlib import asynccontextmanager
+from dataclasses import dataclass
 
 from fastapi import FastAPI
 
@@ -12,30 +13,63 @@ from common.logging import logger
 from spark_executor import app as spark_executor_app
 
 
-@asynccontextmanager
-async def lifespan(app: "FastAPI"):
-    logger.info(f"Startup {app.title}")
-    spark_executor_mcp = init_mcp_server(spark_executor_app)
-    spark_executor_mcp.mount_http(app, mount_path="/spark-executor-mcp")
-    logger.info(
-        f"MCP server mounted at /spark-executor-mcp "
-        f"(source: {spark_executor_app.title} v{spark_executor_app.version})"
-    )
+# --- MCP service registry ---
+#
+# Each McpService entry becomes one mount under the root app. The lifespan
+# iterates this list at startup, so adding a new MCP service is a one-line
+# append here — no edits to the lifespan function needed.
+#
+# Example (future):
+#   McpService(name="metrics",  app=metrics_app,  mount_path="/metrics-mcp"),
+#   McpService(name="catalog",  app=catalog_app,  mount_path="/catalog-mcp"),
 
-    # Enumerate every MCP tool + its description so operators can see at a
-    # glance what the server exposes (useful for debugging / change review).
-    tools = spark_executor_mcp.tools
-    logger.info(f"MCP tools registered ({len(tools)}):")
+
+@dataclass(frozen=True)
+class McpService:
+    """One MCP-exposed FastAPI sub-app to be mounted under the root."""
+
+    name: str          # short identifier used in startup logs (e.g. "spark_executor")
+    app: FastAPI       # the sub-app whose routes will be exposed as MCP tools
+    mount_path: str    # HTTP path under the root app (e.g. "/spark-executor-mcp")
+
+
+MCP_SERVICES: list[McpService] = [
+    McpService(
+        name="spark_executor",
+        app=spark_executor_app,
+        mount_path="/spark-executor-mcp",
+    ),
+]
+
+
+def _log_mcp_tools(mcp_server, service: McpService) -> None:
+    """Log every tool the MCP server exposes, with its description.
+
+    fastapi-mcp appends an auto-generated "### Responses: ..." block to
+    every tool description; we strip that here so the startup log stays
+    scannable. Each tool is logged on one line: `  N. tool_name  -  desc`.
+    """
+    tools = mcp_server.tools
+    logger.info(f"MCP tools registered ({service.name}, {len(tools)}):")
     for i, tool in enumerate(tools, 1):
         desc = (tool.description or "").strip()
-        # fastapi-mcp appends an auto-generated "### Responses: ..." block to
-        # every tool description; drop it so the startup log stays scannable.
         if "### Responses" in desc:
             desc = desc.split("### Responses", 1)[0].rstrip()
-        # Collapse internal newlines / extra spaces into one line.
         desc = " ".join(desc.split())
         logger.info(f"  {i:>2}. {tool.name}  -  {desc}")
 
+
+@asynccontextmanager
+async def lifespan(app: "FastAPI"):
+    logger.info(f"Startup {app.title}")
+    for svc in MCP_SERVICES:
+        logger.info(
+            f"Mounting MCP service {svc.name!r} at {svc.mount_path} "
+            f"(source: {svc.app.title} v{svc.app.version})"
+        )
+        mcp = init_mcp_server(svc.app)
+        mcp.mount_http(app, mount_path=svc.mount_path)
+        _log_mcp_tools(mcp, svc)
     yield
     logger.info(f"Shutdown {app.title}")
 
@@ -43,6 +77,11 @@ async def lifespan(app: "FastAPI"):
 app = FastAPI(title="Main App", lifespan=lifespan)
 
 
+@app.get("/health")
+def health_check():
+    return {"status": "ok"}
+
+
 if __name__ == "__main__":
     import uvicorn