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
+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}