fix: drop file tools 1 MB cap + add YarnError -> 502 handler

Two follow-ups to the previous error-handling fix (16fe011):

1) Drop the 1 MB cap on read_job_file and update_job_file.

   The cap was originally there to keep MCP responses bounded, but it
   blocks legitimate use of large PySpark scripts and large update
   payloads. With the new model (LLM composes scripts via
   write_job_file and edits them via read+update), a fixed 1 MB
   cap is more hindrance than protection. MCP response size is
   already bounded by the JSON transport and httpx; the tool itself
   doesn't need a second limit.

   Changes:
   - tools/job_file.py: delete MAX_FILE_BYTES constant, drop the two
     size checks in read_job_file and update_job_file, update
     module docstring.
   - tests/unit/test_job_file.py: delete test_update_rejects_
     oversized_content and test_update_rejects_1mb_plus_1_byte
     (the two tests that asserted the cap), replace with
     test_update_accepts_content_larger_than_former_1mb_cap.
   - server.py: drop "Caps reads at 1 MB" and "Caps writes at 1 MB"
     from the two route descriptions.

2) Add YarnError -> 502 handler.

   yarn_client wraps every httpx call: on connect / TLS / timeout /
   4xx / 5xx / parse failure it raises YarnError. Previously this
   was unhandled, so all six external job tools (get_external_job_*,
   list_applications, plus anything else that hits YARN) returned
   500 "Internal Server Error" with no detail — the LLM couldn't
   tell whether the cluster was down or the request was bad.

   Same fix as 16fe011 (which did this for fetch_url directly).
   The new handler returns HTTP 502 Bad Gateway with the YarnError
   message in the response detail. 502 because the MCP service is
   acting as a gateway to YARN — 502 is the standard status for
   "upstream didn't respond correctly".

   Changes:
   - server.py: import YarnError, add @app.exception_handler
     returning 502 + the YarnError message.
   - tests/integration/test_mcp_routes.py: new test asserts that
     when get_application_status raises YarnError, /get_job_status
     returns 502 with the YarnError message in the detail.

   Note: ValueError (request was bad) is still 400, KeyError (job
   not in JobStore) is still 404. The three handlers form a clean
   3-way classification of tool-layer errors.

Tests: 401 passed (was 401, +1 YarnError test, -2 cap tests = net -1).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Claude
2026-07-09 18:06:50 +08:00
co-authored by Claude Fable 5
parent a5e6d1ffe1
commit 3b698e1bdc
4 changed files with 71 additions and 41 deletions
+21 -7
View File
@@ -6,6 +6,7 @@
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from spark_executor.core.yarn_client import YarnError
from spark_executor.tools.connections import (
delete_connection,
get_connection,
@@ -83,6 +84,20 @@ async def _valueerror_handler(_request: Request, exc: ValueError) -> JSONRespons
return JSONResponse(status_code=400, content={"detail": str(exc)})
@app.exception_handler(YarnError)
async def _yarnerror_handler(_request: Request, exc: YarnError) -> JSONResponse:
"""YARN RM unreachable or returned an error.
Surfaces as HTTP 502 Bad Gateway (we are a gateway to YARN). The
detail includes whatever yarn_client put in the YarnError message
(host:port unreachable, HTTP status from YARN, parse failure,
etc). The LLM can act on this to distinguish "YARN is down" from
"the request was bad" (the latter would be a 400 from
_valueerror_handler instead).
"""
return JSONResponse(status_code=502, content={"detail": str(exc)})
@app.get("/health")
def health_check():
return {"status": "ok"}
@@ -513,10 +528,9 @@ def _write_job_file(req: WriteJobFileRequest):
summary="Read the contents of an existing PySpark script",
description=(
"Returns the text content of an existing script file at the given "
"path. Caps reads at 1 MB. Typical use: after write_job_file "
"returns a path, call read_job_file on that path to inspect what "
"was actually written, before deciding to prepare_submit_job or "
"update_job_file."
"path. No size cap. Typical use: after write_job_file returns a "
"path, call read_job_file on that path to inspect what was actually "
"written, before deciding to prepare_submit_job or update_job_file."
),
)
def _read_job_file(req: ReadJobFileRequest):
@@ -531,9 +545,9 @@ def _read_job_file(req: ReadJobFileRequest):
"Replaces the entire content of an existing script file. Path must "
"be under SPARK_EXECUTOR_JOBS_DIR (the dir write_job_file writes "
"to) — protects against overwriting host-mounted configs or other "
"non-script files. Caps writes at 1 MB. Typical use: read_job_file, "
"edit the content (LLM or human), update_job_file, then "
"prepare_submit_job with the same path."
"non-script files. No size cap. Typical use: read_job_file, edit "
"the content (LLM or human), update_job_file, then prepare_submit_job "
"with the same path."
),
)
def _update_job_file(req: UpdateJobFileRequest):
+8 -19
View File
@@ -16,7 +16,10 @@ Safety:
- update_job_file: must be under SPARK_EXECUTOR_JOBS_DIR
(settings.jobs_dir) so the agent cannot overwrite host-mounted
configs or arbitrary files on the container FS.
- 1 MB cap on both read and write payloads to keep MCP responses bounded.
- No size cap on read or write payloads — the caller controls the
file size. MCP response size is bounded by httpx and the JSON
transport; for very large files the LLM should split via
write_job_file and read in chunks via read_job_file + update.
"""
import os
from pathlib import Path
@@ -24,9 +27,6 @@ from pathlib import Path
from common.config import settings
from common.logging import logger
MAX_FILE_BYTES = 1 * 1024 * 1024 # 1 MB
class ScriptFileError(ValueError):
"""Raised when read/update fails. -> HTTP 400 via the FastAPI ValueError
handler in server.py.
@@ -68,16 +68,10 @@ def _check_writable(script_path: str) -> None:
def read_job_file(script_path: str) -> dict[str, object]:
"""Return the text content of an existing script file.
Caps the read at 1 MB to keep MCP responses bounded; raises
ScriptFileError (-> 400) if the file is missing or too large.
No size cap. Raises ScriptFileError (-> 400) if the file is missing.
"""
_check_readable(script_path)
size = os.path.getsize(script_path)
if size > MAX_FILE_BYTES:
raise ScriptFileError(
f"Script is too large to read back ({size} bytes > {MAX_FILE_BYTES} "
f"byte cap). Edit it via a host volume mount instead."
)
logger.debug(f"read_job_file enter script_path={script_path} size={size}")
with open(script_path, encoding="utf-8") as f:
content = f.read()
@@ -88,17 +82,12 @@ def read_job_file(script_path: str) -> dict[str, object]:
def update_job_file(script_path: str, content: str) -> dict[str, object]:
"""Overwrite an existing script file with new content.
Restricted to paths under settings.jobs_dir. Caps writes at 1 MB.
Raises ScriptFileError (-> 400) if the path is missing, outside
the allowed dir, or the content is too large.
Restricted to paths under settings.jobs_dir. No size cap.
Raises ScriptFileError (-> 400) if the path is missing or outside
the allowed dir.
"""
_check_writable(script_path)
encoded_size = len(content.encode("utf-8"))
if encoded_size > MAX_FILE_BYTES:
raise ScriptFileError(
f"content is too large ({encoded_size} bytes > {MAX_FILE_BYTES} "
f"byte cap). Split the script into multiple files."
)
logger.debug(
f"update_job_file enter script_path={script_path} "
f"new_bytes={encoded_size}"
+35
View File
@@ -8,6 +8,7 @@ from fastapi.testclient import TestClient
from spark_executor.core import connection_store, pending_store
from spark_executor.core.connection_store import ConnectionStore
from spark_executor.core.pending_store import PendingStore
from spark_executor.models import Connection
from spark_executor.server import app
from spark_executor.tools import connections, submit
@@ -544,3 +545,37 @@ def test_update_pending_job_route_rejects_submitted(tmp_path, monkeypatch):
)
assert r.status_code == 400
assert "only PENDING submissions can be updated" in r.json()["detail"]
# --- YarnError -> 502 handler ---
def test_yarn_error_returns_502_with_detail(tmp_path, monkeypatch):
"""When a YARN REST call raises YarnError (host unreachable, YARN
returned 4xx/5xx, parse failure), the route must surface HTTP 502
with the YarnError message in the detail — not 500 'Internal
Server Error' with no info. Mirrors the same fix for fetch_url.
"""
from unittest.mock import patch
from spark_executor.core.yarn_client import YarnError
connection_store.store.save(
Connection(
name="prod",
master="yarn",
yarn_rm_url="http://ccam1:8088",
)
)
c = TestClient(app)
with patch(
"spark_executor.tools.status.get_application_status",
side_effect=YarnError("YARN connection failed: ConnectError: Connection refused"),
):
r = c.post(
"/get_job_status",
json={"job_id": "a1b2c3d4e5f6"},
)
assert r.status_code == 502
detail = r.json()["detail"]
assert "YARN connection failed" in detail
assert "Connection refused" in detail
+7 -15
View File
@@ -118,25 +118,17 @@ def test_update_rejects_missing_file(tmp_path: Path):
update_job_file(str(tmp_path / "nope.py"), "x\n")
def test_update_rejects_oversized_content(tmp_path: Path, monkeypatch):
"""Cap at 1 MB so the MCP response stays bounded."""
def test_update_accepts_content_larger_than_former_1mb_cap(tmp_path: Path):
"""The 1 MB cap was removed (commit a5e6d1f and following). Content
larger than the old cap must now be accepted; size is bounded by
the MCP transport, not the tool."""
config.settings.jobs_dir = str(tmp_path)
p = tmp_path / "big.py"
p.write_text("# small\n")
# Synthesize 2 MB of content (don't actually write 2 MB to disk)
big = "x" * (2 * 1024 * 1024)
with pytest.raises(ScriptFileError, match="too large"):
update_job_file(str(p), big)
def test_update_rejects_1mb_plus_1_byte(tmp_path: Path):
"""Exactly at the boundary: 1 MB + 1 byte must be rejected."""
config.settings.jobs_dir = str(tmp_path)
p = tmp_path / "x.py"
p.write_text("# t\n")
just_over = "x" * (1024 * 1024 + 1)
with pytest.raises(ScriptFileError, match="too large"):
update_job_file(str(p), just_over)
out = update_job_file(str(p), big)
assert out["bytes_written"] == len(big)
assert p.read_text() == big
# --- round-trip: write -> read -> update -> read ---