fix: map KeyError->404 and ValueError->400 for MCP clients
Tool functions raise KeyError for unknown ids and ValueError for invalid
state transitions. Without handlers, FastAPI returned a bare 500
"Internal Server Error" to MCP clients — useless for an agent that needs
to know whether the id was wrong vs the cluster was down.
KeyError -> 404 with the message ("Unknown job_id: foo").
ValueError -> 400 with the message (e.g. "pending_id p_xyz is in status
'CANCELLED', not PENDING").
This commit is contained in:
@@ -3,7 +3,8 @@
|
|||||||
@Time :2026/6/24
|
@Time :2026/6/24
|
||||||
@Author :tao.chen
|
@Author :tao.chen
|
||||||
"""
|
"""
|
||||||
from fastapi import FastAPI
|
from fastapi import FastAPI, Request
|
||||||
|
from fastapi.responses import JSONResponse
|
||||||
|
|
||||||
from spark_executor.tools.connections import (
|
from spark_executor.tools.connections import (
|
||||||
delete_connection,
|
delete_connection,
|
||||||
@@ -34,6 +35,23 @@ from spark_executor.tools.submit import (
|
|||||||
app = FastAPI(title="Spark Executor", version="0.0.1", description="Spark Executor MCP Server")
|
app = FastAPI(title="Spark Executor", version="0.0.1", description="Spark Executor MCP Server")
|
||||||
|
|
||||||
|
|
||||||
|
# --- Exception handlers: translate tool-layer errors into proper HTTP statuses ---
|
||||||
|
#
|
||||||
|
# Tool functions raise KeyError for "unknown id" (job_id, pending_id, connection
|
||||||
|
# name) and ValueError for invalid state transitions (e.g. confirming a
|
||||||
|
# CANCELLED pending). Without these handlers FastAPI would map them to a bare
|
||||||
|
# 500 "Internal Server Error" which is useless to MCP clients.
|
||||||
|
|
||||||
|
@app.exception_handler(KeyError)
|
||||||
|
async def _keyerror_handler(_request: Request, exc: KeyError) -> JSONResponse:
|
||||||
|
return JSONResponse(status_code=404, content={"detail": str(exc)})
|
||||||
|
|
||||||
|
|
||||||
|
@app.exception_handler(ValueError)
|
||||||
|
async def _valueerror_handler(_request: Request, exc: ValueError) -> JSONResponse:
|
||||||
|
return JSONResponse(status_code=400, content={"detail": str(exc)})
|
||||||
|
|
||||||
|
|
||||||
@app.get("/health")
|
@app.get("/health")
|
||||||
def health_check():
|
def health_check():
|
||||||
return {"status": "ok"}
|
return {"status": "ok"}
|
||||||
|
|||||||
@@ -110,3 +110,69 @@ def test_list_connections_works_with_empty_body():
|
|||||||
r = c.post("/list_connections", json={})
|
r = c.post("/list_connections", json={})
|
||||||
assert r.status_code == 200
|
assert r.status_code == 200
|
||||||
assert r.json() == []
|
assert r.json() == []
|
||||||
|
|
||||||
|
|
||||||
|
# --- Exception handlers: KeyError -> 404, ValueError -> 400 ---
|
||||||
|
|
||||||
|
def test_unknown_job_id_returns_404():
|
||||||
|
c = TestClient(app)
|
||||||
|
r = c.post("/get_job_status", json={"job_id": "missing"})
|
||||||
|
assert r.status_code == 404
|
||||||
|
assert "missing" in r.json()["detail"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_unknown_pending_id_returns_404():
|
||||||
|
c = TestClient(app)
|
||||||
|
r = c.post("/get_pending_job", json={"pending_id": "p_nope"})
|
||||||
|
assert r.status_code == 404
|
||||||
|
assert "p_nope" in r.json()["detail"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_unknown_connection_name_returns_404():
|
||||||
|
c = TestClient(app)
|
||||||
|
r = c.post("/get_connection", json={"name": "nope"})
|
||||||
|
assert r.status_code == 404
|
||||||
|
assert "nope" in r.json()["detail"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_unknown_connection_in_prepare_returns_404():
|
||||||
|
c = TestClient(app)
|
||||||
|
r = c.post(
|
||||||
|
"/prepare_submit_job",
|
||||||
|
json={"connection": "nope", "script_path": "/tmp/x.py"},
|
||||||
|
)
|
||||||
|
assert r.status_code == 404
|
||||||
|
assert "nope" in r.json()["detail"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_confirm_non_pending_returns_400():
|
||||||
|
"""A CANCELLED pending should refuse confirm; FastAPI should surface ValueError as 400."""
|
||||||
|
c = TestClient(app)
|
||||||
|
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
||||||
|
prep = c.post(
|
||||||
|
"/prepare_submit_job",
|
||||||
|
json={"connection": "prod", "script_path": "/tmp/x.py"},
|
||||||
|
).json()
|
||||||
|
pid = prep["pending_id"]
|
||||||
|
c.post("/cancel_pending_job", json={"pending_id": pid})
|
||||||
|
r = c.post("/confirm_submit_job", json={"pending_id": pid})
|
||||||
|
assert r.status_code == 400
|
||||||
|
assert "CANCELLED" in r.json()["detail"]
|
||||||
|
|
||||||
|
|
||||||
|
def test_cancel_already_submitted_returns_400():
|
||||||
|
"""Cancelling a SUBMITTED pending should refuse with ValueError -> 400."""
|
||||||
|
c = TestClient(app)
|
||||||
|
c.post("/save_connection", json={"name": "prod", "master": "yarn"})
|
||||||
|
prep = c.post(
|
||||||
|
"/prepare_submit_job",
|
||||||
|
json={"connection": "prod", "script_path": "/tmp/x.py"},
|
||||||
|
).json()
|
||||||
|
pid = prep["pending_id"]
|
||||||
|
# Simulate a SUBMITTED state by mutating the pending directly
|
||||||
|
p = pending_store.store.get(pid)
|
||||||
|
p.status = "SUBMITTED"
|
||||||
|
pending_store.store.save(p)
|
||||||
|
r = c.post("/cancel_pending_job", json={"pending_id": pid})
|
||||||
|
assert r.status_code == 400
|
||||||
|
assert "SUBMITTED" in r.json()["detail"]
|
||||||
|
|||||||
Reference in New Issue
Block a user