151 lines
4.7 KiB
Python
151 lines
4.7 KiB
Python
"""脚本执行器：subprocess 沙箱版。
|
||
|
||
Phase 5 简化实现：在 worker 进程内直接 subprocess 执行脚本字符串。
|
||
- 支持 python / shell / node
|
||
- 超时自动 kill
|
||
- 每行 stdout/stderr 作为一条 TaskLog
|
||
|
||
（Celery worker 版本为未来优化项；当前流程在 FastAPI 后台任务里执行即可。）
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import shutil
|
||
import sys
|
||
import tempfile
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any, AsyncIterator
|
||
|
||
|
||
def _resolve_lang_cmd() -> dict[str, list[str]]:
|
||
"""启动时探测各语言解释器。Python 直接用当前 venv 的解释器，不依赖 PATH。"""
|
||
node = shutil.which("node") or "node"
|
||
bash = shutil.which("bash") or "/bin/bash"
|
||
return {
|
||
"python": [sys.executable, "-u"], # -u 保证 stdout 无缓冲
|
||
"node": [node],
|
||
"shell": [bash],
|
||
}
|
||
|
||
|
||
LANG_CMD: dict[str, list[str]] = _resolve_lang_cmd()
|
||
|
||
|
||
@dataclass
|
||
class ScriptChunk:
|
||
level: str # INFO / ERROR
|
||
message: str
|
||
|
||
|
||
@dataclass
|
||
class ScriptResult:
|
||
returncode: int
|
||
stdout: str
|
||
stderr: str
|
||
status: str # SUCCEEDED / FAILED / TIMEOUT
|
||
output: dict[str, Any]
|
||
|
||
|
||
async def run_script(
|
||
*,
|
||
language: str,
|
||
code: str,
|
||
params: dict[str, Any],
|
||
timeout_seconds: int = 60,
|
||
on_chunk: callable | None = None,
|
||
) -> ScriptResult:
|
||
if language not in LANG_CMD:
|
||
raise ValueError(f"不支持的脚本语言: {language}")
|
||
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
p = Path(tmpdir)
|
||
ext = {"python": ".py", "node": ".js", "shell": ".sh"}[language]
|
||
script_path = p / f"entry{ext}"
|
||
script_path.write_text(code, encoding="utf-8")
|
||
# 通过环境变量 LOOM_PARAMS 传入参数（JSON），脚本可 json.loads 读
|
||
env = {"LOOM_PARAMS": json.dumps(params, ensure_ascii=False, default=str)}
|
||
|
||
cmd = LANG_CMD[language] + [str(script_path)]
|
||
import os as _os
|
||
|
||
proc = await asyncio.create_subprocess_exec(
|
||
*cmd,
|
||
cwd=str(p),
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
env={
|
||
**env,
|
||
"PATH": _os.environ.get("PATH", "/usr/local/bin:/usr/bin:/bin"),
|
||
},
|
||
)
|
||
|
||
stdout_chunks: list[str] = []
|
||
stderr_chunks: list[str] = []
|
||
|
||
async def _pipe(stream, is_err: bool) -> None:
|
||
if stream is None:
|
||
return
|
||
async for raw in stream:
|
||
try:
|
||
line = raw.decode("utf-8", errors="replace").rstrip("\n")
|
||
except Exception: # noqa: BLE001
|
||
line = str(raw)
|
||
(stderr_chunks if is_err else stdout_chunks).append(line)
|
||
if on_chunk:
|
||
await on_chunk(ScriptChunk(level=("ERROR" if is_err else "INFO"), message=line))
|
||
|
||
try:
|
||
await asyncio.wait_for(
|
||
asyncio.gather(_pipe(proc.stdout, False), _pipe(proc.stderr, True)),
|
||
timeout=timeout_seconds,
|
||
)
|
||
returncode = await proc.wait()
|
||
except asyncio.TimeoutError:
|
||
proc.kill()
|
||
await proc.wait()
|
||
return ScriptResult(
|
||
returncode=-1,
|
||
stdout="\n".join(stdout_chunks),
|
||
stderr="\n".join(stderr_chunks) + "\n[timeout killed]",
|
||
status="TIMEOUT",
|
||
output={"timeout": True, "timeout_seconds": timeout_seconds},
|
||
)
|
||
|
||
status = "SUCCEEDED" if returncode == 0 else "FAILED"
|
||
return ScriptResult(
|
||
returncode=returncode,
|
||
stdout="\n".join(stdout_chunks),
|
||
stderr="\n".join(stderr_chunks),
|
||
status=status,
|
||
output={
|
||
"returncode": returncode,
|
||
"stdout_tail": "\n".join(stdout_chunks[-20:]),
|
||
},
|
||
)
|
||
|
||
|
||
async def stream_lines(
|
||
*, language: str, code: str, params: dict[str, Any], timeout_seconds: int = 60
|
||
) -> AsyncIterator[tuple[str, str]]:
|
||
"""供外部需要逐行消费时使用（当前 run_script 已支持 callback，这个保留扩展位）。"""
|
||
chunks: asyncio.Queue[tuple[str, str] | None] = asyncio.Queue()
|
||
|
||
async def _on_chunk(c: ScriptChunk) -> None:
|
||
await chunks.put((c.level, c.message))
|
||
|
||
task = asyncio.create_task(
|
||
run_script(
|
||
language=language, code=code, params=params,
|
||
timeout_seconds=timeout_seconds, on_chunk=_on_chunk,
|
||
)
|
||
)
|
||
task.add_done_callback(lambda _: chunks.put_nowait(None))
|
||
while True:
|
||
item = await chunks.get()
|
||
if item is None:
|
||
break
|
||
yield item
|
||
await task
|