126 lines
4.6 KiB
Python
126 lines
4.6 KiB
Python
"""Standalone Playwright screenshot worker (fresh process, Windows-safe).
|
||
|
||
Invoked by fidelity_loop.capture_real_app_screenshot — never import uvicorn here.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import base64
|
||
import json
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
|
||
def _force_proactor() -> None:
|
||
if sys.platform == "win32":
|
||
asyncio.set_event_loop_policy(asyncio.WindowsProactorEventLoopPolicy())
|
||
|
||
|
||
async def _shot_async(payload: dict) -> tuple[bytes | None, str, str | None]:
|
||
from playwright.async_api import async_playwright
|
||
|
||
draft = payload.get("draft") or {}
|
||
rows = payload.get("rows") or []
|
||
web = (os.getenv("WEB_BASE") or "http://127.0.0.1:5173").rstrip("/")
|
||
goto_timeout = int(os.getenv("FIDELITY_GOTO_TIMEOUT_MS") or "45000")
|
||
slug = str(((draft.get("meta") or {}).get("slug")) or "preview")
|
||
preview_id = str(payload.get("preview_id") or "preview")
|
||
inject_payload = {
|
||
"ok": True,
|
||
"id": preview_id,
|
||
"blueprint": draft,
|
||
"rows": rows,
|
||
"resource": payload.get("resource") or "records",
|
||
}
|
||
inject_js = (
|
||
"try{localStorage.removeItem('ajz_session');}catch(e){}"
|
||
"window.__AJZ_PREVIEW__ = "
|
||
+ json.dumps(inject_payload, ensure_ascii=False)
|
||
+ ";"
|
||
)
|
||
target = f"{web}/#/preview/{preview_id}"
|
||
|
||
def fmt(err: BaseException) -> str:
|
||
msg = str(err).strip() or repr(err)
|
||
return msg.replace("\n", " ")[:500]
|
||
|
||
try:
|
||
async with async_playwright() as p:
|
||
browser = await p.chromium.launch(headless=True)
|
||
try:
|
||
page = await browser.new_page(viewport={"width": 1600, "height": 1100})
|
||
await page.add_init_script(inject_js)
|
||
|
||
async def fulfill(route):
|
||
await route.fulfill(
|
||
status=200,
|
||
content_type="application/json; charset=utf-8",
|
||
body=json.dumps(inject_payload, ensure_ascii=False).encode("utf-8"),
|
||
)
|
||
|
||
await page.route("**/api/v1/preview/**", fulfill)
|
||
await page.route("**/ai/api/v1/preview/**", fulfill)
|
||
await page.goto(target, wait_until="domcontentloaded", timeout=goto_timeout)
|
||
await page.evaluate(inject_js)
|
||
try:
|
||
await page.wait_for_selector(
|
||
".sf-shell-bar, .sf-root, .gen-app-ops, .gen-app, .sf-line-wrap",
|
||
timeout=45000,
|
||
)
|
||
except Exception:
|
||
try:
|
||
snip = (await page.locator("body").inner_text(timeout=2000))[:240]
|
||
except Exception:
|
||
snip = page.url
|
||
raise RuntimeError(
|
||
f"preview shell missing. page={page.url} text={snip!r}"
|
||
) from None
|
||
login_btns = page.locator("button").filter(has_text="登录")
|
||
shell = page.locator(".sf-shell-bar, .sf-root, .gen-app")
|
||
if await login_btns.count() and not await shell.count():
|
||
raise RuntimeError("preview fell back to login")
|
||
await page.wait_for_timeout(2500)
|
||
if await page.locator(".gen-app").count():
|
||
png = await page.locator(".gen-app").first.screenshot(type="png")
|
||
else:
|
||
png = await page.screenshot(type="png", full_page=True)
|
||
finally:
|
||
await browser.close()
|
||
return (
|
||
png,
|
||
f"已截取草稿预览（/#/preview/{preview_id}，注入数据无需再拉 AI）参与对比",
|
||
slug,
|
||
)
|
||
except Exception as e: # noqa: BLE001
|
||
return None, f"真页面截图失败: preview={fmt(e)}（目标 {target}）", slug
|
||
|
||
|
||
def main() -> int:
|
||
_force_proactor()
|
||
inp = Path(os.environ["FIDELITY_SHOT_INPUT"])
|
||
outp = Path(os.environ["FIDELITY_SHOT_OUTPUT"])
|
||
payload = json.loads(inp.read_text(encoding="utf-8"))
|
||
png, note, slug = asyncio.run(_shot_async(payload))
|
||
out_dir = os.getenv("FIDELITY_SHOT_DIR") or ""
|
||
if png and out_dir:
|
||
d = Path(out_dir)
|
||
d.mkdir(parents=True, exist_ok=True)
|
||
(d / f"preview_{payload.get('preview_id') or 'x'}.png").write_bytes(png)
|
||
outp.write_text(
|
||
json.dumps(
|
||
{
|
||
"png_b64": base64.b64encode(png).decode("ascii") if png else None,
|
||
"note": note,
|
||
"slug": slug,
|
||
},
|
||
ensure_ascii=False,
|
||
),
|
||
encoding="utf-8",
|
||
)
|
||
return 0 if png else 2
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|