From cc7eafaa16912b3952c0a8a7a41ff600510f3c5e Mon Sep 17 00:00:00 2001
From: Claude <claude@anthropic.com>
Date: Wed, 24 Jun 2026 14:50:51 +0800
Subject: [PATCH] feat: register 12 MCP tools on FastAPI app

---
 spark_executor/server.py             | 107 ++++++++++++++++++++++++++-
 tests/integration/__init__.py        |   5 ++
 tests/integration/test_mcp_routes.py |  33 +++++++++
 3 files changed, 143 insertions(+), 2 deletions(-)
 create mode 100644 tests/integration/__init__.py
 create mode 100644 tests/integration/test_mcp_routes.py

diff --git a/spark_executor/server.py b/spark_executor/server.py
index 6b69e74..fe18745 100644
--- a/spark_executor/server.py
+++ b/spark_executor/server.py
@@ -1,12 +1,115 @@
 # coding=utf-8
 """
 @Time   :2026/6/24
-@Author :tao.chen 
+@Author :tao.chen
 """
 from fastapi import FastAPI
 
+from spark_executor.tools.connections import (
+    delete_connection,
+    get_connection,
+    list_connections,
+    save_connection,
+)
+from spark_executor.tools.kill import kill_job
+from spark_executor.tools.logs import get_job_logs
+from spark_executor.tools.status import get_job_status
+from spark_executor.tools.submit import (
+    cancel_pending_job,
+    confirm_submit_job,
+    get_pending_job,
+    list_pending_jobs,
+    prepare_submit_job,
+)
+
 app = FastAPI(title="Spark Executor", version="0.0.1", description="Spark Executor MCP Server")
 
+
 @app.get("/health")
 def health_check():
-    return {"status": "ok"}
\ No newline at end of file
+    return {"status": "ok"}
+
+
+# MCP tool routes. fastapi-mcp discovers these and registers them as MCP tools.
+
+# --- Pending submission flow (two-step submit) ---
+
+@app.post("/prepare_submit_job")
+def _prepare_submit_job(connection: str, script_path: str, queue: str = "default",
+                        executor_memory: str = "4G", executor_cores: int = 2,
+                        num_executors: int = 2):
+    return prepare_submit_job(
+        connection=connection,
+        script_path=script_path,
+        queue=queue,
+        executor_memory=executor_memory,
+        executor_cores=executor_cores,
+        num_executors=num_executors,
+    )
+
+
+@app.post("/confirm_submit_job")
+def _confirm_submit_job(pending_id: str):
+    return confirm_submit_job(pending_id=pending_id)
+
+
+@app.post("/list_pending_jobs")
+def _list_pending_jobs():
+    return list_pending_jobs()
+
+
+@app.post("/get_pending_job")
+def _get_pending_job(pending_id: str):
+    return get_pending_job(pending_id)
+
+
+@app.post("/cancel_pending_job")
+def _cancel_pending_job(pending_id: str):
+    return cancel_pending_job(pending_id)
+
+
+# --- Spark job tools ---
+
+@app.post("/get_job_status")
+def _get_job_status(job_id: str):
+    return get_job_status(job_id)
+
+
+@app.post("/get_job_logs")
+def _get_job_logs(job_id: str, tail_chars: int = 5000):
+    return get_job_logs(job_id, tail_chars=tail_chars)
+
+
+@app.post("/kill_job")
+def _kill_job(job_id: str):
+    return kill_job(job_id)
+
+
+# --- Connection management tools ---
+
+@app.post("/save_connection")
+def _save_connection(name: str, master: str, deploy_mode: str = "cluster",
+                     yarn_rm_url: str | None = None,
+                     spark_conf: dict[str, str] | None = None):
+    return save_connection(
+        name=name,
+        master=master,
+        deploy_mode=deploy_mode,
+        yarn_rm_url=yarn_rm_url,
+        spark_conf=spark_conf,
+    )
+
+
+@app.post("/list_connections")
+def _list_connections():
+    return list_connections()
+
+
+@app.post("/get_connection")
+def _get_connection(name: str):
+    return get_connection(name)
+
+
+@app.post("/delete_connection")
+def _delete_connection(name: str):
+    return delete_connection(name)
diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py
new file mode 100644
index 0000000..36f7e1b
--- /dev/null
+++ b/tests/integration/__init__.py
@@ -0,0 +1,5 @@
+# coding=utf-8
+"""
+@Time   :2026/6/24
+@Author :tao.chen
+"""
diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py
new file mode 100644
index 0000000..dd1ded9
--- /dev/null
+++ b/tests/integration/test_mcp_routes.py
@@ -0,0 +1,33 @@
+# coding=utf-8
+from fastapi.testclient import TestClient
+
+from spark_executor.server import app
+
+
+def test_health_still_present():
+    c = TestClient(app)
+    r = c.get("/health")
+    assert r.status_code == 200
+    assert r.json() == {"status": "ok"}
+
+
+def test_twelve_tool_routes_registered():
+    paths = {r.path for r in app.routes}
+    for path in (
+        # pending-submission flow (5)
+        "/prepare_submit_job",
+        "/confirm_submit_job",
+        "/list_pending_jobs",
+        "/get_pending_job",
+        "/cancel_pending_job",
+        # job lifecycle (3)
+        "/get_job_status",
+        "/get_job_logs",
+        "/kill_job",
+        # connection management (4)
+        "/save_connection",
+        "/list_connections",
+        "/get_connection",
+        "/delete_connection",
+    ):
+        assert path in paths, f"missing MCP tool route: {path}"