From dd197019f937697c66ea03df72336561fbcd9bac Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 25 Jun 2026 10:41:19 +0800 Subject: [PATCH] fix: validate script_path exists + tell agent to call generate_job_file Two problems with the prior prepare_submit_job flow: 1. The agent could pass a script_path that only existed in its own context (LLM-generated code not yet on disk) or a path on the host filesystem that's invisible inside the container. The new check surfaces this as a 400 with a clear remediation hint instead of letting spark-submit fail later with an opaque FileNotFoundError -> 500. 2. The container-isolation issue: any path the agent gives is interpreted inside the container. The two ways a file can legitimately exist there are (a) generate_job_file(code=...) just wrote it to SPARK_EXECUTOR_JOBS_DIR (the default ./data/jobs/ is the only gitignored dir that survives restarts), or (b) a host dir was mounted via -v. The error message spells both out so an agent can self-correct. Implementation: - submit.py: new _check_script_path() that raises ValueError (-> 400) when the path is missing, empty, or a directory. Called in both prepare_submit_job and confirm_submit_job (defense in depth). - prepare_submit_job checks the connection FIRST (KeyError -> 404) before the script (ValueError -> 400), so an agent with both problems sees the more fundamental 'unknown connection' error first. - server.py / requests.py: route description and Pydantic field description spell out the generate_job_file pattern so an LLM reading the tool schema learns the right next step. 8 new tests; 8 existing tests adjusted to create real files (they used synthetic /tmp/*.py paths that don't exist). --- spark_executor/server.py | 8 ++- spark_executor/tools/requests.py | 13 +++- spark_executor/tools/submit.py | 40 +++++++++++ tests/integration/test_mcp_routes.py | 42 ++++++++--- tests/unit/test_logging.py | 8 ++- tests/unit/test_submit_tool.py | 102 ++++++++++++++++++++------- 6 files changed, 173 insertions(+), 40 deletions(-) diff --git a/spark_executor/server.py b/spark_executor/server.py index daf5aee..211447f 100644 --- a/spark_executor/server.py +++ b/spark_executor/server.py @@ -73,7 +73,13 @@ def health_check(): "Snapshot the named Connection's master / deploy_mode / spark_conf / " "yarn_rm_url into a PendingSubmission record and persist it. " "Does NOT invoke spark-submit. Returns pending_id for use with " - "confirm_submit_job (the user-second-confirmation step)." + "confirm_submit_job (the user-second-confirmation step).\n\n" + "REQUIRED PATTERN for LLM-generated code: call generate_job_file(code=...) " + "first, then pass the returned script_path here. Direct submission with a " + "synthetic path (one that only exists in the agent's context) will be " + "rejected with HTTP 400 — the script must exist inside the container's " + "filesystem. For pre-existing files, mount the host directory into the " + "container and pass the in-container path." ), ) def _prepare_submit_job(req: PrepareSubmitJobRequest): diff --git a/spark_executor/tools/requests.py b/spark_executor/tools/requests.py index 2242a68..12ce3cc 100644 --- a/spark_executor/tools/requests.py +++ b/spark_executor/tools/requests.py @@ -26,7 +26,18 @@ class SaveConnectionRequest(BaseModel): class PrepareSubmitJobRequest(BaseModel): connection: str - script_path: str + script_path: str = Field( + ..., + description=( + "Absolute path to the PySpark script inside the container's " + "filesystem. Must point at an existing regular file. For " + "LLM-generated code, call generate_job_file(code=...) first " + "and pass the returned script_path here. For pre-existing " + "files on the host, mount them via a docker volume and pass " + "the in-container path. Returns 400 with a remediation hint " + "if the path is missing or not a file." + ), + ) queue: str = "default" executor_memory: str = "4G" executor_cores: int = 2 diff --git a/spark_executor/tools/submit.py b/spark_executor/tools/submit.py index 1518774..793c0b5 100644 --- a/spark_executor/tools/submit.py +++ b/spark_executor/tools/submit.py @@ -3,6 +3,7 @@ @Time :2026/6/24 @Author :tao.chen """ +import os import secrets import uuid from datetime import datetime @@ -20,6 +21,27 @@ from spark_executor.core.spark_submit import ( from spark_executor.models import Job, PendingSubmission, SubmitResult +def _script_path_error(path: str) -> ValueError: + """Instructive error for an invalid script_path. The MCP client receives + this via the ValueError -> 400 exception handler in server.py.""" + return ValueError( + f"script_path does not exist or is not a file: {path!r}. " + f"Two ways to fix this:\n" + f" 1. (Recommended) Call generate_job_file(code=...) first to write " + f"the PySpark source to disk, then pass the returned script_path.\n" + f" 2. If the file already exists on the host, mount it into the " + f"container (e.g. -v /host/path:/app/scripts:ro in docker run) and " + f"pass the in-container path here." + ) + + +def _check_script_path(script_path: str) -> None: + """Verify script_path points at an existing regular file. Raises ValueError + (-> 400 via the FastAPI exception handler) with an instructive message.""" + if not script_path or not os.path.isfile(script_path): + raise _script_path_error(script_path) + + def _new_pending_id() -> str: return "p_" + secrets.token_hex(6) @@ -39,9 +61,23 @@ def prepare_submit_job( f"queue={queue} executor_memory={executor_memory} executor_cores={executor_cores} " f"num_executors={num_executors}" ) + # Order of checks matters for the error the agent sees: + # 1. Unknown connection -> 404 (KeyError -> 404 in server.py) + # 2. Missing script file -> 400 (ValueError -> 400) + # Connection is checked first because it is a more fundamental problem + # (the agent is asking about a cluster that doesn't exist), and the + # agent shouldn't have to fix the script path only to learn the + # connection name is wrong. conn = conn_store.get(connection) if conn is None: raise KeyError(f"Unknown connection: {connection}") + # Fail fast: a non-existent path is the most common agent mistake (it + # generated the code in its own context but forgot to call + # generate_job_file first, or its path refers to the host filesystem + # which is invisible inside the container). Better to surface this with + # a 400 + clear remediation than to let spark-submit fail later with + # an opaque FileNotFoundError -> 500. + _check_script_path(script_path) pending_id = _new_pending_id() pending = PendingSubmission( @@ -85,6 +121,10 @@ def confirm_submit_job(*, pending_id: str) -> SubmitResult: raise ValueError( f"pending_id {pending_id} is in status {pending.status!r}, not PENDING" ) + # Defense in depth: re-verify the script still exists. A user could + # delete the file between prepare and confirm (or an external cleanup + # job could remove it). 400 via the ValueError -> 400 handler. + _check_script_path(pending.script_path) cmd = build_spark_submit_command( master=pending.master, diff --git a/tests/integration/test_mcp_routes.py b/tests/integration/test_mcp_routes.py index 4495f33..c94bf2d 100644 --- a/tests/integration/test_mcp_routes.py +++ b/tests/integration/test_mcp_routes.py @@ -76,12 +76,15 @@ def test_save_connection_accepts_dict_spark_conf_in_body(): assert r.json() == {"name": "prod", "status": "SAVED"} -def test_prepare_submit_job_works_via_body(): +def test_prepare_submit_job_works_via_body(tmp_path): c = TestClient(app) c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + # The script must exist (Stage 2 cleanup) — write a real file first. + script = tmp_path / "demo.py" + script.write_text("print('hi')\n") r = c.post( "/prepare_submit_job", - json={"connection": "prod", "script_path": "/tmp/demo.py", "queue": "research"}, + json={"connection": "prod", "script_path": str(script), "queue": "research"}, ) assert r.status_code == 200, r.text body = r.json() @@ -91,12 +94,29 @@ def test_prepare_submit_job_works_via_body(): assert body["parameters"]["master"] == "yarn" # snapshotted from connection -def test_list_and_get_pending_job_roundtrip_via_body(): +def test_prepare_rejects_nonexistent_script_path_with_400(): + """MCP clients should see a clean 400 with a remediation hint, not a 500.""" c = TestClient(app) c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + r = c.post( + "/prepare_submit_job", + json={"connection": "prod", "script_path": "/nope/does_not_exist.py"}, + ) + assert r.status_code == 400 + detail = r.json()["detail"] + # The message must guide the agent to the right next step + assert "does not exist" in detail + assert "generate_job_file" in 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"}) + script = tmp_path / "a.py" + script.write_text("print('hi')\n") prep = c.post( "/prepare_submit_job", - json={"connection": "prod", "script_path": "/tmp/a.py"}, + json={"connection": "prod", "script_path": str(script)}, ).json() pid = prep["pending_id"] @@ -104,7 +124,7 @@ def test_list_and_get_pending_job_roundtrip_via_body(): assert any(p["pending_id"] == pid for p in listed) got = c.post("/get_pending_job", json={"pending_id": pid}).json() - assert got["script_path"] == "/tmp/a.py" + assert got["script_path"] == str(script) assert got["status"] == "PENDING" @@ -164,13 +184,15 @@ def test_unknown_connection_in_prepare_returns_404(): assert "nope" in r.json()["detail"] -def test_confirm_non_pending_returns_400(): +def test_confirm_non_pending_returns_400(tmp_path): """A CANCELLED pending should refuse confirm; FastAPI should surface ValueError as 400.""" c = TestClient(app) c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + script = tmp_path / "x.py" + script.write_text("print('hi')\n") prep = c.post( "/prepare_submit_job", - json={"connection": "prod", "script_path": "/tmp/x.py"}, + json={"connection": "prod", "script_path": str(script)}, ).json() pid = prep["pending_id"] c.post("/cancel_pending_job", json={"pending_id": pid}) @@ -179,13 +201,15 @@ def test_confirm_non_pending_returns_400(): assert "CANCELLED" in r.json()["detail"] -def test_cancel_already_submitted_returns_400(): +def test_cancel_already_submitted_returns_400(tmp_path): """Cancelling a SUBMITTED pending should refuse with ValueError -> 400.""" c = TestClient(app) c.post("/save_connection", json={"name": "prod", "master": "yarn"}) + script = tmp_path / "x.py" + script.write_text("print('hi')\n") prep = c.post( "/prepare_submit_job", - json={"connection": "prod", "script_path": "/tmp/x.py"}, + json={"connection": "prod", "script_path": str(script)}, ).json() pid = prep["pending_id"] # Simulate a SUBMITTED state by mutating the pending directly diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py index 44b746b..9913abe 100644 --- a/tests/unit/test_logging.py +++ b/tests/unit/test_logging.py @@ -42,14 +42,16 @@ def test_save_connection_emits_info_log(log_capture): assert "connection saved" in text -def test_prepare_submit_job_emits_debug_and_info(log_capture): +def test_prepare_submit_job_emits_debug_and_info(log_capture, tmp_path): connections.save_connection(name="prod", master="yarn", deploy_mode="cluster") + script = tmp_path / "demo.py" + script.write_text("print('hi')\n") log_capture.truncate(0); log_capture.seek(0) - submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py", queue="research") + submit.prepare_submit_job(connection="prod", script_path=str(script), queue="research") text = log_capture.getvalue() assert "DEBUG|prepare_submit_job enter" in text assert "INFO|prepare_submit_job ok" in text - assert "script_path=/tmp/j.py" in text + assert f"script_path={script}" in text assert "queue=research" in text diff --git a/tests/unit/test_submit_tool.py b/tests/unit/test_submit_tool.py index 378aa22..d44eb74 100644 --- a/tests/unit/test_submit_tool.py +++ b/tests/unit/test_submit_tool.py @@ -26,6 +26,14 @@ def _fresh(tmp_path: Path, monkeypatch): )) +@pytest.fixture +def real_script(tmp_path: Path) -> Path: + """A real .py file on disk that satisfies prepare_submit_job's path check.""" + p = tmp_path / "demo.py" + p.write_text("print('hi from test')\n") + return p + + def _last_pending_id() -> str: pending = submit.pending_store.list_all() assert pending, "no pending submission was created" @@ -34,16 +42,16 @@ def _last_pending_id() -> str: # --- prepare_submit_job --- -def test_prepare_does_not_invoke_spark_submit(monkeypatch): +def test_prepare_does_not_invoke_spark_submit(monkeypatch, real_script): called = [] monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd)) - out = submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + out = submit.prepare_submit_job(connection="prod", script_path=str(real_script)) assert called == [] # never called assert out["status"] == "PENDING" assert out["pending_id"].startswith("p_") -def test_prepare_persists_pending_with_snapshot(monkeypatch): +def test_prepare_persists_pending_with_snapshot(monkeypatch, real_script): monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) connection_store.store.save( Connection( @@ -54,7 +62,7 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch): spark_conf={"spark.sql.shuffle.partitions": "200"}, ) ) - submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py", queue="research") + submit.prepare_submit_job(connection="prod", script_path=str(real_script), queue="research") p = submit.pending_store.get(_last_pending_id()) assert p is not None assert p.connection == "prod" @@ -63,18 +71,57 @@ def test_prepare_persists_pending_with_snapshot(monkeypatch): assert p.yarn_rm_url == "http://rm:8088" assert p.spark_conf == {"spark.sql.shuffle.partitions": "200"} assert p.queue == "research" - assert p.script_path == "/tmp/j.py" + assert p.script_path == str(real_script) -def test_prepare_raises_for_unknown_connection(): +def test_prepare_raises_for_unknown_connection(real_script): with pytest.raises(KeyError, match="missing"): - submit.prepare_submit_job(connection="missing", script_path="/tmp/j.py") + submit.prepare_submit_job(connection="missing", script_path=str(real_script)) -def test_prepare_snapshots_connection_at_prepare_time(monkeypatch): +# --- script_path validation --- + +def test_prepare_rejects_nonexistent_script_path(tmp_path): + missing = str(tmp_path / "does_not_exist.py") + with pytest.raises(ValueError, match="does not exist or is not a file"): + submit.prepare_submit_job(connection="prod", script_path=missing) + + +def test_prepare_rejects_directory_as_script_path(tmp_path): + """A directory is not a file, even if it exists.""" + with pytest.raises(ValueError, match="does not exist or is not a file"): + submit.prepare_submit_job(connection="prod", script_path=str(tmp_path)) + + +def test_prepare_rejects_empty_script_path(): + with pytest.raises(ValueError, match="does not exist or is not a file"): + submit.prepare_submit_job(connection="prod", script_path="") + + +def test_prepare_error_message_mentions_generate_job_file(tmp_path): + """The agent must be told to call generate_job_file first.""" + with pytest.raises(ValueError, match="generate_job_file"): + submit.prepare_submit_job( + connection="prod", script_path=str(tmp_path / "x.py") + ) + + +def test_confirm_rejects_if_script_was_deleted_after_prepare(tmp_path, monkeypatch): + """Defense in depth: file present at prepare, gone by confirm -> 400.""" + script = tmp_path / "demo.py" + script.write_text("print('hi')\n") + submit.prepare_submit_job(connection="prod", script_path=str(script)) + pid = _last_pending_id() + # Simulate the file being deleted (e.g. cleanup job) between prepare and confirm + script.unlink() + with pytest.raises(ValueError, match="does not exist or is not a file"): + submit.confirm_submit_job(pending_id=pid) + + +def test_prepare_snapshots_connection_at_prepare_time(monkeypatch, real_script): """Editing the connection between prepare and confirm must NOT silently retarget.""" monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) - submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + submit.prepare_submit_job(connection="prod", script_path=str(real_script)) p = submit.pending_store.get(_last_pending_id()) # Now mutate the connection connection_store.store.save( @@ -88,8 +135,8 @@ def test_prepare_snapshots_connection_at_prepare_time(monkeypatch): # --- confirm_submit_job --- -def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch): - submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") +def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch, real_script): + submit.prepare_submit_job(connection="prod", script_path=str(real_script)) pid = _last_pending_id() fake_proc = type("P", (), { "returncode": 0, @@ -100,7 +147,7 @@ def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch): cmd = m.call_args.args[0] assert "yarn" in cmd assert "cluster" in cmd - assert cmd[-1] == "/tmp/j.py" + assert cmd[-1] == str(real_script) assert result.application_id == "application_17400000001" # pending updated p = submit.pending_store.get(pid) @@ -116,8 +163,8 @@ def test_confirm_raises_for_unknown_pending_id(): submit.confirm_submit_job(pending_id="missing") -def test_confirm_refuses_non_pending_status(monkeypatch): - submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") +def test_confirm_refuses_non_pending_status(monkeypatch, real_script): + submit.prepare_submit_job(connection="prod", script_path=str(real_script)) pid = _last_pending_id() # Mark it CANCELLED first p = submit.pending_store.get(pid) @@ -127,9 +174,9 @@ def test_confirm_refuses_non_pending_status(monkeypatch): submit.confirm_submit_job(pending_id=pid) -def test_confirm_marks_failed_on_spark_submit_error(monkeypatch): +def test_confirm_marks_failed_on_spark_submit_error(monkeypatch, real_script): from spark_executor.core.spark_submit import SparkSubmitError - submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + submit.prepare_submit_job(connection="prod", script_path=str(real_script)) pid = _last_pending_id() def _raise(_cmd): raise SparkSubmitError("boom") @@ -147,20 +194,23 @@ def test_list_pending_jobs_empty(): assert submit.list_pending_jobs() == [] -def test_list_pending_jobs_returns_all(monkeypatch): +def test_list_pending_jobs_returns_all(monkeypatch, tmp_path): + a = tmp_path / "a.py" + b = tmp_path / "b.py" + a.write_text("a\n"); b.write_text("b\n") monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) - submit.prepare_submit_job(connection="prod", script_path="/tmp/a.py") - submit.prepare_submit_job(connection="prod", script_path="/tmp/b.py") + submit.prepare_submit_job(connection="prod", script_path=str(a)) + submit.prepare_submit_job(connection="prod", script_path=str(b)) out = submit.list_pending_jobs() - assert {p["script_path"] for p in out} == {"/tmp/a.py", "/tmp/b.py"} + assert {p["script_path"] for p in out} == {str(a), str(b)} assert all(p["status"] == "PENDING" for p in out) # --- get_pending_job --- -def test_get_pending_job_returns_dump(monkeypatch): +def test_get_pending_job_returns_dump(monkeypatch, real_script): monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) - submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + submit.prepare_submit_job(connection="prod", script_path=str(real_script)) pid = _last_pending_id() out = submit.get_pending_job(pid) assert out["pending_id"] == pid @@ -175,9 +225,9 @@ def test_get_pending_job_unknown_raises(): # --- cancel_pending_job --- -def test_cancel_pending_job_marks_cancelled(monkeypatch): +def test_cancel_pending_job_marks_cancelled(monkeypatch, real_script): monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None) - submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") + submit.prepare_submit_job(connection="prod", script_path=str(real_script)) pid = _last_pending_id() out = submit.cancel_pending_job(pid) assert out == {"pending_id": pid, "status": "CANCELLED"} @@ -189,8 +239,8 @@ def test_cancel_pending_job_unknown_raises(): submit.cancel_pending_job("missing") -def test_cancel_pending_job_refuses_submitted(monkeypatch): - submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py") +def test_cancel_pending_job_refuses_submitted(monkeypatch, real_script): + submit.prepare_submit_job(connection="prod", script_path=str(real_script)) pid = _last_pending_id() p = submit.pending_store.get(pid) p.status = "SUBMITTED"