feat: add generate_job_file MCP tool (Stage 2 Task 22)

Exposes a single new MCP tool: generate_job_file(code) -> {script_path}.

Wiring follows the Stage 1 conventions (Pydantic body model for the
route, loguru DEBUG/INFO logging, summary+description for the startup
log, MCP arg-pattern via the body model so tools/call roundtrips long
code strings without 422-ing on FastAPI query length limits).

Flow:
  1. LLM calls generate_job_file(code=...)          -> {script_path}
  2. LLM calls prepare_submit_job(connection=...,
                                  script_path=...) -> {pending_id}
  3. User reviews the file + pending record
  4. LLM calls confirm_submit_job(pending_id=...)  -> spark-submit runs

The output directory is SPARK_EXECUTOR_JOBS_DIR (default ./data/jobs/,
gitignored, persists across container restarts via the existing
./data volume mount in docker-compose.yml).

3 new tests:
  - Unit: env-var override, default fallback in tmp cwd
  - Integration: end-to-end body call with a > FastAPI-query-limit code
    string (the canary test that would have caught the Stage 1
    query-params-422 bug)

Live MCP smoke verified: tools/list shows 14 tools (13 from Stage 1 +
the new one), tools/call generate_job_file returns the absolute path
under ./data/jobs/.
This commit is contained in:
Claude
2026-06-25 10:19:59 +08:00
parent 9340061b30
commit 994c2d67ab
5 changed files with 97 additions and 1 deletions
+19
View File
@@ -12,11 +12,13 @@ from spark_executor.tools.connections import (
list_connections,
save_connection,
)
from spark_executor.tools.generate import generate_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,
GetJobLogsRequest,
JobIdRequest,
PendingIdRequest,
@@ -201,3 +203,20 @@ def _get_connection(req: ConnectionNameRequest):
)
def _delete_connection(req: ConnectionNameRequest):
return delete_connection(req.name)
# --- LLM-driven PySpark generation (Stage 2) ---
@app.post(
"/generate_job_file",
summary="Write LLM-generated PySpark code to disk",
description=(
"Takes a PySpark code string 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."
),
)
def _generate_job_file(req: GenerateJobFileRequest):
return generate_job_file(req.code)