Files
mcp-server/spark_executor/tools/requests.py
T
Claude dd197019f9 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).
2026-06-25 10:41:19 +08:00

73 lines
2.0 KiB
Python

# coding=utf-8
"""
@Time :2026/6/24
@Author :tao.chen
Pydantic request models for the FastAPI route layer. The underlying tool
functions in tools/*.py still take keyword arguments; these models exist only
so fastapi-mcp can call the routes via tools/call (which sends args as a
JSON body) without 422-ing on dict-typed parameters like spark_conf.
"""
from pydantic import BaseModel, Field
class EmptyRequest(BaseModel):
"""Used for tools that take no arguments (list_connections, list_pending_jobs)."""
pass
class SaveConnectionRequest(BaseModel):
name: str
master: str
deploy_mode: str = "cluster"
yarn_rm_url: str | None = None
spark_conf: dict[str, str] | None = None
class PrepareSubmitJobRequest(BaseModel):
connection: 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
num_executors: int = 2
class PendingIdRequest(BaseModel):
pending_id: str
class JobIdRequest(BaseModel):
job_id: str
class GetJobLogsRequest(BaseModel):
job_id: str
tail_chars: int = 5000
class ConnectionNameRequest(BaseModel):
name: str
class GenerateJobFileRequest(BaseModel):
code: str = Field(
...,
description=(
"Full PySpark source code to write to disk. Will be passed verbatim "
"to spark-submit after the agent calls prepare_submit_job on the "
"returned path."
),
)