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,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"]
|
||||
Reference in New Issue
Block a user