ClaudeandClaude
deafb5a26b
- fetch_url: remove the 1MB body cap and the streaming helper. The manual redirect loop with allowlist re-check (the SSRF fix) is kept intact; the body is now read in full via resp.text. Drop the now- meaningless `truncated` field from FetchUrlResult and the tests that asserted on it. Switched from httpx.stream() back to httpx.get() for the redirect loop — cleaner without the body cap. - datetime: replace deprecated datetime.utcnow() with datetime.now(timezone.utc) in submit.py (3 sites) and core/job_writer.py (1 site). Update the stale comment in core/pending_store.py that referenced the old call. - Clean up an unused `from common.config import settings` import in submit.py that ruff flagged. Co-Authored-By: Claude <noreply@anthropic.com>
77 lines
2.8 KiB
Python
77 lines
2.8 KiB
Python
# coding=utf-8
|
|
"""
|
|
@Time :2026/6/24
|
|
@Author :tao.chen
|
|
|
|
Writes LLM-generated PySpark code to disk so the agent can then call
|
|
prepare_submit_job(script_path=...) on the returned path.
|
|
|
|
Resolution order for the output directory:
|
|
1. explicit jobs_dir argument (used in tests)
|
|
2. SPARK_EXECUTOR_JOBS_DIR environment variable (operator override)
|
|
3. DEFAULT_JOBS_DIR constant (./data/jobs/ — gitignored, persists across
|
|
container restarts when ./data is mounted)
|
|
"""
|
|
import os
|
|
import secrets
|
|
from datetime import datetime, timezone
|
|
|
|
from common.config import settings
|
|
from common.logging import logger
|
|
|
|
ENV_JOBS_DIR = "SPARK_EXECUTOR_JOBS_DIR"
|
|
|
|
|
|
def resolve_jobs_dir(jobs_dir: str | None = None) -> str:
|
|
"""Pick the effective jobs directory in priority order: arg > settings.jobs_dir.
|
|
|
|
`settings.jobs_dir` itself is computed in `common/config.py` as:
|
|
SPARK_EXECUTOR_JOBS_DIR env var (if set) > <data_dir>/jobs
|
|
"""
|
|
if jobs_dir is not None:
|
|
return jobs_dir
|
|
return settings.jobs_dir
|
|
|
|
|
|
def _not_utf8_error(exc: UnicodeEncodeError) -> ValueError:
|
|
"""Standard 400 error when the LLM-supplied code string cannot be
|
|
encoded as UTF-8 (typically because it contains a lone surrogate
|
|
from a broken decode earlier in the pipeline)."""
|
|
return ValueError(
|
|
f"code is not valid UTF-8 and cannot be written as a .py file. "
|
|
f"This usually means the LLM produced a string with a lone "
|
|
f"surrogate codepoint (U+D800..U+DFFF) — re-generate the code "
|
|
f"as plain Unicode. Underlying error: {exc}"
|
|
)
|
|
|
|
|
|
def write_job_file(code: str, jobs_dir: str | None = None) -> str:
|
|
"""Write `code` to <jobs_dir>/job_<timestamp>_<rand>.py; return abs path.
|
|
|
|
Creates the directory if missing. Filename is unique per call (timestamp
|
|
down to the second + 3 random bytes) so concurrent agents cannot collide.
|
|
|
|
Raises ValueError (-> 400) if `code` cannot be encoded as UTF-8 (e.g.
|
|
contains a lone surrogate). The check is cheap and runs before any
|
|
file I/O, so a bad payload never produces a half-written .py.
|
|
"""
|
|
try:
|
|
code.encode("utf-8")
|
|
except UnicodeEncodeError as exc:
|
|
logger.warning(f"write_job_file rejected: code is not valid UTF-8: {exc}")
|
|
raise _not_utf8_error(exc) from exc
|
|
|
|
effective_dir = resolve_jobs_dir(jobs_dir)
|
|
os.makedirs(effective_dir, exist_ok=True)
|
|
stamp = datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S")
|
|
name = f"job_{stamp}_{secrets.token_hex(3)}.py"
|
|
path = os.path.join(effective_dir, name)
|
|
abs_path = os.path.abspath(path)
|
|
logger.debug(
|
|
f"write_job_file enter jobs_dir={effective_dir} code_bytes={len(code)}"
|
|
)
|
|
with open(abs_path, "w", encoding="utf-8") as f:
|
|
f.write(code)
|
|
logger.info(f"write_job_file ok script_path={abs_path} bytes={len(code)}")
|
|
return abs_path
|