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
+50 -2
View File
@@ -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"})