Auto-Finance/src/infra/browser/actor.py

"""浏览器会话的常驻持有者与独占锁

两件事必须同时成立，填报才可信：

1. **线程亲和**：Playwright 同步 API 底层是 greenlet，browser/page 对象必须在创建它的那个线程里使用。
   MCP 服务器是多线程的（每个工具调用可能在不同线程），所以所有浏览器操作统一提交到这个单线程 actor。
   与 COM 的 actor 分开：COM 套间要求不同，混用会踩坑。
2. **进程独占**：宿主可能同时存在多代 MCP 子进程（重连窗口、subagent/teammate 各一条 transport，
   实测 4 个 PID 并存）。每个进程各开一个 Chromium，Agent 看到的截图就不是它正在操作的那个页面。
"""

from __future__ import annotations

import atexit
import json
import os
import threading
import time
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any

from ... import get_logger

log = get_logger("browser.actor")

LOCK_FILENAME = ".browser.lock"
# 心跳超过这个秒数没更新就视为持有者已死，可被接管
STALE_AFTER_SECONDS = 180

_executor: ThreadPoolExecutor | None = None
_lock = threading.Lock()
# case_id -> {"bot": BaseBot, "steps": [str, ...], "kind": str|None}
_sessions: dict[str, dict[str, Any]] = {}


def _pool() -> ThreadPoolExecutor:
    global _executor
    if _executor is None:
        _executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="playwright")
    return _executor


def run_browser[T](fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
    """在专用 Playwright 线程上串行执行 fn 并等待结果。"""
    with _lock:
        return _pool().submit(fn, *args, **kwargs).result()


# ------------------------------------------------------------------
# 会话注册表（只在 actor 线程里读写，不加锁也不竞争）
# ------------------------------------------------------------------


def session(case_id: str) -> dict[str, Any] | None:
    """取会话的浏览器上下文；未打开返回 None。"""
    return _sessions.get(case_id)


def put_session(case_id: str, bot: Any, kind: str | None = None) -> dict[str, Any]:
    ctx: dict[str, Any] = {"bot": bot, "steps": [], "kind": kind, "opened_at": time.time()}
    _sessions[case_id] = ctx
    return ctx


def drop_session(case_id: str) -> dict[str, Any] | None:
    return _sessions.pop(case_id, None)


def list_sessions() -> list[str]:
    return sorted(_sessions)


def mark_step(case_id: str, step: str) -> None:
    ctx = _sessions.get(case_id)
    if ctx is not None and step not in ctx["steps"]:
        ctx["steps"].append(step)


# ------------------------------------------------------------------
# 磁盘锁
# ------------------------------------------------------------------


def lock_path(case_dir: Path) -> Path:
    return case_dir / LOCK_FILENAME


def read_lock(case_dir: Path) -> dict[str, Any] | None:
    """读锁文件；损坏或不存在返回 None（损坏不阻断，按无锁处理）。"""
    path = lock_path(case_dir)
    if not path.exists():
        return None
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception as exc:
        log.warning("浏览器锁文件损坏，按无锁处理: %s (%s)", path, exc)
        return None
    return data if isinstance(data, dict) else None


def _write_lock(case_dir: Path, url: str = "") -> dict[str, Any]:
    payload = {
        "pid": os.getpid(),
        "started_at": round(time.time(), 3),
        "heartbeat": round(time.time(), 3),
        "url": url,
    }
    lock_path(case_dir).write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    return payload


def acquire(case_dir: Path, url: str = "") -> tuple[bool, dict[str, Any]]:
    """尝试取得该会话的浏览器所有权。返回 (是否取得, 锁信息或占用者信息)。"""
    held = read_lock(case_dir)
    if held is not None:
        other_pid = int(held.get("pid", -1))
        age = time.time() - float(held.get("heartbeat", 0))
        if other_pid != os.getpid() and age < STALE_AFTER_SECONDS:
            return False, {**held, "age_seconds": round(age, 1), "stale": False}
        if other_pid != os.getpid():
            log.warning("接管过期浏览器锁（原持有 pid=%s，%.0fs 无心跳）", other_pid, age)
    return True, _write_lock(case_dir, url)


def touch(case_dir: Path, url: str = "") -> None:
    """刷心跳：每次浏览器工具调用都调一下。"""
    if read_lock(case_dir) is None:
        return
    _write_lock(case_dir, url)


def release(case_dir: Path) -> None:
    held = read_lock(case_dir)
    if held is not None and int(held.get("pid", -1)) == os.getpid():
        try:
            lock_path(case_dir).unlink()
        except OSError:
            pass


# ------------------------------------------------------------------
# 收尾
# ------------------------------------------------------------------


def close_all() -> None:
    """尽力关闭所有浏览器并释放锁（进程退出时调用）。"""
    for case_id in list(_sessions):
        ctx = _sessions.pop(case_id, None)
        if not ctx:
            continue
        bot = ctx.get("bot")
        work_dir = getattr(bot, "work_dir", None)
        try:
            if bot is not None:
                bot.close()
        except Exception as exc:
            log.warning("关闭浏览器失败 %s: %s", case_id, exc)
        if isinstance(work_dir, Path):
            release(work_dir)


atexit.register(close_all)