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:

  generate_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.
  - 1 MB cap on both read and write payloads to keep MCP responses bounded.
"""
import os
from pathlib import Path

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

MAX_FILE_BYTES = 1 * 1024 * 1024  # 1 MB


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

    Caps the read at 1 MB to keep MCP responses bounded; raises
    ScriptFileError (-> 400) if the file is missing or too large.
    """
    _check_readable(script_path)
    size = os.path.getsize(script_path)
    if size > MAX_FILE_BYTES:
        raise ScriptFileError(
            f"Script is too large to read back ({size} bytes > {MAX_FILE_BYTES} "
            f"byte cap). Edit it via a host volume mount instead."
        )
    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. Caps writes at 1 MB.
    Raises ScriptFileError (-> 400) if the path is missing, outside
    the allowed dir, or the content is too large.
    """
    _check_writable(script_path)
    encoded_size = len(content.encode("utf-8"))
    if encoded_size > MAX_FILE_BYTES:
        raise ScriptFileError(
            f"content is too large ({encoded_size} bytes > {MAX_FILE_BYTES} "
            f"byte cap). Split the script into multiple files."
        )
    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}