mcp-server/docs/superpowers/plans/2026-06-24-spark-executor-mcp.md

Spark Executor MCP Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: Build a three-stage MCP server that lets an LLM Agent submit, monitor, fetch logs from, and kill PySpark jobs on YARN — escalating from a synchronous minimum-viable tool set to an async, multi-tenant job platform backed by SQLite.

Architecture: FastAPI app exposes MCP tools (via fastapi-mcp) that wrap shell invocations of spark-submit and yarn. Stage 1 uses synchronous subprocess.run and an in-memory job map. Stage 2 adds an LLM-facing write_job_file helper. Stage 3 refactors submission to asyncio.create_subprocess_exec, persists job state in SQLite, and introduces offset-based log caching.

Tech Stack: Python 3.12+, FastAPI 0.138+, fastapi-mcp 0.4+, Pydantic 2.13+, loguru, httpx (already in pyproject.toml); uv for env mgmt; pytest for tests; SQLite (stdlib sqlite3) for Stage 3 registry.

---

File Structure

Stage 1 (this plan) creates:

spark_executor/
  __init__.py
  server.py                       # FastAPI app, mounts MCP tools
  tools/
    __init__.py
    submit.py                     # submit_job MCP tool (uses a saved Connection)
    status.py                     # get_job_status MCP tool
    logs.py                       # get_job_logs MCP tool
    kill.py                       # kill_job MCP tool
    connections.py                # save/list/get/delete_connection MCP tools
  core/
    __init__.py
    spark_submit.py               # builds & runs spark-submit commands
    yarn_client.py                # wraps yarn application / yarn logs
    log_parser.py                 # extracts application_id from spark-submit output
    job_store.py                  # in-memory dict: job_id -> {application_id, ...}
    connection_store.py           # JSON CRUD over ./data/connections.json
  models.py                       # Pydantic: Job, JobStatus, SubmitResult, Connection
data/
  connections.json                # created on first save_connection (gitignored)
tests/
  __init__.py
  conftest.py
  unit/
    test_log_parser.py
    test_spark_submit.py
    test_yarn_client.py
    test_connection_store.py
    test_connection_tools.py
    test_submit_tool.py
    test_status_tool.py
    test_logs_tool.py
    test_kill_tool.py

Stage 2 adds:

spark_executor/
  tools/
    generate.py                   # write_job_file MCP tool
  core/
    job_writer.py                 # writes PySpark code to /tmp/jobs/job_xxx.py
tests/
  unit/
    test_write_job_tool.py
    test_job_writer.py

Stage 3 adds (in a follow-on plan, not implemented here):

spark_executor/
  core/
    job_registry.py               # SQLite-backed CRUD
    log_cache.py                  # offset-based yarn-logs cache
    async_submitter.py            # asyncio.create_subprocess_exec wrapper
    status_poller.py              # background task updating registry

---

Global Constraints

• Python ≥ 3.12 (per pyproject.toml).
• Module header on every new Python file: # coding=utf-8 + @Time / @Author docstring, matching the existing common/ and spark_executor/ files.
• Logging: import from common.logging import logger (the shared loguru logger). Do not instantiate new loggers.
• MCP tools are registered by decorating FastAPI routes on the app in spark_executor/server.py; fastapi-mcp auto-discovers them. No manual MCP registration code.
• Subprocess calls must capture stdout AND stderr (Spark writes the application_id line to stderr, not stdout). This is true for both spark-submit and yarn application -status.
• Default spark-submit flags (Stage 1, fixed): --master and --deploy-mode are derived from the Connection passed to submit_job (see Task 11). Other params are user-tunable.
• Connection records persist to ./data/connections.json by default. Override with SPARK_EXECUTOR_DATA_DIR env var (the file path becomes <dir>/connections.json). The directory is auto-created on first write. ./data/ is added to .gitignore so credentials/configs are not committed.
• All test commands use the venv: uv run pytest ... (or activate .venv/bin/activate first).
• Every task ends in a commit. Commit messages use conventional-commits prefixes: test:, feat:, chore:, refactor:.
• YAGNI: do not introduce the SQLite registry, async execution, log cache, or ownership filters in Stage 1 tasks. They are out of scope and belong to Stage 3.

---

Stage 1 — Minimum Closed Loop (★☆☆☆☆)

Target: 4 MCP tools (submit_job, get_job_status, get_job_logs, kill_job) backed by synchronous subprocess.run. ~300–500 lines of Python.

---

Task 1: Test infrastructure and Pydantic models

Files:

• Create: tests/__init__.py
• Create: tests/conftest.py
• Create: tests/unit/__init__.py
• Create: spark_executor/models.py
• Create: spark_executor/core/__init__.py
• Create: spark_executor/tools/__init__.py
• Modify: pyproject.toml (add [dependency-groups].dev with pytest)
• Test: tests/unit/test_models.py

Interfaces:

• Produces:

  • class Job(BaseModel) with fields job_id: str, application_id: str, script_path: str, queue: str, submit_time: datetime, connection: str (name of the Connection used to submit)
  • class JobStatus(BaseModel) with fields application_id: str, state: str, raw: str
  • class SubmitResult(BaseModel) with fields job_id: str, application_id: str, tracking_url: str | None
  • class Connection(BaseModel) with fields name: str, master: str, deploy_mode: str = "cluster", yarn_rm_url: str | None = None, spark_conf: dict[str, str] = Field(default_factory=dict)
  • class PendingSubmission(BaseModel) with fields pending_id: str, connection: str, master: str, deploy_mode: str, script_path: str, queue: str, executor_memory: str, executor_cores: int, num_executors: int, spark_conf: dict[str, str], created_at: datetime, status: str = "PENDING" (one of "PENDING" | "SUBMITTED" | "CANCELLED" | "FAILED"), error: str | None = None, job_id: str | None = None, application_id: str | None = None. The master / deploy_mode / spark_conf fields are a snapshot of the source Connection taken at prepare_submit_job time so that editing the connection between prepare and confirm does not silently change the submission target.

• Step 1: Add pytest to dev dependencies

Edit pyproject.toml, appending at the end:

[dependency-groups]
dev = [
    "pytest>=8.0",
    "pytest-mock>=3.14",
]

Then run:

cd /Users/taochen/llm/spark-executor-mcp && uv sync

Expected: uv.lock updated; pytest available in .venv/bin/.

• Step 2: Create package skeleton with empty __init__.py files

mkdir -p tests/unit spark_executor/core spark_executor/tools

Each new __init__.py gets the standard header:

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen 
"""

• Step 3: Write the failing test for models

tests/unit/test_models.py:

# coding=utf-8
from datetime import datetime
from spark_executor.models import Connection, Job, JobStatus, PendingSubmission, SubmitResult


def test_job_roundtrip():
    job = Job(
        job_id="abc123",
        application_id="application_17400000001",
        script_path="/tmp/jobs/job_001.py",
        queue="default",
        submit_time=datetime(2026, 6, 24, 10, 0, 0),
        connection="prod-yarn",
    )
    dumped = job.model_dump()
    assert dumped["job_id"] == "abc123"
    assert dumped["application_id"] == "application_17400000001"
    assert dumped["connection"] == "prod-yarn"


def test_job_status_default_raw():
    s = JobStatus(application_id="application_1", state="RUNNING")
    assert s.raw == ""


def test_submit_result_tracking_url_optional():
    r = SubmitResult(job_id="j1", application_id="application_1", tracking_url=None)
    assert r.tracking_url is None


def test_connection_defaults():
    c = Connection(name="prod", master="yarn")
    assert c.deploy_mode == "cluster"
    assert c.yarn_rm_url is None
    assert c.spark_conf == {}


def test_connection_with_spark_conf():
    c = Connection(
        name="dev",
        master="spark://master:7077",
        deploy_mode="client",
        spark_conf={"spark.sql.shuffle.partitions": "200"},
    )
    assert c.spark_conf["spark.sql.shuffle.partitions"] == "200"


def test_pending_submission_defaults_to_pending_status():
    p = PendingSubmission(
        pending_id="p_1",
        connection="prod",
        master="yarn",
        deploy_mode="cluster",
        script_path="/tmp/j.py",
        queue="default",
        executor_memory="2G",
        executor_cores=2,
        num_executors=2,
        spark_conf={},
        created_at=datetime(2026, 6, 24),
    )
    assert p.status == "PENDING"
    assert p.error is None
    assert p.job_id is None
    assert p.application_id is None


def test_pending_submission_can_record_outcome():
    p = PendingSubmission(
        pending_id="p_2",
        connection="prod",
        master="yarn",
        deploy_mode="cluster",
        script_path="/tmp/j.py",
        queue="default",
        executor_memory="2G",
        executor_cores=2,
        num_executors=2,
        spark_conf={},
        created_at=datetime(2026, 6, 24),
        status="SUBMITTED",
        job_id="j_abc",
        application_id="application_17400000001",
    )
    assert p.status == "SUBMITTED"
    assert p.job_id == "j_abc"
    assert p.application_id == "application_17400000001"

• Step 4: Run test to verify it fails

cd /Users/taochen/llm/spark-executor-mcp && uv run pytest tests/unit/test_models.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.models'.

• Step 5: Implement spark_executor/models.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from datetime import datetime

from pydantic import BaseModel, Field


class Job(BaseModel):
    job_id: str
    application_id: str
    script_path: str
    queue: str
    submit_time: datetime
    connection: str


class JobStatus(BaseModel):
    application_id: str
    state: str
    raw: str = Field(default="")


class SubmitResult(BaseModel):
    job_id: str
    application_id: str
    tracking_url: str | None = None


class Connection(BaseModel):
    name: str
    master: str
    deploy_mode: str = "cluster"
    yarn_rm_url: str | None = None
    spark_conf: dict[str, str] = Field(default_factory=dict)


class PendingSubmission(BaseModel):
    pending_id: str
    connection: str
    master: str
    deploy_mode: str
    script_path: str
    queue: str
    executor_memory: str
    executor_cores: int
    num_executors: int
    spark_conf: dict[str, str] = Field(default_factory=dict)
    created_at: datetime
    status: str = "PENDING"  # PENDING | SUBMITTED | CANCELLED | FAILED
    error: str | None = None
    job_id: str | None = None
    application_id: str | None = None

• Step 6: Run test to verify it passes

uv run pytest tests/unit/test_models.py -v

Expected: 8 passed.

• Step 7: Create conftest.py to ensure spark_executor is importable

tests/conftest.py:

# coding=utf-8
# Ensures the project root is on sys.path when running pytest from any cwd.
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

• Step 7b: Add ./data/ to .gitignore

Create .gitignore (or append, if it exists):

data/

So credentials and config blobs are not committed.

• Step 8: Commit

git add pyproject.toml uv.lock tests spark_executor/models.py spark_executor/core/__init__.py spark_executor/tools/__init__.py .gitignore
git commit -m "chore: add pytest dev-deps and pydantic models (Job, Connection)"

---

Task 2: log_parser — extract application_id and tracking URL from spark-submit output

Spark writes the application_id to stderr in one of two forms:

Submitted application application_17400000001

or

tracking URL: http://rm:8088/proxy/application_17400000001/

Files:

• Create: spark_executor/core/log_parser.py
• Test: tests/unit/test_log_parser.py

Interfaces:

• Produces:

  • def parse_spark_submit_output(stderr: str) -> tuple[str, str | None] — returns (application_id, tracking_url). Raises ValueError if no application_id is found.

• Step 1: Write the failing tests

tests/unit/test_log_parser.py:

# coding=utf-8
import pytest

from spark_executor.core.log_parser import parse_spark_submit_output


def test_parses_submitted_application_line():
    stderr = "Warning: ignoring...\nSubmitted application application_17400000001\n"
    app_id, url = parse_spark_submit_output(stderr)
    assert app_id == "application_17400000001"
    assert url is None


def test_parses_tracking_url_line():
    stderr = "tracking URL: http://rm:8088/proxy/application_17400000002/\n"
    app_id, url = parse_spark_submit_output(stderr)
    assert app_id == "application_17400000002"
    assert url == "http://rm:8088/proxy/application_17400000002/"


def test_prefers_submitted_application_line_over_tracking_url():
    stderr = (
        "tracking URL: http://rm:8088/proxy/application_9999/\n"
        "Submitted application application_1234\n"
    )
    app_id, _ = parse_spark_submit_output(stderr)
    assert app_id == "application_1234"


def test_raises_when_no_application_id_found():
    with pytest.raises(ValueError, match="application_id"):
        parse_spark_submit_output("some unrelated output\n")

• Step 2: Run tests to verify they fail

uv run pytest tests/unit/test_log_parser.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.core.log_parser'.

• Step 3: Implement log_parser.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
import re

_APP_ID_RE = re.compile(r"Submitted application (\S+)")
_TRACKING_URL_RE = re.compile(r"tracking URL:\s+(\S+)")


def parse_spark_submit_output(stderr: str) -> tuple[str, str | None]:
    """Extract (application_id, tracking_url) from spark-submit stderr.

    Raises ValueError if no application_id can be found.
    """
    app_match = _APP_ID_RE.search(stderr)
    if app_match:
        application_id = app_match.group(1)
    else:
        url_match = _TRACKING_URL_RE.search(stderr)
        if not url_match:
            raise ValueError("Could not find application_id in spark-submit output")
        # tracking URL is of the form http://rm:8088/proxy/application_xxx/
        tracking_url = url_match.group(1)
        tail = tracking_url.rstrip("/").rsplit("/", 1)[-1]
        if not tail.startswith("application_"):
            raise ValueError("Could not find application_id in spark-submit output")
        application_id = tail

    url_match = _TRACKING_URL_RE.search(stderr)
    tracking_url = url_match.group(1) if url_match else None
    return application_id, tracking_url

• Step 4: Run tests to verify they pass

uv run pytest tests/unit/test_log_parser.py -v

Expected: 4 passed.

• Step 5: Commit

git add spark_executor/core/log_parser.py tests/unit/test_log_parser.py
git commit -m "feat: add spark-submit output parser"

---

Task 3: spark_submit — build and run a spark-submit command

Files:

• Create: spark_executor/core/spark_submit.py
• Test: tests/unit/test_spark_submit.py

Interfaces:

• Produces:

  • def build_spark_submit_command(*, master: str, deploy_mode: str, script_path: str, queue: str, executor_memory: str, executor_cores: int, num_executors: int, spark_conf: dict[str, str] | None = None) -> list[str] — --master and --deploy-mode come from the Connection chosen by the caller (see Task 11). spark_conf entries are emitted as --conf k=v pairs in declaration order.
  • def run_spark_submit(cmd: list[str]) -> subprocess.CompletedProcess[str] — runs the command, returns the CompletedProcess. Stdout AND stderr captured (text mode, errors="replace").
  • class SparkSubmitError(Exception)

• Step 1: Write the failing tests

tests/unit/test_spark_submit.py:

# coding=utf-8
from unittest.mock import MagicMock, patch

import pytest

from spark_executor.core.spark_submit import (
    SparkSubmitError,
    build_spark_submit_command,
    run_spark_submit,
)


def test_build_command_uses_provided_master_and_deploy_mode():
    cmd = build_spark_submit_command(
        master="yarn",
        deploy_mode="cluster",
        script_path="/tmp/jobs/job_001.py",
        queue="default",
        executor_memory="2G",
        executor_cores=2,
        num_executors=2,
    )
    assert cmd[:2] == ["spark-submit", "--master"]
    assert "yarn" in cmd
    assert "cluster" in cmd
    assert "--queue" in cmd and "default" in cmd
    assert "--executor-memory" in cmd and "2G" in cmd
    assert "--executor-cores" in cmd and "2" in cmd
    assert "--num-executors" in cmd and "2" in cmd
    assert cmd[-1] == "/tmp/jobs/job_001.py"


def test_build_command_supports_standalone_master():
    cmd = build_spark_submit_command(
        master="spark://master:7077",
        deploy_mode="client",
        script_path="/tmp/j.py",
        queue="default",
        executor_memory="1G",
        executor_cores=1,
        num_executors=1,
    )
    assert "spark://master:7077" in cmd
    assert "client" in cmd


def test_build_command_appends_spark_conf_entries():
    cmd = build_spark_submit_command(
        master="yarn",
        deploy_mode="cluster",
        script_path="/tmp/j.py",
        queue="default",
        executor_memory="2G",
        executor_cores=2,
        num_executors=2,
        spark_conf={"spark.sql.shuffle.partitions": "200", "spark.executor.memoryOverhead": "1G"},
    )
    # every k=v is emitted as --conf k=v
    assert "--conf" in cmd
    assert "spark.sql.shuffle.partitions=200" in cmd
    assert "spark.executor.memoryOverhead=1G" in cmd
    # script_path is still last
    assert cmd[-1] == "/tmp/j.py"


def test_run_spark_submit_returns_completed_process(monkeypatch):
    fake = MagicMock()
    fake.returncode = 0
    fake.stderr = "Submitted application application_1\n"
    with patch("spark_executor.core.spark_submit.subprocess.run", return_value=fake) as m:
        result = run_spark_submit(["spark-submit", "/tmp/x.py"])
    assert result is fake
    m.assert_called_once()
    # capture_output=True and text=True must be set
    kwargs = m.call_args.kwargs
    assert kwargs["capture_output"] is True
    assert kwargs["text"] is True


def test_run_spark_submit_raises_on_nonzero_return():
    fake = MagicMock()
    fake.returncode = 1
    fake.stderr = "boom"
    with patch("spark_executor.core.spark_submit.subprocess.run", return_value=fake):
        with pytest.raises(SparkSubmitError, match="spark-submit failed"):
            run_spark_submit(["spark-submit", "/tmp/x.py"])

• Step 2: Run tests to verify they fail

uv run pytest tests/unit/test_spark_submit.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.core.spark_submit'.

• Step 3: Implement spark_submit.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
import subprocess


class SparkSubmitError(Exception):
    """Raised when spark-submit exits with a non-zero return code."""


def build_spark_submit_command(
    *,
    master: str,
    deploy_mode: str,
    script_path: str,
    queue: str,
    executor_memory: str,
    executor_cores: int,
    num_executors: int,
    spark_conf: dict[str, str] | None = None,
) -> list[str]:
    cmd = [
        "spark-submit",
        "--master", master,
        "--deploy-mode", deploy_mode,
        "--queue", queue,
        "--executor-memory", executor_memory,
        "--executor-cores", str(executor_cores),
        "--num-executors", str(num_executors),
    ]
    for key, value in (spark_conf or {}).items():
        cmd.extend(["--conf", f"{key}={value}"])
    cmd.append(script_path)
    return cmd


def run_spark_submit(cmd: list[str]) -> "subprocess.CompletedProcess[str]":
    result = subprocess.run(cmd, capture_output=True, text=True, errors="replace")
    if result.returncode != 0:
        raise SparkSubmitError(
            f"spark-submit failed (rc={result.returncode}): {result.stderr}"
        )
    return result

• Step 4: Run tests to verify they pass

uv run pytest tests/unit/test_spark_submit.py -v

Expected: 5 passed.

• Step 5: Commit

git add spark_executor/core/spark_submit.py tests/unit/test_spark_submit.py
git commit -m "feat: add spark-submit command builder and runner"

---

Task 4: yarn_client — wrap yarn application -status, yarn logs, yarn application -kill

Files:

• Create: spark_executor/core/yarn_client.py
• Test: tests/unit/test_yarn_client.py

Interfaces:

• Produces:

  • def get_application_status(application_id: str) -> tuple[str, str] — returns (state, raw_output). State is the value after State : in the output. Raises YarnError if yarn exits non-zero.
  • def get_application_logs(application_id: str) -> str — returns full yarn logs -applicationId ... stdout. Raises YarnError on failure.
  • def kill_application(application_id: str) -> None — runs yarn application -kill .... Raises YarnError on failure.
  • class YarnError(Exception)

• Step 1: Write the failing tests

tests/unit/test_yarn_client.py:

# coding=utf-8
from unittest.mock import MagicMock, patch

import pytest

from spark_executor.core.yarn_client import (
    YarnError,
    get_application_logs,
    get_application_status,
    kill_application,
)


def _fake_proc(returncode: int, stdout: str = "", stderr: str = ""):
    p = MagicMock()
    p.returncode = returncode
    p.stdout = stdout
    p.stderr = stderr
    return p


def test_status_parses_state_line():
    fake = _fake_proc(
        0,
        stdout="Application Report :\n  State : RUNNING\n  ...\n",
    )
    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
        state, raw = get_application_status("application_1")
    assert state == "RUNNING"
    assert "RUNNING" in raw


def test_status_raises_on_nonzero_return():
    fake = _fake_proc(1, stderr="not found")
    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
        with pytest.raises(YarnError):
            get_application_status("application_x")


def test_logs_returns_stdout():
    fake = _fake_proc(0, stdout="log line 1\nlog line 2\n")
    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake) as m:
        out = get_application_logs("application_1")
    assert out == "log line 1\nlog line 2\n"
    args = m.call_args.args[0]
    assert args[:3] == ["yarn", "logs", "-applicationId"]
    assert args[3] == "application_1"


def test_kill_invokes_yarn_application_kill():
    fake = _fake_proc(0)
    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake) as m:
        kill_application("application_1")
    args = m.call_args.args[0]
    assert args[:3] == ["yarn", "application", "-kill"]
    assert args[3] == "application_1"


def test_kill_raises_on_nonzero_return():
    fake = _fake_proc(1, stderr="denied")
    with patch("spark_executor.core.yarn_client.subprocess.run", return_value=fake):
        with pytest.raises(YarnError):
            kill_application("application_1")

• Step 2: Run tests to verify they fail

uv run pytest tests/unit/test_yarn_client.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.core.yarn_client'.

• Step 3: Implement yarn_client.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
import re
import subprocess


class YarnError(Exception):
    """Raised when a yarn CLI invocation fails."""


_STATE_RE = re.compile(r"State\s*:\s*(\S+)")


def _run(cmd: list[str]) -> "subprocess.CompletedProcess[str]":
    return subprocess.run(cmd, capture_output=True, text=True, errors="replace")


def get_application_status(application_id: str) -> tuple[str, str]:
    proc = _run(["yarn", "application", "-status", application_id])
    if proc.returncode != 0:
        raise YarnError(
            f"yarn application -status failed (rc={proc.returncode}): {proc.stderr}"
        )
    match = _STATE_RE.search(proc.stdout)
    if not match:
        raise YarnError(f"Could not parse YARN state from output: {proc.stdout!r}")
    return match.group(1), proc.stdout


def get_application_logs(application_id: str) -> str:
    proc = _run(["yarn", "logs", "-applicationId", application_id])
    if proc.returncode != 0:
        raise YarnError(
            f"yarn logs failed (rc={proc.returncode}): {proc.stderr}"
        )
    return proc.stdout


def kill_application(application_id: str) -> None:
    proc = _run(["yarn", "application", "-kill", application_id])
    if proc.returncode != 0:
        raise YarnError(
            f"yarn application -kill failed (rc={proc.returncode}): {proc.stderr}"
        )

• Step 4: Run tests to verify they pass

uv run pytest tests/unit/test_yarn_client.py -v

Expected: 5 passed.

• Step 5: Commit

git add spark_executor/core/yarn_client.py tests/unit/test_yarn_client.py
git commit -m "feat: add yarn client (status, logs, kill)"

---

Task 5: job_store — in-memory job_id → Job registry

Files:

• Create: spark_executor/core/job_store.py
• Test: tests/unit/test_job_store.py

Interfaces:

• Produces:
  • class JobStore with methods:
    • put(job: Job) -> None
    • get(job_id: str) -> Job | None
    • list() -> list[Job]

Stage 1 keeps this trivial (a dict behind a class). Stage 3 replaces the implementation with SQLite.

• Step 1: Write the failing test

tests/unit/test_job_store.py:

# coding=utf-8
from datetime import datetime

from spark_executor.core.job_store import JobStore
from spark_executor.models import Job


def _job(jid: str) -> Job:
    return Job(
        job_id=jid,
        application_id=f"application_{jid}",
        script_path="/tmp/j.py",
        queue="default",
        submit_time=datetime(2026, 6, 24),
    )


def test_put_then_get_roundtrip():
    store = JobStore()
    store.put(_job("a"))
    assert store.get("a") is not None
    assert store.get("a").application_id == "application_a"


def test_get_missing_returns_none():
    store = JobStore()
    assert store.get("nope") is None


def test_list_returns_all_jobs():
    store = JobStore()
    store.put(_job("a"))
    store.put(_job("b"))
    assert {j.job_id for j in store.list()} == {"a", "b"}

• Step 2: Run test to verify it fails

uv run pytest tests/unit/test_job_store.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.core.job_store'.

• Step 3: Implement job_store.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from threading import Lock

from spark_executor.models import Job


class JobStore:
    """In-memory job registry. Stage-3 will swap this for SQLite."""

    def __init__(self) -> None:
        self._lock = Lock()
        self._jobs: dict[str, Job] = {}

    def put(self, job: Job) -> None:
        with self._lock:
            self._jobs[job.job_id] = job

    def get(self, job_id: str) -> Job | None:
        with self._lock:
            return self._jobs.get(job_id)

    def list(self) -> list[Job]:
        with self._lock:
            return list(self._jobs.values())

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_job_store.py -v

Expected: 3 passed.

• Step 5: Commit

git add spark_executor/core/job_store.py tests/unit/test_job_store.py
git commit -m "feat: add in-memory JobStore"

---

Task 6: connection_store — JSON-backed Connection CRUD

Files:

• Create: spark_executor/core/connection_store.py
• Test: tests/unit/test_connection_store.py

Interfaces:

• Produces:

  • class ConnectionStore with methods:
    • list_all() -> list[Connection]
    • get(name: str) -> Connection | None
    • save(conn: Connection) -> None — upsert by conn.name, then persist to disk atomically
    • delete(name: str) -> bool — returns True if a connection was removed
  • Module-level singleton store: ConnectionStore, defaulting to ./data/connections.json (overridable via SPARK_EXECUTOR_DATA_DIR env var).
  • class ConnectionNotFound(KeyError) — raised by tool layer when a name is missing.

• Step 1: Write the failing test

tests/unit/test_connection_store.py:

# coding=utf-8
import json
from pathlib import Path

from spark_executor.core import connection_store
from spark_executor.models import Connection


def _conn(name: str) -> Connection:
    return Connection(name=name, master="yarn")


def test_save_then_get_roundtrip(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
    connection_store.store.save(_conn("prod"))
    assert connection_store.store.get("prod") is not None
    assert connection_store.store.get("prod").master == "yarn"


def test_save_persists_to_json_file(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
    connection_store.store.save(_conn("prod"))
    path = tmp_path / "connections.json"
    assert path.is_file()
    payload = json.loads(path.read_text())
    assert "prod" in payload
    assert payload["prod"]["master"] == "yarn"


def test_get_missing_returns_none(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
    assert connection_store.store.get("missing") is None


def test_list_all_returns_empty_when_file_absent(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
    assert connection_store.store.list_all() == []


def test_list_all_returns_saved_connections(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
    connection_store.store.save(_conn("a"))
    connection_store.store.save(_conn("b"))
    names = {c.name for c in connection_store.store.list_all()}
    assert names == {"a", "b"}


def test_save_upserts_existing(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
    connection_store.store.save(_conn("prod"))
    updated = Connection(name="prod", master="spark://new:7077", deploy_mode="client")
    connection_store.store.save(updated)
    assert connection_store.store.get("prod").master == "spark://new:7077"
    assert connection_store.store.get("prod").deploy_mode == "client"


def test_delete_returns_true_when_present(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
    connection_store.store.save(_conn("prod"))
    assert connection_store.store.delete("prod") is True
    assert connection_store.store.get("prod") is None


def test_delete_returns_false_when_absent(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
    assert connection_store.store.delete("nope") is False

• Step 2: Run test to verify it fails

uv run pytest tests/unit/test_connection_store.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.core.connection_store'.

• Step 3: Implement connection_store.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
import json
import os
import tempfile
from pathlib import Path
from threading import Lock

from pydantic import ValidationError

from spark_executor.models import Connection

DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data")
DEFAULT_FILE_NAME = "connections.json"


class ConnectionNotFound(KeyError):
    """Raised when a named connection does not exist."""


class ConnectionStore:
    """JSON-backed CRUD for Connection records.

    File path: <DEFAULT_DATA_DIR>/<DEFAULT_FILE_NAME>.
    Writes are atomic: tempfile + os.replace.
    """

    def __init__(self, data_dir: str | None = None, file_name: str = DEFAULT_FILE_NAME) -> None:
        self._data_dir = data_dir or DEFAULT_DATA_DIR
        self._file_name = file_name
        self._lock = Lock()

    @property
    def path(self) -> Path:
        return Path(self._data_dir) / self._file_name

    def _load(self) -> dict[str, Connection]:
        if not self.path.exists():
            return {}
        raw = json.loads(self.path.read_text(encoding="utf-8"))
        return {name: Connection.model_validate(c) for name, c in raw.items()}

    def _dump(self, records: dict[str, Connection]) -> None:
        os.makedirs(self._data_dir, exist_ok=True)
        payload = {name: c.model_dump() for name, c in records.items()}
        # atomic write: tempfile in same dir, then replace
        fd, tmp_path = tempfile.mkstemp(
            prefix=self._file_name + ".", dir=self._data_dir
        )
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as f:
                json.dump(payload, f, indent=2, ensure_ascii=False)
            os.replace(tmp_path, self.path)
        except Exception:
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)
            raise

    def list_all(self) -> list[Connection]:
        with self._lock:
            return list(self._load().values())

    def get(self, name: str) -> Connection | None:
        with self._lock:
            return self._load().get(name)

    def save(self, conn: Connection) -> None:
        with self._lock:
            records = self._load()
            records[conn.name] = conn
            self._dump(records)

    def delete(self, name: str) -> bool:
        with self._lock:
            records = self._load()
            if name not in records:
                return False
            del records[name]
            self._dump(records)
            return True


# Module-level singleton; replaced in tests.
store = ConnectionStore()

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_connection_store.py -v

Expected: 8 passed.

• Step 5: Commit

git add spark_executor/core/connection_store.py tests/unit/test_connection_store.py
git commit -m "feat: add JSON-backed ConnectionStore"

---

Task 7: save_connection MCP tool

Files:

• Create: spark_executor/tools/connections.py
• Test: tests/unit/test_connection_tools.py::test_save_connection

Interfaces:

• Produces:

  • def save_connection(*, name: str, master: str, deploy_mode: str = "cluster", yarn_rm_url: str | None = None, spark_conf: dict[str, str] | None = None) -> dict[str, str] — returns {"name": name, "status": "SAVED"}. Upserts: re-saving with the same name overwrites.

• Step 1: Write the failing test (this file is shared by all four connection tools)

tests/unit/test_connection_tools.py:

# coding=utf-8
from pathlib import Path

import pytest

from spark_executor.core import connection_store
from spark_executor.tools import connections


@pytest.fixture(autouse=True)
def _fresh_store(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
    monkeypatch.setattr(connections, "store", connection_store.store)
    return connection_store.store


def test_save_connection(_fresh_store):
    out = connections.save_connection(name="prod", master="yarn")
    assert out == {"name": "prod", "status": "SAVED"}
    assert _fresh_store.get("prod") is not None


def test_save_connection_upserts(_fresh_store):
    connections.save_connection(name="prod", master="yarn")
    connections.save_connection(name="prod", master="spark://new:7077", deploy_mode="client")
    assert _fresh_store.get("prod").master == "spark://new:7077"
    assert _fresh_store.get("prod").deploy_mode == "client"


def test_save_connection_with_spark_conf(_fresh_store):
    connections.save_connection(
        name="dev",
        master="yarn",
        spark_conf={"spark.sql.shuffle.partitions": "200"},
    )
    assert _fresh_store.get("dev").spark_conf["spark.sql.shuffle.partitions"] == "200"

• Step 2: Run the new tests to verify they fail

uv run pytest tests/unit/test_connection_tools.py::test_save_connection -v

Expected: ModuleNotFoundError: No module named 'spark_executor.tools.connections'.

• Step 3: Implement connections.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from common.logging import logger
from spark_executor.core.connection_store import ConnectionStore, store
from spark_executor.models import Connection


def save_connection(
    *,
    name: str,
    master: str,
    deploy_mode: str = "cluster",
    yarn_rm_url: str | None = None,
    spark_conf: dict[str, str] | None = None,
) -> dict[str, str]:
    conn = Connection(
        name=name,
        master=master,
        deploy_mode=deploy_mode,
        yarn_rm_url=yarn_rm_url,
        spark_conf=spark_conf or {},
    )
    store.save(conn)
    logger.info(f"save_connection name={name} master={master}")
    return {"name": name, "status": "SAVED"}


def list_connections() -> list[dict[str, object]]:
    return [c.model_dump() for c in store.list_all()]


def get_connection(name: str) -> dict[str, object]:
    conn = store.get(name)
    if conn is None:
        raise KeyError(f"Unknown connection: {name}")
    return conn.model_dump()


def delete_connection(name: str) -> dict[str, str]:
    removed = store.delete(name)
    if not removed:
        raise KeyError(f"Unknown connection: {name}")
    logger.info(f"delete_connection name={name}")
    return {"name": name, "status": "DELETED"}

• Step 4: Run the save tests to verify they pass

uv run pytest tests/unit/test_connection_tools.py -v -k save_connection

Expected: 3 passed (the 3 test_save_connection_* cases). The other tests (list/get/delete) will fail because the implementation file already contains their functions — that is fine, they are filled in by Tasks 8–10 below.

• Step 5: Commit

git add spark_executor/tools/connections.py tests/unit/test_connection_tools.py
git commit -m "feat: add save_connection tool"

---

Task 8: list_connections MCP tool

Files:

• Modify: tests/unit/test_connection_tools.py (append test case)

• Test: tests/unit/test_connection_tools.py::test_list_connections

• Step 1: Append the failing test to tests/unit/test_connection_tools.py

Add at the end of the file:

def test_list_connections_empty(_fresh_store):
    assert connections.list_connections() == []


def test_list_connections_returns_all_saved(_fresh_store):
    connections.save_connection(name="a", master="yarn")
    connections.save_connection(name="b", master="spark://m:7077", deploy_mode="client")
    out = connections.list_connections()
    names = {c["name"] for c in out}
    assert names == {"a", "b"}
    masters = {c["name"]: c["master"] for c in out}
    assert masters == {"a": "yarn", "b": "spark://m:7077"}

• Step 2: Run the new tests to verify they pass

uv run pytest tests/unit/test_connection_tools.py -v -k list_connections

Expected: 2 passed. (list_connections was already implemented in Task 7 Step 3 as part of the same module.)

• Step 3: Commit

git add tests/unit/test_connection_tools.py
git commit -m "test: cover list_connections"

---

Task 9: get_connection MCP tool

Files:

• Modify: tests/unit/test_connection_tools.py (append test case)

• Step 1: Append the failing test to tests/unit/test_connection_tools.py

def test_get_connection_returns_dict(_fresh_store):
    connections.save_connection(name="prod", master="yarn", deploy_mode="cluster")
    out = connections.get_connection("prod")
    assert out["name"] == "prod"
    assert out["master"] == "yarn"
    assert out["deploy_mode"] == "cluster"
    assert out["yarn_rm_url"] is None
    assert out["spark_conf"] == {}


def test_get_connection_unknown_raises(_fresh_store):
    with pytest.raises(KeyError):
        connections.get_connection("missing")

• Step 2: Run the new tests to verify they pass

uv run pytest tests/unit/test_connection_tools.py -v -k get_connection

Expected: 2 passed.

• Step 3: Commit

git add tests/unit/test_connection_tools.py
git commit -m "test: cover get_connection"

---

Task 10: delete_connection MCP tool

Files:

• Modify: tests/unit/test_connection_tools.py (append test case)

• Step 1: Append the failing test to tests/unit/test_connection_tools.py

def test_delete_connection_returns_status(_fresh_store):
    connections.save_connection(name="prod", master="yarn")
    assert connections.delete_connection("prod") == {"name": "prod", "status": "DELETED"}
    assert _fresh_store.get("prod") is None


def test_delete_connection_unknown_raises(_fresh_store):
    with pytest.raises(KeyError):
        connections.delete_connection("missing")

• Step 2: Run the new tests to verify they pass

uv run pytest tests/unit/test_connection_tools.py -v -k delete_connection

Expected: 2 passed.

• Step 3: Commit

git add tests/unit/test_connection_tools.py
git commit -m "test: cover delete_connection"

---

Task 11: pending_store — JSON-backed store for PendingSubmission records

Submission is a two-step flow: prepare_submit_job writes a PendingSubmission here, the user reviews and confirms, and only then does confirm_submit_job actually invoke spark-submit. The pending store sits in ./data/pending_jobs.json (same data_dir as connection_store).

Files:

• Create: spark_executor/core/pending_store.py
• Test: tests/unit/test_pending_store.py

Interfaces:

• Produces:

  • class PendingStore with methods:
    • list_all() -> list[PendingSubmission]
    • get(pending_id: str) -> PendingSubmission | None
    • save(pending: PendingSubmission) -> None — upsert
    • delete(pending_id: str) -> bool
  • Module-level singleton store: PendingStore, defaulting to <SPARK_EXECUTOR_DATA_DIR or ./data>/pending_jobs.json.

• Step 1: Write the failing test

tests/unit/test_pending_store.py:

# coding=utf-8
from datetime import datetime
from pathlib import Path

import pytest

from spark_executor.core import pending_store
from spark_executor.models import PendingSubmission


def _pending(pid: str = "p_1") -> PendingSubmission:
    return PendingSubmission(
        pending_id=pid,
        connection="prod",
        master="yarn",
        deploy_mode="cluster",
        script_path="/tmp/j.py",
        queue="default",
        executor_memory="2G",
        executor_cores=2,
        num_executors=2,
        spark_conf={},
        created_at=datetime(2026, 6, 24),
    )


@pytest.fixture(autouse=True)
def _fresh(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(pending_store, "store", pending_store.PendingStore())


def test_save_then_get_roundtrip():
    pending_store.store.save(_pending())
    out = pending_store.store.get("p_1")
    assert out is not None
    assert out.connection == "prod"
    assert out.master == "yarn"


def test_save_persists_to_json(tmp_path: Path):
    pending_store.store.save(_pending())
    path = tmp_path / "pending_jobs.json"
    assert path.is_file()


def test_get_missing_returns_none():
    assert pending_store.store.get("missing") is None


def test_list_all_empty_when_file_absent():
    assert pending_store.store.list_all() == []


def test_list_all_returns_saved():
    pending_store.store.save(_pending("a"))
    pending_store.store.save(_pending("b"))
    assert {p.pending_id for p in pending_store.store.list_all()} == {"a", "b"}


def test_save_upserts_existing():
    pending_store.store.save(_pending("a"))
    updated = _pending("a")
    updated.status = "SUBMITTED"
    updated.job_id = "j_1"
    pending_store.store.save(updated)
    assert pending_store.store.get("a").status == "SUBMITTED"
    assert pending_store.store.get("a").job_id == "j_1"


def test_delete_returns_true_when_present():
    pending_store.store.save(_pending())
    assert pending_store.store.delete("p_1") is True
    assert pending_store.store.get("p_1") is None


def test_delete_returns_false_when_absent():
    assert pending_store.store.delete("nope") is False

• Step 2: Run test to verify it fails

uv run pytest tests/unit/test_pending_store.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.core.pending_store'.

• Step 3: Implement pending_store.py

pending_store.py is structurally identical to connection_store.py (Task 6) but operates on PendingSubmission records and a different file. Reuse the same atomic-write pattern.

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
import json
import os
import tempfile
from pathlib import Path
from threading import Lock

from spark_executor.models import PendingSubmission

DEFAULT_DATA_DIR = os.environ.get("SPARK_EXECUTOR_DATA_DIR", "./data")
DEFAULT_FILE_NAME = "pending_jobs.json"


class PendingStore:
    """JSON-backed CRUD for PendingSubmission records.

    File path: <DEFAULT_DATA_DIR>/<DEFAULT_FILE_NAME>.
    Atomic writes (tempfile + os.replace), same as ConnectionStore.
    """

    def __init__(self, data_dir: str | None = None, file_name: str = DEFAULT_FILE_NAME) -> None:
        self._data_dir = data_dir or DEFAULT_DATA_DIR
        self._file_name = file_name
        self._lock = Lock()

    @property
    def path(self) -> Path:
        return Path(self._data_dir) / self._file_name

    def _load(self) -> dict[str, PendingSubmission]:
        if not self.path.exists():
            return {}
        raw = json.loads(self.path.read_text(encoding="utf-8"))
        return {pid: PendingSubmission.model_validate(p) for pid, p in raw.items()}

    def _dump(self, records: dict[str, PendingSubmission]) -> None:
        os.makedirs(self._data_dir, exist_ok=True)
        payload = {pid: p.model_dump() for pid, p in records.items()}
        fd, tmp_path = tempfile.mkstemp(
            prefix=self._file_name + ".", dir=self._data_dir
        )
        try:
            with os.fdopen(fd, "w", encoding="utf-8") as f:
                json.dump(payload, f, indent=2, ensure_ascii=False, default=str)
            os.replace(tmp_path, self.path)
        except Exception:
            if os.path.exists(tmp_path):
                os.unlink(tmp_path)
            raise

    def list_all(self) -> list[PendingSubmission]:
        with self._lock:
            return list(self._load().values())

    def get(self, pending_id: str) -> PendingSubmission | None:
        with self._lock:
            return self._load().get(pending_id)

    def save(self, pending: PendingSubmission) -> None:
        with self._lock:
            records = self._load()
            records[pending.pending_id] = pending
            self._dump(records)

    def delete(self, pending_id: str) -> bool:
        with self._lock:
            records = self._load()
            if pending_id not in records:
                return False
            del records[pending_id]
            self._dump(records)
            return True


# Module-level singleton; replaced in tests.
store = PendingStore()

Note: default=str is required on json.dump because created_at is a datetime.

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_pending_store.py -v

Expected: 8 passed.

• Step 5: Commit

git add spark_executor/core/pending_store.py tests/unit/test_pending_store.py
git commit -m "feat: add JSON-backed PendingStore"

---

Task 12: prepare_submit_job MCP tool — record params, do NOT submit

prepare_submit_job is the first half of the submit flow. It snapshots the connection's master / deploy_mode / spark_conf into a PendingSubmission, persists it, and returns the pending_id. No spark-submit is invoked.

Files:

• Create: spark_executor/tools/submit.py (this replaces the direct-submit tool that previously lived here)
• Create: tests/unit/test_submit_tool.py (this file is fresh — the old direct-submit tests are removed)

Interfaces:

• Produces:

  • def prepare_submit_job(*, connection: str, script_path: str, queue: str = "default", executor_memory: str = "2G", executor_cores: int = 2, num_executors: int = 2) -> dict[str, object] — returns {"pending_id": ..., "status": "PENDING", "parameters": {...}}. The full PendingSubmission (minus pending_id, created_at, status) is echoed back in parameters so the user (or agent) can review exactly what will be submitted. Raises KeyError if connection is unknown.
  • Module-level pending_store (singleton) imported from spark_executor.core.pending_store.

• Step 1: Write the failing test

tests/unit/test_submit_tool.py:

# coding=utf-8
from pathlib import Path
from unittest.mock import patch

import pytest

from spark_executor.core import connection_store, pending_store
from spark_executor.core.pending_store import PendingStore
from spark_executor.models import Connection
from spark_executor.tools import submit


@pytest.fixture(autouse=True)
def _fresh(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(connection_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(connection_store, "store", connection_store.ConnectionStore())
    monkeypatch.setattr(pending_store, "DEFAULT_DATA_DIR", str(tmp_path))
    monkeypatch.setattr(pending_store, "store", PendingStore())
    submit.pending_store = pending_store.store
    connection_store.store.save(Connection(name="prod", master="yarn", deploy_mode="cluster"))


def _last_pending_id() -> str:
    pending = submit.pending_store.list_all()
    assert pending, "no pending submission was created"
    return pending[-1].pending_id


def test_prepare_does_not_invoke_spark_submit(monkeypatch):
    called = []
    monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: called.append(cmd))
    out = submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
    assert called == []  # never called
    assert out["status"] == "PENDING"
    assert out["pending_id"].startswith("p_")


def test_prepare_persists_pending_with_snapshot(monkeypatch):
    monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
    connection_store.store.save(
        Connection(
            name="prod",
            master="yarn",
            deploy_mode="cluster",
            spark_conf={"spark.sql.shuffle.partitions": "200"},
        )
    )
    submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py", queue="research")
    p = submit.pending_store.get(_last_pending_id())
    assert p is not None
    assert p.connection == "prod"
    assert p.master == "yarn"
    assert p.deploy_mode == "cluster"
    assert p.spark_conf == {"spark.sql.shuffle.partitions": "200"}
    assert p.queue == "research"
    assert p.script_path == "/tmp/j.py"


def test_prepare_raises_for_unknown_connection():
    with pytest.raises(KeyError, match="missing"):
        submit.prepare_submit_job(connection="missing", script_path="/tmp/j.py")


def test_prepare_snapshots_connection_at_prepare_time(monkeypatch):
    """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")
    p = submit.pending_store.get(_last_pending_id())
    # Now mutate the connection
    connection_store.store.save(
        Connection(name="prod", master="spark://attacker:7077", deploy_mode="client")
    )
    # Snapshot is unchanged
    p2 = submit.pending_store.get(p.pending_id)
    assert p2.master == "yarn"
    assert p2.deploy_mode == "cluster"

• Step 2: Run test to verify it fails

uv run pytest tests/unit/test_submit_tool.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.tools.submit'.

• Step 3: Implement submit.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
import secrets
from datetime import datetime

from common.logging import logger
from spark_executor.core.connection_store import store as conn_store
from spark_executor.core.pending_store import store as pending_store
from spark_executor.models import PendingSubmission


def _new_pending_id() -> str:
    return "p_" + secrets.token_hex(6)


def prepare_submit_job(
    *,
    connection: str,
    script_path: str,
    queue: str = "default",
    executor_memory: str = "2G",
    executor_cores: int = 2,
    num_executors: int = 2,
) -> dict[str, object]:
    """Snapshot connection params and persist a PendingSubmission. Does NOT submit."""
    conn = conn_store.get(connection)
    if conn is None:
        raise KeyError(f"Unknown connection: {connection}")

    pending_id = _new_pending_id()
    pending = PendingSubmission(
        pending_id=pending_id,
        connection=connection,
        master=conn.master,
        deploy_mode=conn.deploy_mode,
        script_path=script_path,
        queue=queue,
        executor_memory=executor_memory,
        executor_cores=executor_cores,
        num_executors=num_executors,
        spark_conf=dict(conn.spark_conf),
        created_at=datetime.utcnow(),
        status="PENDING",
    )
    pending_store.save(pending)
    logger.info(
        f"prepare_submit_job pending_id={pending_id} connection={connection} master={conn.master}"
    )
    return {
        "pending_id": pending_id,
        "status": "PENDING",
        "parameters": pending.model_dump(),
    }

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_submit_tool.py -v

Expected: 4 passed.

• Step 5: Commit

git add spark_executor/tools/submit.py tests/unit/test_submit_tool.py
git commit -m "feat: add prepare_submit_job (no spark-submit invocation)"

---

Task 13: confirm_submit_job MCP tool — actually run spark-submit

confirm_submit_job is the second half of the submit flow. It loads a PendingSubmission by pending_id, builds the spark-submit command from the snapshot (not the live connection), runs it, creates the Job, and marks the pending entry as SUBMITTED (or FAILED).

Files:

• Modify: spark_executor/tools/submit.py (append confirm_submit_job)
• Modify: tests/unit/test_submit_tool.py (append tests)

Interfaces:

• Produces:

  • def confirm_submit_job(*, pending_id: str) -> SubmitResult — raises KeyError if pending_id unknown. Raises ValueError if the pending is not in PENDING status. On SparkSubmitError, the pending is marked FAILED with the error and re-raised.

• Step 1: Append the failing tests to tests/unit/test_submit_tool.py

def test_confirm_invokes_spark_submit_and_marks_submitted(monkeypatch):
    submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
    pid = _last_pending_id()
    fake_proc = type("P", (), {
        "returncode": 0,
        "stderr": "tracking URL: http://rm:8088/proxy/application_17400000001/\n",
    })()
    with patch("spark_executor.tools.submit.run_spark_submit", return_value=fake_proc) as m:
        result = submit.confirm_submit_job(pending_id=pid)
    cmd = m.call_args.args[0]
    assert "yarn" in cmd
    assert "cluster" in cmd
    assert cmd[-1] == "/tmp/j.py"
    assert result.application_id == "application_17400000001"
    # pending updated
    p = submit.pending_store.get(pid)
    assert p.status == "SUBMITTED"
    assert p.application_id == "application_17400000001"
    assert p.job_id is not None


def test_confirm_raises_for_unknown_pending_id():
    with pytest.raises(KeyError, match="missing"):
        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")
    pid = _last_pending_id()
    # Mark it CANCELLED first
    p = submit.pending_store.get(pid)
    p.status = "CANCELLED"
    submit.pending_store.save(p)
    with pytest.raises(ValueError, match="CANCELLED"):
        submit.confirm_submit_job(pending_id=pid)


def test_confirm_marks_failed_on_spark_submit_error(monkeypatch):
    from spark_executor.core.spark_submit import SparkSubmitError
    submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
    pid = _last_pending_id()
    def _raise(_cmd):
        raise SparkSubmitError("boom")
    monkeypatch.setattr(submit, "run_spark_submit", _raise)
    with pytest.raises(SparkSubmitError):
        submit.confirm_submit_job(pending_id=pid)
    p = submit.pending_store.get(pid)
    assert p.status == "FAILED"
    assert "boom" in (p.error or "")

• Step 2: Run the new tests to verify they fail

uv run pytest tests/unit/test_submit_tool.py -v -k confirm

Expected: AttributeError: module 'spark_executor.tools.submit' has no attribute 'confirm_submit_job'.

• Step 3: Append confirm_submit_job to submit.py

import uuid
from datetime import datetime as _datetime

from spark_executor.core.job_store import JobStore
from spark_executor.core.log_parser import parse_spark_submit_output
from spark_executor.core.spark_submit import (
    SparkSubmitError,
    build_spark_submit_command,
    run_spark_submit,
)
from spark_executor.models import Job, SubmitResult

# Module-level job store singleton; replaced in tests.
job_store: JobStore = JobStore()


def confirm_submit_job(*, pending_id: str) -> SubmitResult:
    """Actually invoke spark-submit for a previously-prepared PendingSubmission."""
    pending = pending_store.get(pending_id)
    if pending is None:
        raise KeyError(f"Unknown pending_id: {pending_id}")
    if pending.status != "PENDING":
        raise ValueError(
            f"pending_id {pending_id} is in status {pending.status!r}, not PENDING"
        )

    cmd = build_spark_submit_command(
        master=pending.master,
        deploy_mode=pending.deploy_mode,
        script_path=pending.script_path,
        queue=pending.queue,
        executor_memory=pending.executor_memory,
        executor_cores=pending.executor_cores,
        num_executors=pending.num_executors,
        spark_conf=pending.spark_conf,
    )
    logger.info(f"confirm_submit_job pending_id={pending_id} cmd={cmd}")
    try:
        result = run_spark_submit(cmd)
    except SparkSubmitError as exc:
        pending.status = "FAILED"
        pending.error = str(exc)
        pending_store.save(pending)
        raise

    application_id, tracking_url = parse_spark_submit_output(result.stderr)
    job_id = uuid.uuid4().hex[:12]

    job_store.put(
        Job(
            job_id=job_id,
            application_id=application_id,
            script_path=pending.script_path,
            queue=pending.queue,
            submit_time=_datetime.utcnow(),
            connection=pending.connection,
        )
    )

    pending.status = "SUBMITTED"
    pending.job_id = job_id
    pending.application_id = application_id
    pending_store.save(pending)
    logger.info(
        f"confirm_submit_job pending_id={pending_id} job_id={job_id} application_id={application_id}"
    )
    return SubmitResult(
        job_id=job_id,
        application_id=application_id,
        tracking_url=tracking_url,
    )

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_submit_tool.py -v

Expected: 8 passed (4 prepare + 4 confirm).

• Step 5: Commit

git add spark_executor/tools/submit.py tests/unit/test_submit_tool.py
git commit -m "feat: add confirm_submit_job (two-step submit flow)"

---

Task 14: list_pending_jobs MCP tool

Files:

• Modify: spark_executor/tools/submit.py (append list_pending_jobs)
• Modify: tests/unit/test_submit_tool.py (append tests)

Interfaces:

• Produces:

  • def list_pending_jobs() -> list[dict[str, object]] — returns the model dump of every PendingSubmission in the store.

• Step 1: Append the failing test to tests/unit/test_submit_tool.py

def test_list_pending_jobs_empty():
    assert submit.list_pending_jobs() == []


def test_list_pending_jobs_returns_all(monkeypatch):
    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")
    out = submit.list_pending_jobs()
    assert {p["script_path"] for p in out} == {"/tmp/a.py", "/tmp/b.py"}
    assert all(p["status"] == "PENDING" for p in out)

• Step 2: Run the new tests to verify they fail

uv run pytest tests/unit/test_submit_tool.py -v -k list_pending

Expected: AttributeError: module 'spark_executor.tools.submit' has no attribute 'list_pending_jobs'.

• Step 3: Append list_pending_jobs to submit.py

def list_pending_jobs() -> list[dict[str, object]]:
    return [p.model_dump() for p in pending_store.list_all()]

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_submit_tool.py -v -k list_pending

Expected: 2 passed.

• Step 5: Commit

git add spark_executor/tools/submit.py tests/unit/test_submit_tool.py
git commit -m "feat: add list_pending_jobs"

---

Task 15: get_pending_job MCP tool

Files:

• Modify: spark_executor/tools/submit.py (append get_pending_job)
• Modify: tests/unit/test_submit_tool.py (append tests)

Interfaces:

• Produces:

  • def get_pending_job(pending_id: str) -> dict[str, object] — returns the model dump. Raises KeyError if unknown.

• Step 1: Append the failing test to tests/unit/test_submit_tool.py

def test_get_pending_job_returns_dump(monkeypatch):
    monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
    submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
    pid = _last_pending_id()
    out = submit.get_pending_job(pid)
    assert out["pending_id"] == pid
    assert out["connection"] == "prod"
    assert out["status"] == "PENDING"


def test_get_pending_job_unknown_raises():
    with pytest.raises(KeyError):
        submit.get_pending_job("missing")

• Step 2: Run the new tests to verify they fail

uv run pytest tests/unit/test_submit_tool.py -v -k get_pending

Expected: AttributeError: module 'spark_executor.tools.submit' has no attribute 'get_pending_job'.

• Step 3: Append get_pending_job to submit.py

def get_pending_job(pending_id: str) -> dict[str, object]:
    p = pending_store.get(pending_id)
    if p is None:
        raise KeyError(f"Unknown pending_id: {pending_id}")
    return p.model_dump()

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_submit_tool.py -v -k get_pending

Expected: 2 passed.

• Step 5: Commit

git add spark_executor/tools/submit.py tests/unit/test_submit_tool.py
git commit -m "feat: add get_pending_job"

---

Task 16: cancel_pending_job MCP tool

Files:

• Modify: spark_executor/tools/submit.py (append cancel_pending_job)
• Modify: tests/unit/test_submit_tool.py (append tests)

Interfaces:

• Produces:

  • def cancel_pending_job(pending_id: str) -> dict[str, str] — flips the entry's status to CANCELLED (the record is kept for audit). Refuses to cancel a record that is already SUBMITTED or FAILED (raises ValueError). Unknown id raises KeyError.

• Step 1: Append the failing test to tests/unit/test_submit_tool.py

def test_cancel_pending_job_marks_cancelled(monkeypatch):
    monkeypatch.setattr(submit, "run_spark_submit", lambda cmd: None)
    submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
    pid = _last_pending_id()
    out = submit.cancel_pending_job(pid)
    assert out == {"pending_id": pid, "status": "CANCELLED"}
    assert submit.pending_store.get(pid).status == "CANCELLED"


def test_cancel_pending_job_unknown_raises():
    with pytest.raises(KeyError):
        submit.cancel_pending_job("missing")


def test_cancel_pending_job_refuses_submitted(monkeypatch):
    submit.prepare_submit_job(connection="prod", script_path="/tmp/j.py")
    pid = _last_pending_id()
    p = submit.pending_store.get(pid)
    p.status = "SUBMITTED"
    submit.pending_store.save(p)
    with pytest.raises(ValueError, match="SUBMITTED"):
        submit.cancel_pending_job(pid)

• Step 2: Run the new tests to verify they fail

uv run pytest tests/unit/test_submit_tool.py -v -k cancel_pending

Expected: AttributeError: module 'spark_executor.tools.submit' has no attribute 'cancel_pending_job'.

• Step 3: Append cancel_pending_job to submit.py

def cancel_pending_job(pending_id: str) -> dict[str, str]:
    p = pending_store.get(pending_id)
    if p is None:
        raise KeyError(f"Unknown pending_id: {pending_id}")
    if p.status in ("SUBMITTED", "FAILED"):
        raise ValueError(
            f"pending_id {pending_id} is in status {p.status!r} and cannot be cancelled"
        )
    p.status = "CANCELLED"
    pending_store.save(p)
    logger.info(f"cancel_pending_job pending_id={pending_id}")
    return {"pending_id": pending_id, "status": "CANCELLED"}

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_submit_tool.py -v -k cancel_pending

Expected: 3 passed.

• Step 5: Commit

git add spark_executor/tools/submit.py tests/unit/test_submit_tool.py
git commit -m "feat: add cancel_pending_job"

---

Task 17: get_job_status MCP tool

Files:

• Create: spark_executor/tools/status.py
• Test: tests/unit/test_status_tool.py

Interfaces:

• Produces:

  • def get_job_status(job_id: str) -> JobStatus — looks up the job in the store to get application_id, then calls get_application_status. Raises KeyError if job_id unknown.

• Step 1: Write the failing test

tests/unit/test_status_tool.py:

# coding=utf-8
from datetime import datetime
from unittest.mock import patch

from spark_executor.core.job_store import JobStore
from spark_executor.models import Job
from spark_executor.tools import status


def test_get_job_status_returns_state():
    status.store = JobStore()
    status.store.put(
        Job(
            job_id="abc",
            application_id="application_1",
            script_path="/tmp/j.py",
            queue="default",
            submit_time=datetime(2026, 6, 24),
            connection="prod",
        )
    )
    with patch(
        "spark_executor.tools.status.get_application_status",
        return_value=("RUNNING", "State : RUNNING\n"),
    ):
        out = status.get_job_status("abc")
    assert out.application_id == "application_1"
    assert out.state == "RUNNING"
    assert "RUNNING" in out.raw


def test_get_job_status_raises_for_unknown_job():
    status.store = JobStore()
    import pytest
    with pytest.raises(KeyError):
        status.get_job_status("missing")

• Step 2: Run test to verify it fails

uv run pytest tests/unit/test_status_tool.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.tools.status'.

• Step 3: Implement status.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from common.logging import logger
from spark_executor.core.job_store import JobStore
from spark_executor.core.yarn_client import get_application_status
from spark_executor.models import JobStatus

store = JobStore()


def get_job_status(job_id: str) -> JobStatus:
    job = store.get(job_id)
    if job is None:
        raise KeyError(f"Unknown job_id: {job_id}")
    state, raw = get_application_status(job.application_id)
    logger.info(f"get_job_status job_id={job_id} state={state}")
    return JobStatus(application_id=job.application_id, state=state, raw=raw)

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_status_tool.py -v

Expected: 2 passed.

• Step 5: Commit

git add spark_executor/tools/status.py tests/unit/test_status_tool.py
git commit -m "feat: add get_job_status tool"

---

Task 18: get_job_logs MCP tool (tail to 5000 chars)

Files:

• Create: spark_executor/tools/logs.py
• Test: tests/unit/test_logs_tool.py

Interfaces:

• Produces:

  • def get_job_logs(job_id: str, tail_chars: int = 5000) -> str — looks up application_id from the store, calls get_application_logs, returns the last tail_chars characters (default 5000). Raises KeyError if job_id unknown.

• Step 1: Write the failing test

tests/unit/test_logs_tool.py:

# coding=utf-8
from datetime import datetime
from unittest.mock import patch

from spark_executor.core.job_store import JobStore
from spark_executor.models import Job
from spark_executor.tools import logs


def _seed(job_id="abc", app_id="application_1"):
    logs.store = JobStore()
    logs.store.put(
        Job(
            job_id=job_id,
            application_id=app_id,
            script_path="/tmp/j.py",
            queue="default",
            submit_time=datetime(2026, 6, 24),
            connection="prod",
        )
    )


def test_get_job_logs_tails_to_default_5000():
    _seed()
    big = "x" * 8000 + "\nEND"
    with patch("spark_executor.tools.logs.get_application_logs", return_value=big):
        out = logs.get_job_logs("abc")
    assert out.endswith("END")
    assert len(out) == 5000


def test_get_job_logs_respects_custom_tail_chars():
    _seed()
    with patch(
        "spark_executor.tools.logs.get_application_logs",
        return_value="0123456789",
    ):
        out = logs.get_job_logs("abc", tail_chars=3)
    assert out == "789"


def test_get_job_logs_raises_for_unknown_job():
    _seed()
    import pytest
    with pytest.raises(KeyError):
        logs.get_job_logs("missing")

• Step 2: Run test to verify it fails

uv run pytest tests/unit/test_logs_tool.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.tools.logs'.

• Step 3: Implement logs.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from common.logging import logger
from spark_executor.core.job_store import JobStore
from spark_executor.core.yarn_client import get_application_logs

store = JobStore()


def get_job_logs(job_id: str, tail_chars: int = 5000) -> str:
    job = store.get(job_id)
    if job is None:
        raise KeyError(f"Unknown job_id: {job_id}")
    full = get_application_logs(job.application_id)
    tailed = full[-tail_chars:] if len(full) > tail_chars else full
    logger.info(
        f"get_job_logs job_id={job_id} application_id={job.application_id} "
        f"chars={len(tailed)}"
    )
    return tailed

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_logs_tool.py -v

Expected: 3 passed.

• Step 5: Commit

git add spark_executor/tools/logs.py tests/unit/test_logs_tool.py
git commit -m "feat: add get_job_logs tool with tail"

---

Task 19: kill_job MCP tool

Files:

• Create: spark_executor/tools/kill.py
• Test: tests/unit/test_kill_tool.py

Interfaces:

• Produces:

  • def kill_job(job_id: str) -> dict[str, str] — returns {"job_id": ..., "application_id": ..., "status": "KILLED"}. Raises KeyError if job_id unknown.

• Step 1: Write the failing test

tests/unit/test_kill_tool.py:

# coding=utf-8
from datetime import datetime
from unittest.mock import patch

from spark_executor.core.job_store import JobStore
from spark_executor.models import Job
from spark_executor.tools import kill


def test_kill_job_calls_yarn_kill():
    kill.store = JobStore()
    kill.store.put(
        Job(
            job_id="abc",
            application_id="application_1",
            script_path="/tmp/j.py",
            queue="default",
            submit_time=datetime(2026, 6, 24),
            connection="prod",
        )
    )
    with patch("spark_executor.tools.kill.kill_application") as m:
        result = kill.kill_job("abc")
    m.assert_called_once_with("application_1")
    assert result == {
        "job_id": "abc",
        "application_id": "application_1",
        "status": "KILLED",
    }


def test_kill_job_raises_for_unknown_job():
    kill.store = JobStore()
    import pytest
    with pytest.raises(KeyError):
        kill.kill_job("missing")

• Step 2: Run test to verify it fails

uv run pytest tests/unit/test_kill_tool.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.tools.kill'.

• Step 3: Implement kill.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from common.logging import logger
from spark_executor.core.job_store import JobStore
from spark_executor.core.yarn_client import kill_application

store = JobStore()


def kill_job(job_id: str) -> dict[str, str]:
    job = store.get(job_id)
    if job is None:
        raise KeyError(f"Unknown job_id: {job_id}")
    kill_application(job.application_id)
    logger.info(f"kill_job job_id={job_id} application_id={job.application_id}")
    return {
        "job_id": job_id,
        "application_id": job.application_id,
        "status": "KILLED",
    }

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_kill_tool.py -v

Expected: 2 passed.

• Step 5: Commit

git add spark_executor/tools/kill.py tests/unit/test_kill_tool.py
git commit -m "feat: add kill_job tool"

---

Task 20: Wire all tools into the FastAPI app and verify MCP discovery

fastapi-mcp auto-discovers FastAPI routes and exposes them as MCP tools. We mount the twelve tool functions as POST routes on the existing app in spark_executor/server.py.

Files:

• Modify: spark_executor/server.py

• Test: tests/integration/__init__.py (empty header file)

• Test: tests/integration/test_mcp_routes.py

• Step 1: Create integration test directory

tests/integration/__init__.py (standard header, empty body).

• Step 2: Write the integration test

tests/integration/test_mcp_routes.py:

# coding=utf-8
from fastapi.testclient import TestClient

from spark_executor.server import app


def test_health_still_present():
    c = TestClient(app)
    r = c.get("/health")
    assert r.status_code == 200
    assert r.json() == {"status": "ok"}


def test_twelve_tool_routes_registered():
    paths = {r.path for r in app.routes}
    for path in (
        # pending-submission flow (5)
        "/prepare_submit_job",
        "/confirm_submit_job",
        "/list_pending_jobs",
        "/get_pending_job",
        "/cancel_pending_job",
        # job lifecycle (3)
        "/get_job_status",
        "/get_job_logs",
        "/kill_job",
        # connection management (4)
        "/save_connection",
        "/list_connections",
        "/get_connection",
        "/delete_connection",
    ):
        assert path in paths, f"missing MCP tool route: {path}"

• Step 3: Run test to verify it fails

uv run pytest tests/integration/test_mcp_routes.py -v

Expected: AssertionError: missing MCP tool route: /prepare_submit_job (or similar — only /health is registered).

• Step 4: Add the twelve tool routes to server.py

Replace spark_executor/server.py with:

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from fastapi import FastAPI

from spark_executor.tools.connections import (
    delete_connection,
    get_connection,
    list_connections,
    save_connection,
)
from spark_executor.tools.kill import kill_job
from spark_executor.tools.logs import get_job_logs
from spark_executor.tools.status import get_job_status
from spark_executor.tools.submit import (
    cancel_pending_job,
    confirm_submit_job,
    get_pending_job,
    list_pending_jobs,
    prepare_submit_job,
)

app = FastAPI(title="Spark Executor", version="0.0.1", description="Spark Executor MCP Server")


@app.get("/health")
def health_check():
    return {"status": "ok"}


# MCP tool routes. fastapi-mcp discovers these and registers them as MCP tools.

# --- Pending submission flow (two-step submit) ---

@app.post("/prepare_submit_job")
def _prepare_submit_job(connection: str, script_path: str, queue: str = "default",
                        executor_memory: str = "2G", executor_cores: int = 2,
                        num_executors: int = 2):
    return prepare_submit_job(
        connection=connection,
        script_path=script_path,
        queue=queue,
        executor_memory=executor_memory,
        executor_cores=executor_cores,
        num_executors=num_executors,
    )


@app.post("/confirm_submit_job")
def _confirm_submit_job(pending_id: str):
    return confirm_submit_job(pending_id=pending_id)


@app.post("/list_pending_jobs")
def _list_pending_jobs():
    return list_pending_jobs()


@app.post("/get_pending_job")
def _get_pending_job(pending_id: str):
    return get_pending_job(pending_id)


@app.post("/cancel_pending_job")
def _cancel_pending_job(pending_id: str):
    return cancel_pending_job(pending_id)


# --- Spark job tools ---

@app.post("/get_job_status")
def _get_job_status(job_id: str):
    return get_job_status(job_id)


@app.post("/get_job_logs")
def _get_job_logs(job_id: str, tail_chars: int = 5000):
    return get_job_logs(job_id, tail_chars=tail_chars)


@app.post("/kill_job")
def _kill_job(job_id: str):
    return kill_job(job_id)


# --- Connection management tools ---

@app.post("/save_connection")
def _save_connection(name: str, master: str, deploy_mode: str = "cluster",
                     yarn_rm_url: str | None = None,
                     spark_conf: dict[str, str] | None = None):
    return save_connection(
        name=name,
        master=master,
        deploy_mode=deploy_mode,
        yarn_rm_url=yarn_rm_url,
        spark_conf=spark_conf,
    )


@app.post("/list_connections")
def _list_connections():
    return list_connections()


@app.post("/get_connection")
def _get_connection(name: str):
    return get_connection(name)


@app.post("/delete_connection")
def _delete_connection(name: str):
    return delete_connection(name)

• Step 5: Run integration test to verify it passes

uv run pytest tests/integration/test_mcp_routes.py -v

Expected: 2 passed.

• Step 6: Run the full suite to confirm no regressions

uv run pytest -v

Expected: all unit + integration tests pass (8 + 4 + 5 + 5 + 3 + 8 + 9 + 8 + 15 + 2 + 3 + 2 + 2 = 74).

• Step 7: Smoke-test the running server (manual)

uv run main.py &
sleep 2
curl -s http://localhost:8000/health
curl -s -X POST http://localhost:8000/save_connection \
    -H "Content-Type: application/json" \
    -d '{"name":"prod","master":"yarn"}'
curl -s -X POST http://localhost:8000/prepare_submit_job \
    -H "Content-Type: application/json" \
    -d '{"connection":"prod","script_path":"/tmp/jobs/demo.py"}'
# (no submit happens; pending_id is returned)
curl -s -X POST http://localhost:8000/list_pending_jobs
curl -s -X POST http://localhost:8000/spark-executor-mcp/mcp -H "Content-Type: application/json" \
    -d '{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}'
kill %1

Expected:

• /health returns {"status":"ok"}.

• save_connection returns {"name":"prod","status":"SAVED"}; ./data/connections.json is created.

• prepare_submit_job returns a payload containing pending_id (prefix p_); ./data/pending_jobs.json is created. No spark-submit is invoked.

• list_pending_jobs returns the prepared entry with status: "PENDING".

• tools/list enumerates all 12 tools: prepare_submit_job, confirm_submit_job, list_pending_jobs, get_pending_job, cancel_pending_job, get_job_status, get_job_logs, kill_job, save_connection, list_connections, get_connection, delete_connection.

• Step 8: Commit

git add spark_executor/server.py tests/integration
git commit -m "feat: register 12 MCP tools (pending flow + job lifecycle + connections)"

---

Stage 1 checkpoint

After Task 20 you have a working minimum-closed-loop MCP server with cluster-connection management and a two-step pending submission flow:

• save_connection / list_connections / get_connection / delete_connection — persist to ./data/connections.json.
• prepare_submit_job(connection, ...) — snapshots the connection's master / deploy_mode / spark_conf into a PendingSubmission, writes ./data/pending_jobs.json, returns pending_id. Does NOT invoke spark-submit.
• list_pending_jobs / get_pending_job — read pending entries.
• cancel_pending_job — flips a PENDING entry to CANCELLED.
• confirm_submit_job(pending_id) — only on user second confirmation: runs spark-submit with the snapshot, parses application_id, creates a Job in the in-memory JobStore, marks the pending entry SUBMITTED.
• get_job_status → yarn application -status, parses State : ....
• get_job_logs → yarn logs -applicationId ..., tails to 5000 chars.
• kill_job → yarn application -kill.

The end-to-end agent loop is therefore reviewable:

prepare_submit_job(...)    ->  pending_id
<user reviews pending_id>
confirm_submit_job(pending_id)  ->  job_id, application_id
get_job_status(job_id)     ->  RUNNING | SUCCEEDED | FAILED | KILLED
get_job_logs(job_id)       ->  tail of driver logs
kill_job(job_id)           ->  KILLED

Stop here and demo before continuing. Stage 2 is a small additive change; Stage 3 is a real refactor.

---

Stage 2 — LLM-driven PySpark generation (★★☆☆☆)

Adds one tool, write_job_file, that takes a PySpark code string and writes it to a timestamped file under /tmp/jobs/. The Agent can then call submit_job on the path. No new infrastructure beyond one file.

---

Task 21: job_writer — write PySpark code to disk

Files:

• Create: spark_executor/core/job_writer.py
• Test: tests/unit/test_job_writer.py

Interfaces:

• Produces:

  • def write_job_file(code: str, jobs_dir: str = "/tmp/jobs") -> str — writes code to <jobs_dir>/job_<timestamp>_<rand>.py, creates the directory if needed, returns the absolute path.

• Step 1: Write the failing test

tests/unit/test_job_writer.py:

# coding=utf-8
import os
import tempfile
from pathlib import Path

from spark_executor.core.job_writer import write_job_file


def test_write_job_file_creates_file_with_code(tmp_path: Path):
    out = write_job_file("print('hi')\n", jobs_dir=str(tmp_path))
    assert os.path.isfile(out)
    assert out.startswith(str(tmp_path))
    assert out.endswith(".py")
    with open(out) as f:
        assert f.read() == "print('hi')\n"


def test_write_job_file_creates_jobs_dir_if_missing(tmp_path: Path):
    target = tmp_path / "newdir"
    out = write_job_file("x = 1\n", jobs_dir=str(target))
    assert target.is_dir()
    assert os.path.isfile(out)

• Step 2: Run test to verify it fails

uv run pytest tests/unit/test_job_writer.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.core.job_writer'.

• Step 3: Implement job_writer.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
import os
import secrets
from datetime import datetime


def write_job_file(code: str, jobs_dir: str = "/tmp/jobs") -> str:
    os.makedirs(jobs_dir, exist_ok=True)
    stamp = datetime.utcnow().strftime("%Y%m%d%H%M%S")
    name = f"job_{stamp}_{secrets.token_hex(3)}.py"
    path = os.path.join(jobs_dir, name)
    with open(path, "w", encoding="utf-8") as f:
        f.write(code)
    return os.path.abspath(path)

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_job_writer.py -v

Expected: 2 passed.

• Step 5: Commit

git add spark_executor/core/job_writer.py tests/unit/test_job_writer.py
git commit -m "feat: add job_writer for LLM-generated PySpark"

---

Task 22: write_job_file MCP tool

Files:

• Create: spark_executor/tools/write_job.py
• Test: tests/unit/test_write_job_tool.py
• Modify: spark_executor/server.py (register the new route)
• Modify: tests/integration/test_mcp_routes.py (assert new route registered)

Interfaces:

• Produces:

  • def write_job_file(code: str) -> dict[str, str] — returns {"script_path": "..."}. The code parameter is the full PySpark source to write to disk.

• Step 1: Write the failing test

tests/unit/test_write_job_tool.py:

# coding=utf-8
import os
import tempfile
from pathlib import Path

from spark_executor.tools import generate


def test_write_job_file_writes_and_returns_path(tmp_path: Path, monkeypatch):
    monkeypatch.setattr(generate, "DEFAULT_JOBS_DIR", str(tmp_path))
    out = generate.write_job_file("spark = SparkSession.builder.getOrCreate()\n")
    assert os.path.isfile(out["script_path"])
    with open(out["script_path"]) as f:
        assert f.read() == "spark = SparkSession.builder.getOrCreate()\n"

• Step 2: Run test to verify it fails

uv run pytest tests/unit/test_write_job_tool.py -v

Expected: ModuleNotFoundError: No module named 'spark_executor.tools.generate'.

• Step 3: Implement generate.py

# coding=utf-8
"""
@Time   :2026/6/24
@Author :tao.chen
"""
from common.logging import logger
from spark_executor.core.job_writer import write_job_file

DEFAULT_JOBS_DIR = "/tmp/jobs"


def write_job_file(code: str) -> dict[str, str]:
    path = write_job_file(code, jobs_dir=DEFAULT_JOBS_DIR)
    logger.info(f"write_job_file script_path={path}")
    return {"script_path": path}

• Step 4: Run test to verify it passes

uv run pytest tests/unit/test_write_job_tool.py -v

Expected: 1 passed.

• Step 5: Register the route in server.py

In spark_executor/server.py, add import and route:

from spark_executor.tools.write_job import write_job_file

@app.post("/write_job_file")
def _write_job_file(code: str):
    return write_job_file(code)

• Step 6: Update integration test to assert the new route

In tests/integration/test_mcp_routes.py, add to the loop in test_eight_tool_routes_registered:

        "/write_job_file",

• Step 7: Run full suite

uv run pytest -v

Expected: all tests pass (77 total).

• Step 8: Commit

git add spark_executor/tools/write_job.py spark_executor/server.py tests/unit/test_write_job_tool.py tests/integration/test_mcp_routes.py
git commit -m "feat: add write_job_file MCP tool"

---

Stage 2 checkpoint

You can now demonstrate the full agent loop:

save_connection(name=..., master=..., ...)         ->  {"name":..., "status":"SAVED"}
list_connections()                                  ->  [Connection, ...]
write_job_file(code=...)                         ->  script_path
prepare_submit_job(connection=..., script_path=...) ->  pending_id  (no spark-submit)
<user reviews pending entry; agent may call get_pending_job / list_pending_jobs>
confirm_submit_job(pending_id=...)                  ->  job_id, application_id
get_job_status(job_id=...)                          ->  RUNNING | SUCCEEDED | FAILED | KILLED
get_job_logs(job_id=...)                            ->  tail of driver logs
kill_job(job_id=...)                                ->  KILLED

The submission step is now a two-phase commit: prepare_submit_job snapshots the connection's master / deploy_mode / spark_conf into a PendingSubmission record, and only confirm_submit_job actually invokes spark-submit. Edits to the named Connection between prepare and confirm are intentionally ignored.

---

Stage 3 — Async, registry, log cache, multi-tenant (★★★☆☆)

Note: Stage 3 is a non-trivial refactor. It is broken out here at a planning level only — do not implement it as part of this plan. A separate plan should be authored for Stage 3 once Stage 2 is shipping. This section exists so the Stage-1 / Stage-2 code can be structured to make Stage 3 cheap.

Why Stage 3 needs its own plan

Stage 3 changes the runtime model:

1. submit_job becomes non-blocking: it returns job_id immediately, and spark-submit runs as an asyncio.create_subprocess_exec task in the background.
2. The in-memory JobStore (Task 5) is replaced with a SQLite-backed JobRegistry (spark_executor/core/job_registry.py). Schema: job_id TEXT PK, application_id TEXT, status TEXT, submit_time TEXT, owner TEXT, connection TEXT, script_path TEXT, queue TEXT, params JSON. The connection table mirrors ConnectionStore records.
3. A status poller (spark_executor/core/status_poller.py) runs as a FastAPI lifespan background task, periodically refreshing status for every non-terminal job via get_application_status.
4. get_job_logs no longer calls yarn logs on every request. A LogCache (spark_executor/core/log_cache.py) persists yarn logs output to a per-application file under .log_cache/<application_id>.log and accepts an offset parameter; subsequent calls only fetch logs after that offset.
5. An owner field is added to Job and threaded through every tool, with a filter so a tenant can only see their own jobs. Connection records may also be scoped per-owner at this point.

Implications for Stage 1 / Stage 2 code (so Stage 3 stays cheap)

These are constraints the current plan already respects:

• Tool functions take keyword-only arguments and return a Pydantic model. Stage 3 will add an owner kwarg without changing call sites.
• JobStore is a class with put / get / list, not module-level globals. Stage 3 swaps the implementation; the surface stays the same.
• ConnectionStore is a class with list_all / get / save / delete keyed by name, with a JSON file under SPARK_EXECUTOR_DATA_DIR (default ./data/). Stage 3 may migrate it into the same SQLite DB; the surface is unchanged.
• PendingStore is a class with list_all / get / save / delete keyed by pending_id, with a JSON file in the same data_dir. The two-step prepare → confirm flow is intentionally synchronous and CPU-only; Stage 3 does not need to change it to introduce async submission, because confirm_submit_job already separates "what to submit" (the persisted snapshot) from "where to submit" (the YARN call). The async refactor in Stage 3 will move the YARN call inside confirm_submit_job to a background task; the prepare/confirm shape and the PendingStore interface stay the same.
• All yarn/spark-submit calls go through core/ modules that return data, never raise HTTP errors directly. Stage 3 wraps them in async equivalents.
• Pydantic Job includes submit_time, script_path, queue, and connection already — Stage 3 only needs to add owner and status. SQLite migration is one CREATE TABLE away.
• PendingSubmission carries a snapshot of master / deploy_mode / spark_conf at prepare time, not a reference to a live Connection. This means Stage 3's multi-tenant scoping can be enforced at prepare time (rejecting unauthorised connections) without needing a connection lookup at confirm time.

Stage 3 task outline (deferred)

When the user is ready, a new plan should be written with at minimum:

1. Async confirm_submit_job: asyncio.create_subprocess_exec wrapper, persist Job with status="SUBMITTING" immediately and update to SUBMITTED/FAILED when the background task completes. prepare_submit_job is unchanged.
2. SQLite JobRegistry with init_db() called from main.py lifespan.
3. Status poller task started in lifespan, cancelled on shutdown.
4. LogCache writing per-application logs to disk; get_job_logs(job_id, offset=0).
5. owner parameter plumbed through all twelve tools (5 pending + 3 job + 4 connection); JobStore.list(owner=...) and PendingStore.list_all(owner=...) filters.
6. get_job_status returns cached status from the registry (refreshed by poller), not a live YARN call.

Target size: 1000–2000 lines of Python, including tests.

---

Self-Review Notes (author checklist)

• Spec coverage:
  • Stage 1 covers the job lifecycle tools (get_job_status / get_job_logs / kill_job), application_id extraction (both forms), 5000-char log tail, and YARN CLI invocation. ✔
  • Stage 1 also covers the connection-management requirement: 4 tools (save/list/get/delete) backed by ConnectionStore (JSON at ./data/connections.json, overridable via SPARK_EXECUTOR_DATA_DIR). ✔
  • Stage 1 also covers the two-step pending submission requirement (prepare_submit_job / confirm_submit_job / list_pending_jobs / get_pending_job / cancel_pending_job): parameters are snapshotted into a PendingSubmission and persisted to ./data/pending_jobs.json; spark-submit is only invoked by confirm_submit_job, which is the user-second-confirmation step. ✔
  • Stage 2 covers write_job_file writing to /tmp/jobs/job_xxx.py. ✔
  • Stage 3 is documented as deferred; the Stage-1/Stage-2 code is shaped to make it cheap. ✔
• Placeholder scan: No "TBD" or "implement later" markers. All steps contain code or concrete commands. Stage 3 is explicitly out of scope and named as such.
• Type consistency: Job, JobStatus, SubmitResult, Connection, PendingSubmission Pydantic models defined in Task 1. JobStore.put/get/list is consistent across the three job-lifecycle tools (Tasks 17–19). ConnectionStore.list_all/get/save/delete is consistent across connections.py (Tasks 7–10) and submit.py reads it in Task 12. PendingStore.list_all/get/save/delete is consistent across all five pending tools (Tasks 12–16); confirm_submit_job writes back through it with status transitions. Job.connection is set by confirm_submit_job (Task 13) and consumed by get_job_status / get_job_logs / kill_job (Tasks 17–19).