tao.chen/mcp-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(mcp): rename generate_job_file to write_job_file to match what it does

Browse Source 
The MCP tool was named `generate_job_file` from Stage 2 but it does
NOT generate PySpark code — the calling LLM writes the code in its own
context, and this tool only persists it to a file under
SPARK_EXECUTOR_JOBS_DIR so `spark-submit` can see it. The misleading
`generate_` prefix sent agents (and humans) looking for a code
generator that doesn't exist.

This commit folds three related polish changes into one (split later
with rebase -i if you want them as separate history):

  1. The rename itself:
     - `tools/generate.py`  →  `tools/write_job.py`
     - `generate_job_file`  →  `write_job_file`
     - `GenerateJobFileRequest`  →  `WriteJobFileRequest`
     - `/generate_job_file` route  →  `/write_job_file`
     - `operation_id="generate_job_file"`  →  `operation_id="write_job_file"`
     The internal helper `core.job_writer.write_job_file` (which just
     writes bytes to disk with no SQL guard) is imported with an
     `_write_to_disk` alias to avoid the name collision with the
     MCP-exposed function in the same module.
     The description for the tool now explicitly states 'this tool
     does NOT generate PySpark code. The calling LLM is expected to
     have already written the code; this tool only persists it.'

  2. Skill for LLM agents operating the service
     (`docs/superpowers/skills/spark-executor-mcp-operate/SKILL.md`,
     449 lines). Covers the 16 tools, the two-step prepare/confirm
     flow, the dual-ID contract (job_id vs application_id), the
     PendingSubmission state machine, the Connection profile, the
     job-file workflow, the error reference, common pitfalls, and a
     full end-to-end word-count example.

  3. Default `executor_memory` lowered 4G → 2G
     (`_DEFAULTS_TO_CONFIRM` in `server.py`). Mirrors the matching
     change in `test_mcp_routes.py` and the 5 unit tests that
     reference the default. Aligns with the lighter workloads the
     service is sized for in its current container profile.

Also tracked in git for the first time:
  - `docs/superpowers/plans/2026-06-24-spark-executor-mcp.md`
    (the original Stage 1/2/3 design plan, updated to use the new
    tool name throughout).

Test rename:
  - `tests/unit/test_generate_tool.py`  →  `test_write_job_tool.py`
  - the new test file picks up an extra assertion that the SQL guard
    rejects a `DROP TABLE` statement at write time.

243 tests pass (was 242; +1 new SQL-guard assertion). Zero regressions.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-06-29 19:13:29 +08:00
co-authored by Claude Fable 5
parent 0fef77c0b8
commit 523e6a9c76
18 changed files with 3661 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
spark_executor/server.py
+19 -14
View File
@@ -12,14 +12,14 @@ from spark_executor.tools.connections import (
list_connections,
save_connection,
)
from spark_executor.tools.generate import generate_job_file
from spark_executor.tools.write_job import write_job_file
from spark_executor.tools.job_file import read_job_file, update_job_file
from spark_executor.tools.kill import kill_job
from spark_executor.tools.logs import get_job_logs
from spark_executor.tools.requests import (
ConnectionNameRequest,
EmptyRequest,
GenerateJobFileRequest,
WriteJobFileRequest,
GetJobLogsRequest,
JobIdRequest,
PendingIdRequest,
@@ -46,7 +46,7 @@ app = FastAPI(title="Spark Executor MCP", version="0.0.1", description="Spark Ex
_DEFAULTS_TO_CONFIRM = {
"queue": "default",
"executor_memory": "4G",
"executor_memory": "2G",
"executor_cores": 2,
"num_executors": 2,
}
@@ -90,7 +90,7 @@ def health_check():
"yarn_rm_url into a PendingSubmission record and persist it. "
"Does NOT invoke spark-submit. Returns pending_id for use with "
"confirm_submit_job (the user-second-confirmation step).\n\n"
"REQUIRED PATTERN for LLM-generated code: call generate_job_file(code=...) "
"REQUIRED PATTERN for LLM-generated code: call write_job_file(code=...) "
"first, then pass the returned script_path here. Direct submission with a "
"synthetic path (one that only exists in the agent's context) will be "
"rejected with HTTP 400 — the script must exist inside the container's "
@@ -297,19 +297,24 @@ def _delete_connection(req: ConnectionNameRequest):
# --- LLM-driven PySpark generation (Stage 2) ---
@app.post(
"/generate_job_file",
operation_id="generate_job_file",
summary="Write LLM-generated PySpark code to disk",
"/write_job_file",
operation_id="write_job_file",
summary="Write LLM-authored PySpark code to disk",
description=(
"Takes a PySpark code string and writes it to a timestamped file under "
"Takes a PySpark code string the LLM has already composed in its "
"context and writes it to a timestamped file under "
"SPARK_EXECUTOR_JOBS_DIR (default ./data/jobs/). Returns the absolute "
"path for use as the script_path argument of prepare_submit_job — the "
"two-step pattern means the LLM can produce code, the user can review "
"the resulting file, and only then is the job submitted."
"two-step pattern means the LLM writes the file, the user can review "
"it (via read_job_file), and only then is the job submitted.\n\n"
"Note: this tool does NOT generate PySpark code. The calling LLM is "
"expected to have already written the code; this tool only persists "
"it. Code is also run through the SQL safety policy (SELECT/INSERT "
"only) before being written — forbidden statements cause a 400."
),
)
def _generate_job_file(req: GenerateJobFileRequest):
return generate_job_file(req.code)
def _write_job_file(req: WriteJobFileRequest):
return write_job_file(req.code)
@app.post(
@@ -318,7 +323,7 @@ def _generate_job_file(req: GenerateJobFileRequest):
summary="Read the contents of an existing PySpark script",
description=(
"Returns the text content of an existing script file at the given "
"path. Caps reads at 1 MB. Typical use: after generate_job_file "
"path. Caps reads at 1 MB. Typical use: after write_job_file "
"returns a path, call read_job_file on that path to inspect what "
"was actually written, before deciding to prepare_submit_job or "
"update_job_file."
@@ -334,7 +339,7 @@ def _read_job_file(req: ReadJobFileRequest):
summary="Overwrite an existing PySpark script with new content",
description=(
"Replaces the entire content of an existing script file. Path must "
"be under SPARK_EXECUTOR_JOBS_DIR (the dir generate_job_file writes "
"be under SPARK_EXECUTOR_JOBS_DIR (the dir write_job_file writes "
"to) — protects against overwriting host-mounted configs or other "
"non-script files. Caps writes at 1 MB. Typical use: read_job_file, "
"edit the content (LLM or human), update_job_file, then "
spark_executor/tools/generate.py
-43
View File
@@ -1,43 +0,0 @@
# coding=utf-8
"""
@Time :2026/6/24
@Author :tao.chen
"""
from common.logging import logger
from common.sql_guard import validate_pyspark_code
from spark_executor.core.job_writer import write_job_file
class SqlGuardViolation(ValueError):
"""Raised when generate_job_file receives code that violates the SQL
safety policy. -> HTTP 400 via the FastAPI ValueError handler.
"""
pass
def generate_job_file(code: str) -> dict[str, str]:
"""Write a PySpark code string to disk; return its absolute path.
Validates the code against the SQL safety policy (SELECT/INSERT only)
BEFORE writing, so the agent gets immediate feedback rather than
learning at prepare_submit_job time. Use the returned path as the
`script_path` argument of prepare_submit_job.
The output directory is controlled by the SPARK_EXECUTOR_JOBS_DIR env var
(default: ./data/jobs/).
"""
logger.debug(f"generate_job_file enter code_bytes={len(code)}")
offenses = validate_pyspark_code(code)
if offenses:
logger.warning(
f"generate_job_file rejected: SQL policy violation(s): {offenses}"
)
raise SqlGuardViolation(
f"PySpark code violates SQL safety policy. "
f"Only SELECT and INSERT statements are allowed. "
f"Found forbidden statement(s): {offenses}. "
f"Rewrite the code to use only SELECT/INSERT."
)
path = write_job_file(code)
logger.info(f"generate_job_file ok script_path={path}")
return {"script_path": path}
spark_executor/tools/job_file.py
+4 -4
View File
@@ -6,7 +6,7 @@
Read + update the contents of an existing PySpark script file. These two
tools close the review-and-edit loop:
generate_job_file(code=...) -> {script_path}
write_job_file(code=...) -> {script_path}
read_job_file(script_path=...) -> {content, path} <-- inspect
update_job_file(path, content) -> {path, bytes_written} <-- edit
prepare_submit_job(path) -> {pending_id, ...}
@@ -43,14 +43,14 @@ def _check_readable(script_path: str) -> None:
def _check_writable(script_path: str) -> None:
"""update_job_file is restricted to files under settings.jobs_dir
(the same dir generate_job_file writes to). This prevents the agent
(the same dir write_job_file writes to). This prevents the agent
from overwriting arbitrary host-mounted files or the app's own code.
"""
if not script_path or not os.path.isfile(script_path):
raise ScriptFileError(
f"script_path does not exist or is not a file: {script_path!r}. "
f"update_job_file can only edit existing files. "
f"Use generate_job_file to create a new one."
f"Use write_job_file to create a new one."
)
jobs_root = Path(settings.jobs_dir).resolve()
target = Path(script_path).resolve()
@@ -59,7 +59,7 @@ def _check_writable(script_path: str) -> None:
except ValueError:
raise ScriptFileError(
f"script_path must be under {jobs_root} (the directory "
f"generate_job_file writes to). Got {script_path!r}. "
f"write_job_file writes to). Got {script_path!r}. "
f"This restriction protects host-mounted configs and other "
f"non-script files from being overwritten by the agent."
)
spark_executor/tools/requests.py
+10 -8
View File
@@ -42,7 +42,7 @@ class PrepareSubmitJobRequest(BaseModel):
description=(
"Absolute path to the PySpark script inside the container's "
"filesystem. Must point at an existing regular file. For "
"LLM-generated code, call generate_job_file(code=...) first "
"LLM-generated code, call write_job_file(code=...) first "
"and pass the returned script_path here. For pre-existing "
"files on the host, mount them via a docker volume and pass "
"the in-container path. Returns 400 with a remediation hint "
@@ -54,8 +54,8 @@ class PrepareSubmitJobRequest(BaseModel):
description="YARN queue to submit to. Must be explicitly confirmed by the caller.",
)
executor_memory: str = Field(
default="4G",
description="Executor memory, e.g. '4G'. Must be explicitly confirmed by the caller.",
default="2G",
description="Executor memory, e.g. '2G'. Must be explicitly confirmed by the caller.",
)
executor_cores: int = Field(
default=2,
@@ -102,13 +102,15 @@ class ConnectionNameRequest(BaseModel):
name: str
class GenerateJobFileRequest(BaseModel):
class WriteJobFileRequest(BaseModel):
code: str = Field(
...,
description=(
"Full PySpark source code to write to disk. Will be passed verbatim "
"to spark-submit after the agent calls prepare_submit_job on the "
"returned path."
"Full PySpark source code to write to disk. The LLM is expected "
"to have already written this code in its own context; this tool "
"persists it to a file under SPARK_EXECUTOR_JOBS_DIR so "
"spark-submit can see it. The returned script_path is what you "
"pass to prepare_submit_job."
),
)
@@ -129,7 +131,7 @@ class UpdateJobFileRequest(BaseModel):
description=(
"Absolute path to an existing PySpark script inside the "
"container's filesystem. Must be under SPARK_EXECUTOR_JOBS_DIR "
"(the same dir generate_job_file writes to) — protects against "
"(the same dir write_job_file writes to) — protects against "
"overwriting host-mounted configs or other critical files."
),
)
spark_executor/tools/submit.py
+5 -5
View File
@@ -30,7 +30,7 @@ def _script_path_error(path: str) -> ValueError:
return ValueError(
f"script_path does not exist or is not a file: {path!r}. "
f"Two ways to fix this:\n"
f" 1. (Recommended) Call generate_job_file(code=...) first to write "
f" 1. (Recommended) Call write_job_file(code=...) first to write "
f"the PySpark source to disk, then pass the returned script_path.\n"
f" 2. If the file already exists on the host, mount it into the "
f"container (e.g. -v /host/path:/app/scripts:ro in docker run) and "
@@ -52,7 +52,7 @@ def _sql_guard_error(offenses: list[str]) -> ValueError:
f"Only SELECT and INSERT statements are allowed. "
f"Found forbidden statement(s): {offenses}. "
f"Edit the script and try again. (If the code was generated via "
f"generate_job_file, the generator should have caught this — check "
f"write_job_file, the generator should have caught this — check "
f"for f-string-based SQL injection where the static analysis can't "
f"see the runtime value.)"
)
@@ -67,7 +67,7 @@ def prepare_submit_job(
connection: str,
script_path: str,
queue: str = "default",
executor_memory: str = "4G",
executor_memory: str = "2G",
executor_cores: int = 2,
num_executors: int = 2,
app_name: str,
@@ -92,12 +92,12 @@ def prepare_submit_job(
raise KeyError(f"Unknown connection: {connection}")
# Fail fast: a non-existent path is the most common agent mistake (it
# generated the code in its own context but forgot to call
# generate_job_file first, or its path refers to the host filesystem
# write_job_file first, or its path refers to the host filesystem
# which is invisible inside the container). Better to surface this with
# a 400 + clear remediation than to let spark-submit fail later with
# an opaque FileNotFoundError -> 500.
_check_script_path(script_path)
# SQL safety: re-validate the script even though generate_job_file
# SQL safety: re-validate the script even though write_job_file
# already guards its own output. Catches files written by other means
# (host volume mounts, manual edits).
with open(script_path, encoding="utf-8") as f:
spark_executor/tools/write_job.py
+57
View File
@@ -0,0 +1,57 @@
# coding=utf-8
"""
@Time :2026/6/24
@Author :tao.chen
The MCP-exposed `write_job_file` tool. Despite the historical name
`generate_job_file`, this tool does NOT generate PySpark code — the
calling LLM writes the code, and this tool persists it to disk so
`spark-submit` can see it. The rename (2026-06-29) tightens the naming
to match what the tool actually does, eliminating the "I thought this
would generate code for me" confusion.
It also runs the SQL safety policy (SELECT/INSERT only) on the code
BEFORE writing, so the agent gets immediate feedback if it slipped a
DROP / DELETE / UPDATE past its own code-generation step.
"""
from common.logging import logger
from common.sql_guard import validate_pyspark_code
# Internal helper already named write_job_file (it just writes bytes to
# disk; no SQL guard). Import with an alias to avoid the name collision
# with the MCP-exposed function below.
from spark_executor.core.job_writer import write_job_file as _write_to_disk
class SqlGuardViolation(ValueError):
"""Raised when write_job_file receives code that violates the SQL
safety policy. -> HTTP 400 via the FastAPI ValueError handler."""
pass
def write_job_file(code: str) -> dict[str, str]:
"""Persist a PySpark code string to disk and return its absolute path.
Validates the code against the SQL safety policy (SELECT/INSERT only)
BEFORE writing, so the agent gets immediate feedback rather than
learning at prepare_submit_job time. Use the returned path as the
`script_path` argument of prepare_submit_job.
The output directory is controlled by the SPARK_EXECUTOR_JOBS_DIR env var
(default: ./data/jobs/). Each call writes to a fresh timestamped file
under that directory, so repeated calls never overwrite one another.
"""
logger.debug(f"write_job_file enter code_bytes={len(code)}")
offenses = validate_pyspark_code(code)
if offenses:
logger.warning(
f"write_job_file rejected: SQL policy violation(s): {offenses}"
)
raise SqlGuardViolation(
f"PySpark code violates SQL safety policy. "
f"Only SELECT and INSERT statements are allowed. "
f"Found forbidden statement(s): {offenses}. "
f"Rewrite the code to use only SELECT/INSERT."
)
path = _write_to_disk(code)
logger.info(f"write_job_file ok script_path={path}")
return {"script_path": path}