mcp-server/spark_executor/tools/job_file.py

# 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:

  write_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.
  - No size cap on read or write payloads — the caller controls the
    file size. MCP response size is bounded by httpx and the JSON
    transport; for very large files the LLM should split via
    write_job_file and read in chunks via read_job_file + update.
"""
import os
from pathlib import Path

from common.config import settings
from common.logging import logger

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 write_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 write_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"write_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.

    No size cap. Raises ScriptFileError (-> 400) if the file is missing.
    """
    _check_readable(script_path)
    size = os.path.getsize(script_path)
    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. No size cap.
    Raises ScriptFileError (-> 400) if the path is missing or outside
    the allowed dir.
    """
    _check_writable(script_path)
    encoded_size = len(content.encode("utf-8"))
    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}