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:
@@ -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}
|
||||
@@ -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."
|
||||
),
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user