feat: add read_job_file and update_job_file MCP tools

Closes the review-and-edit loop for LLM-generated PySpark code:

  generate_job_file(code=...)      ->  {script_path}
  read_job_file(script_path=...)   ->  {content, path, size}
  update_job_file(path, content)   ->  {path, bytes_written}
  prepare_submit_job(path)         ->  {pending_id, ...}

Or, in a single edit cycle:
  1. generate    (LLM writes initial draft)
  2. read        (LLM or human inspects)
  3. update      (overwrite with edited version)
  4. prepare     (submit for two-step confirmation)

Safety:
  - read_job_file has no path restriction (read-only; useful for
    inspecting any file the agent can see: scripts, logs/, hadoop-conf/)
  - update_job_file is sandboxed to settings.jobs_dir (the same dir
    generate_job_file writes to). Rejects paths outside that tree,
    including ../-traversal attempts. This protects host-mounted
    configs (/etc/passwd, hadoop-conf/*) from being overwritten by
    the agent.
  - 1 MB cap on both reads and writes so MCP responses stay bounded.

Pydantic body models (ReadJobFileRequest, UpdateJobFileRequest) follow
the Stage 1 pattern so tools/call roundtrips long code strings without
the FastAPI query-length 422.

Tests (13 new):
  - 9 unit tests: read success/missing/empty/dir, update success/outside/
    relative-escape/missing/oversize/1mb+1, full edit cycle round-trip
  - 4 integration tests: read via MCP, missing file 400, write+readback
    via MCP, outside-jobs_dir rejection via MCP

163/146 still pass. Live verified end-to-end: generate -> read v1
-> update -> read v2; update /etc/passwd correctly 400'd with
'script_path must be under ... data/jobs/'.
This commit is contained in:
Claude
2026-06-25 13:18:39 +08:00
parent a1710ec4db
commit 35bff11edf
5 changed files with 386 additions and 2 deletions
+34
View File
@@ -13,6 +13,7 @@ from spark_executor.tools.connections import (
save_connection,
)
from spark_executor.tools.generate import generate_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 (
@@ -23,7 +24,9 @@ from spark_executor.tools.requests import (
JobIdRequest,
PendingIdRequest,
PrepareSubmitJobRequest,
ReadJobFileRequest,
SaveConnectionRequest,
UpdateJobFileRequest,
)
from spark_executor.tools.status import get_job_status
from spark_executor.tools.submit import (
@@ -226,3 +229,34 @@ def _delete_connection(req: ConnectionNameRequest):
)
def _generate_job_file(req: GenerateJobFileRequest):
return generate_job_file(req.code)
@app.post(
"/read_job_file",
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 "
"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."
),
)
def _read_job_file(req: ReadJobFileRequest):
return read_job_file(req.script_path)
@app.post(
"/update_job_file",
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 "
"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 "
"prepare_submit_job with the same path."
),
)
def _update_job_file(req: UpdateJobFileRequest):
return update_job_file(req.script_path, req.content)
+111
View File
@@ -0,0 +1,111 @@
# coding=utf-8
"""
@Time :2026/6/24
@Author :tao.chen
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}
read_job_file(script_path=...) -> {content, path} <-- inspect
update_job_file(path, content) -> {path, bytes_written} <-- edit
prepare_submit_job(path) -> {pending_id, ...}
Safety:
- read_job_file: any existing regular file. Path-existence only.
- update_job_file: must be under SPARK_EXECUTOR_JOBS_DIR
(settings.jobs_dir) so the agent cannot overwrite host-mounted
configs or arbitrary files on the container FS.
- 1 MB cap on both read and write payloads to keep MCP responses bounded.
"""
import os
from pathlib import Path
from common.config import settings
from common.logging import logger
MAX_FILE_BYTES = 1 * 1024 * 1024 # 1 MB
class ScriptFileError(ValueError):
"""Raised when read/update fails. -> HTTP 400 via the FastAPI ValueError
handler in server.py.
"""
pass
def _check_readable(script_path: str) -> None:
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}"
)
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
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."
)
jobs_root = Path(settings.jobs_dir).resolve()
target = Path(script_path).resolve()
try:
target.relative_to(jobs_root)
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"This restriction protects host-mounted configs and other "
f"non-script files from being overwritten by the agent."
)
def read_job_file(script_path: str) -> dict[str, object]:
"""Return the text content of an existing script file.
Caps the read at 1 MB to keep MCP responses bounded; raises
ScriptFileError (-> 400) if the file is missing or too large.
"""
_check_readable(script_path)
size = os.path.getsize(script_path)
if size > MAX_FILE_BYTES:
raise ScriptFileError(
f"Script is too large to read back ({size} bytes > {MAX_FILE_BYTES} "
f"byte cap). Edit it via a host volume mount instead."
)
logger.debug(f"read_job_file enter script_path={script_path} size={size}")
with open(script_path, encoding="utf-8") as f:
content = f.read()
logger.info(f"read_job_file ok script_path={script_path} size={size}")
return {"path": script_path, "content": content, "size": size}
def update_job_file(script_path: str, content: str) -> dict[str, object]:
"""Overwrite an existing script file with new content.
Restricted to paths under settings.jobs_dir. Caps writes at 1 MB.
Raises ScriptFileError (-> 400) if the path is missing, outside
the allowed dir, or the content is too large.
"""
_check_writable(script_path)
encoded_size = len(content.encode("utf-8"))
if encoded_size > MAX_FILE_BYTES:
raise ScriptFileError(
f"content is too large ({encoded_size} bytes > {MAX_FILE_BYTES} "
f"byte cap). Split the script into multiple files."
)
logger.debug(
f"update_job_file enter script_path={script_path} "
f"new_bytes={encoded_size}"
)
with open(script_path, "w", encoding="utf-8") as f:
written = f.write(content)
logger.info(
f"update_job_file ok script_path={script_path} bytes_written={written}"
)
return {"path": script_path, "bytes_written": written}
+29
View File
@@ -70,3 +70,32 @@ class GenerateJobFileRequest(BaseModel):
"returned path."
),
)
class ReadJobFileRequest(BaseModel):
script_path: str = Field(
...,
description=(
"Absolute path to a PySpark script inside the container's "
"filesystem. Must point at an existing regular file."
),
)
class UpdateJobFileRequest(BaseModel):
script_path: str = Field(
...,
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 "
"overwriting host-mounted configs or other critical files."
),
)
content: str = Field(
...,
description=(
"New file content (replaces the file in full; no merge/diff). "
"Maximum 1 MB to keep the MCP response bounded."
),
)