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:

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

Then run:

```bash
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**

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

Each new `__init__.py` gets the standard header:

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

- [ ] **Step 3: Write the failing test for models**

`tests/unit/test_models.py`:

```python
# 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**

```bash
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`**

```python
# 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**

```bash
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`:

```python
# 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**

```bash
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`:

```python
# 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**

```bash
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`**

```python
# 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**

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

Expected: 4 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

```bash
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`**

```python
# 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**

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

Expected: 5 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

```bash
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`**

```python
# 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**

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

Expected: 5 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

```bash
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`**

```python
# 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**

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

Expected: 3 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

```bash
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`**

```python
# 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**

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

Expected: 8 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

```bash
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`**

```python
# 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**

```bash
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**

```bash
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:

```python
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**

```bash
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**

```bash
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`**

```python
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**

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

Expected: 2 passed.

- [ ] **Step 3: Commit**

```bash
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`**

```python
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**

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

Expected: 2 passed.

- [ ] **Step 3: Commit**

```bash
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`:

```python
# 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**

```bash
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.

```python
# 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**

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

Expected: 8 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

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

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

- [ ] **Step 3: Implement `submit.py`**

```python
# 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**

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

Expected: 4 passed.

- [ ] **Step 5: Commit**

```bash
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`**

```python
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**

```bash
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`**

```python
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**

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

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

- [ ] **Step 5: Commit**

```bash
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`**

```python
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**

```bash
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`**

```python
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**

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

Expected: 2 passed.

- [ ] **Step 5: Commit**

```bash
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`**

```python
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**

```bash
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`**

```python
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**

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

Expected: 2 passed.

- [ ] **Step 5: Commit**

```bash
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`**

```python
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**

```bash
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`**

```python
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**

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

Expected: 3 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

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

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

- [ ] **Step 3: Implement `status.py`**

```python
# 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**

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

Expected: 2 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

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

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

- [ ] **Step 3: Implement `logs.py`**

```python
# 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**

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

Expected: 3 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

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

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

- [ ] **Step 3: Implement `kill.py`**

```python
# 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**

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

Expected: 2 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

```bash
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:

```python
# 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**

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

Expected: 2 passed.

- [ ] **Step 6: Run the full suite to confirm no regressions**

```bash
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)**

```bash
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**

```bash
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`:

```python
# 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**

```bash
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`**

```python
# 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**

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

Expected: 2 passed.

- [ ] **Step 5: Commit**

```bash
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`:

```python
# 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**

```bash
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`**

```python
# 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**

```bash
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:

```python
from spark_executor.tools.write_job import write_job_file
```

```python
@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`:

```python
        "/write_job_file",
```

- [ ] **Step 7: Run full suite**

```bash
uv run pytest -v
```

Expected: all tests pass (77 total).

- [ ] **Step 8: Commit**

```bash
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).