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/'.

spark_executor/server.py

@@ -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)

spark_executor/tools/job_file.py

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

spark_executor/tools/requests.py

@@ -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."
+        ),
+    )

tests/integration/test_mcp_routes.py

@@ -33,7 +33,7 @@ def test_health_still_present():
     assert r.json() == {"status": "ok"}
 
 
-def test_thirteen_tool_routes_registered():
+def test_sixteen_tool_routes_registered():
     paths = {r.path for r in app.routes}
     for path in (
         # pending-submission flow (5)
@@ -51,8 +51,10 @@ def test_thirteen_tool_routes_registered():
         "/list_connections",
         "/get_connection",
         "/delete_connection",
-        # LLM-driven PySpark generation (1) — Stage 2
+        # LLM-driven PySpark generation (3) — Stage 2
         "/generate_job_file",
+        "/read_job_file",
+        "/update_job_file",
     ):
         assert path in paths, f"missing MCP tool route: {path}"
 
@@ -144,6 +146,52 @@ def test_generate_rejects_sql_policy_violation_with_400(tmp_path):
     assert list(tmp_path.glob("*.py")) == []
 
 
+# --- Stage 2: read_job_file / update_job_file ---
+
+def test_read_job_file_returns_content_via_mcp(tmp_path):
+    c = TestClient(app)
+    script = tmp_path / "demo.py"
+    script.write_text("print('hello from read_job_file')\n")
+    r = c.post("/read_job_file", json={"script_path": str(script)})
+    assert r.status_code == 200
+    body = r.json()
+    assert body["content"] == "print('hello from read_job_file')\n"
+    assert body["path"] == str(script)
+
+
+def test_read_job_file_404_for_missing_file(tmp_path):
+    c = TestClient(app)
+    r = c.post("/read_job_file", json={"script_path": str(tmp_path / "nope.py")})
+    assert r.status_code == 400
+    assert "does not exist" in r.json()["detail"]
+
+
+def test_update_job_file_writes_and_reads_back(tmp_path):
+    from common import config
+    config.settings.jobs_dir = str(tmp_path)
+    c = TestClient(app)
+    script = tmp_path / "edit.py"
+    script.write_text("v1\n")
+    # Update
+    r1 = c.post("/update_job_file", json={"script_path": str(script), "content": "v2\n"})
+    assert r1.status_code == 200
+    assert r1.json()["bytes_written"] == 3
+    # Read back
+    r2 = c.post("/read_job_file", json={"script_path": str(script)})
+    assert r2.json()["content"] == "v2\n"
+
+
+def test_update_job_file_rejects_paths_outside_jobs_dir(tmp_path):
+    from common import config
+    config.settings.jobs_dir = str(tmp_path / "jobs")
+    c = TestClient(app)
+    other = tmp_path / "elsewhere.py"
+    other.write_text("x\n")
+    r = c.post("/update_job_file", json={"script_path": str(other), "content": "y\n"})
+    assert r.status_code == 400
+    assert "must be under" in r.json()["detail"]
+
+
 def test_list_and_get_pending_job_roundtrip_via_body(tmp_path):
     c = TestClient(app)
     c.post("/save_connection", json={"name": "prod", "master": "yarn"})

tests/unit/test_job_file.py

@@ -0,0 +1,162 @@
+# coding=utf-8
+import os
+from pathlib import Path
+
+import pytest
+
+from common import config
+from spark_executor.tools import job_file
+from spark_executor.tools.job_file import (
+    ScriptFileError,
+    read_job_file,
+    update_job_file,
+)
+
+
+@pytest.fixture(autouse=True)
+def _restore_settings():
+    snapshot = config.Settings(
+        data_dir=config.settings.data_dir,
+        jobs_dir=config.settings.jobs_dir,
+        log_dir=config.settings.log_dir,
+        yarn_resource_manager_url=config.settings.yarn_resource_manager_url,
+        log_level=config.settings.log_level,
+    )
+    yield
+    config.settings.data_dir = snapshot.data_dir
+    config.settings.jobs_dir = snapshot.jobs_dir
+    config.settings.log_dir = snapshot.log_dir
+    config.settings.yarn_resource_manager_url = snapshot.yarn_resource_manager_url
+    config.settings.log_level = snapshot.log_level
+
+
+# --- read_job_file ---
+
+def test_read_returns_existing_content(tmp_path: Path):
+    p = tmp_path / "demo.py"
+    p.write_text("print('hello')\n")
+    out = read_job_file(str(p))
+    assert out["content"] == "print('hello')\n"
+    assert out["path"] == str(p)
+    assert out["size"] == len("print('hello')\n")
+
+
+def test_read_works_for_any_existing_file(tmp_path: Path):
+    """read_job_file has no jobs_dir restriction — it's read-only and the
+    agent needs to inspect arbitrary files (e.g. logs/, /etc/hadoop/conf)."""
+    p = tmp_path / "anything.py"
+    p.write_text("x = 1\n")
+    out = read_job_file(str(p))
+    assert out["content"] == "x = 1\n"
+
+
+def test_read_raises_for_missing_file(tmp_path: Path):
+    with pytest.raises(ScriptFileError, match="does not exist"):
+        read_job_file(str(tmp_path / "nope.py"))
+
+
+def test_read_raises_for_directory(tmp_path: Path):
+    with pytest.raises(ScriptFileError, match="does not exist"):
+        read_job_file(str(tmp_path))
+
+
+def test_read_raises_for_empty_path():
+    with pytest.raises(ScriptFileError, match="does not exist"):
+        read_job_file("")
+
+
+# --- update_job_file ---
+
+def test_update_overwrites_existing_file_in_jobs_dir(tmp_path: Path):
+    """The file must be under settings.jobs_dir (default './data/jobs')."""
+    config.settings.jobs_dir = str(tmp_path)
+    p = tmp_path / "script.py"
+    p.write_text("v1\n")
+    out = update_job_file(str(p), "v2\n")
+    assert out["bytes_written"] == 3
+    assert p.read_text() == "v2\n"
+
+
+def test_update_rejects_paths_outside_jobs_dir(tmp_path: Path):
+    """update_job_file is sandboxed to settings.jobs_dir."""
+    config.settings.jobs_dir = str(tmp_path / "jobs")
+    other = tmp_path / "elsewhere.py"
+    other.write_text("x\n")
+    with pytest.raises(ScriptFileError, match="must be under"):
+        update_job_file(str(other), "y\n")
+
+
+def test_update_rejects_paths_outside_jobs_dir_via_existing_file(tmp_path: Path):
+    """Path validation: even when the file exists, an absolute path OUTSIDE
+    settings.jobs_dir is rejected. (In production, /etc/passwd would exist
+    in the container; here we synthesize the same scenario with a real
+    file outside the configured jobs_dir.)"""
+    config.settings.jobs_dir = str(tmp_path / "jobs")
+    (tmp_path / "jobs").mkdir()
+    outside = tmp_path / "hadoop-core-site.xml"
+    outside.write_text("<config/>\n")
+    with pytest.raises(ScriptFileError, match="must be under"):
+        update_job_file(str(outside), "tampered\n")
+
+
+def test_update_rejects_relative_escape_attempt(tmp_path: Path):
+    """Path traversal via '..' must not bypass the jobs_dir check."""
+    config.settings.jobs_dir = str(tmp_path / "jobs")
+    (tmp_path / "jobs").mkdir()
+    (tmp_path / "etc").mkdir()
+    secret = tmp_path / "etc" / "passwd"
+    secret.write_text("root:x:0:0:...\n")
+    # Try to reach the file via '../etc/passwd' relative to jobs_dir
+    sneaky = str(tmp_path / "jobs" / ".." / "etc" / "passwd")
+    with pytest.raises(ScriptFileError, match="must be under"):
+        update_job_file(sneaky, "tampered\n")
+
+
+def test_update_rejects_missing_file(tmp_path: Path):
+    config.settings.jobs_dir = str(tmp_path)
+    with pytest.raises(ScriptFileError, match="does not exist"):
+        update_job_file(str(tmp_path / "nope.py"), "x\n")
+
+
+def test_update_rejects_oversized_content(tmp_path: Path, monkeypatch):
+    """Cap at 1 MB so the MCP response stays bounded."""
+    config.settings.jobs_dir = str(tmp_path)
+    p = tmp_path / "big.py"
+    p.write_text("# small\n")
+    # Synthesize 2 MB of content (don't actually write 2 MB to disk)
+    big = "x" * (2 * 1024 * 1024)
+    with pytest.raises(ScriptFileError, match="too large"):
+        update_job_file(str(p), big)
+
+
+def test_update_rejects_1mb_plus_1_byte(tmp_path: Path):
+    """Exactly at the boundary: 1 MB + 1 byte must be rejected."""
+    config.settings.jobs_dir = str(tmp_path)
+    p = tmp_path / "x.py"
+    p.write_text("# t\n")
+    just_over = "x" * (1024 * 1024 + 1)
+    with pytest.raises(ScriptFileError, match="too large"):
+        update_job_file(str(p), just_over)
+
+
+# --- round-trip: write -> read -> update -> read ---
+
+def test_full_edit_cycle_via_tools(tmp_path: Path):
+    """Simulate the agent's review-and-edit loop end-to-end."""
+    config.settings.jobs_dir = str(tmp_path)
+    path = tmp_path / "script.py"
+    # 1. Initial state: file doesn't exist (in real flow, generate_job_file
+    # would create it; we just create it directly here for the test)
+    path.write_text("# v1: initial\n")
+    # 2. Agent reads back
+    r1 = read_job_file(str(path))
+    assert r1["content"] == "# v1: initial\n"
+    # 3. Agent edits (LLM or human)
+    edited = "# v2: edited by agent\n" + r1["content"].split("\n", 1)[1]
+    # 4. Agent writes back
+    w = update_job_file(str(path), edited)
+    assert w["bytes_written"] == len(edited)
+    # 5. Read back to verify
+    r2 = read_job_file(str(path))
+    assert r2["content"] == edited
+    assert "v2: edited" in r2["content"]