107 lines
4.6 KiB
Python
107 lines
4.6 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import shutil
|
||
import uuid
|
||
from pathlib import Path
|
||
|
||
|
||
class WorkerError(Exception):
|
||
def __init__(self, message: str, code: str = "worker_error"):
|
||
super().__init__(message)
|
||
self.code = code
|
||
|
||
|
||
class BrowserWorker:
|
||
"""Serialized RPC to the single browser owner; never imports the old application."""
|
||
|
||
def __init__(self, root: Path, home: Path):
|
||
self.root, self.home = root, home
|
||
self.process: asyncio.subprocess.Process | None = None
|
||
self.lock = asyncio.Lock()
|
||
self.stderr_task: asyncio.Task | None = None
|
||
self.closing = False
|
||
self.close_task: asyncio.Task | None = None
|
||
|
||
async def _start(self) -> None:
|
||
if self.process and self.process.returncode is None:
|
||
return
|
||
node = os.environ.get("ASSISTANT_NODE") or shutil.which("node")
|
||
if not node:
|
||
raise WorkerError("找不到 Node.js，请通过项目启动脚本运行后台", "node_not_found")
|
||
script = self.root / "tools" / "learning" / "browser-worker.mjs"
|
||
if not script.is_file():
|
||
raise WorkerError("本机学习执行器尚未就绪", "worker_unavailable")
|
||
environment = {**os.environ, "ASSISTANT_LEARNING_HOME": str(self.home)}
|
||
self.process = await asyncio.create_subprocess_exec(node, str(script), cwd=self.root, env=environment, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, limit=8 * 1024 * 1024)
|
||
if self.closing:
|
||
await self._stop()
|
||
raise WorkerError("学习执行器正在正常退出", "worker_closed")
|
||
process = self.process
|
||
|
||
async def consume_stderr() -> None:
|
||
assert process.stderr
|
||
while await process.stderr.readline():
|
||
pass # Do not forward browser diagnostic text or URLs into business logs.
|
||
|
||
self.stderr_task = asyncio.create_task(consume_stderr())
|
||
|
||
async def call(self, method: str, params: dict) -> dict:
|
||
async with self.lock:
|
||
if self.closing:
|
||
raise WorkerError("学习执行器正在正常退出", "worker_closed")
|
||
await self._start()
|
||
process = self.process
|
||
assert process and process.stdin and process.stdout
|
||
request_id = uuid.uuid4().hex
|
||
try:
|
||
process.stdin.write((json.dumps({"id": request_id, "method": method, "params": params}, ensure_ascii=False) + "\n").encode())
|
||
await process.stdin.drain()
|
||
async with asyncio.timeout(120):
|
||
while True:
|
||
line = await process.stdout.readline()
|
||
if not line:
|
||
raise WorkerError("本机学习浏览器已退出，请重新打开会话", "worker_closed")
|
||
try:
|
||
response = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
if response.get("id") != request_id:
|
||
continue
|
||
if response.get("error"):
|
||
error = response["error"]
|
||
raise WorkerError(str(error.get("message", "学习执行失败")), str(error.get("code", "worker_error")))
|
||
return response.get("result") or {}
|
||
except (TimeoutError, BrokenPipeError, ConnectionResetError):
|
||
await self._stop()
|
||
raise WorkerError("学习执行器响应中断；请核对页面与步骤后重新打开，不会自动重放动作", "worker_timeout") from None
|
||
|
||
async def _stop(self) -> None:
|
||
process, self.process = self.process, None
|
||
if not process:
|
||
return
|
||
if process.stdin:
|
||
process.stdin.close()
|
||
try:
|
||
await asyncio.wait_for(process.wait(), 15)
|
||
except TimeoutError:
|
||
process.terminate()
|
||
try:
|
||
await asyncio.wait_for(process.wait(), 3)
|
||
except TimeoutError:
|
||
process.kill()
|
||
await process.wait()
|
||
if self.stderr_task:
|
||
self.stderr_task.cancel()
|
||
self.stderr_task = None
|
||
|
||
async def close(self) -> None:
|
||
# EOF must reach the child even while an RPC holds the request lock.
|
||
# Node closes its contexts on EOF, unblocking that exact active action.
|
||
self.closing = True
|
||
if self.close_task is None:
|
||
self.close_task = asyncio.create_task(self._stop())
|
||
await asyncio.shield(self.close_task)
|